Cost-free Spectral Estimation for Adaptive Newton--Schulz in Matrix Optimizers

Kristi Topollai; Anna Choromanska

cs.LGSep 27, 2026

Cost-free Spectral Estimation for Adaptive Newton--Schulz in Matrix Optimizers

Authors: Kristi Topollai, Anna Choromanska

Organizations: New York University

Abstract

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\times$ 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 $[\rho,1]$ . (a) Optimal uniform error versus depth. (b,c) Composed responses at $K=5$ and $K=8$ . Shading marks $1\pm\varepsilon_{K}^{\star}(\rho)$ .

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

\log_{10}(\sigma/\|M\|_{F})

and the bottom row the CDF, both conditioned on

\sigma/\|M\|_{F}\geq 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

\tau=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

\tau

( Table 2 ). Ours meets

\tau

in predicted RMS error, the oracle in worst-case error, and fixed PE is the uniform depth that meets

\tau

		RMS error			Worst-case error
	$K$	Fixed PE	Ours	Oracle	Fixed PE	Ours	Oracle
Attention	$5$	$0.667$	$\boldsymbol{0.472}$	$0.653$	$0.976$	$0.982$	$\boldsymbol{0.966}$
	$8$	$0.159$	$\boldsymbol{0.0431}$	$0.104$	$0.222$	$0.353$	$\boldsymbol{0.161}$
