PaperlayerView exact revision on arXiv

Source-linked AI summary

The Newton-Muon Optimizer

Zhehang Du, Weijie Su

arXiv:2604.01472v1math.OCcs.AIcs.LG

TL;DR

Muon’s matrix-gradient orthogonalization lacks a clear design principle, motivating a tractable Newton-style explanation. The paper introduces a triplet quadratic surrogate and derives Newton–Muon with input-second-moment right preconditioning; on a reproduced Modded-NanoGPT run, it reaches the target validation loss in 6% fewer steps and cuts wall-clock time by about 4%.

  • Problem

    Muon performs strongly, but the theoretical mechanism underlying its matrix-gradient orthogonalization remains largely mysterious, while existing isotropic-curvature models are too unspecified for practical LLM optimization.

  • Method

    The paper introduces a triplet quadratic surrogate using gradient G, output-space curvature H, and activation matrix Z, then derives Newton–Muon under an isotropic displacement proxy.

  • Results

    6% fewer steps and about 4% lower wall-clock time let Newton–Muon reach the Muon baseline target validation loss in reproduced Modded-NanoGPT speedrun configurations.

  • Takeaways & Limitations

    Standard Muon can be interpreted as an implicit Newton-type method that omits right preconditioning by the input second moment, while Newton–Muon incorporates that data geometry.

  • Takeaways & Limitations

    Newton–Muon relies on an isotropic displacement proxy that may discard useful weight-distribution information and become inaccurate later in training; the curvature approximation also omits token coupling and uses shared curvature across samples and positions.

Abstract

from arXiv · show

The Muon optimizer has received considerable attention for its strong performance in training large language models, yet the design principle behind its matrix-gradient orthogonalization remains largely elusive. In this paper, we introduce a surrogate model that not only sheds new light on the design of Muon, but more importantly leads to a new optimizer. In the same spirit as the derivation of Newton's method, the surrogate approximates the loss as a quadratic function of the perturbation to a weight matrix $W$ using only three matrices: the gradient $G$, an output-space curvature matrix $H$, and the data matrix $Z$ that stacks the layer inputs. By minimizing this surrogate in one step and adopting a certain isotropic assumption on the weights, we obtain the closed-form update rule (up to momentum and weight decay) $W \leftarrow W - η\cdot \mathrm{msgn}(G(ZZ^\top)^{-1})$, where $η$ is the learning rate and $\mathrm{msgn}(X)=UV^\top$ if $X=USV^\top$ is a compact singular value decomposition. This new optimization method, which we refer to as Newton-Muon, shows that standard Muon can be interpreted as an implicit Newton-type method that neglects the right preconditioning induced by the input second moment. Empirically, on a reproduction of the earliest publicly released Modded-NanoGPT speedrun configuration using Muon for GPT-2 pretraining, Newton-Muon reaches the target validation loss in 6\% fewer iteration steps and reduces wall-clock training time by about 4\%.

1 Introduction

The paper develops a triplet quadratic surrogate that clarifies Muon’s matrix-gradient orthogonalization and derives Newton–Muon by incorporating input-data geometry through right preconditioning. On the Modded-NanoGPT reproduction, Newton–Muon reaches the target validation loss in 6% fewer iterations and reduces wall-clock time by about 4%.

  • Motivation: Muon’s theoretical mechanism remains largely mysterious despite strong empirical performance across neural-network and language-model training.This motivates an implicit Newton-type explanation with per-iteration cost comparable to AdamW or Muon.
  • Surrogate model: The triplet quadratic surrogate approximates loss changes using the gradient G, output-space curvature H, and activation-input matrix Z.Its linear term captures first-order behavior, while its quadratic term has a Kronecker-factored curvature form related to K-FAC.
  • Derivation: Under an isotropic displacement assumption, one-step minimization yields a Newton-type direction proportional to msgn(G(ZZ^T)^-1).The assumption treats the displacement between current and locally optimal weights as directionally unbiased, enabling a closed-form update despite unknown curvature H.
  • Newton–Muon: Newton–Muon right-preconditions G with the inverse input second moment, whereas standard Muon omits this factor when ZZ^T is approximately isotropic.The paper argues that this distinction incorporates anisotropic input-data geometry into the update; implementation uses periodically refreshed estimates and cached damped inverses.
  • Empirical result: 6% fewer iterations and about 4% lower wall-clock time let Newton–Muon reach the Muon baseline’s final validation loss on the reproduced Modded-NanoGPT speedrun.The right-preconditioning adds a 1.8% per-step cost in the Figure 1 comparison.

2 Derivation of Newton–Muon

