Training Transformers for KV-Cache Compressibility¶
会议: NeurIPS 2026
arXiv: 2605.05971
OpenReview: https://openreview.net/forum?id=igMG5UgMAj
论文: https://arxiv.org/abs/2605.05971
代码: 待开源
领域: 模型压缩
关键词: KV缓存压缩, 模型压缩, 持续预训练, 自蒸馏, 长上下文推理
一句话总结¶
针对事后KV缓存压缩受限于冻结预训练表征固有冗余结构的瓶颈，本文形式化界定了Transformer的KV可压缩性并从理论上证明其取决于表征结构，进而提出KV压缩感知训练（KV-CAT），通过训练期自适应KV稀疏化与自蒸馏目标引导模型自发学习可压缩表征。
研究背景与动机¶
大语言模型在代码库全解析、超长文档问答、智能体连续记忆等长上下文场景中的广泛部署，正面临着日益严峻的内存与计算瓶颈。在自回归生成过程中，模型必须在每一层和每个注意力头为所有历史Token保存键（Key）与值（Value）向量。这种KV缓存（KV Cache）的显存占用和解码访问开销随前缀长度呈严格线性增长，已成为限制高吞吐长文本推理服务的头号系统制约。为了化解这一矛盾，近年来学术界提出了大量事后（Post-hoc）压缩方案，从基于注意力权重与衰减启发式的Token驱逐（如 H2O、StreamingLLM），到基于优化驱动的紧凑KV重建（如 Cartridges 与 Attention Matching），试图在前缀推理阶段削减冗余槽位。
然而，现有后处理压缩方法存在一个普遍被忽视的核心假设：它们均在参数完全冻结的预训练模型上执行。这意味着后处理压缩的极限效果，天然被该特定模型在无约束预训练阶段所形成的内部激活表征所禁锢。从信息论与表达能力视角看，即便两套Transformer参数在相同输入序列上能输出完全一致的下一Token预测概率分布，其内部各层KV投影所编码的状态几何结构也可以截然不同。标准因果语言建模目标仅追求无损 dense 前向下的预测精度，对KV向量间是否存在低维压缩投影或多槽位冗余毫无归纳偏置，常常导致模型学到理论上精确却在事后根本无法低秩压缩或稀疏剔除的隐式表征。
本文敏锐地将研究视角从“如何设计更复杂的后处理压缩算法”转向“如何让模型自身学到更易压缩的KV表征”。作者在理论上严格定义了 \((N, \varepsilon, r)\)-KV可压缩性，并证明几乎任何序列到向量的映射任务均同时存在极端可压缩与本质不可压缩的两类Transformer参数解。核心 idea：将KV压缩的归纳偏置前置到模型训练阶段，提出KV压缩感知训练（KV-CAT），在持续预训练中通过轻量级可学习路由对KV槽位实施动态掩蔽，并利用稠密教师模型的输出概率做自蒸馏与锚定，使模型在完全保留无压缩生成精度的同时，表征自然演化出对后续事后压缩高度友好的紧凑结构。
方法详解¶
整体框架¶
KV-CAT 作为轻量级持续预训练（Continued Pretraining, CPT）流程，直接以预训练好的因果语言模型为起点。算法的核心执行流程分为双路并行前向传播：稠密前向通道（Dense Forward Pass）与带掩蔽的稀疏前向通道（Masked Forward Pass）。在掩蔽通道中，模型在层间插入轻量级路由模块（Learned Routers），为连续层组计算动态激活掩码，将低于阈值的KV槽位在自注意力计算中剔除；在稠密通道中，模型保留全部完整KV槽位正常推理。
随后，模型同时计算三项自监督损失：利用稠密通道的预测输出作为无梯度教师，指导掩蔽通道通过KL散度自蒸馏对齐分布（\(\mathcal{L}_{\mathrm{mask}}\)）；施加可控预算正则项以约束路由激活率稳定在目标保留率 \(\rho\)（\(\mathcal{L}_{\mathrm{budget}}\)）；同时在稠密通道保留标准下一Token预测（NTP）交叉熵损失作为锚定项（\(\mathcal{L}_{\mathrm{anchor}}\)）。推理部署时，训练期间引入的路由模块与掩码机制被完全剥离，模型恢复为标准的密集Transformer架构，直接送入下游各类事后KV压缩算法（如 Attention Matching 或基于梯度的KV压实）。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入上下文序列 a"] --> B["双路前向计算架构"]
B --> C["线性路由模块与KV槽位掩蔽<br/>因果线性摘要 + 余弦重要性打分"]
B --> D["密集无掩蔽前向通道<br/>标准因果注意力提取稠密 logits"]
C --> E["稀疏通道概率分布输出<br/>仅保留 active 槽位参与注意力"]
D --> F["稠密锚定损失 L_anchor<br/>标准 NTP 监督锁定原始能力"]
E & D --> G["自蒸馏匹配损失 L_mask<br/>KL(sg[p_dense] || p_mask)"]
C --> H["平滑预算损失 L_budget<br/>引导槽位平均保留率趋近目标 rho"]
G & F & H --> I["联合反向传播与参数更新<br/>优化 Transformer 与路由模块"]关键设计¶
1. KV可压缩性理论界定与构造性反例：表征几何决定压缩上限
为了厘清压缩瓶颈的物理本质，论文首先将模型KV缓存压缩抽象为一个映射族 \(\mathsf{C}=(\mathsf{c}_1, \dots, \mathsf{c}_L)\)，其将长度为 \(n\) 的KV序列压缩为长度为 \(r(n) \le n\) 的紧凑状态 \((\tilde{\bm{K}}, \tilde{\bm{V}})\)。若存在压缩策略 \(\mathsf{C}\) 使得对任意前缀 \(\bm{a}\) 和后缀 \(\bm{b}\)，压缩前缀后的模型输出与原全量输入输出之差小于 \(\varepsilon\)，则称该模型是 \((N, \varepsilon, r)\)-可压缩的。定理 3.1 证明了几乎任何非平凡序列函数 \(f\) 均存在两种参数化实现：一种可用 \(r(n)=1\) 的单键值对实现任意精度拟合，另一种对任意 \(r(n) < n\) 的压缩均存在常数级别下界误差 \(C > 0\)。
在直观的字符直方图统计示例中，传统两层Transformer通常让第二层注意力权重退化为均匀平均（\(\bm{W}_Q = \mathbf{0}\)），使得末尾Token表示为前缀与后缀的算术均值。当把前缀压缩至 \(r(n)\) 槽位后，后缀注意力权重分母变为 \(r(n)+k\)，从而引入由后缀决定的不可消除偏差 \(\tilde{\bm{b}} = (\frac{1}{r(n)+k} - \frac{1}{n+k})\sum \bm{e}_{b_i}\)。相反，若模型在第一层将位置编码保留在表征特定子空间中，压缩策略即可把未归一化直方图与原始前缀长度 \(n\) 封装在单一KV对中，使顶层MLP能够依据后缀位置精确重整化分母。这一理论深刻指出：标准训练并不天然激励模型显式打包此类校准信息，因此必须通过专门机制诱导可压缩表征的生成。
2. 组间共享因果线性路由：轻量化诱导动态稀疏瓶颈
在训练循环中直接运行复杂的二阶优化压缩算法在计算上是不可行的。KV-CAT 设计了极低开销的自适应训练期KV稀疏化策略。作者在Transformer中每隔若干层（如4个连续层组）插入一个轻量级可学习路由模块。为避免路由本身产生额外的KV显存驻留，路由器采用无状态的因果线性注意力（Linear Attention）累积前文摘要： $$ S_t = \sum_{j \le t} k_j r_j^\top, \quad z_t = \sum_{j \le t} k_j, \quad a_t = W_O \frac{q_t^\top S_t}{q_t^\top z_t + \epsilon} $$ 各Token的保留概率 \(p_t \in [0, 1]\) 由其当前隐状态投影与融合摘要特征后的归一化余弦相似度算出：\(u_t = \frac{W_P h_t}{\|W_P h_t\|_2}, w_t = \frac{h_t + \alpha a_t}{\|h_t + \alpha a_t\|_2}, p_t = \frac{1 - \langle u_t, w_t \rangle}{2}\)。随后通过硬阈值 \(m_t = \mathbf{1}\{p_t > \tau\}\)（设定 \(\tau = 0.5\)）生成二值掩码，并在反向传播中采用直通估计器（Straight-Through Estimator, STE）将梯度回传至路由参数。初始化时置 \(W_P = -I, \alpha = 0\)，使得初始阶段所有Token评分均为 1，确保模型平滑从标准稠密状态过渡到稀疏训练。
3. 双向锚定自蒸馏机制：零外部教师实现无损无压缩推理
强行剥离KV槽位极易导致模型原有语言建模能力退化。为了杜绝“为了压缩而损害原始精度”的问题，KV-CAT 舍弃了引入独立冻结外部教师的高显存做法，设计了同模型内生自蒸馏目标。在同一批次数据上，模型同时执行无掩蔽前向与带掩蔽前向，整体训练损失函数为： $$ \mathcal{L}(\theta) = \lambda_{\mathrm{mask}} \mathcal{L}{\mathrm{mask}} + \lambda}} \mathcal{L{\mathrm{budget}} + \lambda $$ 其中，}} \mathcal{L}_{\mathrm{anchor}\(\mathcal{L}_{\mathrm{mask}} = \frac{1}{T-1} \sum_{t=2}^T D_{\mathrm{KL}}\left(\operatorname{sg}[p_\theta^{\mathrm{dense}}(\cdot \mid x_{<t})] \,\middle\|\, p_\theta^{\mathrm{mask}}(\cdot \mid x_{<t})\right)\) 强迫稀疏掩蔽下的输出分布逼近自身密集前向分布（\(\operatorname{sg}\) 为截断梯度）；\(\mathcal{L}_{\mathrm{anchor}} = \frac{1}{T-1}\sum_{t=2}^T -\log p_\theta^{\mathrm{dense}}(x_t \mid x_{<t})\) 在稠密通道上施加标准自回归NTP损失，充当锚点以防表征漂移；\(\mathcal{L}_{\mathrm{budget}}\) 采用基于动态块平衡的双极正则化，严格约束路由的平均激活率收敛于目标比例 \(\rho = 0.5\)。三者协同使得Transformer被迫将全序列的核心语义重组为高能量密度的紧凑表征。
损失函数 / 训练策略¶
训练采用流式 FineWeb-Edu 英文高质量语料库，针对 Qwen2.5-0.5B 和 Qwen2.5-1.5B 展开持续预训练，序列长度统一配置为 1024。损失权重设定为 \(\lambda_{\mathrm{mask}} = 1.0, \lambda_{\mathrm{anchor}} = 1.0, \lambda_{\mathrm{budget}} = 0.1\)，优化器为 AdamW，峰值学习率为 \(10^{-4}\)，总训练Token量为 \(5.24 \times 10^9\)（约 5.24B Tokens）。
实验关键数据¶
主实验¶
在评测阶段，KV-CAT 训练后的模型在完全关闭路由掩码的无压缩基线测试中完好保留了密集能力；随后采用 SOTA 的 Attention Matching（每层8次NNLS迭代）与基于梯度的KV压实（Cartridges 衍生算法），在多种压缩保留率（Keep Ratio: 5%, 10%, 20%, 40%）下评估前缀压缩后的后缀预测保真度。
| 模型架构 | 保留率 (Keep Ratio) | 评估指标 | 原始 Base 模型 | KV-CAT (本文) | 质量提升 / 差距缩小 |
|---|---|---|---|---|---|
| Qwen2.5-0.5B | 5% | \(\Delta\)PPL (↓) / Top-1一致性 (↑) | 0.780 / 88.5% | 0.461 / 91.1% | \(\Delta\)PPL 下降 40.9% (+2.6% Top-1) |
| Qwen2.5-0.5B | 10% | \(\Delta\)PPL (↓) / Top-1一致性 (↑) | 0.662 / 89.3% | 0.422 / 91.4% | \(\Delta\)PPL 下降 36.3% (+2.1% Top-1) |
| Qwen2.5-0.5B | 20% | \(\Delta\)PPL (↓) / Top-1一致性 (↑) | 0.777 / 88.7% | 0.438 / 90.8% | \(\Delta\)PPL 下降 43.6% (+2.1% Top-1) |
| Qwen2.5-0.5B | 40% | \(\Delta\)PPL (↓) / Top-1一致性 (↑) | 0.592 / 90.3% | 0.360 / 92.5% | \(\Delta\)PPL 下降 39.2% (+2.2% Top-1) |
| Qwen2.5-1.5B | 5% | \(\Delta\)PPL (↓) / Top-1一致性 (↑) | 0.573 / 89.2% | 0.292 / 90.9% | \(\Delta\)PPL 下降 49.0% (+1.7% Top-1) |
| Qwen2.5-1.5B | 10% | \(\Delta\)PPL (↓) / Top-1一致性 (↑) | 1.027 / 88.3% | 0.320 / 90.7% | \(\Delta\)PPL 下降 68.8% (3.21× 保留提升) |
| Qwen2.5-1.5B | 20% | \(\Delta\)PPL (↓) / Top-1一致性 (↑) | 1.093 / 87.7% | 0.536 / 90.2% | \(\Delta\)PPL 下降 50.9% (+2.5% Top-1) |
| Qwen2.5-1.5B | 40% | \(\Delta\)PPL (↓) / Top-1一致性 (↑) | 0.742 / 90.0% | 0.259 / 92.4% | \(\Delta\)PPL 下降 65.1% (+2.4% Top-1) |
下表展示了在长文本检索（大海捞针 Passkey 检索）与 LongBench v2 长文档复杂问答任务中，对前缀KV进行压实后的下游端到端表现：
| 评估任务 | 骨干模型 | 压缩保留比例 | Base 基线准确率 | KV-CAT 准确率 | 绝对提升 |
|---|---|---|---|---|---|
| 大海捞针检索 (Needle-in-a-Haystack) | Qwen2.5-0.5B | 30% 保留率 | 23.0% | 34.0% | +11.0% |
| 大海捞针检索 (Needle-in-a-Haystack) | Qwen2.5-0.5B | 50% 保留率 | 28.0% | 47.0% | +19.0% (相对提升 67.9%) |
| 大海捞针检索 (Needle-in-a-Haystack) | Qwen2.5-1.5B | 30% 保留率 | 41.0% | 44.0% | +3.0% |
| 大海捞针检索 (Needle-in-a-Haystack) | Qwen2.5-1.5B | 50% 保留率 | 49.0% | 67.0% | +18.0% (相对提升 36.7%) |
| LongBench v2 长程问答 (7子任务总分) | Qwen2.5-0.5B | 10% 保留率 | 22.2% | 30.3% | +8.1% (相对提升 36.5%) |
| LongBench v2 长程问答 (7子任务总分) | Qwen2.5-0.5B | 20% 保留率 | 22.2% | 30.8% | +8.6% (相对提升 38.7%) |
| LongBench v2 长程问答 (7子任务总分) | Qwen2.5-0.5B | 50% 保留率 | 25.3% | 31.2% | +5.9% (相对提升 23.3%) |
消融实验¶
论文在附录系统对比了训练期间采用不同KV稀疏化策略（可学习Router、随机均匀丢弃Rand、基于注意力权重的静态截断Attn）对原始无压缩基准性能以及下游Attention Matching压缩效果的影响：
| 训练期稀疏化策略 | 无压缩平均准确率 (LightEval 6项任务) | 10%保留率 \(\Delta\)PPL (↓) | 10%保留率 KL散度 (↓) | 10%保留率 Top-1一致性 (↑) |
|---|---|---|---|---|
| Base 原始预训练模型 | 51.5% | 0.662 | 0.0483 | 89.3% |
| KV-CAT (可学习 Router) | 52.2% (完全无损且微增) | 0.422 | 0.0321 | 91.4% |
| KV-CAT (随机丢弃 Rand) | 50.7% (-0.8%) | 0.427 | 0.0356 | 90.9% |
| KV-CAT (注意力截断 Attn) | 50.2% (-1.3%) | 0.495 | 0.0363 | 90.8% |
关键发现¶
- 无压缩基线能力零折损：得益于 \(\mathcal{L}_{\mathrm{anchor}}\) 的约束，KV-CAT 训练后的模型在标准多任务（HellaSwag, WinoGrande, PIQA, OpenBookQA, ARC）上的平均得分（0.5B 52.2 vs 51.5；1.5B 60.1 vs 60.6）与原始模型基本持平，彻底破除了“增加稀疏约束必然伤害通用稠密能力”的顾虑。
- 事后优化收敛速度提升高达 5 倍：在基于梯度的KV压实收敛曲线中，KV-CAT 仅需原模型 \(20\%\) 的优化步数（约 20 步对比 100 步）即可达到相同的困惑度差距 \(\Delta\text{PPL}\)，大幅缩短长上下文缓存构建的前置时延。
- 泛化性跨越任务与压缩算法：无论下游使用解析式的 Attention Matching 还是基于梯度的优化器，无论评估指标是分布级KL散度、精确匹配检索还是高阶多跳问答，KV-CAT 均展现出一致的鲁棒优势。
亮点与洞察¶
- 提出“KV可压缩性是表征属性而非单纯任务属性”的关键认知转变，通过严谨的构造性证明指出即使功能等效的Transformer，其表征对事后压缩的敏感度也存在天壤之别。
- 设计了自给自足的双流自蒸馏目标，仅凭模型自身的稠密前向即可实现高保真教师引导，避免了加载外部冻结大模型所带来的成倍显存消耗。
- 将模型改造与推理机制彻底解耦：KV-CAT 产出的模型依然是绝对标准的稠密Transformer架构，完全不增加推理阶段的代码复杂度与硬件适配成本，对任何现存与未来的后处理压缩工具即插即用。
局限与展望¶
- 持续预训练的计算开销：在数十亿甚至百亿级参数模型上执行 5B Tokens 的持续预训练依然需要可观的高性能GPU算力（实验基于8卡 H100），对于算力受限团队或超大规模模型（70B+）直接落地的成本较高。
- 长序列训练受限：由于实验阶段受限于硬件环境，持续预训练的上下文窗口仅为 1024 Tokens，虽然下游验证了更长序列的迁移性，但在原生超长序列（如 32K~128K）下的训练动态与跨注意力头分配模式仍待进一步探索。
- 极端压缩率下的检索衰减：在极端严苛的超低保留率（如 5%~10%）下，大海捞针测试的绝对召回率仍然偏低（~20%），表明极少数关键事实槽位在粗粒度平均路由下仍有误伤风险。
相关工作与启发¶
- vs 后处理KV压缩（Cartridges [Eyuboglu et al. 2025] / Attention Matching [Zweiger et al. 2026]）：此类方法专注于在固定模型上设计更精细的特征逼近算法；本文证明了“巧妇难为无米之炊”，模型自身表征的不可压缩性是硬天花板，KV-CAT 与其完全正交正向互补。
- vs 训练期长文本高效架构（Mamba [Gu & Dao 2023] / Linear Attention）：高效新架构通常需要完全推翻现有预训练底座且下游精度不及标准自注意力；KV-CAT 保持了标准Transformer的表达上限，且只需轻量微调即可盘活存量开源模型生态。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [首次从表征学习角度系统界定与理论剖析KV可压缩性，范式视角极具启发性]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖了理论反例、自制路由消融、多模型尺度、两种代表性压缩算法及三大主流长文本基准]
- 写作质量: ⭐⭐⭐⭐⭐ [论证逻辑环环相扣，从数学构造平滑过渡到系统工程实现，图表详实自洽]
- 价值: ⭐⭐⭐⭐⭐ [为解决LLM长文本Serving的KV显存危机提供了全新的训练期解法，通用性极强]