PaperlayerView exact revision on arXiv

Source-linked AI summary

Universally Quantized Neural Compression

Eirikur Agustsson, Lucas Theis

arXiv:2006.09952v2stat.MLcs.CVcs.ITcs.LGmath.IT

TL;DR

Additive uniform noise is a popular differentiable approximation to quantization, but training and test losses can differ. This paper implements the uniform noise channel at test time with universal quantization and reports improved performance, especially for low-complexity models.

  • Problem

    Additive uniform noise approximates quantization during training, but gaps between training and test losses remain.

    Source: §1 · pp. 1–2

  • Method

    The paper uses universal quantization to implement the uniform noise channel at test time and introduces soft-rounding to interpolate toward hard quantization.

    Source: §1 · p. 1, §6 · p. 8

  • Results

    A better match between training and test phases improves performance, especially in models of lower complexity.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Universal quantization efficiently communicates samples from a uniform distribution because it assumes a constrained distribution class.

    Source: §6 · p. 8

    The study examines only one-dimensional uniform dither; broader distribution classes remain a direction for future work.

    Source: §6 · p. 9

Abstract

from arXiv · show

A popular approach to learning encoders for lossy compression is to use additive uniform noise during training as a differentiable approximation to test-time quantization. We demonstrate that a uniform noise channel can also be implemented at test time using universal quantization (Ziv, 1985). This allows us to eliminate the mismatch between training and test phases while maintaining a completely differentiable loss function. Implementing the uniform noise channel is a special case of the more general problem of communicating a sample, which we prove is computationally hard if we do not make assumptions about its distribution. However, the uniform special case is efficient as well as easy to implement and thus of great interest from a practical point of view. Finally, we show that quantization can be obtained as a limiting case of a soft quantizer applied to the uniform noise channel, bridging compression with and without quantization.

1 Introduction

Neural lossy compression has advanced, but quantization remains a challenge because it blocks direct backpropagation and creates a potential training–test mismatch. The paper uses universal quantization to implement uniform noise at test time, and reports improved performance, especially for lower-complexity models.

1 Introduction

  • End-to-end neural compression has advanced from barely beating JPEG2000 [4] to outperforming manually designed image-compression schemes [2], while practical challenges remain.The introduction names computational complexity, temporal inconsistencies, and perceptual metrics that are effective yet easy to optimize.Source: §1 · p. 1
  • Quantization produces a discrete representation for transmission but is non-differentiable, so common methods use a differentiable training approximation and quantize at test time.The paper identifies the effect of this training–test mismatch on performance as unclear.Source: §1 · p. 1
  • Existing alternatives communicate information differentiably at training and test time, but tend to be computationally costly or statistically inefficient.The paper frames its approach as bridging quantized and quantizationless compression.Source: §1 · p. 1
  • Universal quantization implements the uniform-noise channel at test time, eliminating the training–test mismatch while retaining differentiability.The paper also shows that noise can be integrated out analytically when calculating gradients, sometimes drastically reducing their variance.Source: §1 · pp. 1–2
  • A better match between training and test phases improves performance, especially for models of lower complexity.The paper also reports that analytically integrating out noise can sometimes drastically reduce gradient variance.Source: §1 · p. 2

2 Related work

Prior work addresses non-differentiable quantization with stochastic, straight-through, and smooth approximations, while another line of work simulates noisy channels using noiseless communication. These approaches motivate the paper’s focus on differentiable compression and efficient sample communication.

2 Related work

  • Stochastic rounding [32], straight-through gradient estimates [31], and annealed smooth vector quantization [1] are prior approaches to handling quantization in end-to-end compression.Theis et al. [31] use hard quantization during training and inference; Agustsson et al. [1] anneal a smooth approximation toward hard quantization.Source: §2 · p. 2
  • Ballé et al. [4] use additive uniform noise during training as an approximation to test-time rounding; under mean-squared error, this is equivalent to a variational autoencoder with a uniform encoder [25] [18].The cited formulation uses a density and uniform noise drawn from U([−0.5, 0.5)D).Source: §2 · p. 2
  • Reverse channel coding communicates samples from a noisy channel using a noiseless one, with the reverse Shannon theorem bounding bits by mutual information.Prior approaches include the likelihood encoder [12], MIRACLE [15], and rejection sampling; MIRACLE was applied to lossy image compression.Source: §2 · p. 2
  • Prior reverse-channel methods select among shared random samples, but their implementations differ in how they select the communicated index.The cited work includes likelihood-based stochastic selection [11] [12] and importance sampling [15].Source: §2 · p. 2
  • Universal quantization had been used as a relaxation of hard quantization [9], but with a non-differentiable loss and straight-through estimates [7]; Roberts applied it directly to grayscale pixels and reported better picture quality than quantization.A cited communication bound has overhead that remains significant for small amounts of information but becomes negligible as mutual information increases.Source: §2 · pp. 2–3

3 Compression without quantization

Communicating arbitrary samples efficiently requires assumptions about their distributions, while the uniform-noise channel admits an efficient implementation through universal quantization. This provides a differentiable alternative for compression without quantization.

