Optimizing Learned Image Compression on Scalar and Entropy-Constraint Quantization
This work addresses a specific bottleneck in learned image compression for improving rate-distortion efficiency, representing an incremental advancement.
The paper tackles the problem of suboptimal performance in learned image compression due to inaccurate quantization modeling during training, proposing a finetuning step that retrains parts of the network on correctly quantized latents. This approach yields additional coding gains, with average bitrate savings of 1-2% on the Kodak test set and up to 2.2% on the TecNick test set.
The continuous improvements on image compression with variational autoencoders have lead to learned codecs competitive with conventional approaches in terms of rate-distortion efficiency. Nonetheless, taking the quantization into account during the training process remains a problem, since it produces zero derivatives almost everywhere and needs to be replaced with a differentiable approximation which allows end-to-end optimization. Though there are different methods for approximating the quantization, none of them model the quantization noise correctly and thus, result in suboptimal networks. Hence, we propose an additional finetuning training step: After conventional end-to-end training, parts of the network are retrained on quantized latents obtained at the inference stage. For entropy-constraint quantizers like Trellis-Coded Quantization, the impact of the quantizer is particularly difficult to approximate by rounding or adding noise as the quantized latents are interdependently chosen through a trellis search based on both the entropy model and a distortion measure. We show that retraining on correctly quantized data consistently yields additional coding gain for both uniform scalar and especially for entropy-constraint quantization, without increasing inference complexity. For the Kodak test set, we obtain average savings between 1% and 2%, and for the TecNick test set up to 2.2% in terms of Bjøntegaard-Delta bitrate.