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

Summary (Overview)

Cost-free spectral estimation from Newton–Schulz: The paper shows that the intermediate Gram matrices formed during Newton–Schulz (NS) iterations already contain sufficient information to estimate the input's singular-value spectrum via inexpensive scalar reductions (traces and Frobenius norms), requiring no additional matrix multiplications.
Spectrum-adaptive orthogonalization: Using these moments, the authors fit a maximum-entropy spectral measure and use it to select a specialized Polar Express (PE) polynomial routine per matrix, replacing the standard fixed conservative routine used for all layers and throughout training.
Theoretical guarantees: They prove (Theorem 1) that broader spectral intervals strictly increase minimax error at fixed depth, and (Theorem 2) that the fitted spectrum predicts spectral averages (like RMS error) with bounded error, even when the spectrum is not uniquely identifiable.
Empirical validation: Offline, spectral estimation reduces polar RMS error by up to 90× on saved MLP momentum matrices at K=5 iterations, and reaches target accuracy with fewer iterations (e.g., 5 vs. 12 for MLP). In GPT pretraining (160M–1B params), it consistently improves validation loss for both Muon and Online KL Shampoo (OKLS) optimizers.
Practical efficiency: The added cost is negligible—scalar reductions add ~3r² work per iteration (a fraction 3/(2n+r) of the iteration), and the fit runs on a fixed grid on CPU, independent of matrix size. For a 1B Muon model, this adds under 1ms of GPU time and 4ms CPU per step.

Introduction and Theoretical Foundation

Background and Motivation

Matrix optimizers (Muon, Shampoo, SOAP, Dion, OKLS) are increasingly popular for neural network training, with Muon adopted in systems like DeepSeek-V4, GLM-5, and Kimi K3. Muon approximately orthogonalizes momentum matrices via a small number of Newton–Schulz iterations, replacing expensive SVDs with hardware-efficient matrix multiplications.

The key observation is that the difficulty of orthogonalization depends strongly on the input's singular-value spectrum, and spectra vary significantly across:

Matrix types (attention vs. MLP)
Layers (deeper layers often have broader spectra)
Training time (spectra evolve as training proceeds)

A conservative routine designed for the hardest case (broadest spectrum) wastes iterations on easier cases (narrow spectra).

Theoretical Foundation

Newton–Schulz iteration. For $M \in \mathbb{R}^{r \times n}$ with $r \le n$ , a quintic iteration takes the form:

A_j = X_j X_j^\top, \quad X_{j+1} = a_j X_j + (b_j A_j + c_j A_j^2) X_j. \tag{1}

If $p_j(x) = a_j x + b_j x^3 + c_j x^5$ , then after $K$ iterations:

X_K = U \operatorname{diag}(P(\gamma_i)) V^\top, \quad P = p_{K-1} \circ \cdots \circ p_0. \tag{2}

Polar Express (PE) chooses coefficients by minimizing worst-case error over a prescribed spectral interval $[\ell_j, u_j]$ :

(a_j^\star, b_j^\star, c_j^\star) \in \arg\min_{a,b,c \in \mathbb{R}} \max_{x \in [\ell_j, u_j]} |1 - (ax + bx^3 + cx^5)|. \tag{3}

The only relevant parameter is the ratio $\rho = \ell_0 / u_0$ (after rescaling the upper endpoint to 1). A broader interval (smaller $\rho$ ) defines a harder problem.

Theorem 1 (Interval width and minimax error). Let $\varepsilon_K^\star(\rho)$ be the minimum uniform error on $[\rho, 1]$ over all compositions of $K$ odd polynomials of degree ≤ 5. Then:

\varepsilon_{K+1}^\star(\rho) = e\!\left(\frac{1 - \varepsilon_K^\star(\rho)}{1 + \varepsilon_K^\star(\rho)}\right), \tag{4}

and for $0 < \rho_1 < \rho_2 < 1$ :

\varepsilon_K^\star(\rho_1) > \varepsilon_K^\star(\rho_2), \quad \varepsilon_K^\star(\rho) \ge \left(\frac{1 - \rho}{1 + \rho}\right)^{5K}. \tag{5}

This establishes that broader intervals require strictly more iterations to reach a given accuracy.

RMS error metric. For normalized nonzero singular values $\gamma_i$ :

E_{\text{RMS}}(P)^2 = \frac{1}{r_+} \sum_{i=1}^{r_+} |1 - P(\gamma_i)|^2, \quad r_+ = \operatorname{rank}(M). \tag{6}

This is the squared relative Frobenius error to the polar factor. Minimizing worst-case error need not minimize this average—allowing larger error in a sparse tail can improve accuracy over the bulk.

Methodology

4.1 Spectral Moments from the Newton–Schulz Computation

The core insight is that the intermediate matrices $A_j = X_j X_j^\top$ and $B_j = A_j^2$ expose even spectral moments of orders 2–8 via scalar reductions, with no additional matrix multiplications.

