Adaptive Newton–Schulz via Spectral Estimation

Newton–Schulz can measure the singular-value distribution of its own input

Kristi Topollai

September 2026 · Paper (arXiv)

This post is a companion to our paper, Cost-free Spectral Estimation for Adaptive Newton–Schulz in Matrix Optimizers^[1].

1Muon and Newton–Schulz

Muon^[2] updates a weight matrix $W$ with the polar factor of its momentum $M$. If $M=U\,\mathrm{diag}(\sigma_i)\,V^\top$ is the SVD of $M$, the polar factor is $\mathrm{polar}(M)=UV^\top$. It is the same matrix with every nonzero singular value set to $1$. The update is

Computing an SVD at every step is too slow. Muon instead runs $K$ Newton–Schulz iterations. It starts from $X_0$, a rescaled copy of $M$ whose singular values $\gamma_i$ lie in $[0,1]$, and repeats

Each iteration costs three matrix multiplications: $A_j$, $A_j^2$, and the product with $X_j$. The iterations keep the singular vectors $U$ and $V$ and change only the singular values. A singular value $x$ of $X_j$ becomes

in $X_{j+1}$. After $K$ iterations each $\gamma_i$ has become $P(\gamma_i)$, where $P=p_{K-1}\circ\dots\circ p_0$ is the response of the routine. So the whole matrix iteration is described by one scalar function. If $P(\gamma_i)\approx1$ for every $i$, then $X_K\approx UV^\top$. Figure 1 shows this on a real momentum matrix.

2The design interval

How Polar Express picks its coefficients. Polar Express^[3] picks the coefficients one iteration at a time. Suppose we know that the singular values of $X_j$ lie in an interval $[\ell_j,u_j]$, the design interval. At the start this is $[\rho,1]$. The coefficients are those of the odd quintic that stays closest to $1$ on the whole interval:

The best polynomial equioscillates. Its error $1-p_j(x)$ reaches $+e_j$ and $-e_j$ in turn at four points: the two ends of the interval and two points inside it. After the iteration all singular values lie in $[1-e_j,\,1+e_j]$, so this becomes the design interval of the next iteration. Its ratio $(1-e_j)/(1+e_j)$ is closer to $1$ than $\ell_j/u_j$ was. Repeating this $K$ times gives the response $P$. Figure 2 shows the construction.

A broader interval needs more iterations. Only the ratio $\rho=\ell_0/u_0$ matters, and a smaller ratio, meaning a broader interval, always gives a larger error for the same number of iterations. When the ratio is very small, one quintic cannot do much. It has to raise the small singular values but stay bounded at the top of the interval, so the error stays close to $1$ for the first iterations. Once the ratio reaches about $10^{-1}$ the error drops fast, and two or three more iterations are enough.

Below $\rho$. The routine gives no guarantee there. For $x$ much smaller than $\ell_j$ an iteration is close to linear, $p_j(x)\approx a_jx$, so singular values far below $\rho$ stay small. This does not matter if no singular values are that small. But if the spectrum ends two decades above $\rho$, the iterations spent growing the interval from $\rho$ up to the true $\sigma_{\min}/\sigma_{\max}$ did no useful work.

Choosing $\rho$. Muon does not know the spectrum of $M$, so it uses one interval for every matrix. Polar Express^[3] recommends $\rho=10^{-3}$. This choice assumes that singular values below $10^{-3}$ do not matter, since the routine leaves them small. That is not always the case: the attention matrices in Section 6 have a tail that reaches $10^{-5}$. In this post the fixed routine uses the more conservative interval $[10^{-5},1]$. It covers such tails, but it spends more iterations on matrices that do not have one. Figure 3 compares the response of a routine with real spectra.

In the matrices we study, MLP spectra end two decades above the conservative lower endpoint, while attention spectra have a tail that reaches it. For MLP matrices, many of the iterations are spent on a part of the interval that holds no singular values.

