DAM : Dynamic Attention Mask for Long-Context Large Language Model Inference Acceleration (2025)

Abstract: Long-context understanding is crucial for many NLP applications, yet transformers struggle with efficiency due to the quadratic complexity of self-attention. Sparse attention methods alleviate this cost but often impose static, predefined masks, failing to capture heterogeneous attention patterns. This results in suboptimal token interactions, limiting adaptability and retrieval accuracy in long-sequence tasks. This work introduces a dynamic sparse attention mechanism that assigns adaptive masks at the attention-map level, preserving heterogeneous patterns across layers and heads. Unlike existing approaches, our method eliminates the need for fine-tuning and predefined mask structures while maintaining computational efficiency. By learning context-aware attention structures, it achieves high alignment with full-attention models, ensuring minimal performance degradation while reducing memory and compute overhead. This approach provides a scalable alternative to full attention, enabling the practical deployment of large-scale Large Language Models (LLMs) without sacrificing retrieval performance. DAM is available at: https://github.com/HanzhiZhang-Ulrica/DAM.

## introduction understanding long contexts is essential for document summarization, question answering, and retrieval-augmented generation. long-context nlp applications power legal analysis, financial reporting, and knowledge graph construction, where maintaining coherence across tokens is critical. however, existing llms struggle with long sequences due to inefficiency in self-attention. llms leverage the transformer architecture, which models long-range token dependencies through self-attention, enabling direct interactions between all tokens. multi-head attention enhances expressivity by capturing multiple token relationships, while positional embeddings preserve order information. dynamic token representations allow contextual meaning to evolve across layers, ensuring coherent and accurate long-term recall. however, transformers process long contexts inefficiently. quadratic complexity makes full attention computationally prohibitive, forcing models to truncate inputs, leading to information loss in longdocument tasks. attempts to mitigate this through fixed context windows or uniform attention fail to distinguish between critical and redundant information. meanwhile, streaming applications suffer from recomputing at every step, as each newly generated token attends to all previous tokens. this results in redundant computation, growing memory usage, and increased latency, making real-time processing impractical. the inability to process long contexts also directly impacts businesses and researchers. legal and financial institutions rely on ai to analyze contracts and reports @xcite , but truncated inputs cause critical details to be lost. ai assistants in customer service forgetting past interactions may fail to maintain coherent conversations. researchers pushing transformer efficiency face skyrocketing computational costs, making largescale deployment unsustainable. without an efficient solution, enterprises must rely on costly and ineffective workarounds like document chunking, which may destroy contextual coherence. sparse attention is an efficiency-driven approach to mitigating long-context inefficiency in generative llms by enforcing structured sparsity. figure 1(b) illustrates static sparse attention methods that reduce computational cost by enforcing fixedspan sliding window and global masks across all heads and input lengths. this approach improves efficiency but sacrifices flexibility, forcing models to rely only on local interactions. as a result, these models fail to adapt to long-range dependencies, reducing accuracy in complex retrieval and reasoning tasks. figure 1(c ) improves flexibility by assigning paper boats sailed across a puddle of stardust drifting toward the moon short sequence long sequence matrix value 0 1 (b) static sparse mask (a) attention map (c) layer-, head-wise sparse mask (d) dam dynamic sparse mask short sequence token connections fully connection local & global connection local & global connection irregular connection different masks to layers and heads, removing the need for fine-tuning. however, it assumes attention structures can be predefined, failing to capture heterogeneous token interactions that emerge dynamically. the result is a rigid sparsity pattern that still requires processing all sequence lengths, making it resource-intensive for long-context applications. this work introduces dam, a novel framework for dynamic sparse attention, as illustrated in figure 1(d ). dam generates adaptive sparse attention masks at the granularity of individual attention maps, thereby capturing both layer-specific structural patterns and input-dependent variations in attention. in contrast to prior approaches that often rely on fixed or globally-defined sparsity patterns, dam preserves the heterogeneity of attention patterns across different layers and heads, leading to improved expressiveness. furthermore, the method eliminates the need for manual, task-specific finetuning of the sparsity structure, while maintaining the computational benefits of sparse attention. our contributions are summarized as follows: • we propose a dynamic sparse attention framework that assigns distinct, adaptive sparse masks to each attention map, preserving heterogeneous patterns across heads and layers. ## related work attention mechanisms enable transformers to model dependencies across sequences but introduce computational challenges at scale. researchers have explored multiple strategies to address these inefficiencies, including kv-caching for faster inference, sparse and hierarchical attention for memory reduction, state-space models for efficient streaming, and hybrid architectures for improved long-term memory tracking. while these approaches enhance scalability, each introduces trade-offs that limit their applicability to long-sequence processing. kv-cache enhances autoregressive decoding by storing key and value representations from previous steps, allowing reuse instead of recomputing attention scores for all tokens @xcite @xcite @xcite @xcite . this reduces redundant computation and accelerates inference but increases memory usage, limiting scalability for long sequences @xcite @xcite . cache management adds complexity, and performance gains depend on reuse efficiency @xcite @xcite . while kv-cache mitigates inefficiencies in autoregressive generation, it does not reduce the fundamental complexity of self-attention. sparse attention reduces token interactions to improve efficiency @xcite @xcite . static sparse attention applies predefined masks across all processed sentences to lower computational cost and improve hardware utilization @xcite @xcite . common approaches include global, sliding window, and random masks, with local attention patterns enabling kv-cache eviction beyond the attention span to reduce memory usage @xcite @xcite . however, static masks remain uniform across layers and heads, ignoring token-specific dependencies. this rigidity leads to information loss in longsequence tasks, where retrieval accuracy relies on adapting attention spans dynamically. other strategies generate distinct masks by leveraging statistical information, defining role-specific constraints for attention heads, or introducing context-dependent sparsity @xcite @xcite . while these approaches increase flexibility, they fail to dynamically capture heterogeneous attention within individual maps and still process all sequence lengths, raising resource costs. to introduce flexibility, some approaches assign different predefined sparse patterns to layers and heads @xcite @xcite . they select masks based on input length, improving adaptability without requiring fine-tuning. however, it assumes optimal attention structures are predefined, missing dynamic token interactions, and require evaluating multiple sequence lengths, thereby increasing computational overhead. this limitation motivates methods to infer sparse structures without exhaustive manual design or repeated inference. ## preliminaries transformer models adopt the scaled dot-product attention mechanism, a core component for capturing relationships between tokens in a sequence @xcite . attention scores are calculated as @xmath0, where @xmath1, respectively. here, n and m denote the number of query and key/value vectors, while d k represents the dimensionality of each key/query vector. the resulting matrix @xmath2, representing the pairwise similarities between queries and keys. the scaling factor 1 √ d k is crucial for maintaining numerical stability during training, preventing the dot products from growing excessively large, which can lead to vanishing gradients during backpropagation. this scaling mitigates issues caused by large variances in the logits, particularly when applying a masking operation. the computation of attention scores for all pairs of tokens has a quadratic time complexity of o(n 2 ) with respect to the sequence length, which becomes computationally expensive for long sequences. sparse attention mechanisms address this computational bottleneck by imposing structured sparsity on the attention matrix. this is achieved through a binary mask m ℓ,@xmath3, 1} n×m for each layer ℓ and head h, defined as: the mask m ℓ,h is applied element-wise to the attention logits s ′ = s ⊙ m ℓ,h , where ⊙ denotes the hadamard product (element-wise multiplication). this effectively prevents attention between specific token pairs. the masked attention logits are then normalized using the softmax function: the output of the attention mechanism is computed as a weighted sum of the values, where @xmath4. while sparse attention mechanisms substantially improve computational efficiency, they inherently restrict the model's ability to learn long-range dependencies by limiting token interactions. a key limitation of many sparse attention approaches is their reliance on fixed sparsity patterns. such patterns are unable to adapt to the dynamic nature of attention, including variations in sequence length and the diversity of attention distributions across different inputs. this rigidity can result in a significant reduction in performance, especially when dealing with long sequences or tasks requiring the modeling of intricate relationships. moreover, predefined sparse attention structures like sliding window @xcite or global attention @xcite often overlook the critical variations in attention patterns that occur across different layers and heads within the network. the optimal set of token interactions evolves across layers, rendering fixed sparsity patterns a bottleneck. this motivates the need for a dynamic, structure-aware sparsity mechanism that adapts to position-wise attention patterns while maintaining compatibility with pretrained transformer architectures. ## dynamic attention masks (dam) this section introduces our proposed dynamic attention mask (dam) mechanism. we first motivate the design by illustrating the dynamic nature of attention patterns across layers and heads in transformer models. then, we detail the architecture of dam, and finally, we describe its integration into the standard transformer framework. prior research has investigated and validated the existence of dynamic attention patterns across attention heads and layers in transformer models @xcite . to visualize these patterns, we analyze attention maps obtained from a llama 3.2 3b instruct model (ai, 2024) evaluated on the multi-news summarization benchmark @xcite . figure 2 presents these attention maps, revealing inconsistencies in the underlying sparse structures across different heads and layers. the top row of figure 2 displays the average attention values across the dataset. while these average maps suggest the presence of dynamic patterns, the patterns themselves are not readily discernible, hindering a deeper understanding and impeding the design of effective sparsity-inducing techniques. ## conclusion we introduce dam, a sparse attention method that dynamically captures heterogeneous token interactions, overcoming the limitations of static and predefined sparsity patterns. dam learns adaptive attention masks that retain crucial dependencies, improving retrieval accuracy while significantly reducing computational cost. experiments across long-context benchmarks demonstrate dam's effectiveness in maintaining full-attention performance with lower memory and compute requirements. by bridging the gap between efficiency and expressivity in sparse attention, dam provides a scalable solution for long-context processing.