Second-Order Muon Done Right: A Principled Marriage of Spectral Geometry and Curvature

S ECOND -O RDER M UON D ONE R IGHT

S ECOND -O RDER M UON D ONE R IGHT: A P RINCIPLED M ARRIAGE OF S PECTRAL G EOMETRY AND C URVATURE Tong Che NVIDIA Research tongc@nvidia.com

arXiv:2608.09763v1 [cs.AI] 10 Aug 2026

A BSTRACT Muon’s polar update is exact for an unweighted spectral geometry. We introduce G O -MUON, which uses a matched data-dependent geometry and reuses it across several optimization steps. Conditioned on any positive-definite left and right maps, its raw update exactly solves the corresponding weighted spectral oracle; this statement is independent of how the maps are estimated or how recently they were refreshed. For softmax cross-entropy, we quantify when the observed-label backward factor approaches the model Fisher and generalized Gauss–Newton factor. We also show that four-step refresh nearly preserves the tracking delay of slowly changing geometry while increasing stationary factor noise, making lazy geometry a compute–statistics tradeoff rather than a denoising mechanism. In a hash-matched Penn Treebank comparison, G O -MUON lowers the geometric mean of three fixed training-loss windows by 0.38% across three paired seeds. On Tiny Shakespeare, a completed three-seed comparison lowers the same trajectory summary by 3.71% against the registered Muon control. Four-step refresh also improves the late PTB training window and reduces mean step time by about 20%. On corrected independent-CLS modular addition, a formal one-seed-per-modulus screen finds that G O -MUON reaches sustained 99% held-out accuracy at update 290 versus 2320 for Muon on modulus 103, and at update 220 versus 4520 on modulus 107.

1

I NTRODUCTION

For a matrix parameter W ∈ Rm×n , Muon applies a polar map to a momentum source M (Jordan et al., 2024). Up to the sign convention of the parameter update, the direction solves max ⟨M, D⟩ = ∥M ∥∗ , D = Polar(M ). (1) ∥D∥op ≤1

This moves every nonzero singular direction equally under an operator-norm update budget. The same construction extends to any declared symmetric positive-definite (SPD) coordinate maps PA , PB . The matched direction D0 = PB Polar(PB M PA )PA is the exact support direction of the ball ∥PB−1 DPA−1 ∥op ≤ 1. This separates two questions that are often conflated: the oracle is exact for the supplied geometry, while the choice and estimation of that geometry determine whether the direction is useful. Fisher semantics, Kronecker factorization, and factor freshness are not assumptions of the oracle statement. G O -MUON instantiates the maps with inverse fourth roots of OAS-regularized activation and backpropagated-gradient second moments. Under softmax cross-entropy these observed-label factors approach model-Fisher/GGN factors when minibatch error vanishes and the predictive distribution approaches the data conditional. The quarter power tempers factor anisotropy relative to a full inverse, while a positive Frobenius graft restores Muon’s update energy. The geometry is refreshed every K = 4 updates. Reuse steps skip activation retention, outputgradient hooks, moment construction, shrinkage, and factor eigendecompositions, while gradient 1

S ECOND -O RDER M UON D ONE R IGHT

K and momentum information remains current. The coefficient βB preserves elapsed-step decay of the retained backward state. It does not reproduce the omitted minibatch averages: under stationary noise the resulting estimator is noisier than a per-step EMA, even though its average lag on a slowly varying signal is nearly unchanged.

Our contributions are: 1. a factor-agnostic formulation of the matched spectral oracle, together with quantitative conditioning, grafted-alignment, and cached-geometry bounds for the implemented direction; 2. a conditional empirical-Fisher-to-Fisher/GGN result for softmax backward factors and an K exact variance–lag characterization of sparse βB refresh; 3. the G O -MUON system, which combines OAS K-FAC-style factors with hook-gated factor refresh, cached inverse fourth roots, and fresh gradient and momentum updates; 4. paired character-language-model experiments showing a 0.38% fixed-window trajectory reduction against Muon in a hash-matched PTB comparison, a 3.71% reduction against the registered Tiny Shakespeare Muon control, and a 0.798× mean-step-time ratio for four-step refresh; and 5. a corrected independent-CLS modular-addition screen in which G O -MUON reaches sustained 99% held-out accuracy at updates 290 and 220, versus 2320 and 4520 for Muon.

2

M ATCHED SPECTRAL GEOMETRY

2.1

A FACTOR - AGNOSTIC ORACLE

n Let PB ∈ Sm ++ , PA ∈ S++ , and define

C = PB M PA ,

qP (D) = ∥PB−1 DPA−1 ∥op .

(2)

For ρ ≥ 0, consider the support problem max ⟨M, D⟩ D

subject to

qP (D) ≤ ρ.

(3)

Theorem 1 (Matched spectral oracle). Let Q(C) be any support polar satisfying Q(C) ∈ arg max ⟨C, U ⟩ . ∥U ∥op ≤1

Then the value of Equation (3) is ρ∥C∥∗ , and one maximizer is Dρ = ρPB Q(C)PA .

(4)

⟨M, D0 ⟩ = ⟨C, Q(C)⟩ = ∥C∥∗ .

(5)

At unit radius, write D0 := D1 ; then

For M ̸= 0, this quantity is strictly positive. Under the conventional minimization definition, −Dρ is the LMO; the optimizer equivalently subtracts the positive-alignment direction Dρ . Proof. Set U = PB−1 DPA−1 , so D = PB U PA . Cyclicity of the trace gives ⟨M, D⟩ = ⟨PB M PA , U ⟩ = ⟨C, U ⟩ . The constraint is ∥U ∥op ≤ ρ. Operator/nuclear duality gives the value ρ∥C∥∗ , and mapping a support polar back gives Equation (4). This is the matrix specialization of the general preconditioned-norm pullback oracle (Veprikov et al., 2025). 2