3Newton–Schulz measures its own input

Each iteration forms $A_j=X_jX_j^\top$ and $B_j=A_j^2$. These two matrices are already in memory, and four numbers follow from them by scalar reductions (sums of squares and inner products), without any extra matrix multiplication:

Let $z_j(x)$ be the value that a singular value $x$ of $X_0$ has after $j$ iterations, so $z_0(x)=x$ and $z_{j+1}=p_j(z_j)$. If $X_0=U\,\mathrm{diag}(\gamma_i)\,V^\top$, then $X_j=UZ_jV^\top$ with $Z_j=\mathrm{diag}(z_j(\gamma_i))$. The eigenvalues of $A_j$ are $z_j(\gamma_i)^2$, so

Here $r$ is the smaller dimension of $M$, and $\mu$ is the empirical spectrum of $X_0$, which puts mass $1/r$ on each $\gamma_i$. Each spectral moment $m_{2k,j}$ is the average of a known measurement function $\phi(x)=z_j(x)^{2k}$ over the singular values. The output gives one more moment, $\|X_K\|_F^2/r$. These identities are exact in exact arithmetic.

Each iteration applies a different polynomial, so the moments from different iterations look at the same spectrum through different functions, and each one is sensitive to a different part of $[0,1]$. We call a Newton–Schulz run that records these moments a measurement pass. Figure 4 shows one.

Before the first iteration, $M$ has to be scaled so that all its singular values are at most $1$. Dividing by the Frobenius norm works, but it overestimates the largest singular value and wastes part of the interval. We use the Gram-power bound $s$ from CANS^[4] instead and start from $X_0=M/(s\|M\|_F)$. The bound is computed from the products of the first iteration and is close to the largest singular value of $M/\|M\|_F$. The measurement pass, the fit and the scoring all use this scale.

4Fitting a spectrum to the moments

4.1 A maximum-entropy moment fit

We describe the spectrum by masses $w_b\ge0$ on fixed grid points $g_1,\dots,g_B\in[0,1]$, spaced logarithmically. The mass $w_b$ is the fraction of singular values at $g_b$. A moment is the average of a known function over the singular values. On the grid, this average becomes a sum over the grid points weighted by $w_b$. With $\Phi_{qb}=\phi_q(g_b)$, the moments are linear in $w$:

$\Phi$ depends only on the routine and the grid. We compute it once by running the scalar recursion $z_{j+1}=p_j(z_j)$ on the grid points. The masses must sum to $1$, and the second moment is known exactly from the scaling. Both enter as constraints:

Step 1: best fit. Find the smallest weighted residual between the predicted and the measured moments,

Here $D=\mathrm{diag}(1/d_q)$, and the row scales $d_q$ bring moments of very different sizes to a common scale.

Step 2: maximum entropy. Many grid measures reach this residual, because a finite set of moments cannot fix every grid weight. Among all measures whose residual is within $\kappa$ of the best, we take the one with maximum entropy:

On a logarithmic grid, maximum entropy gives the flattest spectrum in $\log\sigma$ that agrees with the moments. Both problems are convex, and the entropy makes the answer unique.

Figure 5 shows the effect of the entropy tolerance $\kappa$. Without it, the fit is one of many spiky measures that match the moments equally well. A small $\kappa$ gives a smooth fit with almost the same residual. A large $\kappa$ pushes the fit toward a flat spectrum, and the residual grows.

4.2 What the fitted spectrum can predict

We do not need the exact spectrum. We only need the quantities used to score routines. A finite set of moments may not determine the spectrum, but it can determine an average $\int f\,d\mu$ when $f$ is close to a combination of the measurement functions. The error of the prediction then has three parts: how well the measurement functions approximate $f$, the error in the measured moments, and the residual of the fit (Theorem 2 in the paper^[1]).

We can also see how much the moments leave open. For a linear quantity $c^\top w$, such as a band mass, the mass above the cutoff, or the CDF at a point, two linear programs give the smallest and the largest value that agrees with the moments:

We call this interval the guaranteed range. The maximum-entropy fit $\hat w$ is one point inside it.

5Selecting a Polar Express routine

5.1 The candidate set

A candidate $\pi=(\rho_\pi,K_\pi)$ in the candidate set $\mathcal C$ is a full Polar Express routine: the composition designed for $[\rho_\pi,1]$, with $K_\pi$ iterations and response $P_\pi$. The candidates cover design lower endpoints from $10^{-5}$ to $10^{-1}$ and several depths. We compute the response of every candidate at the grid points once, before training. Each setting of the two sliders in Figure 3 is one candidate. Fixed PE, $[10^{-5},1]$, is the most conservative one.

5.2 Scoring each candidate

We score all candidates on the same interval $R=[x^\star/s,\,1]$. The scoring cutoff $x^\star=10^{-5}$ is fixed relative to the Frobenius norm, and singular values below it are ignored. We measure a routine by its RMS error, the root-mean-square distance of the output singular values from $1$ over the scored singular values. The true error and the error predicted from the fit are

The predicted error is a spectral average of the kind described in Section 4.2. The responses are precomputed, so scoring takes only scalar operations on the fitted spectrum and no matrix products.

5.3 Fixed-depth and target rules

The same scores serve two goals:

The fixed-depth rule picks the routine with the lowest predicted error for a given number of iterations. The target rule picks the shortest routine whose predicted error is at most $\tau$. If no candidate reaches $\tau$, it falls back to fixed PE at the largest allowed depth. Figure 7 shows both choices for one matrix.

Several matrices can also share one routine. The moments are averages over the spectrum, so the average of the moments of a group of matrices, for example all matrices of one type, gives the moments of their combined spectrum. One fit then serves the whole group.

6Spectra across layers and training

The best routine depends on the spectrum, and momentum spectra change across matrix types, across layers, and during training. Figures 8 and 9 show this for the run we study. Attention spectra are broad: most of the mass is above $10^{-3}$, with a tail down to $10^{-5}$. MLP spectra are narrow and lie above $10^{-3}$. One conservative interval has to cover all of these cases, including matrices with almost no mass near its lower end.

After the first part of training the spectra change slowly, so a routine chosen from one measurement stays a good choice for many steps. Matrices of the same type differ more from layer to layer than a single matrix changes over time, which is why we choose a routine for each matrix.

7Periodic measurement

Measure every $T$ steps. Every $T$ training steps, each momentum matrix goes through a measurement pass. This is Newton–Schulz with a measurement routine $S_{\rm meas}$ of depth $K_{\rm meas}$. Its output is the update for that step, and its moments are recorded. If the steps in between use $\bar K$ iterations on average, the average number of iterations per step is

The extra cost of measuring shrinks as $T$ grows.

Fit and select in the background. The fit and the scoring work on a fixed number of moments and a fixed grid, so their cost does not grow with the matrix size. They run on the CPU while the GPU keeps training with the current routines. When the CPU job is done, the new routines are used from the next step on.

8Results

We test the method on Muon momentum matrices saved at several checkpoints of a GPT training run on FineWeb^[5]. They cover all seven matrix types: Q, K, V, O, MLP up, MLP gate and MLP down. We computed the SVD of every matrix offline and use it only to get the exact moments and as the ground truth. The fit only sees the moments, not the singular values. For each number of iterations $K$ we compare three methods:

• fixed PE on $[10^{-5},1]$, the conservative routine of Section 2, which starts from $M/\|M\|_F$ as in Muon;
• ours, the routine picked by spectral estimation;
• the RMS oracle, which picks from the same candidates using the true RMS error.

The difference between ours and the RMS oracle shows how much we lose by not knowing the spectrum. Ours and the RMS oracle use the same scaling and candidates. All three are scored on the same singular values, and we report the RMS error of Section 5.2.

