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

在大型语言模型中，对建模长上下文的需求持续增长，但标准自注意力机制的二次复杂度常常成为瓶颈。尽管现有的稀疏注意力机制提高了效率，但它们可能仍然会遇到静态模式或信息丢失等问题。我们引入了一种可训练的动态掩码稀疏注意力机制——动态掩码注意力（Dynamic Mask Attention，DMA），它有效地利用了内容感知和位置感知稀疏性。DMA 通过两项关键创新实现这一目标：首先，它从值表示中动态生成内容感知稀疏掩码，使模型能够自适应地识别并聚焦于关键信息。其次，它实现了位置感知稀疏注意力计算，有效地跳过不必要的计算区域。这种双重稀疏性设计使得模型能够显著降低重要信息的计算复杂度，同时保留完整信息，在信息保真度和计算效率之间取得了出色的平衡。我们通过全面的实验验证了 DMA 的性能。比较研究表明，在 Chinchilla 缩放定律设置下，DMA 在困惑度方面优于多头注意力、滑动窗口注意力、多头潜在注意力和原生稀疏注意力。此外，在具有挑战性的多查询关联召回任务中，与这些方法相比，DMA 也展现出卓越的性能和效率。至关重要的是，在对一个 1.7B 参数模型的评估中，DMA 在标准基准性能和具有挑战性的“大海捞针”任务中均显著优于多头注意力。这些实验结果突出了其有效平衡模型效率和长上下文建模能力的能力。