PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Inference for Neural Image Compression

Yibo Yang, Robert Bamler, Stephan Mandt

arXiv:2006.04240v4eess.IVcs.LGstat.ML

TL;DR

VAE-based lossy image compression relies on inference networks to produce compressible latent representations, but conventional amortized inference leaves room for improvement. The paper combines iterative inference, discrete stochastic annealing, and lossy bits-back coding to address three approximation gaps. With an established architecture otherwise unchanged, these inference changes achieve new state-of-the-art lossy image-compression performance.

  • Problem

    Conventional amortized inference in VAE-based compression leaves substantial room for improvement, while continuous relaxations and discretization also leave approximation gaps.

    Source: §3.1 · p. 3, §3.2 · p. 4

  • Method

    The paper combines hybrid amortized-iterative inference, stochastic annealing over discrete representations, and a remedy enabling bits-back coding for lossy compression.

    Source: §1 · p. 2

  • Results

    The three proposed inference enhancements significantly improve compression performance and achieve new state-of-the-art lossy image compression using an otherwise unchanged architecture.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Improved inference is a promising direction for compression that is orthogonal to modeling choices.

    Source: §5 · p. 9

    The proposed iterative inference methods incur about a 100× slowdown versus standard encoder-network prediction, although the standalone SGA variant does not affect decompression speed.

    Source: S4 Additional Results · p. 15

Abstract

from arXiv · show

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.

1 Introduction

Conventional amortized inference in VAE-based compression leaves room for improvement. The paper proposes three inference innovations and reports new state-of-the-art lossy image-compression performance using an otherwise unchanged architecture.

1 Introduction

  • Deep learning methods have begun to outperform classical codecs on image compression, and image compression is also a stepping stone toward video codecs.The passage notes that improved video codecs could reduce global internet traffic.Source: §1 · p. 1
  • Neural compression uses VAE inference to map images to compressible latent representations that a generative network maps back to reconstructed images.Compression can be viewed as inference over a specified decoding procedure.Source: §1 · p. 1
  • Iterative inference builds on a connection between compression procedures and variational-inference methods for closing the amortization gap.The paper connects an iterative compression procedure [Campos et al., 2019] to the broader VI literature.Source: §1 · p. 2
  • Stochastic annealing performs inference directly over discrete latent points, addressing the discretization required for entropy coding.Existing methods typically relax discretization during inference and discretize afterward.Source: §1 · p. 2
  • A proposed remedy to the decoder-side inference obstacle makes bits-back coding applicable to lossy compression for the first time.Bits-back coding had been considered incompatible with lossy compression because the decoder lacks the exact undistorted input [Habibian et al., 2019].Source: §1 · p. 2
  • The three inference innovations significantly improve compression performance and achieve new state-of-the-art lossy image compression on an otherwise unchanged architecture.The evaluation compares against a wide range of baselines and ablations.Source: §1 · p. 2

2 Background: Lossy Neural Image Compression as Variational Inference

The background describes lossy image compression as a VAE process in which latent variables are encoded and decoded, with training balancing rate and distortion. Its relaxed training objective is connected to variational inference.

Generative Model

  • A hierarchical VAE models an image through latents y and hyperlatents z, with a generative network mapping the latents to the image distribution.The image likelihood is Gaussian, and a second network supplies the conditional prior parameters for y given z.Source: §2 · p. 2

Compression and Decompression

  • The encoder maps an image to continuous latent estimates, rounds them to discrete values, and entropy-codes them using the hyperprior and conditional prior.The decoder recovers the latent values and reconstructs the image from y.Source: §2 · p. 3

Model Training

  • Training minimizes a rate-distortion objective that trades bitrate against reconstruction error, while uniform noise replaces rounding to permit gradient-based optimization.The relaxed objective samples from width-one uniform distributions centered on the continuous latent estimates.Source: §2 · p. 3

Connection to Variational Inference

  • The relaxed compression objective is the negative evidence lower bound of variational inference under the stated identification of λ with the likelihood variance.The paper distinguishes using VI to compress data with a trained model from variational EM, which trains a model.Source: §2 · p. 3

3 Novel Inference Techniques for Data Compression

The paper identifies amortization, discretization, and marginalization gaps in VAE-based compression and proposes inference-time methods to address them. These methods refine amortized inference, optimize discrete representations with annealing, and apply bits-back coding to lossy compression.

3.1 Amortization Gap and Hybrid Amortized-Iterative Inference

  • At compression time, hybrid amortized-iterative inference initializes local variational parameters from trained inference networks, then minimizes the NELBO over those parameters without changing model training.The approach reframes compression-time iterative inference as a way to refine the amortized network's output.Source: §3.1 · pp. 3–4

