PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kristi Topollai, Anna Choromanska

arXiv:2609.33047v1cs.LG

TL;DR

Matrix optimizers often use the same Newton–Schulz routine despite spectra varying across inputs, even though those spectra affect approximation difficulty. This paper estimates spectra from measurements already produced by Newton–Schulz and selects a routine for each matrix. The approach improves fixed-budget approximation and improves two matrix optimizers in GPT pretraining.

  • Problem

    Matrix-function approximation difficulty varies with the input spectrum, but existing routines typically use the same computation across matrices.

    Source: §1 · p. 1

  • Method

    The method extracts spectral measurements from Newton–Schulz intermediate matrices at negligible cost and uses them to select a polynomial routine for each matrix.

    Source: §1 · p. 2

  • Results

    Spectral estimation improves polar approximation at a fixed iteration budget or reaches the same accuracy with fewer iterations, and improves Muon and OKLS in GPT pretraining.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Newton–Schulz can reuse measurements from its existing computation to specialize routines without additional matrix multiplications.

    Source: §1 · p. 2

    Finitely many moments may not identify a distribution, so the fitted spectrum is useful when it predicts the spectral averages used to score candidate routines.

    Source: §4.3 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

Matrix-optimizer routines face different spectra across matrix types, layers, and training time, so a single conservative routine can waste iterations. The paper uses measurements already present in Newton–Schulz iterations to select spectrum-matched routines, improving approximation and pretraining outcomes.

1 INTRODUCTION

  • Matrix optimizers use matrix structure in gradients, momentum, or preconditioning statistics; Muon approximately orthogonalizes momentum matrices before using them as updates.Examples include Shampoo, SOAP, Dion, and Muon.Source: §1 · p. 1
  • Approximation difficulty depends on the input spectrum, which varies across matrix types, layers, and training time; one routine can waste iterations on easier cases.Attention spectra are often substantially broader than MLP spectra.Source: §1 · p. 1
  • At K = 5, spectral estimation cuts mean polar RMS error 90× versus fixed Polar Express on 540 saved MLP matrices, and improves Muon and OKLS pretraining.The error is scored above 10^-5 in Frobenius-normalized coordinates.Source: Figure 1 · p. 2, §1 · p. 2
  • Newton–Schulz intermediate matrices provide spectral measurements through scalar reductions, without additional matrix multiplications.The measurements can be used to estimate the spectrum and select a better-matched routine.Source: §1 · p. 2
  • The method selects a polynomial routine for each matrix instead of applying one conservative choice to all matrices.It is evaluated offline on saved momentum matrices and end-to-end in GPT pretraining.Source: §1 · p. 2

2 MATRIX FUNCTIONS IN OPTIMIZATION

Matrix optimizers apply matrix-valued transformations to gradients, momentum, or preconditioning statistics. Newton–Schulz approximates these transformations through polynomial iterations, whose response and quality depend on the singular-value spectrum and chosen design interval.

2.1 FROM MATRIX UPDATES TO SPECTRAL TRANSFORMATIONS

  • Muon’s ideal update is the polar factor UV^⊤, which preserves singular-vector directions while mapping every nonzero singular value to one.Shampoo and SOAP typically use eigendecomposition or related dense linear algebra for preconditioners, while OKLS uses coupled Newton–Schulz iterations for inverse square roots (Gupta et al., 2018; Anil et al., 2020).Source: §2.1 · p. 2
  • Direct SVDs or eigendecompositions are expensive to evaluate repeatedly during training, whereas Newton–Schulz uses short matrix-multiplication sequences that map efficiently to GPUs.The paper focuses on polar approximation; coupled variants also provide routes to inverse roots and fractional powers.Source: §2.1 · p. 3

