Reproduce “When and Why Grouping Attention Heads Accelerates Muon Optimization” on OpenResearch

About this paper

Hongtao Zhang, Wenjie Zhou, Wei Chen, Xueqi Cheng

This work introduces a "gain–cost trade-off" criterion to analytically determine the optimal granularity for applying the Muon optimizer to multi-head attention mechanisms. Researchers at the Chinese Academy of Sciences propose Group Muon, an algorithm that applies orthogonalization to intermediate groups of attention heads, achieving a 0.18% lower validation loss on GPT-2 Small compared to a full-matrix Muon baseline.

Key insights

• The optimal granularity for Muon application is determined by a balance between "group-wise whitening gain" (first-order benefit from increased nuclear norm sum) and "grouping-induced norm cost" (second-order cost from increased Frobenius norm of whitened updates).
• The most effective grouping granularity changes throughout training; finer grouping is generally better in early stages, while moderately coarse grouping performs better in middle and late stages.
• Random grouping of Query and Key projections with intermediate group sizes often yields better performance than adjacent or interval grouping, suggesting benefits from decorrelating head interactions.

Summary by alphaXiv.