Second-Order Muon Done Right: A Principled Marriage of Spectral Geometry and Curvature
2026-08-10 • Artificial Intelligence
Artificial Intelligence
AI summaryⓘ
The authors introduce GO-MUON, a new method that updates parameters exactly based on a special geometry matching the data, and reuses this geometry for several optimization steps. They show that under certain mathematical conditions, their update solves a weighted optimization problem exactly, regardless of how the weights are computed or updated. They specifically analyze how this works for softmax cross-entropy loss, studying the relationship between different factors used in optimization. Additionally, they find that updating the geometry every four steps balances computation and noise, making it a tradeoff rather than a way to reduce noise.
spectral geometryGO-MUONweighted spectral oraclesoftmax cross-entropymodel Fishergeneralized Gauss–Newtonoptimization stepsdata-dependent geometrylazy geometryparameter update
Authors
Tong Che
Abstract
Muon's polar update is exact for an unweighted spectral geometry. We introduce GO-MUON, which uses a matched data-dependent geometry and reuses it across several optimization steps. Conditioned on any positive-definite left and right maps, its raw update exactly solves the corresponding weighted spectral oracle; this statement is independent of how the maps are estimated or how recently they were refreshed. For softmax cross-entropy, we quantify when the observed-label backward factor approaches the model Fisher and generalized Gauss--Newton factor. We also show that four-step refresh nearly preserves the tracking delay of slowly changing geometry while increasing stationary factor noise, making lazy geometry a compute--statistics tradeoff rather than a denoising mechanism.