3 Compression without quantization

  • General noisy-channel simulation requires a number of random samples exponential in DKL[qZ|x || pZ], creating a trade-off between computational cost and statistical overhead.Because mutual information lower-bounds average KL divergence, reducing computational cost can conflict with transmitting enough information to limit statistical overhead.Source: §3 · p. 3
  • If RP ≠ NP, no algorithm can approximate an arbitrary q within total variation distance 1/12 in time polynomial in DKL[q || p], given unlimited samples from p.The result motivates making assumptions about the distributions when designing efficient sample-communication algorithms.Source: §3 · p. 3

3.1 Uniform noise channel

  • The additive uniform-noise channel is computationally and statistically efficient to simulate, unlike general noisy-channel approaches.The paper focuses on uniform noise while noting that other noise distributions, such as Gaussian, are possible.Source: §3.1 · p. 3, §3.2 · p. 4

3.2 Universal quantization

  • Universal quantization uses random-offset quantization whose distribution matches adding uniform noise, allowing it to implement the differentiable channel at deployment instead of switching to hard quantization.Training adds uniform noise; deployment uses universal quantization to remove the training–test mismatch.Source: §3.2 · p. 3
  • The encoder transmits K = ⌊Y − U⌉ and the decoder uses U to compute K + U; the conditional coding cost is H[K | U] = I[Y, Y + U] = h[Y + U].This cost is differentiable when the density of Y is differentiable, and universal quantization has complexity linear in the number of coefficients rather than exponential in the number of bits.Source: §3.2 · pp. 3–4
  • The paper assumes independent noise across coefficients; using a shared offset is not equivalent to adding independent or dependent training noise.Dependent noise would fail to create an information bottleneck because a zero coefficient could reveal the noise and thereby the other coefficients.Source: §3.2 · p. 4

4 Compression with quantization

The paper treats uniform noise as a differentiable communication channel and shows how soft rounding can recover hard quantization as a limit. It also develops expected-gradient calculations to address variance from soft rounding.

4 Compression with quantization

  • Universal quantization can be suboptimal for mean-squared error under some conditions with a fixed encoder and decoder, but flexible encoders and decoders let quantization arise as a limit while retaining a differentiable loss.This frames the section’s bridge between differentiable uniform-noise communication and quantization.Source: §4 · p. 4

4.1 Simulating quantization with uniform noise

  • Applying soft rounding before and after the uniform-noise channel recovers hard quantization as α tends to infinity, allowing α to be annealed during training while keeping the loss differentiable.The parameter α controls the soft-rounding function’s approximation fidelity.Source: §4.1 · pp. 4–5
  • The reconstruction minimizing mean-squared error after observing soft-rounded noisy data is the conditional expectation E[Y | sα(Y) + U = z].The posterior is a truncated prior; the approximation assumes the prior is smooth enough to be approximately uniform within each interval.Source: §4.1 · p. 4

4.2 Reducing the variance of gradients

  • For large α, soft-rounding derivatives are usually near zero but occasionally large, which can give encoder gradients high variance.The paper proposes analytically integrating out uniform noise to address this gradient-variance problem.Source: §4.2 · p. 5
  • The gradient of an expectation over uniform noise can be computed with finite differences, including when the function is not differentiable.For the multivariable approximation, the loss derivative and the derivative through the perturbed input are assumed uncorrelated; this holds, for example, if the loss is locally linear around h(y).Source: §4.2 · pp. 5–6
  • For soft rounding and its inverse, the expected derivative is always 1, matching a straight-through gradient estimate.The paper applies the expectation-based gradient calculation to the rate term and uses it for the distortion term with rα in place of h.Source: §4.2 · p. 6

5 Experiments

Experiments compare uniform-noise and soft-rounding approaches on linear and hyperprior compression models evaluated on Kodak. Across both models, soft rounding improves rate-distortion performance over test-time quantization, while expected gradients help convergence and training stability.

5.1 Models

  • The experiments compare a linear model and a non-autoregressive Mean & Scale Hyperprior model using uniform noise with test-time quantization, uniform noise with universal quantization, and uniform noise with universal quantization plus soft rounding.The hyperprior encodes noisy hyper latents first, then uses them to conditionally encode noisy coefficients; soft rounding was integrated around the predicted mean.Source: §5.1 · pp. 6–7

5.2 Traininig

  • Models were evaluated on Kodak rate-distortion curves, plotting bits per pixel against peak signal-to-noise ratio.Training used 256x256 crops from one million high-resolution JPEG images; soft-rounding α was linearly annealed from 1 to 16 over two million steps.Source: §5.3 · p. 7, §5.2 · p. 7

5.3 Results

  • Soft rounding improves rate-distortion performance over the test-time quantization baseline for the linear model and at low bitrates for the hyperprior model.Test-time quantization itself improves rate-distortion performance relative to using the uniform-noise channel at inference; the soft-rounding gain is weaker for the hyperprior model.Source: §5.3 · p. 7, Figure 2 · p. 7, Figure 3 · p. 8
  • Expected gradients speed convergence for the linear model and improve hyperprior performance and training stability, especially for α > 7.Without expected gradients, hyperprior models with α > 7 did not train stably, particularly at higher bitrates, and performed more poorly.Source: §5.3 · p. 8