The paper derives Newton–Muon from a triplet quadratic surrogate that combines gradient, curvature, and layer-input information. Under an isotropic displacement proxy, the resulting update right-preconditions the gradient by the inverse input second moment, while standard Muon is recovered for isotropic activations.

  • Triplet Quadratic Surrogate: The triplet quadratic surrogate approximates loss changes using the gradient G, curvature matrix H, and activation matrix Z.Its quadratic term retains the matrix structure of the update and captures interactions among H, Z, and G.
  • Newton Derivation: Minimizing the surrogate connects its optimizer to a Newton direction through a Kronecker-factored approximation of parameter-space curvature.The analysis relates the surrogate minimizer to the displacement from the current weights to a nearby reference optimum.
  • Newton–Muon Update: Newton–Muon is implemented by periodically estimating and inverting the input second moment, while reusing the cached inverse between refreshes.The update is always a descent direction under the stated positive-definiteness assumptions.
  • Newton–Muon Update: An isotropic proxy for the displacement second moment yields an observable closed-form direction based on msgn(G(ZZᵀ)^−1).This proxy is explicitly an approximation for the unknown displacement covariance and need not hold literally when dimensions make it rank-deficient.
  • Relation to Muon: Standard Muon is recovered when ZZᵀ is proportional to the identity, whereas anisotropic activations motivate Newton–Muon’s right preconditioning.Right equivariance also depends on transforming Z alongside a rotated gradient; keeping Z fixed generally breaks it.

3 Convergence Analysis of Newton–Muon: Case Study

The case study analyzes Newton–Muon in a single-spike quadratic model where activation anisotropy permits explicit mode-wise dynamics. Its contraction rate is independent of the spike condition number, unlike gradient descent and Muon.

  • Single-Spike Model: The analysis restricts attention to a quadratic model with one spiked eigendirection of ZZᵀ and an isotropic complementary subspace.The structured ground-truth matrix and initialization make the dynamics analytically tractable.
  • Mode-Wise Dynamics: Gradient descent and Muon multiply the displacement by ZZᵀ, so activation anisotropy distorts their effective updates in the spiked model.The case study isolates this distortion through the condition number κ of the activation second moment.
  • Mode-Wise Dynamics: The three optimizers preserve the mode decomposition, reducing matrix updates to scalar recursions for the coefficients.The spike affects only the mixed mode involving the spiked direction and its non-spike component.
  • Convergence Rates: Newton–Muon reaches error ε in O(log(r0/ε)) iterations, while gradient descent and Muon require O(κ log(r0/ε)) iterations under greedy step selection.All three methods decrease the scalar bound geometrically, but only Newton–Muon’s rate is independent of κ.

4 One-Step Analysis of Newton–Muon

The one-step analysis compares update directions using a scale-invariant surrogate score and numerical simulations under spiked activations. Newton–Muon scores higher than Muon when activation and curvature anisotropy are strong and lies closer to Newton’s direction.

  • One-Step Score: The analysis scores each update direction by its best predicted one-step loss decrease after a one-dimensional line search.This produces a scale-invariant criterion for comparing directions.
  • Numerical Study: The numerical study compares gradient descent, Muon, Newton–Muon, and Newton variants using exact SVD or five Newton–Schulz iterations.The simulations use square matrices with specified curvature and spiked activation settings.
  • Numerical Study: Newton–Muon achieves substantially higher scores than Muon when activation and curvature anisotropy are both strong.It is closer to the optimal Newton direction in the reported simulations.
  • Numerical Study: Muon substantially outperforms gradient descent when the curvature is anisotropic.The reported comparisons average results over 1024 independent simulations and include 2.5%–97.5% intervals.

5 LLM Experiments

