OrScale: Orthogonalised Optimization with Layer-Wise Trust-Ratio Scaling

Abstract:Muon fixes the \emph{direction} of every matrix-valued update at the polar factor of its momentum, while each layer's step \emph{magnitude} is addressed only by a static shape correction. We derive a dynamic per-layer scalar by adapting the LARS/LAMB trust-ratio principle to the orthogonalized setting, where the standard denominator candidates---the raw momentum norm or the polar-factor norm---either live in the wrong unit space or carry no update-scale information. The resulting method, \emph{OrScale}, uses the norm of the parameter-space direction actually applied and anchors each layer's ratio at one via a per-layer calibration, so that the Moonlight recipe (tuned for AdamW, shared with Muon via RMS matching) transfers with \emph{no additional sweep}; a component ablation confirms each design choice is individually load-bearing. Theoretically, OrScale retains a nuclear-norm $O(1/\sqrt{T})$ convergence rate for any clipped multiplier and achieves a strict layer-adaptive descent gain $\kappa_{\mathrm{eff}}>1$ under two conditions estimable from standard training diagnostics---a bound that predicts the gain should \emph{grow with architectural heterogeneity}. Experiments confirm the prediction: with every hyperparameter inherited verbatim from the Moonlight recipe, OrScale matches or beats Muon+Moonlight across dense 125M--1.1B FineWeb-Edu pre-training, and on a 16B-A3B mixture-of-experts model---where the logged trust ratios separate cleanly by layer class---the gap widens by an order of magnitude to $0.130$ nats ($3.8\%$ relative) at parity wall-clock cost.

Machine Learning (cs.LG); Computation and Language (cs.CL)

arXiv:2605.07815 [cs.LG]

(or arXiv:2605.07815v2 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2605.07815 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.