Polar Switch: Fast Muon Orthogonalization with Polynomial Circuits
Abstract
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.
Related
- Bi-Lipschitz invariance of the transverse polynomial along polar arcs
- Fast Differentiable SVD on GPU via Polar Decomposition
- Stable and Gaussian Fluctuation Limit of a Lévy-Driven Slow-Fast System with Polynomial Dissipation
- Contour-Based Fast Fractional Integration in a Jacobi Fractional Polynomial Basis
- Simultaneous Clustered Orthogonalization
BibTeX
@misc{du2026polar,
title = {{Polar Switch: Fast Muon Orthogonalization with Polynomial Circuits}},
author = {Zhehang Du and Weijie Su},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.26742},
eprint = {2610.26742},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.26742},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Polar Switch: Fast Muon Orthogonalization with Polynomial Circuits
AU - Zhehang Du
AU - Weijie Su
PY - 2026
DA - 2026/10/19
PB - arXiv
DO - 10.48550/arxiv.2610.26742
UR - https://arxiv.org/abs/2610.26742
AB - 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 2^k+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.
ER -