[Paper Review] Joint Autoregressive and Hierarchical Priors for Learned Image Compression

David Minnen, Johannes Ballé+1 more|arXiv (Cornell University)|Sep 7, 2018

Generative Adversarial Networks and Image SynthesisComputer Science19 references593 citations

TL;DR

The paper extends learned image compression with a joint autoregressive context model and a hyperprior, achieving state-of-the-art rate–distortion performance that outperforms BPG and other codecs on PSNR and MS-SSIM.

ABSTRACT

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.

Motivation & Objective

• Extend GSM-based entropy models with a Gaussian mixture prior.
• Incorporate an autoregressive context model to further reduce entropy.
• Combine autoregressive context with a hyperprior to maximize rate–distortion performance.
• Evaluate variants and quantify trade-offs between context size, distributions, and complexity.
• Demonstrate state-of-the-art results on standard benchmarks (Kodak) and compare to BPG and other codecs.

Proposed method

• Generalized the entropy model from a scale hyperprior to a Gaussian mixture model conditioned on hyperpriors.
• Added an autoregressive context model over latents to predict mean and scale for each latent.
• Combined context model and hyperprior to form a joint entropy model with a conditional Gaussian; trained end-to-end with a rate–distortion objective.
• Modeled latents as Gaussian convolved with a unit uniform during training for differentiability.
• Used a two-part latent pipeline: latents y and hyper-latents z, each with respective entropy costs in the loss function.
• Explored architectural variants (context-only, hyperprior-only, and combined) and analyzed context size and distribution choices.

Experimental results

Research questions

• RQ1 Can autoregressive priors improve compression performance when combined with hierarchical priors in learned image codecs?
• RQ2 What is the impact of using a Gaussian mixture entropy model compared to a scale hyperprior on rate–distortion performance?
• RQ3 How does combining context (autoregressive) with a hyperprior affect bitstream size and reconstruction quality?
• RQ4 What are the practical trade-offs between model complexity, context size, and compression gains in learned image compression?

Key findings

• The combined autoregressive and hyperprior model achieves state-of-the-art rate–distortion performance.
• The combined model yields a 15.8% average reduction in file size over the previous state-of-the-art learned method.
• The reduction corresponds to a 59.8% size reduction over JPEG, and more than 35% reduction compared to WebP and JPEG2000.
• Bitstreams are 8.4% smaller than BPG, the then-current state-of-the-art codec.
• To the authors' knowledge, the model is the first learning-based method to outperform BPG on both PSNR and MS-SSIM.
• A Gaussian mixture entropy model provides benefits over a simpler GSM, without increasing asymptotic complexity.