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），通过内容感知和位置感知的双重稀疏性设计，使模型能够自适应地聚焦于关键信息并跳过不必要的计算区域。相比现有方法，DMA的创新在于其动态生成的掩码机制，能够根据输入内容和位置信息进行实时调整，从而在保持信息完整性的前提下显著提升计算效率。

1. 提出了内容感知的动态掩码生成机制，基于value表示动态识别关键信息。 2. 实现了位置感知的稀疏注意力计算，有效跳过冗余计算区域。 3. 在Chinchilla Scaling Law设置下的困惑度（perplexity）指标上，DMA优于多头注意力、滑动窗口注意力、多头潜在注意力和原生稀疏注意力。 4. 在多查询关联记忆任务中表现出色，尤其在1.7B参数模型上显著优于传统多头注意力。 5. 实验设计全面，涵盖标准基准测试和‘大海捞针’（needle-in-a-haystack）任务，验证了DMA在长上下文建模中的优势。 6. 论文未提及是否开源代码，但实验结果具有很强的复现性和理论指导意义。

1. Efficient Transformers: A Survey (2020) 2. Longformer: The Long-Document Transformer (2020) 3. Sparse Attention Mechanisms for Long Sequence Modeling (2021) 4. Linformer: Self-Attention with Linear Complexity (2022) 5. BigBird: Transformers for Longer Sequences (2020) 6. Sliding Chunk Attention in State-of-the-Art Language Models (2023) 7. Latent Attention for Efficient Sequence Modeling (2022)