What if your AI’s memory didn’t have to balloon with every extra sentence?

University of Oxford, Technion, AITHYRA, and NVIDIA introduce KV-Compression Aware Training (KV-CAT) — a method that forces transformers to learn more compressible key-value caches during training, not just after.

By randomly masking KV slots during continued pretraining, KV-CAT nudges models to rely on fewer memory entries, making post-hoc compression far more effective.

Result: KV-CAT beats existing compression methods on retrieval, long-context QA, and perplexity — all while using the same or smaller memory budget.

Training Transformers for KV Cache Compressibility

Paper: https://arxiv.org/abs/2605.05971

Our report: https://mp.weixin.qq.com/s/LDkb9g1Jfq5wf-oC8aL32w…

Training Transformers for KV Cache Compressibility

Source: https://arxiv.org/html/2605.05971 Yoav Gelberg1Yam Eitan211footnotemark:1Michael Bronstein1,3Yarin Gal1Haggai Maron2,4

1University of Oxford2Technion – Israel Institute of Technology3AITHYRA4NVIDIA

(April 2026)

Abstract

Long-context language modeling is increasingly constrained by the Key–Value (KV) cache, whose memory and decode-time access costs scale linearly with the prefix length. This bottleneck has motivated a range of context-compression methods, from token-level summarization to recent optimization-based KV cache compression methods. These post-hoc methods operate on the KV cache of a fixed pretrained model, so their effectiveness is fundamentally limited by how well the model’s internal representations can be compressed. In this work, we formalize the notion of KV compressibility and show that it is a property of thelearned representations, rather than of the context alone. We prove that almost any sequence-to-vector function admits both highly compressible and inherently non-compressible transformer implementations, highlighting the need to guide transformers toward compressible representations during training. Motivated by this, we proposeKV-CompressionAwareTraining (KV-CAT), a continued pretraining procedure that incentivizes the emergence of compressible representations. We introduce a train-time KV sparsification policy that masks KV slots during training. This forces the model to use fewer KV slots and encourages it to learn representations amenable to post-hoc compression. Empirically, we show thatKV-CATimproves the quality–budget tradeoff of downstream compression methods across retrieval, long-context question answering, and perplexity-based evaluation of compressed-prefix continuation.

1Introduction

Language models (LMs) are increasingly deployed in long-horizon settings, from understanding large codebases, long-form documents, and personal data repositories[41;23;1], to long-form reasoning and continuously learning agents[8;60;35]. Using autoregressive transformer-based LMs in these settings introduces a significant memory bottleneck: the Key–Value (KV) cache. During inference, these models must store the key and value vectors for every token, at every layer and (KV-)attention head. For long sequences, this cache can dominate both memory usage and decoding cost, turning context length into a primary serving bottleneck[59;34;52].

A substantial body of recent work has sought to mitigate this bottleneck, broadly falling into two categories. The first category focuses on designing more efficient alternatives to transformers. Examples include linear attention mechanisms[29;11;20;54], state space models[18;17], sparse attention variants[3;57], and more. These methods significantly reduce the computational cost of long-context language modeling. However, this efficiency typically comes at the cost of empirical performance, and such models still lag behind transformers at scale.

The second category focuses on inference-time interventions applied to a fixed pretrained transformer. These methods either operate on the input context or, more generally, directly on the KV cache. Early approaches are largely heuristic, including textual summarization[46;42], learned token filtering[24], and policies based on attention patterns, recency, heavy-hitter behavior, or layer-wise importance[59;52;34;6]. More recently, optimization-based KV cache compression111These methods are also commonly referred to as KV compaction methods in the literature.methods have emerged as powerful techniques. For example,Eyuboglu et al. [16]use gradient-based optimization to match the distribution induced by the original cache, whileZweiger et al. [61]employ layer-wise objectives to reproduce its attention traces.

The success of KV cache compression suggests that the full KV cache often contains redundancies. However, as most approaches operate on afixed model, they are inherently limited by how well that particular model’s representations can be compressed. Crucially, this compressibility is not determined by the input sequence alone: two transformers can use very different internal representations when processing the same sequence, yet have identical next-token distributions. Consequently, some transformers’ KV caches may be more amenable to compression than others.

This work.In this paper, we study the following question: can transformer LMs be trained in a way that leads to KV caches that are more amenable to post-hoc compression? This shifts the target from the KV cache compression algorithm to the model itself. To address this question, we introduce the notion ofKV-compressibility. Informally, a transformer is KV-compressible if there exists a compression policy that maps the KV cache of long input sequences to a shorter KV cache while preserving the model’s next token distribution.

a1a_{1}a2a_{2}a3a_{3}a4a_{4}a5a_{5}a6a_{6}a7a_{7}a8a_{8}masked\mathrm{masked}dense\mathrm{dense}𝒌1,𝒗1{\bm{k}}_{1}{,}{\bm{v}}_{1}𝒌2,𝒗2{\bm{k}}_{2}{,}{\bm{v}}_{2}𝒌3,𝒗3{\bm{k}}_{3}{,}{\bm{v}}_{3}𝒌4,𝒗4{\bm{k}}_{4}{,}{\bm{v}}_{4}𝒌5,𝒗5{\bm{k}}_{5}{,}{\bm{v}}_{5}𝒌6,𝒗6{\bm{k}}_{6}{,}{\bm{v}}_{6}𝒌7,𝒗7{\bm{k}}_{7}{,}{\bm{v}}_{7}𝒌8,𝒗8{\bm{k}}_{8}{,}{\bm{v}}_{8}𝒌1,𝒗1{\bm{k}}_{1}{,}{\bm{v}}_{1}𝒌2,𝒗2{\bm{k}}_{2}{,}{\bm{v}}_{2}𝒌3,𝒗3{\bm{k}}_{3}{,}{\bm{v}}_{3}𝒌4,𝒗4{\bm{k}}_{4}{,}{\bm{v}}_{4}𝒌5,𝒗5{\bm{k}}_{5}{,}{\bm{v}}_{5}𝒌6,𝒗6{\bm{k}}_{6}{,}{\bm{v}}_{6}𝒌7,𝒗7{\bm{k}}_{7}{,}{\bm{v}}_{7}𝒌8,𝒗8{\bm{k}}_{8}{,}{\bm{v}}_{8}𝒌1,𝒗1{\bm{k}}_{1}{,}{\bm{v}}_{1}𝒌2,𝒗2{\bm{k}}_{2}{,}{\bm{v}}_{2}𝒌3,𝒗3{\bm{k}}_{3}{,}{\bm{v}}_{3}𝒌4,𝒗4{\bm{k}}_{4}{,}{\bm{v}}_{4}𝒌5,𝒗5{\bm{k}}_{5}{,}{\bm{v}}_{5}𝒌6,𝒗6{\bm{k}}_{6}{,}{\bm{v}}_{6}𝒌7,𝒗7{\bm{k}}_{7}{,}{\bm{v}}_{7}𝒌8,𝒗8{\bm{k}}_{8}{,}{\bm{v}}_{8}𝒌1,𝒗1{\bm{k}}_{1}{,}{\bm{v}}_{1}𝒌2,𝒗2{\bm{k}}_{2}{,}{\bm{v}}_{2}𝒌3,𝒗3{\bm{k}}_{3}{,}{\bm{v}}_{3}𝒌4,𝒗4{\bm{k}}_{4}{,}{\bm{v}}_{4}𝒌5,𝒗5{\bm{k}}_{5}{,}{\bm{v}}_{5}𝒌6,𝒗6{\bm{k}}_{6}{,}{\bm{v}}_{6}𝒌7,𝒗7{\bm{k}}_{7}{,}{\bm{v}}_{7}𝒌8,𝒗8{\bm{k}}_{8}{,}{\bm{v}}_{8}𝒌1,𝒗1{\bm{k}}_{1}{,}{\bm{v}}_{1}𝒌1,𝒗1{\bm{k}}_{1}{,}{\bm{v}}_{1}𝒌1,𝒗1{\bm{k}}_{1}{,}{\bm{v}}_{1}𝒌1,𝒗1{\bm{k}}_{1}{,}{\bm{v}}_{1}𝒌2,𝒗2{\bm{k}}_{2}{,}{\bm{v}}_{2}𝒌2,𝒗2{\bm{k}}_{2}{,}{\bm{v}}_{2}𝒌2,𝒗2{\bm{k}}_{2}{,}{\bm{v}}_{2}𝒌2,𝒗2{\bm{k}}_{2}{,}{\bm{v}}_{2}𝒌3,𝒗3{\bm{k}}_{3}{,}{\bm{v}}_{3}𝒌3,𝒗3{\bm{k}}_{3}{,}{\bm{v}}_{3}𝒌3,𝒗3{\bm{k}}_{3}{,}{\bm{v}}_{3}𝒌3,𝒗3{\bm{k}}_{3}{,}{\bm{v}}_{3}𝒌4,𝒗4{\bm{k}}_{4}{,}{\bm{v}}_{4}𝒌4,𝒗4{\bm{k}}_{4}{,}{\bm{v}}_{4}𝒌4,𝒗4{\bm{k}}_{4}{,}{\bm{v}}_{4}𝒌4,𝒗4{\bm{k}}_{4}{,}{\bm{v}}_{4}𝒌5,𝒗5{\bm{k}}_{5}{,}{\bm{v}}_{5}𝒌5,𝒗5{\bm{k}}_{5}{,}{\bm{v}}_{5}𝒌5,𝒗5{\bm{k}}_{5}{,}{\bm{v}}_{5}𝒌5,𝒗5{\bm{k}}_{5}{,}{\bm{v}}_{5}𝒌6,𝒗6{\bm{k}}_{6}{,}{\bm{v}}_{6}𝒌6,𝒗6{\bm{k}}_{6}{,}{\bm{v}}_{6}𝒌6,𝒗6{\bm{k}}_{6}{,}{\bm{v}}_{6}𝒌6,𝒗6{\bm{k}}_{6}{,}{\bm{v}}_{6}𝒌7,𝒗7{\bm{k}}_{7}{,}{\bm{v}}_{7}𝒌7,𝒗7{\bm{k}}_{7}{,}{\bm{v}}_{7}𝒌7,𝒗7{\bm{k}}_{7}{,}{\bm{v}}_{7}𝒌7,𝒗7{\bm{k}}_{7}{,}{\bm{v}}_{7}𝒌8,𝒗8{\bm{k}}_{8}{,}{\bm{v}}_{8}𝒌8,𝒗8{\bm{k}}_{8}{,}{\bm{v}}_{8}𝒌8,𝒗8{\bm{k}}_{8}{,}{\bm{v}}_{8}𝒌8,𝒗8{\bm{k}}_{8}{,}{\bm{v}}_{8}ℒmask\mathcal{L}_{\mathrm{mask}}ℒbudget\mathcal{L}_{\mathrm{budget}}ℒanchor\mathcal{L}_{\mathrm{anchor}}joint update:modelθ\theta+ routers

Figure 1:KV-Compression Aware Training (KV-CAT).A context𝒂{\bm{a}}goes through both masked (left) and dense (right) forward passes. In the masked forward pass, routers (orange) compute masks for groups of consecutive layers, marking KV slots as active (green) or inactive (muted green). In the dense forward pass (blue), all KV slots are kept. The output of the masked forward pass is used to computeℒmask{\mathcal{L}}_{\mathrm{mask}}, router distributions are used to computeℒbudget{\mathcal{L}}_{\mathrm{budget}}, and the outputs of the dense forward pass are used to computeℒanchor{\mathcal{L}}_{\mathrm{anchor}}. These are jointly used to update the parameters.Our theoretical results show that for almost any sequence-to-vector function, there exist transformer implementations whose prefix can be compressed to a single KV pair, as well as implementations for which any non-trivial compression incurs a constant error. To build intuition, we consider a motivating example: character histogram computation. We show that natural transformer implementations of this computation may produce incompressible token representations, whereas more structured alternatives are highly compressible. This motivates treating KV compressibility as an explicit training objective.

Guided by this perspective, we proposeKV-CompressionAwareTraining (KV-CAT), a continued pretraining (CPT) procedure that promotes the emergence of KV-compressible internal representations. Starting from a pretrained transformer,KV-CATintroduces a train-time KV sparsification policy which masks out a constant fraction of the KV slots. The training objective combines a self-distillation loss, which matches the masked model’s distribution to the dense model’s distribution, and an NTP loss applied to the unmasked forward pass to preserve uncompressed model behavior. This exposes the model to the information bottleneck induced by KV cache compression during training, encouraging it to reorganize its representations into a more compressible form while maintaining performance in the uncompressed setting.

To demonstrate thatKV-CATimproves the quality–budget tradeoff of post-hoc KV cache compression, we apply it toQwen2.5models[51], and evaluate state-of-the-art optimization-based compression methods on the KV caches of the resulting checkpoints. We measure performance along three axes: (i) suffix perplexity under prefix compression, (ii) retrieval accuracy from a compressed prefix, and (iii) compressed long-context question answering on LongBench v2[2]tasks. Across a range of compression budgets, model sizes, and compression methods, models trained withKV-CATconsistently achieve a better quality–budget tradeoff than the base model, yielding improvements of up to3.21×\timesin suffix perplexity retention,5×\timesin optimization time,68%in retrieval accuracy, and39%in long-context QA. Importantly, our goal is not to replace post-hoc methods, but to make models more amenable to them, enabling more effective compression.

Contributions.In summary, this paper makes the following contributions:

1. (1)We characterize KV-compressibility as a property of the learned transformer representation, rather than solely of the task or the input sequence.
2. (2)We provide theoretical results showing that almost all sequence-to-vector functions can admit both compressible and non-compressible transformer implementations.
3. (3)We proposeKV-CAT, a training procedure that promotes compressible KV representations.
4. (4)We show thatKV-CATimproves post-hoc KV cache compression at matched cache and optimization budgets across: suffix perplexity retention, retrieval, and long-context QA.

2KV Cache Compression: Problem Formulation

In this section, we theoretically formalize the KV cache compression problem. AKV cache compression policyfor a transformer withLLlayers is a collection of functions𝖢=(𝖼1,…,𝖼L),\mathsf{C}=(\mathsf{c}_{1},\dots,\mathsf{c}_{L}),where each𝖼ℓ\mathsf{c}_{\ell}maps sequences ofnnkey–value pairs to sequences of lengthr​(n)≤nr(n)\leq n. For KV pairs in layerℓ\ell

(𝑲,𝑽)=([𝒌1,…,𝒌n]⊤,[𝒗1,…,𝒗n]⊤),({\bm{K}},{\bm{V}})=([{\bm{k}}_{1},\dots,{\bm{k}}_{n}]^{\top},[{\bm{v}}_{1},\dots,{\bm{v}}_{n}]^{\top}),(1)we write

𝖼ℓ​(𝑲,𝑽)=([𝒌~1,…,𝒌~r​(n)]⊤,[𝒗~1,…,𝒗~r​(n)]⊤),\mathsf{c}_{\ell}({\bm{K}},{\bm{V}})=([\tilde{\bm{k}}_{1},\dots,\tilde{\bm{k}}_{r(n)}]^{\top},[\tilde{\bm{v}}_{1},\dots,\tilde{\bm{v}}_{r(n)}]^{\top}),(2) wherer​(n)r(n)is referred to as thecompression budget. Given a transformer𝖬\mathsf{M}, a compression policy𝖢\mathsf{C}, and a sequence of tokens𝒂=(a1,…,an){\bm{a}}=(a_{1},\dots,a_{n}), we define a compressed model𝖬𝖢,𝒂\mathsf{M}_{\mathsf{C},{\bm{a}}}that shares the parameters of𝖬\mathsf{M}but uses a modified forward pass. For an input sequence𝒃{\bm{b}}, the computation of𝖬𝖢,𝒂​(𝒃)\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})proceeds as in𝖬​([𝒂,𝒃])\mathsf{M}([{\bm{a}},{\bm{b}}]), except for the attention over prefix tokens, which is augmented the following way: at layerℓ\ell, let(𝑲𝒂,𝑽𝒂)({\bm{K}}_{\bm{a}},{\bm{V}}_{\bm{a}})denote the KV cache of𝒂{\bm{a}}, based on the original model𝖬\mathsf{M}and let𝒀{\bm{Y}}be the representations of𝒃{\bm{b}}after layerℓ−1\ell-1of𝖬𝖢,𝒂\mathsf{M}_{\mathsf{C},{\bm{a}}}. We compute

𝑸𝒃=𝒀​𝑾Q,𝑲𝒃=𝒀​𝑾K,𝑽𝒃=𝒀​𝑾V,{\bm{Q}}_{\bm{b}}={\bm{Y}}{\bm{W}}_{Q},\quad{\bm{K}}_{\bm{b}}={\bm{Y}}{\bm{W}}_{K},\quad{\bm{V}}_{\bm{b}}={\bm{Y}}{\bm{W}}_{V},(3)and replace the prefix KV cache(𝑲𝒂,𝑽𝒂)({\bm{K}}_{\bm{a}},{\bm{V}}_{\bm{a}})with its compressed version(𝑲~𝒂,𝑽~𝒂)=𝖼ℓ​(𝑲𝒂,𝑽𝒂)(\tilde{\bm{K}}_{\bm{a}},\tilde{\bm{V}}_{\bm{a}})=\mathsf{c}_{\ell}({\bm{K}}_{\bm{a}},{\bm{V}}_{\bm{a}}). The attention computation then becomes

AttnHead​(𝒀)=softmax​(1dk​𝑸𝒃​[𝑲~𝒂𝑲𝒃]⊤)​[𝑽~𝒂𝑽𝒃].\mathrm{AttnHead}({\bm{Y}})=\mathrm{softmax}\!\left(\frac{1}{\sqrt{d_{k}}}{\bm{Q}}_{\bm{b}}\begin{bmatrix}\tilde{\bm{K}}_{\bm{a}}\\ {\bm{K}}_{\bm{b}}\end{bmatrix}^{\!\top}\right)\begin{bmatrix}\tilde{\bm{V}}_{\bm{a}}\\ {\bm{V}}_{\bm{b}}\end{bmatrix}.(4) Finally, following standard convention in next-token prediction transformers, we define the output of𝖬𝖢,𝒂​(𝒃)\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})to be the representation of the final token. Thus, compression modifies only the prefix KV pairs, leaving the remainder of the computation unchanged. We now formalize the notion of transformerKV-compressibility.

Definition 2.1.

FixN∈ℕN\in{\mathbb{N}},ε>0\varepsilon>0, and letr:ℕ→ℕr:{\mathbb{N}}\to{\mathbb{N}}be a budget function. A transformer𝖬\mathsf{M}is said to be(N,ε,r)(N,\varepsilon,r)-compressible222We includeNNexplicitly, as many natural sequence-to-vector functions require model dimension scaling withNN(e.g.,O​(N)O(N)orO​(log⁡N)O(\log N)) in order to achieve arbitrarily accurate approximation; see e.g.,Sanford et al. [48]; Yehudai et al. [55].if there exists a KV cache compression policy𝖢\mathsf{C}with budgetrrsuch that for every pair of sequences𝒂{\bm{a}},𝒃{\bm{b}}of lengthsnnandkkrespectively, and with a combined lengthn+k≤Nn+k\leq N, it holds that

