💬 NLP

TokenButler: Token Importance is Predictable

TokenButler 是一种轻量级且感知查询的预测器，它通过蒸馏掩蔽因果注意力分布来动态识别关键令牌以进行高效的 KV 缓存管理，在不永久驱逐令牌的情况下实现了接近神谕的检索准确率和显著的延迟降低。

原作者： Yash Akhauri, Ahmed F AbouElhamayed, Yifei Gao, Chi-Chih Chang, Sameh Gobriel, Nilesh Jain, Mohamed S. Abdelfattah

发布于 2026-05-18

📖 1 分钟阅读☕ 轻松阅读

原作者： Yash Akhauri, Ahmed F AbouElhamayed, Yifei Gao, Chi-Chih Chang, Sameh Gobriel, Nilesh Jain, Mohamed S. Abdelfattah

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

想象一下，你正试图与一位非常聪明但略显健忘的图书管理员进行对话。这位图书管理员（即人工智能）阅读过海量的书籍（即训练数据），现在正试图根据你刚刚交给他们的一段非常长的特定故事（即上下文）来回答你的问题。

问题在于，这个故事实在太长了——有时长达数十万字——以至于图书管理员的办公桌（即计算机内存）变得完全杂乱无章。为了跟上进度，图书管理员必须保留一份他们迄今为止读过的每一个单词的实时清单（这被称为KV-Cache）。随着故事变长，这份清单变得过大，无法放在办公桌上，导致一切进程慢如蜗牛。

旧方法：丢弃或分组

为了解决这个问题，以往的方法主要尝试了两种方案，但两者都存在缺陷：

“垃圾桶”方法：一些图书管理员决定，一旦办公桌满了，就直接将清单中的旧单词扔掉。
- 缺陷：想象一下，故事开头提到了一个名叫“齐拉梅尔格罗夫”（Ziramelgrove）的角色。图书管理员当时觉得这个名字不重要，就把它扔掉了。但在 50 页之后，你问：“齐拉梅尔格罗夫是谁？”图书管理员完全不知道那是谁，因为他们把这个名字扔进了垃圾桶。
“盒子”方法：其他图书管理员保留了所有单词，但将它们整理成大盒子（页面）。当他们需要查找某样东西时，就会拿走整个盒子。
- 缺陷：如果重要的单词“齐拉梅尔格罗夫”恰好被分割在两个盒子之间，图书管理员可能会拿错盒子，或者因为他们是把整个盒子作为一个整体来看待，而不是关注盒子里的特定单词，从而完全错过这个词。

新方案：TokenButler

这篇论文介绍了TokenButler，一位智能助手，它能帮助图书管理员决定确切应该保留哪些单词在桌面上，而无需永久丢弃任何内容。

将TokenButler想象成一位站在图书管理员身旁、受过高度训练的观察员。

工作原理：图书管理员不再猜测哪些单词重要，而是由 TokenButler 查看你当前的问题（即“查询”），并预测出长故事中究竟需要哪些特定单词来回答它。
魔法技巧：它无需重新阅读整个故事就能知道这一点。它使用一个微小的、轻量级的“作弊小抄”（一个小预测模型），该模型在训练期间学会了识别模式。它知道，如果你询问 1 万字之前提到的某个特定地点，那么即使 10 秒前它看起来无关紧要，这个地点现在突然变成了宇宙中最重要的事物。

为何更优

这篇论文在一个“词语捉迷藏”游戏中测试了该方法。

测试：故事在早期隐藏了一个秘密地点名称，然后用长时间的数学问题和烹饪技巧来分散读者的注意力，最后才问：“这个地点在哪里？”
结果：“垃圾桶”和“盒子”方法经常失败，因为它们要么扔掉了地点名称，要么在错误的盒子里找不到它。然而，TokenButler 成功地让地点名称保持就绪，几乎每次都能找到它，就像一个“神谕”（完美的预测者）那样运作。

速度与效率

你可能会认为增加一位观察员会拖慢图书管理员的速度。但论文展示了 TokenButler 避免这种情况的两种巧妙方法：

“批处理”技巧：图书管理员不是每写一个单词就请观察员检查一次清单，而是每隔几个单词请观察员检查一次。观察员会说：“保留这些单词”，然后图书管理员在接下来的几步中保留它们。这使得整个过程快得多。
“邻居”技巧：观察员知道重要信息通常成簇出现（如全名或句子）。因此，如果观察员选中了某个特定单词，它也会顺便抓取紧邻的单词，以防万一。这确保了即使重要性发生轻微偏移，他们也不会遗漏任何内容。

核心结论

TokenButler 使计算机能够在不耗尽内存或降低速度的情况下阅读和理解海量故事（高达 100 万字）。它通过学会预测当前问题究竟需要哪些单词来实现这一点，从而保持内存的整洁和快速，同时确保没有任何关键细节被意外丢弃。

在测试中，该方法使计算机在显卡上运行时的速度提高了1.6 倍，而在计算机必须从主处理器借用额外内存时，速度提高了7.6 倍，同时保持了与在桌面上保留每一个单词时完全相同的回答准确性。

技术摘要：TokenButler

问题陈述

随着上下文长度扩展至 64k–1M 个 token，大型语言模型（LLM）面临显著的内存与计算瓶颈。用于存储 token 历史以实现高效解码的键值（KV）缓存随上下文长度线性增长，消耗了大量的 GPU 显存和带宽。

现有的缓解该瓶颈的方案通常分为三类，每类均有其特定的局限性：

