[论文] DashAttention: Differentiable and Adaptive Sparse Hierarchical Attenti...

论文概要

研究领域: NLP 作者: Yuxiang Huang, Nuno M. T. Gonçalves, Federico Alvetreti 发布时间: 2026-05-19 arXiv: 2505.14310

中文摘要

当前的分层注意力方法（如NSA和InfLLMv2）基于粗粒度注意力分数选择top-k相关的键值（KV）块，然后对选中的token应用细粒度softmax注意力。然而top-k操作假设每个查询的相关token数量是固定的，且阻碍了稀疏与稠密阶段之间的梯度流动。本文提出DashAttention（可微自适应稀疏分层注意力），利用自适应稀疏的\(\alpha\)-entmax变换在第一阶段根据当前查询选择可变数量的块，为第二阶段的softmax注意力提供先验，使整个层次结构保持完全可微。与其他分层注意力方法不同，我们证明了DashAttention具有非分散性，从而带来更好的长上下文建模能力。在大型语言模型上的实验表明，DashAttention在75%稀疏度下达到与完整注意力相当的精度，并在高稀疏度范围内比NSA和InfLLMv2具有更好的帕累托前沿。我们还提供了基于Triton的高效GPU感知实现，推理时速度比FlashAttention-3更快。总体而言，DashAttention为长上下文建模提供了一种经济高效的策略。

原文摘要

Current hierarchical attention methods, such as NSA and InfLLMv2, select the top-k relevant key-value (KV) blocks based on coarse attention scores and subsequently apply fine-grained softmax attention on the selected tokens. However, the top-k operation assumes the number of relevant tokens for any query is fixed and it precludes the gradient flow between the sparse and dense stages. In this work, we propose DashAttention (Differentiable and Adaptive Sparse Hierarchical Attention), which leverages the adaptively sparse \(\alpha\)-entmax transformation to select a variable number of blocks according to the current query in the first stage. This in turn provides a prior for the second-stage softmax attention, keeping the entire hierarchy fully differentiable. Contrary to other hierarchical att...

自动采集于 2026-05-20

#论文 #arXiv #NLP #小凯