🤖 machine learning

MomentKV: Closing the Directional Gap in KV Cache Eviction for Long-Context Inference

MomentKV 通过维护紧凑的矩统计量以确保被剔除 Token 的几何正则性，并为其注意力贡献提供闭式修正，从而解决了长上下文 KV 缓存剔除中的方向失配瓶颈，使其在各种基准测试和压缩水平下均显著优于现有方法。

原作者： Yu Li, Binxu Li, Tian Lan

发布于 2026-06-02

📖 1 分钟阅读☕ 轻松阅读

原作者： Yu Li, Binxu Li, Tian Lan

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

想象一下，你正试图记住一个非常长的故事，以便继续进行创作。为了实现这一点，你的大脑保留了一个关于目前所读内容的“草稿纸”（即 KV Cache）。问题在于，随着故事变得越来越长，这张草稿纸会变得巨大，最终占满你的整个桌面，让你无法开展工作。

为了解决这个问题，现有的方法使用了一种“垃圾桶”策略：它们观察故事，挑选出最重要的句子留在桌面上，然后把剩下的扔掉。它们假设如果保留了“最好”的句子，就可以忽略那些垃圾。

论文的大发现：“方向”问题
这篇论文的作者提出了 MOMENTKV，他们意识到这种“垃圾桶”方法中隐藏着一个缺陷。

想象一下，你正试图根据几片留在桌上的叶子来推测风向。

旧方法： 你保留那些吹得最猛的叶子（即最“重要”的叶子），并将其他的扔进垃圾桶。然后，你尝试仅利用桌上的叶子来推测风向。
问题所在： 你扔掉的叶子可能正朝着一个完全不同的方向（与你保留的叶子垂直）。即使你扔掉了 90% 的叶子，如果剩下的 10% 指向北方，而垃圾桶里的叶子指向东方，你的推测将会大错特错。你不能仅仅通过“重新归一化”（重新计算）北方的叶子来修复缺失的东风。错误不在于你扔掉了“多少”，而在于你丢失了方向。

解决方案：MOMENTKV
MOMENTKV 不仅仅是把垃圾扔掉并寄希望于结果，它会在垃圾进入回收站之前，为这些内容保留一份极小的、超紧凑的“摘要笔记”。

可以这样理解：

摘要笔记（矩统计量/Moment Statistics）： 在扔掉一个句子之前，系统会快速计算关于它的几个简单数值：“平均含义是什么？”以及“这个含义通常如何随下一个词发生变化？”它并不存储整个句子，只存储这些少量的“矩”（就像是一个统计学上的指纹）。
更聪明的丢弃策略： 在决定扔掉什么时，系统会询问：“这个句子是否已经被我的摘要笔记很好地代表了？”如果是，那么扔掉它是安全的。如果不是（即它拥有摘要无法捕捉的独特方向），它就会将其保留。这使得“垃圾”在几何上是规则且可预测的。
神奇的修正： 当模型需要写下一个词时，它不仅仅看桌上的句子。它利用来自垃圾桶的那些微小的“摘要笔记”，通过数学手段重建那些缺失的句子本会贡献什么。它本质上是在说：“我知道我把这些扔掉了，但根据我的笔记，我可以推测它们曾将故事推向这个特定的方向，所以我会把那部分补回来。”

为什么有效
论文在两个著名的 AI 模型（LLaMA 和 Qwen）以及两个基准测试（LongBench 和 RULER）上对该方法进行了测试。

结果： MOMENTKV 的表现始终优于所有其他方法，尤其是在“桌面空间”非常小（高强度压缩）的情况下。
类比： 这就像有一位图书管理员，他不仅在书架上保留最受欢迎的书籍，还为地下室里的每本书都保留了一张微小的索引卡。当你提问时，管理员可以使用这些索引卡瞬间召回地下室书籍的精髓，从而给你一个比完全忽略地下室要好得多的答案。

