# When and Why Grouping Attention Heads Accelerates Muon Optimization

arXiv: https://arxiv.org/abs/2605.08933 · HTML: https://mlsanity.com/paper/2605.08933/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 8 of its 14 references that are in the map.
- First posted on arXiv 9 May 2026.

## Builds on (most cited first, 8 of 8)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Dion: Distributed Orthonormalized Updates](https://mlsanity.com/paper/2504.05295/) (2025)
- [MuonBP: Faster Muon via Block-Periodic Orthogonalization](https://mlsanity.com/paper/2510.16981/) (2025)
- [Beyond the Ideal: Analyzing the Inexact Muon Update](https://mlsanity.com/paper/2510.19933/) (2025)
- [Dion2: A Simple Method to Shrink Matrix in Muon](https://mlsanity.com/paper/2512.16928/) (2025)

## Built on by (most cited first, 1 of 1)

- [Hierarchical Muon: Tiled Newton-Schulz Updates for Efficient Muon Optimization](https://mlsanity.com/paper/2606.27216/) (2026)

## Abstract (arXiv)

Muon orthogonalizes matrix updates, but multi-head attention naturally operates at the level of heads. This granularity mismatch raises the question of whether Muon should be applied to the full attention projection, to individual heads, or to intermediate head groups. We study this question through a one-step descent comparison between full-matrix Muon and group-wise Muon. Our analysis reveals a trade-off between the group-wise whitening gain from group-wise updates and the grouping-induced norm cost, an additional update-norm cost caused by replacing full-matrix whitening with group-wise whitening. Motivated by this trade-off, we propose Group Muon, which treats head group size and grouping rule as optimizer hyperparameters. On GPT-2 Small trained on FineWeb, appropriate grouping improves validation loss over both full-QKV Muon and fully head-wise MuonSplit.