Technical Summary: OrScale – Orthogonalised Optimization with Layer-Wise Trust-Ratio Scaling

1. Problem Statement

The Muon optimizer has emerged as a practical tool for large-scale training, utilizing orthogonalised updates to fix the direction of matrix-valued steps via Newton–Schulz iterations (approximating the polar factor). However, Muon suffers from a critical limitation: while it optimizes the direction of updates, it leaves the magnitude controlled by a single, essentially global learning rate.

This global scaling is suboptimal because different layers (e.g., attention projections, MLP up/down-projections, embeddings) exhibit vastly different gradient distributions and effective curvatures. While the Moonlight variant introduced a static per-shape factor to equalize per-element RMS, it cannot adapt to the dynamic drift of layer scales during training.

Previous attempts to port Layer-wise Adaptive Rate Scaling (LARS) and Layer-wise Adaptive Moment Scaling (LAMB)—which successfully decoupled step sizes from global learning rates in SGD and Adam—to Muon have failed. The paper identifies that naive ports collapse due to three specific structural reasons:

1. Degenerate Denominators: Using the Frobenius norm of the orthogonalized direction ($\|Q_\ell\|_F$) as the trust-ratio denominator fails because it is nearly a constant determined by the matrix shape (rank), not the update magnitude.
2. Unit Mismatch & Clip Saturation: Using the raw momentum norm ($\|\tilde{M}_\ell\|_F$) introduces a unit mismatch (gradient units vs. weight units), causing the trust ratio to hit the upper clipping bound ($r_{max}$) almost universally, rendering the adaptation ineffective.
3. Decoupled Weight-Decay Runaway: Combining calibrated denominators with decoupled weight decay creates a positive feedback loop where weight norms grow uncontrollably (10–50$\times$), destabilizing training.

2. Methodology

The authors propose OrScale, a framework built on a single design principle: The denominator of a layer-wise trust ratio must measure the Frobenius norm of the real update direction the optimizer will subtract from the parameters.

Core Principle

Let $D_{\ell,t}$ be the actual parameter-space direction subtracted from weights $W_{\ell,t}$ before the global learning rate is applied. The trust ratio is defined as:
$$r_{\ell,t} = \frac{\|W_{\ell,t}\|_F}{\|D_{\ell,t}\|_F + \epsilon}$$
Crucially, $D_{\ell,t}$ must include both the weight decay term and the orthogonalized direction.

Two Variants

The framework yields two specific implementations:

1. OrScale (General/Vision Default):

   • Update Direction: $D_{\ell,t} = \lambda W_{\ell,t} + Q_{\ell,t}$ (where $\lambda$ is weight decay and $Q$ is the orthogonalized direction).
   • Shape Factor: $s_\ell = 1$ (no static scaling).
   • Calibration: None.
   • Weight Decay: Coupled directly into the trust-ratio-scaled update step.

2. OrScale-LM (LLM Default):

   • Update Direction: $D_{\ell,t} = \lambda W_{\ell,t} + s_\ell Q_{\ell,t}$.
   • Shape Factor: Adopts Moonlight's static factor $s_\ell = 0.2\sqrt{\max(m_\ell, n_\ell)}$ to match AdamW's per-element RMS.
   • Calibration: Introduces a lazy per-layer calibration constant ($c_{\text{denom}, \ell}$) computed once at initialization. This anchors the initial trust ratio to exactly 1 ($r_{\ell,0}=1$).
   • Significance of Calibration: This allows OrScale-LM to inherit the learning rate and weight decay hyperparameters tuned for Moonlight (and by extension, AdamW) without requiring a new hyperparameter sweep. It ensures width-invariance at initialization.

Algorithmic Details

• Coupled Weight Decay: Weight decay is integrated into the trust-ratio calculation ($\lambda W$ is part of the denominator) rather than applied as a separate pull-back. This prevents the runaway feedback loop observed in decoupled variants.
• Efficiency: The method adds negligible overhead (<1% wall-clock time), requiring only one additional scalar per matrix layer and two Frobenius norm calculations per step, which are dominated by the existing Newton–Schulz iterations.