简而言之
目前的 AI 方法丢弃旧的上下文，并寄希望于剩余的部分足够应对。MOMENTKV 意识到，被丢弃部分的方向与这些部分本身一样重要。通过保留一份关于“垃圾”的微小且智能的摘要，并利用它在数学上修复 AI 的记忆，它让模型能够处理更长的故事，而不会迷失方向或出错。

技术摘要：MOMENTKV

问题陈述

基于 Transformer 的语言模型中的自回归解码依赖于键值（KV）缓存，该缓存存储了先前计算的投影以在后续解码步骤中进行复用。随着序列长度扩展到数万甚至数十万个 token，KV 缓存的内存占用量呈线性增长，成为长文本推理的主要瓶颈。

KV 缓存驱逐（KV cache eviction）通过保留固定大小的键值对子集并丢弃其余部分来解决这一问题。虽然现有方法改进了 token 选择策略（例如，基于累积注意力权重或基于值的评分），但它们都有一个共同的根本局限性：驱逐后的推理仅在保留的 token 上进行重归一化，完全丢弃了被驱逐的集合。

作者指出，输出性能下降的主要来源不仅是由于被驱逐 token 产生的残余注意力质量（现有方法已经最小化了这一点），而是由于保留集合与被驱逐集合之间的方向失配（directional mismatch）。实证分析表明，在实践中，被驱逐的 token 通常与保留的 token 近乎正交。因此，即使是极小比例的被驱逐注意力质量也会产生显著的方向误差，而标准的重归一化无法恢复这种误差。这造成了一个“方向间隙（directional gap）”，即仅仅提高 token 选择的精确度会带来收益递减，因为更具选择性的保留会将保留集合进一步集中在特定的查询对齐方向上，从而进一步扩大与被驱逐集合之间的差距。

方法论：MOMENTKV

为了解决这一方向间隙，作者提出了 MOMENTKV，该框架对被驱逐的 token 集维护紧凑、小规模的矩统计量（moment statistics）。这些统计量具有双重用途：指导驱逐决策并修正推理过程中的注意力输出。

1. 紧凑矩统计量

MOMENTKV 并不存储被驱逐的 token，而是为每个注意力头维护四个增量更新的运行总和：

计数 ( $n_e$ ): 被驱逐的 token 数量。
键总和 ( $s_k$ ): $\sum_{i \in E} k_i$ 。
值总和 ( $s_v$ ): $\sum_{i \in E} v_i$ 。
外积和 ( $S$ ): $\sum_{i \in E} v_i k_i^\top$ 。

这些统计量允许在不重新访问被驱逐 token 的情况下，计算出键均值 ( $\bar{k}$ )、值均值 ( $\bar{v}$ ) 和值-键协方差 ( $\tilde{S}$ )。总存储开销为每个头 $O(d^2)$ ，与上下文长度无关。

2. 被驱逐输出的一阶近似

在推理期间，MOMENTKV 推导出了被驱逐注意力输出 $\hat{f}_E(q)$ 的闭式一阶近似。通过围绕平均键对 logits 进行重中心化，并对 softmax 函数应用一阶泰勒展开，作者得出：
$\hat{f}_E(q) = \bar{v} + \frac{1}{n_e \sqrt{d}} \tilde{S} q$

$\bar{v}$ 项作为一个与查询无关的基准，恢复了被驱逐子输出的平均方向。
$\tilde{S} q$ 项作为一个与查询自适应的修正项，捕捉了基准本身无法捕捉的垂直变化（方向内容）。

这种近似允许模型重建被丢失的被驱逐贡献的估计值，有效地恢复了标准驱逐策略所抹除的方向信息。

3. 矩信息驱动的驱逐（Moment-Informed Eviction）