2.2 NEWTON–SCHULZ AND OPTIMAL POLYNOMIAL ITERATIONS

  • Newton–Schulz applies low-degree polynomials to singular values, with the classical update driving normalized positive values toward one.Modern variants use improved scaling and higher-degree iterations (Amsel et al., 2026a).Source: §2.2 · p. 3
  • Polar Express designs polynomial coefficients to minimize worst-case error over a prescribed spectral interval, with its optimality relative to that interval (Amsel et al., 2026a).After rescaling the upper endpoint to one, the interval ratio ρ determines the design problem; broader intervals are harder.Source: §2.2 · p. 3

3 SPECTRUM-DEPENDENT COST OF NEWTON–SCHULZ

The Newton–Schulz iteration budget needed for polar approximation depends on the design interval and on how the spectrum varies across matrices. Narrower, matrix-specific spectral information can identify when a conservative interval spends unnecessary work.

3.1 BROADER DESIGN INTERVALS REQUIRE MORE COMPUTATION

  • As ρ approaches zero, minimax error approaches one, so achieving a fixed target accuracy eventually requires more iterations.At every fixed depth, a broader interval has strictly larger minimax error.Source: §3.1 · p. 3

3.2 TRAINING PRODUCES DIFFERENT SPECTRA ACROSS LAYERS AND TIME

  • Attention spectra are broader than MLP spectra, and both vary across layers and training time in the 160M Muon baseline.Figure 3 compares Q and MLP-up spectra across layers at step 2k and Layer 11 across checkpoints.Source: §3.2 · p. 4, Figure 3 · p. 4

3.3 SPECTRAL INFORMATION IDENTIFIES UNNECESSARY WORK

  • Using each matrix’s exact nonzero spectral range, an oracle finds that narrower MLP spectra reach the same target in far fewer iterations than a common conservative interval.Attention still requires a relatively deep iteration; average error also depends on spectral mass, not only endpoints.Source: §3.3 · pp. 4–5, Figure 4 · p. 4

4 SPECTRAL ESTIMATION FROM NEWTON–SCHULZ

Newton–Schulz iterations expose spectral information that can be measured and fitted cheaply, then used to select a polynomial routine for each matrix. This selection predicts candidate quality without requiring exact spectrum recovery.

4 SPECTRAL ESTIMATION FROM NEWTON–SCHULZ

  • Spectral estimation improves polar approximation at fixed depth or reaches a target accuracy with fewer Newton–Schulz iterations.The procedure measures, fits, and selects routines; between measurement passes it reuses the selected routine as spectra change gradually.Source: §4 · p. 5, §4.4 · p. 7

4.1 MOMENTS FROM THE NEWTON–SCHULZ COMPUTATION

  • Each Newton–Schulz iteration yields four even spectral moments through scalar reductions of matrices it already forms, requiring no additional matrix multiplication.Together with the final output norm, Kmeas measurement iterations provide 4Kmeas + 1 scalar measurements.Source: §4.1 · p. 5

4.2 FITTING A SPECTRUM TO THE MEASURED MOMENTS

  • A maximum-entropy moment fit recovers broad attention and concentrated MLP spectra, with close agreement in cumulative mass.The fitted measures are compared with SVD spectra in Figure 5.Source: §4.2 · p. 6, Figure 5 · p. 6

4.3 WHAT THE FITTED SPECTRUM CAN PREDICT

  • The fitted spectrum need not be unique if it predicts the spectral averages used to score candidate routines.Prediction is exact when the target function lies in the affine span of the measurement functions and the fitted moments are exact.Source: §4.3 · pp. 6–7

4.4 SELECTING A POLAR EXPRESS ROUTINE

  • Scalar scores from precomputed Polar Express responses support both best-routine-at-fixed-depth and shortest-routine-meeting-target selections.If no candidate meets the target, selection falls back to fixed PE at the largest allowed depth; Figure 6 illustrates both objectives.Source: §4.4 · p. 7, Figure 6 · p. 7

5 EXPERIMENTS

Across GPT model sizes, spectral estimation improves the fixed-baseline validation loss for Muon and OKLS. Offline, it reduces RMS error or routine depth, with trade-offs against worst-case error.

