Matrix optimizers such as Muon transform each momentum matrix through an approximate orthogonalization, typically implemented by a small number of Newton-Schulz matrix multiplications. The quality and cost of this approximation depend strongly on the singular-value spectrum of its input, yet existing implementations use the same fixed polynomial routine for every layer and throughout training. We show that this uniform treatment is unnecessary: the computations in the Newton-Schulz method already reveal enough information to make the method adaptive. The Gram matrices formed inside Newton-Schulz iterations yield spectral moments through inexpensive scalar reductions, requiring no additional matrix multiplications. From these moments, we recover an estimate of the empirical singular-value distribution and use it to select a polynomial routine specialized to the current matrix. This turns Newton--Schulz orthogonalization into a spectrum-adaptive procedure that responds to differences across both layers and training time. On saved momentum matrices, spectral estimation substantially reduces orthogonalization error at a fixed iteration budget or reaches the same accuracy with fewer iterations, and in GPT pretraining up to 1B parameters it lowers the validation loss of two matrix optimizers. Our results suggest that matrix-function operations inside optimizers need not be designed for a conservative worst-case spectrum: they can cheaply measure the spectrum they are already processing and specialize computations accordingly.
Figures & tables
Figure 1: Newton–Schulz can measure its input and specialize its own computation. (a) Full singular-value CDFs show different momentum spectra for attention Q and MLP-up at layer 5 , step 16 k. (b) Read moments from existing matrix products, fit the spectrum, and select a routine for subsequent steps. (c) On 540 saved MLP matrices, spectral estimation reduces mean polar RMS error by 90× relative to fixed Polar Express (PE) at K=5 iterations, scored above 10−5 in Frobenius-normalized coordinates.
Figure 2: Wider design intervals require more iterations. Ideal PE on [ρ,1] . (a) Optimal uniform error versus depth. (b,c) Composed responses at K=5 and K=8 . Shading marks 1±εK⋆(ρ) .
Figure 3: Spectra vary across layers and training. CDFs of the 160 M Muon baseline. (a,b) Q and MLP-up across layers at step 2 k. (c,d) Layer 11 across checkpoints. The dotted line marks 10−5 .
Figure 4: A conservative interval spends unnecessary work on narrow spectra. Q and MLP-up at layer 5 , step 16 k. The common interval is [10−7,1] , and the exact interval is each matrix’s nonzero spectral range, as used by an oracle. Labels show the first depth reaching 10−2 .
Figure 5: Moment fits recover useful spectral structure. Four matrices at layer 5 , step 16 k. The top row shows the density in log10(σ/∥M∥F) and the bottom row the CDF, both conditioned on σ/∥M∥F≥10−5 .
Figure 6: One set of predicted scores supports both objectives. MLP-up, layer 5 , step 16 k. (a) Predicted errors, with the best depth- 8 endpoint (circle) and the shortest routine meeting τ=10−2 (star). (b) Predicted and SVD-evaluated errors of the selected routines. Values below 10−6 are shown at that floor.
Figure 7: Offline accuracy and depth. Means over attention and MLP matrices, scored above 10−5 in Frobenius-normalized coordinates. (a–b) Polar RMS error at fixed depth ( Table 2 ). (c–d) Iterations needed to meet a target τ ( Table 2 ). Ours meets τ in predicted RMS error, the oracle in worst-case error, and fixed PE is the uniform depth that meets τ .
RMS error
Worst-case error
K
Fixed PE
Ours
Oracle
Fixed PE
Ours
Oracle
Attention
5
0.667
0.472
0.653
0.976
0.982
0.966
8
0.159
0.0431
0.104
0.222
0.353
0.161
MLP
5
0.626
0.00694
0.0179
0.976
0.0643
0.0266
8
0.162
0.00101
1.13×10−6
0.222
0.00153
1.78×10−6
Table 1: Fixed-depth polar error. Mean per-matrix RMS and worst-case error (minimax oracle).
Table 9
Figure 8: Training loss relative to the fixed baseline. Training loss minus the fixed baseline at K=5 in 50 -step windows, where negative is better.
Appendix figures & tables29 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9: Normalization explains part of the gain. (a) Ratio of each upper bound to the largest singular value, by matrix type. (b) Fixed PE, fixed PE with the Gram-power scale (scaling only), and spectral estimation.
Figure 10: Different spectral measures can match the same measurements. Maximum-entropy and extreme-point grid measures for representative attention and MLP spectra. The top row shows the mass on the estimator grid and the bottom row the conditional CDF on the scoring interval.
Component
Setting
Scale
1.005s8 , rounded up to the grid 10−j/24 , j=0,…,15
Scoring cutoff
x⋆=10−5 in Frobenius-normalized coordinates
Spectral grid
zero, 24 points below the cutoff, 120 in the scoring interval
Measurements
four moments per iteration ( 0 – 9 ) and the final second moment
Entropy tolerance
κ=3×10−4
Row-scale floor
dmin=10−10
Appendix
Table 5: Estimator and candidate settings.
Primary run
Held-out run
Dataset
FineWeb
Nemotron-CC v2 High-Quality
Model
12 layers, width 768
24 layers, width 768 ( 160 M Muon baseline)
Checkpoints
every 2,000 steps, 2 k– 30 k
step 1 , then every 500 steps to 4,500
Matrices
720 attention, 540 MLP
1,680
Matrix types
Q, K, V, O, MLP up, MLP gate, MLP down
Muon momentum
0.95
0.9586
Appendix
Table 6: Offline data.
Figure 11: Fitted and SVD spectra across layers (step 16 k), with the SVD in grey and the moment fit in purple. Both are restricted to σ/∥M∥F≥10−5 and smoothed in log10σ . The numbers give the largest gap between the two CDFs.
Figure 12: Fitted and SVD spectra across training (layer 5 ). Same conventions as Figure 11 .
Regret
MLP misses (%)
Attention
MLP
τ
Rule on the fit
K=5
K=8
K=5
K=8
10−1
10−2
10−3
Minimax
0.0014
0.0038
0.283
0.0069
20.0
18.7
21.1
RMS (ours)
3×10−5
0.00068
0.0036
0.0010
0.0
5.7
5.4
Appendix
Table 7: Scoring rules on the fitted spectrum. Regret is the mean error of the selected routine minus that of the matching oracle, in the rule’s own criterion (worst-case error for minimax, RMS error for ours). Misses are the share of MLP matrices whose true error exceeds τ under the target rule.
RMS error
Worst-case error
K
Fixed PE
Ours
Oracle
RMS oracle
Fixed PE
Ours
Oracle
Attention
4
0.729
0.618
0.718
0.618
0.994
0.995
0.992
5
0.667
0.472
0.653
0.472
0.976
0.982
0.966
6
0.604
0.311
0.583
0.311
0.904
0.934
0.868
7
0.465
0.158
0.402
0.157
0.665
0.774
0.582
Appendix
Table 8: Fixed-depth RMS and worst-case error. Means over attention and MLP matrices. Oracle is the minimax oracle, and the RMS oracle selects by exact RMS error. All methods use the same candidates and scoring interval.
RMS error
Worst-case error
K
Fixed PE
Ours
Oracle
RMS oracle
Fixed PE
Ours
Oracle
Attention
5
0.640
0.153
0.348
0.149
0.976
0.550
0.516
8
0.161
0.0123
0.0288
0.00977
0.222
0.126
0.0444
MLP
5
0.623
0.0315
0.0247
0.0214
0.976
0.0539
0.0373
Appendix
Table 9: Held-out run. Mean errors over all checkpoints with unchanged settings. Oracle is the minimax oracle.
Figure 13: Fewer measurements. The left and center columns use measurements from only the first iterations of the measurement pass, and the right column fewer moment orders at every iteration. Colors give x⋆ , the scoring cutoff and smallest design endpoint. Solid lines show spectral estimation and dashed lines the RMS oracle.
Figure 14: Lower scoring cutoffs. Each row repeats Figure 7 with the RMS oracle at the cutoff x⋆ shown on the left. Panels (a–b) show RMS error at fixed depth and (c–d) the mean depth at target τ , capped at the largest allowed depth (dotted). At the lower cutoffs, almost no attention matrix meets τ≤10−2 within the cap, for either method.
Setting
Value
Attn K=5
Attn K=8
MLP K=5
MLP K=8
Default
Table 5
0.472
0.0431
0.00694
0.00101
Entropy tolerance κ
0 (best fit only)
0.472
0.0437
0.00729
0.000622
10−4
0.472
0.0430
0.00684
0.000777
10−3
0.472
0.0434
0.00862
0.00140
3×10−3
0.472
0.0439
0.0131
0.00216
Grid, scoring interval
60
0.472
0.0435
0.0199
0.00223
Appendix
Table 10: Estimator settings. Mean RMS error of the selected routine, changing one setting at a time. RMS-oracle rows use the corresponding candidate set.
Mean depth
Misses (%)
τ
m
Attn
MLP
Attn
MLP
Max error
10−1
1
7.91
4.23
1.9
0.0
0.124
1.5
8.19
4.44
0.0
0.0
0.0780
10
9.00
6.09
0.0
0.0
0.0567
RMS oracle
7.83
4.02
0.0
0.0
0.0995
10−2
1
9.00
6.09
0.0
5.7
0.0567
Appendix
Table 11: Target rule with a safety margin m . A candidate is accepted if its predicted RMS error is at most τ/m . Misses are the fraction of matrices whose true RMS error exceeds τ .
K
Current
Previous
First
Every fourth
RMS oracle
Oracle
Fixed PE
5
0.275
0.275
0.278
0.275
0.273
0.381
0.650
8
0.0254
0.0257
0.0261
0.0256
0.0246
0.0601
0.160
Appendix
Table 12: Reusing earlier choices. Mean RMS error on checkpoints from step 4 k. Columns give the checkpoint whose choice is used, namely the current one, the previous one, the first one, or one updated only at every fourth checkpoint. Oracle is the minimax oracle.
K
Per type
Ours
RMS oracle
Oracle
5
0.286
0.283
0.282
0.383
8
0.0290
0.0273
0.0266
0.0639
Appendix
Table 13: Per-type versus per-matrix selection. Mean RMS error on the later checkpoints ( 16 k– 30 k), with the per-type endpoints chosen on the earlier ones. Oracle is the minimax oracle.
Method
Extra products
Work that grows with M
Fixed work
Ours
none
≈3r2 per iteration
fit, scoring
SVD (singular values)
–
O(r2n)
–
Gram eigenvalues
MM⊤
O(r2n+r3)
–
Appendix
Table 14: Added cost of each spectrum estimator for M∈Rr×n , r≤n . One Newton–Schulz iteration costs 2r2n+r3 .
Figure 15: GPU time per matrix versus matrix size (one H200, median). Ours denotes the compiled reductions. Circles mark square matrices, triangles wide matrices, and diamonds the matrices of Table 15 . The Newton–Schulz pass runs in bf16, and the SVD and the Gram eigenvalues in fp32. Matrices are random Gaussian.
Model
Matrix
Shape
NS pass
Ours
Gram eig.
SVD
DeepSeek-V4
routed expert
3072×7168
4.9
0.96
49
480
attention output
7168×16384
62
1.1
358
3481
GLM-5
attention output
6144×16384
45
1.1
252
2461
dense FFN
6144×12288
36
1.1
246
2473
Kimi K3
attention output
7168×12288
50
1.1
349
3446
dense FFN
7168×33792
114
1.3
392
3667
Appendix
Table 15: Matrix shapes of recent models trained with Muon ( DeepSeek-AI, 2026 ; GLM-5 Team, 2026 ; Kimi Team, 2026 ) and GPU time per matrix in milliseconds (one H200, median). The NS pass runs ten bf16 iterations, ours is the compiled reductions, and the last two columns compute the eigenvalues of MM⊤ and the singular values. The DeepSeek-V4 attention output is the second matrix of its grouped output projection.
Reduction
Function
\trHj/n , ∥Hj∥F2/n
zj2 , zj4
\trYj/n , ∥Yj∥F2/n
γzj , γ2zj2
⟨Yj,Hj⟩F/n , ∥YjHj∥F2/n
γzj3 , γ2zj6
\trZj/n , ∥Zj∥F2/n
zj/γ , zj2/γ2
⟨Zj,Hj⟩F/n , ∥HjZj∥F2/n
zj3/γ , zj6/γ2
Appendix
Table 16: Measurements from the coupled iteration. Each function is averaged over the distribution of γ .
Figure 16: Inverse square roots. Mean relative RMS error on regularized momentum Gram matrices for the fixed cubic routine, spectral estimation, and the RMS oracle over the same candidates.
Component
Setting
Test matrices
MˉMˉ⊤+εI , ε∈{10−10,10−8}
Endpoints
49 , log-spaced in [10−5,10−1]
Candidate depths
K=4,…,18
Grid
one low point, 150 log-spaced
Entropy tolerance
κ=3×10−4
Measurement depth
15 ( ε=10−10 ) or 13 ( ε=10−8 )
Appendix
Table 17: Inverse-root test settings.
Setting
160 M
300 M
1 B
Layers / width
24 / 768
24 / 1024
24 / 2048
Query heads / KV heads
12 / 3
16 / 4
32 / 8
Head dimension / MLP width
64 / 2304
64 / 3072
64 / 6144
Non-embedding / total params
163.1 M / 262.9 M
289.9 M / 423.0 M
1159.3 M / 1425.6 M
Attention / MLP
grouped-query attention with a per-head sigmoid gate / SwiGLU
Normalization / positions
pre-RMSNorm and final RMSNorm ( ϵ=10−6 ) / RoPE ( θ=104 )
Appendix
Table 18: Model and training settings.
Setting
Muon
OKLS
Peak learning rate
0.01202
0.09434
Momentum β or β1
0.9586
0.9684
Factor decay β2
–
0.9482
Factor diagonal shift ϵ
–
10−9
Weight decay λwd
1.008×10−4
0.0303
Initialization scale σ∗
0.1877
0.07539
Appendix
Table 19: Hidden-matrix optimizer settings. The decay coefficient λwd enters as θ←(1−λwdηt2/ηpeak)θ .
Muon
OKLS
Scoring cutoff
10−5
10−7 (relative)
Candidates
49 PE endpoints in [10−5,10−1] at the run’s depth K
cubic, 12 endpoints per decade in [10−7,1] , at the run’s depth K
Measurement pass
10 iterations, 41 measurements
10 iterations, 94 measurements
Measurement period T
200 ( 500 also at 160 M and 300 M)
200 ( 500 also at 160 M and 300 M)
Fixed baseline
PE[10−5,1]
cubic routine for eigenvalues above 10−7λmax
Oracle
minimax choice from the exact spectrum every 500 steps
Appendix
Table 20: Pretraining routine settings.
Optimizer
Size
K
Fixed
T=200
T=500
Muon
160M
5
2.761
2.751
2.748
Muon
300M
5
2.662
2.655
2.651
OKLS
160M
5
2.687
2.673
2.671
OKLS
300M
5
2.597
2.590
2.591
Appendix
Table 21: Measurement period. Final validation loss of our method with T=200 and T=500 .
Size
PE[10−5,1]
PE[10−3,1]
Ours
(baseline)
(default)
160M
2.761
2.765
2.751
300M
2.662
2.669
2.655
Appendix
Table 22: Practitioner endpoint. Muon, K=5 , final validation loss.
Endpoint
RMS error
Optimizer
K
T
N
Agreement (%)
Distance
Fixed
RMS oracle
Ours
Regret
Muon
4
200
3,864
58.6
0.00
0.6681
0.2226
0.2254
0.0028
Muon
4
500
1,512
61.1
0.00
0.6603
0.2194
0.2219
0.0025
Muon
5
200
3,864
27.3
0.25
0.6359
0.0967
0.1066
0.0099
Muon
5
500
1,512
28.2
0.25
0.6322
0.0724
0.0828
0.0104
OKLS
5
200
4,704
77.8
0.00
0.6267
0.2390
0.2394
0.0004
Appendix
Table 23: Selection quality during training. Agreement is the share of measurements where spectral estimation picks the RMS oracle’s endpoint, distance the median gap between the two endpoints in decades, regret our RMS error minus the RMS oracle’s, and N the number of matrix or factor measurements. 160 M runs.
Figure 17: Selection by matrix type (Muon 160 M, K=5 , T=500 ). (a) Mean RMS error on the recorded exact spectra. (b) Median design endpoint chosen by spectral estimation and by both oracles on the same spectra.
Optimizer
Reductions
Count
Exact
Measured
Muon
all
41
0.0078
0.0090
OKLS
all
94
0.0004
0.0008
OKLS
Hj only
20
0.0019
0.0044
Appendix
Table 24: Moment precision and measurement set. Mean regret against the RMS oracle on saved training matrices. Columns compare moments computed exactly with moments measured during the pass (bf16 for Muon, fp32 for OKLS), and rows the reductions used. The exact columns cover 504 Muon matrices and 1008 OKLS factors, and the measured ones 504 and 56 .
Matrix-valued optimizers such as Muon exploit the spectral structure of neural network updates through Newton--Schulz orthogonalization, but their near-flattening of the singular spectrum discards relative magnitude information across gradient modes. We introduce Soren (Spectral Orthogonal Reshaping), a matrix-valued optimizer that preserves the singular subspaces of the gradient while applying a bounded, monotone sigmoid transformation to its singular values. This smoothly compresses dominant modes without fully flattening the spectrum. We interpret Soren as a positive-definite preconditioned gradient method and establish convergence guarantees under relative smoothness and metric Polyak--Łojasiewicz geometry. To avoid explicit singular value decomposition, we further develop a finite-depth Soft Newton--Schulz (SNS) polynomial realization of the sigmoid spectral map and characterize how its spectral approximation affects the induced convergence geometry. Experiments across LLM pre-training, supervised fine-tuning, and direct preference optimization demonstrate the effectiveness and robustness of Soren against established optimizers.
Dat Phi Van, Ngo Vu Minh, Tuc Nguyen +3
Center for AI Research, VinUniversity, Hanoi, Vietnam · Indiana University, Bloomington, IN, USA · Monash University, Clayton, VIC, Australia
Muon has recently emerged as a strong optimizer for large-scale deep learning, where it reshapes gradient updates through approximate orthogonalization and has been reported to outperform Adam and AdamW in large language model training. Its empirical success has motivated a growing body of theoretical work that interprets Muon as steepest descent under the spectral norm. Yet it remains unclear which of Muon's advantages stem from its update rule itself and which are artifacts of the scale, architecture, and data of modern deep networks. In this work, we isolate the optimizer from these confounding factors by studying Muon on a simple, well-understood, and spectrally structured problem: low-rank matrix factorization. Through a controlled comparison against carefully tuned adaptive baselines, we find that Muon does not consistently outperform AdamW in this setting and that several previously reported advantages are sensitive to hyperparameter choices. Our results provide a more nuanced picture of when spectrum-aware orthogonalization is beneficial and argue for evaluating modern optimizers on controlled problems in addition to end-to-end benchmarks.
Ali Parviz, Gal Mishne, Alex Cloninger
†. Halicioğlu Data Science Institute, UC San Diego. · ∗. Department of Mathematics, UC San Diego.
We show that for tall matrix parameters, like projection matrices in the MLP layers, the Muon update can have row norms that are arbitrarily non-uniform. This can lead to a self-reinforcing feedback loop whereby neurons receive persistently small updates and eventually do not contribute meaningfully to network outputs. This problem is effectively mitigated by an additional row normalization step, but current methods do this in a way that moves the Muon update geometry away from the polar factor of the momentum matrix, which we find is undesirable. We propose Aurora, an optimizer that enforces row-uniformity of matrix parameter updates while respecting Muon's polar factor geometry. Aurora outperforms Muon in our pre-training experiments and, when combined with existing methods, achieves state-of-the-art performance among spectral optimizers on the optimizer track of the modded-nanoGPT speedrun. Additionally, we find that Aurora's empirical gains over Muon scale with the MLP expansion factor, suggesting that Aurora may allow for effective training of very wide MLP layers.