Muon, a 2024-born optimizer for language models, differs from the 2015-born spectral gradient descent mainly in its fast, GPU-friendly Newton--Schulz polynomial for orthogonalizing the gradient matrix. Accordingly, a line of research proposes faster polynomial iterations for gradient orthogonalization, such as Polar Express and Gram Newton--Schulz, both built on quintic polynomials. In this paper, we introduce Polar Switch, a new approach to polynomial gradient orthogonalization. Using simple polynomial circuits, Polar Switch reaches a given high polynomial degree at a lower FLOP cost than quintic iterations. We define a circuit as a short sequence of symmetric matrix products on a few square buffers, and circuits with k products form a restricted family of odd polynomials of degree up to 2k+1+1. We choose the sequence of products in each circuit, their coefficients, and how the circuits are composed by numerically minimizing a worst-case bound on the polar approximation error. This minimization is subject to a FLOP budget that accounts for symmetric matrix multiplication kernels and a rounding constraint for numerical stability in half precision. Our two recommended Polar Switch methods have 15.5--18.7% lower runtime than Polar Express on square matrices and 8.6--16.2% lower runtime than Gram Newton--Schulz on rectangular matrices with aspect ratios 3.25, 3.5, and 8, with comparable approximation accuracy. In pretraining Modded-NanoGPT #91, Qwen3-0.6B, and Gemma 3-1B, Polar Switch reaches final validation losses similar to those of established orthogonalization methods. Across the more than 14 million matrices that each of the two recommended methods orthogonalized in these runs, the estimated output spectral norms never exceed the designed bounds by more than 3%, demonstrating their numerical stability.
Figures & tables
Method
Memory
FLOPs
Final interval
Jordan [ 6 ]
n2(2α+2)
(15α+5)n3
[0.562544,1.437456]
Polar Express [ 15 ]
n2(2α+2)
(15α+5)n3
[0.858872,1.141128]
Gram Newton–Schulz [ 16 ]
n2(2α+4)
(6α+14)n3
[0.787021,1.212979]
Polar Switch 2-55
n2(2α+2)
(6α+10)n3
[0.872263,1.106236]
Polar Switch 2-64
n2(2α+2)
(6α+10)n3
[0.897110,1.081213]
Polar Switch 2-54
n2(2α+2)
(6α+9)n3
[0.771715,1.201363]
Table 1: Memory in matrix elements, FLOPs, and final singular value intervals for X∈Rn×m with n≤m and α=m/n . In exact arithmetic, each method maps singular values in [0.001,1] into its final interval. Jordan and Gram Newton–Schulz multiply the output of their final iteration by a constant that centers this interval at one. Section A.2 gives the iterations of Polar Switch .
Model
Parameters
Tokens
Total FLOPs
Seeds
Total runs
Modded-NanoGPT #91
637,656,487
343,408,640
(2.413–2.418)×1017
20
260
Qwen3-0.6B
596,049,920
12,000,428,032
(5.148–5.155)×1019
1
13
Gemma 3-1B
999,885,952
22,000,173,056
(1.360–1.363)×1020
1
10
Table 2: LLM training setups. Of the Polar Switch methods, the Gemma 3-1B setup uses only the methods with two outer iterations. Total FLOPs are estimated for one run and depend on the orthogonalization method.
Figure 1: Final validation loss of Modded-NanoGPT #91 (left), and of Qwen3-0.6B and Gemma 3-1B (right). On the left, markers show runs with the same 20 seeds, vertical ticks show means, and bars span the minimum and maximum. Triangles mark baselines, circles mark PS2 , and squares mark PS3 . On the right, all methods use the same seed, and a dash marks methods not trained on Gemma 3-1B.
Modded-NanoGPT #91
Qwen3-0.6B
Gemma 3-1B
Method
Bound
3–30
30–300
> 300
3–30
30–300
> 300
3–30
30–300
> 300
Jordan
1.4375
0
0
0
0
0
0
0
0
0
PE
1.1412
0
0
0
0
0
0
0
0
0
PE (BF16)
1.1412
0
0
0
0
0
0
0
0
0
GNS
1.2130
4
0
0
0
0
0
0
0
0
PS2-55
1.1063
0
0
0
0
0
0
0
0
0
Table 3: Number of matrices whose spectral norm estimate exceeds the bound in column Bound, the upper endpoint in Table 1 rounded up to four decimals, by 3–30%, 30–300%, and more than 300%, out of 2,683,200 for Modded-NanoGPT #91, 4,486,244 for Qwen3-0.6B, and 7,637,084 for Gemma 3-1B per method. PE (BF16) uses the bound of PE . A dash marks methods not trained on Gemma 3-1B.
Figure 2: Median runtime in milliseconds for single matrices (top) and batches (bottom). Colors show median runtime ratios relative to PE (left) and GNS (right, rectangular shapes only). Values below one indicate lower runtime. Column labels give tensor dimensions, with the batch size first.
Figure 3: Polar factor error ( 10 ) on synthetic matrices with minimum singular value σmin=10−3 (left) and σmin=10−4 (middle), 1,500 matrices each, and on 33,592 matrices from one Modded-NanoGPT #91 run with PE (right). Percentages give the fraction of pairs of a matrix and another method of the panel in which the method has the smaller error, and a tie counts as half.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Method
Iteration t
ct
bt
at
Jordan
1–4
3.4445
−4.775
2.0315
5
4.117967776458191
−5.708606802899656
2.4286983706996126
PE
1
8.156554524902461
−22.483292925577953
15.878769915207457
2
4.042929935166734
−2.8089174659087113
0.500017845105131
3
3.891667802292659
−2.772484153217684
0.506064817850339
4
3.2857536577556523
−2.3681294933425354
0.4644902423300304
Appendix
Table 4: Baseline coefficients for pt(x)=ctx+btx3+atx5 , including the safety factor of each method. Jordan and GNS are centered at the last iteration.
Figure 4: Diagonal drift (top) and directional mixing (bottom), defined in ( 13 ) and ( 14 ), with the panels and conventions of Figure 3 .
Figure 5: Runtime ratios relative to PE . Values below one indicate lower runtime. GNS appears only for rectangular shapes, because it falls back to PE for square matrices. Markers show medians across five runs, and bars span the minimum and maximum.
Muon optimizers improve neural-network training by replacing ill-conditioned momentum updates with approximately semi-orthogonal updates. This motivates a practical question: how much orthogonalization does Muon actually require? We study this question using a relaxed cubic Newton--Schulz schedule derived directly for Muon's low precision singular value band. The resulting five-step cubic construction uses ten dominant matrix multiplications, compared with fifteen for five quintic Newton--Schulz iterations. The cubic schedule is not intended as a more accurate polar solver; instead, it is a principled low-cost variant that lets us probe the relation between polar accuracy, spectral shaping, and training quality. Across synthetic diagnostics, NanoGPT ablations, and training experiments on hybrid MoE/Mamba models, we find that training quality is not governed monotonically by polar-decomposition accuracy: truncated Polar Express, Muon-Jordan, cubic Newton--Schulz, and an explicit FP32 SVD polar factor can reach nearly indistinguishable final loss on GPT-2 Small, and cubic5 matches the Muon-Jordan quintic update within about 10−3 validation loss on hybrid MoE/Mamba models with one billion to four billion parameters. These results support cubic5 as a practical low-cost Muon orthogonalization variant, with empirical evidence of training-quality parity in the settings tested.
Muon is an optimizer that computes updates using the polar factor of the momentum matrix and has shown strong empirical performance across a range of training settings. A key component of Muon is the Newton-Schulz iteration used to compute this polar factor. Although this avoids the cost of an exact singular value decomposition, it remains expensive in practice because it is applied at every optimization step. At the same time, the momentum matrix changes smoothly over training, suggesting strong temporal correlation in the corresponding polar factors. In this paper, we exploit this structure and propose CacheMuon, a temporal preconditioning method that reuses information from previous optimization steps to approximate the polar factor at the current step. This reduces redundant orthogonalization computation across iterations. We analyze CacheMuon as an inexact Muon update, with error controlled by fresh-solver error and cache staleness. Empirically, CacheMuon provides a controllable quality-efficiency frontier: conservative thresholds closely match fresh Muon on language-model and vision training while reducing orthogonalization FLOPs, whereas more aggressive thresholds yield larger arithmetic savings at the cost of modest validation-quality degradation.
Bishnu Dev, Sushil Bohara, Martin Takáč +1
Mohamed bin Zayed University of Artificial Intelligence, Abu Dhabi, UAE · Mohamed bin Zayed University of Artificial Intelligence
Orthogonalized-update optimizers such as Muon improve training of matrix-valued parameters, but existing extensions typically either rescale updates after orthogonalization or use heavier whitening-based preconditioners before it. We introduce {\method}, a lightweight family of pre-orthogonalization equilibration schemes for Muon with three forms: two-sided row/column normalization (RC), row normalization (R), and column normalization (C). By rebalancing the momentum matrix before finite-step Newton--Schulz orthogonalization, {\method} improves the geometry seen by orthogonalization. We show that finite-step orthogonalization is governed by the input spectrum, especially stable rank and condition number, and that row/column normalization acts as a zeroth-order surrogate for whitening. For hidden matrix weights, R is the default variant. Theoretically, {\method} (R) retains the standard O(T−1/4) Muon-type nonconvex stationarity guarantee with decoupled weight decay and a horizon-free diminishing learning-rate schedule, and extends it to finite-step NS5 up to an explicit inexactness constant. In LLaMA2 pretraining on C4, {\method} (R) consistently outperforms Muon on 130M, 350M, and 1B models, with faster convergence and lower validation perplexity. The code is available at the MuonEq codebase.
Da Chang, Qiankun Shi, Lvgang Zhang +5
Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences · Pengcheng Laboratory · University of Chinese Academy of Sciences +4