💬 NLP

DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention

DashAttention 引入了一种基于 $\alpha$ -entmax 的完全可微且自适应的稀疏分层注意力机制，用于动态选择可变数量的 KV 块，从而在长上下文建模精度和推理速度方面超越了 NSA 和 InfLLMv2 等现有方法。

原作者： Yuxiang Huang, Nuno M. T. Gonçalves, Federico Alvetreti, Lei Li, Xu Han, Edoardo M. Ponti, André F. T. Martins, Marcos V. Treviso

发布于 2026-05-19

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Yuxiang Huang, Nuno M. T. Gonçalves, Federico Alvetreti, Lei Li, Xu Han, Edoardo M. Ponti, André F. T. Martins, Marcos V. Treviso

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

想象一下，你正试图阅读一部多达 10 万页的百科全书，只为回答一个问题。

问题：“全量”与“前 K 项”的两难困境
当前的 AI 模型（大型语言模型）通常通过两种方式处理这种情况，但两者都存在缺陷：

“全量阅读”方法（全注意力机制）： 模型试图阅读百科全书中的每一个单词以寻找答案。这种方法准确，但极其缓慢且昂贵，就像为了找一道食谱而试图读完整本书。
“挑选前 5 项”方法（Top-K 稀疏注意力）： 模型快速扫描目录，挑选它认为最相关的前 5 章，并忽略其余部分。这很快，但过于僵化。如果答案实际上在第 6 章怎么办？或者如果答案需要阅读 20 个分散的页面怎么办？此外，一旦模型选定了这 5 章，它就无法从被忽略的部分中“学习”，使得训练过程变得笨拙。

解决方案：DashAttention
本文的作者提出了一种名为DashAttention的新方法。把它想象成一位聪明、自适应的图书管理员，他不仅不固定挑选书籍的数量，而是根据问题的复杂程度来决定需要抽取多少本书。

以下是 DashAttention 的工作原理，通过一个简单的类比分为三个阶段：

阶段 0：“章节摘要”（局部块摘要）

模型并非立即查看每一个单词，而是首先将庞大的文本分解为小的“块”（就像章节一样）。

旧方法： 过去它只是取章节中所有单词的平均值（就像说“这一章主要关于猫”）。
DashAttention 方法： 它使用一个微小的、经过学习的“读者”来扫描章节并撰写一个智能、细致的摘要。这就像一位人类图书管理员阅读一章后，写下一个能捕捉精髓而非仅仅是平均值的 2 句话摘要。关键在于，这个摘要是灵活的；如果模型开始训练，它会随着时间的推移学会撰写更好的摘要。

阶段 1：“自适应守门人”（Entmax 路由）

现在，模型拥有一份章节摘要列表。它需要决定哪些章节需要详细阅读。

旧方法（Top-K）： 模型有一条严格规则：“总是恰好挑选 5 章。”如果问题很简单，它浪费时间去阅读 5 章；如果问题很难，它因为上限为 5 而错过重要信息。
DashAttention 方法： 模型使用一种特殊的数学工具，称为 $\alpha$ -entmax。想象一位守门人，他审视问题和摘要。
- 如果问题很简单（“法国的首都是什么？”），守门人说：“只需要 1 章”，并将其余部分锁在外面。
- 如果问题很复杂（“追踪跨越三大洲的贸易路线历史”），守门人说：“好的，我们需要 15 章”，并打开大门。
- 魔力所在： 这位守门人是“可微分的”。这意味着模型可以学习如何成为一名更好的守门人。如果在训练期间它挑选了错误的章节，它会收到信号来调整其守门策略。这不是一个生硬的“是/否”开关，而是一个平滑的、可学习的旋钮。

阶段 2：“深入挖掘”（先验诱导的稀疏 Softmax）

最后，模型阅读守门人选定的特定章节。

它利用守门人（阶段 1）投出的“票”来引导对选定文本的详细阅读。
它确保即使跳过了书中的大部分内容，也不会丢失故事的连贯性。它填补了空白，使最终答案的准确性与阅读整本书一样，但速度快得多。

为什么这更好？（结果）

本文声称 DashAttention 在三个关键领域胜出：

