Improving Inference for Neural Image Compression

TL;DR

This paper enhances neural image compression by addressing inference approximation gaps in hierarchical VAEs, applying iterative inference, stochastic annealing, and bits-back coding to achieve state-of-the-art results.

Contribution

It introduces novel methods to close inference gaps in VAEs for image compression, including the first use of bits-back coding in this context.

Findings

Achieved new state-of-the-art compression performance.

Demonstrated effectiveness of iterative inference and stochastic annealing.

Validated improvements through extensive experiments.

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.