跳转至

Adaptive Mass-Segmented KV Compression for Long-Form Reasoning

会议: NeurIPS 2026
arXiv: 2605.23200
代码: https://github.com/EIT-NLP/AdaptiveMassSegment
领域: LLM 效率
关键词: KV 缓存压缩、区域配额、注意力质量分布、自适应切段、长程推理

一句话总结

AMS 不替换 token 重要性评分器，而是先用注意力质量分布自适应切段并分配保留配额，再在段内评分筛选，从而缓解长推理中连续区域被清空的问题；Math500 上 AMS-Expected 在 256-token 缓存预算下比 AdaKV-ExpE2 高 16.0 个百分点。

研究背景与动机

大语言模型生成长思维链时，历史 token 的 key/value 会逐步累积。即使模型权重不变，每一步解码仍要读取越来越长的缓存，显存占用与访存开销随之增长。TOVA、Expected Attention、KeyDiff、R-KV 等免训练压缩方法用注意力、几何变化或冗余程度判断哪些 token 值得留下，再周期性删除其余 KV。问题在于，评分器回答的是“这个 token 有多重要”，而全局 Top-k 默认把所有位置放在同一个竞争池里：若一组局部峰值足够突出，一整段较分散的中间推导可能一个位置也保不住。

这种失败并不等同于删除了几个低分词。一个中间推导区可能保存着原题约束、临时结论及其依据；整体消失后，后续生成可能改写题意、推翻正确结论，或者反复回到已做过的推导。作者把连续区域遭到严重清空的现象称为区域清空（Region Wipe-out）。固定长度分块可以减少直接的 token 竞争，但缓存重要性密度并不均匀，也会随着解码变化：同样长的两个区间，可能一个包含密集的有效推导，另一个主要是冗余叙述。只按固定位置分块，仍可能把宝贵预算给错区域。

AMS 的切入点是把“哪里应该有记忆容量”与“该区域内留下哪些位置”拆开。前者用近期注意力形成的空间分布决定，后者继续交给原有评分器；跨压缩事件的历史信号只平滑区域分配，不改写基础评分规则。核心 idea：先为注意力质量分布定义的自适应区域分配可行的保留配额，再让 token 在各自区域内竞争，避免全局高分尖峰独占缓存。

方法详解

整体框架

AMS 是解码阶段的缓存选择包装层，不是新模型，也不需要训练。模型先正常处理完整提示，生成过程中按固定间隔触发压缩；每次在各层、各 KV head 的当前缓存上独立执行选择，输出长度为目标预算的紧凑 KV。需要区分两个输入：近期注意力使用量负责构建质量分布（quality mass），基础评分器负责给段内候选排序，两者并不要求相同。

一次压缩先构建质量分布，并可用 EMA credit 混入历史；然后沿当前缓存顺序按累计质量切段，修正过短或过长的段；接着安排段级预算和必须保留位置；最后进行段内 Top-k、预算校正及 gather。这里的“区域”是当前缓存中的连续位置，不是模型读懂后划出的完整语义步骤，压缩后也不再与原始生成位置一一等长对应。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["当前 KV + 近期注意力"] --> B["质量分布与历史平滑"]
    B --> C["累计质量自适应切段"]
    C --> D["段级配额与保留优先级"]
    D --> E["段内选择与物理压紧"]
    F["基础评分器"] --> E
    E --> G["紧凑 KV → 继续解码"]
    E -->|保留索引对齐历史状态| B

图中只有推理数据流与跨压缩事件的状态复用，没有训练监督。基础评分器可以是 TOVA、Expected Attention、TriAttention、KeyDiff 或 R-KV；AMS 改变的是评分结果如何被用于选择，而不是要求这些评分器都变成注意力评分。

关键设计

1. 质量分布与历史平滑：用区域使用密度安排容量，而不冒充新的 token 评分器