S ECOND -O RDER M UON D ONE R IGHT

Polar selection. If C has full rank k = min(m, n), its support polar is unique. At rank deficiency, the support maximizer is not unique: the canonical partial polar sets the null-space completion to zero, whereas the compact full-SVD polar used in our implementation supplies a semi-orthogonal completion. Both solve Equation (3), but the completion can affect a later Frobenius graft. For M = 0, every feasible point is optimal; the implementation defines D0 = 0 and skips the graft. Corollary 2 (Cached geometry). For any cached SPD maps PeA , PeB and any M ̸= 0, D E M, PeB Q(PeB M PeA )PeA = ∥PeB M PeA ∥∗ > 0. (6) Corollary 2 certifies exactness for the cached ball, not closeness to the current ball. With momentum, M is the source passed to the polar map; current-gradient alignment and finite-step loss change are separate properties. A scale-free comparison with current maps is possible. For −1 e −1 e Θt = κ2 (PB,t PB ) κ2 (PA,t PA ),

(7)

Proposition 7 shows that the cached direction, rescaled to its current weighted radius, achieves at least a Θ−1 t fraction of the current oracle value. This is conditional on coordinate-map drift; the present experiments do not measure Θt . 2.2

W HY THE MAP - BACK MATTERS

Conditioning only before the polar map need not preserve source alignment. For example, take   −1 −4 A = B = diag(1/9, 9), M= . −2 −6 The inside-only score and its polar are  −9 Z = B −1/2 M A−1/2 = −2

 −4 , −2/3

Polar(Z) = √

 1 −25 949 −18

 −18 , 25

so

17 ⟨M, Polar(Z)⟩ = − √ < 0. (8) 949 A positive scalar graft cannot change this sign. The matched map-back instead optimizes in transformed coordinates and returns to the original tangent coordinates, yielding the positive identity in Equation (5).

3

K-FAC- ROOT INSTANTIATION

3.1

O BSERVED - LABEL FACTORS AND THEIR INTERPRETATION

For a linear layer y = W x, the implementation forms minibatch second moments N X b= 1 A xi x⊤ i , N i=1

N X b= 1 B δi δ ⊤ , N i=1 i

δi =

∂ℓi . ∂yi

(9)

These are activation and observed-label backpropagated-gradient factors. K-FAC approximates an ba expected Kronecker block by a product of expected factors (Martens & Grosse, 2015). Calling B Fisher factor requires an additional relation between the data conditional and the model predictive distribution (Kunstner et al., 2019). For softmax cross-entropy, that relation can be quantified exactly. Fix an input x, write p = pθ (· | x) and q = q(· | x), and let δz = p − ey be the logit gradient. Define Bq,z = Ey∼q [δz δz⊤ ],

BF,z = Ey∼p [δz δz⊤ ] = Diag(p) − pp⊤ .

For logits as natural parameters, BF,z is also the output GGN factor (Martens, 2020). 3

S ECOND -O RDER M UON D ONE R IGHT

Proposition 3 (Observed-label factor mismatch). Let r = q − p. Then Bq,z − BF,z = Diag(r) − pr⊤ − rp⊤ ,

∥Bq,z − BF,z ∥2 ≤ 3∥q − p∥2 .

(10)

If δl = Jl⊤ δz , where Jl is independent of y, then ∥Bq,l − BF,l ∥2 ≤ 3∥Jl ∥22 ∥q − p∥2 .

(11)

Proof. Expanding Ey∼q [(p − ey )(p − ey )⊤ ] and subtracting Diag(p) − pp⊤ gives the identity. Then ∥ Diag(r)∥2 ≤ ∥r∥2 and ∥pr⊤ ∥2 = ∥p∥2 ∥r∥2 ≤ ∥r∥2 . The hidden-layer result follows by congruence with Jl . Corollary 4 (Conditional factor consistency). Consider a sequence θj for which pθj (· | x) → q(· | x) bj,l is consistent for Bq,l (θj ), then and supj ∥Jl (θj )∥2 < ∞. If a minibatch estimator B p

bj,l − BF,l (θj )∥2 − ∥B →0 as both the effective sample size and j tend to infinity. Corollary 4 requires realizability or predictive consistency; it is not a property of arbitrary training trajectories. It also applies only to the backward factor: the full block still incurs the K-FAC approximation E[(xx⊤ ) ⊗ (δδ ⊤ )] ̸= E[xx⊤ ] ⊗ E[δδ ⊤ ] in general. Moreover, the softmax Fisher always has the logit-shift null direction, and its remaining eigenvalues can vanish as p approaches a point mass. Absolute factor convergence therefore does not imply relative inverse-root accuracy near interpolation; the effective floor and conditioning bounds below must be included when analyzing the resulting geometry. 3.2

Q UARTER - POWER GEOMETRY AND F ROBENIUS GRAFT

Let A ≻ 0, B ≻ 0 denote the effective regularized factors. For 0 ≤ α ≤ 1/2, define Cα = B −α M A−α , Dα,0 = B

−α

(12) −α

Q(Cα )A , √ k b Dα = Dα,0 , ∥Dα,0 ∥F

(13) k = min(m, n).

(14)

The cases α = 0, 1/4, 1/2 give Muon, the G O -MUON quarter-power geometry, and full-factor spectral map-back, respectively. The last is a spectral oracle in full-factor coordinates, not a naturalgradient step. For α = 1/4, ∥B 1/4 DA1/4 ∥op ≤ 1 is the operator-norm analogue of the square-root K-FAC quadratic geometry, since vec(D)⊤ (A1/2 ⊗ B 1/2 ) vec(D) = ∥B 1/4 DA1/4 ∥2F .

(15)