3.2 Discretization Gap and Stochastic Gumbel Annealing (SGA)

  • SGA initializes continuous proxies from the inference networks, stochastically rounds coordinates to integers, and anneals temperature toward zero while optimizing the rate-distortion objective.Gradients through Bernoulli rounding samples are propagated with the Gumbel-softmax trick.Source: §3.2 · pp. 4–5
  • SGA closes the discretization gap and achieves the lowest true rate-distortion loss among the compared methods; deterministic annealing also closes the gap but converges to a worse solution.MAP converges to a non-integer solution, STE diverges, and Uniform Noise produces a consistently negative discretization gap.Source: §3.2 · pp. 5–6

3.3 Marginalization Gap and Lossy Bits-Back Coding

  • The marginalization gap arises because the ideal code for y alone uses an intractable marginal prior, so standard compression transmits hyperlatent z as well; the gap is the information in z given y.The paper relates this overhead to the information content of the hyperlatent under the ideal posterior P(ẑ|ŷ).Source: §3.3 · p. 6
  • Lossy bits-back coding uses a two-stage inference procedure: jointly optimize the discrete representation and hyperlatent parameters, then refit those parameters reproducibly from ŷ so the decoder can recover side information.The method encodes side information in the choice of ẑ using an entropy model; the decoder reproduces that model from ŷ.Source: §3.3 · pp. 6–7

4 Experiments

The experiments apply two inference variants to an established hyperprior model and compare them with neural and classical codecs. Both variants improve substantially over the base model and outperform prior state of the art; ablations attribute gains to iterative inference and stochastic rounding, with a smaller gain from bits-back coding.

Results

  • Over 15% BD rate savings on Kodak and 20% on Tecnick; both proposed variants outperform the Base Hyperprior and previous state-of-the-art methods.The Kodak comparison uses PSNR and includes rate-distortion curves and BD rate savings relative to BPG.Source: §4 · pp. 7–8, Figure 4 · p. 8

Proposed Methods

  • SGA changes only compression-time inference in the Base Hyperprior model; SGA+BB adds bits-back coding and supports image compression with side information.SGA combines hybrid amortized-iterative inference with Stochastic Gumbel Annealing; SGA+BB uses a more flexible hyperlatent inference distribution.Source: §4 · p. 7, Table 1 · p. 8

Baselines and Results

  • The comparison includes the Base Hyperprior, neural compression methods, BPG, and JPEG 2000; the proposed variants also improve visual quality at comparable bitrates.The qualitative comparison reports sharper hair and face details than the Base Hyperprior and avoids ringing artifacts shown for BPG.Source: §4 · pp. 7–8, Figure 6 · p. 9

Ablations

  • Iterative inference and SGA contribute approximately equally to performance gains, while bits-back coding contributes less; bits-back without the other improvements hurts performance.The authors suggest the latter result may reflect the constraint that the hyperlatent posterior conditions on the reconstructed image rather than the original image.Source: §4 · pp. 8–9, Figure 5 · p. 8

5 Discussion

The discussion presents improved inference as a promising compression direction distinct from modeling choices, while noting limits and possible extensions of the current approach.

Discussion

  • Improved inference is described as orthogonal to modeling choices, including autoregressive priors that can harm decoding efficiency.The authors say bits-back coding gave relatively minor benefits in the present VAE but may reach its potential in more hierarchical architectures; iterative inference during training may also improve performance.Source: §5 · p. 9
  • The authors identify more hierarchical architectures and iterative amortized inference during training as possible paths to greater gains.These are proposed directions, not results demonstrated in this work.Source: §5 · p. 9
  • In the present VAE, lossy bits-back coding provided relatively minor benefits.The discussion suggests it may be more effective in more hierarchical architectures.Source: §5 · p. 9

Broader Impacts

The broader-impact discussion links better compression to growing digital data needs and identifies both potential uses and current neural-compression drawbacks.

Broader Impacts

  • Improved image and video compression is presented as increasingly important as digital life expands.The discussion cites video streaming as a large share of web traffic and notes emerging data types such as medical images, LIDAR, graphics, and DNA sequences.Source: Broader Impacts · p. 9
  • Learnable compression could avoid the cost of designing custom codecs for diverse data, while better compression of data and models could support machine learning on local devices.The authors suggest local processing may address privacy concerns.Source: Broader Impacts · p. 9
  • Neural compression currently has higher computation, storage, and energy costs and lacks robustness compared with hand-engineered codecs.Improved encoding may lower overall resource consumption when its added cost is amortized across many file transfers or decoding operations.Source: Broader Impacts · p. 9

S1 Stochastic Annealing

The supplement illustrates stochastic annealing on a scalar integer-optimization example using a continuous proxy and stochastic rounding.