5.1 ORTHOGONALIZATION OF REAL SPECTRA

  • For a target accuracy, spectral estimation selects shorter routines than fixed PE, especially for MLP matrices.Figure 7 reports mean iterations needed to meet τ; spectral estimation uses predicted RMS error, while the oracle uses worst-case error.Source: §5.1 · p. 8, Figure 7 · p. 8, Table 2 · p. 8
  • At fixed depth, spectral estimation has lower RMS error than the minimax oracle on attention through K = 9 and MLP through K = 6, while the oracle has lower worst-case error at every depth.Figure 7 summarizes mean errors over attention and MLP matrices; Table 1 reports mean per-matrix RMS and worst-case error.Source: Figure 7 · p. 8, §5.1 · p. 8, Table 1 · p. 8

5.2 LLM PRETRAINING

  • Spectral estimation improves on the fixed baseline for both Muon and OKLS at every model size, matching or improving periodic exact-SVD or eigendecomposition selection.Table 3 reports final validation loss at K = 5 after 20B tokens, with T = 200 for spectral estimation.Source: Table 3 · p. 9, §5.2 · p. 9, Figure 8 · p. 9
  • At 160M parameters, five adaptive Muon iterations reach lower validation loss than the fixed baseline at any depth, while fixed-baseline loss is unchanged from six to eight iterations.Table 4 reports final validation loss across ordinary-step depths.Source: §5.2 · p. 9, Table 4 · p. 9

6 CONCLUSION

Newton–Schulz iterations already provide measurements that can specialize their polynomial routines. The method improves approximation at fixed depth or reduces iterations to a target, and useful adaptation requires prediction rather than exact spectrum recovery.

6 CONCLUSION

  • Measurements from Newton–Schulz iterations support spectrum-adapted polynomial routines, improving fixed-depth approximation or reducing iterations to target accuracy.The approach is evaluated on saved momentum spectra and end-to-end GPT pretraining; its fit uses a fixed number of scalars.Source: §6 · p. 9

A.1 PROOF OF PROPOSITION 1

The proof sections connect Newton–Schulz polynomial responses to polar-approximation error, establish how interval width affects minimax error, and justify moment-based spectrum fitting and candidate selection.

A.1 PROOF OF PROPOSITION 1

  • The singular-value iteration identity yields Equation (2), and the polar-factor error decomposes into scored singular values and values below the scoring cutoff.The score covers only the first term; in offline data, an average 3.8% of attention singular values lie below the cutoff, versus none for MLP matrices.Source: §A.1 · p. 13

A.2 PROOF OF THEOREM 1

  • The proof derives strict monotonicity of minimax error with interval width and a lower bound for compositions of K quintics.The lower bound is E ≥ ((1−ρ)/(1+ρ))^(5K).Source: §A.2 · p. 13

A.3 PROOF OF THEOREM 2 AND ACCURACY OF SELECTION

  • The fitted spectral measure exists and is unique, while the score compares estimated error integrals relative to scored mass to select a candidate.The entropy objective selects a unique representative when multiple grid distributions match the measurements.Source: §A.3 · p. 14

B METHOD DETAILS

The method normalizes each input, estimates its spectrum from Newton–Schulz measurements, and selects a candidate routine periodically. Offline comparisons show close spectrum fits and favorable RMS-error results, including on a held-out run.

B.1 INPUT NORMALIZATION

  • The Gram-power scale is nearly tight across matrix types, and better normalization alone accounts for only part of spectral selection’s gain.The first-iteration products provide both the scale and moments; rescaling leaves the polar factor unchanged.Source: §B.1 · p. 14, Figure 9 · p. 15

B.2 CANDIDATES AND FITTING

  • Candidates are complete polynomial routines specified by a design endpoint and depth, with routines failing positivity checks removed.Maximum-entropy fitting chooses one representative among feasible spectra even when distinct measures match the finite measurements.Source: §B.2 · pp. 14–15, Figure 10 · p. 15

