Muse: Representation Geometry of Muon Beyond Normalized Momentum

2026-07-15

AI Summary

This work addresses the unclear influence of parameter representation on the optimization geometry of Muon-style optimizers, which depend on the form of parameter blocks prior to orthogonalization. The authors incorporate parameter representation into the geometric design of the optimizer and propose the Muse family, unifying momentum rules and the Newton–Schulz backend across diverse representations—native, near-square, elongated, and vectorized—and revealing an induced polar-coordinate steepest descent geometry. Theoretical analysis demonstrates that Frobenius isometric representations govern the number of singular channels, retraction scaling, and non-convex stochastic convergence bounds via the short dimension, while linking the ratio of nuclear to squared Frobenius norms with early-stage dissipation. Experiments on LLaMA2-130M/600M pretraining show that balanced non-native representations match native performance, whereas shortening the short dimension degrades scaling capacity and singular channel support, yielding behavior akin to normalized momentum.

Abstract

Muon-style optimizers apply a polar map to matrix momentum, but their updates also depend on the representation of each parameter block before orthogonalization. We study this representation choice as a form of optimizer geometry and introduce {\method}, a family of Muon-style optimizers that shares the same momentum rule and Newton--Schulz backend across native, nearest-square, skinny, and vector representations. Each Frobenius-isometric representation induces a distinct polar steepest-descent geometry, in which the shorter matrix dimension determines the number of supported singular channels, the pullback scaling, and the constants in stochastic nonconvex convergence bounds. In a teacher--student model, curvature collapse and an isotropic Marchenko--Pastur spectral profile connect early-stage dissipation to the represented nuclear-to-squared-Frobenius norm ratio. Pretraining experiments on LLaMA2-130M and LLaMA2-600M, together with fixed-momentum diagnostics, show that balanced non-native representations can match the performance of the native representation, whereas reducing the shorter dimension weakens the scaling and singular-channel support, leading to behavior that increasingly resembles normalized momentum.

Problem

• optimizer geometry
• representation choice
• polar steepest-descent
• singular channels
• normalized momentum

Innovation

• optimizer geometry
• polar map
• Frobenius-isometric representation
• singular channels
• momentum normalization