PaperlayerView exact revision on arXiv

Source-linked AI summary

Muon is Not That Special: Random or Inverted Spectra Work Just as Well

Zakhar Shumaylov, Nathaël Da Costa, Peter Zaika, Bálint Mucsányi, Alex Massucco, Yoav Gelberg, Carola-Bibiane Schönlieb, Yarin Gal, Philipp Hennig

arXiv:2605.11181v1cs.LGcs.AIcs.NAmath.NAmath.OCstat.ML

TL;DR

The paper asks whether precise geometry explains spectral optimizers’ performance, given that reported advantages can diminish with proper baseline tuning and vary with batch size. It tests this account with Schatten-based Freon, randomized-spectrum Kaon, and an analysis of local alignment and descent potential. Freon’s reported GPT-2 optima enter a regime outside unitarily invariant norms, while Kaon closely matches Muon despite using random spectra.

  • Problem

    The paper asks whether exact geometric formulations or the LMO framework explain spectral optimizers’ performance, amid benchmarks where their advantages diminish with proper baseline tuning and are sensitive to batch size.

    Source: §1 · p. 2

  • Method

    The paper introduces Schatten-based Freon and randomized-spectrum Kaon, then analyzes spectral updates through local batch-gradient alignment and directional descent potential.

    Source: §1 · p. 2

  • Results

    Freon’s optimal GPT-2 exponents concentrate in c ∈[1/2, 1], outside unitarily invariant norms, while Kaon closely matches Muon’s performance.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The paper concludes that spectral-optimizer performance is better understood through alignment, descent potential, and step-size choices than through a precise target geometry.

    Source: §4 · p. 9

    For p < 1, Freon’s general update is not steepest descent under the p-quasi-norm constraint, but is better viewed as preconditioned spectral descent.

    Source: §2.4 · p. 5

Abstract

from arXiv · show

The recent empirical success of the Muon optimizer has renewed interest in non-Euclidean optimization, typically justified by similarities with second-order methods, and linear minimization oracle (LMO) theory. In this paper, we challenge this geometric narrative through three contributions, demonstrating that precise geometric structure is not the key factor affecting optimization performance. First, we introduce Freon, a family of optimizers based on Schatten (quasi-)norms, powered by a novel, provably optimal QDWH-based iterative approximation. Freon naturally interpolates between SGD and Muon, while smoothly extrapolating into the quasi-norm regime. Empirically, the best-performing Schatten parameters for GPT-2 lie strictly within the quasi-norm regime, and thus cannot be represented by any unitarily invariant LMO. Second, noting that Freon performs well across a wide range of exponents, we introduce Kaon, an absurd optimizer that replaces singular values with random noise. Despite lacking any coherent geometric structure, Kaon matches Muon's performance and retains classical convergence guarantees, proving that strict adherence to a precise geometry is practically irrelevant. Third, having shown that geometry is not the primary driver of performance, we demonstrate it is instead controlled by two local quantities: alignment and descent potential. Ultimately, each optimizer must tune its step size around these two quantities. While their dynamics are difficult to predict a-priori, evaluating them within a stochastic random feature model yields a precise insight: Muon succeeds not by tracking an ideal global geometry, but by guaranteeing step-size optimality.

1 Introduction

The introduction questions whether precise geometric structure explains spectral optimizers’ performance, given that their advantages can diminish with baseline tuning and vary with batch size. It introduces Freon, Kaon, and a local account based on gradient alignment and descent potential.

1 Introduction

  • The introduction asks whether exact spectral geometry matters, noting that reported spectral-optimizer advantages diminish with proper baseline tuning and are sensitive to batch size.It frames this question against the prevailing explanation based on linear minimization oracles and strict geometric preconditioning.Source: §1 · pp. 1–2
  • Freon’s best GPT-2 exponents concentrate in c ∈[1/2, 1], outside the range of any unitarily invariant norm.Freon interpolates between SGD at c = 0 and Muon at c = 1/2; the paper says optimal exponents vary significantly between batches.Source: §1 · p. 2
  • Kaon replaces gradient singular values with noise yet closely matches Muon’s performance, challenging the need for a precise target spectrum.The paper also reports that Kaon retains classical convergence guarantees.Source: §1 · p. 2
  • The proposed local account decomposes loss change into batch gradient alignment and directional descent potential, whose tradeoff requires step sizes tuned to exploit descent potential.The account is mechanistic only in quadratic models.Source: §1 · p. 2

