Rate-Distortion Optimization for Transformer Inference

TL;DR

This work proposes a principled rate-distortion framework for compressing intermediate transformer representations to enable multi-device inference. It introduces a transformer-based auto-regressive entropy model with a hyper-prior and defines the $\mathcal{V}$-entropy gap to quantify rate-entropy mismatch under modeling constraints, along with PAC-style generalization bounds. The approach is validated on language-model benchmarks and vision tasks, where it achieves substantial rate-distortion gains over Fourier-basis and direct-access baselines, and demonstrates practical speedups under constrained links. The analysis links rate to target representation statistics via covariance and Rademacher complexity, offering theoretical insight and practical guidance for designing learnable codecs for machines.

Abstract

Transformers achieve superior performance on many tasks, but impose heavy compute and memory requirements during inference. This inference can be made more efficient by partitioning the process across multiple devices, which, in turn, requires compressing its intermediate representations. In this work, we introduce a principled rate-distortion-based framework for lossy compression that learns compact encodings that explicitly trade off bitrate against accuracy. Experiments on language benchmarks show that the proposed codec achieves substantial savings with improved accuracy in some cases, outperforming more complex baseline methods. We characterize and analyze the rate-distortion performance of transformers, offering a unified lens for understanding performance in representation coding. This formulation extends information-theoretic concepts to define the gap between rate and entropy, and derive some of its bounds. We further develop probably approximately correct (PAC)-style bounds for estimating this gap. For different architectures and tasks, we empirically demonstrate that their rates are driven by these bounds, adding to the explainability of the formulation.

Figures (6)

• Figure 1: Architecture overview of the proposed codec. The AE and AD blocks correspond to arithmetic encoders and decoders, respectively. They use the probability distributions provided by the entropy models to encode their target representation into a bitstream and decode it back. The dotted line separates two devices, with bitstreams (gray blocks) connecting them.
• Figure 2: Architecture diagram of the different entropy models for the target representation $Y$. The direct-access entropy model replaces the proposed $g_y$ with a series of transformer blocks that combine the hyper-prior and the target representation. Q, K, V corresponds to the query, key, and value embeddings in an attention mechanism. The Fourier basis method is (a) with a different $g_w$.
• Figure 3: Rate-distortion performance for GPT-2. The rate is measured in bits-per-token (BPT). Perplexity is the exponent of the classification cross-entropy loss, used as distortion. Uncompressed is a model with no quantization or rate penalty ($\lambda = 0$). The proposed method outperforms the other entropy models. Its rate-distortion performance decreases with the split point.
• Figure 4: Rate, covariance determinant, and Rademacher complexity estimates at different split points, for GPT-2 Small, Pythia 160M, ViT B/16, and ResNet 34. Some axes are min-max scaled per method to facilitate comparison. The logarithmic scale is further scaled by $1/2D$. The Rademacher complexity and covariance determinant strongly correlate with rate. Only in ResNets, the rate-distortion performance increases with the split point.
• Figure 5: Rate-performance for GPT-2 evaluated on the the LAMBADA language task. The rate is measured in bits-per-token (BPT). Uncompressed is a model with no quantization or rate penalty ($\lambda = 0$).

Theorems & Definitions (23)

• Definition 1
• Theorem 1
• Theorem 2
• Definition 2
• Theorem 3
• Lemma 1
• proof
• Theorem 3
• proof
• Theorem 4