KV缓存压缩的风险

# KV缓存压缩的风险  
来源：https://arxiv.org/html/2607.01520  
Lukas Haverbeck  
亚琛工业大学  
lukas\.haverbeck@rwth\-aachen\.de  

&  
Carmen Amo Alonso  
斯坦福大学  
camoalon@stanford\.edu  

Andres Felipe Posada\-Moreno  
亚琛工业大学  
andres\.posada@dsme\.rwth\-aachen\.de  

&  
Sebastian Trimpe  
亚琛工业大学  
trimpe@dsme\.rwth\-aachen\.de  

&  
Marco Pavone  
斯坦福大学  
pavone@stanford\.edu  

###### 摘要  
Transformer对长序列的推理成本高昂，因为softmax注意力会反复读取一个庞大的KV缓存。解决这一瓶颈的主流方法是*KV缓存压缩（KV cache compression）*，即用紧凑的摘要替换完整的缓存。尽管其在实际中很重要，但这种摘要的设计在很大程度上仍依赖经验性实验。在理论方面，现有结果表明KV缓存压缩在最坏情况下是不可能实现的，但很少能为在可实现精确压缩的机制下设计算法提供系统性指导。我们通过刻画KV缓存压缩的*极小极大风险（minimax risk）*来弥合这一差距，该风险以缓存的固有可压缩性来表征，揭示了精确压缩何时以及如何成为可能。这些结果在因果掩码（causal masking）下为KV缓存压缩提供了新颖的设计原则，这些原则能高效映射到预填充（prefill）和自回归解码（autoregressive decoding），同时实现极小极大最优风险。我们将这些原则实例化为一个实用算法，并在LongBench上的针对性实验中报告了有希望的性能。总体而言，我们的结果为具有理论保证的实用KV缓存压缩提供了一条原则性途径。  