Measurement pass (Algorithm):

Require: X_0, coefficients (a_j, b_j, c_j)
for j = 0, ..., K_meas - 1 do
    A_j ← X_j X_j^⊤                      (1)
    B_j ← A_j A_j                        (2)
    (m_{2,j}, m_{4,j}) ← (tr A_j, ‖A_j‖_F²)/r
    (m_{6,j}, m_{8,j}) ← (⟨A_j, B_j⟩_F, ‖B_j‖_F²)/r
    C_j ← b_j A_j + c_j B_j
    X_{j+1} ← a_j X_j + C_j X_j          (3)
end for
Read ‖X_{K_meas}‖_F²/r

(1)–(3) are existing matrix products; purple marks added scalar reductions.

Proposition 1 (Spectral moments from the iteration). In exact arithmetic:

m_{2k,j} = \frac{\operatorname{tr} A_j^k}{r} = \frac{1}{r} \sum_{i=1}^r z_j(\gamma_i)^{2k} = \int z_j(x)^{2k} d\mu(x), \quad k = 1, \ldots, 4. \tag{9}

where $z_0(x) = x$ , $z_{j+1}(x) = p_j(z_j(x))$ , and $\mu = \frac{1}{r}\sum_{i=1}^r \delta_{\gamma_i}$ is the empirical scaled spectrum. This gives $4K_{\text{meas}} + 1$ scalar measurements total.

4.2 Fitting a Spectrum via Maximum Entropy

The method recovers a discrete spectral measure from the measurements using a maximum-entropy moment fit. With masses $w_b$ on grid points $g_b \in [0,1]$ and $\Phi_{qb} = \phi_q(g_b)$ :

\Delta = \left\{ w \ge 0 : \mathbf{1}^\top w = 1, \sum_b w_b g_b^2 = \frac{1}{rs^2} \right\}. \tag{10}

The two-stage optimization is:

\xi^\star = \min_{w \in \Delta} \|D(\Phi w - \hat{y})\|_2, \quad \hat{w} = \arg\max_{w \in \Delta} \left\{ -\sum_b w_b \log w_b : \|D(\Phi w - \hat{y})\|_2 \le \xi^\star + \kappa \right\}. \tag{11}

Both problems are convex. The normalization uses the CANS Gram-power upper bound $s_8 = \|\bar{A}^2\|_F^{1/4}$ , which safely bounds the spectrum from above.

4.3 Prediction Guarantee

Theorem 2 (Prediction from the fitted spectrum). Let $\mu$ be the true spectral measure, $\phi_1, \ldots, \phi_L$ the measurement functions with observations $\hat{y}_q$ satisfying $|\hat{y}_q - \int \phi_q d\mu| \le \omega_q$ . Let $\hat{w}$ solve (11) and define $\hat{\mu} = \sum_b \hat{w}_b \delta_{g_b}$ . For any bounded measurable $f$ and affine combination $\psi_\alpha = \alpha_0 + \sum_{q=1}^L \alpha_q \phi_q$ with $\beta_\alpha(f) = \sup_{x \in [0,1]} |f(x) - \psi_\alpha(x)|$ :

\left|\int f d\hat{\mu} - \int f d\mu\right| \le 2\beta_\alpha(f) + \sum_{q=1}^L |\alpha_q| \omega_q + (\xi^\star + \kappa) \|D^{-1}\alpha\|_2. \tag{12}

The three terms capture: how well $f$ is represented by the measured functions, error in observations, and the fit residual.

4.4 Routine Selection

The fitted spectrum scores a precomputed set of PE routines. Each candidate $\pi = (\rho_\pi, K_\pi)$ has true and predicted squared RMS errors:

E_R(P_\pi)^2 = \frac{\int_R (1 - P_\pi(x))^2 d\mu(x)}{\mu(R)}, \quad \hat{E}_\pi^2 = \frac{\sum_{g_b \in R} \hat{w}_b (1 - P_\pi(g_b))^2}{\sum_{g_b \in R} \hat{w}_b}. \tag{13}

Two selection objectives:

Fixed-depth rule: picks the best routine at a given depth $K$ : $\hat{\pi}_K \in \arg\min_{\pi \in \mathcal{C}, K_\pi = K} \hat{E}_\pi$
Target rule: picks the shortest routine predicted to meet accuracy $\tau$ : $\hat{\pi}_\tau \in \arg\min_{\pi \in \mathcal{C}, \hat{E}_\pi \le \tau} K_\pi$

Measurements are taken every $T$ steps (spectra change gradually), with fitting/selection running in the background.

Empirical Validation / Results

5.1 Offline Orthogonalization of Real Spectra

Evaluated on saved Muon momentum spectra from a FineWeb run. Baselines: fixed conservative PE on $[10^{-5}, 1]$ and a minimax oracle (exact spectrum).

Table 1: Fixed-depth polar error (mean per-matrix RMS and worst-case error):