‖𝖬​([𝒂,𝒃])−𝖬𝖢,𝒂​(𝒃)‖<ε.\|\mathsf{M}([{\bm{a}},{\bm{b}}])-\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})\|<\varepsilon.(5)

A more formal treatment of these terms and definitions is provided in AppendixA.1.

3Motivation and Theoretical Results

We begin with a theoretical analysis of transformer KV-compressibility, motivated by the following central question:

Can the same sequence-to-vector function admit transformer implementations with different levels of KV-compressibility?

The following theorem shows that for almost any sequence-to-vector function, there exist transformer architectures that admit both fully compressible and inherently non-compressible implementations. The proof is deferred to AppendixA.3.

Theorem 3.1.

LetAAbe a finite alphabet, and letf:⋃n≤NAn→ℝdoutf:\bigcup_{n\leq N}A^{n}\to\mathbb{R}^{d_{\text{out}}}be a sequence-to-vector function. Suppose that there exists a sequence𝐚{\bm{a}}of lengthnnand two sequences𝐛1,𝐛2{\bm{b}}^{1},{\bm{b}}^{2}of lengthk≤N−nk\leq N-nsuch that

f​([𝒂,𝒃1])≠f​([𝒂,𝒃2]).f([{\bm{a}},{\bm{b}}^{1}])\neq f([{\bm{a}},{\bm{b}}^{2}]).(6)Then for everyε,C>0\varepsilon,C>0:

1. 1.There exists a transformer that approximatesffand is(N,ε,1)(N,\varepsilon,1)-compressible.
2. 2.There exists a transformer of the same architecture, that approximatesffbut is not(N,C,r)(N,C,r)-compressible for any budget function satisfyingr​(n)<nr(n)<n.

We note that, as KV cache compression methods perform well in practice, the worst-case behavior implied by this theorem may not typically arise; understanding in what scenarios this occurs is an interesting direction for future work. Nevertheless, the result highlights that compressibility can vary across transformer implementations, suggesting the potential benefit of guiding transformers toward more compressible representations at train time. Before turning to training methods, we build intuition through a simple motivating example. We show that even for natural sequence-to-vector functions, simple transformer implementations can yield non-compressible representations, while more structured, highly compressible implementations exist.

Motivating example: histogram computation.LetA=[m]A=[m]be a finite alphabet. Given a sequence𝒂=(a1,…,an)∈An{\bm{a}}=(a_{1},\dots,a_{n})\in A^{n}, the histogram function computes the empirical distribution of symbols,

fhist​(𝒂)=(n1n,…,nmn),f_{\mathrm{hist}}({\bm{a}})=\left(\tfrac{n_{1}}{n},\dots,\tfrac{n_{m}}{n}\right),(7)whereni=|{j:aj=i}|n_{i}=|\{j:a_{j}=i\}|. This is a natural sequence-to-vector function that depends only on aggregate statistics of the input. A straightforward way to implement this function with a 2-layer transformer is as follows. First, take the embedding map to be333We typically assume token embeddings of the formemb​(ai,i)=𝒖ai+𝒑i\mathrm{emb}(a_{i},i)={\bm{u}}_{a_{i}}+{\bm{p}}_{i}. The embedding map used above can be obtained from this representation by a minor adjustment to the first attention layer; see AppendixA.2for details.emb​(ai,i)=𝒆ai\mathrm{emb}(a_{i},i)={\bm{e}}_{a_{i}}. We take the first transformer layer to be the identity444This can be done by setting𝑾O=𝟎{\bm{W}}_{O}=\mathbf{0}and choosing the feedforward network to be the identity., and then apply a second attention layer with uniform attention weights (by setting𝑾Q=𝟎{\bm{W}}_{Q}=\mathbf{0}), So that the representation of the final token becomes the average of all previous token embeddings, obtaining

𝖬​(𝒂)=1n​∑i=1n𝒆ai=fhist​(𝒂).\mathsf{M}({\bm{a}})=\frac{1}{n}\sum_{i=1}^{n}{\bm{e}}_{a_{i}}=f_{\mathrm{hist}}({\bm{a}}).(8)This gives a natural implementation of the histogram function with a very simple architecture (see AppendixA.2for details). Interestingly, we find that this implementation is not compressible:

Proposition 3.2.

For anyN∈ℕN\in{\mathbb{N}}and any budget functionrrwithr​(N−1)<N−1r(N-1)<N-1there exists a constantC>0C>0such that the above transformer implementation offhistf_{\mathrm{hist}}is not(N,C,r)(N,C,r)-compressible.

Proof sketch..

Recall that in the compressed setting the input is[𝒂,𝒃][{\bm{a}},{\bm{b}}]and we want to compress the KV caches associated with the prefix𝒂{\bm{a}}. We begin by observing that since in the first layer𝑾O=𝟎{\bm{W}}_{O}=\mathbf{0}, compressing the prefix of the first layer has no effect on the representation of the tokens of𝒃{\bm{b}}. Additionally, since in the second layer𝑾Q=𝟎{\bm{W}}_{Q}=\mathbf{0}, the attention scores are independent of the keys, and therefore the attention weights are uniform across all tokens. Thus, for a prefix sequence𝒂{\bm{a}}of lengthnnand a suffix𝒃{\bm{b}}of lengthkk, the compressed model satisfies, for any compression policy𝖢\mathsf{C},

𝖬𝖢,𝒂​(𝒃)=1r​(n)+k​(∑i=1r​(n)𝒗~i+∑i=1k𝒆bi).\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})=\frac{1}{r(n)+k}\Big(\sum_{i=1}^{r(n)}\tilde{\bm{v}}_{i}+\sum_{i=1}^{k}{\bm{e}}_{b_{i}}\Big).(9)Similarly, the full model computes

𝖬​([𝒂,𝒃])=1n+k​(∑i=1n𝒆ai+∑i=1k𝒆bi).\mathsf{M}([{\bm{a}},{\bm{b}}])=\frac{1}{n+k}\Big(\sum_{i=1}^{n}{\bm{e}}_{a_{i}}+\sum_{i=1}^{k}{\bm{e}}_{b_{i}}\Big).(10)For notational convenience, define

𝒗~=1r​(n)+k​∑i=1r​(n)𝒗~i,𝒂~=1n+k​∑i=1n𝒆ai,𝒃~=(1r​(n)+k−1n+k)​∑i=1k𝒆bi.\tilde{\bm{v}}=\frac{1}{r(n)+k}\sum_{i=1}^{r(n)}\tilde{\bm{v}}_{i},\quad\tilde{\bm{a}}=\frac{1}{n+k}\sum_{i=1}^{n}{\bm{e}}_{a_{i}},\quad\tilde{\bm{b}}=\Big(\frac{1}{r(n)+k}-\frac{1}{n+k}\Big)\sum_{i=1}^{k}{\bm{e}}_{b_{i}}.(11) Then the difference between the compressed and full outputs can be written as

‖𝖬𝖢,𝒂​(𝒃)−𝖬​([𝒂,𝒃])‖=‖𝒗~−𝒂~+𝒃~‖.\|\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})-\mathsf{M}([{\bm{a}},{\bm{b}}])\|=\left\|\tilde{\bm{v}}-\tilde{\bm{a}}+\tilde{\bm{b}}\right\|.(12)Both𝒗~\tilde{\bm{v}}and𝒂~\tilde{\bm{a}}depend only on the prefix𝒂{\bm{a}}, while𝒃~\tilde{\bm{b}}depends on the suffix𝒃{\bm{b}}and can vary over a range of values. For𝖬𝖢,𝒂​(𝒃)\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})toε\varepsilon-approximate𝖬​([𝒂,𝒃])\mathsf{M}([{\bm{a}},{\bm{b}}])uniformly over all suffixes𝒃{\bm{b}}, Equation12requires that the fixed term𝒗~−𝒂~\tilde{\bm{v}}-\tilde{\bm{a}}beε\varepsilon-close to every possible value of𝒃~\tilde{\bm{b}}. For sufficiently smallε\varepsilon, this is impossible. Hence, the error is bounded below by a constant.

∎

For a formal proof of Proposition3.2, see AppendixA.2. We now show that a slight modification of the above implementation yields highly compressible KV representations.

Proposition 3.3.

For everyN∈ℕN\in{\mathbb{N}}andε>0\varepsilon>0, there exists afhistf_{\mathrm{hist}}transformer implementation with the same architecture as above that is(N,ε,r)(N,\varepsilon,r)-compressible withr​(n)=1r(n)=1.

Proof sketch.

We assume an injective token embedding of the formemb​(ai,i)=𝒖ai+𝒑i\mathrm{emb}(a_{i},i)={\bm{u}}_{a_{i}}+{\bm{p}}_{i}where𝒑i{\bm{p}}_{i}are positional embeddings. As in the previous construction, we set𝑾O=𝟎{\bm{W}}_{O}=\mathbf{0}in the first transformer layer, but now choose the feedforward network to approximately satisfy

ρ1​(𝒖ai+𝒑i)=[𝒆ai⊤,𝒑i⊤,𝟎⊤,𝟎⊤]⊤\rho_{1}({\bm{u}}_{a_{i}}+{\bm{p}}_{i})=[{\bm{e}}_{a_{i}}^{\top},{\bm{p}}_{i}^{\top},\mathbf{0}^{\top},\mathbf{0}^{\top}]^{\top}(13)That is, each token representation is partitioned into four blocks: the first encodes token identity, the second retains positional information, and the last two are auxiliary slots used later in the construction. We then apply a second attention layer with𝑾Q=𝟎{\bm{W}}_{Q}=\mathbf{0}, so attention is uniform. With the residual connection, the final token representation can be made to contain the empirical average in the third block while preserving the positional encoding𝒑n{\bm{p}}_{n}, which encodes the sequence length. The final feedforward networkρ2\rho_{2}is designed to behave differently depending on the available information. In the uncompressed setting, the value of the final block remains𝟎\mathbf{0}, in which case it simply returns the empirical average (from the third slot). In the compressed setting, however, an optimal compression policy can provide additional positional information in the final vector block, which can then be used to correct for the distortion introduced by compression. Specifically, we define a compression policy that maps the prefix to a single KV pair, whose value stores theunnormalizedhistogram together with the prefix length (encoded via𝒑n{\bm{p}}_{n}in the last block). When processing a suffix, attention averages this summary with the suffix tokens, producing a mis-scaled estimate of the histogram. Crucially, the model now has access to both the original prefix length and the total sequence length (via positional encodings). The final feedforward network uses this information to re-normalize the average and recover the correct histogram. This establishes(N,ε,r)(N,\varepsilon,r)-compressibility withr​(n)=1r(n)=1.

∎

For a formal proof of Proposition3.3, see AppendixA.2. Note that while incorporating positional information to correct attention weights under compression is conceptually straightforward, a standard transformer trained only on full (uncompressed) sequences is not explicitly encouraged to learn such a mechanism. Since computing the histogram depends only on token identities and is invariant to their positions, the model can, in principle, learn to ignore positional information. This suggests that obtaining compressible solutions may benefit from a training procedure that encourages compression.

4Training for KV Cache Compressibility

Table 1:KV-CATmatches base model performance.We test ourKV-CAT-trainedQwen2.5checkpoints on a variety of QA benchmarks without any KV cache compression applied. Values are normalized multiple-choice accuracy (%) with 1000 examples per task.As shown in the previous section, training transformers to solve sequence-to-vector tasks may admit both compressible and non-compressible solutions. This naturally raises the following question:

Can we guide transformers toward more compressible representations during training?

To answer this question, we introduce KV-Compression Aware Training (KV-CAT), a continued pretraining procedure that explicitly encourages compressible internal representations.

Train-time KV sparsification policy.Our goal is to train transformers whose representations are more amenable to SOTA optimization-based KV cache compression methods. Ideally, one would incorporate such methods directly into the training loop, adapting the compression policy as the model evolves; however, this is computationally prohibitive. A practical alternative is to introduce a simple and scalable training-time KV sparsification policy, such as random KV slot dropping, attention-based filtering, or a lightweight parameterized policy that adapts jointly with the model. In our ablations (AppendixD.1), we find that a parameterized policy provides the strongest performance, in terms of both downstream compression and retention of uncompressed accuracy. We hypothesize that this is because adaptive policies better align with the test-time behavior of optimization-based compressors, whereas fixed policies may be too rigid. Accordingly, in our experiments, we implementKV-CATusing lightweight learned routers, described below, while emphasizing that the framework is general and can accommodate a range of sparsification policies.

Routing mechanism.Starting from a pretrained decoder-only transformer, we insert learnedroutersbetween consecutive layers and train them jointly with the model. At layerℓ\ell, the router takes as input the token representations from layerℓ−1\ell-1and outputs a scalar score in[0,1][0,1]for each token, indicating its importance. These scores are thresholded using a (non-trainable) hyperparameterτ\tauto produce a binary mask: tokens with scores aboveτ\tauare active in attention, while the rest are masked out. For efficiency, routers are implemented as linear attention modules. Routers are initialized with all tokens active (i.e., all scores are set to11), so training begins from the standard dense transformer and gradually learns to mask tokens, using a budget loss (defined below). To further reduce overhead, routers are shared across groups of consecutive layers. See AppendixBfor additional details.

Training objective.We train the transformer using an augmented next-token prediction (NTP) objective. For a sequence𝒂=(a1,…,an){\bm{a}}=(a_{1},\dots,a_{n}), letpθmask(⋅∣𝒂<i)p_{\theta}^{\mathrm{mask}}(\cdot\mid{\bm{a}}_{<i})denote the model distribution when masking is applied to the KV cache, andpθdense(⋅∣𝒂<i)p_{\theta}^{\mathrm{dense}}(\cdot\mid{\bm{a}}_{<i})denote the distribution under the standard (unmasked) forward pass. During training, we optimize

ℒ​(θ)=λmask​ℒmask+λbudget​ℒbudget+λanchor​ℒanchor.\mathcal{L}(\theta)=\lambda_{\mathrm{mask}}\mathcal{L}_{\mathrm{mask}}+\lambda_{\mathrm{budget}}\mathcal{L}_{\mathrm{budget}}+\lambda_{\mathrm{anchor}}\mathcal{L}_{\mathrm{anchor}}.(14)The first term is given by

ℒmask=1n∑i=1nDKL(sg[pθdense(⋅∣a<i)]∥pθmask(⋅∣a<i)),\mathcal{L}_{\mathrm{mask}}=\frac{1}{n}\sum_{i=1}^{n}D_{\mathrm{KL}}\!\left(\operatorname{sg}[p_{\theta}^{\mathrm{dense}}(\cdot\mid a_{<i})]\,\middle\|\,p_{\theta}^{\mathrm{mask}}(\cdot\mid a_{<i})\right),(15)wheresg\operatorname{sg}denotes the stop-gradient operation. This term trains the masked forward pass to match the dense model distribution via self-distillation. When the KV-sparsification policy is learnable, we include a budget loss designed to maintain a desirable retention rate. FollowingHwang et al. [22], in our implementation, we use:

ℒbudget=ρ−1​F​G+(1−ρ)−1​(1−F)​(1−G),F=1L​∑i=1Lmi,G=1L​∑i=1Lqi,\mathcal{L}_{\mathrm{budget}}=\rho^{-1}FG+(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},(16)whereqiq_{i}andmim_{i}denote the router’s score and binary mask for tokenii, respectively, andρ\rhois the target retention rate. The loss is averaged across router layers, and gradients are taken with respect to the scoresqiq_{i}only. The final term

ℒanchor=1n​∑i=1n−log⁡pθdense​(ai∣a<i),\mathcal{L}_{\mathrm{anchor}}=\frac{1}{n}\sum_{i=1}^{n}-\log p_{\theta}^{\mathrm{dense}}(a_{i}\mid a_{<i}),(17)applies standard NTP to the output of the dense forward pass, preserving its distribution and providing a stable teacher for distillation.555We note that a natural alternative to our training objective is to use distillation from a frozen base model for both the masked and dense forward passes. We do not use this variant since it requires keeping an additional frozen teacher model in memory and running a third forward pass, significantly increasing training cost.See Figure1for an illustration of this training procedure.

Inference.At evaluation time, we use the unmasked forward pass, on top of which we can apply standard KV cache compression methods. Thus, the output ofKV-CATis a standard transformer whose representations are trained to be more amenable to compression.

5Empirical Evaluation

Figure 2:KV-CATspeeds up gradient-based KV cache compression.We plot the gap in suffix perplexity under full/compressed-prefix inference throughout gradient-based KV cache optimization. Each panel fixes a different KV keep ratio. Across ratios, theKV-CATcheckpoint achieves a comparableΔ\DeltaPPL in fewer optimization steps than the base model, yielding up to a5×\timesspeedup.Table 2:Needle retrieval from compressed haystack.Exact-match accuracy (%) over 100 examples per keep ratio.We evaluate the effect ofKV-CATon the performance of downstream optimization-based KV cache compression techniques. Our empirical study is guided by four central questions:(Q1)DoesKV-CATtraining preserve base model performance in uncompressed settings?(Q2)DoesKV-CATimprove compression quality under a fixed KV-optimization budget?(Q3)DoesKV-CATimprove retrieval from a KV-compressed context?(Q4)Do these gains transfer to KV-compressed long-context question answering? Across all experiments, we compare the original pretrained model to the model obtained viaKV-CATcontinued pretraining, applying the same downstream compression method at matched KV retention and optimization budgets. We provide full experimental details in AppendixC, and additional results in AppendixD.

KV-CATtraining setup.We applyKV-CATtoQwen2.5-0.5BandQwen2.5-1.5Bcheckpoints via continued pretraining on FineWeb-Edu[44]. In both cases, we train all model parameters, and introduce four learned routers shared across four layer groups. Each model is trained on a total of5.24×1095.24\times 10^{9}tokens with a max learning rate of10−410^{-4}. See further details in AppendixC.1.

Table 3:Suffix completion under prefix compression.Each example uses a 768-token prefix and a 256-token suffix. Prefix KVs are compressed using Attention Matching. We report: (1)Δ\DeltaPPL, the difference in suffix perplexity under the full prefix vs. compressed prefix; (2) KL, the divergence between the two token distributions; and (3) Top-1, percentage of tokens on which the two distributions’ top-1 agree.KV cache compression methods.Across experiments, we use two post-hoc KV cache compression procedures: Attention Matching[61]and a gradient-based KV cache optimization method adapted fromEyuboglu et al. [16]. Attention Matching constructs a compact prefix cache layer-by-layer, approximating the dense model’s attention traces. The gradient-based method instead directly optimizes the compact prefix KV cache to match the dense model’s logits on suffix tokens. Both methods compress the KV cache of a given prefix<text>by leveraging supervision from a suffix, which together form areconstruction sequence. For the retrieval and QA evaluations, we followZweiger et al. [61], and use a reconstruction sequence of the form<text><instruction><text>, where<instruction>is set to “Reproduce the preceding passage verbatim.”, and the loss is applied only on the second repeated<text>span. Additional implementation details are given in AppendixC.

