Mixture of Attention
Optimizing LLM Inference Efficiency with Heterogeneous Sliding-Window Lengths.
Overview
Short attention spans can still preserve long-context behavior.
Uniform sliding windows give every head the same local budget. MoA assigns each head its own elastic span rule so local heads stay cheap and global heads keep enough context.
Uniform spans misallocate budget. Some heads mainly need local context, while others require broad context to preserve retrieval and understanding.
MoA searches per-head elastic rules. Attention span becomes a head-specific function of input length under a target density constraint.
The method is training-free. A calibration pass profiles attention influence, then an optimizer selects the best heterogeneous span plan.
The Discovery
Attention is heterogeneous in two ways.
Profiling pre-trained LLMs reveals two facts about attention heads. Together they explain why a single uniform sliding window cannot serve a whole model.
Different heads have different spans.
At a fixed input length, some heads concentrate in a narrow band near the diagonal, while others spread across the entire input. A uniform window truncates the global heads and breaks long-range retrieval, yet still overspends compute and KV cache on the local ones. At 25% attention density and 8k input length, heterogeneous elastic rules raise retrieval accuracy from 25% to 98%.
Spans scale differently as the input grows.
At short lengths, many heads look identically dense, and the difference only appears once the input stretches. Some heads keep a fixed local span no matter how long the input gets, while others expand in proportion to it. MoA models each head's span as a linear rule S = α + β · N and searches α and β per head.
Local, fixed
S = 1.6k (β = 0)
40% of input
Medium, fixed
S = 5.6k (β = 0)
100% of input
Global, elastic
S = 0.9 · N (β > 0)
90% of input
Illustrative N×N causal attention matrices. Blue cells are kept by each head's rule; faint cells are masked. Heads B and C look alike at 4k. By 16k, one has stayed put while the other has kept growing.
Method
Search a heterogeneous attention-span plan.
MoA constructs a search space of elastic span rules, profiles attention influence on long-context calibration data, and selects the optimal rule for each head.
Elastic Search Space
Each attention head selects from rules that map input length to a hardware-friendly sliding-window span.
Influence Profiling
First-order gradient profiling estimates how masking attention positions changes the prediction loss.
Automatic Optimization
A density-constrained optimizer finds the best per-head plan, then validates it at an unseen sequence length.
Calibration
Useful profiling must test long-context dependence.
MoA calibrates on multi-document summarization and uses the dense model's own summaries as references. Profiling then measures how compression affects predictions that depend on distant context.
Long-range content
Summaries depend on evidence across documents.
Dependent output loss
Compute loss on the generated summary.
Model-aligned reference
Use dense-model summaries as supervision.
Weak Profiling
RedPajama context
MoA Profiling
MultiNews + dense reference
Experiments
MoA improves the accuracy-efficiency frontier.
Experiments cover long-context retrieval, LV-Eval, LongBench, perplexity, throughput, memory, and search-space ablations.
Overall Long-Context Performance
At 50% density, MoA stays close to full attention across retrieval accuracy, LV-Eval, LongBench, and perplexity. It reduces maximum relative drops to 5% on LV-Eval and 3% on LongBench, while the average retrieval drop stays below 1%.
For Vicuna-7B, effective context reaches up to 3.9× the average attention span across heads (Figure 7). Effective context is the longest input with retrieval accuracy above 90%.
Runtime and Memory Efficiency
MoA reduces attention computation and bounds KV cache size, allowing larger batches. A dedicated CUDA kernel supports its heterogeneous windows. The throughput comparison uses the largest supported batch size for each method, rather than a fixed batch size.
Search Space Ablation
Starting from a uniform mask, MoA progressively adds layer heterogeneity, head heterogeneity, and elastic rules. Head heterogeneity and elasticity deliver the largest retrieval improvements at 25% density.
Demo
From calibration to efficient long-context inference.
Calibrate once, search a model-specific attention plan, then use it for long-context decoding within the chosen KV cache budget.
Profile an LLM with a model-aligned long-context calibration set.
Search per-layer, per-head elastic window rules under the density budget.
Serve with a static, sparse, heterogeneous KV cache plan.
Citation
Cite our work.
If you find MoA useful for your research, please consider citing our paper.
@inproceedings{fu2025moa,
title={Mixture of Attention Spans: Optimizing LLM Inference Efficiency with Heterogeneous Sliding-Window Lengths},
author={Fu, Tianyu and Huang, Haofeng and Ning, Xuefei and Zhang, Genghan and Chen, Boju and Wu, Tianqi and Wang, Hongyi and Huang, Zixiao and Li, Shiyao and Yan, Shengen and Dai, Guohao and Yang, Huazhong and Wang, Yu},
booktitle={Conference on Language Modeling},
year={2025},
url={https://arxiv.org/abs/2406.14909}
}