3. Key Contributions

• Algorithm: A unified, parameter-light recipe (Algorithm 1) that successfully integrates layer-wise trust-ratio scaling into Muon. The authors identify the unique "workable corner" in the design space (Real-update-direction denominator + Coupled WD).
• Failure Analysis: A rigorous mapping of the Muon trust-ratio design space (Table 1) that analytically and empirically characterizes three failure modes (degenerate denominator, clip saturation, and runaway weight growth), demonstrating why they fail and why the proposed solution succeeds.
• Theory:
  • Convergence: Proves a nuclear-norm $O(1/\sqrt{T})$ nonconvex convergence guarantee.
  • Adaptive Gain: Establishes a layer-adaptive descent constant $\kappa_{\text{eff}} > 1$ under measurable layer heterogeneity, strictly improving over Muon's homogeneous constant.
  • Separation: Proves a strict separation from MuTrust/MuScale under conditions of clip saturation.
  • Calibration: Demonstrates that OrScale-LM preserves $\mu$P-style learning-rate transfer at initialization.
• Empirics: Comprehensive evaluation across vision and language model tasks.

4. Experimental Results

Vision (CIFAR-10 / DavidNet)

• Performance: OrScale achieved 94.05% $\pm$ 0.08 validation top-1 accuracy, ranking first among five baselines.
• Improvement: This represents a +0.35 point improvement over standard Muon (93.70%) and +0.30 over Muon+Moonlight (93.75%).
• Stability: OrScale demonstrated a wider stable learning rate window compared to Muon, consistent with theoretical predictions of layer-adaptive gains.
• Ablation: A failure-mode ablation study confirmed that variants corresponding to the identified failure modes (e.g., using $\|Q_\ell\|_F$ or raw momentum) exhibited the predicted symptoms (e.g., 99.5–100% clip saturation or weight-norm runaway) and failed to match OrScale's performance.

Large Language Models (FineWeb-Edu Pre-training)

• Scaling: Evaluated across four model scales (125M to 1.1B parameters).
• Comparison: OrScale-LM outperformed Muon+Moonlight at three of four scales (125M, 545M, 1.1B) and beat AdamW at every scale.
• Loss Reduction: At the 1.1B scale, OrScale-LM achieved a final validation cross-entropy of 2.6251, compared to 2.6360 for Muon+Moonlight and 2.7304 for AdamW.
• Efficiency: The advantage was preserved across the compute range, with OrScale-LM showing a clean win in the regime where compute matters most for LLMs.

5. Significance and Claims

The paper claims that OrScale successfully bridges the gap between the spectral geometry of Muon and the layer-wise adaptivity of LARS/LAMB. Its primary significance lies in:

1. Solving the Magnitude Problem: It provides the first time-varying, layer-wise magnitude control for Muon that is robust to structural failure modes.
2. Practical Transferability: Through the OrScale-LM variant and lazy calibration, it allows practitioners to transfer learning rates tuned for AdamW or Moonlight directly to OrScale-LM without re-tuning, lowering the barrier to adoption.
3. Theoretical Soundness: It offers a rigorous theoretical justification for why previous hybrid attempts failed and proves that the proposed method yields strictly better convergence constants under heterogeneity.
4. Empirical Validation: It demonstrates that these theoretical improvements translate to tangible gains in both small-scale vision tasks and large-scale LLM pre-training, establishing OrScale as a state-of-the-art optimizer for matrix-valued layers.

The authors maintain modesty regarding the LLM results, noting that the 1.1B scale advantage is comparable to single-seed noise at smaller scales, and thus the results should be viewed as directional evidence of the method's efficacy rather than a definitive statistical superiority in all regimes. However, the consistent improvement across scales and the successful ablation of failure modes strongly support the proposed design principle.