Training Transformers to Be Compression‑Friendly: A New Memory‑Discard Paradigm

As AI vendors push context windows to millions of tokens, the KV‑Cache required by Transformers grows linearly, consuming tens to hundreds of gigabytes of GPU memory and slowing inference, turning the context‑window race into a memory‑war.

The authors highlight that two models producing identical outputs can differ drastically in KV‑Cache compressibility. Using a simple letter‑frequency task, they contrast a naïve implementation that averages token representations—fragile to compression—with a structured implementation that records sequence length, allowing the compressed cache to be recalibrated and retain zero error.

To encourage the latter, the joint Oxford‑NVIDIA team proposes KV‑CAT (KV‑Compression Aware Training). The method adds lightweight router modules that, at each training step, decide which KV slots to keep (target ~50% retention). Training proceeds with two forward passes per batch—full‑cache and compressed‑cache—combined with three loss terms: a self‑distillation loss aligning compressed‑cache outputs to full‑cache outputs, an anchoring loss preserving next‑token prediction ability, and a router loss enforcing the 50% retention target.

After training, the router is disabled; the model remains a standard Transformer but its internal representations are inherently compression‑friendly, allowing any downstream KV‑compression technique to work more effectively.

Experiments on Qwen2.5 models (0.5B and 1.5B parameters) show that baseline capabilities are unchanged (±0.5 % on six language‑understanding benchmarks). Under equal compression budgets, KV‑CAT reduces perplexity gaps by up to 3.21× and cuts required optimization steps by up to 5×. In a long‑text retrieval test, accuracy rises from 28 % to 47 % (0.5B) and 49 % to 67 % (1.5B). On LongBench v2, compressed KV‑CAT models improve average accuracy by up to 39 % across seven long‑document QA tasks.

The approach adds extra pre‑training cost and router complexity, and current results are limited to sub‑2B‑parameter models; scaling to hundred‑billion‑parameter LLMs remains an open question.

Nevertheless, training‑time preparation for inference, by making models naturally compressible, is presented as a valuable direction for future large‑model engineering.