2 Background and Methods

The section introduces LMO-based steepest descent, then examines whether precise geometric constraints explain spectral optimizers’ behavior. It presents TruncatedSGD, Kaon, and Freon, alongside convergence results and methods for computing their updates.

2.1 Limitations of the LMO Framework

  • On WikiText-2, TruncatedSGD improves over SGD but does not close Muon’s gap, while Kaon and Freon with c ≈2/3 closely match Muon.The reported optimal Freon c lies outside the proper Schatten-norm range (c > 1/2).Source: Figure 4 · p. 4
  • Only order-preserving updates can be steepest descent under unitarily invariant norms, while Freon with c > 1/2, TruncatedSGD, and Kaon fall outside that class.The section also states that LMOs cannot remove bad large or intermediate singular directions, while sufficient-descent and boundedness conditions permit convergence without an LMO.Source: §2.1 · p. 4

2.2 TruncatedSGD

  • TruncatedSGD zeros the top p% of singular values; its improvement over SGD shows that suppressing large singular values aids stability, but is insufficient to match Muon.The authors note that full SVD makes TruncatedSGD impractical and use it as a pedagogical test of whether large singular values destabilize SGD.Source: §2.2 · pp. 4–5

2.3 Kaon

  • Kaon replaces gradient singular values with uniform random noise yet closely matches Muon, and its stochastic spectral descent has an almost-sure O(1/K) convergence guarantee.The paper characterizes Kaon primarily as a pedagogical construction rather than an optimizer based on a target geometry.Source: §2.3 · p. 5

2.4 Freon

  • Freon’s update is parameterized by dual exponent q, with c = 1 − q/2; for p ∈ (0,1), it extends analytically but is not steepest descent under the p-quasi-norm constraint.For p < 1, the p-Schatten quantity is a quasi-norm because the triangle inequality fails. The LMO collapses to the singular vector of largest singular value.Source: §2.4 · p. 5

2.4.1 Norm Hyperparameter Transfer

  • Mean Schatten-norm scaling ties the update constraint to the spectral ball, keeping the baseline maximum step size O(1) as rank and p vary.The scaling uses r^1/p so the normalized p-Schatten norm is bounded by the spectral norm.Source: §2.4.1 · p. 6

2.4.2 Particular case of p = 0 (c = 1)

  • At p = 0 (c = 1), Freon updates are equivariant under scaling and orthogonal layerwise symmetries, leaving the network’s optimization trajectory invariant.The passage identifies this as a property of the layerwise updates W_l − η(G_lG_l⊤)^−1G_l.Source: §2.4.2 · p. 6

2.4.3 Efficient Computation: From Polar to Hölder Express

  • The rational QDWH-based approach gives optimal approximations with doubly exponential convergence, avoiding the practical instability reported for polynomial iterations away from a/b = 1/2.The passage describes rational iterations as more stable and the approximation as optimal at every step.Source: §2.4.3 · pp. 6–7

3 Theoretical and Empirical Observations

The section analyzes optimization through local gradient alignment and directional descent potential, then tests their behavior in random-feature and GPT-2 settings. Its observations show that local and global landscapes can differ, while practical performance depends on the tradeoff between these quantities and step-size behavior.

3.1 The Two Quantities

  • The exact local Taylor expansion isolates batch gradient alignment γk and directional descent potential Φk as factors governing the optimizer's trajectory.The expansion also defines λk as the local directional curvature term, which appears multiplied by the learning rate.Source: §3.1 · p. 7
  • When λk is approximately constant, the quadratic loss-change expression is minimized at its optimal step size, yielding ∆f = −1/2 Φkγk^2.The derivation requires smoothness but no LMO or restriction on the direction; λk is not practically calculable because the intermediate point is unknown.Source: §3.1 · p. 7

3 Theoretical and Empirical Observations

  • Overall, performance depends on alignment, directional descent potential, and learning rate rather than the strict form of the update direction.The analysis is exact in the random-feature model and empirical for a GPT-2 training run.Source: §3.1 · p. 7

