Rate-Distortion Optimization for Transformer Inference
This work addresses efficiency challenges in deploying large models for practitioners, though it is incremental as it builds on existing compression and information theory concepts.
The paper tackles the problem of high compute and memory requirements in Transformer inference by introducing a rate-distortion framework for lossy compression of intermediate representations, achieving substantial savings with improved accuracy on language benchmarks.
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.