| | 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⁻⁶ | 0.222 | 0.00153 | 1.78×10⁻⁶ |

Key findings:

90× reduction in mean polar RMS error for MLP matrices at K=5 relative to fixed PE
Spectral estimation has lower RMS error than the minimax oracle on attention through K=9 and on MLP through K=6 (because it optimizes RMS, not worst-case)
With target accuracy, spectral estimation picks shorter routines (e.g., MLP needs 5 vs. 12 iterations for τ=10⁻²)

5.2 LLM Pretraining

GPT-style decoder-only models (160M, 300M, 1B non-embedding params) trained on 20B tokens of Nemotron-CC v2 High-Quality. Both Muon (polar factor) and OKLS (inverse square roots via coupled NS) were tested.

Table 3: Pretraining across optimizers and sizes (final validation loss, K=5, T=200):

| | Muon | | | OKLS | | | |---|---|---|---|---|---| | Size | Fixed | Oracle | Ours | Fixed | Oracle | Ours | | 160M | 2.761 | 2.752 | 2.751 | 2.687 | 2.675 | 2.673 | | 300M | 2.662 | 2.654 | 2.655 | 2.597 | 2.591 | 2.590 | | 1B | 2.455 | 2.450 | 2.449 | 2.449 | 2.439 | 2.436 |

Table 4: Muon 160M across depths (final validation loss):

K	Fixed	Ours
4	2.808	2.791
5	2.761	2.751
6	2.754	2.752
8	2.752	2.752

Key findings:

Spectral estimation improves on the fixed baseline for both optimizers at every model size
It matches or improves on periodic selection from exact SVD/eigendecomposition (oracle)
With only 5 iterations, our method reaches lower loss than the fixed baseline at any depth—adaptation can replace extra iterations

Cost Analysis

Table 14: Added cost of each spectrum estimator for $M \in \mathbb{R}^{r \times n}$ , $r \le n$ (one NS iteration costs $2r^2n + r^3$ ):

Method	Extra products	Work that grows with M	Fixed work
Ours	none	≈ 3r² per iteration	fit, scoring
SVD (singular values)	–	O(r²n)	–
Gram eigenvalues $MM^\top$	–	O(r²n + r³)	–

For the 1B Muon model (N=168, K=5, T=200, ~40s/step): under 1ms of reductions, 0.025 iterations, and 4ms CPU per step.

Theoretical and Practical Implications

Theoretical Implications

Newton–Schulz is self-measuring: The intermediate matrices of the iteration are not just computational byproducts—they encode spectral information that can be extracted at negligible cost. This reframes the problem: matrix-function operations inside optimizers need not be designed for a conservative worst-case spectrum.
Error bounds for spectral prediction: Theorem 2 provides a rigorous bound on the error of predicting spectral averages from finitely many moments. The bound decomposes into three interpretable terms: representation error, observation error, and fit residual. This connects to the broader theory of moment problems (Krein–Nudel'man) and spectral density estimation.
Optimality of spectrum-adaptive design: Theorem 1 formalizes that wider design intervals strictly degrade minimax error at fixed depth, providing theoretical justification for adapting the routine to the actual spectrum rather than covering all possible cases.

Practical Implications

Immediate cost savings: The method requires no additional matrix multiplications, and the fit runs on a fixed grid on CPU, independent of matrix size. This makes it applicable at frontier model scales where exact spectra are prohibitively expensive.
Applicable to multiple optimizers: The same measure-fit-select procedure works for both Muon (polar factor) and OKLS (inverse square roots), suggesting broad applicability to any Newton–Schulz-based matrix optimizer.
Adaptation replaces computation: The results show that per-matrix spectral adaptation can achieve the same accuracy with fewer iterations, or better accuracy at the same cost—effectively substituting cheap measurement for expensive computation.
Future directions: The method can be combined with other optimizations (Gram NS, communication-efficient variants like Dion), and the grouped selection idea could benefit mixture-of-experts models.

Conclusion

The paper demonstrates that Newton–Schulz orthogonalization does not need to treat the spectrum it processes as unknown—the iteration itself produces the measurements needed to specialize its own computation. Key contributions:

Spectral moments at negligible cost: The Gram matrices formed inside NS iterations yield spectral moments through scalar reductions, with no additional matrix multiplications.
Spectrum-adaptive selection: A maximum-entropy fit recovers enough spectral structure to select a specialized polynomial routine for each matrix, responding to differences across layers and training time.
Comprehensive validation: Offline, spectral estimation reduces polar RMS error by up to 90× at fixed iteration budget or reaches the same accuracy with fewer iterations. In GPT pretraining up to 1B parameters, it improves validation loss for both Muon and OKLS.

The results suggest a broader principle: matrix-function operations inside optimizers can cheaply measure the spectrum they are already processing and specialize computations accordingly, rather than being designed for a conservative worst-case spectrum. Useful adaptation does not require recovering the spectrum exactly—only predicting which candidate routine will work well.