🤖 machine learning

Fast KV Compaction via Attention Matching

本文介绍了“注意力匹配”，这是一种在潜在空间中通过求解具有闭式解的子问题来快速高效地压缩语言模型 KV 缓存的方法，在最小化质量损失的前提下实现高达 50 倍的压缩，从而克服了以往基于优化方法的速度的局限性。

原作者： Adam Zweiger, Xinghong Fu, Han Guo, Yoon Kim

发布于 2026-05-28

📖 1 分钟阅读☕ 轻松阅读

原作者： Adam Zweiger, Xinghong Fu, Han Guo, Yoon Kim

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

想象一下，你正试图记住一个非常长的故事，以便稍后回答关于它的问题。在人工智能（AI）的世界里，这种“记忆”被称为KV Cache（键值缓存）。随着故事变长，这个记忆文件会变得巨大，占满计算机的硬盘并拖慢一切。

通常，当内存变得过大时，AI 系统会尝试通过总结故事来解决这个问题。它们丢弃细节，只保留简短的摘要。但这就像试图只通过阅读封底来记住一部复杂的悬疑小说：你会失去线索、情节转折以及回答具体问题的能力。

另一种称为“ cartridges（弹匣）”的方法，试图通过对每一个故事进行海量的数学训练，来创建一个微小而完美的记忆版本。它行之有效，但速度太慢且成本太高，就像每次你想移动一件家具时，都要雇佣一支建筑师团队重新设计房屋一样。

本文介绍了一种名为Attention Matching（注意力匹配）的新方法，它更快。以下是其工作原理，使用简单的类比说明：

1. 问题：“太长”的书架

将 AI 的记忆想象成一个拥有数千本书（token）的书架。当你提出一个问题时，AI 会查看所有书籍以找到相关的那几本。如果书架太满，AI 就会不堪重负。

旧方法（总结）： 扔掉 90% 的书，只保留一张摘要便条。你节省了空间，但再也找不到具体细节了。
旧方法（Cartridges）： 尝试将整个图书馆重写为一本完美、微小的书。它很准确，但需要数天才能写完。

2. 解决方案：“荧光笔与翻译器”（注意力匹配）

这种方法不像旧方法那样扔掉书籍或重写整个图书馆，而是像一位聪明的图书管理员，瞬间完成两件事：

步骤 A：荧光笔（选择键）
图书管理员查看故事并问道：“如果我要问关于这个故事的问题，我会查看哪些页面？”他们识别出最重要的页面（键），并只保留这些页面。
步骤 B：翻译器（调整值与偏差）
这里是魔法所在。如果你只保留几页，故事会感觉“变轻”了，因为你移除了缺失页面的重量。为了解决这个问题，图书管理员会给保留的页面添加一个特殊的偏差（微小的权重调整）。
- 类比： 想象你有一个装着 100 块重石头的背包。你需要背着它，但只能拿 5 块石头。如果你只挑 5 块石头，背包就太轻了。所以，你在每块石头上附加一个“魔法重量”，使得这 5 块石头的总重量和重要性感觉起来就像原来的 100 块一样。

3. 如何在无需缓慢训练的情况下工作

本文声称，这种方法不像"Cartridges"方法那样花费数小时训练新模型，而是使用数学捷径（闭式解）。

这就像使用公式来解谜题，而不是尝试每一种可能的拼图组合。
它精确计算出如何调整“权重”（偏差）和“值”（内容），以便当 AI 查看这个经过压缩的小型记忆时，它能获得与看到完整原始故事完全相同的“感觉”或结果。

4. 结果：快速且准确

作者在长文档（如医疗记录或长篇文章）上测试了这种方法，并与其他方法进行了比较。

速度： 他们可以在几秒钟内将内存缩小50 倍。
质量： 与会丢失准确性的总结不同，这种方法保持了 AI 回答问题几乎与拥有完整记忆时一样好的能力。
权衡： 它位于“帕累托前沿”上，意味着它在速度和质量之间提供了最佳平衡。它比缓慢的训练方法快得多，比快速的总结方法准确得多。

5. 一个特殊功能：“非均匀”压缩

