Trainable Dynamic Mask Sparse Attention

• URL: http://arxiv.org/abs/2508.02124v5
• Date: Mon, 20 Oct 2025 08:48:15 GMT
• Title: Trainable Dynamic Mask Sparse Attention
• Authors: Jingze Shi, Yifan Wu, Yiran Peng, Bingheng Wu, Liangdong Wang, Guang Liu, Yuyu Luo,
• Abstract summary: We introduce a trainable dynamic mask sparse attention mechanism, a method that merges the advantages of both position-aware and content-aware approaches.<n>We demonstrate that the introduced dynamic mask and sparse weights do not obstruct gradients, supporting end-to-end training.
• Score: 11.506985057671015
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: The increasing demand for long-context modeling in large language models (LLMs) is bottlenecked by the quadratic complexity of the standard self-attention mechanism. The community has proposed sparse attention to mitigate this issue. However, position-aware sparse attention methods rely on static sparse structures that lack adaptability to diverse query contexts, while content-aware sparse attention methods depend on heuristic key-value selection, hindering full differentiability. We introduce a trainable dynamic mask sparse attention mechanism, a method that merges the advantages of both position-aware and content-aware approaches. Dynamic Mask Attention (DMA) achieves this through three key innovations: First, it leverages value vector representations to generate content-aware dynamic masks, enabling the model to adaptively identify and attend to critical information. Second, it computes position-aware sparse weights in a hardware-friendly manner, efficiently skipping unnecessary computational regions. Finally, we demonstrate that the introduced dynamic mask and sparse weights do not obstruct gradients, supporting end-to-end training. We have validated the performance of DMA through comprehensive experiments. A large body of experimental evidence shows that DMA consistently holds a Pareto advantage over state-of-the-art sparse attention baselines in tasks including scaling laws, multi-query associative recall, standard benchmarks, and needle in a haystack tests, while also delivering up to a 10x overall speedup. These results highlight its ability to effectively balance model efficiency with long-context modeling capabilities. Our computational kernel code is now open-source at https://github.com/SmallDoges/flash-dmattn to encourage further research and application by the community.

Related papers

• Elastic Attention: Test-time Adaptive Sparsity Ratios for Efficient Transformers [42.80120203718226]
  We propose Elastic Attention, which allows the model to adjust its overall sparsity based on the input.<n>Within only 12 hours of training on 8xA800 GPUs, our method enables models to achieve both strong performance and efficient inference.
  arXiv  Detail & Related papers  (2026-01-24T08:22:07Z)
• MTraining: Distributed Dynamic Sparse Attention for Efficient Ultra-Long Context Training [23.925430484357975]
  MTraining is a distributed methodology for training Large Language Models with ultra-long contexts.<n>MTraining integrates a dynamic sparse training pattern, balanced sparse ring attention, and hierarchical sparse ring attention.<n>MTraining achieves up to a 6x higher training throughput while preserving model accuracy.
  arXiv  Detail & Related papers  (2025-10-21T17:25:32Z)
• Less is More: Empowering GUI Agent with Context-Aware Simplification [62.02157661751793]
  We propose a context-aware framework for building an efficient and effective GUI Agent, termed SimpAgent.<n>With the above components, SimpAgent reduces 27% FLOPs and achieves superior GUI navigation performances.
  arXiv  Detail & Related papers  (2025-07-04T17:37:15Z)
• DAM: Dynamic Attention Mask for Long-Context Large Language Model Inference Acceleration [12.172968576254469]
  We introduce a dynamic sparse attention mechanism that assigns adaptive masks at the attention-map level.<n>By learning context-aware attention structures, it achieves high alignment with full-attention models, ensuring minimal performance degradation.<n>This approach provides a scalable alternative to full attention, enabling the practical deployment of large-scale Large Language Models.
  arXiv  Detail & Related papers  (2025-06-06T20:24:36Z)
• Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free [81.65559031466452]
  We conduct experiments to investigate gating-augmented softmax attention variants.<n>We find that a simple modification-applying a head-specific sigmoid gate after the Scaled Dot-Product Attention (SDPA)-consistently improves performance.
  arXiv  Detail & Related papers  (2025-05-10T17:15:49Z)
• Dynamic Attention Mechanism in Spatiotemporal Memory Networks for Object Tracking [8.040709469401257]
  We propose a differentiable dynamic attention mechanism that adaptively channel adjusts attention weights by analyzing spatial attention weights.<n>A lightweight gating network that autonomously allocates computational resources based on target motion states, prioritizes high-discriminability features in challenging scenarios.
  arXiv  Detail & Related papers  (2025-03-21T00:48:31Z)
• Transformer Meets Twicing: Harnessing Unattended Residual Information [2.1605931466490795]
  Transformer-based deep learning models have achieved state-of-the-art performance across numerous language and vision tasks.<n>While the self-attention mechanism has proven capable of handling complex data patterns, it has been observed that the representational capacity of the attention matrix degrades significantly across transformer layers.<n>We propose the Twicing Attention, a novel attention mechanism that uses kernel twicing procedure in nonparametric regression to alleviate the low-pass behavior of associated NLM smoothing.
  arXiv  Detail & Related papers  (2025-03-02T01:56:35Z)
• DSMoE: Matrix-Partitioned Experts with Dynamic Routing for Computation-Efficient Dense LLMs [70.91804882618243]
  This paper proposes DSMoE, a novel approach that achieves sparsification by partitioning pre-trained FFN layers into computational blocks.<n>We implement adaptive expert routing using sigmoid activation and straight-through estimators, enabling tokens to flexibly access different aspects of model knowledge.<n>Experiments on LLaMA models demonstrate that under equivalent computational constraints, DSMoE achieves superior performance compared to existing pruning and MoE approaches.
  arXiv  Detail & Related papers  (2025-02-18T02:37:26Z)