MLP	$5$	$0.626$	$\boldsymbol{0.00694}$	$0.0179$	$0.976$	$0.0643$	$\boldsymbol{0.0266}$
	$8$	$0.162$	$0.00101$	$\boldsymbol{1.13\!\times\!10^{-6}}$	$0.222$	$0.00153$	$\boldsymbol{1.78\!\times\!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

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.005\,s_{8}$ , rounded up to the grid $10^{-j/24}$ , $j=0,\ldots,15$
Scoring cutoff	$x_{\star}=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	$\kappa=3\times 10^{-4}$
Row-scale floor	$d_{\min}=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 $\sigma/\|M\|_{F}\geq 10^{-5}$ and smoothed in $\log_{10}\sigma$ . 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		$\tau$
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)	$\boldsymbol{3\!\times\!10^{-5}}$	$\boldsymbol{0.00068}$	$\boldsymbol{0.0036}$	$\boldsymbol{0.0010}$	$\boldsymbol{0.0}$	$\boldsymbol{5.7}$	$\boldsymbol{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

\tau

under the target rule.

Attention
	RMS error				Worst-case error
$K$	Fixed PE	Ours	Oracle	RMS oracle	Fixed PE	Ours	Oracle
$4$	$0.729$	$\boldsymbol{0.618}$	$0.718$	$\boldsymbol{0.618}$	$0.994$	$0.995$	$\boldsymbol{0.992}$
$5$	$0.667$	$\boldsymbol{0.472}$	$0.653$	$\boldsymbol{0.472}$	$0.976$	$0.982$	$\boldsymbol{0.966}$
$6$	$0.604$	$\boldsymbol{0.311}$	$0.583$	$\boldsymbol{0.311}$	$0.904$	$0.934$	$\boldsymbol{0.868}$
$7$	$0.465$	$0.158$	$0.402$	$\boldsymbol{0.157}$	$0.665$	$0.774$	$\boldsymbol{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.

Attention
	RMS error				Worst-case error
$K$	Fixed PE	Ours	Oracle	RMS oracle	Fixed PE	Ours	Oracle
$5$	$0.640$	$0.153$	$0.348$	$\boldsymbol{0.149}$	$0.976$	$0.550$	$\boldsymbol{0.516}$
$8$	$0.161$	$0.0123$	$0.0288$	$\boldsymbol{0.00977}$	$0.222$	$0.126$	$\boldsymbol{0.0444}$
MLP
$5$	$0.623$	$0.0315$	$0.0247$	$\boldsymbol{0.0214}$	$0.976$	$0.0539$	$\boldsymbol{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_{\star}

, 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_{\star}

shown on the left. Panels (a–b) show RMS error at fixed depth and (c–d) the mean depth at target

\tau

, capped at the largest allowed depth (dotted). At the lower cutoffs, almost no attention matrix meets

\tau\leq 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 $\kappa$	$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\times 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.

$\tau$	$m$	Attn	MLP	Attn	MLP	Max error
		Mean depth		Misses (%)
$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

\tau/m

. Misses are the fraction of matrices whose true RMS error exceeds

\tau

$K$	Current	Previous	First	Every fourth	RMS oracle	Oracle	Fixed PE
5	$0.275$	$0.275$	$0.278$	$0.275$	$\boldsymbol{0.273}$	$0.381$	$0.650$
8	$0.0254$	$0.0257$	$0.0261$	$0.0256$	$\boldsymbol{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$	$\boldsymbol{0.282}$	$0.383$
8	$0.0290$	$0.0273$	$\boldsymbol{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	$\approx 3r^{2}$ per iteration	fit, scoring
SVD (singular values)	–	$O(r^{2}n)$	–
Gram eigenvalues	$MM^{\top}$	$O(r^{2}n+r^{3})$	–

Appendix

Table 14: Added cost of each spectrum estimator for

M\in\mathbb{R}^{r\times n}

r\leq n

. One Newton–Schulz iteration costs

2r^{2}n+r^{3}

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\times 7168$	$4.9$	$\boldsymbol{0.96}$	$49$	$480$
	attention output	$7168\times 16384$	$62$	$\boldsymbol{1.1}$	$358$	$3481$
GLM-5	attention output	$6144\times 16384$	$45$	$\boldsymbol{1.1}$	$252$	$2461$
	dense FFN	$6144\times 12288$	$36$	$\boldsymbol{1.1}$	$246$	$2473$
Kimi K3	attention output	$7168\times 12288$	$50$	$\boldsymbol{1.1}$	$349$	$3446$
	dense FFN	$7168\times 33792$	$114$	$\boldsymbol{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^{\top}

and the singular values. The DeepSeek-V4 attention output is the second matrix of its grouped output projection.

Reduction	Function
$\tr H_{j}/n$ , $\\|H_{j}\\|_{F}^{2}/n$	$z_{j}^{2}$ , $z_{j}^{4}$
$\tr Y_{j}/n$ , $\\|Y_{j}\\|_{F}^{2}/n$	$\gamma z_{j}$ , $\gamma^{2}z_{j}^{2}$
$\langle Y_{j},H_{j}\rangle_{F}/n$ , $\\|Y_{j}H_{j}\\|_{F}^{2}/n$	$\gamma z_{j}^{3}$ , $\gamma^{2}z_{j}^{6}$
$\tr Z_{j}/n$ , $\\|Z_{j}\\|_{F}^{2}/n$	$z_{j}/\gamma$ , $z_{j}^{2}/\gamma^{2}$
$\langle Z_{j},H_{j}\rangle_{F}/n$ , $\\|H_{j}Z_{j}\\|_{F}^{2}/n$	$z_{j}^{3}/\gamma$ , $z_{j}^{6}/\gamma^{2}$

Appendix

Table 16: Measurements from the coupled iteration. Each function is averaged over the distribution of

\gamma

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	$\bar{M}\bar{M}^{\top}+\varepsilon I$ , $\varepsilon\in\{10^{-10},10^{-8}\}$
Endpoints	$49$ , log-spaced in $[10^{-5},10^{-1}]$
Candidate depths	$K=4,\ldots,18$
Grid	one low point, $150$ log-spaced
Entropy tolerance	$\kappa=3\times 10^{-4}$
Measurement depth	$15$ ( $\varepsilon=10^{-10}$ ) or $13$ ( $\varepsilon=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 ( $\epsilon=10^{-6}$ ) / RoPE ( $\theta=10^{4}$ )

Appendix

Table 18: Model and training settings.

Setting	Muon	OKLS
Peak learning rate	$0.01202$	$0.09434$
Momentum $\beta$ or $\beta_{1}$	$0.9586$	$0.9684$
Factor decay $\beta_{2}$	–	$0.9482$
Factor diagonal shift $\epsilon$	–	$10^{-9}$
Weight decay $\lambda_{\rm wd}$	$1.008\times 10^{-4}$	$0.0303$
Initialization scale $\sigma_{*}$	$0.1877$	$0.07539$

Appendix

Table 19: Hidden-matrix optimizer settings. The decay coefficient

\lambda_{\rm wd}

enters as

\theta\leftarrow(1-\lambda_{\rm wd}\eta_{t}^{2}/\eta_{\rm peak})\theta

	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	$\mathrm{PE}[10^{-5},1]$	cubic routine for eigenvalues above $10^{-7}\lambda_{\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$	$\boldsymbol{2.748}$
Muon	300M	$5$	$2.662$	$2.655$	$\boldsymbol{2.651}$
OKLS	160M	$5$	$2.687$	$2.673$	$\boldsymbol{2.671}$
OKLS	300M	$5$	$2.597$	$\boldsymbol{2.590}$	$2.591$

Appendix

Table 21: Measurement period. Final validation loss of our method with

T=200

and

T=500

Size	$\mathrm{PE}[10^{-5},1]$	$\mathrm{PE}[10^{-3},1]$	Ours
	(baseline)	(default)
160M	$2.761$	$2.765$	$\boldsymbol{2.751}$
300M	$2.662$	$2.669$	$\boldsymbol{2.655}$

Appendix

Table 22: Practitioner endpoint. Muon,

K=5

, final validation loss.

Optimizer	$K$	$T$	$N$	Agreement (%)	Distance	Fixed	RMS oracle	Ours	Regret
				Endpoint		RMS error
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

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	$H_{j}$ 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 Optimization

Matrix optimization develops methods for optimizing matrix-valued variables with structure such as low rank, sparsity, positivity, or orthogonality.

See all 58 papers

Orthogonal Witness Control for Muon Optimization via Sigmoid Spectral Reshaping
Sep 27, 2026Dat Phi Van, Ngo Vu Minh, Tuc Nguyen +3Neural Network Optimization Spectral Optimization
Center for AI Research, VinUniversity, Hanoi, Vietnam · Indiana University, Bloomington, IN, USA · Monash University, Clayton, VIC, Australia
Reassessing Muon for Matrix Factorization
Jul 14, 2026Ali Parviz, Gal Mishne, Alex CloningerDeep Learning Optimization Muon Optimizer
†. Halicioğlu Data Science Institute, UC San Diego. · ∗. Department of Mathematics, UC San Diego.
Aurora: A Leverage-Aware Spectral Optimizer
Jun 26, 2026Alec Dewulf, Dhruv Pai, Li Yang +2Neural Network Optimization Muon Optimizer
Tilde Research

Sep 27, 2026J/K move · Enter open · S save

Orthogonal Witness Control for Muon Optimization via Sigmoid Spectral Reshaping

Dat Phi Van, Ngo Vu Minh, Tuc Nguyen, Thin Nguyen+2

Center for AI Research, VinUniversity, Hanoi, Vietnam · Indiana University, Bloomington, IN, USA · Monash University, Clayton, VIC, Australia

Neural Network Optimization Spectral Optimization Muon Optimizer Nonconvex Optimization Newton-Schulz Iteration

Open paper AskPDF arXiv

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.

Abstract

Matrix Optimization

Explore similar work

Orthogonal Witness Control for Muon Optimization via Sigmoid Spectral Reshaping

Reassessing Muon for Matrix Factorization

Aurora: A Leverage-Aware Spectral Optimizer