[Paper Review] Universally Quantized Neural Compression

Eirikur Agustsson, Lucas Theis|arXiv (Cornell University)|Jun 17, 2020

Neural Networks and ApplicationsComputer Science32 references45 citations

TL;DR

The authors propose using universal quantization to implement a differentiable, train-test matched communication channel for neural image compression, including a soft-rounding approach to interpolate between uniform noise and hard quantization; they demonstrate empirical gains on linear and hyperprior models.

ABSTRACT

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.

Motivation & Objective

• Motivate reducing train-test mismatch in end-to-end neural compression by removing non-differentiable quantization.
• Introduce universal quantization as an efficient, differentiable test-time channel.
• Bridge quantization and soft quantization to allow smooth transition and potential gains when annealing to hard quantization.
• Investigate gradient estimation techniques to stabilize training with uniform noise and quantization.
• Evaluate the impact of the approach on simple linear and hyperprior models.

Proposed method

• Use additive uniform noise as a differentiable approximation during training.
• Replace test-time quantization with universal quantization to maintain differentiability and match training/test phases.
• Derive the bit-cost as h[Y+U], where U is uniform noise, for the discrete index K conditioned on U.
• Introduce a soft rounding function s_alpha that interpolates between no rounding and hard rounding.
• Propose an analytical gradient integration technique to reduce gradient variance when computing expectations over U.
• Apply a hyperprior and a linear model to evaluate the approach on Kodak data with RD curves.

Experimental results

Research questions

• RQ1 Can universal quantization serve as an efficient, differentiable test-time channel for neural compression?
• RQ2 Does aligning training and testing by using universal quantization improve rate-distortion performance over traditional train-time noise approximations?
• RQ3 Can soft-rounding provide benefits by interpolating between uniform noise and hard quantization while maintaining differentiability?
• RQ4 Do gradient estimation techniques (expected gradients) stabilize training and improve convergence for soft rounding?
• RQ5 What is the impact of the approach on simple linear models and more complex hyperprior models?

Key findings

• Universal quantization yields a differentiable channel that eliminates train-test mismatch without sacrificing efficiency.
• Quantization can be viewed as a limiting case of universal quantization through soft rounding with annealing.
• Soft rounding and expected gradients improve RD performance, especially at low bitrates and with hyperprior models.
• Expected gradients stabilize training and speed convergence for larger soft-rounding parameters.
• Across models, UN + UQ often improves over UN + Q baseline, and adding SR can surpass test-time quantization performance at certain rates.