arXiv:2508.02124cs.AIcs.CL2025-08被引 5

提出可训练的动态掩码稀疏注意力,高效处理长文本且支持端到端训练。

Trainable Dynamic Mask Sparse Attention

  • 用值向量生成内容感知的动态掩码,自适应聚焦关键信息。
  • 在硬件友好的方式下计算位置感知稀疏权重,实现高达10倍加速。
  • 模型可端到端训练,适合需要长上下文建模的场景。

大语言模型对长上下文建模的需求日益增长,但标准自注意力机制的二次复杂度成为瓶颈。社区虽已提出稀疏注意力缓解此问题,但位置感知方法依赖静态稀疏结构,缺乏对查询上下文的适应性;内容感知方法依赖启发式键值选择,难以实现完全可微。本文提出可训练的动态掩码稀疏注意力(Dynamic Mask Attention, DMA),融合两类方法优势。其三大创新为:首先,利用值向量表示生成内容感知动态掩码,使模型自适应识别并关注关键信息;其次,以硬件友好方式计算位置感知稀疏权重,高效跳过冗余计算区域;最后,证明动态掩码与稀疏权重不阻碍梯度传播,支持端到端训练。实验验证表明,DMA在缩放定律、多查询关联召回、标准基准测试及针堆中找针任务中均显著优于现有最优稀疏注意力基线,同时实现最高达10倍的整体加速。结果表明其能有效平衡模型效率与长上下文建模能力。相关计算核代码已开源(https://github.com/SmallDoges/flash-dmattn),欢迎社区进一步研究与应用。

原文摘要 · Abstract (English)

The increasing demand for long-context modeling in large language models (LLMs) is bottlenecked by the quadratic complexity of the standard self-attention mechanism. The community has proposed sparse attention to mitigate this issue. However, position-aware sparse attention methods rely on static sparse structures that lack adaptability to diverse query contexts, while content-aware sparse attention methods depend on heuristic key-value selection, hindering full differentiability. We introduce a trainable dynamic mask sparse attention mechanism, a method that merges the advantages of both position-aware and content-aware approaches. Dynamic Mask Attention (DMA) achieves this through three key innovations: First, it leverages value vector representations to generate content-aware dynamic masks, enabling the model to adaptively identify and attend to critical information. Second, it computes position-aware sparse weights in a hardware-friendly manner, efficiently skipping unnecessary computational regions. Finally, we demonstrate that the introduced dynamic mask and sparse weights do not obstruct gradients, supporting end-to-end training. We have validated the performance of DMA through comprehensive experiments. A large body of experimental evidence shows that DMA consistently holds a Pareto advantage over state-of-the-art sparse attention baselines in tasks including scaling laws, multi-query associative recall, standard benchmarks, and needle in a haystack tests, while also delivering up to a 10x overall speedup. These results highlight its ability to effectively balance model efficiency with long-context modeling capabilities. Our computational kernel code is now open-source at https://github.com/SmallDoges/flash-dmattn to encourage further research and application by the community.

稀疏注意力长序列建模可训练掩码效率优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。