Search papers, labs, and topics across Lattice.
This paper investigates the impact of representation choices in Muon-style optimizers, focusing on how different parameter block representations affect the optimizer's geometry and convergence properties. By introducing a family of Muon-style optimizers that maintain a consistent momentum rule and backend across various representations, the authors reveal that Frobenius-isometric representations lead to distinct polar steepest-descent geometries. Key findings indicate that balanced non-native representations can achieve performance parity with native representations, while reducing the shorter dimension diminishes scaling and singular-channel support, ultimately resembling normalized momentum behavior.
Balanced non-native representations can match the performance of native representations in Muon-style optimizers, challenging assumptions about representation necessity.
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.
