PaperlayerView exact revision on arXiv

Source-linked AI summary

MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration

Da Chang, Qiankun Shi, Lvgang Zhang, Yu Li, Ruijie Zhang, Yao Lu, Yongxiang Liu, Ganzhao Yuan

arXiv:2603.28254v2cs.LGstat.ML

TL;DR

MuonEq addresses the gap between post-orthogonalization scaling and computationally heavier whitening-based preconditioners for Muon. It equilibrates momentum before finite-step Newton–Schulz using RC, R, or C, and the default R variant retains Muon-type theory while consistently outperforming Muon in the reported LLaMA2 pretraining experiments.

  • Problem

    Existing Muon extensions either rescale updates after orthogonalization or use heavier whitening-based preconditioners before it, leaving lightweight pre-orthogonalization geometry correction to study.

  • Method

    MuonEq applies lightweight row/column equilibration to the momentum matrix before finite-step Newton–Schulz, using RC, R, or C, with R default for hidden weights.

  • Results

    MuonEq (R) consistently outperforms Muon across tested LLaMA2 pretraining model sizes and token budgets, with faster convergence and lower validation perplexity.

  • Takeaways & Limitations

    Row/column equilibration provides a lightweight zeroth-order whitening surrogate, while R offers the preferred geometry and standard Muon-type nonconvex stationarity guarantee for hidden matrix weights.

  • Takeaways & Limitations

    Comparisons involving row/column terms require caution because their meanings vary by layout and task.

Abstract

from arXiv · show

