Title: Dynamic Mask Attention: End-to-End Trainable Content-aware Sparse Attention
Primary area: foundation or frontier models, including LLMs
Keywords: Trainable Sparse Attention, Dynamic Mask Attention

Abstract: Self-attention's computational cost, which scales quadratically with sequence length, creates a fundamental bottleneck for long-context modeling in LLMs, limiting applications such as document understanding, multi-turn reasoning, and code generation. Sparse attention has been proposed to mitigate this issue. Early content-agnostic designs, such as sliding-window attention, successfully reduce computational complexity. However, their reliance on fixed patterns means they can miss important but distant tokens, and are unable to adapt to diverse queries. Recent content-aware methods improve adaptivity but typically rely on hard binary masks or heuristic KV selection, which introduce runtime overhead and fail to be fully differentiable. We propose Dynamic Mask Attention (DMA), a trainable content-aware sparse attention mechanism with head-wise specialization. DMA generates content-driven dynamic masks with continuous importance weights based on value representations, enabling both expressiveness and end-to-end optimization.  Extensive experiments demonstrate that DMA consistently outperforms state-of-the-art sparse attention baselines across pre-training and downstream tasks, reducing perplexity, improving accuracy, and delivering substantial long-sequence speedups of up to 10 times. Our implementation is released in an anonymous repository and has been integrated into widely used frameworks, supporting adoption in large-scale LLM training.