Preserving uncompressed performance (Q1).We evaluate the effect ofKV-CATon the performance of the modelwith no KV cache compression applied. We compare the base model and theKV-CAT-trained checkpoint on six standard multiple-choice benchmarks using the LightEval harness[19]. As shown in Table1,KV-CATclosely matches the base model, with an average gain of 0.7 accuracy points forQwen2.5-0.5Band a drop of 0.5 points forQwen2.5-1.5B. These results indicate thatKV-CATpreserves standard dense behavior rather than trading it off for compressibility.

Compression under fixed KV-optimization budget (Q2).We evaluate compression quality at different keep ratios, under a fixed KV-optimization budget. To do this, we sample held–out prefix–suffix pairs from FineWeb and directly optimize a compacted prefix KV cache to produce the corresponding suffix. We optimize the KV slots via Attention Matching using 256 query states and 8 NNLS iterations across all keep ratios. We report three different similarity metrics between the suffix next token distribution under the full/compressed prefix, averaged over 128 pairs per keep ratio. As seen in Table3, under a fixed optimization budget, theKV-CAT-trained model consistently outperforms the base model across compression ratios, model sizes, and similarity measures, and achieves up to3.21×\timesbetter retention of suffix perplexity. We additionally evaluate gradient-based KV cache optimization on the same prefix–suffix pairs. As seen in Figure2, theKV-CAT-trained model achieves better suffix perplexity retention throughout optimization, requiring up to5×\timesfewer optimization steps to match the base model’s performance. For additional details, see AppendixC.3and for cross-domain evaluation see AppendixD.2.

Table 4:Long-form QA after context compression.We evaluate the models on seven tasks from LongBench v2. The KV caches of the questions’ contexts are compressed with a gradient-based method.Retrieval from a compressed context (Q3).We conduct a needle-in-a-haystack experiment in which we first optimize a compact KV cache for the haystack prefix, and then evaluate exact-match passkey retrieval from the resulting compressed cache. We compress the haystack using the gradient-based method described earlier, training the compacted KVs for 100 steps on the reconstruction sequence. We find thatKV-CATimproves retrieval accuracy over the base model, particularly at moderate compression budgets.KV-CATincreases mean retrieval accuracy by 6.4 points forQwen2.5-0.5Band 5.2 points forQwen2.5-1.5B(Table2), achieving an 11–19 point improvement at 30–50 percent keep ratios. See additional details in AppendixC.4.

Compressed long-context question answering (Q4).We evaluate on seven LongBench v2[2]tasks, where the context of each question is compressed using the reconstruction sequence before model prediction. As shown in Table4, compared to the base model, theKV-CAT-trained model achieves better average accuracy across all KV retention ratios, with an average improvement of up to39%. These results demonstrate thatKV-CATyields consistent downstream gains, even when applied to LLMs using only an NTP objective and no post-training. See additional details in AppendixC.5.

6Related Work

Token-level context compression and retrieval.One approach to long-context inference reduces the number of input tokens before prompting, rather than modifying the model’s hidden state. Retrieval-augmented generation keeps the prompt short by selecting only relevant external text[32], while prompt-compression methods prune or rewrite the provided context in token space[24;25;33;42;46]. These methods are complementary to latent KV cache compression.

Learned memory and soft-token compression.Another line of work uses learned memory or soft tokens as compact substitutes for long contexts. Recurrent Memory Transformers add memory tokens that carry information between segments[5]; AutoCompressors adapt language models to map earlier segments into summary vectors that act as soft prompts[10]; and gist-tokens train models to compress prompts into reusable learned tokens[40]. DuoAttention instead learns which attention heads require full retrieval over the KV cache, and which can run with a streaming-style state[53]. These methods modify the model’s runtime interface, architecture, or cache policy. Another direction[9;7;36]uses hypernetworks to encode long contexts into model parameters, typically via LoRA adapters, rather than storing them in the prompt or KV cache.

Post-hoc KV cache compression.Another line of work focuses on reducing the KV cache of pretrained transformers at inference time. Token-selection and eviction methods keep a subset of original KV slots using attention, recency, or query-dependent importance signals[59;43;34;6;50]. KVzip selects keys using reconstruction-style objectives over the context[30]. Other post-hoc methods move beyond choosing a subset of the original tokens. Cartridges optimize a compact latent KV cache for each context[16], Attention Matching constructs compact keys and values to preserve attention behavior[61], and Lexico represents KV vectors with sparse codes over learned universal dictionaries[31]. Our work is orthogonal to these compressors: we ask whether the model can be trained so that the same post-hoc methods work better.

7Conclusion

In this work, we argue that KV cache compressibility is not only a property of the task or input, but also of the representation learned by the transformer. Theoretically, we show that almost any sequence-to-vector function admits both highly compressible and inherently non-compressible implementations. We also provide a motivating example which further suggests that standard training may favor simpler but non-compressible solutions over more structured, compressible ones, thereby motivating compression-aware training. Guided by these insights, we introduceKV-CAT, a compression-aware training procedure that promotes compressible representations. By incorporating a train-time KV sparsification policy with a fixed retention budget into the forward pass,KV-CATexposes the model to compression constraints during training, encouraging representations that remain effective under cache compression. Empirically, we demonstrate that models trained withKV-CATconsistently improve the quality–budget tradeoff of state-of-the-art post-hoc compression methods, at matched cache budgets over a range of tasks. Overall, our results suggest a complementary approach to post-hoc compression methods: training models to become compressible.

8Acknowledgements

YG is supported by the UKRI Engineering and Physical Sciences Research Council (EPSRC) CDT in Autonomous and Intelligent Machines and Systems (grant reference EP/S024050/1). MB is partially supported by the EPSRC Turing AI World-Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub No. EP/Y028872/1. HM is supported by the Israel Science Foundation through a personal grant (ISF 264/23) and an equipment grant (ISF 532/23).

References

• Arora and Ré [2022]Simran Arora and Christopher Ré.Can foundation models help us achieve perfect secrecy?arXiv preprint arXiv:2205.13722, 2022.
• Bai et al. [2025]Yushi Bai, Shangqing Tu, Jiajie Zhang, Hao Peng, Xiaozhi Wang, Xin Lv, Shulin Cao, Jiazheng Xu, Lei Hou, Yuxiao Dong, et al.Longbench v2: Towards deeper understanding and reasoning on realistic long-context multitasks.InProceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 3639–3664, 2025.
• Beltagy et al. [2020]Iz Beltagy, Matthew E Peters, and Arman Cohan.Longformer: The long-document transformer.arXiv preprint arXiv:2004.05150, 2020.
• Bisk et al. [2020]Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi.PIQA: Reasoning about physical commonsense in natural language.InProceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 7432–7439, 2020.doi:10.1609/aaai.v34i05.6239.URLhttps://ojs.aaai.org/index.php/AAAI/article/view/6239.
• Bulatov et al. [2022]Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev.Recurrent memory transformer, 2022.
• Cai et al. [2025]Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Yucheng Li, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Junjie Hu, and Wen Xiao.PyramidKV: Dynamic kv cache compression based on pyramidal information funneling, 2025.
• Charakorn et al. [2026]Rujikorn Charakorn, Edoardo Cetin, Shinnosuke Uesaka, and Robert Tjarko Lange.Doc-to-lora: Learning to instantly internalize contexts.arXiv preprint arXiv:2602.15902, 2026.
• Chen et al. [2025]Qiguang Chen, Libo Qin, Jinhao Liu, Dengyun Peng, Jiannan Guan, Peng Wang, Mengkang Hu, Yuhang Zhou, Te Gao, and Wanxiang Che.Towards reasoning era: A survey of long chain-of-thought for reasoning large language models.arXiv preprint arXiv:2503.09567, 2025.
• Chen et al. [2024]Tong Chen, Hao Fang, Patrick Xia, Xiaodong Liu, Benjamin Van Durme, Luke Zettlemoyer, Jianfeng Gao, and Hao Cheng.Generative adapter: Contextualizing language models in parameters with a single forward pass.arXiv preprint arXiv:2411.05877, 2024.
• Chevalier et al. [2023]Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen.Adapting language models to compress contexts, 2023.
• Choromanski et al. [2020]Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al.Rethinking attention with performers.arXiv preprint arXiv:2009.14794, 2020.
• Clark et al. [2018]Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord.Think you have solved question answering? try ARC, the AI2 reasoning challenge.arXiv preprint arXiv:1803.05457, 2018.URLhttps://arxiv.org/abs/1803.05457.
• Cohan et al. [2018]Arman Cohan, Franck Dernoncourt, Doo Soon Kim, Trung Bui, Seokhwan Kim, Walter Chang, and Nazli Goharian.A discourse-aware attention model for abstractive summarization of long documents.InProceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers), pages 615–621, New Orleans, Louisiana, 2018. Association for Computational Linguistics.doi:10.18653/v1/N18-2097.URLhttps://aclanthology.org/N18-2097/.
• Cybenko [1989]George Cybenko.Approximation by superpositions of a sigmoidal function.Mathematics of control, signals and systems, 2(4):303–314, 1989.
• Eitan [2021]Yam Eitan.The centered convex body whose marginals have the heaviest tails.arXiv preprint arXiv:2110.14382, 2021.
• Eyuboglu et al. [2025]Sabri Eyuboglu, Ryan Ehrlich, Simran Arora, Neel Guha, Dylan Zinsley, Emily Liu, Will Tennien, Atri Rudra, James Zou, Azalia Mirhoseini, and Christopher Ré.Cartridges: Lightweight and general-purpose long context representations via self-study, 2025.
• Gu and Dao [2023]Albert Gu and Tri Dao.Mamba: Linear-time sequence modeling with selective state spaces.arXiv preprint arXiv:2312.00752, 2023.
• Gu et al. [2021]Albert Gu, Karan Goel, and Christopher Ré.Efficiently modeling long sequences with structured state spaces.arXiv preprint arXiv:2111.00396, 2021.
• Habib et al. [2023]Nathan Habib, Clémentine Fourrier, Hynek Kydlíček, Thomas Wolf, and Lewis Tunstall.Lighteval: A lightweight framework for llm evaluation.https://github.com/huggingface/lighteval, 2023.GitHub repository.
• Horn et al. [2017]Alex Horn, Ali Kheradmand, and Mukul Prasad.Delta-net: Real-time network verification using atoms.In14th USENIX Symposium on Networked Systems Design and Implementation (NSDI 17), pages 735–749, 2017.
• Hornik [1991]Kurt Hornik.Approximation capabilities of multilayer feedforward networks.Neural networks, 4(2):251–257, 1991.
• Hwang et al. [2025]Sukjun Hwang, Brandon Wang, and Albert Gu.Dynamic chunking for end-to-end hierarchical sequence modeling.arXiv preprint arXiv:2507.07955, 2025.
• Islam et al. [2023]Pranab Islam, Anand Kannappan, Douwe Kiela, Rebecca Qian, Nino Scherrer, and Bertie Vidgen.Financebench: A new benchmark for financial question answering.arXiv preprint arXiv:2311.11944, 2023.
• Jiang et al. [2023]Huiqiang Jiang, Qianhui Wu, Chin-Yew Lin, Yuqing Yang, and Lili Qiu.LLMLingua: Compressing prompts for accelerated inference of large language models, 2023.
• Jiang et al. [2024]Huiqiang Jiang, Qianhui Wu, Xufang Luo, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu.LongLLMLingua: Accelerating and enhancing llms in long context scenarios via prompt compression, 2024.
• Karlin and Studden [1966a]Samuel Karlin and William J Studden.Optimal experimental designs.The Annals of Mathematical Statistics, 37(4):783–815, 1966a.
• Karlin and Studden [1966b]Samuel Karlin and William J Studden.Tchebycheff systems: With applications in analysis and statistics.(No Title), 1966b.
• Karlin and Ziegler [1966]Samuel Karlin and Zvi Ziegler.Chebyshevian spline functions.Siam Journal on Numerical Analysis, 3(3):514–543, 1966.
• Katharopoulos et al. [2020]Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret.Transformers are rnns: Fast autoregressive transformers with linear attention.InInternational conference on machine learning, pages 5156–5165. PMLR, 2020.
• Kim et al. [2025a]Jang-Hyun Kim, Jinuk Kim, Sangwoo Kwon, Jae W Lee, Sangdoo Yun, and Hyun Oh Song.Kvzip: Query-agnostic kv cache compression with context reconstruction.arXiv preprint arXiv:2505.23416, 2025a.
• Kim et al. [2025b]Junhyuck Kim, Jongho Park, Jaewoong Cho, and Dimitris Papailiopoulos.Lexico: Extreme KV cache compression via sparse coding over universal dictionaries.InProceedings of the 42nd International Conference on Machine Learning, volume 267 ofProceedings of Machine Learning Research, pages 30672–30687, 2025b.
• Lewis et al. [2020]Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela.Retrieval-augmented generation for knowledge-intensive NLP tasks.InAdvances in Neural Information Processing Systems, volume 33, pages 9459–9474, 2020.
• Li et al. [2023]Yucheng Li, Bo Dong, Frank Guerin, and Chenghua Lin.Compressing context to enhance inference efficiency of large language models.InProceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 6342–6353, 2023.doi:10.18653/v1/2023.emnlp-main.391.
• Li et al. [2024a]Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen.SnapKV: Llm knows what you are looking for before generation, 2024a.
• Li et al. [2024b]Zhuoling Li, Xiaogang Xu, Zhenhua Xu, SerNam Lim, and Hengshuang Zhao.Larm: Large auto-regressive model for long-horizon embodied intelligence.arXiv preprint arXiv:2405.17424, 2024b.
• Liu et al. [2026]Yewei Liu, Xiyuan Wang, Yansheng Mao, Yoav Gelbery, Haggai Maron, and Muhan Zhang.Shine: A scalable in-context hypernetwork for mapping context to lora in a single pass.arXiv preprint arXiv:2602.06358, 2026.
• Merity et al. [2016]Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher.Pointer sentinel mixture models, 2016.
• Micchelli and Pinkus [1977]CA Micchelli and Allan Pinkus.Moment theory for weak chebyshev systems with applications to monosplines, quadrature formulae and best one-sided lˆ1-approximation by spline functions with fixed knots.SIAM Journal on Mathematical Analysis, 8(2):206–230, 1977.
• Mihaylov et al. [2018]Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal.Can a suit of armor conduct electricity? a new dataset for open book question answering.InProceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2381–2391, Brussels, Belgium, 2018. Association for Computational Linguistics.doi:10.18653/v1/D18-1260.URLhttps://aclanthology.org/D18-1260/.
• Mu et al. [2023]Jesse Mu, Xiang Lisa Li, and Noah Goodman.Learning to compress prompts with gist tokens, 2023.
• Nam et al. [2024]Daye Nam, Andrew Macvean, Vincent Hellendoorn, Bogdan Vasilescu, and Brad Myers.Using an llm to help with code understanding.InProceedings of the IEEE/ACM 46th International Conference on Software Engineering, pages 1–13, 2024.
• Okcular [2025]Emre Okcular.Context Engineering - Short-Term Memory Management with Sessions from OpenAI Agents SDK, September 2025.
• Oren et al. [2024]Matanel Oren, Michael Hassid, Nir Yarden, Yossi Adi, and Roy Schwartz.Transformers are multi-state RNNs.InProceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 18724–18741, 2024.doi:10.18653/v1/2024.emnlp-main.1043.
• Penedo et al. [2024]Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, Thomas Wolf, et al.The fineweb datasets: Decanting the web for the finest text data at scale.Advances in Neural Information Processing Systems, 37:30811–30849, 2024.
• Rae et al. [2019]Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap.Compressive transformers for long-range sequence modelling.arXiv preprint, 2019.URLhttps://arxiv.org/abs/1911.05507.
• Rajasekaran et al. [2025]Prithvi Rajasekaran, Ethan Dixon, Carly Ryan, and Jeremy Hadfield.Effective context engineering for ai agents, September 2025.
• Sakaguchi et al. [2020]Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi.WinoGrande: An adversarial winograd schema challenge at scale.InProceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 8732–8740, 2020.doi:10.1609/aaai.v34i05.6399.URLhttps://ojs.aaai.org/index.php/AAAI/article/view/6399.
• Sanford et al. [2023]Clayton Sanford, Daniel J Hsu, and Matus Telgarsky.Representational strengths and limitations of transformers.Advances in Neural Information Processing Systems, 36:36677–36707, 2023.
• Sap et al. [2019]Maarten Sap, Hannah Rashkin, Derek Chen, Ronan Le Bras, and Yejin Choi.Social IQa: Commonsense reasoning about social interactions.InProceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 4463–4473, Hong Kong, China, 2019. Association for Computational Linguistics.doi:10.18653/v1/D19-1454.URLhttps://aclanthology.org/D19-1454/.
• Tang et al. [2024]Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han.QUEST: Query-aware sparsity for efficient long-context LLM inference.InProceedings of the 41st International Conference on Machine Learning, volume 235 ofProceedings of Machine Learning Research, pages 47901–47911, 2024.
• Team [2024]Qwen Team.Qwen2.5: A party of foundation models, September 2024.URLhttps://qwenlm.github.io/blog/qwen2.5/.
• Xiao et al. [2024]Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis.Efficient streaming language models with attention sinks, 2024.
• Xiao et al. [2025]Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han.Duoattention: Efficient long-context LLM inference with retrieval and streaming heads.InInternational Conference on Learning Representations, 2025.URLhttps://arxiv.org/abs/2410.10819.
• Yang et al. [2024]Songlin Yang, Jan Kautz, and Ali Hatamizadeh.Gated delta networks: Improving mamba2 with delta rule.arXiv preprint arXiv:2412.06464, 2024.
• Yehudai et al. [2024]Gilad Yehudai, Haim Kaplan, Guy Dar, Royi Rassin, Asma Ghandeharioun, Mor Geva, and Amir Globerson.When can transformers count to n?arXiv preprint arXiv:2407.15160, 2024.
• Zaheer et al. [2017]Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola.Deep sets.Advances in neural information processing systems, 30, 2017.
• Zaheer et al. [2020]Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al.Big bird: Transformers for longer sequences.Advances in neural information processing systems, 33:17283–17297, 2020.
• Zellers et al. [2019]Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi.HellaSwag: Can a machine really finish your sentence?InProceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791–4800, Florence, Italy, 2019. Association for Computational Linguistics.doi:10.18653/v1/P19-1472.URLhttps://aclanthology.org/P19-1472/.
• Zhang et al. [2023]Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen.H2​OH_{2}O: Heavy-hitter oracle for efficient generative inference of large language models, 2023.
• Zheng et al. [2026]Junhao Zheng, Chengming Shi, Xidi Cai, Qiuke Li, Duzhen Zhang, Chenxing Li, Dong Yu, and Qianli Ma.Lifelong learning of large language model based agents: A roadmap.IEEE Transactions on Pattern Analysis and Machine Intelligence, 2026.
• Zweiger et al. [2026]Adam Zweiger, Xinghong Fu, Han Guo, and Yoon Kim.Fast kv compaction via attention matching, 2026.

Appendix ATheory

A.1Transformers and KV Cache Compression

Notation.For a setAA, letA∗A^{*}denote the set of all finite sequences overAA. We denote by𝒆i{\bm{e}}_{i}theii-th standard (one hot encoded) basis vector. Forn∈ℕn\in{\mathbb{N}}, we write[n]={1,…,n}[n]=\{1,\dots,n\}.

