Source-linked AI summary
Conditional Probability Models for Deep Image Compression
Fabian Mentzer, Eirikur Agustsson, Michael Tschannen, Radu Timofte, Luc Van Gool
TL;DR
Learned image-compression autoencoders must optimize distortion against latent-representation entropy, while discretization makes bitrate optimization difficult. This paper jointly trains the autoencoder with a lightweight 3D-CNN context model that estimates entropy, producing a simple system with state-of-the-art MS-SSIM performance and competitive learned-compression results. Future work could investigate heavier context models to improve compression and enable lossy natural-image sampling.
Problem
Learning image compression requires navigating the distortion–entropy trade-off, but discretizing latent representations makes bitrate optimization non-differentiable.
Method
The paper jointly trains a convolutional autoencoder and lightweight 3D-CNN context model, using learned conditional probabilities to estimate latent entropy during optimization.
Results
The resulting system achieves state-of-the-art MS-SSIM performance, outperforms BPG and JPEG2000 across four testing sets, and is competitive with a leading learned-compression method.
Takeaways & Limitations
Context modeling can be incorporated directly into autoencoder training to yield a simple, effective image-compression system.
Takeaways & Limitations
The paper identifies heavier and more powerful context models as future work that could further improve compression and enable lossy natural-image sampling.
Abstract
from arXiv · showhide
Deep Neural Networks trained as image auto-encoders have recently emerged as a promising direction for advancing the state-of-the-art in image compression. The key challenge in learning such networks is twofold: To deal with quantization, and to control the trade-off between reconstruction error (distortion) and entropy (rate) of the latent image representation. In this paper, we focus on the latter challenge and propose a new technique to navigate the rate-distortion trade-off for an image compression auto-encoder. The main idea is to directly model the entropy of the latent representation by using a context model: A 3D-CNN which learns a conditional probability model of the latent distribution of the auto-encoder. During training, the auto-encoder makes use of the context model to estimate the entropy of its representation, and the context model is concurrently updated to learn the dependencies between the symbols in the latent representation. Our experiments show that this approach, when measured in MS-SSIM, yields a state-of-the-art image compression system based on a simple convolutional auto-encoder.
1. Introduction
Deep neural image-compression autoencoders face a rate-distortion optimization challenge because discretizing latent representations complicates bitrate optimization. The paper addresses this by jointly training a context model with the autoencoder to estimate latent entropy, achieving state-of-the-art MS-SSIM performance with a simple system.
- Image compression balances reconstruction distortion d against bitrate R by minimizing d + βR.The parameter β controls the trade-off between the two objectives.
- Optimizing the bitrate is difficult because discretizing the latent representation into symbols is non-differentiable.After discretization, entropy H is the information-theoretic measure of bitrate.
- The paper jointly trains a convolutional autoencoder and context model, using the model’s learned conditional distribution to estimate latent entropy during optimization.The context model learns dependencies among latent symbols while the autoencoder uses its estimate to navigate the rate-distortion trade-off.
- The proposed system combines a standard convolutional autoencoder with a 3D-CNN-based context model.The approach directly incorporates context modeling into entropy optimization rather than using it only after training.
- The system achieves state-of-the-art image-compression performance when evaluated with MS-SSIM.The authors characterize it as simple and straightforward to implement with standard deep-learning toolboxes.
2. Related work
Prior learned image-compression work uses autoencoders, recurrent networks, perceptual or pixelwise objectives, progressive strategies, and context models. PixelRNN and PixelCNN additionally model natural-image distributions for lossless compression and image generation.
- Deep image-compression systems commonly use autoencoders or recurrent neural networks.Training objectives include mean-squared error and perceptual metrics such as MS-SSIM.
- Recent approaches also explore progressive encoding and decoding, adversarial training, multi-scale image decompositions, and generalized divisive normalization.
- Context models have long been used for entropy estimation in engineered compression and are also used by recent DNN-based lossy compression methods.Their roles include adaptive arithmetic coding and entropy-term estimation.
- PixelRNN and PixelCNN model natural-image distributions in lossless settings and can support learned lossless compression and image generation.PixelRNN is recurrent while PixelCNN is convolutional.
3. Proposed method
The proposed system jointly trains a convolutional auto-encoder and a causal context model to estimate latent entropy and optimize the rate-distortion trade-off. It combines finite-center quantization, conditional probability modeling with a masked 3D-CNN, and optional spatial bit allocation through an importance map.
- The compression system comprises an encoder, quantizer, and decoder, mapping an image through a quantized latent representation to a reconstructed image.The encoder produces z, the quantizer maps coordinates to learned centers C, and the decoder reconstructs the image from the quantized representation.
- The auto-encoder minimizes distortion plus β times latent entropy, with β controlling the trade-off between reconstruction quality and bitrate.The entropy H is the cost of encoding the quantized latent representation, and large or small β favors lower or higher entropy, respectively.
- Soft quantization provides backward-pass gradients while nearest-neighbor assignments use learned finite centers for the quantized latent representation.The forward pass uses hard assignments, while differentiable soft quantization is used only for gradient computation, avoiding an annealing schedule.
- A context model factorizes the latent distribution into conditional probabilities over previous symbols and estimates each symbol distribution at every 3D location.The model outputs probabilities for the centers in C, with raster-scan causality requiring each prediction to depend only on preceding latent values.
- A masked 3D-CNN generalizes PixelCNN-style causal modeling to the multichannel latent feature volume while respecting dependencies between symbols.The masking is applied to convolutional filters in every layer, enabling a lightweight context model for latent volumes with many channels.
- The context model’s cross-entropy loss estimates coding cost, allowing the auto-encoder and context model to be trained concurrently for rate-distortion optimization.The context model learns latent dependencies, while its estimated entropy adjusts the encoder and decoder; an importance map can also allocate bits spatially and support multiple rates in one model.
4. Experiments
The experiments evaluate a convolutional auto-encoder with a lightweight 3D-CNN context model across multiple image datasets and compression settings. The method outperforms traditional codecs in MS-SSIM, while visual and ablation studies examine context modeling and importance maps.
- Experimental setup: Models use a convolutional auto-encoder and a four-layer 3D-CNN context model, with latent channel count K varied across models.The auto-encoder architecture is described in Figure 2, while the context model is described in Figure 3.
- Experimental setup: All models are trained and tested using MS-SSIM-based distortion, with β controlling the rate-distortion trade-off.Training uses β = 10, with six epochs and approximately 24 hours per model on a single GPU.
- Rate-distortion results: The method outperforms BPG, JPEG, and JPEG2000 across Kodak from 0.3 bpp to 0.9 bpp and across the other reported testing sets.The Kodak comparison also reports performance comparable to Rippel & Bourdev.
- Visual comparisons: At very strong compression on kodim21, the network preserves clouds well, while BPG preserves high frequencies better and JPEG2000 produces block artifacts.JPEG breaks down at the tested rate.
- Context-model ablation: 10 % rate reduction is achieved by the context model even though the auto-encoder was optimized using a uniform prior.The ablation compares context modeling for a fixed auto-encoder and reports the reduction for the proposed context model.
- Importance-map study: An importance-map model compresses an ImageNetTest image to 0.463 bpp, compared with 0.504 bpp for a similar-rate model without an importance map.The importance map dynamically changes the number of channels used at different spatial locations.
5. Discussion
The proposed system combines a convolutional auto-encoder with a lightweight 3D-CNN context model trained concurrently, achieving strong compression performance across multiple datasets and against context-model baselines.
- The convolutional auto-encoder and lightweight 3D-CNN context model are trained concurrently, with the context model incorporated into entropy optimization.
- The method outperforms BPG and JPEG2000 in MS-SSIM across ImageNetTest, Kodak, B100, and Urban100.
- The method achieves performance competitive with the current state-of-the-art learned compression method from [14].
- The 3D-CNN context model is significantly more powerful than first-order histogram and second-order one-step-prediction baselines.
- The importance map learns to condense image information into fewer latent channels without explicit supervision and is trained jointly without optimization difficulties.
6. Conclusions
The paper introduces joint training of a lossy image-compression auto-encoder and lightweight context model through an entropy loss, reaching competitive deep-compression performance while identifying heavier context models as future work.
- The paper proposes jointly learning a lossy image-compression auto-encoder and lightweight context model by incorporating the context model into the auto-encoder’s entropy loss.
- The resulting system achieves performance competitive with the current state of the art in deep image compression.
- Future work could explore heavier and more powerful context models to further improve compression performance.
Conditional Probability Models for Deep Image Compression – Suppl. Material
The supplementary material details the masked 3D-CNN context model, its joint-training benefit, and visual comparisons showing strengths and weaknesses relative to established codecs.
- 3D probability classifier: A 3D-CNN introduces a depth dimension and slides over latent channels, enabling the context model to process W × H × K feature maps.
- 3D probability classifier: Masked convolutions enforce causality by conditioning each latent symbol only on previously encoded spatial and depth positions.
- On the benefit of 3DCNN and joint training: The masked 3D-CNN with joint training reduces bpp by 22% relative to the separately trained 3D-CNN when comparing mean points.
- Non-realistic images: On Manga1097, the approach preserves text sufficiently for legibility and preserves facial fine texture better than BPG, although BPG renders text more crisply.
- Visual examples: Across visual examples, the approach can over-blur high frequencies, while BPG discards some low-contrast high frequencies that the approach preserves.The paper relates this difference to BPG’s MSE optimization and the proposed method’s MS-SSIM optimization.