💬 NLP

EntmaxKV: Support-Aware Decoding for Entmax Attention

EntmaxKV 是一种支持感知的稀疏解码框架，它利用 $\alpha$ -entmax 注意力的精确稀疏性，在推理前选择性地加载键值（KV）缓存页面，从而显著减少内存流量，在长上下文生成中实现大幅加速，同时保持与全缓存基线相当的准确性。

原作者： Gonçalo Duarte, Miguel Couceiro, Marcos V. Treviso

发布于 2026-05-22

📖 1 分钟阅读☕ 轻松阅读

原作者： Gonçalo Duarte, Miguel Couceiro, Marcos V. Treviso

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

以下是EntmaxKV论文的通俗解释，辅以生动的类比。

核心难题：“无限图书馆”的瓶颈

想象你是一位图书管理员（即 AI 模型），正在撰写一个故事。每写一个新句子，你都需要回顾之前写下的所有内容，以确保新句子合乎逻辑。

在标准 AI 中，这种“记忆”就像一座不断增长的图书馆。每当你写下一个新词，图书馆就会多出一本书。

问题所在：当故事变得非常长（达到数百万字）时，图书馆会变得极其庞大。
瓶颈：为了写出下一个词，图书管理员必须跑向书架，取出图书馆里的每一本书，阅读书脊，并判断哪些是相关的。即使 99% 的书都无关紧要，管理员仍必须物理地搬运它们。这耗费了巨大的时间和能量，拖慢了一切进程。

旧方案："Softmax"（“人人有票”法）

当前的 AI 模型使用一种称为Softmax的方法。

工作原理：当图书管理员查看图书馆时，Softmax 会给每一本书分配一张微小的、非零的“票”（概率分数）。即使是一本关于“如何烤面包”的书，如果故事讲的是“太空旅行”，它也会获得一张微小的票。
缺陷：因为每一本书都有票，管理员无法直接忽略那些无关的书。他们必须将所有书都加载到房间里进行检查。如果试图为了节省时间而跳过无关的书，就会意外丢弃 Softmax 分配的那些微小票数，从而破坏数学计算并毁掉故事。
结果：这就像试图在一堆干草中找一根针，即使你知道针只在一个极小的角落，你也必须查看每一根干草。

新构想："Entmax"（“精确零”法）

作者引入了一种新的数学工具，称为 $\alpha$ -entmax。

魔法技巧：与 Softmax 不同，Entmax 非常严格。如果一本书不相关，它获得的票数就是精确的零。这不是“微小”，而是什么都没有。
优势：如果一本书的票数为零，它对故事就毫无贡献。你可以直接将其扔掉，而完全不会改变结果。
目标：目标不再是试图近似处理整个“干草堆”，而是找到特定的“针”（即支持集）。如果你能找出那些拥有非零票数的少数几本书，就不需要查看图书馆的其余部分。

解决方案：EntmaxKV（“聪明图书管理员”）

论文提出了EntmaxKV，这是一个利用“精确零”特性来加速的系统。其工作原理如下：

1. “盒子”检查（查询感知页面评分）

想象图书馆的书不是散放的，而是存放在盒子（页面）中。

在图书管理员打开盒子阅读内部书籍之前，他们会先查看盒子上的标签。
标签包含了对盒内书籍的“摘要”（最小和最大分数）。
管理员会问：“这个盒子里有可能包含相关的书吗？”
如果答案是“不”（该盒子肯定不相关），管理员就永远不会打开这个盒子。他们节省了走向书架并将其取出的时间。

2. “高斯”猜测（高斯感知选择器）

有时，仅凭盒子标签还不足以 100% 确定。作者增加了一个巧妙的猜测游戏。

他们查看盒内书籍的平均值和分布范围。
他们使用统计猜测（就像天气预报）来估算该盒子中任何一本书可能获得的最高分数。
如果“天气预报”显示该盒子中最好的书仍然无聊到无关紧要，他们就会跳过这个盒子。这使他们能够更激进地跳过无关的盒子，而不会错过好书。

3. “精确”搜索（支持集恢复）