静态策略（如 StreamingLLM）： 使用固定规则（如滑动窗口）来限制缓存容量。这些策略缺乏查询感知能力，无法适应动态的输入需求。
自适应淘汰（如 H2O、SnapKV）： 基于注意力分数等指标，永久丢弃被认为低重要性的 token。这可能导致丢失在共指语境（即早期提及的实体在很久之后被引用）中后期变得关键的 token，因为重要性往往是非单调的。
自适应动态策略（如 Quest）： 保留完整缓存，但在解码时获取 token 的子集（页）。虽然它们避免了永久淘汰，但其依赖粗粒度的页级检索，在信息密集且上下文丰富的任务中表现不佳，因为关键信息可能跨越页边界，或者需要细粒度的选择。

此外，许多现有方法依赖不准确的 token 重要性代理指标（如近期性或聚合注意力分数），或者在计算每个 token 的完整注意力分数时产生高昂的计算开销。

方法论：TokenButler

作者提出了 TokenButler，这是一种细粒度、查询感知的预测器，旨在识别关键 token，而无需永久淘汰它们或依赖粗粒度的分页。

核心架构

TokenButler 通过在固定深度步长（生成层）预测低维重要性查询，并将其与后续层（消费层）中真实 KV 缓存键的学习投影相结合来运行。

查询预测： 一个轻量级多层感知机（MLP）接收来自“生成”层的隐藏状态，并为每个注意力头生成 $G$ 个特定槽位的查询向量。这些向量是低维的（ $d' \ll E$ ）。
键缓存投影： 对于每个消费层，实际的高维 KV 缓存键使用一个学习到的矩阵投影到相同的低维空间。
评分： 将预测的查询与投影后的键进行点积，以低成本生成 token 重要性分数，从而在固定预算下实现动态的 per-token 选择。

训练

蒸馏： TokenButler 通过蒸馏冻结的基础 LLM 的掩码因果注意力分布进行训练。目标是最小化教师注意力 logits 与学生（TokenButler）预测 logits 之间的交叉熵损失。
效率： 为了在训练期间处理长上下文，损失函数在查询位置的子采样集上计算，重点关注后期上下文区域和最终 token，将复杂度从 $O(L^2)$ 降低到 $O(RL)$。
开销： 该预测器为基础模型增加了不到 1% 的参数开销。

推理策略

TokenButler 将 KV 缓存划分为三个缓冲区以优化内存访问：

汇缓冲区（Sink Buffer）： 前 $S$ 个 token（始终保留）。
局部窗口缓冲区（Local Window Buffer）： 最近 $N$ 个 token 的循环缓冲区（始终保留）。
重要缓冲区（Important Buffer）： 由 TokenButler 根据预测分数动态填充。

优化措施：

批处理： 键的投影以 $N$ 个 token 的批次进行，而非在每个解码步执行，从而利用时间局部性。
预测间隔与邻居获取： 为了分摊预测器的成本，系统并非在每个步长运行预测器，而是每隔 $N$ 步运行一次（预测间隔）。为了降低在间隔之间重要性发生变化的 token 被遗漏的风险，系统采用“邻居获取”机制，将选中的 token 扩展为其空间邻居，确保覆盖多 token 实体或推理链。

主要贡献

轻量级预测器： 一种用于 token 重要性的学习指标，通过注意力蒸馏训练，以最小的参数开销（<1%）提供细粒度估计。
合成基准： 提出了一种新颖的共指检索任务，证明现有的淘汰和基于页的方法即使在短上下文长度（300 个 token）下，也会因永久 token 丢失或页缺失而失效。TokenButler 在此任务上实现了接近 oracle 的准确率。
长上下文性能： 在 RULER 和 LongBench 基准测试（64K 上下文）中，与最先进的稀疏化和压缩方法相比，TokenButler 取得了具有竞争力或更优的结果。
效率提升：
- 利用预测间隔，相比稠密注意力（Dense Attention），GPU 上速度提升高达 1.6 倍。
- 相比带有 CPU 卸载的稠密注意力，延迟降低高达 7.6 倍，使得在显存受限的系统上实现高效解码成为可能。
- 即使将预测器调用次数减少 16 倍，准确率仍保持在每步基线的 ≈1.1% 以内。

实验结果

合成共指任务： TokenButler 在 Llama-3.1-8B 和 Qwen2.5-7B 上实现了约 80-95% 的准确率和覆盖率，显著优于基于淘汰的方法（H2O、SnapKV，其准确率降至 50% 以下）和基于页的方法（Quest，其遭受部分页缺失的影响）。
RULER 基准（64K 上下文）： TokenButler 在 Llama-3.1-8B 和 Qwen2.5-7B 上匹配或超过了 SnapKV、xKV 和 Quest 等基线的性能，并在 Qwen2.5-7B 上取得了最佳平均分。
LongBench： TokenButler 在各种数据集上展示了优于以往 KV 稀疏化和压缩方法的准确率。
推理模型（AIME24）： 在 DeepSeek-R1-Distill-Llama-8B 上，TokenButler 实现了 30.0% 的准确率，接近稠密解码基线（33.3%），并在相同的 8K token 预算下优于其他稀疏化方法。
延迟： 在 CPU 卸载场景（256K+ 上下文）中，TokenButler 将延迟从每个 token 约 3.2 秒（稠密）降低至约 0.6 秒（TokenButler）。

意义

该论文指出，token 重要性具有高度动态性和查询依赖性，使得静态启发式和粗粒度检索不足以应对复杂、长上下文的任务。TokenButler 证明，一种学习的、细粒度的选择机制可以有效替代这些启发式方法。通过保留完整 KV 缓存同时仅选择性访问最相关的 token，TokenButler 缓解了内存带宽瓶颈，且避免了永久淘汰带来的质量下降。所提出的预测间隔和邻居获取策略进一步表明，此类学习预测器的计算成本可以被有效分摊，使其成为在有限硬件上进行实时长上下文推理的实用解决方案。