Orthogonalized-update optimizers such as Muon improve training of matrix-valued parameters, but existing extensions typically either rescale updates after orthogonalization or use heavier whitening-based preconditioners before it. We introduce {\method}, a lightweight family of pre-orthogonalization equilibration schemes for Muon with three forms: two-sided row/column normalization (RC), row normalization (R), and column normalization (C). By rebalancing the momentum matrix before finite-step Newton--Schulz orthogonalization, {\method} improves the geometry seen by orthogonalization. We show that finite-step orthogonalization is governed by the input spectrum, especially stable rank and condition number, and that row/column normalization acts as a zeroth-order surrogate for whitening. For hidden matrix weights, R is the default variant. Theoretically, {\method} (R) retains the standard $\widetilde{\mathcal O}(T^{-1/4})$ Muon-type nonconvex stationarity guarantee with decoupled weight decay and a horizon-free diminishing learning-rate schedule, and extends it to finite-step NS5 up to an explicit inexactness constant. In LLaMA2 pretraining on C4, {\method} (R) consistently outperforms Muon on 130M, 350M, and 1B models, with faster convergence and lower validation perplexity. The code is available at the \href{https://github.com/MaeChd/muon-eq}{MuonEq codebase}.

1 Introduction

MuonEq addresses the limitations of post-orthogonalization scaling and costly whitening by equilibrating Muon’s momentum matrix before finite-step Newton–Schulz orthogonalization. It uses lightweight RC, R, and C variants, with R selected for hidden matrix weights and supported by both theory and experiments.

  • Prior work: Post-orthogonalization methods rescale an already constructed update, while whitening-based pre-orthogonalization methods modify geometry with substantially higher state and computational overhead.MuonEq targets the same pre-orthogonalization opportunity with a lighter intervention.
  • Motivation: Finite-step Newton–Schulz behavior depends on the input spectrum, especially stable rank and condition number.This motivates improving matrix geometry before orthogonalization rather than focusing only on the exact polar factor.
  • Motivation: Classical row/column equilibration offers a lightweight approximation to optimal diagonal preconditioning and acts as a zeroth-order surrogate for whitening.Row and column normalization are also motivated as steepest-descent directions under corresponding operator geometries.
  • Contribution: MuonEq introduces RC, R, and C variants that rebalance the momentum matrix before Newton–Schulz using on-the-fly row and column squared norms.R is the default for hidden matrix weights, while RC and C are retained for analysis and ablation.
  • Theory: MuonEq (R) retains the standard e O(T^-1/4) nonconvex stationarity guarantee while accounting for decoupled weight decay, horizon-free schedules, and finite-step NS5.The theory explicitly extends the guarantee to the finite-step setting.
  • Empirical result: Across model sizes and token budgets, the default R variant consistently outperforms Muon in LLaMA2 pretraining.The supplied contribution passage reports this as the principal empirical finding.

2 Method

MuonEq inserts diagonal equilibration between momentum formation and orthogonalization, using RC, R, or C to reshape the matrix seen by Newton–Schulz. The method preserves Muon’s orthogonalization and optimizer state while exposing a trade-off between spectral correction and preprocessing bias.

  • Muon baseline: Muon’s practical update orthogonalizes momentum with a fixed number of Newton–Schulz iterations rather than exact polar decomposition.The momentum may be standard or Nesterov momentum, and the implementation uses finite-step orthogonalization.
  • Pre-orthogonalization equilibration: MuonEq inserts diagonal preconditioning before orthogonalization, producing an equilibrated momentum matrix that is then orthogonalized into the update.This changes the transient input to Newton–Schulz rather than the persistent optimizer state.
  • Equilibration modes: The mode s selects RC, R, or C, with row and column reductions computed from the current momentum and a stability term ε.R and RC use row reductions, while C and RC use column reductions.
  • Design distinction: MuonEq leaves Muon’s orthogonalization routine and persistent optimizer state unchanged, unlike methods that use matrix-valued or adaptive preconditioners.Its distinction from Muon+ and related methods is that it rebalances the input rather than rescales the output.
  • Trade-off: RC can provide stronger spectral correction, but preprocessing introduces a bias that must be traded against finite-step Newton–Schulz error.The decomposition separates approximation error after preprocessing from the bias caused by preprocessing itself.
  • Algorithm: MuonEq’s training loop forms momentum, applies diagonal preprocessing, computes NS5, and updates parameters with decoupled weight decay.The implementation includes optional Nesterov momentum and mode-dependent row/column scaling.
  • Default choice: R is the default because row-sided geometry is identified as natural for hidden Transformer weights, whereas column-sided geometry is more aligned with embeddings.RC and C remain available for analysis and ablation.

3 Analysis

The analysis links finite-step Newton–Schulz behavior to input spectral geometry and explains MuonEq normalization as a lightweight whitening surrogate. It then establishes convergence guarantees for RC and default R, including finite-step NS5 effects.

  • 3.1 Why spectral geometry matters: Stable rank controls the onset of Newton–Schulz’s linear-to-nonlinear transition, while condition number controls its width across singular directions.
  • 3.1 Why spectral geometry matters: Figure 1 compares finite-step relative Frobenius error and condition numbers before and after normalization on controlled random Gaussian matrices.Two-sided row/column normalization yields the smallest error and most consistent spectral compression.
  • 3.1 Why spectral geometry matters: Figure 2 evaluates module-wise median and mean orthogonalization errors on LLaMA2-130M training trajectories, with two-sided normalization decaying fastest.The figure reports results at 1%, 10%, 50%, and 100% of training with 25%–75% shaded ranges.
  • 3.2 Row/column normalization as a whitening surrogate: Row normalization removes marginal scale mismatch so the leading whitening correction acts on residual off-diagonal Gram error.This supports interpreting row/column normalization as a zeroth-order whitening surrogate rather than heuristic rescaling.
  • 3.3 Convergence analysis: MuonEq analyzes RC and default R separately: RC provides stronger two-sided correction, while R has cleaner one-sided geometry and the main stochastic guarantee.The framework uses row and column squared norms computed from the current momentum before Newton–Schulz.
  • 3.3 Convergence analysis: R retains the standard Muon-type eO(T^-1/4) stationarity rate while covering decoupled weight decay, a horizon-independent diminishing schedule, and finite-step NS5 through ε_ns.The NS5 result transfers the T^-1/4 dependence with explicit (1 ± ε_ns) constants; R needs no χ_ε-type condition.

4 Experiments

Experiments evaluate MuonEq on LLaMA2/C4 and controlled ablations, comparing model-scale behavior, resource metrics, and pre-orthogonalization variants. R consistently performs strongest among static variants and is less sensitive to Newton–Schulz iteration count.

  • Main Results: LLaMA2 experiments compare MuonEq and Muon variants with AdamW and MARS-AdamW across 130M, 350M, and 1B models on C4.Figure 3 reports training loss versus tokens, validation loss versus tokens, and validation loss versus wall-clock time; Table 1 summarizes validation perplexity, memory, and time.
  • Main Results: MuonEq and MuonEq-Nes consistently outperform Muon and Muon-Nes in the LLaMA2/C4 comparison.The reported comparison includes best validation perplexity, peak memory usage, and training time per step.
  • Ablation Study: R is the strongest static variant across CIFAR-10 with ResNet-18 and FineWeb with GPT2-small.Results averaged over three random seeds show higher CIFAR-10 test accuracy and lower FineWeb validation perplexity than Muon-Nes.
  • Ablation Study: At K = 5, Figure 4 evaluates validation perplexity and step time across Newton–Schulz iteration counts and learning-rate/momentum settings.The heatmaps use 2.6B tokens, while the iteration sweep uses 10.5B tokens.
  • Ablation Study: R remains stable across a wider range of Newton–Schulz iteration counts while maintaining comparable step time to Muon-Nes.For tested hidden-weight settings, R offers a favorable trade-off between spectral correction and preprocessing bias.

5 Conclusion

MuonEq concludes that lightweight equilibration improves the geometry presented to finite-step orthogonalization without full whitening. The default R variant combines a zeroth-order whitening interpretation with empirical gains and a Muon-type theoretical guarantee.

  • Conclusion: MuonEq reshapes Newton–Schulz inputs through lightweight RC, R, or C equilibration, with R used by default for hidden weights.The variants compute row or column squared norms from the current momentum without adding persistent optimizer state beyond Muon.
  • Conclusion: R provides cleaner one-sided geometry and supports the main O~(T^-1/4) stationarity guarantee, while RC provides stronger two-sided correction.The paper frames both as zeroth-order whitening surrogates rather than full whitening.
  • Conclusion: The paper situates MuonEq within orthogonalized-update, Muon theory, structured-preconditioning, and approximate-orthogonalization research.These areas include post-orthogonalization scaling, convergence guarantees, matrix-aware preconditioning, and faster polar routines.
  • Conclusion: The proof passages establish finite-step orthogonalization error bounds through singular-value recurrences, Frobenius-norm preservation, and case analysis.They conclude by combining cases and setting α = ∥G∥F.
  • Conclusion: The paper also uses Sylvester equations and Fréchet derivatives to analyze matrix square-root and inverse-square-root perturbations.The relevant lemmas establish existence, uniqueness, integral representations, and norm bounds.

D Proofs of Proposition 3.2

The proof analyzes column/right and row/left normalization through Sylvester equations associated with diagonal factors. Multiplication by the matrix factor and norm bounds then establish the corresponding perturbation controls.

  • Proof of Proposition 3.2: The column/right statement introduces Lc through the Sylvester equation DcLc + LcDc = DcCcDc.The proof then multiplies on the left by M and applies the Sylvester-equation bound.
  • Proof of Proposition 3.2: The row/left statement analogously introduces Lr through DrLr + LrDr = DrCrDr.The proof multiplies on the right by M before invoking the same bound.

E Normalization and whitening

The appendix formalizes how row, column, and two-sided normalization remove marginal scale mismatch before whitening or orthogonalization. First-order expansions show that the remaining correction depends on residual Gram errors, while assumptions establish bounded normalization factors for the momentum iterates.

  • One-sided expansions: For column-normalized inputs, finite-step orthogonalization is expanded around the identity column Gram matrix.The expansion uses a first-order correction in the column Gram error plus a second-order remainder.
  • One-sided expansions: For row-normalized inputs, the analogous expansion is centered on the identity row Gram matrix.The correction is proportional to the row Gram error, with higher-order terms controlled when that error is small.
  • Normalization–whitening estimates: Row and column normalization remove marginal scale terms before whitening, leaving correction of the residual Gram error.For one-sided normalization, the leading correction is governed by the remaining Gram mismatch.
  • Two-sided normalization: Two-sided normalization centers whitening at the identity Gram matrix and leaves only residual mismatch after diagonal equilibration.It no longer needs to correct marginal row or column scales.
  • Technical bounds: The appendix assumes bounded momentum entries and derives bounded row and column normalization factors for the stochastic analysis.The bounds are obtained from the momentum recurrence and row- and column-wise squared norms.

G Proofs of Proposition 3.4

These proofs establish deterministic properties of row-normalized inputs and finite-step NS5. Row normalization controls the update geometry, while the NS5 approximation preserves singular directions and incurs a trajectory-wise inexactness bounded below one.

  • Proposition 3.4: Row normalization makes every nonzero row of the preconditioned momentum have unit Euclidean norm, while zero rows remain zero.The Moore–Penrose pseudoinverse preserves zero rows during normalization.
  • Proposition 3.4: The normalized update has a positive alignment with the original momentum, bounded by 1/√m times its Frobenius norm.The proof sums rowwise contributions after normalization.
  • NS5 error: The trajectory-wise initial NS5 polar-approximation error is strictly below one after pre-scaling.The finite horizon makes the maximum error over the trajectory remain below one.
  • NS5 structure: NS5 preserves the singular vectors of its input and produces a diagonal singular-value transformation of the exact polar factor.For row-normalized trajectory inputs, the transformed singular values remain within the NS5 inexactness interval.
  • Update control: The proofs combine row-normalized update bounds with decoupled weight decay under a restriction on the decay parameter.The resulting inequalities support the subsequent nonconvex convergence analysis.

I Proofs of Theorem 3.5 and Corollary 3.6

The convergence proof controls the row-normalized NS5 update through momentum error, descent, and stochastic-gradient bounds. It yields the finite-step result while recovering the exact-polar and zero-weight-decay cases as specializations.

  • Theorem 3.5: The finite-step guarantee includes an explicit NS5 inexactness constant alongside the horizon-dependent convergence terms.The displayed bound contains C_ns terms and a T^(-1/4) contribution.
  • Theorem 3.5: The proof decomposes the stochastic update into gradient, momentum, and NS5 approximation errors.The argument applies descent inequalities and controls the resulting terms using the preceding lemmas.
  • Theorem 3.5: Summing the descent inequality over T iterations produces a bound on cumulative expected gradient norms.The proof uses the lower bound f(X_T+1) ≥ f⋆ before substituting the auxiliary estimates.
  • Corollary 3.6: Setting ε_ns to zero recovers the exact-polar result, and additionally setting ρ = 0 recovers the zero-weight-decay proof.These specializations connect the finite-step theorem to the corresponding idealized analyses.

J More Results

The additional experiments examine learning-rate and hyperparameter sensitivity, spectral diagnostics, and NS5 error decomposition across model settings. They show stronger spectral correction from RowColNorm but support row normalization as the practical default for hidden weights because two-sided correction also increases preprocessing bias.

  • Pre-NS momentum analysis: RowColNorm consistently improves singular-value entropy and stable rank over Direct, ColNorm, and RowNorm in pre-NS momentum matrices.The diagnostics cover multiple Transformer modules, layers, and training stages.
  • NS5 bias decomposition: At 1% and 10% of training, RowColNorm is usually lowest or near-lowest in NS5 approximation error but has the largest preconditioning bias.By 50% and 100%, approximation-error differences shrink while the additional two-sided bias remains visible.
  • Sensitivity analysis: MuonEq-Nes with row normalization shows a broader low-perplexity region than Muon-Nes across learning-rate and momentum sweeps, especially at K = 5.This indicates reduced sensitivity to hyperparameter choice in the reported GPT2-small/FineWeb experiment.
  • Experimental setup: The LLaMA2 comparisons keep architecture, dataset, and training recipe fixed across AdamW, MARS-AdamW, Muon, Muon-Nes, and MuonEq variants.The main comparisons use LLaMA2-130M, 350M, and 1B trained on C4.

K.2 Ablation Study Details

The ablation compares RC, R, and C under matched optimizer settings to isolate pre-orthogonalization geometry. R is the default for hidden matrix weights, while matrix conventions explain apparent differences with prior normalization results.

  • Experimental setup: Ablations use CIFAR-10 with ResNet-18 and FineWeb pretraining of GPT2-small up to 10.5B tokens.Experiments run on 4 RTX Pro6000 GPUs, with results averaged over three random seeds.
  • Variants: RC, R, and C apply two-sided, row, and column normalization before orthogonalization, respectively, while leaving optimizer states and scalar hyperparameters unchanged.R is the default variant.
  • Evaluation: The protocol keeps the training budget and evaluation procedure fixed, reporting CIFAR-10 test accuracy and FineWeb validation perplexity.This isolates the effect of the diagonal map before orthogonalization.
  • Geometry comparison: RC targets stronger two-sided spectral correction, whereas R and C test the two one-sided geometries under the same optimizer configuration.R is more relevant for hidden matrix weights, while C is its column-sided companion.
  • Matrix conventions: Normalization direction depends on matrix layout: column normalization under θ ∈ R^din×dout corresponds to row normalization of stored W = θᵀ ∈ R^dout×din.This convention explains why prior row-normalization findings must be compared carefully.

L Limitations

MuonEq targets matrix-valued hidden weights, but its default row-normalized form is not claimed optimal for every parameter block or architecture. The auxiliary RC convergence analysis is narrower than the practical RC experiments, while downstream risks remain tied to model and dataset use.

  • Scope: MuonEq’s row-normalized default should not be treated as universally optimal across parameter blocks or architectures.Broader implementation variants and training regimes remain open directions.
  • Theory boundary: The RC convergence guarantee assumes exact polar updates, disabled Nesterov momentum, no decoupled weight decay, and sufficiently large ε.It therefore does not formally cover practical RC ablations using finite-step Newton–Schulz and ε = 10^-8.
  • Societal scope: MuonEq introduces no new data or deployment pipeline, so its direct societal risks are inherited mainly from downstream models and datasets.Lower training costs could support both beneficial and harmful model development, requiring downstream safeguards.
Loading 2603.28254v2…