Reproduce “Training Transformers for KV Cache Compressibility” on OpenResearch

About this paper

Yoav Gelberg, Yam Eitan, Michael Bronstein, Yarin Gal, Haggai Maron

KV-Compression Aware Training (KV-CAT) is a continued pretraining method that guides transformer language models to develop Key-Value (KV) cache representations inherently more amenable to post-hoc compression. This approach enhances the effectiveness of existing compression techniques, leading to improvements in perplexity, retrieval accuracy, and long-context question answering performance without degrading baseline model quality.

Key insights

• KV-compressibility is a characteristic of the specific learned internal representations within a transformer, not solely an intrinsic property of the task or input sequence.
• Theoretical analysis and a motivating example (histogram computation) demonstrate that different transformer implementations of the same function can have vastly different inherent compressibility.
• Models can be explicitly trained to produce more compressible internal representations, thereby improving the efficiency of post-hoc compression methods.

Summary by alphaXiv.