Muse: Representation Geometry of Muon Beyond Normalized Momentum

Muon-style optimizers have emerged as powerful alternatives to AdamW for training large language models by applying a polar map to matrix-valued momentum. However, because the polar map is computed after reshaping parameter blocks into matrices, the choice of representation (e.g., native layout vs. flattened vectors) is not merely a relabeling but a fundamental design choice that alters the optimizer's geometry. This paper investigates how different matrix representations influence the polar update and overall training performance.

The authors introduce Muse, a family of Muon-style optimizers that share the same momentum rule and Newton-Schulz backend but vary the matrix representation of parameter blocks. The representations tested include the native layout, nearest-square flattening, various 'skinny' layouts (reducing the shorter dimension), and the vector endpoint (which reduces Muon to normalized SGD with momentum). The researchers analyze the geometry induced by these representations, specifically focusing on how the shorter matrix dimension (the 'short side') dictates singular channel support, pullback scaling, and convergence bounds. They also utilize a teacher-student model to study curvature collapse and the spectral profile of these representations.

The study establishes that the representation provided to the polar map is a critical component of Muon's geometry. The authors prove that the short side of the represented matrix determines the rank of the polar update and the worst-case stationarity constants. In controlled experiments, they find that representations with larger short sides (like the native or nearest-square layouts) provide higher nuclear support, leading to more effective optimization. As the short side is reduced, the optimizer's behavior increasingly resembles normalized momentum, which the authors show is suboptimal for these tasks. Pretraining experiments on LLaMA2-130M and 600M models confirm that balanced, non-native representations can match the performance of the native layout, while aggressive short-side reduction degrades performance.