thu-nics/MoA

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.

MoA overview: attention patterns and elastic rules

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.

Observation 01

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. Matching each head to its own span raises content retrieval accuracy from 25% to 98%.

Attention matrices of three Vicuna-7B heads: a local head, a medium-span head, and a global head — Attention matrices from three Vicuna-7B heads, each averaged over 256 LongEval items: local, medium, and global spans coexist in one model.

Observation 02

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.

Attention matrix per head

Input length 4k

4k 8k 12k 16k

Head A

Local, fixed

S = 1.6k (β = 0)

40% of input

Head B

Medium, fixed

S = 5.6k (β = 0)

100% of input

Head C

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.

MoA workflow — MoA profiles attention influence, optimizes per-head elastic rules, and deploys the resulting heterogeneous sliding-window mask.

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 profiles on long summaries and uses the dense model's own output as the reference, so gradients track the behavior compression must preserve.

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

0.25 retrieval / 4.95 PPL

MoA Profiling

MultiNews + dense reference

0.95 retrieval / 3.96 PPL

MoA calibration dataset ablation — Table 2 isolates the useful profiling ingredients: long-range dependence and dense-model alignment.

Experiments

MoA improves the accuracy-efficiency frontier.

Experiments cover long-context retrieval, LV-Eval, LongBench, perplexity, throughput, memory, and search-space ablations.

1.5-7.1x

Retrieval gain over uniform-window baselines across Vicuna and Llama3.

6.6-8.2x

Decode throughput speedup over FlashAttention2 at 50% density.

1.7-1.9x

Decode throughput speedup over vLLM with minimal performance impact.

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%.

MoA comparative performance results — Table 4 from the COLM 2025 version compares MoA with H2O, InfLLM, and StreamingLLM across Vicuna-7B, Vicuna-13B, Llama3-8B, and Llama3-70B.

Runtime and Memory Efficiency

The throughput gains come from static KV-cache size, reduced attention computation, increased batch capacity, and a CUDA kernel for heterogeneous attention.

MoA decode throughput comparison — At 50% density, MoA reaches 6.6-8.2x throughput over FlashAttention2 and 1.7-1.9x over vLLM on A100-80GB.

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.

MoA ablation study — Elastic heterogeneous rules improve retrieval accuracy from 0.25 to 0.98 at 8k input length under a 25% density budget.

Demo

A plug-and-play compression pipeline for long-context inference.

The implementation exposes MoA as a practical attention backend for long-context decoding under controlled KV-cache density.

01

Profile an LLM with a model-aligned long-context calibration set.

02

Search per-layer, per-head elastic window rules under the density budget.

03

Serve with a static, sparse, heterogeneous KV-cache plan.

Team

Meet the researchers.

Tianyu Fu^1,2,* Haofeng Huang^1,2,* Xuefei Ning^1,* Genghan Zhang³ Boju Chen¹ Tianqi Wu^1,2 Hongyi Wang^1,2 Zixiao Huang^1,2 Shiyao Li^1,2 Shengen Yan^1,2 Guohao Dai^2,4 Huazhong Yang¹ Yu Wang¹

¹Tsinghua University ²Infinigence-AI ³Stanford University ⁴Shanghai Jiao Tong University
*Equal contribution. Genghan Zhang made contributions while at Tsinghua University.

Citation

Cite our work.

If you find MoA useful for your research, please consider citing our paper.

BibTeX

@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}
}