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

在大型语言模型中，对长上下文建模的需求不断增加，但标准自注意力的二次复杂度成为了显著瓶颈。虽然现有稀疏注意力机制提升了效率，但常存在静态模式和信息丢失等限制。本文提出了一种可训练的动态掩码稀疏注意力机制，通过三大关键创新解决这些挑战。首先，它利用值向量动态生成内容感知的稀疏掩码，使模型能够自适应识别和聚焦关键信息。其次，实施了位置感知的稀疏注意力计算，有效跳过不必要的计算区域。最后，确保引入的动态掩码和稀疏权重不阻碍梯度传递，从而支持端到端训练。这种双重稀疏设计使得模型在显著降低计算复杂度的同时保留完整信息，实现效率与性能的良好平衡。通过全面实验验证了动态掩码注意力的性能。对比研究表明，我们的方法在多个任务中持续实现帕累托优势，包括扩展规律、多查询关联回忆、通用基准和针尖找麦秸测试，最高加速可达10倍。这些结果凸显了其在模型效率与长上下文建模能力间的有效平衡。我们的计算内核已开源于https://github.com/SmallDoges/flash-dmattn，便于社区进一步研究和应用。