3.2 Exact Asymptotics in a Random Feature Model

  • In the aligned random-feature asymptotic model, γSGD ≤ γMuon when Λ is non-increasing, with equality only if Λ = cI or Σ = cI.The analysis assumes Gaussian mean-zero batched activations and a proportional limit with covariance aligned to the gradient's right singular vectors.Source: §3.2 · p. 8

3.3 Empirical Observations

  • With exact per-step line search in random-feature regression, GD reduces loss more than spectral descent but needs impractically oscillatory steps; optimal SpecGD steps remain constant.Figure 5 also tracks training loss, effective step sizes, and the exponent c used by adaptive-exponent methods.Source: §3.3 · p. 8, Figure 5 · p. 8
  • Greedily choosing c to maximize directional potential at each step gives rapid initial descent but harms long-term dynamics; a fixed-exponent rule improves overall performance.The rule assumes a power-law decay and is reported to overcome limitations of spectral descent for mean-zero activations.Source: §3.3 · p. 8
  • In GPT-2, Freon at c = 1 gives the best single-batch local decrease, yet optimal exponents fluctuate across batches in the full-data landscape.This shows that local single-batch models do not capture the complexity of global training dynamics.Source: §3.3 · p. 9, Figure 6 · p. 9
  • GPT-2 optimization rankings shift during training: SGD variants become suboptimal as Freon maintains good validation-loss decrease, while updates favoring mid-spectrum gradients trade alignment for greater descent potential.The observed alignment shift is associated with updates targeting mid-tier rather than extreme singular values, such as c ≈ 3/4.Source: §3.3 · p. 9

4 Conclusion

The conclusion argues that precise geometric structure is not necessary to explain spectral-optimizer performance, which instead depends on local alignment, descent potential, and step-size choices. The accompanying NanoGPT tuning details describe the sensitivity and dense-validation runs used for Figures 2 and 3.

4 Conclusion

  • Freon and Kaon challenge the geometric account of Muon: Freon can use updates outside norm regimes, while randomized spectra in Kaon perform comparably despite lacking coherent geometry.The conclusion attributes performance to the interplay of batch gradient alignment, directional descent potential, and step-size choices.Source: §4 · p. 9

A NanoGPT Tuning Details for Figures 2 and 3

  • The NanoGPT sensitivity sweep covered four optimizers, seven plotted scale factors, and three seeds, totaling 84 runs; the plotted sweep omits ρ = 100.Each optimizer's matrix-update and base learning rates were tuned, and the plotted scales jointly multiply the selected pair.Source: A NanoGPT Tuning Details for Figures 2… · p. 15
  • Dense-validation curves use three seeds and validation every eight optimizer steps, plotting Muon at ρ = 1, Kaon at ρ = 3, and both Freon variants at ρ = 1.The plotted lines average the three seeds, with shaded regions showing ±2 standard deviations.Source: A NanoGPT Tuning Details for Figures 2… · p. 15

B GPT2 Training Details

The GPT-2 experiments use a fixed WikiText-2 training setup and compare SGD variants, Muon, Kaon, and Freon configurations. Freon sweeps exponents spanning the SGD and Muon cases.

B GPT2 Training Details

  • Each GPT-2 run trains a 124M-parameter, 12-layer model on WikiText-2 with GELU, FlashAttention, context length 512, batch size 8, and a 5% warm-up schedule.The combined Figure 4 reports final validation loss across 8–10 learning rates per optimizer configuration.Source: B GPT2 Training Details · p. 15
  • SGD uses momentum 0.9, while TruncatedSGD sweeps truncation levels from 0% to 10% and rescales nonzero-truncation updates to match the untruncated gradient's 2-norm.At 0% truncation, the method recovers vanilla SGD.Source: B GPT2 Training Details · p. 15
  • Muon and Kaon use momentum 0.95 and five matrix iterations; Freon uses the same momentum and iteration count while sweeping c from 0 to 1.The swept Freon values include 0, 1/4, 1/3, 1/2, 2/3, 3/4, and 1; c = 0 recovers gradient descent and c = 1/2 recovers Muon.Source: B GPT2 Training Details · p. 15