B.3 PERIODIC MEASUREMENT

  • Periodic measurement produces the update and moments, while fitting and selection run in the background and replace the active routines when ready.The previous routine remains in use during fitting.Source: §B.3 · pp. 15–16

C.2 FITTED SPECTRA ACROSS LAYERS AND TRAINING

C.3 FIXED-DEPTH ERRORS AND THE CHOICE OF SCORE

  • Spectral estimation nearly matches the RMS oracle on attention and is close at shallow depths on MLP, while beating the minimax oracle in RMS error through K = 9 and K = 6, respectively.It has higher worst-case error than the minimax oracle.Source: §C.3 · p. 17

C.4 HELD-OUT RUN

  • On the held-out run, spectral estimation stays close to the RMS oracle on attention and below fixed PE on both attention and MLP matrices.On attention it also has lower worst-case error than fixed PE; the minimax oracle has lower errors on MLP matrices.Source: §C.4 · p. 17, Figure 12 · p. 18

C.5 ABLATIONS

Ablations examine measurement quantity, scoring cutoff, target selection, and reuse or grouping of routine choices. More measurements help mostly on MLP matrices, while isolated small singular values can still cause target misses.

C.5 ABLATIONS

  • More measurement iterations and moment orders both help, mostly on MLP matrices; lowering the scoring cutoff makes the problem harder for every method.At lower cutoffs, almost no attention matrix meets τ ≤ 10^-2 within the allowed depth cap for either method.Source: §C.5 · p. 18, Figure 13 · p. 20, Figure 14 · p. 20

C.6 TARGET RULE

  • A small safety margin removes all attention target misses at little extra depth, but some MLP-down matrices still miss because isolated small singular values are not detected by moments.Each such singular value adds a large error.Source: §C.6 · p. 21, Table 11 · p. 21

C.7 REUSE AND PER-TYPE SELECTION

  • Reusing earlier design endpoints changes error very little, and all tested reuse variants remain far below fixed PE.The comparison evaluates saved matrices while recomputing their normalization.Source: §C.7 · p. 21, Table 12 · p. 21
  • Per-matrix selection has lower error than per-type selection at both tested depths, especially at larger depth, without requiring exact SVDs from the same run.Grouping can serve a pooled spectrum because averaging group measurements gives the moments of that spectrum.Source: §C.7 · p. 21, Table 13 · p. 22

D COST

The measurement reductions add little computation relative to Newton–Schulz iterations, and periodic amortization keeps per-step overhead small. The same spectral-estimation approach also improves cubic inverse-root selection on regularized momentum Gram matrices.

D COST

  • The moment reductions add about 3r^2 operations per iteration, while fitting and scoring use a fixed-size grid and measurement set independent of matrix size.On one H200, reductions take 1–3 ms at tested sizes and fitting takes 4 ms of CPU time per matrix.Source: D COST · p. 22, Table 14 · p. 22, Figure 15 · p. 22
  • For the 1B Muon model, periodic measurement adds under 1 ms of reductions, 0.025 iterations, and 4 ms of CPU time per step.The accounting uses N = 168 hidden matrices, K = 5, T = 200, and roughly 40 seconds per step.Source: D COST · p. 22

E INVERSE SQUARE ROOTS

  • Selecting a cubic routine from the fitted spectrum improves on the fixed cubic routine at every tested depth for regularized momentum Gram matrices, with largest gains on concentrated spectra.The coupled iteration computes inverse square roots and its measurements are obtained from reductions of matrices the iteration already forms.Source: E INVERSE SQUARE ROOTS · p. 23, Figure 16 · p. 24

F.1 SETUP

  • The GPT setup applies Muon or OKLS to hidden matrices and AdamW to embeddings, output heads, normalization gains, and attention-gate projections, with the data order held fixed across runs.The Muon comparison uses Polar Express rather than CANS in this setup.Source: §F.1 · pp. 24–25, Table 18 · p. 25, Table 19 · p. 25

