Understanding Rate-Distortion Performance in Distributed Transformer Inference

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. We study compressibility of transformer's intermediate representations via learned compression through the lens of rate-distortion and the theory of usable information. Our study reveals that, unlike convolutional models, deeper representations in transformers become more difficult to compress. The reason for this behavior is twofold: first, the complexity of representation increases as we move deeper into the transformer, necessitating a higher rate; and second, the higher complexity of representations worsens the generaliztion bound for learned entropy estimates, which further compromises compression performance. Through a combination of experiments and theory, we characterize and analyze the compressibility of transformer representations, derive bounds on the achievable rate of learned codecs applied to these representations, and offer a unified lens for understanding rate-distortion performance in representation coding.