6 Conclusion

Universal quantization connects uniform dither to efficient sample communication, enabling a differentiable bridge between uniform noise and hard quantization. The paper reports improved performance, especially in lower-complexity models, while its study is limited to one-dimensional uniform dither.

6 Conclusion

  • Universal quantization connects efficient sample communication with uniformly dithered quantization, which works efficiently because it targets a constrained class of distributions.The conclusion also notes that soft-rounding can interpolate between the two approaches.Source: §6 · p. 8
  • The empirical findings use MSE, and the linear model’s blocking artefacts under quantization are almost completely masked by universal quantization, at the expense of grain.The conclusion notes that perceptual training losses may be more sensitive to training–test mismatch.Source: §6 · pp. 8–9
  • The study covers only one-dimensional uniform dither; the authors point to generalizations and richer efficiently communicable distributions as future directions.Source: §6 · p. 9

Broader Impact

The broader-impact discussion links compression rates to information access while noting image-level bit-allocation trade-offs and the need for representative training sets. Appendix A lists the paper’s notation for random variables, distributions, and information measures.

Broader Impact

  • Improved compression rates could make information more broadly available, with image and video compression potentially affecting many people given video’s share of global IP traffic.The passage reports that about 79% of global IP traffic is video.Source: Broader Impact · p. 9
  • Assigning fewer bits to one image entails assigning more to others, so training sets need to be representative; generative misrepresentation is outside this paper’s scope.Source: Broader Impact · p. 9

Appendix A: Notation

  • Appendix A’s notation distinguishes scalar and vector random variables, continuous densities and discrete mass functions, and entropy, expectation, divergence, and distance measures.Source: Appendix A · p. 11

Appendix B: Computational complexity of reverse channel coding

The appendix proves that communicating samples from arbitrary distributions cannot generally be done in polynomial time under RP ≠ NP. It relates this hardness to approximate RBM simulation and rules out algorithms whose runtime is polynomial in the KL divergence from the source distribution.

Appendix B: Computational complexity of reverse channel coding

  • The hardness argument draws on the result that approximately simulating an RBM is computationally hard under RP ≠ NP.The cited theorem concerns an efficiently evaluatable representation of a distribution close to the RBM distribution.Source: Appendix B · p. 11
  • Under RP ≠ NP, no algorithm can achieve the stated approximate sample-communication task for arbitrary q in time polynomial in D_KL[q || p].The algorithm receives unlimited i.i.d. samples from p and must output a distribution within total variation distance 1/12 of q.Source: Appendix B · p. 11
  • A polynomial-in-KL sampling algorithm for an RBM with a uniform source would run in polynomial time in the RBM dimension, contradicting the stated hardness result.Source: Appendix B · p. 11

Appendix C: Generalizations of universal quantization

The appendix discusses lattice and scale-mixture routes for generalizing universal quantization beyond its simple uniform distribution. It also establishes that the soft-rounding function is differentiable everywhere.

Appendix C: Generalizations of universal quantization

  • The main approach communicates samples from a simple uniform distribution; the appendix discusses lattice quantizers as one generalization.The example visualizes lattice points, Voronoi cells, coefficients y, and realizations of y + U.Source: Appendix C · p. 12, Figure 4 · p. 12
  • For certain lattices in high dimensions, the Voronoi-cell noise U is approximately Gaussian, suggesting a way to simulate an additive white Gaussian noise channel.Source: Appendix C · p. 12
  • A uniform scale mixture can implement noise such as Gaussian noise, but transmitting Z alone requires more bits than desired; treating (Z, S) as the message uses only as many bits as the information transmitted.The construction assumes S is positive and known to both encoder and decoder.Source: Appendix C · p. 12

Appendix D: Differentiability of soft-rounding

  • The soft-rounding function is differentiable everywhere because its left and right derivatives agree at zero, and its periodic shift relation extends differentiability to other integers.Source: Appendix D · p. 13

Appendix E: Adapting density models

The appendix adapts density models for noisy encoder outputs, first modeling independent components through their cumulative distribution and then extending the model to invertible transforms. Making the density depend on α lets it adapt when α changes during training.

Appendix E: Adapting density models

  • For independent components of f(x), the rate model needs only the density of each component after adding uniform noise.The model is parameterized through the cumulative distribution of Y.Source: Appendix E · p. 13
  • For an invertible transform s, the density model extends to s(Y) + U using the transformed cumulative distribution.The passage gives c_s(Y) = c_Y(s^-1(y)).Source: Appendix E · p. 13
  • Making the density model depend explicitly on α lets it adapt automatically when α changes during training.Source: Appendix E · p. 13

Appendix G: Qualitative results

The appendix compares low-bitrate Kodak reconstructions across three approaches. Linear-model outputs show blocking with UN + Q and UN + UQ + SR, while UN + UQ largely masks blocking at the cost of grain; hyperprior UN + Q tends to show grid artefacts absent from UN + UQ + SR.

Appendix G: Qualitative results

Related papers

Loading 2006.09952v2…