Summary (Overview)

• Core contribution: The paper introduces KV-Compression Aware Training (KV-CAT), a continued pretraining procedure that trains transformer language models to produce Key-Value (KV) caches that are more amenable to post-hoc compression methods.
• Theoretical foundation: The authors formally define KV-compressibility and prove that almost any sequence-to-vector function admits both highly compressible and inherently non-compressible transformer implementations, demonstrating that compressibility is a property of learned representations rather than the input alone.
• Method: KV-CAT inserts lightweight learned routers between layer groups that mask KV slots during training, combined with a self-distillation loss matching masked to dense model outputs, a budget loss controlling retention rate, and a standard NTP anchor loss.
• Key results: KV-CAT improves suffix perplexity retention by up to 3.21×, speeds up gradient-based KV compression by up to 5×, improves retrieval accuracy by up to 68%, and improves long-context QA by up to 39% across QWEN2.5-0.5B and QWEN2.5-1.5B models.
• Crucially, KV-CAT preserves uncompressed model performance (within ±0.7 accuracy points on standard benchmarks), meaning compressibility is gained without sacrificing dense behavior.

Introduction and Theoretical Foundation

Background and Motivation

Long-context language modeling faces a significant memory bottleneck from the Key-Value (KV) cache, which stores key and value vectors for every token, layer, and attention head. For long sequences, this cache dominates both memory usage and decoding cost. Existing mitigation strategies fall into two categories:

1. Alternative architectures: Linear attention, state space models, and sparse attention variants—these typically trade off empirical performance for efficiency.
2. Inference-time KV cache compression: Optimization-based methods that operate on a fixed pretrained model, including gradient-based approaches [16] and attention-matching methods [61].

The key insight is that post-hoc compression methods are fundamentally limited by how well the model's internal representations can be compressed. Two transformers can produce identical next-token distributions while using very different internal representations, meaning some models' KV caches are inherently more compressible than others.

Formal Definition of KV-Compressibility

A KV cache compression policy for a transformer with $L$ layers is a collection of functions $\mathsf{C} = (\mathsf{c}_1, \ldots, \mathsf{c}_L)$, where each $\mathsf{c}_\ell$ maps sequences of $n$ key-value pairs to sequences of length $r(n) \leq n$. For KV pairs in layer $\ell$:

$$(\boldsymbol{K}, \boldsymbol{V}) = ([\boldsymbol{k}_1, \dots, \boldsymbol{k}_n]^{\top}, [\boldsymbol{v}_1, \dots, \boldsymbol{v}_n]^{\top})$$

The compressed attention computation becomes:

$$\operatorname{AttnHead}(\boldsymbol{Y}) = \operatorname{softmax}\left(\frac{1}{\sqrt{d_k}} \boldsymbol{Q}_b \begin{bmatrix} \tilde{\boldsymbol{K}}_a \\ \boldsymbol{K}_b \end{bmatrix}^{\top}\right) \begin{bmatrix} \tilde{\boldsymbol{V}}_a \\ \boldsymbol{V}_b \end{bmatrix}$$

Definition 2.1: A transformer $M$ is $(N, \varepsilon, r)$-compressible if there exists a KV cache compression policy $\mathsf{C}$ with budget $r$ such that for every pair of sequences $\boldsymbol{a}, \boldsymbol{b}$ with combined length $n + k \leq N$:

$$\left\| \mathrm{M}([\boldsymbol{a}, \boldsymbol{b}]) - \mathrm{M}_{\mathrm{C},\boldsymbol{a}}(\boldsymbol{b}) \right\| < \varepsilon$$

Key Theoretical Result

Theorem 3.1: Let $f$ be a sequence-to-vector function over a finite alphabet. If there exist sequences $\boldsymbol{a}, \boldsymbol{b}^1, \boldsymbol{b}^2$ such that $f([\boldsymbol{a}, \boldsymbol{b}^1]) \neq f([\boldsymbol{a}, \boldsymbol{b}^2])$, then:

