Abstract:Sparse attention can effectively mitigate the significant memory and throughput demands of Large Language Models (LLMs) in long contexts. Existing methods typically employ a uniform sparse attention mask, applying the same sparse pattern across different attention heads and input lengths. However, this uniform approach fails to capture the diverse attention patterns inherent in LLMs, ignoring their distinct accuracy-latency trade-offs. To address this challenge, we propose the Mixture of Attention (MoA), which automatically tailors distinct sparse attention configurations to different heads and layers. MoA constructs and navigates a search space of various attention patterns and their scaling rules relative to input sequence lengths. It profiles the model, evaluates potential configurations, and pinpoints the optimal sparse attention compression plan. MoA adapts to varying input sizes, revealing that some attention heads expand their focus to accommodate longer sequences, while other heads consistently concentrate on fixed-length local contexts. Experiments show that MoA increases the effective context length by $3.9\times$ with the same average attention span, boosting retrieval accuracy by $1.5-7.1\times$ over the uniform-attention baseline across Vicuna-{7B,13B}, and Llama3-{8B,70B} models. Moreover, MoA narrows the capability gaps between sparse and dense models, reducing the maximum relative performance drop from $9\%-36\%$ to within $5\%$ across two long-context understanding benchmarks. MoA achieves a $1.2-1.4\times$ GPU memory reduction, boosting decode throughput by $6.6-8.2\times$ and $1.7-1.9\times$ compared to FlashAttention2 and vLLM, with minimal impact on performance. Our code is available at \url{<a class="link-external link-https" href="https://github.com/thu-nics/MoA" rel="external noopener nofollow">this https URL</a>}.

COLAFormer: Communicating local–global features with linear computational complexity

COLA-Net: Collaborative Attention Network for Image Restoration

Efficient Attention: Attention with Linear Complexities

ELA: Efficient Local Attention for Deep Convolutional Neural Networks

Efficient Learnable Collaborative Attention for Single Image Super-Resolution

Local-Global Attention: An Adaptive Mechanism for Multi-Scale Feature Integration

What Limits the Performance of Local Self-attention?

Non-Local Neural Networks With Grouped Bilinear Attentional Transforms

Short-Long Convolutions Help Hardware-Efficient Linear Attention to Focus on Long Sequences

MoA: Mixture of Sparse Attention for Automatic Large Language Model Compression

Locally Enhanced Self-Attention: Combining Self-Attention and Convolution as Local and Context Terms

Interpreting and Improving Attention From the Perspective of Large Kernel Convolution

Axial-DeepLab: Stand-Alone Axial-Attention for Panoptic Segmentation

FocusLLaVA: A Coarse-to-Fine Approach for Efficient and Effective Visual Token Compression

CAT: Learning to Collaborate Channel and Spatial Attention from Multi-Information Fusion

Efficient conformer-based speech recognition with linear attention

Adversarial Complementary Learning for Weakly Supervised Object Localization

Localization, balance and affinity: a stronger multifaceted collaborative salient object detector in remote sensing images

Conformer: Local Features Coupling Global Representations for Visual Recognition

Efficient Non-Local Contrastive Attention for Image Super-Resolution

Conformer: Local Features Coupling Global Representations for Recognition and Detection