Approaching Rate-Distortion Limits in Neural Compression with Lattice Transform Coding

2024-03-12

International Conference on Learning Representations

AI Summary

Existing neural compression methods employ scalar quantization (integer rounding plus entropy coding), leading to suboptimal quantization regions—deviating from optimal vector quantization—on sources with intrinsic dimension greater than one, thereby limiting rate-distortion (RD) performance. This paper introduces lattice quantization into the latent space of neural compression for the first time, proposing the Lattice Transform Coding (LTC) framework, which jointly optimizes the neural encoder, lattice quantizer, and entropy model. LTC enables the encoder to asymptotically achieve the RD function’s theoretical lower bound and supports block-wise coding for additional gains. Experiments on i.i.d. vector sources demonstrate that LTC significantly outperforms standard neural codecs in RD performance, closely approaching the theoretical optimum of vector quantization, while maintaining tractable computational complexity. The core contribution lies in embedding structured lattice quantization into an end-to-end trainable neural compression architecture, thereby bridging the gap between practical neural codecs and fundamental RD limits.

Technology Category

• Machine Learning: Learning on the Edge & Model Compression
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Neural compression has brought tremendous progress in designing lossy compressors with good rate-distortion (RD) performance at low complexity. Thus far, neural compression design involves transforming the source to a latent vector, which is then rounded to integers and entropy coded. While this approach has been shown to be optimal in a one-shot sense on certain sources, we show that it is highly sub-optimal on i.i.d. sequences, and in fact always recovers scalar quantization of the original source sequence. We demonstrate that the sub-optimality is due to the choice of quantization scheme in the latent space, and not the transform design. By employing lattice quantization instead of scalar quantization in the latent space, we demonstrate that Lattice Transform Coding (LTC) is able to recover optimal vector quantization at various dimensions and approach the asymptotically-achievable rate-distortion function at reasonable complexity. On general vector sources, LTC improves upon standard neural compressors in one-shot coding performance. LTC also enables neural compressors that perform block coding on i.i.d. vector sources, which yields coding gain over optimal one-shot coding.

Problem

• Addressing sub-optimal quantization in neural compression methods
• Improving vector quantization mismatch with lattice transform coding
• Enhancing rate-distortion performance beyond standard latent rounding

Innovation

• Lattice quantization replaces scalar quantization
• LTC recovers optimal vector quantization complexity
• Framework integrates near-optimal information-theoretic designs