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

TL;DR

This paper introduces a dynamic attention mask mechanism for long-context LLMs that adaptively assigns attention patterns, reducing computational costs while maintaining high performance without the need for fine-tuning.

Contribution

It proposes a novel dynamic sparse attention method that learns context-aware masks, outperforming static methods and enabling efficient long-sequence processing without fine-tuning.

Findings

Achieves high alignment with full-attention models.

Reduces memory and compute overhead significantly.

Maintains retrieval accuracy in long-sequence tasks.

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.