F.2 MATRIX-FUNCTION ROUTINES

  • Muon and OKLS use separate Newton–Schulz routines; OKLS runs its coupled inverse-root iteration in fp32 because it drifts in bf16.Spectral measurements are collected every T steps, and exact spectra recorded during Muon runs are used only to evaluate selection quality, not to select routines.Source: §F.2 · p. 25, Table 20 · p. 26

F.3 ADDITIONAL PRETRAINING RESULTS

Changing the measurement period from T = 200 to T = 500 improves on the fixed baseline in every tested configuration, though neither period wins across all configurations. At K = 5, per-matrix selection outperforms both fixed endpoint choices in the reported Muon runs.

Measurement period

  • Both T = 200 and T = 500 improve on the fixed baseline in every tested configuration, but neither period is better in all configurations.Table 21 reports final validation loss for the two measurement periods.Source: §F.3 · p. 26, Table 21 · p. 26

Practitioner endpoint

  • Per-matrix selection improves on both fixed routines at K = 5, while the narrower fixed endpoint only slightly improves on the baseline at 160M and matches it at 300M.The fixed baseline uses [10−5, 1], compared with [10−3, 1] for the practitioner endpoint.Source: §F.3 · p. 26, Table 22 · p. 26

F.4 SELECTION QUALITY DURING TRAINING

During training, spectral estimation selects routines with error close to the RMS oracle and far below the fixed baseline, while measured moments preserve nearly the same selection quality as exact moments. The related work positions this approach as using already-formed Newton–Schulz products to choose routines for current matrices.

F.4 SELECTION QUALITY DURING TRAINING

  • Spectral estimation’s error is close to the RMS oracle’s and far below the fixed baseline’s; its endpoint choices vary by matrix type.On attention Q and O, the minimax rule chooses smaller endpoints than the RMS rule to protect the worst direction; for other types, spectral estimation chooses more conservative endpoints than the RMS oracle.Source: §F.4 · p. 26, Table 23 · p. 27, Figure 17 · p. 27
  • Measured moments give nearly the same selection quality as exact moments, while omitting the OKLS Yj and Zj reductions produces clearly worse choices.Table 24 compares measured moments (bf16 for Muon and fp32 for OKLS) with exact moments and compares the reductions used.Source: §F.4 · p. 26, Table 24 · p. 27

G RELATED WORK

  • Unlike interval-based polynomial routines that typically use one interval for every matrix, this method chooses the interval and depth from each matrix’s spectrum.The polynomial families are retained, but their choices adapt to the matrices being processed.Source: G RELATED WORK · p. 27
  • Unlike methods that reduce iteration or step cost, spectral estimation chooses which iterations to run and can be combined with those methods.The cited approaches include Gram Newton–Schulz, Dion and Dion3, MuonBP, and CacheMuon.Source: G RELATED WORK · pp. 27–28
  • Unlike AMO’s early per-type routine choice and PRISM’s randomized sketches, this method uses products already formed by Newton–Schulz to select a complete routine for later steps.It uses no random probes.Source: G RELATED WORK · p. 28
  • Spectral estimation reduces orthogonalization error at a fixed iteration count, complementing analyses that relate inexact orthogonalization to convergence, learning rate, and momentum.Those analyses include Shulgin et al. (2026), Kim & Oh (2026), Choudhury et al. (2026), Qian & Richtárik (2026), and Do et al. (2026).Source: G RELATED WORK · p. 28
  • Moment-based density estimates need not identify the full distribution to predict averages of functions approximated well by the measured polynomials.Here, Newton–Schulz already forms the matrices whose traces provide the moments, avoiding probe vectors and extra matrix products.Source: G RELATED WORK · p. 28

Related papers

Loading 2609.33047v1…