For matrices𝑿∈ℝn×d{\bm{X}}\in{\mathbb{R}}^{n\times d}and𝒀∈ℝn×d′{\bm{Y}}\in{\mathbb{R}}^{n\times d^{\prime}}, we denote their row-wise concatenation (along the feature dimension) by

[𝑿,𝒀]∈ℝn×(d+d′).[{\bm{X}},{\bm{Y}}]\in{\mathbb{R}}^{n\times(d+d^{\prime})}.(18) For matrices𝑿∈ℝn×d{\bm{X}}\in{\mathbb{R}}^{n\times d}and𝒁∈ℝn′×d{\bm{Z}}\in{\mathbb{R}}^{n^{\prime}\times d}, we denote their column-wise concatenation (along the sequence dimension) by

[𝑿𝒁]∈ℝ(n+n′)×d.\begin{bmatrix}{\bm{X}}\\ {\bm{Z}}\end{bmatrix}\in{\mathbb{R}}^{(n+n^{\prime})\times d}.(19) When convenient, we sometimes write[𝑿⊤,𝒁⊤]⊤[{\bm{X}}^{\top},{\bm{Z}}^{\top}]^{\top}instead of[𝑿𝒁].\begin{bmatrix}{\bm{X}}\\ {\bm{Z}}\end{bmatrix}.

For vectors𝒗,𝒖∈ℝn{\bm{v}},{\bm{u}}\in{\mathbb{R}}^{n}, we treat them as column vectors inℝn×1{\mathbb{R}}^{n\times 1}and write

[𝒗,𝒖]∈ℝn×2.[{\bm{v}},{\bm{u}}]\in{\mathbb{R}}^{n\times 2}.(20)

Definition A.1(Attention Head).

Letdin,dout∈ℕd_{\mathrm{in}},d_{\mathrm{out}}\in\mathbb{N}. Anattention headis a function

AttnHead:(ℝdin)∗→(ℝdout)∗\mathrm{AttnHead}:({\mathbb{R}}^{d_{\mathrm{in}}})^{*}\to({\mathbb{R}}^{d_{\mathrm{out}}})^{*}(21)defined as follows. For a sequence𝑿=[𝒙1,…,𝒙n]⊤∈ℝn×din{\bm{X}}=[{\bm{x}}_{1},\dots,{\bm{x}}_{n}]^{\top}\in\mathbb{R}^{n\times d_{\mathrm{in}}},

𝑸=𝑿​𝑾Q,𝑲=𝑿​𝑾K,𝑽=𝑿​𝑾𝑽,{\bm{Q}}={\bm{X}}{\bm{W}}_{Q},\quad{\bm{K}}={\bm{X}}{\bm{W}}_{K},\quad{\bm{V}}={\bm{X}}{\bm{W}}_{\bm{V}},(22)where𝑾Q,𝑾K,𝑾V∈ℝdin×dout{\bm{W}}_{Q},{\bm{W}}_{K},{\bm{W}}_{V}\in\mathbb{R}^{d_{\mathrm{in}}\times d_{\mathrm{out}}}. The output is

AttnHead​(X)=softmax​(𝑸​𝑲⊤dk)​𝑽,\mathrm{AttnHead}(X)=\mathrm{softmax}\!\left(\frac{{\bm{Q}}{\bm{K}}^{\top}}{\sqrt{d_{k}}}\right){\bm{V}},(23)where the softmax is applied row-wise.

Definition A.2(Transformer Block).

Fixh,din,dout,dff∈ℕh,d_{\mathrm{in}},d_{\mathrm{out}},d_{\mathrm{ff}}\in{\mathbb{N}}. Atransformer blockis a function

Block:(ℝdin)∗→(ℝdout)∗\mathrm{Block}:({\mathbb{R}}^{d_{\mathrm{in}}})^{*}\to({\mathbb{R}}^{d_{\mathrm{out}}})^{*}(24)defined as follows. For a sequence

𝑿=[𝒙1,…,𝒙n]⊤∈ℝn×din,{\bm{X}}=[{\bm{x}}_{1},\dots,{\bm{x}}_{n}]^{\top}\in{\mathbb{R}}^{n\times d_{\mathrm{in}}},(25)we define

Block​(𝑿)=FFN​([head1​(𝑿),…,headh​(𝑿)]​𝑾O+𝑿),\mathrm{Block}({\bm{X}})=\mathrm{FFN}\!\left([\mathrm{head}_{1}({\bm{X}}),\dots,\mathrm{head}_{h}({\bm{X}})]{\bm{W}}_{O}+{\bm{X}}\right),(26)where eachheadi\mathrm{head}_{i}is an attention head with input and output dimensionsdin,doutd_{\mathrm{in}},d_{\mathrm{out}}respectively, and𝑾O∈ℝh⋅dout×dout{\bm{W}}_{O}\in{\mathbb{R}}^{h\cdot d_{\mathrm{out}}\times d_{\mathrm{out}}}, andFFN\mathrm{FFN}is a 2-layer feed forward network applied row-wise, i.e.,

FFN​(𝑿)i=FFN​(𝒙i),i=1,…,n,\mathrm{FFN}({\bm{X}})_{i}=\mathrm{FFN}({\bm{x}}_{i}),\quad i=1,\dots,n,(27)with

FFN​(𝒙)=σ​(𝒙​𝑾1+𝒃1)​𝑾2+𝒃2,\mathrm{FFN}({\bm{x}})=\sigma({\bm{x}}{\bm{W}}_{1}+{\bm{b}}_{1}){\bm{W}}_{2}+{\bm{b}}_{2},(28)where𝑾1∈ℝdout×dff{\bm{W}}_{1}\in{\mathbb{R}}^{d_{\mathrm{out}}\times d_{\mathrm{ff}}},𝑾2∈ℝdff×dout{\bm{W}}_{2}\in{\mathbb{R}}^{d_{\mathrm{ff}}\times d_{\mathrm{out}}},𝒃1∈ℝdff{\bm{b}}_{1}\in{\mathbb{R}}^{d_{\mathrm{ff}}}, and𝒃2∈ℝdout{\bm{b}}_{2}\in{\mathbb{R}}^{d_{\mathrm{out}}}.

Note.In the theoretical analysis, we omit layer normalization, causal masking, and the residual connection following the FFN for simplicity, as is standard in prior theoretical treatments (See e.g.Sanford et al. [2023], Yehudai et al. [2024]). Our results extend to the full architecture with only minor modifications. All empirical evaluations are conducted with these components included.

Definition A.3(Token Embedding).

LetAAbe a finite alphabet. Atoken embeddingfunction is a function

emb:A×ℕ→ℝdmodel.\mathrm{emb}:A\times{\mathbb{N}}\to{\mathbb{R}}^{d_{\mathrm{model}}}.(29)Given a sequence𝒂=(a1,…,an)∈A∗{\bm{a}}=(a_{1},\dots,a_{n})\in A^{*}, its embedding is the sequence

emb​(𝒂)=[emb​(a1,1),…,emb​(an,n)]⊤∈ℝn×dmodel.\mathrm{emb}({\bm{a}})=[\mathrm{emb}(a_{1},1),\dots,\mathrm{emb}(a_{n},n)]^{\top}\in{\mathbb{R}}^{n\times d_{\mathrm{model}}}.(30)

Definition A.4(Transformer).

LetAAbe a finite alphabet. Atransformeris a tuple

𝖬=(Block1,…,BlockL,emb),\mathsf{M}=(\mathrm{Block}_{1},\dots,\mathrm{Block}_{L},\mathrm{emb}),(31) where eachBlockℓ:(ℝdℓ−1)∗→(ℝdℓ)∗\mathrm{Block}_{\ell}:({\mathbb{R}}^{d_{\ell-1}})^{*}\to({\mathbb{R}}^{d_{\ell}})^{*}is a transformer block, andemb:A∗→(ℝd0)∗\mathrm{emb}:A^{*}\to({\mathbb{R}}^{d_{0}})^{*}is a token embedding function.

The transformer𝖬\mathsf{M}defines a function

𝖬:A∗→ℝdL\mathsf{M}:A^{*}\to{\mathbb{R}}^{d_{L}}(32)as follows. For𝒂=(a1,…,an)∈A∗{\bm{a}}=(a_{1},\dots,a_{n})\in A^{*}, let

𝑿(0)=emb​(𝒂),𝑿(ℓ)=Blockℓ​(𝑿(ℓ−1)),ℓ=1,…,L.{\bm{X}}^{(0)}=\mathrm{emb}({\bm{a}}),\qquad{\bm{X}}^{(\ell)}=\mathrm{Block}_{\ell}({\bm{X}}^{(\ell-1)}),\quad\ell=1,\dots,L.(33)Then

𝖬​(𝒂)=𝑿n(L),\mathsf{M}({\bm{a}})={\bm{X}}^{(L)}_{n},(34)i.e., the output is the representation of the final token.

Definition A.5(KV Cache Compression).

AKV cache compression policyis a tuple𝖢=(𝖼1,…,𝖼L)\mathsf{C}=(\mathsf{c}_{1},\dots,\mathsf{c}_{L}), where each

𝖼ℓ:(ℝdℓ×ℝdℓ)∗→(ℝdℓ×ℝdℓ)∗\mathsf{c}_{\ell}:({\mathbb{R}}^{d_{\ell}}\times{\mathbb{R}}^{d_{\ell}})^{*}\to({\mathbb{R}}^{d_{\ell}}\times{\mathbb{R}}^{d_{\ell}})^{*}(35)maps a sequence of key–value pairs to a shorter sequence.

Concretely, for

(𝑲,𝑽)=([𝒌1,…,𝒌n]⊤,[𝒗1,…,𝒗n]⊤),({\bm{K}},{\bm{V}})=\bigl([{\bm{k}}_{1},\dots,{\bm{k}}_{n}]^{\top},[{\bm{v}}_{1},\dots,{\bm{v}}_{n}]^{\top}\bigr),(36)we write

𝖼ℓ​(𝑲,𝑽)=(𝑲~,𝑽~)=([𝒌~1,…,𝒌~r​(n)]⊤,[𝒗~1,…,𝒗~r​(n)]⊤),\mathsf{c}_{\ell}({\bm{K}},{\bm{V}})=(\tilde{{\bm{K}}},\tilde{{\bm{V}}})=\bigl([\tilde{{\bm{k}}}_{1},\dots,\tilde{{\bm{k}}}_{r(n)}]^{\top},[\tilde{{\bm{v}}}_{1},\dots,\tilde{{\bm{v}}}_{r(n)}]^{\top}\bigr),(37)wherer​(n)≤nr(n)\leq n. The functionr​(⋅)r(\cdot)is called thecompression budget.

Given a transformer𝖬=(emb,Block1,…,BlockL)\mathsf{M}=(\mathrm{emb},\mathrm{Block}_{1},\dots,\mathrm{Block}_{L}), a compression policy𝖢\mathsf{C}, and acontextsequence𝒂=(a1,…,an){\bm{a}}=(a_{1},\dots,a_{n}), we define the associatedcompressed transformer𝖬𝖢,𝒂\mathsf{M}_{\mathsf{C},{\bm{a}}}.

For an input sequence𝒃=(b1,…,bk){\bm{b}}=(b_{1},\dots,b_{k}), the output𝖬𝖢,𝒂​(𝒃)\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})is obtained by running the forward pass of𝖬\mathsf{M}on the concatenated sequence[𝒂,𝒃]=(a1,…,an,b1,…,bk)[{\bm{a}},{\bm{b}}]=(a_{1},\dots,a_{n},b_{1},\dots,b_{k}), with the following modifications.

For eachℓ\elland each attention head at theℓ\ell-th block, after computing keys and values, the compression function𝖼ℓ\mathsf{c}_{\ell}is applied to the KV pairs corresponding to the context tokens. The compressed KV pairs replace the original ones in the attention computation.

More precisely, let(𝑲𝒂,𝑽𝒂)({\bm{K}}_{\bm{a}},{\bm{V}}_{\bm{a}})denote the KV cache produced by the prefix𝒂{\bm{a}}at some attention head ofBlockℓ\mathrm{Block}_{\ell}of the original model𝖬\mathsf{M}, and let𝒀∈ℝk×dℓ−1{\bm{Y}}\in{\mathbb{R}}^{k\times d_{\ell-1}}be the input corresponding to𝒃{\bm{b}}in the forward pass computation of𝖬𝖢,𝒂\mathsf{M}_{\mathsf{C},{\bm{a}}}after blockℓ−1\ell-1. We first compute

𝑸𝒃=𝒀​𝑾Q,𝑲𝒃=𝒀​𝑾K,𝑽𝒃=𝒀​𝑾V.{\bm{Q}}_{\bm{b}}={\bm{Y}}{\bm{W}}_{Q},\quad{\bm{K}}_{\bm{b}}={\bm{Y}}{\bm{W}}_{K},\quad{\bm{V}}_{\bm{b}}={\bm{Y}}{\bm{W}}_{V}.(38)We then compress the context KV pairs:

(𝑲~𝒂,𝑽~𝒂)=𝖼ℓ​(𝑲𝒂,𝑽𝒂).(\tilde{{\bm{K}}}_{\bm{a}},\tilde{{\bm{V}}}_{\bm{a}})=\mathsf{c}_{\ell}({\bm{K}}_{\bm{a}},{\bm{V}}_{\bm{a}}).(39)The resulting attention computation is

AttnHead​(𝒀)=softmax​(1dk​𝑸𝒃​[𝑲~𝒂𝑲𝒃]⊤)​[𝑽~𝒂𝑽𝒃].\mathrm{AttnHead}({\bm{Y}})=\mathrm{softmax}\!\left(\frac{1}{\sqrt{d_{k}}}{\bm{Q}}_{\bm{b}}\begin{bmatrix}\tilde{{\bm{K}}}_{\bm{a}}\\ {\bm{K}}_{\bm{b}}\end{bmatrix}^{\top}\right)\begin{bmatrix}\tilde{{\bm{V}}}_{\bm{a}}\\ {\bm{V}}_{\bm{b}}\end{bmatrix}.(40)

Definition A.6.

[KV compressibility] LetN∈ℕN\in{\mathbb{N}},ε>0\varepsilon>0, and letr:ℕ→ℕr:{\mathbb{N}}\to{\mathbb{N}}be a budget function. A transformer𝖬\mathsf{M}is said to be(N,ε,r)(N,\varepsilon,r)-compressible if there exists a KV cache compression policy𝖢\mathsf{C}with budgetrrsuch that for every pair of sequences𝒂,𝒃{\bm{a}},{\bm{b}}of lengthsnnandkkwith combined length satisfyingn+k≤Nn+k\leq N, it holds that

‖𝖬​([𝒂,𝒃])−𝖬𝖢,𝒂​(𝒃)‖<ε.\|\mathsf{M}([{\bm{a}},{\bm{b}}])-\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})\|<\varepsilon.(41)

We includeNNexplicitly in the definition above, as many natural sequence-to-vector functions require model dimension scaling withNN(e.g.,O​(N)O(N)orO​(log⁡N)O(\log N)) in order to achieve arbitrarily accurate approximation; seeSanford et al. [2023], Yehudai et al. [2024]for examples.

A.2Motivating example: histogram computation

In this section we provide formal proofs of Propositions3.2and3.3.

Proof of Proposition3.2.

The construction of𝖬=(emb,Block1,Block2)\mathsf{M}=(\mathrm{emb},\mathrm{Block}_{1},\mathrm{Block}_{2})is straightforward. First, forj∈[m]j\in[m]andi∈[N]i\in[N], define

emb​(j,i)=𝒖j+𝒑i,\mathrm{emb}(j,i)={\bm{u}}_{j}+{\bm{p}}_{i},(42)where𝒖j,𝒑i∈ℝm{\bm{u}}_{j},{\bm{p}}_{i}\in{\mathbb{R}}^{m}encode token value and token position respectively, and assume thatemb\mathrm{emb}is injective over[m]×[N][m]\times[N].

Let𝑿∈ℝn×m{\bm{X}}\in{\mathbb{R}}^{n\times m}. We defineBlock1​(ℝm)∗→(ℝ2​m)∗\mathrm{Block}_{1}({\mathbb{R}}^{m})^{*}\to({\mathbb{R}}^{2m})^{*}by

Block1​(𝑿)i≈ρ1​(𝑿i)\mathrm{Block}_{1}({\bm{X}})_{i}\approx\rho_{1}({\bm{X}}_{i})(43)whereρ1:ℝm→ℝ2​m\rho_{1}:{\mathbb{R}}^{m}\to{\mathbb{R}}^{2m}is applied row wise𝑿{\bm{X}}and satisfies for eachj,i∈[m]×[N]j,i\in[m]\times[N]:

ρ1​(𝒖j+𝒑i)=[𝒆j𝟎].\rho_{1}({\bm{u}}_{j}+{\bm{p}}_{i})=\begin{bmatrix}{\bm{e}}_{j}\\ \mathbf{0}\end{bmatrix}.(44) Block1\mathrm{Block}_{1}can be implemented by zeroing out all attention projections and relying on the residual connection together with a feedforward network that approximatesρ1\rho_{1}to arbitrary precision (this can be obtained as a two-layer feedforward network is a universal approximator of continuous functions on compact sets[Cybenko,1989, Hornik,1991]).

Let𝑿∈ℝn×2​m{\bm{X}}\in{\mathbb{R}}^{n\times 2m}, and write each row as𝒙i{\bm{x}}_{i}. We define

Block2​(𝑿)i=1n​∑j=1n𝒙j\mathrm{Block}_{2}({\bm{X}})_{i}=\frac{1}{n}\sum_{j=1}^{n}{\bm{x}}_{j}(45)i.e., each output token is replaced by the average of all input tokens. This operation can be implemented using a single attention head by setting

𝑾Q=𝑾K=𝟎,{\bm{W}}_{Q}={\bm{W}}_{K}=\mathbf{0},(46)so that all attention weights are uniform, and choosing

𝑾V=𝑾O=(𝟎𝟎𝑰𝟎).{\bm{W}}_{V}={\bm{W}}_{O}=\begin{pmatrix}\mathbf{0}&\mathbf{0}\\ {\bm{I}}&\mathbf{0}\end{pmatrix}.(47)This choice extracts the first half of each vector and places it in the second half before averaging. After applying a residual connection, the intermediate representation of theii-th token is then

[𝒙i1n​∑j=1n𝒙j]\begin{bmatrix}{\bm{x}}_{i}\\ \frac{1}{n}\sum_{j=1}^{n}{\bm{x}}_{j}\par\end{bmatrix}(48) Finally, the feedforward network is taken to be

FFN​(𝒙,𝒚)=𝒚,\mathrm{FFN}({\bm{x}},{\bm{y}})={\bm{y}},(49) resulting in a final representation of1n​∑i=1n𝒙i\frac{1}{n}\sum_{i=1}^{n}{\bm{x}}_{i}. For input sequence𝒂{\bm{a}}of lengthnnwe thus have

