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

TL;DR

OrScale enhances neural network training by orthogonalizing updates with layer-wise trust ratios, improving convergence guarantees and empirical performance on image and language models.

Contribution

Introduces OrScale, a trust-ratio extension of Muon that stabilizes layer updates and provides theoretical guarantees and empirical improvements.

Findings

OrScale achieves state-of-the-art results on CIFAR-10 with CIFAR-10/DavidNet.

OrScale-LM outperforms Muon+Moonlight and AdamW on language model pre-training.

Theoretically, OrScale guarantees nonconvex convergence with layer-adaptive descent.

Abstract

Muon improves neural-network training by orthogonalizing matrix-valued updates, but it leaves each layer's update magnitude controlled mostly by a global learning rate. We introduce OrScale, a trust-ratio extension of Muon built on a simple rule: the denominator of a layer-wise ratio should measure the Frobenius norm of the actual parameter-space direction that will be applied. This yields OrScale for general matrix layers and OrScale-LM for language models, where Moonlight shape scaling is combined with one-time per-layer calibration so every trust ratio starts at one. We analyze why three natural Muon-LAMB hybrids fail through shape-degenerate denominators, raw-momentum clip saturation, and decoupled weight-decay runaway, and show that the real-update-direction denominator with coupled weight decay avoids these failures. Theoretically, OrScale admits an O(1/sqrt(T)) nonconvex convergence guarantee in a nuclear-norm criterion, a strict layer-adaptive descent gain under measurable layer heterogeneity, and calibration properties that preserve muP-style learning-rate transfer at initialization. Empirically, OrScale ranks first on CIFAR-10/DavidNet across three seeds, improving Muon from 93.70% to 94.05% validation top-1, and OrScale-LM improves FineWeb-Edu pre-training versus Muon+Moonlight at three of four scales from 125M to 1.1B parameters while outperforming AdamW at every scale.