AoH: Autonomy-of-Heads

Paper: AoH paper

AoH (Autonomy-of-Heads), a training-free method that classifies attention heads from the frozen query-key projection geometry.

Demo

aoh-demo-0804-under-10mb.mp4

Method

For each query head, AoH measures the effective rank of the query-key interaction matrix using a small head-dimension Gram matrix. Low-rank heads are classified as retrieval heads and retain full attention; higher-rank heads are classified as streaming heads and can use a local attention window.

Both multi-head attention (MHA) and grouped-query attention (GQA) models are supported. For GQA, per-query-head scores are aggregated within each KV group using the mean, minimum, or maximum score.

Installation

Python 3.10 or newer is required.

pip install -e .

Install plotting dependencies when visualization is needed:

pip install -e ".[visualization]"

Install LongBench evaluation dependencies before running the result scripts:

pip install -e ".[evaluation]"

Model weights are loaded through Hugging Face Transformers. Access to gated models must be configured separately with Hugging Face.

Command Line

python -m aoh \
  --model_path Qwen/Qwen2.5-7B \
  --device cuda:0 \
  --sparsity 0.5 \
  --agg mean \
  --output_dir outputs/qwen25-7b

The installed aoh command provides the same interface:

aoh --model_path Qwen/Qwen2.5-7B --output_dir outputs/qwen25-7b

Useful options:

• --sparsity: fraction of heads classified as streaming in each layer.
• --agg {mean,min,max}: GQA score aggregation strategy.
• --binary: use binary projection weights for accelerated rank estimation.
• --torch_dtype {auto,float32,float16,bfloat16}: checkpoint loading dtype; auto preserves the stored dtype and is the default. AoH still performs the Q/K rank computation in float32.

Run python -m aoh --help for the complete interface.

Python API

from aoh import AoHClassifier

classifier = AoHClassifier(
    model_path="Qwen/Qwen2.5-7B",
    device="cuda:0",
    agg="mean",
)
labels = classifier.classify(sparsity=0.5)
classifier.save("outputs/qwen25-7b")

labels has shape (num_layers, num_kv_heads) for GQA models and (num_layers, num_heads) for MHA models. A value of 1 denotes a retrieval head and 0 denotes a streaming head.

Outputs

Each run writes the following files to --output_dir:

• head_classification.npy: binary retrieval/streaming labels.
• eff_rank.npy: per-query-head effective ranks.
• eff_rank_kv.npy: effective ranks aggregated by KV group.
• results.json: model and run metadata.

Examples

python examples/qwen25_7b.py
python examples/qwen3.py

The examples load full model weights. Adjust the model name, device, and output path in each script for your environment.

Results

bash scripts/run_result_eval.sh

We evaluate AoH on LongBench, a long-context understanding suite covering 21 tasks across six categories. The evaluated models are Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen2.5-7B. Unless otherwise stated, streaming heads retain 128 sink tokens and a 256-token recent window, and the main experiments use 50% per-layer sparsity.

Efficiency

bash scripts/run_efficiency_eval.sh

AoH reduces KV-cache memory nearly proportionally to the fraction of streaming heads, and its latency benefits become more pronounced at longer contexts where attention and cache access dominate runtime. As shown in Figure 5 and Table 11, at 75% sparsity, AoH achieves up to 3.24x prefill speedup and 9.14x decode speedup at 256K context, while reducing KV-cache memory by up to 3.98x on Llama-3.1-8B-Instruct-128K.

Visualization

After running AoH:

python visualization/plot_head_classification.py outputs/qwen25-7b
python visualization/plot_effective_rank.py outputs/qwen25-7b

Run the Demo

After setting up the environment, you can run the following script to execute the PassKey retrieval demo on the Llama-3.1-8B-Instruct-128K model. The demo is designed to run on a single A100 GPU and supports a context length of up to 128K tokens.

bash scripts/demo.sh

Citation

If you find AoH useful or relevant to your project and research, please kindly cite our paper:

@misc{yang2026autonomy,
  title={Autonomy-of-Heads: Data-Free Sparse Attention from Frozen Query-Key Geometry},
  author={Yang, Yehan and Shang, Junyuan and Li, Yang and Zhao, Guanqun and Wang, Shuohuan and Yu, Dianhai},
  year={2026},
  eprint={2608.06849},
  archivePrefix={arXiv},
  primaryClass={cs.CL},
  url={https://arxiv.org/abs/2608.06849}
}