Trainable Dynamic Mask Sparse Attention
Authors: name: Jingze Shi, name: Yifan Wu, name: Bingheng Wu, name: Yiran Peng, name: Liangdong Wang, name: Guang Liu, name: Yuyu Luo

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.