[Paper Review] Variational image compression with a scale hyperprior

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

Image and Signal Denoising Methods15 references1,085 citations

TL;DR

An end-to-end variational autoencoder for image compression that introduces a scale hyperprior to model latent dependencies, achieving state-of-the-art MS-SSIM and strong PSNR performance among ANN-based methods.

ABSTRACT

We describe an end-to-end trainable model for image compression based on variational autoencoders. The model incorporates a hyperprior to effectively capture spatial dependencies in the latent representation. This hyperprior relates to side information, a concept universal to virtually all modern image codecs, but largely unexplored in image compression using artificial neural networks (ANNs). Unlike existing autoencoder compression methods, our model trains a complex prior jointly with the underlying autoencoder. We demonstrate that this model leads to state-of-the-art image compression when measuring visual quality using the popular MS-SSIM index, and yields rate-distortion performance surpassing published ANN-based methods when evaluated using a more traditional metric based on squared error (PSNR). Furthermore, we provide a qualitative comparison of models trained for different distortion metrics.

Motivation & Objective

• Develop an end-to-end variational image compression model using a hyperprior to capture spatial dependencies in latent codes.
• Train the hyperprior jointly with the autoencoder to minimize rate–distortion loss.
• Evaluate performance on PSNR and MS-SSIM against existing ANN-based methods and conventional codecs.
• Analyze effects of training with different distortion metrics on reconstruction quality and artifacts.

Proposed method

• Use an analysis transform ga(x; φg) to produce latent y and a synthesis transform gs(ˆy; θg) to reconstruct x.
• Introduce a scale hyperprior by modeling latent scales with z and predicting ˜σ via hs to hierarchically condition the prior p˜y|˜z.
• Model q(˜y, ˜z | x) as a product of uniform variates to enable tractable inference and backpropagation.
• Train with a differentiable surrogate using additive uniform noise in place of quantization and optimize a KL-divergence based objective that includes rate and distortion terms.
• Employ a non-parametric, fully factorized density model for the priors p˜y|ψ and p˜z|ψ to enable end-to-end learning.
• Encode ˆz (side information) via arithmetic coding and use it at the decoder to retrieve ˆσ and accurately decode ˆy.

Experimental results

Research questions

• RQ1 Does a hyperprior on latent scales improve rate–distortion performance compared to a factorized prior?
• RQ2 How does conditioning the latent prior on spatially varying scales affect compression efficiency?
• RQ3 What is the impact of optimizing for different distortion metrics (PSNR vs MS-SSIM) on visual quality and artifacts?
• RQ4 How much side information (ˆz) is required to gain the benefits of the hyperprior?

Key findings

• The hyperprior model achieves state-of-the-art MS-SSIM on Kodak, outperforming conventional codecs and prior ANN methods for that metric.
• When optimized for squared error, the factorized-prior model yields strong PSNR performance and approaches or surpasses some conventional codecs and ANN methods, with the hyperprior providing additional gains.
• Side information for the hyperprior remains a small fraction of total bitrate (well below 0.1 bpp even at high rates), yet yields notable rate–distortion improvements.
• Training the same model with MS-SSIM loss yields different artifact patterns than training with squared error, highlighting qualitative differences across distortion metrics.
• The hyperprior consistently improves over the factorized-prior model across rate–distortion tradeoffs, underscoring the importance of matching the prior to the marginal.