## 1 引言  
Transformer已成为跨多个领域处理和生成长序列的主流架构[14 (https://arxiv.org/html/2607.01520#bib.bib14),5 (https://arxiv.org/html/2607.01520#bib.bib5),22 (https://arxiv.org/html/2607.01520#bib.bib22),4 (https://arxiv.org/html/2607.01520#bib.bib4)]，这得益于它们对token之间复杂交互的出色建模能力。在推理过程中，这些交互通过存储所有先前token的*键-值（KV）缓存*来中介，每个新token通过softmax注意力从该缓存计算得出[23 (https://arxiv.org/html/2607.01520#bib.bib23)]。因此，KV缓存充当了模型对过去的工作记忆。虽然这种对过去数据的持久访问使Transformer功能强大，但其成本也很高。随着序列的增长，缓存也随之增长，每个新步骤都必须读取越来越大的历史数据。对于长上下文应用，这在内存使用和运行时间上都造成了重大瓶颈。一种有吸引力的补救措施是*KV压缩*：用紧凑的摘要替换长的键-值对历史，近似保持原始缓存的注意力输出。¹¹该术语也指存储张量的数值压缩。我们专注于减少键-值对的数量。这样的摘要只有在未来注意力无法可靠地将压缩后的缓存与完整历史区分开时才有用。这使得决定过去哪些区别可以安全地遗忘变得具有挑战性，因为token的重要性并非token本身固有的。在一个上下文中以及对于一组未来查询可以安全丢弃的token，在另一种情况下可能至关重要，而如果注意力以相同方式使用多个不同token，则它们可能是可互换的。因此，压缩的安全性仅限于摘要能够保留未来查询实际可以使用的上下文。  

尽管具有实际重要性，KV压缩缺乏系统性的设计理论，在很大程度上仍然由经验探索所指导[20 (https://arxiv.org/html/2607.01520#bib.bib20),27 (https://arxiv.org/html/2607.01520#bib.bib27),24 (https://arxiv.org/html/2607.01520#bib.bib24),19 (https://arxiv.org/html/2607.01520#bib.bib19),10 (https://arxiv.org/html/2607.01520#bib.bib10),1 (https://arxiv.org/html/2607.01520#bib.bib1)]，使用未来相关性的启发式代理，如近因性、累积的注意力权重和注意力汇（attention sinks）。这些方法表明通常可以实现显著的压缩，但并未解释缓存的哪些属性使压缩安全、可实现的误差应如何随压缩预算变化，或者最优摘要应保留哪些信息。相反，理论结果表明，在最坏情况下精确的KV压缩可能是不可能的[16 (https://arxiv.org/html/2607.01520#bib.bib16)]，并确定了在哪些结构假设下压缩是可能的[18 (https://arxiv.org/html/2607.01520#bib.bib18),11 (https://arxiv.org/html/2607.01520#bib.bib11),26 (https://arxiv.org/html/2607.01520#bib.bib26)]。这使得与实际算法设计最相关的中等机制悬而未决：缓存既不是对抗性的，也不是由固定模型类别规定的。  

为了弥合这一差距，我们基于缓存与未来查询通过softmax注意力的内在交互，对KV缓存的可压缩性进行了分级描述。具体来说，我们刻画了*KV压缩的极小极大风险*，询问：对于固定的压缩预算，给定缓存被未来查询探测的方式，误差有多少是不可避免的？除了量化比率之外，该刻画还确定了最优摘要必须保留的内容。它捕捉了从容易机制到困难机制的连续谱：在容易机制中，许多过去的token实际上是冗余的，一个小的摘要可以准确保持注意力；在困难机制中，缓存包含许多可单独检索的信息片段，任何小的摘要都必须丢失某些内容。此外，它还为算法实现极小极大最优风险提供了具体的设计标准。  

具体而言，我们做出以下贡献：  
(1) 我们将KV压缩重新表述为测度的稀疏逼近，严格推广了token驱逐（token eviction），并形式化了其极小极大风险。  
(2) 我们证明了该风险的紧的上界和下界，这些界以捕捉未来查询如何探测缓存的内在复杂度度量来表示，并展示了有无访问未来查询分布的算法之间的清晰分离。  
(3) 我们推导了在因果掩码预填充和自回归解码过程中实现高效KV压缩的设计原则，具有极小极大最优的压缩风险。  
(4) 我们将这些原则实例化为一个具体的KV压缩算法，并通过在LongBench上的针对性实验进行评估，显示出有希望的性能。  

总体而言，我们的结果将KV压缩推向具有理论保证的实用方法。  

## 2 相关工作  
大量近期工作通过根据未来相关性的经验代理选择token来减少KV缓存大小。这些代理包括近因性和注意力汇[24 (https://arxiv.org/html/2607.01520#bib.bib24),27 (https://arxiv.org/html/2607.01520#bib.bib27)]、累积或持续的注意力权重[20 (https://arxiv.org/html/2607.01520#bib.bib20),27 (https://arxiv.org/html/2607.01520#bib.bib27),1 (https://arxiv.org/html/2607.01520#bib.bib1)]，以及在预填充或跨层期间观察到的注意力模式[19 (https://arxiv.org/html/2607.01520#bib.bib19),10 (https://arxiv.org/html/2607.01520#bib.bib10)]。这些方法的强经验性能表明，真实的Transformer缓存通常包含大量冗余。同时，这些代理本身并不能解释缓存的哪些属性使压缩安全、可实现的误差应如何随压缩预算变化，或者最优摘要必须保留哪些信息。我们的工作通过广泛研究KV压缩算法的极小极大风险来解决这一差距，旨在为超越经验探索的实用压缩提供信息。  

现有理论主要确定了两个端点：对于尖锐token检索的不可行性结果[3 (https://arxiv.org/html/2607.01520#bib.bib3),16 (https://arxiv.org/html/2607.01520#bib.bib16)]，以及在额外结构下的正向逼近保证[26 (https://arxiv.org/html/2607.01520#bib.bib26),18 (https://arxiv.org/html/2607.01520#bib.bib18),11 (https://arxiv.org/html/2607.01520#bib.bib11),2 (https://arxiv.org/html/2607.01520#bib.bib2)]。这些结果确定了重要的限制和可处理的机制，但并未从无损压缩到最坏情况不可压缩性给出分级刻画。我们的上界除了有界值之外，不需要数据的先验结构模型，这比有界键和查询更弱。相反，我们有意识地通过缓存与softmax注意力本身的交互来刻画压缩风险。我们的下界假设在精神上接近于先前硬度结果中考虑的尖锐注意力机制。然而，我们使用它们来刻画即使在这种尖锐注意力下压缩何时仍然可能，而不仅仅是建立不可行性。  

最后，我们将KV压缩的极小极大风险形式化为测度的稀疏逼近。基于测度的softmax注意力和Transformer观点并不新鲜，并在多个设置中产生了有用的视角[15 (https://arxiv.org/html/2607.01520#bib.bib15),9 (https://arxiv.org/html/2607.01520#bib.bib9)]。据我们所知，我们的重新表述对于KV压缩是新的，它将token驱逐从子集选择严格推广到稀疏重新加权。  

## 3 问题设置  
我们研究KV压缩的问题：用紧凑的摘要替换单个softmax注意力头所看到的上下文，同时保持其在未来查询上的输出，以减少内存使用和计算成本。由于softmax注意力仅通过其对键-值对的加权依赖于上下文，我们将KV压缩形式化为对token上的概率测度的稀疏逼近。我们的目标是刻画这个逼近问题可实现的误差，作为上下文在查询分布下的*内在可压缩性*的函数，我们将其形式化为一个依赖于数据的复杂度度量下的极小极大风险。  

### 3.1 KV压缩作为概率测度的稀疏逼近  
在其开创性工作中，Vaswani等人[23 (https://arxiv.org/html/2607.01520#bib.bib23)]将softmax注意力定义为  
Att⁡\(q∣k1,...,kn,v1,...,vn\)=∑i=1nexp⁡\(⟨q,ki⟩/dk\)vi∑i=1nexp⁡\(⟨q,ki⟩/dk\)，  
其中 \(q \in \mathcal{Q}\) 是一个*查询*，\(k_i \in \mathcal{K}\) 是*键*，\(v_i \in \mathcal{V}\) 是*值*，分别位于各自的空间 \(\mathcal{K}, \mathcal{Q} \subseteq \mathbb{R}^{d_k}\) 和 \(\mathcal{V} \subseteq \mathbb{R}^{d_v}\)。我们记 \(\mathcal{X} \coloneqq \mathcal{K} \times \mathcal{V}\) 为键-值空间，并假设值在范数上以常数 \(V>0\) 有界，这在实践中通常不是限制性的。集合 \((k_i, v_i)_{i=1}^n\) 称为*KV缓存*。我们研究*KV压缩*的问题：找到该集合的一个紧凑表示，近似保持 (1) 中的注意力映射。由于 (1) 仅通过加权组合依赖于KV缓存，我们将KV缓存表示为一个有限支撑的概率测度 \(P\)，并定义  
\[
\operatorname{Att}\bigl(q\,\big|\,P\bigr) \coloneqq \int_{\mathcal{X}} a_k(q\mid P) \, v \, \mathrm{d}P(k,v), \qquad a_k(q\mid P) \coloneqq \frac{\kappa(q,k)}{\int_{\mathcal{X}} \kappa(q,k') \, \mathrm{d}P(k',v')}
\]  
其中核函数 \(\kappa(q,k) \coloneqq \exp(-\|q-k\|_2^2/(2\sqrt{d_k}))\) 是 \(\mathbb{R}^{d_k}\) 上的高斯核，尺度参数为 \(\sqrt{d_k}\)。当 \(P \propto \sum_{i=1}^n \exp(\|k_i\|^2/(2\sqrt{d_k})) \cdot \delta_{(k_i,v_i)}\) 时，该定义与 (1) 一致。在本文其余部分，我们使用 (2) 中定义的注意力，使得注意力仅通过*上下文测度* \(P\) 依赖于KV缓存。  

从这一视角形式化KV压缩，我们引入符号 \(\mathcal{P}(\mathcal{S})\) 表示希尔伯特空间 \(\mathcal{H}\) 的子集 \(\mathcal{S} \subseteq \mathcal{H}\) 上的概率测度，用 \(\mathcal{P}_{\mathrm{fin}}(\mathcal{S})\) 表示有限支撑的测度，用 \(\mathcal{P}_K(\mathcal{S})\) 表示支撑点最多为 \(K\) 的测度。KV压缩于是变成以下问题。  

**问题陈述 I（KV压缩）**  
给定预算 \(K \geq 1\)，上下文测度 \(P \in \mathcal{P}_{\mathrm{fin}}(\mathcal{X})\)，以及查询分布 \(\nu \in \mathcal{P}(\mathcal{Q})\)，找到压缩上下文测度 \(\hat{P} \in \mathcal{P}_K(\mathcal{X})\)，使得均方误差  
\[
\mathcal{E}_{P,\nu}(\hat{P}) \coloneqq \mathbb{E}_{q\sim\nu} \bigl\| \operatorname{Att}(q\mid P) - \operatorname{Att}(q\mid \hat{P}) \bigr\|_2^2
\]  
最小化。  

3.1节严格推广了token驱逐[20,27,24,19,10]，后者寻求一个小的、未加权的 \(K\) 个token子集来近似保持注意力映射。在我们的设置中，这些token形成 \(\hat{P}\) 的支撑集。然而，我们允许任意的稀疏重新加权，据我们所知，这为KV压缩提供了新的视角。  

### 3.2 把握上下文的内在可压缩性  
为了理解精确压缩何时可行，我们确定了 \(P\) 的变化如何影响softmax注意力。由于 \(P\) 和 \(\hat{P}\) 都是概率测度，它们的差 \(\Delta \coloneqq \hat{P} - P\) 是一个符号测度，总质量为零。因此压缩可以看作对原始上下文的重新加权，其误差由这种重新加权如何扰动softmax分子和归一化因子决定。事实上，直接计算表明，对于每个查询 \(q \in \mathcal{Q}\)，  
\[
\operatorname{Att}(q\mid \hat{P}) - \operatorname{Att}(q\mid P) = \frac{\int_{\mathcal{X}} a_k(q\mid P) \bigl(v - \operatorname{Att}(q\mid P)\bigr) \, \mathrm{d}\Delta(k,v)}{1 + \int_{\mathcal{X}} \bigl(a_k(q\mid P) - 1\bigr) \, \mathrm{d}\Delta(k,v)}.
\]  
因此，每次压缩扰动都通过 \(\Delta\) 的两个线性泛函进入注意力：一个中心化的分子项和一个中心化的重新归一化项。这确定了每个 \((k,v) \in \mathcal{X}\) 贡献于压缩误差时所用的查询索引系数对。