

























Abstract: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 \textbf{group-wise whitening gain} from group-wise updates and the \textbf{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 \textbf{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.
| Comments: | 16 pages, 4 figures |
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2605.08933 [cs.LG] |
| (or arXiv:2605.08933v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2605.08933 arXiv-issued DOI via DataCite (pending registration) |
From: Wenjie Zhou [view email]
[v1]
Sat, 9 May 2026 13:06:31 UTC (813 KB)
此内容由惯性聚合(RSS阅读器)自动聚合整理,仅供阅读参考。 原文来自 — 版权归原作者所有。