1. There exists a transformer that approximates $f$ and is $(N, \varepsilon, 1)$-compressible (compressible to a single KV pair).
2. There exists a transformer of the same architecture that approximates $f$ but is not $(N, C, r)$-compressible for any budget function with $r(n) < n$.

This proves that compressibility varies dramatically across implementations of the same function, motivating compression-aware training.

Motivating Example: Histogram Computation

The histogram function computes the empirical distribution of symbols:

$$f_{\mathrm{hist}}(\boldsymbol{a}) = \left(\frac{n_1}{n}, \dots, \frac{n_m}{n}\right)$$

• Non-compressible implementation (Proposition 3.2): A simple 2-layer transformer with identity first layer and uniform attention ($W_Q = \boldsymbol{0}$) yields incompressible representations. The error is bounded below by a constant for any non-trivial compression.
• Compressible implementation (Proposition 3.3): By structuring token representations to include positional information in auxiliary slots, a compression policy can store the unnormalized histogram plus prefix length in a single KV pair, and the final feedforward network can re-normalize using positional encodings.

This example shows that standard training (which doesn't encourage such structured representations) may naturally favor simpler but non-compressible solutions.

Methodology

KV-Compression Aware Training (KV-CAT)

KV-CAT is a continued pretraining procedure with three key components:

1. Train-time KV Sparsification Policy

• Lightweight learned routers are inserted between groups of consecutive layers.
• At layer $\ell$, the router takes token representations from layer $\ell-1$ and outputs a scalar importance score in $[0,1]$ per token.
• Scores are thresholded with a hyperparameter $\tau$ to produce a binary mask: active tokens participate in attention, masked tokens are dropped.
• Routers are implemented as linear attention modules and initialized with all scores set to 1 (all tokens active), so training starts from the standard dense transformer.
• Routers are shared across groups of consecutive layers to reduce overhead.

2. Training Objective

The total loss combines three terms:

$$\mathcal{L}(\theta) = \lambda_{\mathrm{mask}} \mathcal{L}_{\mathrm{mask}} + \lambda_{\mathrm{budget}} \mathcal{L}_{\mathrm{budget}} + \lambda_{\mathrm{anchor}} \mathcal{L}_{\mathrm{anchor}}$$

• Mask loss (self-distillation): Trains the masked forward pass to match the dense model's distribution:

$$\mathcal{L}_{\mathrm{mask}} = \frac{1}{n} \sum_{i=1}^{n} D_{\mathrm{KL}}\left(\mathrm{sg}[p_{\theta}^{\mathrm{dense}}(\cdot | a_{<i})] \| p_{\theta}^{\mathrm{mask}}(\cdot | a_{<i})\right)$$

• Budget loss: Controls the retention rate, following Hwang et al. [22]:

$$\mathcal{L}_{\mathrm{budget}} = \rho^{-1} F G + (1-\rho)^{-1} (1-F)(1-G), \quad F = \frac{1}{L} \sum_{i=1}^{L} m_i, \quad G = \frac{1}{L} \sum_{i=1}^{L} q_i$$

where $q_i$ and $m_i$ are the router's score and binary mask for token $i$, and $\rho$ is the target retention rate.

• Anchor loss: Standard NTP on the dense forward pass to preserve uncompressed behavior:

$$\mathcal{L}_{\mathrm{anchor}} = \frac{1}{n} \sum_{i=1}^{n} -\log p_{\theta}^{\mathrm{dense}}(a_i | a_{<i})$$

3. Inference

At evaluation time, the unmasked forward pass is used, and standard KV cache compression methods are applied on top. KV-CAT produces a standard transformer whose representations are trained to be more amenable to compression.

Experimental Setup

• Models: QWEN2.5-0.5B and QWEN2.5-1.5B
• Data: FineWeb-Edu, $5.24 \times 10^9$ tokens, max learning rate $10^{-4}$
• Routers: Four learned routers shared across four layer groups
• Compression methods evaluated:
  • Attention Matching [61]: constructs compact prefix cache layer-by-layer, approximating dense attention traces
  • Gradient-based KV cache optimization [16]: directly optimizes compact prefix KV cache to match dense logits on suffix tokens

Empirical Validation / Results

Q1: Preserving Uncompressed Performance

KV-CAT closely matches base model performance on standard benchmarks (Table 1):

Q2: Compression Quality Under Fixed Optimization Budget

Using Attention Matching with 256 query states and 8 NNLS iterations (Table 3):

• QWEN2.5-0.5B: At 5% keep ratio, $\Delta PPL$ improves from 0.780 (base) to 0.461 (KV-CAT); at 40% keep, from 0.592 to 0.360.
• QWEN2.5-1.5B: At 10% keep ratio, $\Delta PPL$ improves from 1.027 to 0.320—a 3.21× improvement in perplexity retention.
• KL divergence and Top-1 agreement consistently improve across all keep ratios and model sizes.

For gradient-based optimization (Figure 2), KV-CAT achieves comparable $\Delta PPL$ in up to 5× fewer optimization steps than the base model.

Q3: Retrieval from Compressed Context

Needle-in-a-haystack experiment (Table 2), with gradient-based compression over 100 optimization steps:

KV-CAT achieves 11–19 point improvements at 30–50% keep ratios, with mean retrieval accuracy gains of 6.4 points (0.5B) and 5.2 points (1.5B).

Q4: Long-Context Question Answering

LongBench v2 evaluation (Table 4) with gradient-based compression:

The largest gains appear in Legal (27.3 → 48.5 at 10% keep) and Table QA (11.1 → 27.8 at 10% keep) subdomains.

Theoretical and Practical Implications

Theoretical Implications

1. Compressibility is a representational property: The paper formally establishes that KV compressibility is determined by the learned internal representations, not just the task or input. This reframes the KV compression problem from purely algorithmic to also a training problem.

2. Existence of compressible implementations: Theorem 3.1 guarantees that for almost any function, a highly compressible transformer implementation exists. This means the ceiling for compression is not fundamentally limited by the task—only by the specific model trained.

3. Training matters: The histogram example demonstrates that standard training objectives (NTP) do not naturally encourage compressible representations. Models may converge to simpler but non-compressible solutions without explicit pressure toward compressibility.

Practical Implications

1. Complementary to post-hoc methods: KV-CAT is orthogonal to compression algorithms—it makes models more amenable to existing compressors, meaning improvements compound with advances in compression methods.

2. Practical efficiency gains: The 5× speedup in optimization time means KV-CAT-trained models can achieve the same compression quality with significantly less compute at inference time, or better quality at the same compute budget.

3. No performance sacrifice: KV-CAT preserves (or slightly improves) dense model performance, making it a safe drop-in continued pretraining procedure.

4. Broad applicability: Gains transfer across retrieval, perplexity-based evaluation, and long-context QA, and across model sizes (0.5B and 1.5B), suggesting the approach scales.

Conclusion

This work establishes that KV cache compressibility is a learnable property of transformer representations, not merely a property of the input or task. The theoretical analysis proves that almost any sequence-to-vector function admits both compressible and non-compressible implementations, highlighting the need for compression-aware training.

KV-CAT addresses this by:

• Introducing train-time KV sparsification with learned routers
• Using self-distillation to preserve dense behavior while exposing the model to compression constraints
• Maintaining anchor NTP loss to prevent degradation of uncompressed performance

Empirical results demonstrate consistent improvements across compression ratios, model sizes, tasks, and compression methods, with gains of up to 3.21× in perplexity retention, 5× in optimization speed, 68% in retrieval accuracy, and 39% in long-context QA.

Future directions identified by the authors include:

• Understanding when worst-case non-compressible behaviors arise in practice
• Exploring more sophisticated train-time sparsification policies
• Investigating whether the approach scales to larger models and other architectures