更智能的选择： 与僵化的“前 5 项”规则不同，DashAttention 具有适应性。它在困难问题上投入更多“脑力”，在简单问题上投入更少。这使其在寻找干草堆中的特定针（检索任务）方面表现更好。
无“分散”： 在长文本中，标准 AI 模型往往会变得“分心”，将注意力 spread 得太薄，就像光束太宽而无法看清任何东西。DashAttention 保持光束聚焦，确保模型即使在海量文本中也能保持敏锐。
速度： 因为它跳过了不相关部分的阅读，所以速度极快。
- 作者为计算机芯片（GPU）构建了一个专用版本，在处理非常长的文本时，其运行速度比当前的行业标准（FlashAttention-3）快 3.36 倍。
- 它达到了与阅读整本书相同的准确性，但仅使用了25% 的算力（75% 的稀疏度）。

总结

DashAttention 就像从一位僵化、受规则束缚、总是挑选 5 本书的图书管理员，升级为一位高度智能、自适应的助手。这位助手会阅读目录，针对具体问题决定确切需要多少章，然后仅深入挖掘那些章节。它比以前的方法更快、更聪明、学习得更好，使 AI 能够处理海量信息而不会感到不知所措或速度变慢。

技术摘要：DashAttention

问题陈述

当前分层注意力方法（如 NSA 和 InfLLMv2）通过基于粗略注意力分数选择前 $k$ 个相关键值（KV）块，随后对选中的令牌执行细粒度 softmax 注意力，从而解决长上下文建模的计算瓶颈。然而，这些方法存在两个主要局限性：

固定预算与缺乏灵活性：前 $k$ 个（top- $k$ ）操作假设任何查询的相关令牌数量是固定的，无法适应不同令牌、注意力头和层之间异质的信息需求。
不可微性：硬性的 top- $k$ 选择切断了粗略路由阶段与细粒度令牌级注意力之间的梯度流。这阻碍了模型学习如何调整路由分数以改变块成员资格，导致稀疏选择与最终损失之间出现脱节。
分散性：近期的理论研究表明，标准 softmax 注意力在长上下文中可能是“分散的”，即注意力分布的香农熵随序列长度对数增长，使得长程依赖建模变得困难。虽然 top- $k$ 缓解了这一问题，但现有的在选前通过 softmax 聚合头的分层方法可能会重新引入分散性。

方法：DashAttention

DashAttention（可微且自适应稀疏分层注意力）是一种旨在完全可微且自适应稀疏的多阶段机制。它用 $\alpha$ -entmax 变换替代了硬性的 top- $k$ 路由，使模型能够根据输入动态学习相关块的数量。

该架构包含三个阶段：

阶段 0：局部块摘要

上下文被划分为大小为 $B$ 的连续块。DashAttention 不使用简单的平均池化或可学习的 MLP（这些方法难以从预训练模型中适应），而是采用可学习的局部注意力机制。一个可学习的查询向量 $\bar{q}$ 通过局部缩放点积注意力（SDPA）关注每个块内的键。

初始化： $\bar{q}$ 初始化为零，导致该机制在训练开始时表现为均匀平均池化。
适应：随着训练进行， $\bar{q}$ 学习对块内的令牌进行加权，生成更具表现力的摘要。
效率：一旦块完全生成，其摘要即被固定并缓存，确保推理期间无需重新计算。

阶段 1：Entmax 块路由

对于每个查询，模型计算与块摘要的 logits。这些 logits 不使用 top- $k$ 截断，而是使用 $\alpha$ -entmax 进行变换：
$\hat{w} = \alpha\text{-entmax}(\gamma \bar{z})$

自适应稀疏性： $\alpha$ -entmax 生成一个稀疏概率向量，其中非零元素（支持集）的数量由输入分数本身决定，而非固定的超参数。
可微性：与 top- $k$ 不同， $\alpha$ -entmax 是完全可微的。这使得梯度能够反向传播通过路由决策到达块摘要阶段（阶段 0）。
GQA 处理：在分组查询注意力（GQA）设置中，共享 KV 头的查询头的路由概率被平均，以生成共享的稀疏支持集。

阶段 2：先验诱导的稀疏 Softmax