作者先收集最近 \(W\) 个解码 query 对各缓存位置的注意力使用量。实现中沿序列做小核一维平均池化，以削弱局部噪声；近期后缀因因果掩码只被部分 query 看见，缺失观测用窗口最大分数填补，避免新生成位置仅因被观测次数少而吃亏。这些使用量经过非负化和加小常数后归一化：

\[ m_i=\frac{\max(u_i,0)+\epsilon}{\sum_{v=1}^{T}(\max(u_v,0)+\epsilon)},\qquad \epsilon>0. \]

这里 \(u_i\) 是使用量，\(m_i\) 是总和为 1 的质量份额。它不是“该 token 包含多少正确推理信息”的真值，也不是基础评分器 \(g_i\) 的另一种名字：质量分布决定区间边界和预算，\(g_i\) 决定给定预算下段内哪些位置胜出。比如配合几何评分器时，区域容量仍来自注意力使用量，而区域内竞争可以依据 key 的几何特征；AMS 的可插拔性正来自这种职责分离。

当前 query 可能短暂集中在某一段，若每次都按瞬时分布重分预算，边界和保留集合容易抖动。AMS 在每层、每个 KV head 上维护历史 credit，先更新其 EMA，再与当前质量混合：

\[ c\leftarrow\lambda c+(1-\lambda)m^{(e)},\qquad m_{\mathrm{used}}=\operatorname{normalize}\!\left(\beta m^{(e)}+(1-\beta)\operatorname{normalize}(c)\right). \]

这只改变后续切段与配额所用的质量，不改变基础评分器的分数。默认 \(\lambda=0.9\)、\(\beta=0.9\)，意味着主要依据当前使用量，同时让持续有用的历史位置对分配产生影响；不能把它解释成 90% 的最终质量都来自历史。

credit 必须跟 KV 的实际位置保持对齐：缓存压紧后，旧 credit 应随同一保留索引保留下来，新 token 加入时也需要对应的新状态，否则同一向量下标会指向不同历史 token。原文明确给出 EMA 更新与跨事件复用，但没有完整展开 credit 的压缩后重索引及新位置初始化规则；这属于复现时需要核实的状态管理细节，不能自行补成作者的精确实现。历史平滑也无法恢复已经被删除的 KV。

2. 累计质量自适应切段：在高密度区域切得细，在低密度区域覆盖得长

AMS 沿当前缓存序列求质量前缀和，对 \(\Delta,2\Delta,3\Delta,\ldots\) 这些累计质量阈值，寻找首次达到阈值的位置作为切点。也就是说，它不是每隔固定数量 token 切一刀，而是每积累大致相同的注意力质量切一刀：高密度区域很快达到阈值，段更短；低密度区域需要走过更多位置，段更长。核心边界规则是：

\[ C(t)=\sum_{u=1}^{t}m_u,\qquad b_k=\min\{t\mid C(t)\geq k\Delta\}. \]

初始段大致具有相同质量，因此高密度区域会产生更多段。由于每段随后都有最低配额，细分本身就让这些区域获得更密的结构保护；即使初始各段质量接近，最低配额按 token 长度看也不是均匀分配。遇到单位置尖峰或长低密度区间时，作者再用 split/merge 启发式修正段长：超过最大长度的段拆成近似等长子段，过短的相邻段合并。修正后各段的质量不必继续相同，因此仍要重新统计质量用于预算分配。

默认目标段质量为 0.1，最小/最大段长为 16/256。长度约束避免给极短段无意义地重复设保底，也避免一个低密度区域延伸得过长、只有极少位置能代表它。原文是启发式处理，而不是解析出的最优语义分段；不能宣称这些边界等于每个推理步骤的起止位置。

3. 段级配额与保留优先级：保底有预算前提，必须保留集合先于普通竞争

给定切好的段，作者先为第 \(i\) 段安排不超过段长的最低配额，再把剩余容量按段质量分配。忽略 must-keep 的额外修正时，主文写出的分配形式为：