Replacing the Frobenius norm by the operator norm is a design choice. Under the conditions of Proposition 3, this admits a root-Fisher/GGN interpretation; the oracle itself requires only SPD factors. Proposition 5 (Grafted conditioning and source alignment). Let Q be the compact full-SVD support polar used by the implementation and let KAB = κ2 (A)κ2 (B). Let κ+ denote the ratio of largest to smallest nonzero singular value. For M ̸= 0, D E α b α ≥ K −α ∥M ∥∗ , b α ) ≤ KAB M, D κ+ (D . (16) AB 4

S ECOND -O RDER M UON D ONE R IGHT

The first bound compares equal-Frobenius-radius source alignment with Muon’s value ∥M ∥∗ . At α = 1/4, the retained fraction is at least [κ(A)κ(B)]−1/4 . This is a worst-case conditioning certificate, not a pointwise ordering of Muon, quarter-power, and half-power directions. Appendix A gives the proof. The positive graft preserves the oracle ray and sign: √ D E k bα = M, D ∥Cα ∥∗ > 0. ∥Dα,0 ∥F It does√not preserve the original unit-radius LMO. Its weighted spectral radius is the source-dependent value k/∥Dα,0 ∥F . The grafted direction is invariant to independent positive scalar rescalings of A and B. With the canonical polar—and with the compact polar at full rank—it is also equivariant under orthogonal changes of the input and output bases: replacing (M, B, A) by (U M V ⊤ , U BU ⊤ , V AV ⊤ ) sends b α to U D b α V ⊤ . At rank deficiency, a full-SVD null-space completion can break this equivariance. D General affine reparameterization invariance does not follow because matrix powers and the polar map are not equivariant under arbitrary congruences (Luk & Grosse, 2018). 3.3

OAS FACTORS AND NUMERICAL SPD MAPS

Finite-minibatch factors can be singular. For a d × d empirical second moment S built from N samples, we use Oracle Approximating Shrinkage (OAS) (Chen et al., 2010): (1 − 2/d) tr(S 2 ) + tr(S)2 , (17) (N + 1 − 2/d)(tr(S 2 ) − tr(S)2 /d) tr(S) OAS(S) = (1 − λ)S + λ I. (18) d When λ > 0 and tr(S) > 0, d − (d − 1)λ κ2 (OAS(S)) ≤ . (19) λ The implementation additionally floors eigenvalues at roundoff scale before taking inverse fourth roots. The resulting numerical SPD matrices, rather than the potentially singular raw moments, are the factors to which Theorem 1 and Proposition 5 apply. λ(S, N ) = clip[0,1]

Writing Bteff and B̄teff for the current and cached factors after shrinkage and the root routine’s eigenvalue floor, Fisher-factor error has four distinct sources: bt − Bq,t ∥2 + ∥Bq,t − BF,t ∥2 ∥B̄teff − BF,t ∥2 ≤ ∥B | {z } | {z } finite sample

model mismatch

bt ∥2 + ∥B̄teff − Bteff ∥2 . + ∥Bteff − B |

{z

shrinkage/floor

}

|

{z

(20)

}

EMA and staleness

The decomposition separates the conditional Fisher interpretation from the regularization and temporal choices used by the optimizer. 3.4

S PARSE REFRESH IS A VARIANCE – LAG TRADEOFF

At refresh step t, we compute bt ), At = OAS(A

bt ). Bt = OAS(B

(21)

The first Bt initializes B̄. Later refreshes use K K B̄t = βB B̄t−K + (1 − βB )Bt .

(22)

The exponent applies K elapsed steps of decay to the retained state. It does not reproduce K per-step EMA observations. To characterize the difference, let Xt be a centered post-OAS factor observation and compare Zt = βB Zt−1 + (1 − βB )Xt 5

S ECOND -O RDER M UON D ONE R IGHT

with the refresh-boundary process K K Yj = βB Yj−1 + (1 − βB )XjK .

Proposition 6 (Sparse-EMA covariance and lag). If vec(Xt ) is weakly stationary with finite covariance and both recursions are in steady state, then at refresh boundaries Cov(vec Yj ) ⪰ Cov(vec ZjK ).

(23)

For temporally IID observations with covariance Σ, K 1 − βB Σ, K 1 + βB 1 − βB Cov(vec Z) = Σ. 1 + βB

Cov(vec Y ) =

(24) (25)

For an affine signal observed without noise, the phase-averaged lag of the held estimator is LK =

K KβB K −1 + . K 2 1 − βB

(26)

At βB = 0.97, K = 4, the sparse-to-per-step variance ratio is 3.99537, while L4 = 32.3714 steps versus L1 = 32.3333. Thus sparse refresh nearly preserves average slow-signal lag while retaining about one quarter as many independent factor observations. Sampling occurs before the refreshdomain EMA, so high-frequency variation can alias. The activation factor At has no EMA and is simply sampled and held. Holding one geometry within each four-step block creates blockwise geometry coherence: the preconditioner does not rotate between minibatches inside the block. Whether this coherence stabilizes the optimizer trajectory is an empirical mechanism hypothesis, not a consequence of Proposition 6. 3.5

C OMPLETE UPDATE

With K = 4, βB = 0.97, momentum coefficient β, and a Nesterov-style source, one optimizer step is: if t mod K = 0 :

At , Bt ← OAS K-FAC statistics,

(27)

K K B̄t ← βB B̄t−K + (1 − βB )Bt , −1/4 −1/4 PA ← At , PB ← B̄t ;

(28)

Mt = βMt−1 + (1 − β)Gt , ct = (1 − β)Gt + βMt , M ct PA , Ct = PB M D0,t = PB Polar(Ct )PA , p min(m, n) , Dt = D0,t ∥D0,t ∥F p Wt+1 = Wt − ηt max(1, m/n)Dt .

(29) (30) (31) (32) (33) (34) (35)

The learning rate ηt follows a cosine schedule. Embeddings, normalization parameters, and other non-matrix parameters use AdamW (Kingma & Ba, 2015). The implementation gates its forward hooks before the model forward pass. On refresh steps it retains linear-layer inputs and registers output-gradient hooks. On reuse steps the hook returns immediately, bypassing activation storage, delta hooks, moment formation, OAS, and factor eigendecomposition. Momentum, the central polar map, grafting, and the parameter update execute every step. 6

S ECOND -O RDER M UON D ONE R IGHT

3.6

C OST

For a layer of shape m × n, a refresh computes moments and matrix functions of the m × m and n × n factors. Reuse eliminates those operations and their telemetry while retaining the matrix products and polar map needed for the direction. A refresh fraction of 1/K amortizes factor collection and decomposition across K updates. The experiments measure the resulting savings end to end.

4

R ELATED WORK

4.1

M UON AND NORM - BASED MATRIX OPTIMIZATION

Muon applies an approximate polar map to momentum matrices, using Newton– Schulz iterations to make matrix updates approximately semi-orthogonal (Jordan et al., 2024). The norm-based account of matrix optimization identifies this map as steepest descent under an operator-norm constraint, whose dual objective is the nuclear norm (Bernstein & Newhouse, 2024). Modular norm optimization extends this viewpoint from individual matrices to composed neural-network modules (Large et al., 2024); Scion develops norm-constrained linear minimization oracles more generally (Pethick et al., 2025); and Gluon connects practical Muon-like updates with convergence theory for LMO-based optimization (Riabinin et al., 2025). Subsequent systems work implements Muon at language-model scale (Liu et al., 2025). Several variants refine Muon’s normalization while retaining a polar or orthogonalization core. Muon+ adds an additional normalization step (Zhang et al., 2026a), while MuonEq uses inexpensive equilibration before orthogonalization (Chang et al., 2026). These methods modify the matrix presented to the polar map. G O -MUON instead instantiates the general preconditioned-norm pullback oracle: the same SPD maps transform the source into polar coordinates and map the support direction back. Its positive Frobenius graft restores Muon’s update energy while retaining that oracle ray. 4.2

K RONECKER CURVATURE AND STRUCTURED PRECONDITIONING

K-FAC approximates a layerwise model-Fisher block by the Kronecker product of an activation factor and a backward factor (Martens & Grosse, 2015). It has been scaled through distributed factor formation, inversion, and communication (Osawa et al., 2019; Pauloski et al., 2020); SINGD gives an inverse-free, numerically stable structured natural-gradient formulation (Lin et al., 2024). Full Gauss– Newton experiments on small language models demonstrate accurate layerwise curvature as a strong optimization oracle (Abreu et al., 2025). G O -MUON uses observed-label second moments rather than model-sampled Fisher factors. Proposition 3 quantifies when the backward factor approaches the Fisher/GGN factor; the distinction remains essential away from predictive consistency (Kunstner et al., 2019). K-FAC’s natural-gradient update retains invariance to affine activation reparameterizations under its metric construction (Luk & Grosse, 2018). The quarter-power spectral oracle retains only the orthogonal and scalar equivariances stated in Section 3. Shampoo uses Kronecker-factored matrix moments and inverse matrix powers for large-scale preconditioning (Gupta et al., 2018); scalable Shampoo amortizes these matrix operations in distributed training (Anil et al., 2020). SOAP runs an Adam-like method in Shampoo’s evolving eigenbasis (Vyas et al., 2024). The general theory of preconditioned matrix norms makes explicit how transformations inside and outside an LMO determine the corresponding steepest-descent geometry (Veprikov et al., 2025). Theorem 1 is a Kronecker matrix instance of this general pullback, not a new duality principle. In G O -MUON, OAS-regularized activation and backward second moments select the geometry; Shampoo uses gradient moments. OAS was introduced as an automatic shrinkage estimator for finite-sample covariance matrices (Chen et al., 2010); here it supplies isotropic shrinkage before the numerical eigenvalue floor. 4.3

C URVATURE - AWARE M UON VARIANTS

Mousse applies a quarter-power matched polar map with a Frobenius graft to Shampoo-style gradientGram geometry (Zhang et al., 2026b). It updates factors every step, refreshes matrix functions periodically, and uses an inverse-power exponent below one quarter in its reported language-model 7

S ECOND -O RDER M UON D ONE R IGHT

Table 1: Relationship to representative structured matrix optimizers. “Factor skip” denotes steps that skip accumulation of the underlying second-moment factors. Method

Statistical geometry

Muon K-FAC

None Activation/backward Fisher factors Scalable Shampoo Gradient Gram moments SOAP Mousse FISMO G O -MUON

Matrix update

Curvature amortization

Polar(M ) Approximate natural gradient Kronecker inverse-power preconditioning Adam in the Shampoo basis

— Periodic inverse; later systems also skip factors Periodic roots; optional factor skip Periodic eigenbasis refresh

Shampoo moments and eigenbasis Shampoo gradient moments Root-metric matched polar with graft Fisher/K-FAC factors Fisher-metric matched spectral step OAS observed-label Quarter-power matched activation/backward polar with graft moments

Factors frequent, matrix functions periodic Factors updated in its main algorithm Hook-gated factors and roots every K; fresh source every step

experiments. The quarter-power matched-map structure and graft are shared. G O -MUON differs in its OAS activation/backward factor field, exact quarter-power instance, and hook-gated schedule that skips factor collection on reuse steps. FISMO derives matched map-back for a Fisher-structured spectral trust region, yielding inverse half powers inside and outside the polar map (Xu et al., 2026). G O -MUON uses the square root of the K-FAC block, yielding inverse fourth powers when the observed-label factors admit the conditional Fisher interpretation. Both are instances of matched spectral map-back at different exponents. Corollary 2 gives positive source alignment for any cached SPD maps, independent of their Fisher semantics. Other recent methods insert cheaper curvature information at different points in the Muon pipeline. Newton–Muon corrects the input side using activation moments before orthogonalization (Du & Su, 2026); Muon2 uses Adam-like second moments before orthogonalization and couples this with efficient polar computation (Liu et al., 2026); and MALT uses lightweight diagonal preconditioning (Wu et al., 2026). Empirical studies of whitening optimizers similarly ask which parts of whitening, momentum, and normalization account for their gains (Frans et al., 2025). G O -MUON makes this orientation change coordinate-consistent: the preconditioned matrix enters the polar map and returns through the matching tangent map. The exact oracle is general; the OAS K-FAC-root field is the G O -MUON instantiation. Table 1 compares update construction and curvature reuse across representative structured matrix optimizers. 4.4

A MORTIZED AND STALE CURVATURE

Prior second-order methods separate the frequency of gradient updates from the frequency of curvature work. The original online K-FAC algorithm updates moving-average factors from minibatches while reusing a factorization for multiple iterations (Martens & Grosse, 2015). Distributed K-FAC subsequently moved inverse computation and factor-statistic computation to asynchronous workers, explicitly accepting stale curvature in exchange for throughput (Ba et al., 2017). Large-scale implementations made both timescales configurable: they reduce factor formation as well as eigendecomposition frequency (Osawa et al., 2019; Pauloski et al., 2020). KAISA applies this separation to Transformer linear layers, including BERT, with distinct factor-update and eigendecompositionupdate intervals (Pauloski et al., 2021). These systems schedule factor updates and eigendecompositions independently, including gated activation and backward capture at Transformer scale. The Shampoo lineage makes a related separation. Original Shampoo accumulates matrix moments more often than it computes inverse roots (Gupta et al., 2018). Scalable Shampoo performs roots asynchronously and also exposes a separate interval for skipping gradient-statistic accumulation (Anil et al., 2020). Distributed Shampoo typically updates factors every step while refreshing inverse roots 8

S ECOND -O RDER M UON D ONE R IGHT

Table 2: Training-loss trajectories against Muon controls. Tiny Shakespeare uses three fresh paired seeds; PTB uses three hash-matched paired seeds. “Traj. GM” is the geometric mean of the three window means. Lower is better. Corpus

Method

Matrix LR 201–300 501–600 901–1000 Traj. GM

Tiny Shakespeare Muon+Q/K control Tiny Shakespeare G O -MUON (lazy K = 4)

0.018 0.080

1.70001 1.63052

1.50427 1.43613

1.34825 1.31465

1.51072 1.45471

Penn Treebank Penn Treebank

0.020 0.080

1.46773 1.45689

1.28084 1.28167

1.17560 1.17000

1.30257 1.29757

Muon G O -MUON (lazy K = 4)

periodically (Shi et al., 2023). SOAP treats eigenbasis frequency as an explicit systems parameter, while later analysis separates stale eigenvalues from stale eigenvectors and develops adaptive refresh criteria (Vyas et al., 2024; Eschenhagen et al., 2025). Together, these methods separate factor, matrix-function, and gradient-update timescales. G O -MUON couples quarter-power matched geometry to a two-timescale sample-and-hold estimator. A reuse step skips factor decompositions, activation retention, output-gradient hooks, moment construction, and shrinkage. Equation (6) guarantees strictly positive alignment with the supplied K momentum source for every cached SPD pair. The βB coefficient tracks elapsed decay, while Proposition 6 shows that sparse observation raises stationary factor variance. The contribution is the measured compute–statistics tradeoff, not denoising from staleness.

5

E XPERIMENTS

5.1

O PTIMIZATION AGAINST M UON CONTROLS

We train three-layer causal Transformers of width 128 with four attention heads, SwiGLU width 384, sequence length 128, and 2,048 tokens per update for 1,000 updates. We report mean pre-update training cross-entropy in three fixed 100-update windows and their geometric mean. The PTB comparison uses three hash-matched paired seeds. Here Muon is the standard update D = Polar(M ), implemented with the same exact polar map and scalar schedule as G O -MUON, at learning rate 0.020. Ranking the recorded G O -MUON sweep by the three-window geometric mean selects learning rate 0.080. This three-window ranking was introduced after the original late-window rule, which selected 0.16, so we treat the hash-matched comparison as exploratory. The Tiny Shakespeare comparison uses three fresh paired seeds; its frozen campaign control is Muon with the campaign’s Q/K partner metric, headwise Q/K treatment, and registered scalar schedule. No validation or test split is read. On Tiny Shakespeare, G O -MUON is lower in all nine paired seed–window comparisons. Its reductions in the early, middle, and late windows are 4.09%, 4.53%, and 2.49%, respectively, giving a 3.71% reduction in the trajectory geometric mean. Peak learning rate 0.080 was selected by a broader sweep on tuning seed 17401, where it obtains window losses 1.62170, 1.43461, and 1.31396, with trajectory geometric mean 1.45132 versus 1.51073 for its paired campaign control—a 3.93% reduction. Table 2 evaluates that frozen selection on fresh seeds 17402–17404. On PTB, the three-window-selected G O -MUON lowers the early and late windows by 0.74% and 0.48%. The middle window is effectively tied: Muon is lower by 0.00083. G O -MUON lowers the three-window trajectory geometric mean by 0.38% and wins that per-seed summary on all three matched pairs. 5.2

G ROKKING SPEED ON CORRECTED MODULAR ADDITION

We next measure the onset of grokking rather than terminal loss. For modulus p, the dataset contains all p2 ordered pairs and uses a deterministic 50/50 split, with the training count rounded down, into disjoint train and test subsets. The model is a three-layer causal Transformer of width 64 with four attention heads, SwiGLU width 192, and RoPE. Its input is [a, b, CLS]: the p-row numerical embedding is tied to the p-class output head, while CLS is an independent learned vector and is not an output class. Thus the read-side control token is not also an unused output class. 9

S ECOND -O RDER M UON D ONE R IGHT

Table 3: Grokking onset on corrected independent-CLS modular addition, using one screening seed per modulus. “Grok” is the first of five consecutive evaluations with held-out accuracy at least 99%; evaluations occur every ten updates. “Delay” subtracts the first update with 99% training accuracy. Lower is better for Train, Grok, and Delay; Speedup is the Muon/G O -MUON ratio of Grok steps. The endpoint measures onset, while retention is separate. Task

Method

Matrix LR

Train 99%

Grok 99%

Delay

Speedup

mod 103 mod 103

Muon G O -MUON

0.160 0.080

66 64

2320 290

2254 226

1.0× 8.0×

mod 107 mod 107

Muon G O -MUON

0.160 0.080

65 68

4520 220

4455 152

1.0× 20.5×

Table 4: PTB character-LM training loss over four paired seeds. “Wins” counts seeds where lazy K = 4 is lower than refresh 1. Update window

Refresh 1

Lazy K = 4

Lazy wins

201–300 501–600 901–1000

1.45513 1.28378 1.17238

1.45549 1.28259 1.16917

2/4 2/4 4/4

Training is full batch for 5,000 updates. Both optimizers use momentum 0.8, 50-step linear warmup followed by a constant learning rate, and identical AdamW updates for the tied numerical embedding, CLS, and norm parameters. G O -MUON uses K = 4 and βB = 0.97. A frozen screen over learning rates {0.08, 0.16} for G O -MUON and {0.16, 0.32} for Muon selects 0.080 and 0.160, respectively, under the grokking criterion below. The screen uses one deterministic split and initialization seed for each modulus. We evaluate the held-out subset every ten updates and define the grokking step as the first of five consecutive evaluations with test accuracy at least 99%. Table 3 also reports, as a memorization diagnostic, the first update with 99% training accuracy and the resulting grokking delay. Training-set fitting occurs at nearly the same update for the two optimizers; the separation appears in the transition to held-out accuracy. On mod 103, G O -MUON groks at update 290 versus 2320 for Muon, using one eighth as many updates. On mod 107, the corresponding crossing is 220 versus 4520, or 1/20.5 as many updates. The mod-107 G O -MUON trajectory later leaves the 99% region, with below-threshold evaluations occurring as late as update 4470 before it returns at update 4480. This is a post-grok retention event and does not redefine the time-to-grok endpoint. 5.3

M ATCHED MAP - BACK ON P ENN T REEBANK

We first isolate matched map-back against the registered halfspace baseline. The baseline uses an inside-only inverse-half-power score and a spectral half-space safeguard; the matched arm uses Equation (4) with α = 1/4 and the Frobenius graft. Both use the PTB training characters only, four paired seeds, peak matrix learning rate 0.08, momentum 0.8, and backward EMA 0.97. In the preregistered updates 201–300 window, mean pre-update training loss is 1.48882 for halfspace and 1.45513 for matched map-back, a relative reduction of 2.26%. The matched arm is lower on all four paired seeds. The allocation was requeued before the final preregistered window; in every available common 50-step terminal window, matched map-back remains lower on all four seeds. We do not impute the missing 901–1000 endpoint. 5.4

L AZY REFRESH ON P ENN T REEBANK

We next compare per-step quarter-power geometry with G O -MUON at K = 4 under a fresh, complete 1,000-update paired run. Initialization, minibatch stream, scalar controls, and schedule are identical within each seed. Table 4 reports the preregistered mean pre-update training-loss windows; no validation or test split is read by this protocol. 10

S ECOND -O RDER M UON D ONE R IGHT

The early window is operationally tied. Lazy refresh is lower by 0.00119 in the middle window and 0.00320 in the late window, winning all four late pairs. Its geometric-mean step-time ratio is 0.798×, a reduction of 20.2%. Both arms record zero negative momentum-source and current-gradient alignments. These results show that the noisier sparse factor estimator can coexist with a favorable training trajectory; they do not identify blockwise coherence as the cause.

6

C ONCLUSION

The first-principles object in G O -MUON is an SPD-weighted spectral oracle: C = P B M PA ,

D0 = PB Q(C)PA ,

⟨M, D0 ⟩ = ∥C∥∗ .

This identity holds for every nonzero source and every supplied SPD pair, without Fisher or freshness assumptions. G O -MUON chooses PA = A−1/4 and PB = B̄ −1/4 from OAS observed-label activation/backward factors; under predictive consistency, the backward factor approaches the softmax Fisher/GGN factor. The Frobenius graft retains the oracle ray while quarter powers limit worst-case sensitivity to factor conditioning. Four-step refresh is a separate systems and statistical choice. It preserves elapsed decay and slowsignal lag, increases stationary factor variance, and removes factor capture and matrix functions from three of every four updates. On the frozen PTB training protocol, the three-window-selected G O -MUON improves the hash-matched trajectory summary over Muon, while the refresh trade reduces mean step time and improves the late training window. The corrected modular-addition screen shows a distinct empirical pattern on its screening seeds: training-set fitting time is nearly unchanged, while sustained-99% grokking onset is reached at updates 290 and 220, versus 2320 and 4520 for Muon. The resulting picture is explicit: matched map-back provides coordinate-consistent spectral optimization, factor estimation chooses the geometry, and lazy refresh amortizes its cost.

R EFERENCES Natalie Abreu, Nikhil Vyas, Sham Kakade, and Depen Morwani. The potential of second-order optimization for LLMs: A study with full gauss–newton. arXiv preprint arXiv:2510.09378, 2025. URL https://arxiv.org/abs/2510.09378. Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018, 2020. URL https://arxiv. org/abs/2002.09018. Jimmy Ba, Roger Grosse, and James Martens. Distributed second-order optimization using kroneckerfactored approximations. In International Conference on Learning Representations, 2017. URL https://openreview.net/forum?id=SkkTMpjex. Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. arXiv preprint arXiv:2409.20325, 2024. URL https://arxiv.org/abs/2409.20325. Da Chang, Qiankun Shi, Lvgang Zhang, Yu Li, Ruijie Zhang, Yao Lu, Yongxiang Liu, and Ganzhao Yuan. MuonEq: Balancing before orthogonalization with lightweight equilibration. arXiv preprint arXiv:2603.28254, 2026. URL https://arxiv.org/abs/2603.28254. Yilun Chen, Ami Wiesel, Yonina C. Eldar, and Alfred O. Hero. Shrinkage algorithms for MMSE covariance estimation. IEEE Transactions on Signal Processing, 58(10):5016–5029, 2010. URL https://arxiv.org/abs/0907.4698. Zhehang Du and Weijie Su. The newton-muon optimizer. arXiv preprint arXiv:2604.01472, 2026. URL https://arxiv.org/abs/2604.01472. Runa Eschenhagen, Aaron Defazio, Tsung-Hsien Lee, Richard E. Turner, and HaoJun Michael Shi. Purifying shampoo: Investigating shampoo’s heuristics by decomposing its preconditioner. In Advances in Neural Information Processing Systems, volume 38, 2025. URL https://papers.nips.cc/paper_files/paper/2025/hash/ f14f4eda29a74c02c803699a09529bb9-Abstract-Conference.html. 11

S ECOND -O RDER M UON D ONE R IGHT

Kevin Frans, Pieter Abbeel, and Sergey Levine. What really matters in matrix-whitening optimizers? arXiv preprint arXiv:2510.25000, 2025. URL https://arxiv.org/abs/2510.25000. Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In Proceedings of the 35th International Conference on Machine Learning, pp. 1842–1850, 2018. URL https://proceedings.mlr.press/v80/gupta18a.html. Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks. Online write-up and software, 2024. URL https://kellerjordan.github.io/posts/muon/. Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. URL https://arxiv.org/abs/1412. 6980. Frederik Kunstner, Lukas Balles, and Philipp Hennig. Limitations of the empirical fisher approximation for natural gradient descent. In Advances in Neural Information Processing Systems, volume 32, 2019. URL https://arxiv.org/abs/1905.12558. Tim Large, Yang Liu, Minyoung Huh, Hyojin Bahng, Phillip Isola, and Jeremy Bernstein. Scalable optimization in the modular norm. In Advances in Neural Information Processing Systems, 2024. URL https://arxiv.org/abs/2405.14813. Wu Lin, Felix Dangel, Runa Eschenhagen, Kirill Neklyudov, Agustinus Kristiadi, Richard E. Turner, and Alireza Makhzani. Structured inverse-free natural gradient descent: Memory-efficient & numerically-stable KFAC. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 29974–29991, 2024. URL https://proceedings.mlr.press/v235/lin24f.html. Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for LLM training. arXiv preprint arXiv:2502.16982, 2025. URL https://arxiv.org/abs/2502.16982. Ziyue Liu, Ruijie Zhang, Zhengyang Wang, Yequan Zhao, Yupeng Su, Zi Yang, and Zheng Zhang. Muon2 : Boosting muon via adaptive second-moment preconditioning. arXiv preprint arXiv:2604.09967, 2026. URL https://arxiv.org/abs/2604.09967. Kevin Luk and Roger Grosse. A coordinate-free construction of scalable natural gradient. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2999–3008, 2018. URL https://arxiv.org/abs/ 1808.10340. James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1–76, 2020. URL https://arxiv.org/abs/1412.1193. James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In Proceedings of the 32nd International Conference on Machine Learning, pp. 2408– 2417, 2015. URL https://arxiv.org/abs/1503.05671. Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Large-scale distributed second-order optimization using kronecker-factored approximate curvature for deep convolutional neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12359–12367, 2019. URL https://openaccess. thecvf.com/content_CVPR_2019/html/Osawa_Large-Scale_Distributed_ Second-Order_Optimization_Using_Kronecker-Factored_Approximate_ Curvature_for_Deep_CVPR_2019_paper.html. J. Gregory Pauloski, Zhao Zhang, Lei Huang, Weijia Xu, and Ian T. Foster. Convolutional neural network training with distributed K-FAC. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1–12, 2020. doi: 10.1109/SC41405.2020.00098. URL https://arxiv.org/abs/2007.00784. 12

S ECOND -O RDER M UON D ONE R IGHT

J. Gregory Pauloski, Qi Huang, Lei Huang, Shivaram Venkataraman, Kyle Chard, Ian Foster, and Zhao Zhang. KAISA: An adaptive second-order optimizer framework for deep neural networks. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. Association for Computing Machinery, 2021. doi: 10.1145/3458817. 3476152. URL https://arxiv.org/abs/2107.01739. Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher. Training deep learning models with norm-constrained LMOs. arXiv preprint arXiv:2502.07529, 2025. URL https://arxiv.org/abs/2502.07529. Artem Riabinin, Egor Shulgin, Kaja Gruntkowska, and Peter Richt ’arik. Gluon: Making muon & scion great again! bridging theory and practice of LMO-based optimizers for LLMs. arXiv preprint arXiv:2505.13416, 2025. URL https://arxiv.org/ abs/2505.13416. Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data-parallel PyTorch implementation of the distributed shampoo optimizer for training neural networks at-scale. arXiv preprint arXiv:2309.06497, 2023. URL https://arxiv.org/abs/2309.06497. Andrey Veprikov, Arman Bolatov, Aleksandr Bogdanov, Samuel Horv ’ath, Aleksandr Beznosikov, Martin Tak ’a v c, and Slavomir Hanzely. Preconditioned norms: A unified framework for steepest descent, quasi-newton and adaptive methods. arXiv preprint arXiv:2510.10777, 2025. URL https: //arxiv.org/abs/2510.10777. Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. SOAP: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409.11321, 2024. URL https://arxiv.org/abs/2409.11321. Tongle Wu, Huanyu Dong, Ying Sun, and Ziye Ma. MALT: Lightweight curvature-aware muon via diagonal preconditioning. arXiv preprint arXiv:2608.05088, 2026. URL https://arxiv. org/abs/2608.05088. Chenrui Xu, Wenjing Yan, and Ying-Jun Angela Zhang. FISMO: Fisher-structured momentumorthogonalized optimizer. arXiv preprint arXiv:2601.21750, 2026. URL https://arxiv. org/abs/2601.21750. Ruijie Zhang, Yequan Zhao, Ziyue Liu, Zhengyang Wang, and Zheng Zhang. Muon+: Towards better muon via one additional normalization step. arXiv preprint arXiv:2602.21545, 2026a. URL https://arxiv.org/abs/2602.21545. Yechen Zhang, Shuhao Xing, Junhao Huang, Kai Lv, Yunhua Zhou, Xipeng Qiu, Qipeng Guo, and Kai Chen. Mousse: Rectifying the geometry of muon with curvature-aware preconditioning. arXiv preprint arXiv:2603.09697, 2026b. URL https://arxiv.org/abs/2603.09697.

A

P ROOFS AND ADDITIONAL GUARANTEES

A.1

P ROOF OF P ROPOSITION 5

Write PB = B −α , PA = A−α , and let Q be the compact full-SVD support polar, so ∥Q∥F = Nuclear-norm multiplication inequalities give ∥PB M PA ∥∗ ≥ σmin (PB )σmin (PA )∥M ∥∗ , while ∥PB QPA ∥F ≤ Substitution into

√

k∥PB ∥op ∥PA ∥op .

D E √k∥P M P ∥ B A ∗ b M, Dα = ∥PB QPA ∥F 13

√

k.

S ECOND -O RDER M UON D ONE R IGHT

gives D E b α ≥ σmin (PB )σmin (PA ) ∥M ∥∗ = K −α ∥M ∥∗ . M, D AB σmax (PB )σmax (PA ) The nonzero singular values of a semi-orthogonal Q are all one. Singular-value product inequalities therefore give σmax (PB QPA ) ≤ σmax (PB )σmax (PA ), σk (PB QPA ) ≥ σmin (PB )σmin (PA ). α Their ratio is at most KAB . The positive graft does not change this ratio.

For completeness, Equation (19) follows by writing µ = tr(S)/d. Since S ⪰ 0, λmin (OAS(S)) ≥ λµ, A.2

λmax (OAS(S)) ≤ [d − (d − 1)λ]µ.

A SCALE - FREE CACHED - ORACLE BOUND

Proposition 7 (Cached oracle under coordinate drift). Let qP be the current weighted norm from e be any positive rescaling of an exact unit-radius oracle for cached maps Equation (2), and let D e e PB , PA . Define Θ by Equation (7). Then D E e M, D ≥ Θ−1 . (36) e qP (D)∥P M P ∥ B A ∗ Thus, after rescaling to its actual radius in the current ball, the cached direction achieves at least a Θ−1 fraction of the current optimum. The ratio is unchanged by the Frobenius graft. Proof. Let qe(D) = ∥PeB−1 DPeA−1 ∥op ,

L = PB−1 PeB ,

R = PeA PA−1 .

For U = PeB−1 DPeA−1 , qP (D) = ∥LU R∥op , and hence a qe(D) ≤ qP (D) ≤ b qe(D), where a = σmin (L)σmin (R),

b = σmax (L)σmax (R).

The corresponding unit-ball inclusions imply that the cached support value e h(M ) = ∥PeB M PeA ∥∗ e e satisfies h(M ) ≥ a∥PB M PA ∥∗ . An exact cached unit oracle has pairing h(M ) and current radius at most b. Its normalized efficiency is therefore at least a/b. Since κ2 (R) = κ2 (PA−1 PeA ), a/b = Θ−1 . Positive rescaling cancels from the ratio. The bound uses the singular condition numbers of the actual coordinate changes. A simpler factorlevel power law requires additional commutativity or functional-calculus assumptions and is not used here. A.3

P ROOF OF P ROPOSITION 6

Let FX (dω) be the matrix-valued spectral measure of the centered vector process vec(Xt ). At refresh boundaries, the sparse and per-step transfer functions with respect to the original time index are GK (ω) =

K 1 − βB , K −iKω 1 − βB e

H1 (ω) =

1 − βB . 1 − βB e−iω

Factorizing 1 − (βB e−iω )K gives PK−1 r β |GK (ω)| = P r=0 B ≥ 1. K−1 |H1 (ω)| β r e−irω r=0

14

B

S ECOND -O RDER M UON D ONE R IGHT

Therefore

Z π Cov(vec Yj ) − Cov(vec ZjK ) =

(|GK |2 − |H1 |2 )FX (dω) ⪰ 0.

−π

For IID observations, the standard AR(1) variance formula immediately yields Equation (25). For an affine signal xt = a + vt, the steady refresh-boundary estimate is Yj = xjK − v

K KβB . K 1 − βB

Holding this value through phase s ∈ {0, . . . , K − 1} adds s steps of delay. Averaging over phases gives Equation (26).

15

Related documents

• Hamiltonian-Type Spectral Functionals for Time-Dependent Perturbed Hodge--de Rham Operators#1061311
• A remark on the fourth order Q curvature on manifolds with dimension at least 5#799308
• A remark on the fourth order Q curvature on manifolds with dimension at least 5#797885
• The Segre invariant and Brill-Noether theory on Hirzebruch surfaces#781062
• The Segre invariant and Brill-Noether theory on Hirzebruch surfaces#779687