S1 Stochastic Annealing

  • The illustration starts from a scalar optimization problem over integer values.It is presented as a simple conceptual example of stochastic annealing.Source: S1 Stochastic Annealing · p. 13
  • A continuous proxy x and one-hot rounding direction r define stochastic rounding, with r modeled as a temperature-dependent Bernoulli variable.The method’s stochastic optimization problem is then introduced.Source: S1 Stochastic Annealing · p. 13
  • The supplement plots the rounding-up probability and SGA objective as functions of x at different temperatures.Figure S1 illustrates the probability of rounding up over the interval (0, 1).Source: S1 Stochastic Annealing · p. 13, Figure S1 · p. 13

S2 Model Architecture and Training

The supplement specifies the Base Hyperprior architecture, the modifications used for its bits-back variant, and the models’ training setup. It also describes the temperature-dependent stochastic annealing objective.

S2 Model Architecture and Training

  • Figure S2 shows ℓτ interpolating the discrete optimization problem between integers and approaching round(x)2 as τ → 0.The figure presents the objective at various temperatures.Source: Figure S2 · p. 13
  • The Base Hyperprior follows Minnen et al. [2018] but omits Context Prediction and Entropy Parameters, and doubles hyper-synthesis output channels to predict the conditional latent prior’s mean and log scale.The model is mostly based on Ballé et al.’s implementation.Source: S2 Model Architecture and Training · p. 13
  • For lossy bits-back, the hyper inference network predicts Gaussian q(z|x)’s mean and diagonal log variance, while the hyperprior uses a flexible density without uniform convolution.These are the two stated modifications to the Base Hyperprior.Source: S2 Model Architecture and Training · pp. 13–14
  • Models trained on CLIC-2018 crops used minibatches of eight 256 × 256 patches and λ values from 0.001 to 0.08, with channel counts raised to 256 at λ = 0.04 and 0.08.Training lasted 2 million steps except λ = 0.001 (1 million) and λ = 0.08 (3 million).Source: S2 Model Architecture and Training · p. 14

S3 Hyper-Parameters of Various Methods Considered

The supplement reports optimization settings for hybrid amortized-iterative inference, stochastic annealing, lossy bits-back, and alternative discretization methods. It also describes how a warm fixed-temperature phase affects the annealing schedule.

S3 Hyper-Parameters of Various Methods Considered

  • SGA used Adam at learning rate 0.005 with c = 0.001 and 2000 iterations; slower decay generally needed more iterations but could yield slightly better solutions.The exponential temperature schedule was τ(t) = min(exp{−ct}, 0.5).Source: S3 Hyper-Parameters of Various Methods… · p. 14
  • An initial optimization phase at τ0 = 0.5 for 700 steps made SGA results robust to decay factors; c = 0.002 and 0.001 performed comparably to 0.0005 with faster convergence.The fixed-temperature phase lets the rate-distortion objective roughly converge before annealing; Figure S3 compares schedules.Source: S3 Hyper-Parameters of Various Methods… · p. 14, Figure S3 · p. 14
  • SGA+BB jointly optimized µy, µz, and σ²z with BBVI for 2000 iterations, then used a reproducible BBVI encoder-decoder subroutine with the same random seed.The optimization used Adam; µy used the SGA temperature schedule, while the reproducible subroutine used learning rate 0.003.Source: S3 Hyper-Parameters of Various Methods… · pp. 14–15
  • Alternative discretization methods received 2000 Adam iterations, but STE had convergence issues even at a tiny learning rate; MAP and STE used early stopping for reported rate-distortion results.Figure 3 comparisons used the Base Hyperprior trained with λ = 0.0025.Source: S3 Hyper-Parameters of Various Methods… · p. 15

S4 Additional Results

Additional Kodak and Tecnick results report BD rate savings for SGA and SGA+BB against the Base Hyperprior and BPG. The supplement also compares image quality and notes the encoder-time cost of iterative inference.

S4 Additional Results

  • 17% (SGA+BB) and 15% (SGA) BD rate savings on Kodak versus Base Hyperprior; the savings were also 17% and 15%, respectively, versus BPG.These are the reported Kodak comparisons.Source: S4 Additional Results · p. 15
  • On Tecnick, SGA+BB and SGA achieve BD-rate savings of 20% and 19%, respectively, versus Base Hyperprior, and 22% and 21%, respectively, versus BPG.Figure S6 shows BD-rate savings relative to BPG as a function of PSNR, computed from Tecnick rate-distortion curves; higher values are better.Source: S4 Additional Results · p. 15, Figure S6 · p. 17
  • Iterative inference ran about 100× slower than standard encoder prediction, but standalone SGA did not change decompression speed.Experiments used a Titan RTX GPU with 24GB RAM.Source: S4 Additional Results · p. 15
  • The Kodak visual comparison shows SGA improving Base Hyperprior quality at similar bit rates, while BPG exhibits blocking, geometric, and ringing artifacts.The artifacts occur near the rooftop and around the lighthouse contour in the pictured crop.Source: Figure S4 · p. 16
Loading 2006.04240v4…