\[ q_i^{\min}=\min(q_{\min},L_i),\qquad T_{\mathrm{rem}}=T_{\mathrm{keep}}-\sum_i q_i^{\min},\qquad q_i=q_i^{\min}+\operatorname{round}\!\left(T_{\mathrm{rem}}\frac{M_i}{\sum_j M_j}\right). \]

\(L_i\) 是段长，\(M_i\) 是该段质量之和。默认每段最低配额为 1。取整后总数可能不等于目标预算，某段也可能被分到超过自身长度的容量，因此还需裁剪和重新分配。这一层控制的是“某区域有几个名额”，并不指定留下该段的开头、结尾或完整句子。

“每段至少留一个位置”首先要求总预算足以支付所有最低配额。加入必须保留位置后，更应看这些位置占用容量后，未被它们覆盖的区域保底是否仍可满足；不能只检查总段数。AMS 默认保护开头的 4 个 sink token，并保留近期后缀。主文说明实现先插入 must-keep，再安排剩余段级预算；若剩余预算为负，会缩短近期后缀，同时优先保护 sink。任意小预算下，既保留全部 sink、全部近期后缀、又满足每段保底，不可能无条件同时成立。

段内结果与 must-keep 取并集后，重复位置只算一次。超预算时删除低分的非 must-keep 位置，欠预算时从尚未选择的位置补高分 token。原文也允许必要时下调段级配额，因此“保底防止区域清空”是可行预算下的设计目标，不是对所有超限修正后的任意区域都成立的强定理。尤其不能据此推导“所有推理信息都保留”：一个代表 token 与完整的约束、推导和证明显然不是同一回事。

4. 段内选择与物理压紧：既改变竞争范围，也真正删除未保留的 KV

基础评分器为每个缓存位置给出分数，张量形状为 \([B,H_{kv},T]\)。AMS 按段内配额局部取高分位置，然后执行上述并集、去重和预算修正，将保留索引排序，用相同索引 gather key 和 value。不同 KV head 可以选择不同的历史位置，但每个 head 仍得到固定长度的紧凑缓存；后续解码继续读取这些实际留下的 KV，而不是读取完整缓存再把不需要的位置遮住。

这种 gather-and-compact 与 mask-only 有本质差别。本文中的 AdaKV-ExpE2 实现仅在注意力掩码上施加有效预算，未把全部淘汰 KV 物理释放；AMS 则把实际缓存长度缩到预算。因此“相同有效长度”的准确率比较，不自动意味着“相同物理显存”的系统比较。这个区别属于本文使用的实现路径，不应外推为所有 AdaKV 实现都不能压紧。

附录 E 进一步把选择策略和分页存储布局分离。AMS 输出逐 head 的保留索引，runtime 分配新块，按 head 将原始位置的 KV 复制到新物理槽，替换请求的 block table，再释放旧块。压紧后的同一槽在不同 head 上可以容纳来自不同原始位置的 KV；选择已经落实到内容里，稳态注意力不需要额外逐 head 间接索引。RoPE 已编码在缓存 key 中，但新 token 的逻辑位置仍必须接着原始生成进度，不能把压紧长度误当作下一步位置。

这里的“兼容 vLLM”是系统接口与运行时路径的验证。主实验走 HuggingFace/KVPress，补充代码有参考适配器与 vLLM 风格 runtime hooks；作者明确说它不是上游 vLLM 补丁，也不是完整优化的生产服务性能基准。稳态读取路径不加额外索引，不代表压缩事件、GPU 拷贝、临时新旧块共存或批处理调度完全没有开销。

一个完整示例

以下是说明预算逻辑的构造例子，不是论文新增实验。假设一个 KV head 当前有 16 个位置，最终保留预算为 8；经过切段与段长修正，形成 4 个长度为 4 的段，每段质量都是 0.25。最低配额为 1，先花掉 4 个名额，剩余 4 个名额按质量均分，最终每段各有 2 个名额。

