🤖 machine learning

CompilerKV: Risk-Adaptive KV Compression via Offline Experience Compilation

CompilerKV 引入了一种离线编译的 KV 保留策略，该策略利用跨提示的规律性，以高效的 $O(1)$ 查找取代嘈杂的在线估计，在极端内存约束下，相较于现有的仅预填充压缩方法，实现了最先进的性能与更优越的可扩展性。

原作者： Ning Yang, Chengzhi Wang, Yibo Liu, Baoliang Tian, Haijun Zhang

发布于 2026-05-21

📖 1 分钟阅读☕ 轻松阅读

原作者： Ning Yang, Chengzhi Wang, Yibo Liu, Baoliang Tian, Haijun Zhang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试记住一个非常长的故事，以便稍后能回答关于它的问题。你的大脑（AI 模型）拥有一块有限的“思维草稿纸”空间（称为KV 缓存），用于在对话过程中保留该故事最重要的部分。

如果故事很短，你可以记住所有内容。但如果故事长达 10 万字，你的草稿纸就会溢出。你必须丢弃一些部分以腾出空间。问题在于：一旦你丢弃了某些内容，就永远无法找回它。 如果你不小心扔掉了反派揭示其秘密计划的那部分，整个故事就会分崩离析。

旧方法：即时猜测

之前的方法试图通过实时观察故事来决定保留什么。它们会说：“这个词的注意力分数很高，所以它一定很重要！”或者“这个句子看起来很复杂，保留它！”

本文作者认为，这就像试图通过观看单帧充满噪点的画面来评判整部电影。由于故事如此之长且数据杂乱无章，仅观察单个提示（即单个故事）会导致糟糕的猜测。你可能会保留一个华丽但无用的词，却丢弃了一个安静却关键的线索。

新方法：COMPILERKV（“赛前策略”）

作者引入了COMPILERKV。他们不再临场猜测，而是像教练在重大比赛前研究比赛录像一样，预先“编译”出一套策略。

其工作原理可分解为三个简单步骤：

1. “噪声过滤器”（稳定化效用）

想象你正身处一个嘈杂的房间。有些人正在大喊大叫（瞬态尖峰），而有些麦克风本身音量就更大（尺度偏差）。

解决方案：COMPILERKV 不仅仅听谁喊得最响。它会平滑噪声并标准化音量。它会问：“这个人真的在说重要的话，还是仅仅声音大？”这防止了 AI 被暂时的噪声所误导。

2. “专家地图”（头异质性表）

在 AI 内部，有许多不同的“脑细胞”（称为注意力头）。有些是寻找事实的专家（像图书管理员），而另一些则只是嘈杂的闲聊。

解决方案：作者离线研究了成千上万个故事，并意识到：某些脑细胞始终可靠，而另一些则始终充满噪声。 他们创建了一张永久的“专家地图”。
类比：系统不再向整个团队征求意见，而是知道：“如果‘图书管理员’细胞说某个词很重要，我们就保留它，即使‘话痨’细胞说它是垃圾。”这赋予了可靠的专家“否决权”，以保存关键信息。

3. “风险计”（风险自适应阈值）

并非所有故事都一样。有些很简单（如食谱）；有些则复杂且令人困惑（如悬疑小说）。

解决方案：系统会检查当前故事的“风险”程度。
- 低风险（简单故事）：“我们可以激进一点。丢弃 90% 的文本；我们会没事的。”
- 高风险（复杂故事）：“这很危险！要保守一点。保留 95% 的文本以防万一。”
类比：这就像为旅行打包。如果你知道天气完美，你就轻装上阵。如果有风暴预警，你就多带装备。COMPILERKV 会根据当前提示的“风暴”程度，自动调整保留内容的多少。

为什么这很重要

该论文声称，通过进行这种“赛前研究”（离线编译）而不是“即时猜测”（在线启发式），他们获得了更好的结果：