• Core Context Aware Transformers for Long Context Language Modeling [50.774702091154204]
  We propose a plug-and-play Core Context Aware (CCA) Attention for efficient long-context modeling.<n>Our method automatically focuses and strengthens core context while diminishing redundancy during the learning process.<n>Our method is able to replace the self-attention module in existing Large Language Models with minimal fine-tuning cost.
  arXiv  Detail & Related papers  (2024-12-17T01:54:08Z)
• Unsupervised Representation Learning of Complex Time Series for Maneuverability State Identification in Smart Mobility [0.0]
  In smart mobility, MTS plays a crucial role in providing temporal dynamics of behaviors such as maneuver patterns. In this work, we aim to address challenges associated with modeling MTS data collected from a vehicle using sensors. Our goal is to investigate the effectiveness of two distinct unsupervised representation learning approaches in identifying maneuvering states in smart mobility.
  arXiv  Detail & Related papers  (2024-08-26T15:16:18Z)
• Bringing Masked Autoencoders Explicit Contrastive Properties for Point Cloud Self-Supervised Learning [116.75939193785143]
  Contrastive learning (CL) for Vision Transformers (ViTs) in image domains has achieved performance comparable to CL for traditional convolutional backbones. In 3D point cloud pretraining with ViTs, masked autoencoder (MAE) modeling remains dominant.
  arXiv  Detail & Related papers  (2024-07-08T12:28:56Z)
• Interpreting and Improving Attention From the Perspective of Large Kernel Convolution [51.06461246235176]
  We introduce Large Kernel Convolutional Attention (LKCA), a novel formulation that reinterprets attention operations as a single large- Kernel convolution.<n>LKCA achieves competitive performance across various visual tasks, particularly in data-constrained settings.
  arXiv  Detail & Related papers  (2024-01-11T08:40:35Z)
• Understanding Self-attention Mechanism via Dynamical System Perspective [58.024376086269015]
  Self-attention mechanism (SAM) is widely used in various fields of artificial intelligence. We show that intrinsic stiffness phenomenon (SP) in the high-precision solution of ordinary differential equations (ODEs) also widely exists in high-performance neural networks (NN) We show that the SAM is also a stiffness-aware step size adaptor that can enhance the model's representational ability to measure intrinsic SP.
  arXiv  Detail & Related papers  (2023-08-19T08:17:41Z)
• Sparse Modular Activation for Efficient Sequence Modeling [94.11125833685583]
  Recent models combining Linear State Space Models with self-attention mechanisms have demonstrated impressive results across a range of sequence modeling tasks. Current approaches apply attention modules statically and uniformly to all elements in the input sequences, leading to sub-optimal quality-efficiency trade-offs. We introduce Sparse Modular Activation (SMA), a general mechanism enabling neural networks to sparsely activate sub-modules for sequence elements in a differentiable manner.
  arXiv  Detail & Related papers  (2023-06-19T23:10:02Z)
• Object Segmentation by Mining Cross-Modal Semantics [68.88086621181628]
  We propose a novel approach by mining the Cross-Modal Semantics to guide the fusion and decoding of multimodal features. Specifically, we propose a novel network, termed XMSNet, consisting of (1) all-round attentive fusion (AF), (2) coarse-to-fine decoder (CFD), and (3) cross-layer self-supervision.
  arXiv  Detail & Related papers  (2023-05-17T14:30:11Z)
• Rethinking Attention Mechanism in Time Series Classification [6.014777261874646]
  We promote the efficiency and performance of the attention mechanism by proposing our flexible multi-head linear attention (FMLA) We propose a simple but effective mask mechanism that helps reduce the noise influence in time series and decrease the redundancy of the proposed FMLA. We conduct extensive experiments on 85 UCR2018 datasets to compare our algorithm with 11 well-known ones and the results show that our algorithm has comparable performance in terms of top-1 accuracy.
  arXiv  Detail & Related papers  (2022-07-14T07:15:06Z)
• Robust Person Re-Identification through Contextual Mutual Boosting [77.1976737965566]
  We propose the Contextual Mutual Boosting Network (CMBN) to localize pedestrians. It localizes pedestrians and recalibrates features by effectively exploiting contextual information and statistical inference. Experiments on the benchmarks demonstrate the superiority of the architecture compared the state-of-the-art.
  arXiv  Detail & Related papers  (2020-09-16T06:33:35Z)
• KSM: Fast Multiple Task Adaption via Kernel-wise Soft Mask Learning [49.77278179376902]
  Deep Neural Networks (DNN) could forget the knowledge about earlier tasks when learning new tasks, and this is known as textitcatastrophic forgetting. Recent continual learning methods are capable of alleviating the catastrophic problem on toy-sized datasets. We propose a new training method called textit- Kernel-wise Soft Mask (KSM), which learns a kernel-wise hybrid binary and real-value soft mask for each task.
  arXiv  Detail & Related papers  (2020-09-11T21:48:39Z)
• Optimization-driven Machine Learning for Intelligent Reflecting Surfaces Assisted Wireless Networks [82.33619654835348]
  Intelligent surface (IRS) has been employed to reshape the wireless channels by controlling individual scattering elements' phase shifts. Due to the large size of scattering elements, the passive beamforming is typically challenged by the high computational complexity. In this article, we focus on machine learning (ML) approaches for performance in IRS-assisted wireless networks.
  arXiv  Detail & Related papers  (2020-08-29T08:39:43Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.