矩统计量同时也指导驱逐过程本身。对于每个保留的 token $j$ ，算法计算一个矩残差（moment residual）：
$r_j = v_j - \bar{v} - \frac{1}{n_e \sqrt{d}} \tilde{S} k_j$
该残差衡量了 token $j$ 在超出累积摘要所能捕捉的信息之外，携带了多少新的方向信息。驱逐评分定义为注意力权重与残差范数的乘积： $\text{score}(j) = \alpha_j \cdot \|r_j\|$ 。

低残差的 token 能被矩模型很好地预测，因此会被优先驱逐。
高残差的 token（携带独特的方向信息）即使其注意力权重较低，也会被保留。

这创造了一个相互增强的循环：基于矩信息的驱逐使被驱逐集合保持“几何正则性”（低 logit 离散度 $\sigma$ ），从而收紧了推理修正的近似界限。反过来，准确的修正又为未来的驱逐决策提供了更好的残差。

4. 归一化修正推理（Normalization-Corrected Inference）

最终的注意力输出被重建为保留输出 ( $f_R$ ) 与近似被驱逐输出 ( $\hat{f}_E$ ) 的凸组合：
$\hat{f}(q) = \hat{w}_R f_R(q) + (1 - \hat{w}_R) \hat{f}_E(q)$
混合权重 $\hat{w}_R$ 是通过利用 Jensen 不等式得到的被驱逐部分的配分函数下界来估计的。这确保了模型在修正方向的同时，能够考虑到被驱逐质量的大小。

核心贡献

识别方向间隙： 本文确定了 KV 缓存驱逐的主要误差来源是保留子输出与被驱逐子输出之间的近乎正交性，而非仅仅是由于被驱逐的注意力质量大小。
紧凑矩统计量： 引入了一种内存高效的方法，通过每个头仅需 $O(d^2)$ 的存储空间来追踪被驱逐 token 的几何特性（均值和协方差）。
闭式修正： 推导出了被驱逐注意力输出的一阶近似，能够恢复均值方向和依赖于查询的变化。
双重角色机制： 提出了一个统一的框架，其中矩统计量同时优化驱逐选择（通过识别冗余 token）和修正推理输出（通过恢复丢失的方向信息）。

实验结果

作者在 LongBench 和 RUL_ER 基准测试上评估了 LLaMA-3.1-8B-Instruct 和 Qwen3-4B-Instruct。

性能： 在所有缓存预算 ( $L \in \{128, 256, 512, 1024\}$ ) 下，MOMENTKV 的表现始终优于所有基线模型（H2O, SnapKV, PyramidKV, Ada-KV）。
激进压缩： 在激进压缩场景下观察到了最大的提升。在 LongBench 上使用 LLaMA-3.1-8B 且 $L=128$ 时，MOMENTKV 比最强的基线（Ada-KV）平均得分提高了 1.35，保留了 94.1% 的全量缓存性能。
检索任务： 在测试精确 token 检索能力的 RULER 上，MOMENTKV 分别比 LLaMA 和 Qwen 的 Ada-KV 实现了 +3.4 和 +3.5 的比例增益，验证了其在检索密集型设置中的有效性。
消融研究：
- 归一化修正推理 (NC) 的贡献（+2.41）显著高于矩信息驱逐 (MI) 单独的作用（+0.57），证实了恢复被驱逐方向是性能提升的主要驱动力。
- 两者的结合产生了协同效应（总增益为 +3.17），超过了单个增益之和。
效率： 该方法每 token 仅增加 1.6 ms 的延迟，并需要 4.1 MB 的辅助内存（约为保留缓存的 6.4%），同时相比全量缓存减少了超过 65% 的单 token 延迟。

重要意义

本文认为，进一步提升长文本推理性能的路径不仅在于更好的 token 选择，更在于恢复被当前驱逐策略所抹除的方向信息。MOMENTKV 证明了通过维护被驱逐集的紧凑统计量，可以弥合方向间隙，将驱逐带来的“不可约成本”转化为可修正的近似误差。这种方法为在不牺牲固定预算 KV 缓存效率的前提下部署长文本语言模型提供了一个实用且可扩展的解决方案。