一旦管理员只选中了有希望的盒子，他们就会打开它们并运行Entmax数学计算。

因为 Entmax 会给不相关的项目分配零，数学计算会自然地忽略选中盒子内的垃圾。
结果：如果管理员选对了盒子，故事就是100% 完美的，与阅读了整个图书馆完全一样。他们只是没有浪费时间处理垃圾。

为何重要（实验结果）

论文将这种方法与旧的"Softmax"方法进行了测试，发现：

更少错误：当你尝试使用旧的 Softmax 方法跳过书籍时，你不可避免地会丢弃一些重要的“微小票数”，从而导致错误。只要 EntmaxKV 找到了正确的盒子，它丢弃的重要信息为零。
速度：在非常长的故事（100 万字）中，EntmaxKV 比标准方法快3.36 倍，比未使用此跳过技巧的标准 Entmax 方法快5.43 倍。
准确性：它在仅使用极少内存流量的情况下，保持了高质量的故事（低“困惑度”）。

类比总结

旧方法（Softmax）：你有百万封邮件。你必须阅读每一封邮件的主题行，以决定回复哪些，因为即使是垃圾邮件也有微小的可能性是重要的。
EntmaxKV：你有一个智能过滤器。它首先查看发件人和主题行元数据。它立即识别出 99% 的邮件肯定是垃圾邮件（零概率）。它在未打开的情况下直接删除它们。它只打开那 1% 可能重要的邮件。因为过滤器是完美的，你永远不会错过真正的邮件，但却节省了数小时的时间。

论文的主要主张：通过切换到一种为不相关数据生成“精确零”的数学系统，并在加载数据之前先检查元数据，我们可以使 AI 在处理长任务时速度大幅提升，同时不损失准确性。

技术摘要：EntmaxKV——面向 Entmax 注意力的支持感知解码

1. 问题陈述

大语言模型（LLM）中的长上下文解码日益受到内存流量而非计算能力的瓶颈限制。在自回归生成过程中，每个新 token 都必须关注一个随上下文长度线性增长的键值（KV）缓存。在每一步将该缓存从高带宽内存（HBM）读取到片上 SRAM 的操作主导了延迟，通常占每个 token 生成时间的一半以上。

现有的稀疏解码方法试图通过选择 token、块或页的子集来缓解这一问题。然而，这些方法主要是为softmax注意力设计的。softmax 的一个根本局限性在于其“密集尾部”：无论相关性如何，它都会为每个 token 分配严格为正的概率。因此，在 softmax 下对 KV 缓存的任何截断本质上都会丢弃非零概率质量，使得稀疏解码成为一个近似问题，除非读取完整缓存，否则误差不可避免。

2. 方法：EntmaxKV

作者提出了EntmaxKV，这是一种原生适用于 $\alpha$ -entmax注意力的稀疏解码框架。与 softmax 不同， $\alpha$ -entmax（当 $\alpha > 1$ 时）产生的概率分布具有精确的零值。这一特性将稀疏解码的目标从“近似密集尾部”转变为**“支持恢复”**：如果所选 token 子集包含了真实支持（即具有非零概率的 token 集合），则稀疏解码输出在数学上与完整缓存输出完全相同。

EntmaxKV 通过包含三个核心组件的分页解码流水线来实现这一目标：

A. 查询感知的页评分

KV 缓存被划分为页。在加载完整的键和值之前，系统利用每页存储的轻量级元数据对页进行评分：

边界框界限：坐标方向上的最小和最大键向量（ $k_{min}, k_{max}$ ）允许针对给定查询 $q$ ，对页中任意 token 的最大可能注意力分数确定一个确定性上界。
高斯统计：为了采用更自适应的方法，系统存储页内键的一阶和二阶矩（均值和方差）。这使得可以将页中 token 分数的分布近似为高斯随机变量。

B. 候选选择策略

提出了两种策略来选择加载哪些页：