C Limitations

The paper identifies limits in the scope of its empirical validation and theory, the predictive use of its diagnostic framework, and the current implementation of Freon.

C Limitations

  • Empirical validation covers only a random feature model and GPT-2, while rigorous analysis of alignment and descent potential is limited to the random-feature setting.How these dynamics extend to other modalities or non-quadratic landscapes remains open.Source: C Limitations · p. 16
  • Exact asymptotic results rely on mean-zero activations, excluding activations such as ReLU.Source: C Limitations · p. 16
  • The framework is diagnostic rather than predictive: stochastic optimization quantities can be evaluated reliably only post-factum, not a priori.Source: C Limitations · p. 16
  • Batch-to-batch variation in optimal Schatten exponents prevented a stable, greedily adaptive schedule, and Freon’s implementation may need further efficiency engineering.The authors say it remains unclear whether extensive engineering is necessary given similar behavior across optimizers.Source: C Limitations · p. 16

D.1 Preconditioned Spectral Descent

The section gives convergence conditions for preconditioned and random spectral descent, characterizes when spectral updates correspond to unitarily invariant steepest descent, and develops equivariance and rational-approximation results relevant to Freon.

D.1 Preconditioned Spectral Descent

  • A preconditioned spectral update is steepest descent under a unitarily invariant matrix norm only if its spectral mapping preserves the order of singular values.Source: §D.1 · p. 16
  • Preconditioned spectral descent converges in the lim-inf gradient sense when its updates satisfy sufficient-descent and boundedness conditions and the associated divergence condition holds.The deterministic result states lim inf_k ||G_k||_* = 0.Source: §D.1 · pp. 17–18
  • For random spectral descent with independently sampled positive entries, the minimum squared dual-gradient norm achieves an almost-sure O(1/K) rate.Source: §D.1 · p. 18

D.2 Equivariance of Freon (c = 1)

  • Freon with c = 1 is equivariant under the stated network symmetries, with broader general-linear symmetries available under full-row or full-column rank conditions.Source: §D.2 · pp. 19–20

E.1 Fundamental Properties of the a/b-Method

  • The a/b-method uses a class of generalized rational functions, and the section establishes existence of optimal approximants for the stated zero-deficiency continuous functions.The passage cautions that optimal approximants are not generally guaranteed without the stated function-class conditions.Source: §E.1.1 · pp. 21–22

E.1.1 Existence and Characterization of the Optimal Approximants.

  • An approximant with at least n + d + 2 extremal points is optimal, and the section characterizes optimal approximants through this condition.Source: §E.1.1 · p. 22

E.1.2 Optimality

The section establishes optimality for a composition of rational approximants built through successive interval minimizations, and extends the result to polynomial approximants.

E.1.2 Optimality

  • Composing the stagewise approximants R_t yields an optimal R* for the stated minimax approximation problem.Source: §E.1.2 · p. 23
  • The construction chooses each stage’s interval endpoints from the minimum and maximum of its approximant over the current interval.Source: §E.1.2 · p. 23
  • The same conclusions extend to the polynomial class P_n(r).Source: §E.1.2 · p. 24
  • For the matrix approximation result, the composed approximant gives a best approximant of G for every T.Source: §E.1.2 · p. 24

E.1.3 Approximation Errors

The section gives an approximation-error result for the iterative method and compares numerical stability across implementations and tested conditions.

E.1.3 Approximation Errors

  • Theorem E.8 supplies an approximation-error estimate for X_T when T ≥ 0 and σ_min(G) ≥ l, with l ∈ (0, 1).Source: §E.1.3 · p. 25

E.3 Instability of Polynomial Iterations

  • Coupled_chol remained stable near machine precision across tested matrix sizes, condition numbers, and dtypes, while direct and coupled_pq diverged under high conditioning.Regularization delayed but did not eliminate divergence; coupled_chol required none, with its Cholesky structure maintaining positive definiteness.Source: §E.3 · p. 25

E.4 Iterative Methods for (GG⊤)−a/b G