假设必须保留的开头和结尾位置已分别纳入首段和末段的两个名额，预算修正后仍每段保留 2 个位置。基础评分器此时在各段内选高分候选，而不是让 16 个位置直接竞争 8 个名额。即便第二段的最高分低于其他段里的多个尖峰，它仍有代表位置；全局 Top-8 则可能把第二段全部删除。

下一轮新 token 加入后，AMS 会基于当前压紧缓存及新增位置重新估计质量、重新切段。它不把原来的四段永久锁住，也不会因为示例中保留了第二段两个位置，就保证第二段承载的全部语义仍完整可用。

损失函数 / 训练策略

方法没有训练损失、不更新模型权重，是免训练的解码时策略。主实验默认每生成 512 个 token 触发一次压缩，目标有效缓存长度为 256、512 或 1024，隐藏状态缓冲区为 256；质量估计的近期 query 窗口为 128。间隔压缩意味着缓存会在事件之间增长，不能将目标长度理解成任意时刻物理缓存都严格等于该值。

AMS 新增质量归一化、前缀和、段长修正与配额分配。附录 H 将质量处理和前缀和的开销描述为 \(O(H_{kv}T)\)，但这不是整个压缩器的完整成本：基础评分、段内选择、KV 搬运与 runtime 调度仍需单独计入。

实验关键数据

主实验

下表摘取原文表 1、表 5 的 Math500 pass@1，均使用 DeepSeek-R1-Distill-Qwen-7B。列名是压缩后的目标有效缓存长度，数值单位为 %；Full KV 为无压缩参照，不属于任一固定预算。

方法 256 512 1024
Full KV 52.80 52.80 52.80
TOVA 29.20 44.60 48.80
AMS-TOVA 36.40 48.40 53.40
AdaKV-ExpE2 32.60 46.60 53.40
AMS-Expected 48.60 54.00 54.20
TriAttention 55.00 56.80 57.20
AMS-TriAttention 56.20 60.60 63.60

AMS-Expected 相对 AdaKV-ExpE2 的增益是 16.0、7.4、0.8 个百分点；AMS-TOVA 相对 TOVA 是 7.2、3.8、4.6 个百分点。AMS-Expected 在 512/1024 预算下略高于表中的 Full KV，但这只说明该固定评测中存在这样的结果，不能当作压缩必然优于完整上下文的规律。更强的 TriAttention 本身已经高于 Full KV，接入 AMS 后还能提高 1.2、3.8、6.4 个百分点，支持“分配层可以补充强评分器”的解释。

跨任务结果也不是全面胜出：原文表 6 中 AMS-Expected 的 RepoBench-P 得分为 27.44，Full KV 为 24.04；但 NIAH 为 0.2209，仍低于 Full KV 的 0.3751。跨骨干方面，OpenThinker3-7B 在 512 预算下 AMS-Expected 为 45.4%，无压缩参照为 45.8%；32B Math500 在同预算下为 43.00%，无压缩为 47.80%（表 3、表 4）。

消融实验

下表来自原文表 7：Math500、7B 骨干、TOVA 评分器，指标为 pass@1（%）。

配置 512 1024 说明
完整 AMS 48.4 53.4 自适应切段、质量配额、EMA
去掉质量加权配额 46.0 52.8 分别下降 2.4、0.6 个百分点
去掉 EMA credit 48.0 50.8 分别下降 0.4、2.6 个百分点
固定长度段 48.0 50.4 分别下降 0.4、3.0 个百分点
全局 head-only Top-k 42.8 49.2 分别下降 5.6、4.2 个百分点

全局 head-only Top-k 消融并不等于直接复用表 1 的裸 TOVA 数值：两处 512 结果分别是 42.8 和 44.6，不能混为同一配置。此处最有力的证据是，在作者的消融控制下去掉区域配额损失最大；EMA 与自适应边界的价值则在 1024 预算下更明显。

原文表 8 的系统数据另列如下。峰值显存包含运行时的整体分配，不是纯 KV 字节数；显存与时间列的预算不同，不能交叉配对计算吞吐。