可移植性：他们创建的“专家地图”适用于不同的 AI 模型。它就像一本适用于不同类型大脑的通用规则手册。
节省内存：他们可以在极小的内存预算下（仅原始文本的 1.5%）让 AI 运行，同时仍能正确回答问题。
处理长文本：在针对巨大上下文（高达 128,000 字）的测试中，COMPILERKV 表现强劲，而其他方法则崩溃了。例如，在“大海捞针”测试（在巨大文本中寻找一个特定事实）中，COMPILERKV 找到了针，成功率为 89%，而次优方法的成功率仅为 42%。

总结

COMPILERKV 不再试图在瞬间做出完美决策。相反，它利用一种预先计算且具备风险意识的策略来决定保留什么。这就像赌徒下注与象棋特级大师之间的区别：后者已根据多年的经验计算出了最佳走法。其结果是，AI 能够记住海量文本，而不会耗尽内存或遗忘重要部分。

技术摘要：COMPILERKV

1. 问题陈述

大型语言模型（LLMs）因键值（KV）缓存随上下文长度线性增长而面临内存瓶颈。仅预填充（Prefill-only）的 KV 压缩已成为一种易于部署的解决方案，即在预填充阶段结束时选择一部分 token 并将其冻结，用于整个解码过程。该方法将运行时内存使用转化为固定成本，且无需解码时的干预。

然而，该论文指出，现有的仅预填充方法（如 SnapKV、PyramidKV、HeadKV）存在根本性的统计失配。这些方法试图利用单个、含噪的 prompt 的统计信息，在线估计校正信号——具体而言是每头可靠性（哪些注意力头对检索至关重要）和prompt 级压缩敏感度（特定 prompt 可被压缩的程度）。作者认为这是错误的统计单元：

E1（单 prompt 噪声）： 每 prompt 的注意力分数具有重尾分布且存在尺度偏差；大数值并不保证语义重要性，微小的估计误差会导致不可逆的选择错误。
E2（头的异质性）： 注意力头在功能上是特化的。有些是“检索关键”的，而另一些则始终含噪。简单的聚合方式会让含噪头主导分数。
E3（prompt 级风险方差）： 不同的 prompt 具有不同的压缩极限。固定的保留规则在高熵或高困惑度 prompt 上会失效，导致灾难性的信息丢失。

现有方法为每个 prompt 重新在线估计这些信号，未能利用这些信号表现出高度的跨 prompt 规律性而非 prompt 内部稳定性这一事实。

2. 方法论：COMPILERKV

作者提出了COMPILERKV，这是一种 KV 保留策略，它从校准语料库中离线编译校正表，而非在线估计。核心洞察在于：头的可靠性排名以及从风险特征到安全阈值的映射是稳定的架构先验，可以学习一次并重复使用。

该框架包含三个在预填充阶段协同工作的组件：

A. 稳定化的 Token 效用（解决 E1）

为了在不学习的情况下消除瞬态注意力尖峰和头间尺度偏差，该方法计算一个无参数的效用分数：

窗口聚合注意力： 通过在 prompt 尾部的观察窗口 $\Omega$ 上聚合注意力，平滑瞬态尖峰。
头相对归一化： 通过头的均值范数对 token 分数进行归一化，消除尺度偏差。
公式： $u^{(l,h)}_t = \alpha_t \cdot \rho^{(l,h)}_t$ ，其中 $\alpha_t$ 是窗口质量， $\rho$ 是归一化的值范数。

B. 头异质性表（解决 E2）

COMPILERKV 不是平等对待所有头或在线重新估计重要性，而是编译一个冻结表 $W_{head}[l, h]$ ，其中包含连续的可靠性权重。

离线编译： 在校准语料库上通过**保守 Q 学习（CQL）**学习。状态为 $(l, h, B_l)$ ，动作为离散化的可靠性权重。
机制： 在推理过程中，token 的效用会根据其头的可靠性进行重缩放。关键在于，聚合使用加权最大池化（ $\hat{u}^{(l)}_t = \max_h (u^{(l,h)}_t \cdot W_{head}[l, h])$ ）。这赋予了检索关键头“否决权”：如果单个关键头强烈关注某个 token，无论含噪头是否达成共识，该 token 都会被保留。

C. 风险自适应阈值门（解决 E3）

为了处理具有不同压缩极限的 prompt，风险自适应阈值门将 prompt 级风险特征映射到保留阈值。

