如果你的AI的记忆不必随着每增加一个句子而膨胀会怎样？牛津大学、以色列理工学院、AITHYRA和NVIDIA引入了KV压缩感知训练（KV-CAT）——一种在训练期间（而不仅仅是事后）强制Transformer学习更可压缩的键值缓存的方法。通过在持续预训练期间随机屏蔽KV槽，KV-CAT促使模型依赖更少的记忆条目，使得事后压缩更加有效。结果：KV-CAT在检索、长上下文QA和困惑度上超越了现有的压缩方法，同时使用相同或更小的内存预算。

训练Transformer以获得KV缓存可压缩性
论文：https://arxiv.org/abs/2605.05971
我们的报告：https://mp.weixin.qq.com/s/LDkb9g1Jfq5wf-oC8aL32w…

训练Transformer以获得KV缓存可压缩性

来源：https://arxiv.org/html/2605.05971

Yoav Gelberg¹, Yam Eitan²¹, Michael Bronstein¹,³, Yarin Gal¹, Haggai Maron²,⁴
¹ 牛津大学
² 以色列理工学院
³ AITHYRA
⁴ NVIDIA

（2026年4月）

摘要

长上下文语言建模日益受到键值（KV）缓存的制约，其内存和解码时间访问成本随前缀长度线性增长。这一瓶颈催生了一系列上下文压缩方法，从词元级摘要到最近的基于优化的KV缓存压缩方法。这些事后方法作用于固定预训练模型的KV缓存，因此其有效性从根本上受限于模型内部表示的可压缩程度。在本文中，我们形式化了KV可压缩性的概念，并表明它是学习到的表示的一个属性，而非仅仅取决于上下文。我们证明，几乎任何序列到向量的函数都同时存在高度可压缩和本质上不可压缩的Transformer实现，这突显了在训练过程中引导Transformer朝向可压缩表示的必要性。受此启发，我们提出了KV压缩感知训练（KV-CAT），这是一种鼓励可压缩表示涌现的持续预训练过程。我们引入了一种训练时KV稀疏化策略，在训练期间屏蔽KV槽。这迫使模型使用更少的KV槽，并鼓励它学习适合事后压缩的表示。实验表明，KV-CAT在检索、长上下文问答和压缩前缀续写的困惑度评估中，改善了下游压缩方法的质量-预算权衡。

1 引言

语言模型越来越多地部署在长期场景中，从理解大型代码库、长篇文档和个人数据仓库[41; 23; 1]，到长形式推理和持续学习智能体[8; 60; 35]。在这些场景中使用自回归Transformer语言模型引入了显著的内存瓶颈：键值（KV）缓存。在推理过程中，这些模型必须为每个词元、每一层和每个（KV）注意力头存储键和值向量。对于长序列，这个缓存可能在内存使用和解码成本上占据主导，使上下文长度成为主要的服务瓶颈[59; 34; 52]。

大量近期工作试图缓解这一瓶颈，大致分为两类。第一类专注于设计比Transformer更高效的替代方案。例子包括线性注意力机制[29; 11; 20; 54]、状态空间模型[18; 17]、稀疏注意力变体[3; 57]等。这些方法显著降低了长上下文语言建模的计算成本。然而，这种效率通常以牺牲经验性能为代价，并且此类模型在大规模时仍落后于Transformer。

第二类专注于对固定预训练Transformer应用推理时干预。这些方法要么对输入上下文进行操作，要么更一般地，直接对KV缓存进行操作。早期方法主要是启发式的，包括文本摘要[46; 42]、学习的词元过滤[24]，以及基于注意力模式、新近性、重点行为或逐层重要性的策略[59; 52; 34; 6]。最近，基于优化的KV缓存压缩方法已成为强大的技术。例如，Eyuboglu等人[16]使用基于梯度的优化来匹配原始缓存诱导的分布，而Zweiger等人[61]使用逐层目标来复现其注意力痕迹。KV缓存压缩的成功表明，完整的KV缓存通常包含冗余。然而，由于大多数方法作用于固定模型，它们从根本上受限于该特定模型表示的可压缩程度。关键在于，这种可压缩性并非仅由输入序列决定：两个Transformer在处理相同序列时可以使用非常不同的内部表示，却产生完全相同的下一个词元分布。因此，某些Transformer的KV缓存可能比其他更适合压缩。

本文工作。 在本文中，我们研究以下问题：能否以一种导致KV缓存更适合事后压缩的方式来训练Transformer语言模型？这将目标从KV缓存压缩算法转向模型本身。为了回答这个问题，我们引入了KV可压缩性的概念。非正式地说，如果存在一个压缩策略，将长输入序列的KV缓存映射到更短的KV缓存，同时保持模型的下一个词元分布，那么这个Transformer是KV可压缩的。

