🤖 machine learning

PolyKV: Heterogeneous Retention and Allocation for KV Cache Compression

PolyKV 是一个层级化的 KV 缓存优化框架，它通过将每个 Transformer 层动态路由到最合适的压缩策略并分配非均匀的缓存预算，从而提升长文本大语言模型的推理性能，其表现显著优于现有的统一单一策略基准方案。

原作者： Chao Fei, Panos Kalnis

发布于 2026-06-16

📖 1 分钟阅读☕ 轻松阅读

原作者： Chao Fei, Panos Kalnis

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

想象一下，你正试图记住一个非常长的故事，以便稍后回答相关问题。在人工智能（AI）的世界里，这种“记忆”被称为 KV Cache。随着故事变得越来越长，这种记忆会占用计算机硬盘上的大量空间，从而拖慢运行速度。

为了解决这个问题，科学家们通常使用一种“垃圾桶”策略：他们扔掉他们认为不重要的故事部分以节省空间。然而，大多数 AI 模型对大脑的每一个部分都使用相同的垃圾桶规则。它们假设处理故事开头的脑区与处理中间或结尾的脑区需要被完全一致地对待。

“PolyKV” 这篇论文指出，这就像是在打包行李时，无论每个隔层里实际装的是什么，都扔掉同类型的物品（例如，扔掉所有的袜子）。这种做法过于僵化。

以下是 PolyKV 如何通过简单的类比改变游戏规则的：

1. 问题所在：并非一刀切

把 AI 模型想象成一个由 30 个不同专家组成的团队，他们共同协作来理解一个故事。

专家 A（第 1 层）可能擅长记住角色名字。
专家 B（第 15 层）可能擅长理解情感基调。
专家 C（第 30 层）可能擅长预测下一个词。

目前，大多数 AI 模型会对所有 30 位专家说：“你们的笔记本里只能记 5 条笔记。”这就是“统一预算”（Uniform Budget）。

专家 A 可能需要 20 条笔记才能做好工作。
专家 C 可能只需要 2 条笔记。
通过强制每个人都只有 5 条，专家 A 会感到应接不暇（并导致出错），而专家 C 则拥有了并未使用的多余空间。

2. 解决方案：PolyKV（聪明的管理者）

PolyKV 引入了一位聪明的管理者，他在工作开始前会单独观察每一位专家。它为每位专家做出两个具体的决策：

决策 A：扔掉什么？（驱逐模式/Eviction Pattern）
- 专家 A 可能需要保留开头的几句话和结尾的几句话。
- 专家 B 可能需要保留最热门的句子（即“重量级选手”）。
- PolyKV 会询问每位专家：“你最依赖哪种类型的笔记？”并根据其需求给出一个定制化的保留规则。
决策 B：给多少空间？（预算/Budget）
- 如果专家 A 对信息丢失非常敏感，PolyKV 会给他们一本大的笔记本。
- 如果专家 C 比较强健且不需要太多信息，PolyKV 会给他们一个小笔记本。
- 所有笔记本的总大小保持不变，但分配过程会根据需求进行公平分配。

3. 它是如何工作的：“彩排”（离线校准/Offline Calibration）

你可能会问：“管理者如何在不减慢实际工作速度的情况下，知道每位专家需要什么？”

PolyKV 在正式开始工作前，会在一小段文本样本上进行一次快速彩排（称为“离线校准”）。

它观察每位专家在信息被移除时的反应。
它学习到：“哦，如果移除开头，专家 A 会感到困惑，但专家 B 并不在乎。”
它根据这次彩排写下一份固定的计划。
当正式工作开始时，管理者只需执行该计划即可。实际对话过程中无需进行思考，因此保持了高效。

4. 结果：更好的记忆，同样的空间

研究人员在两个流行的 AI 模型（LLaMA 和 Qwen）上使用标准的内存限制进行了测试。

旧方法： 对所有人使用单一规则，AI 在长故事测试中的得分约为 36.35。
PolyKV 方法： 通过为每个专家定制规则和笔记本大小，得分跃升至 37.79。

这听起来可能很小，但在 AI 领域，这是一个巨大的胜利。这意味着 PolyKV 恢复了“垃圾桶”法与“完美记忆”法（后者使用的空间过多而不切实际）之间 54.5% 的性能差距。

5. 它的优势与局限

优势： PolyKV 非常擅长需要理解大局观或上下文的任务，例如回答关于长文档的问题或总结故事。它知道如何为大脑的每个部分保留“重要”的部分。
局限： 它在处理“大海捞针”（在海量文本中寻找一个特定的、微小的细节）或编程任务时有时会遇到困难。这表明虽然 PolyKV 擅长通用理解，但它有时会扔掉一些在彩排阶段被专家认为不重要、但在随后却变得至关重要的“针”。

总结

PolyKV 就像是从一个每个工人获得完全相同的工具和工作空间的工厂流水线，升级到了一个定制化的工作室，在这里每个工人都能获得完成其特定工作所需的特定工具和桌面大小。通过这样做，整个团队的工作效率更高，记忆更深，并且仍能保持在相同的总空间内。

技术摘要：PolyKV：用于 KV 缓存压缩的异构保留与分配

问题陈述

大语言模型（LLM）的长文本推理面临着由于键值（KV）缓存增长而导致的显著内存瓶颈。虽然现有的 KV 缓存驱逐方法通过保留特定 token 子集有效地减少了内存使用，但它们通常在所有 Transformer 层中应用统一的压缩策略和统一的缓存预算。