𝖬​(𝒂)≈1n​∑i=1n𝒆ai=(1n​∑i=1n𝟏ai=1,1n​∑i=1n𝟏ai=2,…,1n​∑i=1n𝟏ai=m)=fhist​(𝒂),\mathsf{M}({\bm{a}})\approx\frac{1}{n}\sum_{i=1}^{n}{\bm{e}}_{a_{i}}=\left(\frac{1}{n}\sum_{i=1}^{n}\mathbf{1}_{a_{i}=1},\;\frac{1}{n}\sum_{i=1}^{n}\mathbf{1}_{a_{i}=2},\;\dots,\;\frac{1}{n}\sum_{i=1}^{n}\mathbf{1}_{a_{i}=m}\right)=f_{\mathrm{hist}}({\bm{a}}),(50)Thus𝖬\mathsf{M}approximatesfhistf_{\mathrm{hist}}to arbitrary precision. Now let𝖢=(𝖼1,𝖼2)\mathsf{C}=(\mathsf{c}_{1},\mathsf{c}_{2})be any compression policy with budget functionr​(⋅)r(\cdot)satisfyingr​(N−1)<N−1r(N-1)<N-1. Fix a prefix sequence𝒂{\bm{a}}of lengthnnand a suffix sequence𝒃{\bm{b}}of lengthkk, and consider the computation of𝖬𝖢,𝒂​(𝒃)\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}}).

First, observe that in the first block all attention projection matrices (and in particular𝑾O{\bm{W}}_{O}) are zero. Hence,𝖼1\mathsf{c}_{1}has no effect on the representations of the tokens of𝒃{\bm{b}}, and after the first block the representations are exactly as in the uncompressed model.

Next, consider the second block. Let(𝑲𝒂,𝑽𝒂)({\bm{K}}_{\bm{a}},{\bm{V}}_{\bm{a}})denote the KV cache corresponding to the prefix𝒂{\bm{a}}computed durning the computation of𝖬​(𝒂)\mathsf{M}({\bm{a}}), and let(𝑲~𝒂,𝑽~𝒂)=𝖼2​(𝑲𝒂,𝑽𝒂)(\tilde{{\bm{K}}}_{\bm{a}},\tilde{{\bm{V}}}_{\bm{a}})=\mathsf{c}_{2}({\bm{K}}_{\bm{a}},{\bm{V}}_{\bm{a}})be the compressed cache. Since𝑾Q=0{\bm{W}}_{Q}=0, all attention scores are uniform for both𝖬\mathsf{M}and𝖬𝖢,𝒂\mathsf{M}_{\mathsf{C},{\bm{a}}}and so

𝖬𝖢,𝒂​(𝒃)=1r​(n)+k​(∑i=1r​(n)𝒗~i+∑i=1k𝒆bi).\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})=\frac{1}{r(n)+k}\left(\sum_{i=1}^{r(n)}\tilde{\bm{v}}_{i}+\sum_{i=1}^{k}{\bm{e}}_{b_{i}}\right).(51) 𝖬​([𝒂,𝒃])=1n+k​(∑i=1n𝒆ai+∑i=1k𝒆bi).\mathsf{M}([{\bm{a}},{\bm{b}}])=\frac{1}{n+k}\left(\sum_{i=1}^{n}{\bm{e}}_{a_{i}}+\sum_{i=1}^{k}{\bm{e}}_{b_{i}}\right).(52) For notational convenience, define

𝒗~=1r​(n)+k​∑i=1r​(n)𝒗~i,𝒂~=1n+k​∑i=1n𝒆ai,𝒃~=∑i=1k𝒆bi.\tilde{\bm{v}}=\frac{1}{r(n)+k}\sum_{i=1}^{r(n)}\tilde{\bm{v}}_{i},\quad\tilde{\bm{a}}=\frac{1}{n+k}\sum_{i=1}^{n}{\bm{e}}_{a_{i}},\quad\tilde{\bm{b}}=\sum_{i=1}^{k}{\bm{e}}_{b_{i}}.(53) Then the difference between the compressed and full outputs can be written as

‖𝖬𝖢,𝒂​(𝒃)−𝖬​([𝒂,𝒃])‖=‖𝒗~−𝒂~+(1r​(n)+k−1n+k)​𝒃~‖.\|\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})-\mathsf{M}([{\bm{a}},{\bm{b}}])\|=\left\|\tilde{\bm{v}}-\tilde{\bm{a}}+\left(\frac{1}{r(n)+k}-\frac{1}{n+k}\right)\tilde{\bm{b}}\right\|.(54) We now derive a lower bound. First, choose𝒃=(1,1,…,1){\bm{b}}=(1,1,\dots,1)and examine the first coordinate. This yields

‖𝖬𝖢,𝒂​(𝒃)−𝖬​([𝒂,𝒃])‖≥|𝒗~1−𝒂~1+k​(r​(n)−n)(n+k)​(r​(n)+k)|.\|\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})-\mathsf{M}([{\bm{a}},{\bm{b}}])\|\geq\left|\tilde{\bm{v}}_{1}-\tilde{\bm{a}}_{1}+\frac{k(r(n)-n)}{(n+k)(r(n)+k)}\right|.(55)Next, choosing𝒃=(2,2,…,2){\bm{b}}=(2,2,\dots,2)removes the contribution of the last term, giving

‖𝖬𝖢,𝒂​(𝒃)−𝖬​([𝒂,𝒃])‖≥|𝒗~1−𝒂~1|.\|\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})-\mathsf{M}([{\bm{a}},{\bm{b}}])\|\geq|\tilde{\bm{v}}_{1}-\tilde{\bm{a}}_{1}|.(56) Combining the two inequalities, we conclude that for any choice of𝒗~\tilde{\bm{v}}, there exists a𝒃{\bm{b}}such that

‖𝖬𝖢,𝒂​(𝒃)−𝖬​([𝒂,𝒃])‖≥k​(n−r​(n))2​(n+k)​(r​(n)+k).\|\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})-\mathsf{M}([{\bm{a}},{\bm{b}}])\|\geq\frac{k(n-r(n))}{2(n+k)(r(n)+k)}.(57)Finally, takingn=N−1n=N-1andk=1k=1yields

‖𝖬𝖢,𝒂​(𝒃)−𝖬​([𝒂,𝒃])‖≥N−1−r​(N−1)2​(N)​(r​(N−1)+1)=C>0,\|\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})-\mathsf{M}([{\bm{a}},{\bm{b}}])\|\geq\frac{N-1-r(N-1)}{2(N)(r(N-1)+1)}=C>0,(58)which completes the proof. ∎

Note:for the weakest compression caser​(n)=n−1r(n)=n-1, the constant above isC=O​(N−2)C=O(N^{-2}), and so asNNgrows the approximation bound becomes weaker, but for the extreme compression caser​(n)=cr(n)=cfor somec∈ℕc\in{\mathbb{N}},CCcan be chosen independently ofNN.

proof of Proposition3.3.

We construct𝖬=(emb,Block1,Block2)\mathsf{M}=(\mathrm{emb},\mathrm{Block}_{1},\mathrm{Block}_{2})similarly to the above construction, with a few modifications. Like before, forj∈[m]j\in[m]andi∈[N]i\in[N], define

emb​(j,i)=𝒖j+𝒑i,\mathrm{emb}(j,i)={\bm{u}}_{j}+{\bm{p}}_{i},(59)where𝒖j,𝒑i∈ℝm{\bm{u}}_{j},{\bm{p}}_{i}\in{\mathbb{R}}^{m}encode token value and token position respectively, and assume thatemb\mathrm{emb}is injective over[m]×[N][m]\times[N].

Let𝑿∈ℝn×m{\bm{X}}\in{\mathbb{R}}^{n\times m}. We defineBlock1​(ℝm)∗→(ℝ4​m)∗\mathrm{Block}_{1}({\mathbb{R}}^{m})^{*}\to({\mathbb{R}}^{4m})^{*}by

Block1​(𝑿)i≈ρ1​(𝑿i)\mathrm{Block}_{1}({\bm{X}})_{i}\approx\rho_{1}({\bm{X}}_{i})(60)whereρ1:ℝm→ℝ4​m\rho_{1}:{\mathbb{R}}^{m}\to{\mathbb{R}}^{4m}is applied row wise𝑿{\bm{X}}and satisfies for eachj,i∈[m]×[N]j,i\in[m]\times[N]:

ρ1​(𝒖j+𝒑i)=[𝒆j𝒑i𝟎𝟎].\rho_{1}({\bm{u}}_{j}+{\bm{p}}_{i})=\begin{bmatrix}{\bm{e}}_{j}\\ {\bm{p}}_{i}\\ \mathbf{0}\\ \mathbf{0}\end{bmatrix}.(61) Where each of the vector blocks above is of sizemm. That is,Block1\mathrm{Block}_{1}is identical to the construction in the proof of Proposition3.2, except that it preserves the positional component. As before,Block1\mathrm{Block}_{1}can be implemented by zeroing out all attention projections and relying on the residual connection together with a feedforward network that approximatesρ1\rho_{1}to arbitrary precision. This follows from the universal approximation property of two-layer feedforward networks for continuous functions on compact sets[Cybenko,1989, Hornik,1991].

We constructBlock2\mathrm{Block}_{2}to be composed of a single attention head with

𝑾Q=𝑾K=𝟎,𝑾O=𝑰{\bm{W}}_{Q}={\bm{W}}_{K}=\mathbf{0},{\bm{W}}_{O}={\bm{I}}(62)so that all attention weights are uniform, and choosing

𝑾V=(𝑰𝟎𝟎𝟎𝟎𝑰𝟎𝟎𝑰𝟎𝟎𝟎𝟎𝟎𝟎𝟎).{\bm{W}}_{V}=\begin{pmatrix}{\bm{I}}&\mathbf{0}&\mathbf{0}&\mathbf{0}\\ \mathbf{0}&{\bm{I}}&\mathbf{0}&\mathbf{0}\\ {\bm{I}}&\mathbf{0}&\mathbf{0}&\mathbf{0}\\ \mathbf{0}&\mathbf{0}&\mathbf{0}&\mathbf{0}\end{pmatrix}.(63) This choice extracts the first m-length block of each vector and places it in the third block before averaging, keeping the first and second vector blocks the same, and zeroing out the last vector block. After applying a residual connection, the intermediate representation of theii-th tokenaia_{i}is then

[𝒆ai𝒑i1n​∑j=1n𝒆j𝟎].\begin{bmatrix}{\bm{e}}_{a_{i}}\\ {\bm{p}}_{i}\\ \frac{1}{n}\sum_{j=1}^{n}{\bm{e}}_{j}\\ \mathbf{0}\end{bmatrix}.(64) finally, we choose the last feed forward network to approximate a mapρ2\rho_{2}satisfying for alli,j,k∈[N]i,j,k\in[N]:

ρ2​(𝒘,𝒙,𝒚,𝒛)={𝒚if​𝒛=0j−i+1j​𝒚if​𝒛=𝒑ik+1​and​𝒙=𝒑j.\rho_{2}({\bm{w}},{\bm{x}},{\bm{y}},{\bm{z}})=\begin{cases}{\bm{y}}&\text{ if }{\bm{z}}=0\\ \frac{j-i+1}{j}{\bm{y}}&\text{ if }{\bm{z}}=\frac{{\bm{p}}_{i}}{k+1}\text{ and }{\bm{x}}={\bm{p}}_{j}.\end{cases}(65) Such a function can be realized to arbitrary precision by a feedforward network: the above specification is defined on a union of disjoint compact sets, and therefore admits a continuous extension to a compact domain, which can in turn be approximated by a standard feedforward network via universal approximation.

Equations64and65imply that, in the uncompressed setting, for any sequence𝒂{\bm{a}}of lengthn≤Nn\leq N, we recover as before

𝖬​(𝒂)≈1n​∑i=1n𝒆ai=(1n​∑i=1n𝟏ai=1,1n​∑i=1n𝟏ai=2,…,1n​∑i=1n𝟏ai=m)=fhist​(𝒂),\mathsf{M}({\bm{a}})\approx\frac{1}{n}\sum_{i=1}^{n}{\bm{e}}_{a_{i}}=\left(\frac{1}{n}\sum_{i=1}^{n}\mathbf{1}_{a_{i}=1},\;\frac{1}{n}\sum_{i=1}^{n}\mathbf{1}_{a_{i}=2},\;\dots,\;\frac{1}{n}\sum_{i=1}^{n}\mathbf{1}_{a_{i}=m}\right)=f_{\mathrm{hist}}({\bm{a}}),(66) and so𝖬\mathsf{M}approximatesfhistf_{\mathrm{hist}}. We now define a compression policy𝖢=(𝖼1,𝖼2)\mathsf{C}=(\mathsf{c}_{1},\mathsf{c}_{2}). We first note that, like before, SINCE𝑾O=𝟎{\bm{W}}_{O}=\mathbf{0}in the first block,𝖼1\mathsf{c}_{1}does not effect the forward pass of𝖬𝖢,𝒂\mathsf{M}_{\mathsf{C},{\bm{a}}}.

Given a prefix𝒂{\bm{a}}of lengthnnwith KV cache(𝑲𝒂,𝑽𝒂)({\bm{K}}_{\bm{a}},{\bm{V}}_{\bm{a}}), we compress it into a single KV pair:

𝖼2​(𝑲𝒂,𝑽𝒂)=(𝟎,[∑i=1n𝒆ai𝟎𝟎𝒑n]).\mathsf{c}_{2}({\bm{K}}_{\bm{a}},{\bm{V}}_{\bm{a}})=\Bigl(\mathbf{0},\begin{bmatrix}\sum_{i=1}^{n}{\bm{e}}_{a_{i}}\\ \mathbf{0}\\ \mathbf{0}\\ {\bm{p}}_{n}\end{bmatrix}\Bigr).(67)That is, we store a single value vector containing (i) the unnormalized histogram of the prefix and (ii) its length encoded via𝒑n{\bm{p}}_{n}. Now consider processing a suffix𝒃=(b1,…,bk){\bm{b}}=(b_{1},\dots,b_{k}). Since the compressed prefix consists of a single KV entry with zero key, attention again produces a uniform average over the compressed prefix and suffix tokens. The representation of the final tokenbkb_{k}becomes

[𝒆bi𝒑n+k1k+1​(∑i=1n𝒆ai+∑i=1k𝒆bi)1k+1​𝒑n].\begin{bmatrix}{\bm{e}}_{b_{i}}\\ {\bm{p}}_{n+k}\\ \frac{1}{k+1}\Bigl(\sum_{i=1}^{n}{\bm{e}}_{a_{i}}+\sum_{i=1}^{k}{\bm{e}}_{b_{i}}\Bigr)\\ \frac{1}{k+1}{\bm{p}}_{n}\end{bmatrix}.(68) Applyingρ2\rho_{2}and using equation65, we obtain

𝖬𝖢,𝒂​(𝒃)≈1+(n+k)−nn+k⋅1k+1​(∑i=1n𝒆ai+∑i=1k𝒆bi)=1n+k​∑i=1n+k𝒆ci,\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})\approx\frac{1+(n+k)-n}{n+k}\cdot\frac{1}{k+1}\Bigl(\sum_{i=1}^{n}{\bm{e}}_{a_{i}}+\sum_{i=1}^{k}{\bm{e}}_{b_{i}}\Bigr)=\frac{1}{n+k}\sum_{i=1}^{n+k}{\bm{e}}_{c_{i}},(69)wherecic_{i}ranges over[𝒂,𝒃][{\bm{a}},{\bm{b}}]. This matches the histogram of the concatenated sequence (up toε\varepsilon), completing the proof.

∎

A.3Compressible Transformers for General Functions

In this section, we present a formal proof of Theorem3.1. The proof is divided into two lemmas: the first (LemmaA.7) establishes the existence of compressible transformer approximations, and the second (LemmaA.8demonstrates the existence of non-compressible ones. The theorem then follows immediately.

Lemma A.7.

LetAAbe a finite alphabet, and letf:⋃n≤NAn→ℝdoutf:\bigcup_{n\leq N}A^{n}\to\mathbb{R}^{d_{\text{out}}}be any sequence-to-vector function. Then for everyε>0\varepsilon>0There exists a transformer such that

1. 1.(Approximation) For every sequence𝒂=(a1,…,an)∈An{\bm{a}}=(a_{1},\dots,a_{n})\in A^{n}withn≤Nn\leq N, ‖f​(𝒂)−𝖬​(𝒂)‖<ε.\|f({\bm{a}})-\mathsf{M}({\bm{a}})\|<\varepsilon.(70)
2. 2.(Maximal compressibility) There exists a KV cache compression policy𝖢\mathsf{C}with compression budget such that for every prefix𝒂∈An{\bm{a}}\in A^{n}and suffix𝒃∈Ak{\bm{b}}\in A^{k}withk+n≤Nk+n\leq N, ‖𝖬​([𝒂,𝒃])−𝖬𝖢,𝒂​(𝒃)‖<ε.\|\mathsf{M}([{\bm{a}},{\bm{b}}])-\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})\|<\varepsilon.(72)

Proof.

We assume that the token embedding mapemb:A×ℕ→ℝd0\mathrm{emb}:A\times{\mathbb{N}}\to{\mathbb{R}}^{d_{0}}is defined by

emb​(a,i)=𝒖a+𝒑i\mathrm{emb}(a,i)={\bm{u}}_{a}+{\bm{p}}_{i}(73) where‖𝒑i‖=‖𝒑j‖\|{\bm{p}}_{i}\|=\|{\bm{p}}_{j}\|for alli,j∈[N]i,j\in[N], and𝒖a,𝒑i∈ℝd0{\bm{u}}_{a},{\bm{p}}_{i}\in{\mathbb{R}}^{d_{0}}for alla∈Aa\in A. We further assume injectivity of the embedding function, that is

(a,i)≠(b,j)⇒emb​(a,i)≠emb​(b,j).(a,i)\neq(b,j)\;\Rightarrow\;\mathrm{emb}(a,i)\neq\mathrm{emb}(b,j).(74) Such embeddings can be obtained, for example, using a learned lookup table for tokens combined with a positional encoding (e.g., sinusoidal) with sufficiently large frequencies.

Let𝒖∅∈ℝd0{\bm{u}}_{\emptyset}\in{\mathbb{R}}^{d_{0}}be a padding vector such that𝒗∅≠emb​(a,i){\bm{v}}_{\emptyset}\neq\mathrm{emb}(a,i)for alla∈Aa\in Aandi∈ℕi\in{\mathbb{N}}. For a sequence𝒂=(a1,…,an)∈A∗{\bm{a}}=(a_{1},\dots,a_{n})\in A^{*}withn≤Nn\leq N, define

U𝒂={𝑷​[emb​(a1,1),…,emb​(an,n),𝒖∅,…,𝒖∅]⊤∣𝑷∈ℝN×N​is a permutation matrix}⊂ℝN×d0.U_{{\bm{a}}}=\left\{{\bm{P}}[\mathrm{emb}(a_{1},1),\dots,\mathrm{emb}(a_{n},n),{\bm{u}}_{\emptyset},\dots,{\bm{u}}_{\emptyset}]^{\top}\mid{\bm{P}}\in{\mathbb{R}}^{N\times N}\text{ is a permutation matrix}\right\}\subset{\mathbb{R}}^{N\times d_{0}}.(75)That is, we embed the sequence, pad it to lengthNNusing𝒖∅{\bm{u}}_{\emptyset}, and then apply an arbitrary permutation to the order of the elements.