The section develops iterative schemes targeting the fixed point (GG⊤)^−a/bG, then compares their numerical stability and computational costs. Coupled-Chol is the only method reported stable across all tested conditions and precisions; rational-iteration parameters are set by schedules in b.

E.4 Iterative Methods for (GG⊤)−a/b G

  • Coupled-Chol remains stable across all tested conditions and precisions, while direct and coupled_pq degrade or diverge at high condition numbers in low precision.The section identifies Coupled-Chol as the final method with the forward stability needed for reliable convergence.Source: §E.4 · p. 26

E.4.1 Direct Iteration

  • Direct iteration recomputes its polynomial argument from the current iterate, with separate formulations for the general, polar, and Freon cases.The Freon formulation uses a mixed product with the frozen initial transpose to avoid a matrix square root and condition-number squaring.Source: §E.4.1 · pp. 26–27

E.4.2 M-Accumulator Iteration

  • M-Accumulator tracks a preconditioner applied to the frozen gradient at convergence, making it cheapest per step when n ≫ m.At convergence, the preconditioner approaches (GG⊤)^−a/b.Source: §E.4.2 · pp. 27–28

E.4.3 Coupled Iteration

  • Coupled iteration tracks the argument matrix as a state, trading one additional G-MM per step for one fewer S-MM than M-Accumulator.The section says this trade-off is favorable only when m ≈ n.Source: §E.4.3 · p. 28

E.4.4–E.4.6 Coupled-ab, Dual-ab, and Coupled-Dual Iterations

  • Coupled-ab initializes its argument at the raw Gram matrix, Dual-ab uses a factored representation, and Coupled-Dual runs two coupled masters with a cross-stabilizer.These variants respectively avoid fractional-power initialization, avoid explicitly forming GG⊤ at each step, and avoid fractional matrix powers inside the loop.Source: §E.4.4 · p. 28, §E.4.5 · pp. 28–29, §E.4.6 · p. 29

E.4.7 Rational-Chol (QDWH-style) Iteration

  • Rational-Chol applies a rational function targeting x^−1/b and uses block QR with a Cholesky factor to avoid condition-number squaring.For a/b = 1/2, it reduces to QDWH for polar decomposition.Source: §E.4.7 · p. 29

E.4.8 Coupled-Chol Iteration

  • Coupled-Chol generalizes QDWH and improves forward stability by exponentiating the well-conditioned single-step update rather than the accumulated preconditioner.Its costs match Rational-Chol, and the section attributes its superior numerical stability to Wt ≈ I.Source: §E.4.8 · p. 30

E.5 Picking Parameters for Freon Iterations: Cushion and l

  • The rational Newton–Schulz iteration uses b-dependent schedules δ(b) = (1.84 × 10^−8)^1/b and l(b) = (10^−11)^2/b to limit coefficient blow-up and stabilize the fifth iteration.The cushion δ stabilizes denominator coefficients, while l sets the spectral range the iteration must cover.Source: §E.5 · p. 30

F.1 Kaon

Kaon repeatedly applies a polynomial map to gradient singular values instead of using Muon’s exact polar factorization. Its iterates explore a broad stationary distribution, while its matrix-multiplication cost matches Muon’s.

Stationary Distribution

  • At λ = 4.1, Kaon’s singular-value iterates explore a chaotic stationary distribution supported on (0, 1.175), rather than converging to a fixed point.Trajectories also do not settle into periodic orbits; the stochasticity is described as intrinsic to Kaon.Source: §F.1 · p. 30
  • Kaon applies f(x) = λx(1 − x^2)^2 iteratively to gradient singular values as a cheap alternative to Muon’s exact polar factorization.The map’s distributions shift and spread across the tested λ values from 3.5 to 4.25.Source: §F.1 · p. 30

Cost of Kaon

  • Kaon matches Muon’s matrix-multiplication cost, with both methods costing T(4r^2s + 2r^3) FLOPs to leading order.Both use the same number T of steps and perform the same listed matrix operations per step.Source: §F.1 · p. 30

G Random Feature Asymptotics

This material develops random-feature asymptotics for local optimization quantities and describes the random-feature experiments and an approximation for selecting a Schatten exponent. It also gives a diagonal-case comparison of SGD and Muon alignment.

