Attenzione Sparsa Dinamica con Maschera Addestrabile

Trainable Dynamic Mask Sparse Attention

August 4, 2025

Autori: Jingze Shi, Yifan Wu, Bingheng Wu, Yiran Peng, Liangdong Wang, Guang Liu, Yuyu Luo

cs.AI

Abstract

Nei grandi modelli linguistici, la necessità di modellare contesti lunghi è in costante aumento, ma la complessità quadratica del meccanismo standard di self-attention spesso rappresenta un collo di bottiglia. Sebbene i meccanismi esistenti di attenzione sparsa abbiano migliorato l'efficienza, possono comunque incontrare problemi come pattern statici o perdita di informazioni. Introduciamo un meccanismo di attenzione sparsa con maschera dinamica addestrabile, Dynamic Mask Attention (DMA), che sfrutta efficacemente la sparsità consapevole del contenuto e della posizione. DMA raggiunge questo obiettivo attraverso due innovazioni chiave: in primo luogo, genera dinamicamente maschere sparse consapevoli del contenuto dalle rappresentazioni dei valori, consentendo al modello di identificare e concentrarsi in modo adattivo sulle informazioni critiche. In secondo luogo, implementa un calcolo di attenzione sparsa consapevole della posizione che salta efficacemente le regioni di calcolo non necessarie. Questo design a doppia sparsità consente al modello di ridurre significativamente la complessità computazionale delle informazioni importanti mantenendo al contempo l'informazione completa, raggiungendo un eccellente equilibrio tra fedeltà delle informazioni ed efficienza computazionale. Abbiamo verificato le prestazioni di DMA attraverso esperimenti completi. Studi comparativi dimostrano che DMA supera l'attenzione multi-testa, l'attenzione a finestra scorrevole, l'attenzione latente multi-testa e l'attenzione sparsa nativa in termini di perplessità nelle impostazioni della Chinchilla Scaling Law. Inoltre, in impegnativi compiti di richiamo associativo multi-query, DMA dimostra prestazioni e efficienza superiori rispetto a questi metodi. Crucialmente, nella valutazione di un modello con 1,7 miliardi di parametri, DMA supera significativamente l'attenzione multi-testa sia nelle prestazioni standard dei benchmark che nel complesso compito di "ago in un pagliaio". Questi risultati sperimentali evidenziano la sua capacità di bilanciare efficacemente l'efficienza del modello e la capacità di modellare contesti lunghi.

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.