By injectivity of the embeddings, the sets{U𝒂:𝒂∈A∗,|𝒂|≤N}\{U_{{\bm{a}}}:{\bm{a}}\in A^{*},\,|{\bm{a}}|\leq N\}are pairwise disjoint. Consequently, there exists a continuous function

f¯:ℝN×d0→ℝdout\bar{f}:{\mathbb{R}}^{N\times d_{0}}\to{\mathbb{R}}^{d_{\mathrm{out}}}(76)such that for every𝒂∈A∗{\bm{a}}\in A^{*}with|𝒂|≤N|{\bm{a}}|\leq Nand every[𝒗1,…,𝒗N]⊤∈U𝒂[{\bm{v}}_{1},\dots,{\bm{v}}_{N}]^{\top}\in U_{{\bm{a}}},

f¯​(𝒗1,…,𝒗N)=f​(𝒂).\bar{f}({\bm{v}}_{1},\dots,{\bm{v}}_{N})=f({\bm{a}}).(77) Without loss of generality, we may assume thatf¯\bar{f}is permutation-invariant. That is, for every permutationσ∈SN\sigma\in S_{N}and every𝑽=[𝒗1,…,𝒗N]⊤∈ℝN×d0{\bm{V}}=[{\bm{v}}_{1},\dots,{\bm{v}}_{N}]^{\top}\in{\mathbb{R}}^{N\times d_{0}},

f¯​(𝒗1,…,𝒗N)=f¯​(𝒗σ​(1),…,𝒗σ​(N)).\bar{f}({\bm{v}}_{1},\dots,{\bm{v}}_{N})=\bar{f}({\bm{v}}_{\sigma(1)},\dots,{\bm{v}}_{\sigma(N)}).(78) Indeed, iff¯\bar{f}is not permutation-invariant, we can define

f~​(𝒗1,…,𝒗N)=1|SN|​∑σ∈SNf¯​(𝒗σ​(1),…,𝒗σ​(N)).\tilde{f}({\bm{v}}_{1},\dots,{\bm{v}}_{N})=\frac{1}{|S_{N}|}\sum_{\sigma\in S_{N}}\bar{f}({\bm{v}}_{\sigma(1)},\dots,{\bm{v}}_{\sigma(N)}).(79)Thenf~\tilde{f}is continuous and permutation-invariant. Moreover, for every sequence𝒂∈A∗{\bm{a}}\in A^{*}and every𝑽∈U𝒂{\bm{V}}\in U_{{\bm{a}}},

f~​(𝑽)=f¯​(𝑽),\tilde{f}({\bm{V}})=\bar{f}({\bm{V}}),(80)sinceU𝒂U_{{\bm{a}}}is closed under permutations.

Becausef¯\bar{f}is continuous and permutation-invariant, it follows fromZaheer et al. [2017]that there exists a pair of functionsϕ:ℝd0→ℝd1\phi:{\mathbb{R}}^{d_{0}}\to{\mathbb{R}}^{d_{1}}andρ:ℝd1→ℝdout\rho:{\mathbb{R}}^{d_{1}}\to{\mathbb{R}}^{d_{\mathrm{out}}}such that

f¯​(𝒗1,…,𝒗N)=ρ​(∑i=1Nϕ​(𝒗i))\bar{f}({\bm{v}}_{1},\dots,{\bm{v}}_{N})=\rho\!\left(\sum_{i=1}^{N}\phi({\bm{v}}_{i})\right)(81) (Note that some functions required1=O​(N)d_{1}=O(N)). We further assume without loss of generality thatϕ​(𝒗∅)=0\phi({\bm{v}}_{\emptyset})=0. We now construct a compressible transformer𝖬=(emb,Block1,Block2)\mathsf{M}=(\mathrm{emb},\mathrm{Block}_{1},\mathrm{Block}_{2})approximatingffusing this decomposition. The construciton of this transformer cloesly resembles the one used in the proof of Proposition3.3.

For the first blockBlock1\mathrm{Block}_{1}, we set all attention projection matrices to zero,

𝑾Q=𝑾K=𝑾V=𝑾O=𝟎,{\bm{W}}_{Q}={\bm{W}}_{K}={\bm{W}}_{V}={\bm{W}}_{O}=\mathbf{0},(82)so that the block reduces to its feedforward component.

Letϕ~:ℝd0→ℝ4​d0\tilde{\phi}:{\mathbb{R}}^{d_{0}}\to{\mathbb{R}}^{4d_{0}}be a continuous function satisfying

ϕ~​(emb​(a,i))=[ϕ​(emb​(a,i))𝒑i𝟎𝟎].\tilde{\phi}(\mathrm{emb}(a,i))=\begin{bmatrix}\phi(\mathrm{emb}(a,i))\\ {\bm{p}}_{i}\\ \mathbf{0}\\ \mathbf{0}\end{bmatrix}.(83)for alla∈Aa\in Aandi∈[N]i\in[N]where of the four vector blocks above are of sized0d_{0}. That is,ϕ~\tilde{\phi}computesϕ\phion the embedding while preserving positional information and padding with additional zeroes. We choose the feedforward mapFFN1\mathrm{FFN}_{1}in this block to approximateϕ~\tilde{\phi}to accuracyε′>0\varepsilon^{\prime}>0, to be specified later (This is possible as 2-layer feed forwand networks are universal approximators of continuous functions on compact sets[Cybenko,1989, Hornik,1991]). Thus, for any sequence𝒂=(a1,…,an){\bm{a}}=(a_{1},\dots,a_{n}),

Block1​(emb​(𝒂))≈[ϕ~​(emb​(a1,1))⋮ϕ~​(emb​(an,n))].\mathrm{Block}_{1}(\mathrm{emb}({\bm{a}}))\approx\begin{bmatrix}\tilde{\phi}(\mathrm{emb}(a_{1},1))\\ \vdots\\ \tilde{\phi}(\mathrm{emb}(a_{n},n))\end{bmatrix}.(84) For the second blockBlock2\mathrm{Block}_{2}, we again use a single attention head, and set

𝑾Q=𝑾K=𝟎,𝑾O=𝑰{\bm{W}}_{Q}={\bm{W}}_{K}=\mathbf{0},{\bm{W}}_{O}={\bm{I}}(85) The value projection𝑾V{\bm{W}}_{V}is chosen to extract theϕ\phi-component from the output ofϕ~\tilde{\phi}, and copy it to the thirdd1d_{1}-sized vector block, keep the first2​d12d_{1}coordinates the same, and zero out the remaining coordinates. That is:

𝑾V=(𝑰𝟎𝟎𝟎𝟎𝑰𝟎𝟎𝑰𝟎𝟎𝟎𝟎𝟎𝟎𝟎).{\bm{W}}_{V}=\begin{pmatrix}{\bm{I}}&\mathbf{0}&\mathbf{0}&\mathbf{0}\\ \mathbf{0}&{\bm{I}}&\mathbf{0}&\mathbf{0}\\ {\bm{I}}&\mathbf{0}&\mathbf{0}&\mathbf{0}\\ \mathbf{0}&\mathbf{0}&\mathbf{0}&\mathbf{0}\end{pmatrix}.(86) Since𝑾Q=𝑾K=𝟎{\bm{W}}_{Q}={\bm{W}}_{K}=\mathbf{0}, all attention scores are identical, and hence the attention weights are uniform. Applying the attention update inBlock2\mathrm{Block}_{2}to the outputs ofBlock1\mathrm{Block}_{1}and using a residual connection, the representation at positioniibecomes approximately

[ϕ~​(emb​(ai,i))𝒑i1n​∑j=1nϕ​(emb​(aj,j))𝟎].\begin{bmatrix}\tilde{\phi}(\mathrm{emb}(a_{i},i))\\ {\bm{p}}_{i}\\ \frac{1}{n}\sum_{j=1}^{n}\phi(\mathrm{emb}(a_{j},j))\\ \mathbf{0}\end{bmatrix}.(87) We now choose the feedforward network ofBlock2\mathrm{Block}_{2}to approximate up toε′′\varepsilon^{\prime\prime}(to be chosen later) a continuous functionρ~\tilde{\rho}satisfying

ρ~​(𝒘,𝒙,𝒚,𝒛)={ρ​(j⋅𝒚)if​𝒛=0​and​𝒙=𝒑j.ρ​((j−i+1)⋅𝒚)if​𝒛=𝒑ik+1​and​𝒙=𝒑j.\tilde{\rho}({\bm{w}},{\bm{x}},{\bm{y}},{\bm{z}})=\begin{cases}\rho(j\cdot{\bm{y}})&\text{ if }{\bm{z}}=0\text{ and }{\bm{x}}={\bm{p}}_{j}.\\ \rho((j-i+1)\cdot{\bm{y}})&\text{ if }{\bm{z}}=\frac{{\bm{p}}_{i}}{k+1}\text{ and }{\bm{x}}={\bm{p}}_{j}.\end{cases}(88) Intuitively, this map recovers the position indexiifrom the positional embedding𝒑i{\bm{p}}_{i}, rescales the averaged sum as necessary to recover the unweighted sum, and appliesρ\rho. In particular, at the final positionnn, the output of𝖬\mathsf{M}is approximately

ρ​(n⋅1n​∑j=1nϕ​(emb​(aj,j)))=ρ​(∑j=1nϕ​(emb​(aj,j)))=f​(𝒂).\rho\!\left(n\cdot\frac{1}{n}\sum_{j=1}^{n}\phi(\mathrm{emb}(a_{j},j))\right)=\rho\!\left(\sum_{j=1}^{n}\phi(\mathrm{emb}(a_{j},j))\right)=f({\bm{a}}).(89) By choosingε′,ε′′\varepsilon^{\prime},\varepsilon^{\prime\prime}sufficiently small, it follows that𝖬\mathsf{M}approximatesffto arbitrary precision.

We now show that there exists a compression policy withr​(n)=1r(n)=1that satisfies Equation72.

Define𝖢=(𝖼1,𝖼2)\mathsf{C}=(\mathsf{c}_{1},\mathsf{c}_{2})as follows. For any sequence of key–value pairs(𝑲,𝑽)({\bm{K}},{\bm{V}})with𝑽=[𝒗1,…,𝒗n]⊤{\bm{V}}=[{\bm{v}}_{1},\dots,{\bm{v}}_{n}]^{\top},

𝖼1​(𝑲,𝑽)=(𝟎,𝟎),\mathsf{c}_{1}({\bm{K}},{\bm{V}})=(\mathbf{0},\mathbf{0}),(90)and

𝖼2​(𝑲,𝑽)=(𝑲~,𝑽~),\mathsf{c}_{2}({\bm{K}},{\bm{V}})=(\tilde{{\bm{K}}},\tilde{{\bm{V}}}),(91)where𝑲~=𝟎\tilde{{\bm{K}}}=\mathbf{0}and

𝑽~=(𝑰𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎)​∑i=1n𝒗i+(𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝑰𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎)​𝒗n.\tilde{{\bm{V}}}=\begin{pmatrix}{\bm{I}}&\mathbf{0}&\mathbf{0}&\mathbf{0}\\ \mathbf{0}&\mathbf{0}&\mathbf{0}&\mathbf{0}\\ \mathbf{0}&\mathbf{0}&\mathbf{0}&\mathbf{0}\\ \mathbf{0}&\mathbf{0}&\mathbf{0}&\mathbf{0}\end{pmatrix}\sum_{i=1}^{n}{\bm{v}}_{i}\;+\;\begin{pmatrix}\mathbf{0}&\mathbf{0}&\mathbf{0}&\mathbf{0}\\ \mathbf{0}&\mathbf{0}&\mathbf{0}&{\bm{I}}\\ \mathbf{0}&\mathbf{0}&\mathbf{0}&\mathbf{0}\\ \mathbf{0}&\mathbf{0}&\mathbf{0}&\mathbf{0}\end{pmatrix}{\bm{v}}_{n}.(92) clearlyr​(n)=1r(n)=1.

For a prefix𝒂{\bm{a}}of lengthnnand a suffix𝒃{\bm{b}}of lengthkk, we begin by analyzing the forward pass of𝖬𝖢,𝒂​(𝒃)\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}}). After the first block, the representations of the tokensbib_{i}are given by

[ϕ~​(emb​(b1,1+n))⋮ϕ~​(emb​(bk,k+n))].\begin{bmatrix}\tilde{\phi}(\mathrm{emb}(b_{1},1+n))\\ \vdots\\ \tilde{\phi}(\mathrm{emb}(b_{k},k+n))\end{bmatrix}.(93)This follows because the embedding layer of𝖬𝖢,𝒂\mathsf{M}_{\mathsf{C},{\bm{a}}}is identical to that of𝖬\mathsf{M}, applied to the concatenated sequence[𝒂,𝒃][{\bm{a}},{\bm{b}}]. Moreover, the computation ofBlock1\mathrm{Block}_{1}is applied independently to each token representation, and is therefore unaffected by compression. Next, consider the attention computation in the second block. Since𝑾Q=0{\bm{W}}_{Q}=0, all attention scores are zero, and hence the attention weights are uniform, equal to11+k\frac{1}{1+k}. By the definition of𝖼2\mathsf{c}_{2}, it follows that after the attention update inBlock2\mathrm{Block}_{2}, the representation of each tokenbib_{i}is

[ϕ​(emb​(bi,n+i))𝒆n+i11+k​(∑j=1nϕ​(emb​(aj,j))+∑j=1kϕ​(emb​(bj,n+j)))1k+1​𝒑n].\begin{bmatrix}\phi(\mathrm{emb}(b_{i},n+i))\\ {\bm{e}}_{n+i}\\ \frac{1}{1+k}\biggl(\sum_{j=1}^{n}\phi(\mathrm{emb}(a_{j},j))+\sum_{j=1}^{k}\phi(\mathrm{emb}(b_{j},n+j))\biggr)\\ \frac{1}{k+1}{\bm{p}}_{n}\end{bmatrix}.(94)Finally, after applying the feed forward networkFFN\mathrm{FFN}, by choosingε′,ε′′\varepsilon^{\prime},\varepsilon^{\prime\prime}small enough, Equation88implies that the final output of𝖬𝖢,𝒂​(𝒃)\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})is approximately

ρ​((k+n−n+1)⋅1k+1​(∑j=1nϕ​(emb​(aj,j))+∑j=1kϕ​(emb​(bj,n+j))))≈f​([𝒂,𝒃])\rho\!\left((k+n-n+1)\cdot\frac{1}{k+1}\bigl(\sum_{j=1}^{n}\phi(\mathrm{emb}(a_{j},j))+\sum_{j=1}^{k}\phi(\mathrm{emb}(b_{j},n+j))\bigr)\right)\approx f([{\bm{a}},{\bm{b}}])(95) and is in anε\varepsilonapproximator of𝖬𝖢,𝒂​(𝒃)\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}}), completing the proof. Finally, after applying the feedforward networkFFN\mathrm{FFN}, and choosingε′,ε′′\varepsilon^{\prime},\varepsilon^{\prime\prime}sufficiently small, Equation88implies that the final output of𝖬𝖢,𝒂​(𝒃)\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})is approximately

and thus constitutes anε\varepsilon-approximation of𝖬𝖢,𝒂​(𝒃)\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}}), completing the proof.

∎

Lemma A.8.

LetAAbe a finite alphabet, and letf:⋃n≤NAn→ℝdoutf:\bigcup_{n\leq N}A^{n}\to\mathbb{R}^{d_{\text{out}}}be any sequence-to-vector function. Assume additionally that for some prefix sequence𝐚¯\bar{\bm{a}}of lengthn<Nn<Nthere exists two suffix sequences𝐛1{\bm{b}}^{1},𝐛2{\bm{b}}^{2}both of lengthk≤N−nk\leq N-nsuch that

f​([𝒂,𝒃1])≠f​([𝒂¯,𝒃2])f([{\bm{a}},{\bm{b}}^{1}])\neq f([\bar{\bm{a}},{\bm{b}}^{2}])(96) Then for everyε,C>0\varepsilon,C>0There exists a transformer such that

1. 1.(Approximation) For every sequence𝒂=(a1,…,an)∈An{\bm{a}}=(a_{1},\dots,a_{n})\in A^{n}withn≤Nn\leq N, ‖f​(𝒂)−𝖬​(𝒂)‖<ε.\|f({\bm{a}})-\mathsf{M}({\bm{a}})\|<\varepsilon.(97)
2. 2.(Non- compressibility) For every KV cache compression policy𝖢\mathsf{C}with compression budget satisfyingr​(n)<nr(n)<n there exists a suffix𝒃∈Ak{\bm{b}}\in A^{k}withk+n≤Nk+n\leq Nsuch that, ‖𝖬​([𝒂,𝒃])−𝖬𝖢,𝒂​(𝒃)‖>C.\|\mathsf{M}([{\bm{a}},{\bm{b}}])-\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}})\|>C.(98)

Proof.

By the same argument as in the proof of LemmaA.7, there exist functionsϕ:ℝd0→ℝd1\phi:{\mathbb{R}}^{d_{0}}\to{\mathbb{R}}^{d_{1}}andρ:ℝd1→ℝdout\rho:{\mathbb{R}}^{d_{1}}\to{\mathbb{R}}^{d_{\mathrm{out}}}such that for every𝒂=(a1,…,an){\bm{a}}=(a_{1},\dots,a_{n}),

f​(𝒂)=ρ​(∑i=1nϕ​(emb​(ai,i))).f({\bm{a}})=\rho\!\left(\sum_{i=1}^{n}\phi(\mathrm{emb}(a_{i},i))\right).(99) Let𝒂¯\bar{{\bm{a}}}be the prefix from the theorem statement, and let𝒃1,𝒃2∈Ak{\bm{b}}^{1},{\bm{b}}^{2}\in A^{k}be suffixes such that

f​([𝒂¯,𝒃1])≠f​([𝒂¯,𝒃2]).f([\bar{{\bm{a}}},{\bm{b}}^{1}])\neq f([\bar{{\bm{a}}},{\bm{b}}^{2}]).(100) It follows from equations96and99that

∑i=1kϕ​(emb​(bi1,i))≠∑i=1kϕ​(emb​(bi2,i)),\sum_{i=1}^{k}\phi(\mathrm{emb}(b^{1}_{i},i))\;\neq\;\sum_{i=1}^{k}\phi(\mathrm{emb}(b^{2}_{i},i)),(101)where𝒃1=(b11,…,bk1){\bm{b}}_{1}=(b^{1}_{1},\dots,b^{1}_{k})and𝒃2=(b12,…,bk2){\bm{b}}_{2}=(b^{2}_{1},\dots,b^{2}_{k}).

Define

H=conv​({∑i=1kϕ​(emb​(bi,i))|(b1,…,bk)∈Ak})+∑i=1nϕ​(emb​(ai,i)),H=\mathrm{conv}\!\left(\left\{\sum_{i=1}^{k}\phi(\mathrm{emb}(b_{i},i))\;\middle|\;(b_{1},\dots,b_{k})\in A^{k}\right\}\right)\;+\;\sum_{i=1}^{n}\phi(\mathrm{emb}(a_{i},i)),(102)and letDDdenote the diameter ofHH. Sinceρ\rhois continuous andHHis compact, the imageρ​(H)\rho(H)is also compact. Consequently, there exists𝒖0∈ℝdout{\bm{u}}_{0}\in{\mathbb{R}}^{d_{\mathrm{out}}}such that