G Random Feature Asymptotics

  • Under the stated Gaussian random-feature assumptions, the asymptotic analysis yields universal polynomials, with the first three specified in the section.The theorem assumes A = C^1/2 Z, where C is symmetric positive semidefinite with bounded operator norm and Z has iid Gaussian entries.Source: G Random Feature Asymptotics · p. 32

G.1 The Diagonal Case

  • When Λ’s entries are non-increasing, γ_SGD ≤ γ_Muon, with equality if and only if Λ = cI or Σ = cI.This is the stated diagonal-case alignment comparison.Source: §G.1 · p. 36

H.1 Detailed Setup

  • The random-feature experiments use ReLU and SwiGLU settings with different stable ranks, and select c and η jointly by maximizing a quadratic lower bound over a grid.The ReLU run lasts 1,000 iterations and the SwiGLU run 400; both use float64 and a fixed random seed.Source: §H.1 · p. 37

H.2 Exponent Approximation for Optimal Schatten Norm

  • The fixed-exponent approximation assumes power-law decay in gradient singular values and activation projection weights, σ_i ∝ i^−α and k_i ∝ i^−β.Under that assumption, the section derives the exponent relation α(a + 1) = 2aα + β.Source: §H.2 · p. 37

H.3 Optimal C methods

The Optimal C methods choose a Schatten exponent alongside a step size, either by a grid search or by fitting spectral power laws. In the SwiGLU random-feature setting, both methods converge to a smaller exponent than in the ReLU experiment.

H.3 Optimal C methods

  • Both adaptive-exponent methods use directions D_c = (GG⊤)^−cG and the exact optimal step size at each step.The direction is defined through the gradient’s singular values normalized by a power mean.Source: §H.3 · p. 37
  • With SwiGLU features, both Optimal C methods converge to a smaller c, closer to GD, than in the ReLU experiment.The SwiGLU feature matrix has substantially higher stable rank, changing the curvature landscape relative to ReLU.Source: Figure H.1 · p. 38
  • Optimal C [Greedy] searches 41 values of c in [−0.5, 1.5] at each iteration and selects the value minimizing predicted loss decrease.The search evaluates the quadratic predicted loss decrease for each grid value.Source: §H.3 · p. 38
  • Optimal C [Scaling] estimates c analytically by fitting power laws to the gradient singular values and activation projection energies.The fitted exponents describe log σ_i and log k_i as functions of log i.Source: §H.3 · p. 38

I GPT2 Additional Experiments

The additional experiments examine how update direction and learning rate affect loss locally, per layer, and across the full model. They compare gradient alignment and curvature measures, quadratic approximations, and validation-loss evaluations.

Validation and alignment diagnostics

  • The appendix plots validation loss on WikiText-2 at each optimizer’s best learning rate and reports per-layer curvature-alignment and batch-gradient-alignment quantities for candidate directions.The alignment figures cover 48 weight matrices; γk = 1 indicates alignment with the full gradient.Source: Figure I.2 · p. 39, Figure I.1 · p. 39, Figure I.8 · p. 42

Random-feature quadratic analysis

  • The random-feature quadratic analysis evaluates loss change across c and learning rate, identifies the optimal rate for each c, and reports layer-wise optimal c across three training snapshots.c = 0 recovers gradient descent; c = 1/2 recovers Muon, and c = 1 is the pseudoinverse-like endpoint.Source: Figure I.3 · p. 40, Figure I.4 · p. 40, I GPT2 Additional Experiments · p. 40

Local loss comparisons

  • In the single-batch quadratic and actual-loss evaluations, c = 1 achieves the largest per-step decrease, consistent with the GGN approximation.The naive Frobenius random-feature model instead predicts c ≈ 2/3 as optimal.Source: Figure I.5 · p. 40

Validation loss analyses

  • The full-validation analyses compare loss change across c and learning rate globally and per layer, while also measuring each layer’s range of c values within 10% of its best one-step decrease.The figures contrast full-validation per-layer evaluations with evaluations on the same training batch used for the gradient.Source: Figure I.7 · p. 41, Figure I.6 · p. 41, Figure I.10 · p. 44

Related papers

Loading 2605.11181v1…