The LLM experiments compare Newton–Muon with Muon and AdamW across reproduced Modded-NanoGPT benchmark settings, plus ablations and quadratic-score diagnostics. Newton–Muon generally improves over Muon, although gains vary by track and configuration.

  • Benchmark setup: The benchmark evaluates optimizers on short and medium Modded-NanoGPT tracks targeting validation losses of 3.28 and 2.92.Comparisons use publicly logged records and controlled optimizer replacements, with learning-rate tuning when needed.
  • Benchmark setup: The baseline short-track experiment uses Record #4 on one H100, training a 124M GPT-2 model with 3.25B tokens.Newton–Muon uses learning rate 0.0040, β = 0.95, γ = 0.2, and refresh interval k = 32.
  • Ablations: Moderate refresh intervals k ∈ {16, 32, 64} perform best, while very frequent refreshes k ∈ {1, 2} underperform.The refresh ablation varies k and β while fixing γ = 0.2 and learning rate 0.0040.
  • Ablations: Ridge scaling stabilizes the right preconditioner, with γ ∈ [0.02, 0.5] performing similarly and large γ values moving results toward the Muon baseline 3.2793.The learning-rate sweep is relatively flat near the best region, and the best Newton–Muon learning rate is close to Muon’s.
  • Short-track results: On short-track Record #28 using one L40S, Newton–Muon shows a small time-to-loss advantage despite slightly higher total runtime from preconditioning overhead.The comparison uses second-best runs and plots validation loss against both steps and wall-clock time.
  • Medium-track results: On medium-track Record #17, Newton–Muon improves final validation loss only marginally, so wall-clock time is not reported as substantially improved.This setting trains a roughly 455M-parameter model on about 3.12B tokens.
  • Quadratic score: Quadratic-score diagnostics rank Newton–Muon highest, Muon intermediate, and the raw gradient worst across four matrices and two checkpoints.Newton–Muon usually reduces curvature more than alignment relative to Muon, explaining its higher score; NS5 and NS32 are close.
  • Quadratic score: The measured activation second moments are strongly anisotropic, with large diagonal spread, substantial off-diagonal mass, and very large condition numbers.These diagnostics support evaluating input-geometry-aware preconditioning rather than assuming isotropic activations.

6 Discussion

Newton–Muon is introduced through a triplet quadratic surrogate and empirically improves Muon on reproduced Modded-NanoGPT configurations, while the framework has several stated limitations. These concern the isotropic weight proxy, curvature approximation, inverse computation, and single-GPU evaluation scope.

  • The triplet quadratic surrogate provides a local second-order view of Muon and leads to the Newton–Muon optimizer.
  • 6% fewer steps and over 4% lower wall-clock time were reported for reaching target validation loss versus the Muon baseline.
  • Newton–Muon relies on an isotropic weight proxy that may discard displacement information and become inaccurate later in training.
  • The Kronecker curvature approximation omits token coupling and assumes one shared output-space curvature matrix across samples and positions.
  • Structured or polynomial approximations could reduce the memory and computation needed for activation second-moment inversion.
  • Newton–Muon remains to be evaluated in distributed training because the experiments used only a single GPU.

A LLM Experimental Details

The LLM experiments reproduce benchmark configurations while changing the optimizer under controlled settings. Implementation uses layerwise activation second moments, precision-sensitive Cholesky inversion, symmetry-aware computation, and custom Triton kernels.

  • The comparison protocol keeps non-optimizer training settings identical within each reproduced benchmark and reports wall-clock times from the authors’ environment.
  • Activation second moments are formed layerwise, with a block-diagonal approximation used for the MLP contraction to avoid a full 4d × 4d matrix.
  • Accumulation can use bfloat16, but inverse computation and right-preconditioned multiplication are performed in float32 because they are more sensitive.
  • Symmetry of activation second moments and their polynomial functions allows computation of one triangle and reconstruction of the other.
  • The implementation uses custom Triton kernels because PyTorch lacks a simple interface for enforcing triangular computation.
  • Damped Cholesky factorization computes K_γ^-1 from K_γ = ZZ^⊤ + γI_n through K_γ = LL^⊤ and K_γ^-1 = L^-⊤L^-1.

B.3 Polynomial Iteration Inverse

The polynomial inverse constructs an explicit approximation to a damped matrix inverse by repeatedly transforming a scaled residual with matrix polynomials. Interval contraction, coefficient constraints, padding, and noise margins provide a certified approximation strategy.

  • Algorithm 2 takes a positive semidefinite K, damping γ, and SYPP budget B, then outputs an explicit approximation to (K + γI_n)^−1.
  • The method scales K_γ, initializes a residual R_0, and builds X ≈ (I_n − R_0)^−1 as an approximation to K_γ^-1.
  • Each iteration applies a polynomial q_k(R_{k−1}) to update the inverse estimate and transform the residual through R_k = ϕ_k(R_{k−1}).
  • The residual spectrum contracts over iteration-dependent intervals, enabling later polynomials to target tighter ranges and stronger contraction.
  • Polynomial choices minimize worst-case residual bounds under coefficient-magnitude constraints and robustness margins for finite precision and modeling error.
  • Under the stated contraction conditions, the residual norm is bounded by s_k and converges toward zero, yielding X_T close to (I_n − R_0)^−1.
  • Interval padding widens the design domain, while NOISE_ABS models worst-case perturbations in the scalar residual map.

C CIFAR-10 Experiment