dist​(𝒖0,ρ​(H))>C.\mathrm{dist}({\bm{u}}_{0},\rho(H))>C.(103) We construct the transformer𝖬\mathsf{M}similarly to the one constructed in the proof of LemmaA.7, with two modifications to the second block.

First, we replace𝑾O=𝑰{\bm{W}}_{O}={\bm{I}}with

𝑾O=(𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝟎𝑰𝟎𝟎𝟎𝟎𝟎).{\bm{W}}_{O}=\begin{pmatrix}\mathbf{0}&\mathbf{0}&\mathbf{0}&\mathbf{0}\\ \mathbf{0}&\mathbf{0}&\mathbf{0}&\mathbf{0}\\ \mathbf{0}&\mathbf{0}&{\bm{I}}&\mathbf{0}\\ \mathbf{0}&\mathbf{0}&\mathbf{0}&\mathbf{0}\end{pmatrix}.(104)Second, we modify the second feedforward function to approximate a continuous functionρ¯\bar{\rho}satisfying

ρ¯​(𝒘,𝒙,𝒚,𝒛)={ρ​(i⋅𝒚),if​‖𝒙−𝒑i‖<ε′​and​i⋅𝒚∈Hδ2,𝒖0,if​‖𝒙−𝒑i‖<ε′​and​i⋅𝒚∉Hδ,\bar{\rho}({\bm{w}},{\bm{x}},{\bm{y}},{\bm{z}})=\begin{cases}\rho(i\cdot{\bm{y}}),&\text{if }\|{\bm{x}}-{\bm{p}}_{i}\|<\varepsilon^{\prime}\text{ and }i\cdot{\bm{y}}\in H_{\frac{\delta}{2}},\\[5.0pt] {\bm{u}}_{0},&\text{if }\|{\bm{x}}-{\bm{p}}_{i}\|<\varepsilon^{\prime}\text{ and }i\cdot{\bm{y}}\notin H_{\delta},\\ \end{cases}(105) where

δ=(n+kn−1+k−1)​D2,\delta=\left(\frac{n+k}{n-1+k}-1\right)\frac{D}{2},(106)and

Hε:={𝒛∈ℝd:dist​(𝒛,H)<ε}.H_{\varepsilon}:=\{{\bm{z}}\in{\mathbb{R}}^{d}:\mathrm{dist}({\bm{z}},H)<\varepsilon\}.(107) By the same reasoning as in LemmaA.7, Equation97holds.

Let𝖢=(𝖼1,𝖼2)\mathsf{C}=(\mathsf{c}_{1},\mathsf{c}_{2})be any compression policy with compression budgetr​(n)<nr(n)<n. Since the first block of𝖬\mathsf{M}zeroes out all attention heads and depends only on the token embeddings,𝖼1\mathsf{c}_{1}has no effect on the forward pass of𝖬𝖢,𝒂​(𝒃)\mathsf{M}_{\mathsf{C},{\bm{a}}}({\bm{b}}). Thus, after the first block, the representations of the tokensbib_{i}are

[ϕ~​(emb​(b1,1+n))⋮ϕ~​(emb​(bk,k+n))].\begin{bmatrix}\tilde{\phi}(\mathrm{emb}(b_{1},1+n))\\ \vdots\\ \tilde{\phi}(\mathrm{emb}(b_{k},k+n))\end{bmatrix}.(108) In the second block, since𝑾Q=0{\bm{W}}_{Q}=0, all attention scores are identical, and hence the attention weights are uniform, equal to1k+r​(n)\frac{1}{k+r(n)}. It follows that the representation of each tokenbib_{i}after the second block is

ρ¯​(ϕ​(emb​(𝒃i,i+n)),𝒑i+n,1r​(n)+k​(𝒗~+∑j=1kϕ​(emb​(bj,j+n))),𝟎),\bar{\rho}\left(\phi(\mathrm{emb}({\bm{b}}_{i},i+n)),{\bm{p}}_{i+n},\frac{1}{r(n)+k}\left(\tilde{\bm{v}}+\sum_{j=1}^{k}\phi(\mathrm{emb}(b_{j},j+n))\right),\mathbf{0}\right),(109)where the compressed prefix cache is given by

𝖼2​(𝑲𝒂,𝑽𝒂)=(𝑲~𝒂,𝑽~𝒂)=([𝒌~1,…,𝒌~r​(n)]⊤,[𝒗~1,…,𝒗~r​(n)]⊤)\mathsf{c}_{2}({\bm{K}}_{\bm{a}},{\bm{V}}_{\bm{a}})=(\tilde{\bm{K}}_{\bm{a}},\tilde{\bm{V}}_{\bm{a}})=([\tilde{\bm{k}}_{1},\dots,\tilde{\bm{k}}_{r(n)}]^{\top},[\tilde{\bm{v}}_{1},\dots,\tilde{\bm{v}}_{r(n)}]^{\top})(110)and

𝒗~=(∑i=1r​(n)𝒗~i)2​d1:3​d1.\tilde{\bm{v}}=\bigl(\sum_{i=1}^{r(n)}\tilde{\bm{v}}_{i}\bigr)_{2d_{1}:3d_{1}}.(111) Sincen+kr​(n)+k>1\frac{n+k}{r(n)+k}>1by LemmaA.9, there exists a sequence𝒃¯\bar{{\bm{b}}}of lengthkksuch that

dist​(n+kr​(n)+k​(𝒗+∑j=1kϕ​(emb​(b¯j,j+n))),H)>(n+kr​(n)+k−1)​D2≥δ.\mathrm{dist}\!\left(\frac{n+k}{r(n)+k}\left({\bm{v}}+\sum_{j=1}^{k}\phi(\mathrm{emb}(\bar{b}_{j},j+n))\right),H\right)>\left(\frac{n+k}{r(n)+k}-1\right)\frac{D}{2}\;\geq\;\delta.(112) Therefore, by Equation105, the final representation of the tokenb¯k\bar{b}_{k}, which is the output of𝖬𝖢,𝒂¯​(𝒃¯)\mathsf{M}_{\mathsf{C},\bar{{\bm{a}}}}(\bar{{\bm{b}}}), is equal to𝒖0{\bm{u}}_{0}.

On the other hand, by construction,𝖬​([𝒂¯,𝒃¯])∈ρ​(H)\mathsf{M}([\bar{{\bm{a}}},\bar{{\bm{b}}}])\in\rho(H). Hence,

‖𝖬​([𝒂¯,𝒃¯])−𝖬𝖢,𝒂¯​(𝒃¯)‖>C,\|\mathsf{M}([\bar{{\bm{a}}},\bar{{\bm{b}}}])-\mathsf{M}_{\mathsf{C},\bar{{\bm{a}}}}(\bar{{\bm{b}}})\|>C,(113)completing the proof.

∎

Lemma A.9.

Let𝐱1,…,𝐱k∈ℝd{\bm{x}}_{1},\dots,{\bm{x}}_{k}\in\mathbb{R}^{d}be distinct points withk≥2k\geq 2, and let

H:=conv⁡{𝒙1,…,𝒙k}.H:=\operatorname{conv}\{{\bm{x}}_{1},\dots,{\bm{x}}_{k}\}.Letα>1\alpha>1. Then for every𝐯∈ℝd{\bm{v}}\in{\mathbb{R}}^{d}there existsi∈[k]i\in[k]satisfying

dist⁡(α​𝒙i+𝒗,H)≥α−12​diam⁡(H).\operatorname{dist}(\alpha{\bm{x}}_{i}+{\bm{v}},H)\geq\frac{\alpha-1}{2}\operatorname{diam}(H).

Proof.

Let

D:=diam⁡(H)=sup𝒑,𝒒∈H‖𝒑−𝒒‖.D:=\operatorname{diam}(H)=\sup_{{\bm{p}},{\bm{q}}\in H}\|{\bm{p}}-{\bm{q}}\|.Since the points𝒙1,…,𝒙k{\bm{x}}_{1},\dots,{\bm{x}}_{k}are distinct andk≥2k\geq 2, we haveD>0D>0.

Choose𝒑,𝒒∈H{\bm{p}},{\bm{q}}\in Hsuch that‖𝒑−𝒒‖=D\|{\bm{p}}-{\bm{q}}\|=D(such a pair exists sinceHHis compact), and define

𝒖:=𝒑−𝒒‖𝒑−𝒒‖.{\bm{u}}:=\frac{{\bm{p}}-{\bm{q}}}{\|{\bm{p}}-{\bm{q}}\|}.The set{⟨𝒖,𝒙⟩∣𝒙∈H}\{\langle{\bm{u}},{\bm{x}}\rangle\mid{\bm{x}}\in H\}forms a closed interval, and the width ofHHin direction𝒖{\bm{u}}is

max𝒙∈H⁡⟨𝒖,𝒙⟩−min𝒙∈H⁡⟨𝒖,𝒙⟩=D.\max_{{\bm{x}}\in H}\langle{\bm{u}},{\bm{x}}\rangle-\min_{{\bm{x}}\in H}\langle{\bm{u}},{\bm{x}}\rangle=D. For any𝒗∈ℝd{\bm{v}}\in{\mathbb{R}}^{d}, the width ofα​H+𝒗\alpha H+{\bm{v}}in direction𝒖{\bm{u}}is therefore

max𝒙∈α​H+𝒗⁡⟨𝒖,𝒙⟩−min𝒙∈α​H+𝒗⁡⟨𝒖,𝒙⟩=α​D.\max_{{\bm{x}}\in\alpha H+{\bm{v}}}\langle{\bm{u}},{\bm{x}}\rangle-\min_{{\bm{x}}\in\alpha H+{\bm{v}}}\langle{\bm{u}},{\bm{x}}\rangle=\alpha D. Let

ε:=α−12​D.\varepsilon:=\frac{\alpha-1}{2}D.Suppose, for contradiction, that

dist⁡(α​𝒙i+𝒗,H)<εfor all​i=1,…,k.\operatorname{dist}(\alpha{\bm{x}}_{i}+{\bm{v}},H)<\varepsilon\qquad\text{for all }i=1,\dots,k.Since

α​H+𝒗=conv⁡{α​𝒙1+𝒗,…,α​𝒙k+𝒗},\alpha H+{\bm{v}}=\operatorname{conv}\{\alpha{\bm{x}}_{1}+{\bm{v}},\dots,\alpha{\bm{x}}_{k}+{\bm{v}}\},and since theε\varepsilon-neighborhood ofHH,

Hε:={𝒛∈ℝd:dist⁡(𝒛,H)<ε},H_{\varepsilon}:=\{{\bm{z}}\in\mathbb{R}^{d}:\operatorname{dist}({\bm{z}},H)<\varepsilon\},is convex, it follows that

α​H+𝒗⊆Hε.\alpha H+{\bm{v}}\subseteq H_{\varepsilon}.Consequently, the width ofα​H+𝒗\alpha H+{\bm{v}}in direction𝒖{\bm{u}}is strictly smaller than the width ofHεH_{\varepsilon}in direction𝒖{\bm{u}}. This width is at most

Therefore,

α​D<D+2​ε.\alpha D<D+2\varepsilon.But by the definition ofε\varepsilon,

D+2​ε=D+(α−1)​D=α​D,D+2\varepsilon=D+(\alpha-1)D=\alpha D,which is a contradiction. Hence there exists somei∈{1,…,k}i\in\{1,\dots,k\}such that

dist⁡(α​𝒙i+𝒗,H)≥ε.\operatorname{dist}(\alpha{\bm{x}}_{i}+{\bm{v}},H)\geq\varepsilon.∎

Note.We believe that a finer-grained analysis of transformer compressibility may be possible through the theory of Chebyshev systems[Karlin and Studden,1966b], a classical framework with broad applications across both pure and applied mathematics[Karlin and Studden,1966a, Micchelli and Pinkus,1977, Eitan,2021, Karlin and Ziegler,1966]. Chebyshev systems provide tools for controlling the zeros, sign changes, and interpolation structure of linear combinations of functions, which makes them a natural candidate for studying the degrees of freedom available in attention-like mixtures of value functions. We leave a systematic development of this connection to future work.

Appendix BKV-CATimplementation details

Masked attention forward pass.

Instead, at each layer group, the train-time KV sparsification policy defines which previous KV slots are visible to attention. For layerℓ\ell, letmj(ℓ)m_{j}^{(\ell)}be the most recent mask produced by the sparsification policy. The attention for query tokenttis evaluated over

𝒜tℓ={j≤t:mj(ℓ)=1}\mathcal{A}_{t}^{\ell}=\{j\leq t:m_{j}^{(\ell)}=1\}(114)Thus all valid queries are still processed by the transformer, but their access to past KV slots is restricted to the active, unmasked KV slots. The masks are independent across compression points: a token dropped in one layer may be selected again in a later layer. No masking is used at evaluation time in our comparisons.

Router parameterization.

For the main experiments in the paper, we use lightweight learnable router modules, parameterized as linear attention, as our KV sparsification policy. For a sequencex1:Tx_{1:T}, lethtℓ∈ℝdh_{t}^{\ell}\in{\mathbb{R}}^{d}denote the hidden state of tokenttbefore decoder layerℓ\ell. We insert routers at a small set of compression layers𝒞{\mathcal{C}}(e.g., forQwen2.5-0.5Bwe use𝒞={1,7,13,19}{\mathcal{C}}=\{1,7,13,19\}). Each router independently predicts which KV slots are masked for the following group of layers. Given hidden statesHℓ=(h1ℓ,…,hTℓ)H^{\ell}=(h_{1}^{\ell},\ldots,h_{T}^{\ell})and, the router first computes

h~t=LN​(htℓ),qt=ϕ​(WQ​h~t),kt=ϕ​(WK​h~t),rt=WV​h~t,ϕ​(z)=ELU​(z)+1.\tilde{h}_{t}=\mathrm{LN}(h_{t}^{\ell}),\qquad q_{t}=\phi(W_{Q}\tilde{h}_{t}),\quad k_{t}=\phi(W_{K}\tilde{h}_{t}),\quad r_{t}=W_{V}\tilde{h}_{t},\qquad\phi(z)=\mathrm{ELU}(z)+1.(115)It then forms a causal linear-attention summary without storing a router KV cache,

St=∑j≤tkj​rj⊤,zt=∑j≤tkj,at=WO​qt⊤​Stqt⊤​zt+ϵ.S_{t}=\sum_{j\leq t}k_{j}r_{j}^{\top},\qquad z_{t}=\sum_{j\leq t}k_{j},\qquad a_{t}=W_{O}\frac{q_{t}^{\top}S_{t}}{q_{t}^{\top}z_{t}+\epsilon}.(116)The keep probability is computed from a cosine score between a pointwise projection of the current state and the state after the router summary:

ut=WP​htℓ‖WP​htℓ‖2,wt=htℓ+α​at‖htℓ+α​at‖2,pt=1−⟨ut,wt⟩2.u_{t}=\frac{W_{P}h_{t}^{\ell}}{\|W_{P}h_{t}^{\ell}\|_{2}},\qquad w_{t}=\frac{h_{t}^{\ell}+\alpha a_{t}}{\|h_{t}^{\ell}+\alpha a_{t}\|_{2}},\qquad p_{t}=\frac{1-\langle u_{t},w_{t}\rangle}{2}.(117)The binary routing decision is

mtℓ=𝟏​{pt>τ},m_{t}^{\ell}=\mathbf{1}\{p_{t}>\tau\},(118)withτ=0.5\tau=0.5in all reported runs. During training, this hard threshold is optimized with a straight-through estimator: the forward pass uses the binary mask above, while the backward pass passes gradients through the pre-threshold keep probabilityptp_{t}to the router parameters. We initializeWP=−IW_{P}=-Iandα=0\alpha=0, which givespt=1p_{t}=1for every valid token, so training starts from the dense model.

Training objective.

Letpθmask(⋅∣x<t)p_{\theta}^{\mathrm{mask}}(\cdot\mid x_{<t})denote the masked forward pass andpθdense(⋅∣x<t)p_{\theta}^{\mathrm{dense}}(\cdot\mid x_{<t})the same model with masking disabled. We optimize

ℒ=\displaystyle\mathcal{L}={}λmask1T−1∑t=2TDKL(sg[pθdense(⋅∣x<t)]∥pθmask(⋅∣x<t))\displaystyle\lambda_{\mathrm{mask}}\frac{1}{T-1}\sum_{t=2}^{T}D_{\mathrm{KL}}\!\left(\operatorname{sg}\!\left[p_{\theta}^{\mathrm{dense}}(\cdot\mid x_{<t})\right]\,\middle\|\,p_{\theta}^{\mathrm{mask}}(\cdot\mid x_{<t})\right)(119)+λanchor​1T−1​∑t=2T−log⁡pθdense​(xt∣x<t)+λbudget​1|𝒞|​∑c∈𝒞ℬc.\displaystyle+\lambda_{\mathrm{anchor}}\frac{1}{T-1}\sum_{t=2}^{T}-\log p_{\theta}^{\mathrm{dense}}(x_{t}\mid x_{<t})+\lambda_{\mathrm{budget}}\frac{1}{|\mathcal{C}|}\sum_{c\in\mathcal{C}}\mathcal{B}_{c}.wheresg\operatorname{sg}stops gradients through the dense teacher. When using the router based KV sparsification policy, we useλbudget≠0\lambda_{\mathrm{budget}}\neq 0and the budget term follows the load-balancing objective fromHwang et al. [2025]. For target keep rateρ\rho, define

Fc=1T​∑t=1Tmtc,Gc=1T​∑t=1Tptc,ℬc=Fc​Gcρ+(1−Fc)​(1−Gc)1−ρ.F_{c}=\frac{1}{T}\sum_{t=1}^{T}m_{t}^{c},\qquad G_{c}=\frac{1}{T}\sum_{t=1}^{T}p_{t}^{c},\qquad\mathcal{B}_{c}=\frac{F_{c}G_{c}}{\rho}+\frac{(1-F_{c})(1-G_{c})}{1-\rho}.(120)All transformer and router parameters are updated during continued pretraining; in the reported runsρ=0.5\rho=0.5,λmask=1\lambda_{\mathrm{mask}}=1,λanchor=1\lambda_{\mathrm{anchor}}=1, andλbudget=0.1\lambda_{\mathrm{budget}}=0.1.

Appendix CExtended Experimental Details

This section provides additional details for the empirical evaluation described in Section5. We describe the continued-pretraining setup, including the data, architecture hyperparameters, and compression objective, and then give additional details on the post-hoc compression evaluations and experimental setups used throughout the paper. Unless otherwise stated, evaluations compare the base model to theKV-CATcheckpoint with masking disabled. All experiments were run on up to 8 H100 GPUs.

Table 5:KV-CATcontinued-pretraining hyperparameters forQwen2.5checkpoints.### C.1Continued pretraining runs

We train two model sizes, initialized from the publicQwen2.5-0.5BandQwen2.5-1.5B[Team,2024]checkpoints, using continued pretraining on FineWeb-Edu[Penedo et al.,2024]at a context length of 1024. The compression-aware objective is the one described in the method section: a compressed self-distillation loss matched to the model’s dense distribution, an uncompressed next-token prediction anchor, and a budget penalty. All transformer and compression module parameters are updated.

The router layers are inserted at layers 0, 6, 12, and 18 forQwen2.5-0.5Band layers 0, 7, 14, and 21 forQwen2.5-1.5B. Each module is a causal linear-attention boundary predictor (as described in AppendixB) with 64 feature dimensions, a thresholdτ=0.5\tau=0.5, and a target keep rate of 50%. The objective weights areλmask=1\lambda_{\mathrm{mask}}=1,λanchor=1\lambda_{\mathrm{anchor}}=1, andλbudget=0.1\lambda_{\mathrm{budget}}=0.1. We use AdamW with peak learning rate10−410^{-4}, 600 warmup steps, minimum learning rate5×10−65\times 10^{-6}, weight decay 0.01, and gradient clipping at norm 1.0. Both model sizes use a batch of 131,072 tokens per optimizer step, corresponding to 128 sequences of length 1024, and are trained on 8 GPUs. Runs are configured for 40k optimizer steps, which see5.24×1095.24\times 10^{9}tokens.

At evaluation time, the boundary predictors are turned off, and the post-hoc compressor named in each experiment is the only cache-reduction procedure whose quality is measured. Figure3reports the compressed and uncompressed validation NTP loss as well as the KV retention ratio throughout theKV-CATcontinued pretraining forQwen2.5-0.5B. Table5reports all hyperparameters and training setup for both runs.

(a)Dense NTP loss (b)Compressed NTP loss (c)KV retention ratio

Figure 3:Qwen2.5-0.5BKV-CATcontinued pretraining run.We plot dense validation next-token-prediction loss, compressed-path validation next-token-prediction loss, and the realized KV retention ratio over the 5.24B-tokenKV-CATcontinued-pretraining run. The retention trace reports the fraction of KV slots kept by the learned routers, whose budget target is 50%.

C.2No compression QA evaluation

To check that our training procedure does not substantially degrade ordinary model behavior, we evaluate normalized multiple-choice accuracy on a variety of question answering benchmarks using the LightEval harness[Habib et al.,2023]. The evaluation suite contains HellaSwag[Zellers et al.,2019], WinoGrande[Sakaguchi et al.,2020], PIQA[Bisk et al.,2020], Social IQa[Sap et al.,2019], OpenBookQA[Mihaylov et al.,2018], ARC-Easy and ARC-Challenge[Clark et al.,2018]. We use 1000 validation examples per task, except for OpenBookQA, whose validation split contains 500 examples in this setup. The reported average is the unweighted mean over these eight tasks. The purpose of this evaluation is to measure whether the trained checkpoints maintain language modeling performance on standard short-context multiple-choice tasks before we test their behavior under hoc KV cache compression.

C.3Suffix perplexity under prefix KV cache compression

The suffix-prediction experiments evaluate whether the continued-pretrained checkpoints are easier targets for post-hoc prefix KV cache compression. The in-domain evaluation data consists of held-out blocks from FineWeb. Each evaluation example is a 1024-token excerpt split into a 768-token prefix and a 256-token suffix. We first run the model on the full 1024-token sequence and record the suffix logits and per-layer attention traces. The first 768 tokens define the prefix cache to be compressed. The following 256 suffix tokens are never removed or compressed: after a compact prefix cache has been constructed, the same suffix tokens are appended normally and scored under the original transformer layers. This protocol is shared by the Attention Matching table and the gradient based optimization step-curve figures, but the two experiments fit the compact cache differently.

Attention-Matching.

We use the AM-HighestAttnKeys Attention matching variant fromZweiger et al. [2026]. For a keep ratioρ∈{0.05,0.1,0.2,0.4}\rho\in\{0.05,0.1,0.2,0.4\}, Attention Matching constructs a compact prefix cache independently for each layer and key-value head. LetK,V∈ℝL×dK,V\in\mathbb{R}^{L\times d}denote the dense prefix keys and values for one key-value head, whereL=768L=768, and letQ∈ℝM×dQ\in\mathbb{R}^{M\times d}denote suffix query states from the corresponding grouped query heads. We subsample at most 256 suffix queries per key-value head. Attention Matching first selects⌈ρ​L⌉\lceil\rho L\rceilsource keys with the largest root-mean-square attention weight over the sampled suffix queries. The selected keys form the compact keysC1C_{1}. A scalar log-bias vectorβ\betais then fit by two iterations of box-constrained nonnegative least squares so that the compact cache approximately matches the dense attention normalizer. Finally, compact valuesC2C_{2}are fit by ridge least squares to reconstruct the dense attention outputs,

softmax⁡(Q​C1⊤dk+β)​C2≈softmax⁡(Q​K⊤dk)​V.\operatorname{softmax}\!\left(\frac{QC_{1}^{\top}}{\sqrt{d_{k}}}+\beta\right)C_{2}\approx\operatorname{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_{k}}}\right)V.We use ridge coefficient10−410^{-4}with spectral scaling and solve the value regression with a least-squares solver. After constructing the compact prefix cache, we run the model on the same prefix-suffix sequence while replacing the dense prefix cache with the compact cache for all subsequent suffix positions. We report the increase in suffix perplexity relative to the same model’s dense-prefix forward pass, KL divergence to the native dense-prefix logits, and top-1 agreement with the native logits. All metrics are averaged over 128 evaluation examples.

