Attention Sparse Dynamique à Masque Entraînable

Trainable Dynamic Mask Sparse Attention

August 4, 2025

papers.authors: Jingze Shi, Yifan Wu, Bingheng Wu, Yiran Peng, Liangdong Wang, Guang Liu, Yuyu Luo

cs.AI

papers.abstract

Dans les grands modèles de langage, la demande pour modéliser des contextes longs ne cesse d'augmenter, mais la complexité quadratique du mécanisme d'auto-attention standard devient souvent un goulot d'étranglement. Bien que les mécanismes d'attention éparse existants aient amélioré l'efficacité, ils peuvent encore rencontrer des problèmes tels que des motifs statiques ou une perte d'information. Nous introduisons un mécanisme d'attention éparse à masque dynamique et entraînable, appelé Dynamic Mask Attention (DMA), qui exploite efficacement la parcimonie basée sur le contenu et la position. DMA y parvient grâce à deux innovations clés : premièrement, il génère dynamiquement des masques éparses basés sur le contenu à partir des représentations de valeur, permettant au modèle d'identifier et de se concentrer sur les informations critiques de manière adaptative. Deuxièmement, il met en œuvre un calcul d'attention éparse basé sur la position qui saute efficacement les régions de calcul inutiles. Cette conception à double parcimonie permet au modèle de réduire significativement la complexité computationnelle des informations importantes tout en conservant l'intégralité des informations, atteignant un équilibre optimal entre fidélité de l'information et efficacité computationnelle. Nous avons vérifié les performances de DMA à travers des expériences approfondies. Des études comparatives montrent que DMA surpasse l'attention multi-tête, l'attention par fenêtre glissante, l'attention latente multi-tête et l'attention éparse native en termes de perplexité dans les conditions de la loi d'échelle Chinchilla. De plus, dans des tâches complexes de rappel associatif multi-requêtes, DMA démontre également une performance et une efficacité supérieures par rapport à ces méthodes. Enfin, dans l'évaluation d'un modèle de 1,7 milliard de paramètres, DMA surpasse significativement l'attention multi-tête à la fois dans les performances standard sur les benchmarks et dans la tâche difficile de recherche d'une aiguille dans une botte de foin. Ces résultats expérimentaux mettent en évidence sa capacité à équilibrer efficacement l'efficacité du modèle et la capacité de modélisation de contextes longs.

English

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.