可训练动态掩码稀疏注意力

大语言模型（Large Language Models, LLMs）中对长上下文建模的需求不断增长，标准自注意力机制的二次复杂度已成为瓶颈。社区提出稀疏注意力以缓解此问题。然而，基于位置的稀疏注意力方法依赖静态稀疏结构，缺乏对不同查询上下文的适应性；而基于内容的稀疏注意力方法依赖启发式的键值选择，阻碍了完全可微分。我们引入一种可训练的动态掩码稀疏注意力机制，融合了位置感知和内容感知方法的优势。动态掩码注意力（Dynamic Mask Attention, DMA）通过三个关键创新实现：其一，利用值向量表示生成内容感知的动态掩码，使模型能够自适应识别并关注关键信息；其二，以硬件友好的方式计算位置感知稀疏权重，高效跳过不必要的计算区域；最后，我们展示引入的动态掩码和稀疏权重并不阻碍梯度，支持端到端训练。我们通过全面实验验证了 DMA 的性能。大量实验结果表明，DMA 在扩展规律、多查询关联记忆、标准基准测试和针 needle in a haystack 测试等任务中始终保持对最先进稀疏注意力基线的 Pareto 优势，同时实现了最高可达 10 倍的整体加速。这些结果凸显了其在平衡模型效率与长上下文建模能力方面的有效能力。我们的计算内核代码已开源发布于 https://github.com/SmallDoges/flash-dmattn ，旨在鼓励社区进一步的研究与应用。