Trainable Dynamic Mask Sparse Attention

在大型语言模型中，对建模长上下文的需求持续增长，但标准自注意力机制的二次复杂度常常成为瓶颈。尽管现有的稀疏注意力机制在效率上有所提升，但仍可能面临模式固定或信息丢失等问题。我们提出了一种可训练的动态掩码稀疏注意力机制——Dynamic Mask Attention（DMA），它有效地结合了内容感知和位置感知的稀疏性。 DMA通过两项关键技术实现这一目标：首先，它从值表示中动态生成内容感知的稀疏掩码，使模型能够自适应地识别并聚焦于关键信息；其次，它实现了位置感知的稀疏注意力计算，有效地跳过不必要的计算区域。这种双稀疏性设计使模型在显著降低关键信息计算复杂度的同时，保留了完整的信息，从而在信息保真度与计算效率之间实现了出色的平衡。 我们通过全面的实验验证了DMA的性能。对比研究表明，在符合Chinchilla缩放定律设置的情况下，DMA在困惑度（perplexity）指标上优于多头注意力、滑动窗口注意力、多头隐式注意力以及原始稀疏注意力等方法。此外，在具有挑战性的多跳关联记忆任务中，DMA也展现出优于这些方法的表现与效率。尤其关键的是，在评估17亿参数模型时，DMA在标准基准测试性能以及“大海捞针”（needle-in-a-haystack）任务中均显著优于多头注意力机制。这些实验结果突出了DMA在有效平衡模型效率与长上下文建模能力方面的卓越能力。

论文旨在解决大规模语言模型中自注意力机制在建模长上下文时所面临的计算瓶颈问题。由于标准自注意力机制的计算复杂度为二次复杂度，难以高效处理长序列。虽然已有稀疏注意力机制尝试提升效率，但它们往往存在固定的稀疏模式或信息丢失的问题。这是一个当前自然语言处理和模型架构设计领域非常热门且具有挑战性的问题。

论文提出了一种可训练的动态稀疏注意力机制——Dynamic Mask Attention (DMA)，结合了内容感知和位置感知的稀疏性。其核心创新在于：1）通过值向量动态生成内容感知的稀疏掩码，使模型能够自适应地聚焦于关键信息；2）通过位置感知机制跳过不必要的计算区域，进一步提升效率。这种双稀疏设计既保留了信息完整性，又显著降低了计算复杂度，是现有稀疏注意力机制的一个重要改进方向。

1. DMA在Chinchilla Scaling Law设置下，相比多头注意力、滑动窗口注意力、多头潜在注意力和原生稀疏注意力等方法，表现出更低的困惑度（perplexity） 2. 在多查询关联记忆任务中，DMA展现出更高的性能与效率 3. 在1.7B参数模型的评估中，DMA在标准基准测试和‘大海捞针’任务中均显著优于多头注意力 4. 实验设计全面，覆盖语言建模和长序列记忆能力等多个维度 5. 未提及是否开源代码，值得后续关注

1. Longformer: The Long-Document Transformer 2. Sparse Attention Mechanisms for Long Sequence Modeling 3. Linformer: Self-Attention with Linear Complexity 4. BigBird: Transformers for Longer Sequences 5. Sliding Window Attention in Efficient Transformers 6. Routing Networks for Dynamic Sparse Attention