[Paper Review] Improving Inference for Neural Image Compression

Yibo Yang, Robert Bamler+1 more|arXiv (Cornell University)|Jun 7, 2020

Generative Adversarial Networks and Image Synthesis40 references46 citations

TL;DR

The paper identifies three approximation gaps in VAE-based neural image compression and introduces three corresponding inference improvements—hybrid amortized-iterative inference, stochastic Gumbel annealing for discretization, and lossy bits-back coding—achieving state-of-the-art results.

ABSTRACT

We consider the problem of lossy image compression with deep latent variable models. State-of-the-art methods build on hierarchical variational autoencoders (VAEs) and learn inference networks to predict a compressible latent representation of each data point. Drawing on the variational inference perspective on compression, we identify three approximation gaps which limit performance in the conventional approach: an amortization gap, a discretization gap, and a marginalization gap. We propose remedies for each of these three limitations based on ideas related to iterative inference, stochastic annealing for discrete optimization, and bits-back coding, resulting in the first application of bits-back coding to lossy compression. In our experiments, which include extensive baseline comparisons and ablation studies, we achieve new state-of-the-art performance on lossy image compression using an established VAE architecture, by changing only the inference method.

Motivation & Objective

• Motivate improved inference for neural lossy image compression using VAEs.
• Identify and formalize three approximation gaps in current inference approaches.
• Propose and validate methods addressing amortization, discretization, and marginalization gaps.
• Demonstrate performance gains over baselines and ablations on standard benchmarks.

Proposed method

• Analyze the variational autoencoder framework for lossy image compression and its rate-distortion objective.
• Introduce hybrid amortized-iterative inference to reduce the amortization gap during compression.
• Develop Stochastic Gumbel Annealing (SGA) to directly optimize discrete latent representations and close the discretization gap.
• Extend bits-back coding to lossy compression to address the marginalization gap with a two-stage inference procedure and reproducible BBVI on the decoder side.

Experimental results

Research questions

• RQ1 What are the main approximation gaps in conventional VAEs when used for lossy image compression?
• RQ2 Can hybrid amortized-iterative inference, discrete optimization via SGA, and lossy bits-back coding jointly close these gaps and improve compression performance?
• RQ3 How do these methods perform relative to established baselines on standard image compression benchmarks?

Key findings

• The proposed methods yield new state-of-the-art performance on lossy image compression using an established VAE architecture by changing only the inference method.
• SGA closes the discretization gap and achieves the lowest true rate-distortion loss among tested methods with the same initialization.
• The hybrid amortized-iterative approach improves compression by refining the variational parameters during compression time without retraining the model.
• Lossy bits-back coding enables exploiting side information to reduce the marginalization gap and improves overall rate-distortion performance.
• Experiments show significant rate-distortion improvements, with over 15% BD-rate savings on Kodak and over 20% on Tecnick datasets compared to previous methods.