SECOND-ORDER MUON DONE RIGHT: A PRINCIPLED MARRIAGE OF SPECTRAL GEOMETRY AND CURVATURE

GO-MUON addresses the limitations of standard Muon by incorporating a matched spectral geometry. While Muon applies a polar update to momentum, GO-MUON generalizes this by using symmetric positive-definite (SPD) coordinate maps (derived from activation and backpropagated-gradient second moments) to transform the update direction. By applying these maps both before and after the polar operation, the method ensures that the update remains an exact solution to a weighted spectral oracle. To maintain computational efficiency, GO-MUON employs a 'lazy' geometry strategy, refreshing these factors only every four steps while continuing to update gradient and momentum information at every step.

• Matched Spectral Oracle: The paper proves that by using matched left and right maps, the update direction is the exact support direction of the weighted operator-norm ball, regardless of how the maps are estimated.
• Factor Convergence: Under softmax cross-entropy, the observed-label backward factors are shown to approach the model Fisher and generalized Gauss-Newton (GGN) factors as the predictive distribution aligns with the data conditional.
• Compute-Statistics Tradeoff: The four-step refresh schedule is characterized as a tradeoff between stationary factor noise and computational cost. It preserves the tracking delay of slowly changing geometry while reducing mean step time by approximately 20%.
• Empirical Performance: In experiments on Penn Treebank and Tiny Shakespeare, GO-MUON outperformed standard Muon in training loss trajectory. Furthermore, in modular addition tasks, GO-MUON reached high held-out accuracy significantly faster than the control.

Standard second-order or preconditioned optimization methods often struggle with the high computational cost of frequent curvature estimation. GO-MUON provides a mathematically rigorous framework for integrating curvature information into the Muon update pipeline. By demonstrating that one can cache and reuse geometric factors without losing the benefits of a principled spectral oracle, this work offers a practical path for scaling sophisticated optimization techniques to large-scale language models without the prohibitive overhead of per-step curvature computation.