这种统一的设计忽视了 Transformer 层的异构性。先前的研究表明，不同层在注意力模式、表示抽象、上下文角色以及对剪枝的敏感度方面存在差异。因此，单一的全局驱逐规则是一种粗略的抽象，未能考虑到以下事实：

不同层可能偏好不同的驱逐模式（例如，重击者 heavy hitters vs. 注意力汇聚点 attention sinks）。
不同层对 KV 缓存缩减的敏感度不同，这意味着它们需要不同的缓存容量。

目前的自适应方法已开始探索非统一预算，但往往未能将驱逐模式和缓存预算视为耦合的、层级化的设计选择。

方法论：PolyKV

作者提出了 PolyKV，该框架将 KV 缓存驱逐重新定义为一个层级化的设计空间。PolyKV 明确地为每一层及每个推理阶段（预填充 prefill 和解码 decoding）耦合了两个决策：

驱逐什么： 为该层选择一个合适的特定驱逐模式（例如，H2O、TOVA、PyramidKV）。
保留多少： 为该层分配特定的缓存预算。

核心组件

设计空间公式化： PolyKV 定义了一个策略 $\Pi$ ，其中每一层 $l$ 在每个阶段 $\phi$ 被分配一个驱逐模式 $r_{l,\phi}$ 和一个预算 $b_{l,\phi}$ 。总预算受限于一个固定的全局限制 $B$ 。
离线校准： 为了避免穷举式的在线搜索，PolyKV 使用一个小型的校准集来生成一个固定的静态策略。它收集了两类层特定信号：
- 驱逐偏好信号： 衡量哪种驱后模式最能保留该层的 FullKV 计算。指标包括扰动后的激活值与 FullKV 激活值之间的余弦相似度，以及困惑度（PPL）敏感度（由特定模式引起的性能退化）。
- 缓存预算敏感度信号： 估计层受 KV 缩减影响的程度。指标包括注意力熵（弥散的注意力需要更大的预算）和 PPL 敏感度（聚合后的最坏情况退化）。
策略构建：
- 模式选择： 对于每个“层-阶段”对，选择具有最高偏好得分（例如，最高的余弦相似度或最低的 PPL 退化）的驱逐模式。
- 预算分配： 根据敏感度信号按比例分配预算。具有更高敏感度（或熵）的层会获得更大的预算，但受限于最小预算约束（ $b_{min}$ ）。
实现： PolyKV 实现为一个基于 hook 的框架。它在生成流水线的常见点（pre-kernel、post-kernel、缓存更新）注入逻辑，以支持各种驱逐方法的模块化集成，而无需修改底层模型代码。

核心贡献

重构 KV 驱逐： 本文将 KV 缓存驱逐重新定义为一个层级化的设计空间，将驱逐模式的选择与在共享内存约束下的缓存容量分配相结合。
PolyKV 框架： 一个通过在层粒度上组合驱逐模式和预算来使该设计空间具体化的系统。它通过统一接口集成多种驱逐原语，并使用离线校准信号来指导策略构建。
实证验证： 作者证明了在固定的内存预算下，异构的层级化策略优于统一的单策略基线。

实验结果

实验在 LLaMA-3.1-8B 和 Qwen3-8B 上使用 LongBench 和 RULER 基准测试进行。

在 LongBench 上的表现（平均预算 512 个 token）：
- LLaMA-3.1-8B： PolyKV 恢复了最强单策略基线与 FullKV 之间性能差距的 54.5%。
- Qwen3-8B： PolyKV 恢复了差距的 25.7%。
- 在整个预算范围（128–1024 个 token）内，PolyKV 始终优于最强的基线，提升幅度在 1.7% 到 6.4% 之间。
在 RULER 上的表现：
- 结果褒贬不一。Poly-KV 在 LLaMA 上略微超过了基线，但在 Qwen 上表现不佳（下降了 3.98 分）。作者将其归因于 RULER 任务的性质，这些任务依赖于稀疏的证据 token，如果校准信号不能完美识别这些 token，它们可能会被异构策略丢弃。
消融研究：
- 驱逐模式 vs. 预算： 主要的性能增益来自于为不同层分配不同的驱逐模式，特别是在**预填充（prefill）**阶段。
- 预算分配： 非统一的预算分配提供了额外的收益，但其效果更依赖于模型且鲁棒性较低。当总预算足以避免饥饿敏感层时，其效果最为显著。
任务分析： 在需要广泛上下文理解的任务（如问答、少样本学习）中，增益最为显著。在需要精确检索稀疏 token（如 RULER 中的 NIAH）或严格遵循局部形式（如代码生成）的任务中，性能有所下降。

重要性与主张

本文声称，将 KV 缓存驱逐视为一个组合式的、层感知的优化问题，而非单一的全局保留策略，可以获得更优的结果。

异构性是关键： 研究验证了 Transformer 层并非同质的；“一刀切”的驱逐策略是次优的。
静态效率： 通过使用离线校准，PolyKV 在不引入在线策略搜索运行时开销的情况下实现了这些异构收益，使其在实际部署中具有可行性。
局限性： 作者坦诚地指出，目前的校准信号并不总是能保留精确检索任务所需的“稀疏但决定性”的证据 token，这为未来改进信号设计指明了方向。

总之，Poly-KV 证明了通过使驱逐策略和缓存容量与每个 Transformer 层的特定特征相匹配，可以显著降低长上下文 LLM 推理中的内存-质量权衡。