最后阶段在令牌级别细化注意力，但仅限于阶段 1 选中的块内。为了保持可微性并处理从粗略路由到细粒度的过渡，该方法引入了先验诱导的 softmax：

变分公式：标准 softmax 被视为最大化 $z^\top p - KL(p \| u)$ 。DashAttention 将参考分布 $u$ 修改为从阶段 1 路由分数 $w$ 导出的分布 $g$ 。
动态掩码：通过构建 $g$ ，使得当路由分数 $w_{i,j} = 0$ 时 $g_{i,j} = 0$ ，生成的 softmax 注意力自然变得稀疏。零概率块被有效掩码。
对角线处理：对于靠近对角线（可能没有摘要）的块，维护一个单独的分支。超参数 $\sigma$ 控制先验的强度，使模型能够在依赖路由分数和均匀基线之间取得平衡。
实现：该阶段使用修改后的 FlashAttention 内核实现，该内核迭代由位打包掩码定义的活跃块，在 softmax 之前将路由偏差添加到 logits 中。

主要贡献

可微分层注意力：DashAttention 引入了一种端到端可微的稀疏注意力机制。通过使用 $\alpha$ -entmax 进行粗略路由，它消除了 top- $k$ 方法固有的梯度阻塞，使模型能够学习最优的稀疏模式。
自适应稀疏分配：与固定预算方法不同，DashAttention 动态地在令牌、头和层之间分配稀疏性。信息丰富的查询可以关注更多块，而尖锐的查询则关注更少块，从而优化每个实例的计算预算。
非分散性：理论分析表明，DashAttention 使用 $\alpha$ -entmax 进行头聚合，防止了基于 softmax 的分层注意力中常见的分散现象，理论上提高了长上下文建模能力。
高效的 GPU 实现：作者提供了基于 Triton 的实现，包含三个融合内核（每阶段一个）。该设计与现有的高效注意力库（如 FlashAttention）兼容，并支持预填充和增量解码。

实验结果

作者在 MiniCPM-4 模型（1B、3B 和 8B 参数）上评估了 DashAttention，上下文长度为 16K，并与全注意力、NSA 和 InfLLMv2 进行了比较。

任务性能：在 RULER 和 HELMET 基准测试中，DashAttention 始终优于 NSA 和 InfLLMv2，特别是在具有挑战性的长上下文检索任务（如 MK1–MK3）中。在 75% 稀疏度下，其准确率与全注意力相当。
泛化能力：在短上下文通用任务（MMLU、GSM8K 等）上，DashAttention 的表现与全注意力模型相当，未显示出通用能力的显著下降。
效率：
- 预填充：DashAttention 相比 FlashAttention-3 实现了 1.34 倍至 3.09 倍 的加速。
- 解码：在内存受限模式下，相比 FlashAttention-3 实现了高达 3.36 倍 的加速，相比 InfLLMv2 实现了 1.35 倍 的加速。
帕累托前沿：DashAttention 主导了准确率 - 稀疏度权衡曲线，在高稀疏度水平下（例如 ~90% 稀疏度）保持了比基线更高的准确率。
动态行为：消融研究证实，DashAttention 自然地将更多稀疏性分配给中间层，较少分配给早期层，模仿了最优预算分配策略，而无需显式调整。

意义与主张

该论文声称，DashAttention 通过统一稀疏注意力（效率）和稠密注意力（灵活性和可微性）的优势，为长上下文建模提供了一种具有成本效益的策略。

理论优势：它通过确保注意力机制通过 $\alpha$ -entmax 聚合保持非分散性，解决了长上下文建模中的“分散”问题。
实践优势：它通过提供一种完全可微、可训练的方法，弥合了理论稀疏注意力与实际部署之间的差距，该方法可以从预训练的 softmax 模型中适应，而无需进行架构 overhaul。
性能：该方法实现了更优的帕累托前沿，与现有的分层方法相比，在高稀疏度水平下提供了更高的准确率，同时也带来了显著的推理加速。

作者指出了一个局限性：内核尚未集成到 vLLM 或 SGLang 等主要服务框架中，这被确定为未来的工作。他们还建议将其扩展到混合模型架构。