Source-linked AI summary
Joint Autoregressive and Hierarchical Priors for Learned Image Compression
David Minnen, Johannes Ballé, George Toderici
TL;DR
Learned image compression seeks compact bitstreams from quantized latents while preserving reconstruction quality. The paper combines hierarchical and autoregressive priors, and reports state-of-the-art rate–distortion performance, with serial autoregressive decoding as the main practical limitation.
Problem
Existing learned codecs need entropy models that exploit latent structure beyond simpler factorized or hierarchical priors without sacrificing rate–distortion performance.
Method
The paper extends hierarchical entropy modeling with Gaussian mixtures and combines a hyperprior with an autoregressive model over causal latent context.
Results
The combined context-plus-hyperprior model outperforms existing codecs including BPG on Kodak rate–distortion curves for both PSNR and MS-SSIM evaluations.
Takeaways & Limitations
Autoregressive and hierarchical priors are complementary: the former adds causal predictions without potential side-information costs, while the latter reduces residual uncertainty.
Takeaways & Limitations
Autoregressive models are inherently serial and cannot be accelerated through parallelization, making them less practical than hierarchical models.
Abstract
from arXiv · showhide
Recent models for learned image compression are based on autoencoders, learning approximately invertible mappings from pixels to a quantized latent representation. These are combined with an entropy model, a prior on the latent representation that can be used with standard arithmetic coding algorithms to yield a compressed bitstream. Recently, hierarchical entropy models have been introduced as a way to exploit more structure in the latents than simple fully factorized priors, improving compression performance while maintaining end-to-end optimization. Inspired by the success of autoregressive priors in probabilistic generative models, we examine autoregressive, hierarchical, as well as combined priors as alternatives, weighing their costs and benefits in the context of image compression. While it is well known that autoregressive models come with a significant computational penalty, we find that in terms of compression performance, autoregressive and hierarchical priors are complementary and, together, exploit the probabilistic structure in the latents better than all previous learned models. The combined model yields state-of-the-art rate--distortion performance, providing a 15.8% average reduction in file size over the previous state-of-the-art method based on deep learning, which corresponds to a 59.8% size reduction over JPEG, more than 35% reduction compared to WebP and JPEG2000, and bitstreams 8.4% smaller than BPG, the current state-of-the-art image codec. To the best of our knowledge, our model is the first learning-based method to outperform BPG on both PSNR and MS-SSIM distortion metrics.
1 Introduction
Learned image compression maps pixels to quantized latents with neural transforms and entropy models, then optimizes bit rate against reconstruction distortion. This work extends hierarchical priors with Gaussian mixtures and autoregressive context to better model latent structure.
- Transform coding: Learned codecs map pixels into quantized, reduced-dimensionality latents and losslessly compress those latents.CNN-based transforms can learn nonlinear mappings into more compressible latent spaces than traditional linear transforms.
- Rate–distortion optimization: The training objective balances expected bitstream length against reconstructed-image distortion through rate–distortion optimization.The rate term measures cross entropy under a learned entropy model, while distortion can use MSE or MS-SSIM-related objectives.
- Hierarchical priors: Hierarchical priors condition latent entropy models on a quantized hyperprior, enabling end-to-end optimization while exploiting latent structure.The prior work described uses a Gaussian scale mixture whose scale parameters depend on the hyperprior.
- Paper extensions: The paper generalizes hierarchical Gaussian scale mixtures to Gaussian mixtures and adds an autoregressive component to the entropy model.The extensions are evaluated across network architectures while computational capacity is not reduced to optimize complexity.
2 Architecture Details
The proposed compression architecture combines a core autoencoder with a probabilistic latent model that joins a hyperprior and causal context. Quantized latents and hyper-latents are entropy-coded, while decoder-side access is restricted to information available from the bitstream.
- Architecture: The architecture has a core autoencoder and a probabilistic sub-network for entropy coding quantized latents.The probabilistic sub-network combines a Context Model with Hyper Encoder and Hyper Decoder blocks.
- Architecture: The Context Model uses causal latent context, while the hyper-network supplies corrective information for predicting entropy parameters.The Entropy Parameters network combines both sources to generate conditional Gaussian means and scales.
- Decoding constraints: The receiver may use only decoded information, so the context model can access previously decoded latents but not future latents.Arithmetic encoding stores quantized symbols in the compressed file, and decoding begins with previously decoded latents set to zero.
- Entropy model: Each latent is modeled as a Gaussian convolved with unit uniform noise, with mean and scale conditioned on hyperprior information and causal context.The conditioning functions use the hyper-decoder, context model, and entropy-parameters networks.
- Objective: The full squared-error loss includes rates for both quantized latents and hyper-latents in addition to distortion.The hyperprior uses a non-parametric, fully factorized density model, and its contribution is included because both streams enter the bitstream.
- Layer constraints: The model components are constrained where necessary to ensure that the bitstreams alone suffice for image reconstruction.The encoder bottleneck determines the number of latent elements, while the entropy-parameters output has twice the bottleneck channels for Gaussian means and scales.
- Layer constraints: In practice, the autoregressive context uses 5×5 masked convolutions rather than all previous latents.Masked convolutions preserve causal access while limiting the context used by the model.
3 Experimental Results
Experiments on Kodak show that the combined context-plus-hyperprior model achieves the strongest rate–distortion results across PSNR and MS-SSIM evaluations, with favorable visual quality at similar bit rates.
- Rate–distortion performance: The combined context-plus-hyperprior model outperforms existing codecs and learned methods on Kodak PSNR rate–distortion performance, including BPG.The authors identify it as the first learning-based compression model to outperform BPG on PSNR.
- Model comparisons: The combined context-plus-hyperprior configuration performs best among the tested model variants, while the context-only model performs slightly worse than hierarchical alternatives.
- Rate–distortion performance: The combined model also achieves better Kodak MS-SSIM rate–distortion performance than previous methods when optimized for MS-SSIM.When optimized for MSE, it still exceeds all standard codecs and learning-based methods not specifically optimized for MS-SSIM.
- Model comparisons: The baseline combined architecture uses a 5×5 masked context convolution and conditional Gaussian entropy model; alternative distributions and context configurations generally perform worse.Replacing the Gaussian with a logistic distribution increases file size by 0.3%, while other tested changes reduce performance more substantially.
- Visual comparisons: At similar bit rates, the combined method provides higher visual quality than the scale hyperprior baseline and JPEG, while its perceptual quality is closer to BPG.BPG preserves some additional detail but uses up to 9.4% more bits in the comparison.
4 Related Work
Learned image compression evolved from early neural autoencoders to modern models using recurrent architectures, entropy models, context, and side information. The paper also relates these models to Bayesian generative models while emphasizing compression-specific quantization and arithmetic-coding constraints.
- Early neural compression: Early neural image-compression methods used autoencoders with small bottlenecks and uniform or vector quantization, without explicit entropy models.
- Modern learned compression: Later deep-learning approaches introduced recurrent autoencoders, learned entropy models, spatial bit-rate adaptation, and separate models targeting different bit rates.
- Modern learned compression: Modern learned codecs variously use fully factorized entropy models, autoregressive context in code space, or neural-network-predicted side information.
- Connections to generative models: Learned compression shares conceptual links with PixelCNN, variational autoencoders, PixelVAE, and VLAE, but directly optimizes a rate–distortion objective.
- Connections to generative models: Compression-specific requirements include quantization, arithmetic coding, and transitions between continuous differential entropy and discrete Shannon entropy.
5 Discussion
The paper shows that conditional Gaussian mixtures improve rate–distortion performance, while autoregressive and hierarchical priors complement each other by reducing uncertainty from different sources. Their practical benefit is balanced against the serial computation required by autoregressive decoding.
- Hierarchical priors: A conditional Gaussian mixture model outperforms the simpler GSM-based entropy model without increasing asymptotic model complexity.Conditioning both means and scales on the hyperprior improves latent likelihood modeling, while the associated side information remains worthwhile.
- Combined priors: Autoregressive context and hyperprior components are complementary because context avoids extra side information while the hyperprior stores information unavailable from causal context.The hyperprior can be fully known to the decoder from the compressed bitstream and reduce residual uncertainty in autoregressive prediction.
- Internal mechanisms: The combined model reduces prediction error and normalized-latent structure more effectively than the model variants shown in Figure 6.Its smaller predicted scales indicate lower uncertainty, and its normalized latents are closest to i.i.d. Gaussian noise.
- Computational trade-offs: Autoregressive models are inherently serial, preventing the same parallelization-based speedups available to hierarchical models.This makes them less desirable from a practical computational standpoint despite their compression benefits.
- Future directions: The authors are exploring lightweight context models and tighter integration of arithmetic decoding with differentiable autoregression to reduce computational requirements.They also identify more complex strictly hierarchical priors as an alternative way to avoid causality constraints.
A.1 PSNR on Kodak
On Kodak, the MSE-trained method achieves better PSNR rate–distortion performance than all evaluated baseline methods.
- PSNR: The MSE-trained method provides better PSNR rate–distortion performance than all baseline methods on the Kodak image set.Each curve point averages PSNR and bit rate over 24 Kodak images for one standard-codec Q value or learned-method λ.
A.2 MS-SSIM on Kodak
On Kodak, the MS-SSIM-optimized method outperforms all evaluated standard and learned codecs, while the MSE-trained version remains competitive.
- MS-SSIM: The MS-SSIM-trained method provides better MS-SSIM rate–distortion performance than all baseline methods on Kodak.The comparison includes standard codecs and learned methods evaluated with MS-SSIM.
- MS-SSIM: The MSE-trained method has higher MS-SSIM scores than every standard codec and every learned method not specifically optimized for MS-SSIM.The figure reports MS-SSIM scores in decibels for readability.
A.3 SSIM on Kodak
On Kodak and Tecnick, the method generally outperforms the evaluated baselines on structural-similarity metrics and PSNR rate–distortion performance, with one reported SSIM exception.
- SSIM on Kodak: Both MS-SSIM-optimized and MSE-optimized versions outperform all baseline methods according to SSIM, except for one inversion above 1.1 bpp.Above 1.1 bpp, the MSE-optimized version is worse than Ballé et al. when that method is optimized for MS-SSIM.
- PSNR on Tecnick: On Tecnick, the MSE-trained method provides better PSNR rate–distortion performance than all baseline methods.Each curve point averages PSNR and bit rate over 100 Tecnick images at one standard-codec Q value or learned-method λ.
Appendix B Summary of Rate Savings
The appendix compares average PSNR rate savings on Kodak against JPEG and against other codecs over a shared PSNR range. The model saves 8.41% over BPG and 35.52% over JPEG2000.
- Figure 11 reports each method’s average rate savings compared with JPEG (4:2:0), where higher scores indicate larger savings.Scores are calculated from the Kodak PSNR rate-distortion graph over 27.1–39.9 dB.
- 8.41% average rate savings over BPG (4:4:4) are reported for our method on Kodak.The comparison uses the PSNR rate-distortion graph over 27.1–39.9 dB.
- 35.52% average rate savings over JPEG2000 are reported for our method on Kodak.The baseline uses the OpenJPEG implementation and the shared PSNR range is 27.1–39.9 dB.
Appendix C Example Images
Across four Kodak examples at similar bit rates, the proposed method is described as providing the highest visual quality. Its reconstructions are generally sharper than the scale-only model while avoiding artifacts noted for competing codecs.
- Kodak examples: The proposed method provides the highest visual quality on the Kodak 15 image at similar bit rates.BPG shows geometric artifacts, JPEG2000 severe artifacts and blurriness, and JPEG preserves little visual information.
- Kodak examples: The proposed method provides the highest visual quality on the Kodak 19 image at similar bit rates.The scale-only version is blurrier, while BPG shows ringing and geometric artifacts despite preserving extra detail in some regions.
- Kodak examples: 10% fewer bits are used by the proposed method than the scale-only reconstruction while retaining slightly more sharpness on Kodak 20.The comparison highlights sharper details in the propeller and Yosemite Sam nose art.
- Kodak examples: The proposed method provides the highest visual quality on the Kodak 23 image at similar bit rates.The scale-only method has similar quality but is blurrier, while BPG and JPEG2000 exhibit ringing or boundary artifacts.
Appendix D Architecture Comparison
The appendix contrasts fully factorized, hyperprior, autoregressive, and combined architectures. The combined context-plus-hyperprior model uses complementary sources of latent information and achieves better rate-distortion performance than either component alone.
- Architecture variants: The fully factorized model assumes latents are independent and identically distributed and learns a fixed entropy model shared by encoder and decoder.Its compression system is primarily an autoencoder with differentiability requirements for the entropy model.
- Architecture variants: The scale-only hyperprior adapts an image-dependent conditional Gaussian scale mixture to each individual code using transmitted hyperprior information.The Hyper-Encoder forms the hyperlatent representation, and the Hyper-Decoder produces scale parameters.
- Architecture variants: The mean-and-scale hyperprior extends the scale-only variant by using a Gaussian mixture model whose means and scales are predicted from the hyperprior.The Hyper-Decoder transforms the hyperprior into both Gaussian parameters.
- Architecture variants: The context-only model predicts Gaussian-mixture parameters autoregressively without a hyperprior or additional bitstream cost.Because decoding is serial and raster-scan ordered, the model accesses only each code’s causal context.
- Architecture variants: The context-plus-hyperprior model combines the mean-and-scale hyperprior with autoregressive prediction, outperforming either approach alone in rate-distortion performance.The evaluation identifies the autoregressive model and hyperprior as complementary.