8.1 The fit recovers the spectrum at decade scale

The fit is accurate: the decade band masses are usually within a few thousandths of the true values. The guaranteed ranges always contain the true value, but they are wide: the error of the fit is 30 to 140 times smaller than the width of the range.

8.2 Fixed depth

• MLP matrices. Their spectrum is narrow and lies above $10^{-3}$, so a routine designed for that interval brings every singular value close to $1$. With five iterations, our routine has a 90× lower mean RMS error than fixed PE.
• Attention matrices. Their spectrum has a tail down to $10^{-5}$, and a routine with few iterations cannot bring that tail to $1$. Choosing the routine still lowers the error at every depth, but the tail limits how low the error can go.
• Ours and the RMS oracle. On attention matrices the two have nearly the same median error at every depth, within $0.001$. On MLP matrices our mean error is within a few thousandths of the oracle's. The moments are enough to make almost the same choice as the true spectrum.

8.3 Target accuracy

Now the number of iterations is chosen as well. For a target $\tau$, the target rule picks for each matrix the shortest routine whose predicted RMS error is at most $\tau$. Fixed PE has to use one depth for all matrices, the smallest one that reaches $\tau$ on every spectrum.

At $\tau=10^{-1}$, ours and the RMS oracle pick almost the same depths. At the tighter targets, ours uses one or two more iterations than the oracle on MLP matrices. A predicted error of $10^{-2}$ or $10^{-3}$ is close to what a fit at decade resolution can tell apart, so the rule is more cautious there.

9Takeaways

1. Newton–Schulz already forms Gram matrices, and scalar reductions of them give spectral moments of its input. No extra matrix multiplications are needed.
2. A maximum-entropy moment fit turns the moments into a spectrum that is accurate at the scale of a decade. This is enough to predict the RMS error of each candidate Polar Express routine.
3. The fixed-depth rule picks the best routine for a given number of iterations. The target rule picks the shortest routine that is predicted to reach a given accuracy.
4. Spectra differ across matrix types and layers and change slowly during training, so the measurement can be repeated periodically while the fit runs in the background.
5. On saved momentum matrices, spectral estimation lowers the RMS error for a fixed number of iterations, most of all on MLP matrices, or reaches the same accuracy with fewer iterations. The paper^[1] also reports GPT pretraining runs, where adapting the routine to the spectrum lowers the validation loss of Muon and of Online KL-Shampoo.

References

1. K. Topollai, A. Choromanska. Cost-free spectral estimation for adaptive Newton–Schulz in matrix optimizers. arXiv:2609.33047, 2026.
2. K. Jordan, Y. Jin, V. Boza, J. You, F. Cesista, L. Newhouse, J. Bernstein. Muon: An optimizer for hidden layers in neural networks. 2024. kellerjordan.github.io.
3. N. Amsel, D. Persson, C. Musco, R. M. Gower. The Polar Express: Optimal matrix sign methods and their application to the Muon algorithm. ICLR 2026. arXiv:2505.16932.
4. E. Grishina, M. Smirnov, M. Rakhuba. Accelerating Newton–Schulz iteration for orthogonalization via Chebyshev-type polynomials. 2025. arXiv:2506.10935.
5. G. Penedo, H. Kydlíček, L. Ben allal, A. Lozhkov, M. Mitchell, C. Raffel, L. Von Werra, T. Wolf. The FineWeb datasets: Decanting the web for the finest text data at scale. NeurIPS Datasets and Benchmarks, 2024. arXiv:2406.17557.

Citation

To cite the results, please cite the paper^[1]:

To cite this post:

@misc{topollai2026adaptivens,
  title        = {Adaptive {Newton--Schulz} via Spectral Estimation},
  author       = {Topollai, Kristi},
  year         = {2026},
  month        = {September},
  howpublished = {Blog post},
  url          = {https://kristitopollai.github.io/blog/adaptive-muon/}
}