When and Why Grouping Attention Heads Accelerates Muon Optimization

Hongtao Zhang, WenJie Zhou, Wei Chen, Xueqi Cheng

Keywords: Optimization, Training Dynamics, Muon Optimizer

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.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 38