本文还指出，AI 大脑（称为“头”）的各个部分并非同等重要。

类比： 在图书馆中，有些书架存放着最关键的书籍，而其他书架则存放着你很少需要的参考手册。
这种方法会找出哪些书架需要保持充实，哪些可以更激进地清空。它不会以相同的方式对待记忆的每一部分；它会给最重要的部分分配更多空间。

总结

本文提出了一种快速缩小 AI 记忆文件而不丢失细节的方法。它不像总结那样丢弃信息，也不像 Cartridges 那样花费数小时重新训练，而是利用一种数学技巧，保留最重要的部分并正确“加权”，使 AI 表现得仿佛它仍然记得一切。它使 AI 能够处理非常长的对话或文档，而不会耗尽内存或陷入混乱。

技术摘要：基于注意力匹配的快速 KV 压缩

1. 问题陈述

将大型语言模型（LLM）扩展到长上下文目前受限于键值（KV）缓存的内存需求。在自回归 Transformer 架构中，KV 缓存必须保留所有先前 token 的键和值，在长视野场景（例如多会话对话、代理编码）中，每个请求的缓存量往往达到数吉字节（GB）。

现有的减少 KV 缓存大小的方法通常分为两类，但两者均存在显著局限性：

Token 空间方法：诸如 token 驱逐（H2O, SnapKV）、token 合并（KVMerger）和摘要等技术，在高压缩比下性能急剧下降。特别是摘要方法，信息损失严重，损害了需要精确信息提取的任务的下游性能。
潜在空间方法：近期工作如 Cartridges（Eyuboglu 等人，2025）表明，可以在潜在空间中训练紧凑的 KV 缓存以匹配全上下文性能。然而，该方法依赖于针对每个上下文的端到端基于梯度的优化（前缀微调），这极其缓慢且昂贵（每个上下文需要数 GPU 小时）。

开放性问题是如何在潜在空间中实现快速、高质量的上下文压缩，在保留下游模型行为的同时，避免每个上下文进行梯度优化的计算成本。

2. 方法论：注意力匹配

作者提出了一组基于注意力匹配的方法，该方法构建一个紧凑的 KV 缓存 $(C_k, C_v)$ ，以针对一组参考查询复现原始缓存 $(K, V)$ 的注意力输出和注意力质量。

核心目标

目标是找到一个大小为 $t \ll T$ 的紧凑缓存 $(C_k, C_v)$ ，使得对于任意查询 $q$ 和任意附加的固定/未来 token $(K_{fixed}, V_{fixed})$ ，注意力输出得以近似保留。作者利用了连接块上的注意力分解，表明只需针对一组参考查询 $Q_{ref}$ 匹配两个量即可：

局部注意力输出：值的加权和。
注意力质量：未归一化注意力 logits 的总和（softmax 中的分母）。

为了处理紧凑缓存 token 较少（ $t < T$ ）从而无法自然匹配总注意力质量（例如，对于 $q=0$ ，质量为 $T$ 对比 $t$ ）的差异，该方法引入了每 token 标量偏置 $\beta$ 。优化目标为：
$\frac{\exp(qK^\top)V}{\sum \exp(qK^\top)} \approx \frac{\exp(qC_k^\top + \beta)C_v}{\sum \exp(qC_k^\top + \beta)}$
$\sum \exp(qK^\top) \approx \sum \exp(qC_k^\top + \beta)$

算法流程

该框架将问题分解为高效的闭式子问题，在压缩时避免了梯度下降：

参考查询生成 ( $Q_{ref}$ )：
- 重复预填充（Repeat-prefill）：在上下文本身（或其重复版本）上运行预填充过程，以提取查询激活值。
- 自学习（Self-study）：基于上下文生成合成交互，以拓宽查询分布。
- 同策略（On-policy）：按顺序压缩层，并从部分压缩的模型中提取查询，以减少分布偏移。
键选择 ( $C_k$ )：
- 最高注意力键：选择在参考查询中获得最高注意力分数的原始键（通过 RMS 聚合）。
- 正交匹配追踪（OMP）：贪婪地选择键以最好地满足注意力质量匹配目标，随后使用非负最小二乘法（NNLS）拟合权重。
偏置与值拟合 ( $\beta, C_v$ )：
- 拟合 $\beta$ ：通过非负最小二乘法（NNLS）求解，以匹配原始键的注意力质量。
- 拟合 $C_v$ ：通过普通最小二乘法（OLS）求解，以在给定选定键和拟合偏置的情况下匹配注意力输出。

