提出可微分的自适应稀疏注意力,高效建模长文本。
DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention

- 用自适应稀疏的α-entmax选择动态数量的键值块,实现全程可微。
- 在75%稀疏度下达到全注意力精度,高稀疏时性能优于NSA和InfLLMv2。
- 适合需要长上下文建模的大型语言模型,推理速度超FlashAttention-3。
现有层级注意力方法如NSA和InfLLMv2基于粗粒度注意力得分选取前k个键值块,再对选中标记应用细粒度softmax注意力。但top-k操作假设每个查询的相关标记数固定,且阻断了稀疏与稠密阶段间的梯度流。本文提出DashAttention(可微分自适应稀疏层级注意力),在第一阶段利用自适应稀疏的α-entmax变换,根据当前查询动态选择块数;该过程为第二阶段softmax注意力提供先验,使整个层级结构完全可微。与其它方法不同,我们证明DashAttention是非发散的,具备更强的长序列建模能力。大规模语言模型实验表明,其在75%稀疏度下实现与全注意力相当的精度,并在高稀疏场景下超越NSA和InfLLMv2的帕累托前沿。我们还提供了GPU感知的Triton高效实现,在推理时速度提升超过FlashAttention-3。总体而言,DashAttention是一种成本效益高的长上下文建模策略。
原文摘要 · Abstract (English)
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 $α$-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 attention methods, we show that DashAttention is non-dispersive, translating to better long-context modeling ability. Experiments with large language models (LLMs) show that DashAttention achieves comparable accuracy as full attention with 75% sparsity and a better Pareto frontier than NSA and InfLLMv2, especially in high-sparsity regimes. We also provide an efficient, GPU-aware implementation of DashAttention in Triton, which achieves a speedup of up to over FlashAttention-3 at inference time. Overall, DashAttention offers a cost-effective strategy to model long contexts.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。