Trainable Dynamic Mask Sparse Attention

AK's Picks LLM PTSAD Long Context
在大型语言模型中,对建模长上下文的需求持续增长,但标准自注意力机制的二次复杂度常常成为瓶颈。尽管现有的稀疏注意力机制在效率上有所提升,但仍可能面临模式固定或信息丢失等问题。我们提出了一种可训练的动态掩码稀疏注意力机制——Dynamic Mask Attention(DMA),它有效地结合了内容感知和位置感知的稀疏性。 DMA通过两项关键技术实现这一目标:首先,它从值表示中动态生成内容感知的稀疏掩码,使模型能够自适应地识别并聚焦于关键信息;其次,它实现了位置感知的稀疏注意力计算,有效地跳过不必要的计算区域。这种双稀疏性设计使模型在显著降低关键信息计算复杂度的同时,保留了完整的信息,从而在信息保真度与计算效率之间实现了出色的平衡。 我们通过全面的实验验证了DMA的性能。对比研究表明,在符合Chinchilla缩放定律设置的情况下,DMA在困惑度(perplexity)指标上优于多头注意力、滑动窗口注意力、多头隐式注意力以及原始稀疏注意力等方法。此外,在具有挑战性的多跳关联记忆任务中,DMA也展现出优于这些方法的表现与效率。尤其关键的是,在评估17亿参数模型时,DMA在标准基准测试性能以及“大海捞针”(needle-in-a-haystack)任务中均显著优于多头注意力机制。这些实验结果突出了DMA在有效平衡模型效率与长上下文建模能力方面的卓越能力。
In large language models, the demand for modeling long contexts is constantly increasing, but the quadratic complexity of the standard self-attention mechanism often becomes a bottleneck. Although existing sparse attention mechanisms have improved efficiency, they may still encounter issues such as static patterns or information loss. We introduce a trainable dynamic mask sparse attention mechanism, Dynamic Mask Attention, which effectively utilizes content-aware and position-aware sparsity. DMA achieves this through two key innovations: First, it dynamically generates content-aware sparse masks from value representations, enabling the model to identify and focus on critical information adaptively. Second, it implements position-aware sparse attention computation that effectively skips unnecessary calculation regions. This dual-sparsity design allows the model to significantly reduce the computational complexity of important information while retaining complete information, achieving an excellent balance between information fidelity and computational efficiency. We have verified the performance of DMA through comprehensive experiments. Comparative studies show that DMA outperforms multi-head attention, sliding window attention, multi-head latent attention, and native sparse attention in terms of perplexity under Chinchilla Scaling Law settings. Moreover, in challenging multi-query associative recall tasks, DMA also demonstrates superior performance and efficiency compared to these methods. Crucially, in the evaluation of a 1.7B parameter model, DMA significantly outperforms multi-head attention in both standard benchmark performance and the challenging needle-in-a-haystack task. These experimental results highlight its capability to balance model efficiency and long-context modeling ability effectively.