方法 峰值 GB：512 峰值 GB：1024 秒/样本：128 秒/样本：512
StreamingLLM 15.0 14.9 50.1 48.4
TOVA 15.0 14.9 67.1 52.3
PyramidKV 15.3 15.6 67.5 51.6
AdaKV-ExpE2 39.6 39.7 91.8 62.3
AMS-Expected 15.0 14.9 41.4 44.0

关键发现

  • 紧预算更能显示结构保护的价值：256 预算下，KeyDiff 接入 AMS 从 22.80% 到 42.80%（表 5），不是换评分器带来的收益。
  • 附录 C 的区域清空率从 11.3% 降到 8.7%，仍未归零；保留集合 IoU 的提高支持更稳定的历史选择，但不等于证明语义信息无损。
  • 自由生成中，TOVA/AMS-TOVA 的重复率分别为 21.58%/16.18%，平均生成 token 为 2963.6/2799.7，时间为 67.34/63.85 秒（表 9）。时间改善包含少生成重复内容的影响，不是单步 kernel 加速率。
  • 小样本敏感性实验需要克制解读：表 16 中窗口 16/32/64 均为 50.0%，附录的部分实验仅用 Math500 的 10% 子集，不能推出所有窗口与任务都不敏感。

亮点与洞察

  • AMS 的主要贡献是改变竞争规则，而非发明又一个重要性分数。保底配额把区域覆盖作为选择约束，使强评分器仍能负责局部排序。
  • 质量均衡切段把“重要性密度”变成分辨率：高密度区的段更多、更短，最低配额也就更密。这解释了为何切段与分配必须联合理解，而不能仅将 AMS 看作普通分块 Top-k。
  • EMA 应用于预算而非直接修改评分，使历史稳定性与评分器解耦。这个思路可迁移到其他受容量约束的在线记忆选择，但需要同样严格地管理状态与对象身份的对齐。

局限与展望

  • 注意力使用量只是长期效用的代理。低注意力的原题约束或稍后才用到的证据仍可能被淘汰，段内保留少量位置不保证完整推理语义。
  • 最低配额、must-keep 和总预算存在可行性约束；原文的全局修正可能改变段级数量。复现应明确记录无法同时满足保护规则时的处理，而不是声称任何预算都无区域清空。
  • 系统主结果来自 HuggingFace/KVPress，使用 A800 80 GB，默认每任务 2 张 GPU。vLLM 路径是兼容性验证，生产环境中的动态批处理、块分配压力和压缩瞬时峰值仍需专门测量。
  • 原文有表号交叉引用漂移：GSM8K 实际是表 2，通用任务是表 6，显存/延迟是表 8；本笔记按表题及表内值引用。附录表 17 的 metric_main 为未充分解释的原始指标，未将其当作准确率或据此补造实验结果。

相关工作与启发

  • vs TOVA / Expected Attention：这些方法提供 token 级重要性依据；AMS 提供 head 内沿时间轴的区域配额，两者可以组合。AdaKV-ExpE2 还含 head 自适应分配，因此与 AMS-Expected 的对比不只是唯一一个局部开关变化。
  • vs ChunkKV / 固定分块：分块保护局部结构，但切分粒度相对固定。AMS 依赖近期质量分布调整边界，并在段内继续用基础分数选择，不要求整块一起保留。
  • vs ReST-KV / G-KV：相关方法用历史信息稳定淘汰信号；AMS 把历史平滑放在分配层。值得进一步区分“更好的分数”和“更好的预算组织”各自能解释多少收益。

评分

  • 新颖性: 4/5 — 将自适应时间区域配额作为独立包装层，区别于单纯重设评分或 head 预算。
  • 实验充分度: 4/5 — 覆盖多评分器、骨干和任务，仍缺完整生产服务基准与更充分统计不确定性。
  • 写作质量: 3/5 — 核心机制清楚，但表号引用与部分附录指标、状态管理细节需要澄清。
  • 价值: 4/5 — 为免训练 KV 压缩提供可组合的结构保护思路，尤其适合紧预算长推理。