a₁ a₂ a₃ a₄ a₅ a₆ a₇ a₈
masked dense
k₁,v₁ k₂,v₂ k₃,v₃ k₄,v₄ k₅,v₅ k₆,v₆ k₇,v₇ k₈,v₈
k₁,v₁ k₂,v₂ k₃,v₃ k₄,v₄ k₅,v₅ k₆,v₆ k₇,v₇ k₈,v₈
k₁,v₁ k₂,v₂ k₃,v₃ k₄,v₄ k₅,v₅ k₆,v₆ k₇,v₇ k₈,v₈
k₁,v₁ k₁,v₁ k₁,v₁ k₁,v₁ k₂,v₂ k₂,v₂ k₂,v₂ k₂,v₂ k₃,v₃ k₃,v₃ k₃,v₃ k₃,v₃ k₄,v₄ k₄,v₄ k₄,v₄ k₄,v₄ k₅,v₅ k₅,v₅ k₅,v₅ k₅,v₅ k₆,v₆ k₆,v₆ k₆,v₆ k₆,v₆ k₇,v₇ k₇,v₇ k₇,v₇ k₇,v₇ k₈,v₈ k₈,v₈ k₈,v₈ k₈,v₈
ℒₘₐₛₖ ℒₐₙₖₒᵣ ℒₕ
联合更新：模型θ + 路由器

图1：KV压缩感知训练（KV-CAT）。上下文 a 经过掩码（左）和密集（右）前向传递。在掩码前向传递中，路由器（橙色）为连续层组计算掩码，将KV槽标记为活跃（绿色）或不活跃（暗绿色）。在密集前向传递（蓝色）中，保留所有KV槽。掩码前向传递的输出用于计算ℒₘₐₛₖ，路由器分布用于计算ℒₕ，密集前向传递的输出用于计算ℒₐₙₖₒᵣ。这些共同用于更新参数。

我们的理论结果表明，对于几乎任何序列到向量的函数，都存在Transformer实现使其前缀可压缩到单个KV对，也存在实现使得任何非平凡压缩都会引入常数误差。为了建立直觉，我们考虑一个激励性例子：字符直方图计算。我们表明，这种计算的自然Transformer实现可能会产生不可压缩的词元表示，而更有结构的替代方案则高度可压缩。这促使我们将KV可压缩性作为一个明确的训练目标。

基于这一观点，我们提出了KV压缩感知训练（KV-CAT），这是一种持续预训练过程，促进KV可压缩内部表示的涌现。从预训练的Transformer开始，KV-CAT引入了一种训练时KV稀疏化策略，该策略屏蔽固定比例的KV槽。训练目标结合了自蒸馏损失（匹配掩码模型的分布与密集模型的分布）和应用于未掩码前向传递的NTP损失，以保持未压缩的模型行为。这使得模型在训练过程中暴露于KV缓存压缩所导致的的信息瓶颈，鼓励其在保持未压缩设置性能的同时，将其表示重新组织为更可压缩的形式。

为了证明KV-CAT改善了事后KV缓存压缩的质量-预算权衡，我们将其应用于Qwen2.5模型[51]，并在所得检查点的KV缓存上评估最先进的基于优化的压缩方法。我们沿着三个轴测量性能：(i) 前缀压缩下的后缀困惑度，(ii) 从压缩前缀的检索准确率，以及(iii) 在LongBench v2[2]任务上的压缩长上下文问答。在一系列压缩预算、模型大小和压缩方法中，使用KV-CAT训练的模型比基础模型持续获得更好的质量-预算权衡，在后缀困惑度保留方面提升高达3.21倍，优化时间提升5倍，检索准确率提升68%，长上下文QA提升39%。重要的是，我们的目标不是取代事后方法，而是使模型更适合它们，从而实现更有效的压缩。

贡献。 总结来说，本文做出了以下贡献：

1. 我们将KV可压缩性刻画为学习到的Transformer表示的一个属性，而不仅仅是任务或输入序列的属性。
2. 我们提供了理论结果，表明几乎所有序列到向量的函数都可以对应可压缩和不可压缩的Transformer实现。
3. 我们提出了KV-CAT，一种促进可压缩KV表示的训练过程。
4. 我们展示了KV-CAT在匹配的缓存和优化预算下，在后缀困惑度保留、检索和长上下文QA方面改善了事后KV缓存压缩。

2 KV缓存压缩：问题形式化

在本节中，我们从理论上形式化KV缓存压缩问题。对于具有L层的Transformer，一个KV缓存压缩策略是一组函数 C = (c₁, …, cₗ)，其中每个 c_l 将n个键值对的序列映射到长度 r(n) ≤ n 的序列。对于层l中的KV对 (K, V) = ([k₁, …, kₙ]ᵀ, [v₁, …, vₙ]ᵀ)，(1) 我们写作 c_l(K, V) = ([k̃₁, …, k̃_{r(n)}]ᵀ, [ṽ₁, …, ṽ_{r(n)}]ᵀ)，(2) 其中 r(n) 被称为压缩预算。给定一个Transformer M、一个压缩策略 C 和一个词元序列 a = (a₁, …, aₙ)，我们定义一个压缩模型 M_{C,a}，它共享 M 的参数但使用修改的前向传递。对于输入序列 b，M_{C,a}(b) 的计算过程与 M([a,b]) 相同，除了对前缀词元的注意力，其增强方式如下：在层l，令 (Kₐ, Vₐ) …