Top-k 页选择：一种固定预算方法，选择基于边界框元数据得出的上界分数最高的 $k$ 页。
高斯感知 Entmax 选择：一种动态方法，利用页级高斯统计估计全局 $\alpha$ -entmax 阈值（ $\tau$ ）。它求解一个近似的归一化方程以找到 $\hat{\tau}$ ，然后选择估计最大分数超过该阈值（经安全裕度 $\Delta$ 调整）的页。这种方法根据当前查询的特定分数分布调整 token 预算。

C. 稀疏 Entmax 注意力

一旦选定候选页，系统仅加载相应的 KV 条目。它在该缩减集合上计算 $\alpha$ -entmax 注意力权重。如果所选集合包含了 $\alpha$ -entmax 分布的完整支持，则输出是精确的。

3. 理论分析

该论文提供了关于截断引入误差的理论框架：

误差界：输出误差 $\|o - \tilde{o}\|$ 被限制在 $2B\delta$ 以内，其中 $B$ 是值向量范数的上界， $\delta$ 是被丢弃 token 的总概率质量。
支持恢复：对于 softmax，只要丢弃任何 token， $\delta > 0$ ，从而保证存在非零误差。对于 $\alpha$ -entmax，如果所选集合 $I_{keep}$ 包含支持集 $S$ （即 $S \subseteq I_{keep}$ ），则 $\delta = 0$ ，误差完全消失。
比较：论文证明，在相同的保留集合下， $\alpha$ -entmax 丢弃的概率质量远少于 softmax，特别是在候选集成功捕获 entmax 支持时。

4. 关键结果

实验使用 10 亿参数模型（NAPE）在长上下文基准（RULER、Passkey Retrieval、PG19）和效率基准上进行。

近似质量：在匹配的 KV 预算下，EntmaxKV 始终比基于 softmax 的稀疏解码丢弃更少的概率质量并保留更多的支持 token。Entmax 的相对输出误差显著更低，随着支持的恢复趋近于零。
下游性能：
- Passkey Retrieval：EntmaxKV 在长上下文（高达 10 万 token）中保持完美的检索准确率，仅使用不到 0.1% 的 KV 缓存，而基于 softmax 的方法迅速退化。
- RULER 基准：EntmaxKV 显著优于 softmax 稀疏解码，特别是在 16k 上下文长度下，保持了接近完整缓存基线的性能。
- 语言建模（PG19）：使用 1024-token 预算的 Entmax 稀疏解码，其困惑度低于使用更大预算（例如 8192 个 token）的 softmax 解码。
效率：
- EntmaxKV 将瓶颈从完整缓存的内存流量转移到了轻量级的页选择上。
- 在 100 万 token 的上下文长度下，高斯感知变体相比完整缓存 softmax（FlashDecoding）实现了3.36 倍的加速，相比完整缓存 entmax 解码实现了5.43 倍的加速。
- 与 Top-k 变体相比，高斯感知选择器在序列长度扩展方面表现更好，避免了显式的 Top-k 瓶颈。

5. 意义与主张

该论文声称，EntmaxKV 通过利用 $\alpha$ -entmax 的精确稀疏性，从根本上改变了稀疏解码的范式。

从近似到恢复：主要贡献是将稀疏解码重新定义为支持恢复问题。作者认为，如果选择机制成功识别出支持，该方法就不是近似计算，而是精确计算，从而消除了 softmax 截断固有的“密集尾部”误差。
实际效率：这项工作弥合了稀疏注意力的理论优势与自回归解码的实际约束之间的差距。通过在加载 KV 页之前利用稀疏性，EntmaxKV 减少了内存流量，同时在捕获支持时保持了注意力机制的精确性。
自适应选择：引入高斯感知选择器允许系统根据分数分布动态调整候选预算，为固定预算的 Top-k 选择提供了更稳健的替代方案。

作者将 EntmaxKV 定位为 $\alpha$ -entmax 模型的原生解决方案，区别于试图将稀疏性强加于密集 softmax 分布的先前工作。他们强调，虽然最近的核函数已使 $\alpha$ -entmax 能够大规模训练，但 EntmaxKV 是首个通过实现既高效又在支持恢复下理论上精确的稀疏推理，来解决自回归解码瓶颈的工作。