Optimizing Learned Image Compression On Scalar and Entropy-Constraint Quantization

2024-10-27

International Conference on Information Photonics

AI Summary

To address the train-inference inconsistency caused by non-differentiable quantization in learned image compression, this paper proposes a two-stage training paradigm: first, end-to-end pretraining using standard differentiable quantization approximations; second, fine-tuning the decoder and entropy model under strict entropy-constrained quantization (e.g., Trellis-Coded Quantization), with supervision directly from the quantized latent variables. This work is the first to embed exact entropy-constrained quantization into the fine-tuning stage, eliminating modeling biases inherent in noise injection or rounding-based approximations and effectively bridging the gap between quantization distortion and rate-distortion optimization. Experiments on Kodak and TecNick demonstrate average BD-rate reductions of 1.0%–2.0%, up to 2.2%, with zero increase in inference complexity.

Technology Category

• Machine Learning: Learning on the Edge & Model Compression
• Computer Vision: Learning & Optimization for CV
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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{o}ntegaard-Delta bitrate.

Problem

• Optimizing quantization in learned image compression networks
• Addressing suboptimal performance due to quantization noise approximation
• Improving coding gain for entropy-constraint quantization methods

Innovation

• Fine-tuning training on quantized latents
• Retraining with entropy-constraint quantizers
• Improving coding gain without complexity increase