Gradient-based KV cache compression.

In this setting, we directly optimize a small set of continuous key and value vectors for each example while keeping all model parameters fixed. The resulting optimization curves measure how much information about the prefix can be preserved in a compact KV cache when the cache itself is allowed to adapt to the example. The dense model is first run on the full prefix-suffix sequence to produce teacher logits on the suffix. For a target keep ratio

ρ∈{0.05,0.1,0.15,0.2,0.25,0.3,0.35,0.4,0.5},\rho\in\{0.05,0.1,0.15,0.2,0.25,0.3,0.35,0.4,0.5\},we allocate a compact prefix cache with

m=⌈ρ​L⌉m=\left\lceil\rho L\right\rceilKV slots per layer and key-value head, whereL=768L=768is the dense prefix length. The compact cache consists of learnable key and value tensors

K~ℓ,h,V~ℓ,h∈ℝm×d\widetilde{K}_{\ell,h},\widetilde{V}_{\ell,h}\in\mathbb{R}^{m\times d}for each layerℓ\elland key-value headhh. We initializeK~ℓ,h\widetilde{K}_{\ell,h}andV~ℓ,h\widetilde{V}_{\ell,h}from the firstmmdense prefix KV states for that layer and head. The compact cache is optimized per example using the suffix teacher distribution from the dense full-prefix forward pass. Letzt⋆z_{t}^{\star}denote the dense teacher logits at suffix positiontt, and letzt​(K~,V~)z_{t}(\widetilde{K},\widetilde{V})denote the logits obtained when the dense prefix cache is replaced by the compact cache. We minimize

ℒKV=1|𝒮|∑t∈𝒮DKL(softmax(zt⋆)∥softmax(zt(K~,V~))),\mathcal{L}_{\mathrm{KV}}=\frac{1}{|\mathcal{S}|}\sum_{t\in\mathcal{S}}D_{\mathrm{KL}}\!\left(\operatorname{softmax}(z_{t}^{\star})\,\middle\|\,\operatorname{softmax}(z_{t}(\widetilde{K},\widetilde{V}))\right),where𝒮\mathcal{S}is the set of suffix prediction positions. The optimization updates only the compact KV tensors; model weights are never updated. We use AdamW for 100 gradient steps with learning rate10−210^{-2}, weight decay0, and gradient clipping at norm1.01.0. We record intermediate compact caches after steps

{0,1,2,5,10,20,50,100}\{0,1,2,5,10,20,50,100\}to produce the optimization-step curves. After each recorded optimization step, we evaluate next-token prediction on the same suffix using the compact cache in place of the dense prefix cache. The reported perplexity is computed against the ground-truth suffix tokens.

This evaluation should be interpreted as a controlled probe of the cache-construction problem in post-hoc compression. For each example, the compact prefix cache is fit using supervision derived from the same prefix–suffix sequence on which it is evaluated: Attention Matching uses suffix attention/query traces to construct the compact cache, while the gradient-based method uses suffix teacher logits. Thus, the comparison asks whether the trained checkpoints make the per-example compact-cache optimization problem easier under matched cache budgets and matched supervision, rather than testing transfer from a cache fit on one suffix to unseen suffixes (which is tested in the next experiments).

C.4Needle-in-a-haystack retrieval under KV cache compression

We evaluate retrieval under prefix KV cache compression using a deterministic needle-in-a-haystack task. Each example contains neutral filler text, passkey-like distractor numbers, a single six-digit passkey, and a final query asking for the passkey. The prompt length before the final query is 1024 tokens. The model is evaluated using the compressed prompt and the query.

Example construction.

To generate each example, we sample a six-digit stringa∈{000000,…,999999}a\in\{\texttt{000000},\ldots,\texttt{999999}\}, which defines the passkey stored in the prefix and the answer expected at evaluation. For example, ifa=483920a=\texttt{483920}, the needle inserted into the haystack is

Memory record: special_passkey=483920.The final query is

Memory record: special_passkey=and the answer is exactly

The passkey is inserted at relative depths

d∈{0,0.25,0.5,0.75,1},d\in\{0,0.25,0.5,0.75,1\},whered=0.0d=0.0places the needle at the beginning of the haystack andd=1d=1places it at the end. In addition to the answer passkey, examples contain distractors: passkey-like six-digit strings that are explicitly not equal to the answer. These distractors are inserted into the sequence so that the model cannot solve the task by copying an arbitrary number. For example, if the answer is483920, a possible distractor sentences is or

Ignore the misleading special passkey candidate 672041.During example construction, the haystack is assembled from filler text before and after the needle. At each filler-sentence draw, with probability0.350.35we insert a distractor sentence; otherwise we insert a neutral sentence unrelated to the retrieval key. Distractor codes are sampled uniformly from the six-digit code space and rejected if they match the answer. The needle is then inserted at the requested relative depth in the haystack, and the query is appended after the haystack. For each compression ratio, we evaluate 20 instances per depth totaling at a 100 examples.

Training the compact KV cache.

For each example, we train a new compact prefix KV cache while keeping all model weights frozen. Let<haystack>denote the haystack prefix, including the filler text, needle and any distractors, but excluding the final query and answer. We construct an reconstruction sequence

<haystack><instruction><haystack>,\texttt{<haystack><instruction><haystack>},where the instruction is

Reproduce the preceding passage verbatim.Preserve every word, digit, and punctuation mark.The compact cache replaces the KV states for the first<haystack>. The loss is applied only on the second<haystack>, so the optimizer must store enough information in the compact cache to reconstruct the original prefix, including the passkey and distractors.

For a keep ratioρ\rho, the compact cache contains

m=⌈ρ​L⌉m=\lceil\rho L\rceilKV slots per layer and KV head, whereLLis the haystack length. The compact keys and values are initialized by selectingmmrandom haystack KV states and then optimized directly with AdamW. As in AppendixC.3, we use a KL objective between the suffix logits (computed on the second copy of<haystack>) produced by the dense model and the ones produced when attending to the compact cache. For each example, this objective is optimized by running AdamW for 100 steps with learning rate10−210^{-2}, weight decay0, and gradient clipping at norm1.01.0.

Evaluation.

At evaluation time, the haystack prefix is represented by the optimized compact KV cache, while the final query is appended normally and remains uncompressed. We then perform greedy answer generation and compare the generated completion against the six-digit passkey. In the table, we report native-success-conditioned exact-match retrieval accuracy: for each model, keep ratio, and compact-cache optimization step, we report the percentage of those examples for which the compressed-prefix model also generates exactly the passkey.

C.5Long-form question answering under KV cache compression

We evaluate long-context question answering under post-hoc KV cache compression on LongBench v2. Because the compression procedure optimizes a separate compact KV cache for each prompt, this evaluation is substantially more expensive than ordinary decoding. We therefore report results on a fixed subset of seven LongBench v2 subdomains: Academic, Agent history QA, Knowledge graph reasoning, Legal, Many-shot learning, New language translation, and Table QA. This subset contains 221 examples in total.

For each LongBench example, we construct a multiple-choice prompt from the context, question, and four answer choices. The prefix is

Please read the following text and answer the question below.followed by the context from the benchmark. The suffix is

What is the correct answer to this question: <question>followed by the four choices(A),(B),(C), and(D), and the answer-format prefix

The correct answer is (.\texttt{The correct answer is (}.We score the four continuationsA),B),C), andD)by total log-likelihood and predict the answer with the highest score. No free-form generation is used.

For each example, model, and keep ratio∈{0.1,0.2,0.5}\in\{0.1,0.2,0.5\}, we optimize a compact KV cache for example’s context while keeping all model weights frozen. The optimization uses a reconstruction sequence as in AppendixC.4. As common in LongBench evaluations, if the reconstruction sequence length is larger than the context length of the model (32,768 tokens) we truncate the context using middle truncation (keeping a prefix and suffix of equal length)[Bai et al.,2025]. The KC objective is a KL divergence between the dense model’s next-token distribution and the compact-cache model’s next-token distribution on all tokens of the repeated prefix. The compressed caches are initialized from the first prefix KV slots. We optimize each compact cache for 100 AdamW steps with learning rate10−210^{-2}, weight decay 0, gradient clipping norm 1.0, and gradient checkpointing enabled. At evaluation time, the optimized compact KV cache replaces the full KV cache of the context prefix. The question, answer choices, and answer prefix are then processed normally, without compression. We report accuracy, grouped by subdomain and averaged over the 221 examples.

Appendix DAdditional Experimental Results

This section collects supplementary results that are complementary to Section5. In SectionD.1, we ablate the choice of the train-time sparsification policy used inKV-CAT. In SectionD.2, we test whether the Attention Matching performance gains we observe in Section5generalize to prefix–suffix from additional text corpora.

D.1Train-time KV sparsification policy ablation

We use this ablation to choose the train-time KV sparsification policy used in the main experiments. Starting fromQwen2.5-0.5B, we train three checkpoints with the same self-distillation and dense-anchor losses and the same 50% target keep rate, varying only the policy that selects active KV slot during the masked forward pass.Randuniformly samples an exact-count 50% subset of valid KV slots at each sparsification point.Attnis an H2O-inspired[Zhang et al.,2023]attention-mass baseline: it computes dense causal attention at the routed layer, sums the attention mass received by each source token over heads and valid query positions, and keeps the top 50% source KV slots.Routeruses the lightweight learned linear-attention routers described in AppendixB. BecauseRandandAttnsatisfy the target keep rate by construction, we setλbudget=0\lambda_{\mathrm{budget}}=0for those runs; forRouter, we use the budget regularizer described in AppendixB. In all comparisons below, the trained checkpoints are evaluated in the unmasked forward pass mode, with the training-time sparsification mechanism disabled, and post-hoc compression is applied afterward.

Table 6:KV-CATimplemented with aRouterpolicy retains base model performance better than the fixedRandandAttnpolicies.#### Uncompressed performance.

Routeris the only sparsification policy in this ablation that preserves dense downstream accuracy at or above the base model average. Its average score is 52.2, compared to 51.5 for the base model, 50.7 forRand, and 50.2 forAttn. The fixed policies still retain much of the base model’s zero-shot performance after continued pretraining.

Attention Matching Evaluation.

All checkpoints are evaluated with the same Attention Matching procedure described in AppendixC.3. Table7shows that all three train-time sparsification policies Attention Matching performance compared to the original base checkpoint for every reported metric and budget. This further supports theKV-CATframework, demonstrating that the downstream compression improvements are not tied exclusively to learned router parameterization. Among the policies,Routergives the best result for every reported keep ratio and metric. The margin overRandis sometimes small, especially at 10% and 40% keep ratio, butRouteris consistently at least as good under compression and is clearly stronger on dense task retention. We therefore useRouterin the main experiments as the default train-time sparsification policy.

Table 7:Attention-matching prefix compression on FineWeb-Edu for the base model and threeKV-CATmodels with theRouter,Rand, andAttnKV sparsification policies. We report degradation relative to each model’s native dense-prefix forward pass. LowerΔ\DeltaPPL and KL are better; higher top- 1 agreement is better.

D.2Cross-domain attention matching evaluation

We test whether the Attention Matching optimization gains observed forKV-CATtransfer to prefix–suffix pair from additional textual domains. We evaluate Attention Matching KV cache compression on validation examples from three held-out corpora: WikiText-103[Merity et al.,2016], PG-19[Rae et al.,2019], and arXiv[Cohan et al.,2018]. For each example, we take a 768-token prefix and a 256-token suffix, compress only the prefix KV cache with the same Attention Matching procedure used in AppendixC.3, and then score suffix next-token prediction against each model’s own dense full-prefix distribution.

Across all three held-out domains, the compression-aware checkpoint remains a better target for the same post-hoc compressor. It reduces KL and increases top-1 agreement for every corpus and keep ratio. It also lowersΔ\DeltaPPL in 11 of 12 settings; the only exception is arXiv at 40% keep ratio, whereΔ\DeltaPPL increases even though KL and top-1 agreement still improve. We therefore interpret this table as evidence thatKV-CATimproves cross-domain compressibility under Attention Matching.

Table 8:Attention Matching cross-domain completion under prefix-cache compression.We compare the baseQwen2.5-1.5Bmodel and theKV-CAT-trainedQwen2.5-1.5Bcheckpoint, on held-out prefix–suffix pair from different domains. Each example contains a 768-token prefix and a 256-token suffix; Attention Matching compresses only the prefix KV cache. Metrics are computed relative to each model’s native dense full-prefix suffix distribution, so lowerΔ\DeltaPPL and KL and higher top-1 agreement are better.

D.3Additional KV keep-ratio curves

Figure4extends the main-body optimization curves in Figure2with three additional KV keep ratios: 15%, 25%, and 35%.

Figure 4:KV-CATspeeds up gradient-based KV cache compression.We plot the gap in suffix perplexity under full/compressed-prefix inference throughout gradient-based KV cache optimization. Each panel fixes a different KV keep ratio. Across ratios, theKV-CATcheckpoint achieves a comparableΔ\DeltaPPL in fewer optimization steps than the base model, yielding up to a5×\timesspeedup.

Appendix ELimitations

Our approach relies on continued pretraining with additional routing modules, introducing non-trivial training overhead and potentially limiting applicability in low-resource settings. The routing mechanism and training objective also add implementation complexity, which may hinder integration into existing production systems. Finally, our evaluation is limited to a small set of model sizes and tasks, and it remains to be seen how well the method generalizes to larger-scale models or different domains.

Appendix FBroader Impact

Broader Impact.

This work is primarily methodological, aiming to improve the efficiency of transformer-based language models by enabling more effective KV cache compression. On the positive side, our approach can reduce memory and compute requirements for long-context inference, making large models more accessible and energy-efficient, and enabling their deployment in resource-constrained settings. Potential negative impacts are indirect: by improving the efficiency of existing LLMs, our method may lower the cost of deploying such models at scale, which could amplify both beneficial and harmful downstream applications. However, our method does not introduce new capabilities or alter model behavior beyond improving compatibility with compression, and therefore does not create new misuse vectors beyond those already present in the underlying models. Mitigation strategies largely align with existing practices for responsible LLM deployment, including monitoring, access control, and adherence to usage policies of the underlying models.