Research map / Architectures and efficiency / Mixture of experts
MoH: Multi-Head Attention as Mixture-of-Head Attention
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Mixture-of-Head attention (MoH), which treats attention heads as MoE experts with per-token top-K routing and weighted summation, reducing activated heads to 50-90% while matching or beating multi-head attention on ViT, DiT, LLMs, and LLaMA3-8B continue-tuning.
Problem
Standard multi-head attention activates all heads for every token even though many heads are redundant, raising computational cost without a corresponding accuracy benefit.
Approach
Rewrite multi-head attention in summation form and treat heads as MoE experts: a router selects top-K heads per token, combines them with weighted (rather than plain) summation, and designates always-on shared heads, with a two-stage routing that balances shared and routed heads. Pretrained MHA models like LLaMA3-8B can be continue-tuned into MoH with a parameter-free norm-based router and straigh
Key results
- MoH-ViT-B achieves 84.9% ImageNet-1K top-1 with 75% of heads, versus TransNeXt-B at 84.8% with 100% of heads.
- MoH-LLaMA3-8B achieves 64.0% average across 14 benchmarks, outperforming LLaMA3-8B (61.6%) using 75% of heads.
- MoH-DiT-XL/2 at 90% heads reaches FID 8.56 at 7M steps versus 9.62 for DiT-XL/2 with 100% heads.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace TransNeXt-B attention with MoH and activate 75% of heads | TransNeXt-B · ImageNet-1K | accuracy: 84.8 -> 84.9 | Helped |
| Replace DiT-S/2 attention with MoH and activate 90% of heads | DiT-S/2 · ImageNet-1K 256×256, 400K steps | FID: 68.40 -> 67.25 | Helped |
| Replace DiT-B/2 attention with MoH and activate 90% of heads | DiT-B/2 · ImageNet-1K 256×256, 400K steps | FID: 43.47 -> 43.40 | Helped |
| Replace DiT-L/2 attention with MoH and activate 90% of heads | DiT-L/2 · ImageNet-1K 256×256, 400K steps | FID: 23.33 -> 23.17 | Helped |
| Replace DiT-XL/2 attention with MoH and activate 90% of heads | DiT-XL/2 · ImageNet-1K 256×256, 7,000K steps, cfg=1.25 | FID: 3.22 -> 2.94 | Helped |
| Train LLM-S with MoH and activate 50% of heads | LLM-S · LLM benchmarks, 100B tokens | average accuracy: 43.9 -> 45.4 | Helped |
| Continue-tune LLaMA3-8B into MoH and activate 75% of heads | LLaMA3-8B · 14 benchmarks | average accuracy: 61.6 -> 64.0 | Helped |
| Activate 50% of attention heads with MoH | Multi-Head Attention · Attention inference, sequence length 256 | time: 0.360 -> 0.225 | Helped |
| Activate 50% of attention heads with MoH | Multi-Head Attention · Attention inference, sequence length 512 | time: 1.376 -> 0.863 | Helped |
| Add shared heads (ablation) | MoH-ViT-S · ImageNet-1K, 100 epochs, 75% heads | accuracy: 75.6 -> 78.3 | Helped |
| Add two-stage routing (ablation) | MoH-ViT-S · ImageNet-1K, 100 epochs, 75% heads | accuracy: 78.3 -> 78.6 | Helped |
| Vary shared-head ratio from 13.9% to 74.0% (ablation) | MoH-ViT-S · ImageNet-1K, 100 epochs, 75% heads | accuracy: 78.6; 78.4 | No clear effect |
Surprising findings
- For MoH-DiT with only 75% of heads activated, performance is worse than DiT, which the authors attribute to dense prediction requiring fine-grained attention.
- For the small MoH-LLM-S, activating 50% of heads outperforms activating 75%, which the authors suggest is a regularization effect.
Limitations
- Activation rate still 50-90%, not lower.
- Largest model tested is 8B parameters.
- Only evaluated on vision and language; multimodal and audio not explored.
Experimental setup
Datasets: ImageNet-1K RedPajama Dolma Pile
Benchmarks: ImageNet-1K SciQ PIQA WinoGrande OpenbookQA LogiQA TruthfulQA MMLU CEVAL CMMLU GSM8K HellaSwag
Metrics: top-1 accuracy FID sFID IS precision recall exact match average accuracy
Baselines: TransNeXt DiT LLaMA3-8B DeiT Swin Transformer MoA LLM baseline (vanilla MHA, same architecture)
Builds on: multi-head attention mixture-of-experts TransNeXt DiT Megatron
Closest papers by meaning
- SAS: Simulated Attention Score (NeurIPS 2025 · similarity 0.88)
- Mixture-of-Top-$k$ Attention: Efficient Attention as Scalable Fast Weights (NeurIPS 2026 · similarity 0.87)
- Unveiling and Harnessing Hidden Attention Sinks: Enhancing Large Language Models without Training through Attention Calibration (ICML 2024 · similarity 0.87)
- A Statistical Theory of Gated Attention through the Lens of Hierarchical Mixture of Experts (NeurIPS 2026 · similarity 0.86)
- Improving Transformers with Dynamically Composable Multi-Head Attention (ICML 2024 · similarity 0.86)
- MoBA: Mixture of Block Attention for Long-Context LLMs (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.