风险信号： 使用在观察窗口上计算的注意力熵（结构风险）和局部困惑度（语义风险）。
离线编译： 查找表（LUT）将离散化的风险箱 $(b_{ent}, b_{ppl})$ 映射到保留阈值 $\tau^{(l)}$ 。这也是通过 CQL 编译的。
保守偏差： CQL 目标函数惩罚稀有、高风险状态（高熵/高困惑度）下的激进阈值，使表格偏向于更安全的保留，而这通常是在线启发式方法失效的地方。
弹性选择： 最终选择保留所有效用高于阈值 $\tau^{(l)}$ 的 token。如果候选集超过预算 $B_l$ ，则限制为前 $B_l$ 个 token。如果低于预算，则保留所有候选（弹性欠保留）。

D. 理论依据

论文提供了形式化分析，表明近似误差由被丢弃的尾部质量控制（命题 4.1）。它论证了有限表校准误差是在prompt 级别（跨独立 prompt）控制的，而不是假设单个 prompt 的注意力行内部是独立的，从而降低了方差。

3. 主要贡献

重构问题： 指出仅预填充压缩中校正的统计单元应是跨 prompt 尺度，而非单个 prompt。
离线编译： 提出了一种离线编译头可靠性和风险自适应阈值的方法，将其转化为可移植的架构先验。
耦合策略： 提出联合优化，其中头可靠性权重和风险阈值针对相同的下游奖励进行校准，而不是作为独立模块。
可移植性： 证明了编译后的表格可以在不相交的语料库之间转移，甚至在不同模型家族之间（例如从 LLaMA 到 Mistral）转移，且性能下降极小。

4. 实验结果

该方法在四个 7B–8B 骨干模型（LLaMA-3、Mistral、Qwen2、InternLM）上进行了评估，涵盖了 LongBench、RULER 和 Needle-in-a-Haystack (NIAH) 基准测试。

LongBench（512-token 预算）： COMPILERKV 在所有四个骨干模型上均实现了压缩 SOTA，平均比最强基线（如 ChunkKV、HeadKV）高出**+1.67 分**（95% CI [+1.08, +2.37]）。在某些骨干模型上，其性能略超过 FullKV，作者将其归因于去噪而非信息增益。
压力区间：
- RULER（128k 上下文）： 在固定的 512/32k 缓存比例下，COMPILERKV 保持约 73 的分数，显著优于 SnapKV（约 38）和 KVZip（约 66.5），同时低于 FullKV（约 79）。
- NIAH（32k）： 准确率达到 0.89，而 SnapKV 为 0.42。
- 批处理： 在 32k 输入且批量大小为 16 时，COMPILERKV 仍然可行，而 FullKV 则因内存不足（OOM）而失败。
可迁移性：
- 跨语料库： 头可靠性排名在不相交的校准语料库之间显示出高度的 Spearman 相关性（ $\bar{\rho} \approx 0.90$ ）。
- 跨模型： 模型间的直接表格转移（例如从 LLaMA-3 到 Mistral）相比针对目标模型的编译，仅损失0.4–0.8 个 LongBench 分数。
效率： 在线开销极小：每层进行 $O(1)$ 次表查找，并对候选池进行一次排序。在 32k 上下文下，其解码吞吐量相比 FullKV 最高可达2.45 倍。

5. 意义与主张

论文声称，COMPILERKV 通过纠正保留策略计算中的统计失配，缩小了与固定预算下 FullKV 的差距。通过将校正信号的计算从在线、含噪、单 prompt 估计转移到离线、编译、跨 prompt 先验，该方法稳定了选择过程，使其免受噪声和异质性的影响。

作者强调，其新颖性不仅仅在于“头感知”（先前的方法如 HeadKV 也具备），而在于连续头可靠性和 prompt 风险阈值的联合离线编译。这将头可靠性从一种调优技巧转变为一种可重用的科学对象，可以编译一次并部署到不同的模型和上下文中，且在线开销为 $O(1)$ 。这项工作表明，对于仅预填充压缩而言，“正确”的信号是模型的架构属性，而非输入 prompt 的瞬态特征。