非均匀压缩

作者观察到不同的注意力头对压缩表现出不同的敏感度。他们提出了一种非均匀压缩调度，其中不同的头和层被分配不同的保留率。该调度使用基于每头敏感度曲线的贪婪交换算法，在每次模型上预先计算一次，从而允许在不同上下文中稳健应用。

分块压缩

为了支持极长上下文（例如 60k+ token），该方法独立地对输入的连续分块应用压缩。默认使用基于 KV 的分块方法（预填充完整上下文、切片 KV 状态、压缩并合并），以保留全局位置语义和跨分块交互。

3. 主要贡献

注意力匹配公式：一种针对潜在空间压缩的新颖目标，直接优化注意力行为的保留而非输出似然，自然地分解为高效的线性代数问题。
闭式解：该方法消除了压缩过程中端到端基于梯度优化的需求，取而代之的是最小二乘和 NNLS 求解器。
速度 - 质量帕累托前沿：该方法在数秒内（单 GPU）实现了高达 50 倍的压缩比，比 Cartridges 快几个数量级，同时保持相当或更优的下游准确率。
非均匀调度：一种特定于模型的预计算预算分配策略，优先为对压缩最敏感的注意力头分配容量。
偏置机制：引入每 token 标量偏置以纠正注意力质量差异，使得在键集减少的情况下仍能实现精确的质量匹配。

4. 结果

该方法在 QuALITY（长文档问答）、LongHealth（密集临床记录）、QASPER（生成式问答）和 LongBench v2 上进行了评估。

性能与速度：在 QuALITY 上，50 倍压缩比下，注意力匹配方法描绘了帕累托前沿，优于 token 驱逐基线（H2O+, KVzip）和摘要方法。它们匹配或超过了 Cartridges 的准确率，同时速度快两个数量级（分钟级对比小时级）。
高压缩比：在极端比率下（高达 100 倍），该方法保留了显著的下游准确率（例如，LongHealth 上约为 60%，而摘要基线约为 30%）。
泛化性：该方法泛化到生成式任务（QASPER）和检索密集型基准（RULER），其中最快变体（AM-HighestAttnKeys-fast）在 RULER 上超越了最先进的 KVzip。
效率：在 Gemma-3-12B 上压缩 60k token 的上下文，根据键选择策略（OMP 对比最高注意力）的不同，耗时约 100–600 秒，主要受查询生成和键选择步骤主导。
在线压缩：AIME 的初步结果显示，重复的轨迹中途压缩保留了推理性能，将推理深度与物理内存限制解耦。

5. 意义与主张

本文主张，注意力匹配提供了一种实用、高效的 KV 缓存压缩原语，弥合了快速但有损的 Token 空间方法与缓慢但高质量的潜在空间方法之间的差距。

实用性：通过避免每个上下文的梯度下降，该方法使得高保真压缩在延迟和计算成本受限的现实世界部署中变得可行。
系统集成：该方法与现有推理引擎（如 FlashAttention、变长 KV 打包）兼容，并且可以在关键路径之外执行（解耦压缩）。
局限性：作者承认，虽然他们的方法比 Cartridges 快得多，但仍比简单的启发式方法慢。此外，在极端压缩比下（例如 LongHealth 上的 100 倍），基于梯度的方法如 Cartridges 可能仍然表现更优，因为它们能够搜索更广泛的表示空间，而不局限于原始键。
未来方向：本文建议未来的工作集中在优化查询生成、超越子集选择直接优化紧凑键，以及将压缩集成到推理引擎中以支持在线、长视野代理场景。

总之，这项工作确立了通过直接匹配注意力行为可以实现快速、高质量的上下文压缩，为长上下文 LLM 中的内存瓶颈提供了一种可扩展的解决方案。