When and Why Grouping Attention Heads Accelerates Muon Optimization

Hongtao Zhang; Wenjie Zhou; Wei Chen; Xueqi Cheng

cs.LGMay 9, 2026

When and Why Grouping Attention Heads Accelerates Muon Optimization

Authors: Hongtao Zhang, Wenjie Zhou, Wei Chen, Xueqi Cheng

Organizations: State Key Laboratory of AI Safety, Institute of Computing Technology, Chinese Academy of Sciences · School of Advanced Interdisciplinary Sciences, University of Chinese Academy of Sciences · University of Chinese Academy of Sciences

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.

Explore similar work

Muown: Row-Norm Control for Muon Optimization
May 11, 2026Kai Lion, Florian Hübler, Bingcong Li +2Muon Optimizer Muon
Reassessing Muon for Matrix Factorization
Jul 14, 2026Ali Parviz, Gal Mishne, Alex CloningerMuon Optimizer Muon
Muown Implicitly Performs Angular Step-size Decay
Jun 22, 2026Florian Hübler, Kai Lion, Antonio Orvieto +1Muon Optimizer Muon

May 11, 2026 · cs.LGJ/K move · Enter open · S save

Muown: Row-Norm Control for Muon Optimization

Kai Lion, Florian Hübler, Bingcong Li, Antonio Orvieto+1

Department of Computer Science, ETH Zurich, Switzerland · Department of Mathematics, School of Computation, Information and Technology; Technical University of Munich, Germany · ELLIS Institute Tübingen, MPI-IS, Tübingen AI Center, Germany

Muon has emerged as a strong competitor to AdamW for language model pre-training, yet its behavior at scale is sensitive to weight decay. Recent work has observed that, for Muon without decoupled weight decay, the spectral norm of weight matrices drifts upward over training. Through a decomposition of the spectral norm into a row-magnitude factor and a row-coherence factor, we identify the former as the empirical driver of this drift under Muon, while the latter remains well-behaved along the trajectory. Motivated by this diagnosis, we introduce Muown, a drop-in replacement for Muon that treats the row-magnitude vector as an explicit optimizer variable, updating it under the

\ell_\infty

geometry induced by the decomposition, while applying Muon unchanged to the remaining direction component. We prove that Muown attains the optimal non-convex rates in both deterministic and stochastic regimes under a dual norm aligned with the underlying geometries and with a stochastic noise coefficient that empirically remains below that of Muon throughout training. Across GPT-style pre-training on FineWeb-Edu with model sizes from 124M up to 2.7B parameters, Muown improves perplexity over Muon, SOAP, AdamW, and Lion. It also widens the plateau of near-optimal learning rates across model scales, reduces sensitivity to weight decay, and avoids the spectral norm drift at negligible step-time overhead when appropriately sharded.

Muon Optimizer Muon Model Weights Optimizer Design

Open paper Ask about it PDF

Abstract

Explore similar work

Muown: Row-Norm Control for Muon Optimization

Reassessing Muon for Matrix Factorization

Muown Implicitly Performs Angular Step-size Decay