The CIFAR-10 experiment compares AdamW, Muon, and Newton–Muon on a residual MLP under a shared training schedule. The supplied passages specify the dataset/model setup and final optimization configurations, while Table 5 concerns polynomial inverse plans.

  • The experiment uses CIFAR-10 with 45,000 training examples, 5,000 validation examples, and a standard 10,000-image test set.
  • The model is a 32-hidden-layer residual MLP with width 512, LayerNorm, GELU, residual additions, and a 512-to-10 classifier.
  • Table 5 specifies polynomial plans using CMAX=32, INTERVAL_PAD_REL=0.001, and NOISE_ABS=0.001, with total SYPP cost and certified residual bounds reported.
  • All three methods train for 100 epochs with batch size 4096 on an A100 GPU using linear warmup followed by cosine decay.
  • Muon and Newton–Muon optimize hidden-layer weight matrices while AdamW handles the remaining parameters; pure AdamW optimizes all parameters.

D Kronecker-Factored Curvature

The paper derives a Kronecker-factored curvature approximation from token-level output perturbations and uses it in a triplet quadratic surrogate. Under simplifying assumptions, the surrogate supports explicit score analyses for several update directions, while Newton–Muon is evaluated numerically.

  • Triplet Surrogate: The triplet surrogate combines the gradient matrix G, output-space curvature H, and activation matrix Z while preserving matrix structure.The activation matrix collects layer inputs, and the surrogate's quadratic term is induced by a Kronecker-factored curvature approximation.
  • Kronecker-Factored Approximation: The parameter-space curvature is approximated by retaining diagonal-token contributions and replacing varying token curvatures with their average H.This approximation is expected to be more accurate when the loss is not token-coupled.
  • Score Analysis: Under the quadratic surrogate, the analysis derives explicit score expressions for gradient descent, Muon, and Newton directions.The Newton direction is proportional to H^-1G(ZZ^T)^-1, whereas Muon uses the matrix-sign direction.
  • Score Analysis: Newton–Muon is evaluated numerically rather than analyzed theoretically in this section.Its predicted score is compared with Muon in the numerical study.

E.2 Isotropic Baseline Numerical Study

The isotropic baseline study compares update directions under simplifying assumptions on activations, curvature, and displacement. It finds that Muon and Newton–Muon benefit most from strongly anisotropic curvature, while their theoretical predictions coincide under isotropic inputs.

  • Setup: The baseline assumes isotropic Gaussian activations and displacement, a square matrix setting, and curvature generated by a fixed spectrum with a random orthogonal basis.The study varies curvature anisotropy through the eigenvalues of H and uses numerical spectral approximations for several scores.
  • Theoretical Comparison: In the isotropic-curvature asymptotic case, Muon achieves about 0.72 of the Newton score, while gradient descent coincides with Newton.The equality case corresponds to H being proportional to the identity.
  • Setup: Under isotropic inputs, Newton–Muon and Muon have identical theoretical predictions, so empirical differences arise from finite-sample deviations of ZZ^T/N from the identity.The simulations use 1024 independent runs and report mean scores with 2.5%–97.5% intervals.
  • Numerical Study: In the baseline with anisotropic curvature, Muon and Newton–Muon substantially improve over raw gradient descent, with Newton–Muon slightly ahead of Muon.When curvature is more uniform, the score gaps narrow, although Newton–Muon remains slightly better than Muon.
  • Numerical Study: With smaller data, Newton–Muon's advantage over Muon becomes much larger, suggesting stronger compensation for activation anisotropy when the empirical second moment is noisier.The smaller-data setting uses N = 1024 and p = 0.3.

F Non-Isotropic Assumption

The paper explores replacing the isotropic weight-displacement assumption with estimated non-isotropic covariance structure. These practical variants do not yield consistent improvements over Newton–Muon and can underperform standard Muon.

  • Covariance Factorization: The non-isotropic approach estimates a covariance factor M for the displacement covariance ΣW = MM^T and uses its factorized form to construct updates.The factorization avoids explicitly constructing ΣW^1/2.
  • Practical Estimators: The authors test momentum-based and diagonal estimators of non-isotropic displacement structure using update statistics.The momentum variant treats accumulated parameter updates as a coarse proxy for W − W⋆.
  • Results: The tested non-isotropic variants do not yield consistent improvements over Newton–Muon.The diagonal covariance variant is inexpensive but in some cases performs worse than standard Muon.
  • Interpretation: The paper argues that biased covariance estimates can reinforce update misalignment and slow training, motivating the isotropic proxy as a more robust geometric choice.The isotropic proxy avoids assumptions about the unknown displacement distribution while retaining the matrix-sign geometry.
Loading 2604.01472v1…