PaperlayerView exact revision on arXiv

Source-linked AI summary

The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm

Noah Amsel, David Persson, Christopher Musco, Robert M. Gower

arXiv:2505.16932v5cs.LGcs.AIcs.CLmath.NAmath.OC

TL;DR

Deep learning needs fast, GPU-friendly polar-decomposition methods because Muon uses the polar factor while prioritizing throughput over high precision. Polar Express adaptively composes minimax odd polynomials, with finite-precision safeguards for bfloat16, and consistently improves GPT-2 validation loss over prior Muon matrix-sign methods across learning rates and training lengths.

  • Problem

    Deep-learning applications need matrix-function methods optimized for GPU throughput rather than the precision priorities of classical computational science.

  • Method

    Polar Express adaptively selects minimax polynomial updates whose composition optimally approximates polar(M) in worst-case spectral-norm error, with modifications for bfloat16 execution.

  • Results

    Polar Express within Muon consistently achieves lower GPT-2 validation loss than recent matrix-sign alternatives across learning rates, including experiments extended from 1 to 10 billion FineWeb tokens.

  • Takeaways & Limitations

    Five to six iterative polynomial updates provide the recommended Muon implementation instead of exact SVD-based polar computation, which doubles training-step runtime without improving optimization quality.

  • Takeaways & Limitations

    The method requires guessed lower singular-value bounds, although it converges for any positive guess and an order-of-magnitude error delays convergence by only a few iterations.

Abstract

from arXiv · show

Computing the polar decomposition and the related matrix sign function has been a well-studied problem in numerical analysis for decades. Recently, it has emerged as an important subroutine within the Muon optimizer for training deep neural networks. However, the requirements of this application differ sharply from classical settings: deep learning demands GPU-friendly algorithms that prioritize high throughput over high precision. We introduce Polar Express, a new method for computing the polar decomposition. Like Newton-Schulz and other classical polynomial methods, our approach uses only matrix-matrix multiplications, making it very efficient on GPUs. Inspired by earlier work of Chen & Chow and Nakatsukasa & Freund, Polar Express adapts the update rule at each iteration by solving a minimax optimization problem. We prove that this strategy minimizes error in a worst-case sense, allowing Polar Express to converge as rapidly as possible both in the early iterations and asymptotically. We also address finite-precision issues, making it practical to use in bfloat16. When integrated into Muon, our method yields consistent improvements in validation loss for a GPT-2 model trained on one to ten billion tokens from the FineWeb dataset, outperforming recent alternatives across a range of learning rates.

1 INTRODUCTION

Deep learning uses matrix functions in optimizers but prioritizes GPU throughput over the high precision typical of computational science. Polar Express targets this setting by approximating polar(M) with adaptive polynomial updates and improves Muon training outcomes.

  • Matrix sign and polar-decomposition computations have become important subroutines in deep-learning optimizers, including Muon.
  • In GPT-2 training on 1 billion FineWeb tokens, muon-PolarExp achieved final validation loss 3.340, versus 3.399 for muon-You and 3.398 for muon-Jordan, all at lr = 0.02.
  • Muon replaces the momentum step −M_t with −polar(M_t), whose direction is steepest descent under the spectral norm.
  • Polar Express dynamically adapts polynomial updates, optimizing worst-case approximation error while emphasizing rapid early progress and later accuracy.
  • The method uses GPU-friendly matrix operations and finite-precision modifications, addressing deep learning’s throughput and bfloat16 requirements.

2 APPROXIMATIONS BY COMPOSITIONS OF POLYNOMIALS

The paper restricts GPU computation to matrix linear combinations and products, representing iterative polar approximations as compositions of odd polynomials. It then chooses the composition that minimizes worst-case spectral-norm error over bounded singular values.

  • GPU-friendly computation is restricted to matrix linear combinations and matrix-matrix products, enabling odd polynomial evaluation without explicit SVDs.
  • Composing T degree-d polynomials yields degree d^T while requiring only O(Td) matrix-matrix products.
  • Polar Express selects the composition that minimizes worst-case spectral-norm error for given singular-value bounds, degree, and iteration count.
  • The matrix approximation problem reduces to uniformly approximating the constant function 1 over the singular-value interval [ℓ, u].

3 THE POLAR EXPRESS

Polar Express greedily solves a sequence of minimax odd-polynomial problems, propagating singular-value bounds and precomputing updates. The resulting method has rapid convergence, supports bfloat16 execution, and is reported to outperform degree-5 alternatives in convergence experiments.

  • 3.1 GREEDY IS OPTIMAL: Each iteration chooses an odd degree-d polynomial that minimizes the maximum error |p(x) − 1| over the current interval [ℓ_t, u_t].
  • 3.1 GREEDY IS OPTIMAL: The greedy composition is globally optimal, with new bounds computed as ℓ_t+1 = p_t(ℓ_t) and u_t+1 = 2 − ℓ_t+1.
  • 3.1 GREEDY IS OPTIMAL: Polar Express converges at least as fast as same-degree Newton-Schulz; degree 3 and degree 5 achieve quadratic and cubic convergence, respectively.
  • 3.1 GREEDY IS OPTIMAL: When σ_min = ℓ, Polar Express is about twice as fast as Newton-Schulz.
  • 3.5 THE ALGORITHM: Remez or closed-form procedures compute the polynomials offline, while the online stage applies them in bfloat16 using matrix products.
  • 3.5 THE ALGORITHM: With proper tuning, Polar Express outperforms other degree-5 polynomial methods at every iteration on synthetic and GPT-2 gradient matrices.
  • 3.5 THE ALGORITHM: The recommended deep-learning configuration uses degree d = 5, T = 5 or 6, and ℓ_1 = 10^-3.

4 NUMERICAL EXPERIMENTS

The experiments evaluate Polar Express within Muon across model sizes, training durations, iteration counts, and weight-decay settings. Polar Express consistently achieves lower validation loss than the compared methods, while five or six iterations preserve optimization quality without SVD’s runtime cost.

  • Training GPT-2: 4.197, 3.639, 3.629, and 3.588 are the best final validation losses for AdamW, muon-Jordan, muon-You, and muon-PolarExp, respectively, on GPT-2-Small.These results use 1 billion FineWeb tokens without weight decay; the reported learning rates differ by method.
  • Training GPT-2: Across GPT-2 training runs, muon-PolarExp achieves better validation loss than muon-Jordan and muon-You across learning rates and epochs.The three methods use equally expensive degree-5 polynomial iterations, so the validation-loss advantage also appears in training steps and wall-clock comparisons.
  • Ablations: 2.921, 2.919, and 2.913 are the best final validation losses for muon-Jordan, muon-You, and muon-PolarExp on GPT-2-Large trained with 10 billion tokens.Polar Express retains a small but consistent advantage at the longer training duration.
  • Ablations: Five or six Polar Express iterations provide the standard Muon trade-off: fewer iterations worsen final validation loss, while more iterations or exact SVD do not improve optimization quality.The SVD doubles the runtime of each training step, whereas changing the iterative count does not meaningfully affect runtime.
  • Ablations: Adding weight decay of 0.1 produces results similar to the no-weight-decay experiments, with muon-PolarExp again outperforming the other methods.The weight-decay experiments keep the other training settings unchanged.

B RELATED WORK

Related work spans classical polynomial and rational iterations, adaptive optimal approximations, and Muon-specific GPU-oriented methods. Polar Express adopts adaptive approximation with polynomials, avoiding inverses and QR factorizations while establishing global optimality for the composed update.

  • Classical iterations: Classical Newton-Schulz and Padé methods can converge slowly initially, while Newton and other rational methods require explicit inverses.Rescaling can accelerate rational iterations, but Newton-Schulz may diverge when the largest singular value greatly exceeds the smallest.
  • Adaptive rational methods: Adaptive rational methods choose approximations over an interval containing the singular values, with Zolotarev functions providing optimal formulas for each degree.For degree q = 1, the optimal rational approximation coincides with dynamically weighted Halley’s method.
  • Adaptive polynomial methods: Polar Express transfers the adaptive-approximation paradigm to odd polynomials, changing the polynomial at each iteration rather than using a fixed Newton-Schulz update.This polynomial choice avoids QR factorizations and uses only GPU-friendly matrix-matrix multiplications in low-precision arithmetic.
  • Adaptive polynomial methods: The paper proves that selecting an optimal polynomial at every iteration produces a globally optimal composed polynomial under the stated supremum-norm criterion.This gives the adaptive polynomial approach a theoretical foundation connected to the Zolo-pd paradigm.
  • Application within Muon: Muon favors polynomial iterations because they use half-precision matrix products rather than inverses or QR decompositions, reflecting its preference for speed over accuracy.This application setting makes many rational-based numerical-analysis methods unsuitable despite their convergence properties.
  • Application within Muon: Unlike related adaptive polynomial proposals, this work addresses finite-precision considerations and establishes global optimality of the composed polynomial.The comparison is stated explicitly for concurrent work and earlier adaptive approaches.

C PROOF OF THEOREM 3.1

The proof establishes that composing iteration-specific odd polynomials yields the optimal worst-case approximation over the target interval. It proceeds from the equioscillation characterization for one step and extends optimality inductively across T steps.

  • Base case: The proof uses Chebyshev’s equioscillation theorem to characterize the unique optimal odd polynomial approximation.Optimality is equivalent to alternating extremal errors at q + 2 points.
  • Base case: The base-case argument shows that the error has exactly q stationary points in the interval and that the lower endpoint is a maximum.The stationary-point count and endpoint analysis force the required alternation structure.
  • Inductive step: Theorem 3.1 defines interval updates from each polynomial’s minimum and maximum, producing a composed polynomial p⋆ that is optimal.The bounds evolve through ℓt+1 = pt(ℓt) and ut+1 = 2 − ℓt+1.
  • Inductive step: The induction compares any alternative composition with the constructed composition and derives a contradiction from the one-step optimality of each pt.Input rescaling preserves the odd-polynomial space, allowing the inductive comparison to proceed.
  • Conclusion: The resulting error after T steps is determined by the updated lower endpoint, with max error equal to 1 − ℓT+1.This follows from the endpoint extremum and the recurrence for the interval bounds.

D PROOF OF THEOREM 3.3

The convergence proof applies Theorem 3.1 to a normalized matrix and transfers scalar polynomial error bounds to the matrix approximation. It concludes that degree 3 and degree 5 updates converge quadratically and cubically, respectively.

  • Convergence rate: For d = 3 and d = 5, the method converges quadratically and cubically, respectively.These rates are the explicit degree-dependent consequence highlighted by the theorem.
  • Algorithmic realization: Algorithm 1 returns the composed polynomial approximation p⋆(M) by applying the iteration polynomials to M.The construction uses odd polynomials formed from Padé-based updates.

E PROOF OF EQUIVALENCE BETWEEN (5) AND (6)

The proof establishes equivalence between the scalar minimax polynomial problem and the matrix approximation problem by reducing spectral-norm error to the worst scalar error on the singular-value interval.

  • Reduction: Unitary invariance of the spectral norm reduces the matrix objective for a fixed polynomial to a scalar approximation problem.The reduction preserves the relevant worst-case error over [ℓ, u].
  • Sharpness: A rank-compatible test matrix attaining the worst scalar error shows that the scalar bound is sharp for rectangular matrices.Choosing M = x⋆UVᵀ gives polar(M) = UVᵀ and realizes the selected scalar error.

F REMEZ ALGORITHM

The Remez section explains how Polar Express obtains optimal update polynomials when closed forms are unavailable. It gives a closed-form cubic, an iterative degree-5 procedure, and a limiting Padé fallback near a degenerate interval.

  • Degree-3 case: The degree-3 update is obtained by enforcing equioscillation at the positive local extremum and both interval endpoints.Writing p(x) = ax + bx^3 reduces the conditions to three equations.
  • General degree: For larger odd degrees, Remez iteratively refines trial equioscillation points and polynomials toward the optimal approximation.The method typically reaches double-precision accuracy in about 10 iterations.
  • Degree-5 algorithm: The degree-5 implementation fixes the interval endpoints and initializes two interior trial points before repeatedly solving for coefficients and refining extrema.The coefficient system is linear for fixed trial points, while the next points come from extrema of the error polynomial.
  • Implementation: In experiments, Algorithm 2 never required more than five iterations to converge.The algorithm stops when successive maximum-error estimates differ by approximately double-precision machine epsilon.
  • Numerical boundary: When ℓ/u approaches 1, the Vandermonde system becomes ill-conditioned and the optimal polynomial approaches a scaled Padé approximant.For sufficiently close endpoints, the algorithm returns the scaled Newton-Schulz polynomial up to double-precision error.

G FINITE PRECISION CONSIDERATIONS

Finite-precision arithmetic introduces two failure modes for Polar Express: singular values can exceed the assumed upper bound, and non-monotonic polynomial updates can cause precision loss or sign errors. Safety factors, cushioning, and normalization stabilize the method for practical use.

  • Numerical stability: Round-off can push singular values above the current upper bound, causing iterative growth instead of convergence.The authors address this by replacing p_t(x) with p_t(x/1.01).
  • Numerical stability: The optimal polynomial oscillates across the singular-value interval, potentially mapping values near the upper bound down to the lower bound.This non-monotonicity can cause a singular vector to change sign if p_t(σ_i) becomes negative.
  • Numerical stability: A negative polynomial image can make the iteration converge to the polar factor of the wrong matrix.Unlike Newton-Schulz, unscaled Newton, or QDWH, Polar Express is affected by this loss of precision.
  • Stabilization changes: The method uses cushioning when the lower bound is much smaller than the upper bound, slightly slowing convergence to prevent problematic updates.When ℓ_t < u_t/10, the update is selected as though ℓ_t = u_t/10.
  • Stabilization changes: The implementation also normalizes by ∥M∥F + 10^-2 rather than ∥M∥F alone.This modification is inherited from the original Muon implementation.

H.1 CONVERGENCE OF POLAR EX P R E S S AND ITS IMPACT ON MUON

On GPT-2 gradient matrices, Polar Express converges rapidly and outperforms competing degree-5 polynomial methods under both Frobenius-norm and cosine-similarity evaluations. Muon performance saturates after roughly five or six iterations despite incomplete convergence of the underlying polar approximation.

  • Stabilization: The degree-5 polynomial curve can suffer numerical issues by mapping singular values near 0.8 close to zero and amplifying values above 1.The figure contrasts this behavior with a stabilized version.
  • Convergence: Polar Express outperforms other degree-5 polynomial methods on GPT-2 gradient matrices.The tests use gradients from two randomly initialized GPT-2 layers and evaluate Frobenius norm and cosine similarity.
  • Impact on Muon: Muon performance does not improve beyond five or six Polar Express iterations, although five iterations do not fully converge to the exact polar factor.The resulting update directions can still differ meaningfully from the SVD-computed polar(M).
  • Impact on Muon: Small singular values below 10^-4σmax do not affect Muon performance, while values below 10^-3σmax have only a minor effect.The experiment compares exact, truncated, and reverse treatments of the small singular values.
  • Convergence: Polar Express converges in six iterations under relative Frobenius norm and five iterations under cosine similarity.The comparison ignores singular values below 10^-3 and uses a truncated polar factor.

H.2 TRAINING GPT-2

Additional GPT-2 experiments show that weight decay has little effect on the comparative results, while extending training to 10 billion tokens preserves Polar Express’s advantage but narrows the gap as training loss converges. Image-classification experiments do not establish a consistent advantage.

  • GPT-2 training: The GPT-2 experiments report validation loss and training loss across learning rates, iteration counts, and training time.The additional metrics are shown for the 1-billion-token experiments.
  • GPT-2 training: Weight decay of 0.1 has little effect on the GPT-2 comparative results.The rerun results are reported in Figure 12.
  • GPT-2 training: Polar Express still outperforms baselines across all 10-billion-token conditions, but the gap shrinks as training loss converges.The extended runs cover GPT-2 training on FineWeb data.
  • Image classification: Vision Transformer experiments achieved the best training and validation loss with Muon using Polar Express, but this did not consistently yield the best validation accuracy.The authors state that these experiments do not show a consistent advantage for Polar Express.

I INITIALIZATION FOR MATRICES WITH LARGE SPECTRAL GAPS

For matrices with large spectral gaps, the paper proposes spectrum-aware initialization that targets outlying singular values separately from the spectral tail. This extra initialization can accelerate early progress when the largest singular value is well separated.

  • Motivation: Large spectral gaps allow the polynomial to focus separately on outlying singular values and the remaining spectral tail.The standard construction is optimal over the full interval when no additional spectral information is assumed.
  • Initialization procedure: Subspace iteration estimates leading singular values and supplies bounds for the remaining tail.For k selected singular values, the estimates provide lower bounds for the leaders and an upper bound for σk+1 through σn.
  • Experimental context: The experiments report GPT-2 training curves across learning rate, iteration count, and training time for 1-billion- and 10-billion-token settings.The figures use five iterations of each method to compute polar(M).
  • Experimental context: On CIFAR-10 with ResNet-20, muon-PolarExp reaches validation accuracy 0.893 at learning rate 0.001.The figure compares this result with SGD, AdamW, and other Muon variants.
  • Initialization procedure: The spectrum-aware initializer constructs an odd polynomial that approximates the constant function on intervals containing the outliers and on the tail.The k = 1 construction uses a lower bound z for the largest singular value after Frobenius normalization.
  • Results: For a power-law spectrum with σj(M) = j^-5, the extra initialization benefits both Newton-Schulz and Polar Express.The initialization costs about one degree-5 iteration and is counted as an additional iteration for comparison.

J FAST POLYNOMIAL ITERATION FOR RECTANGULAR MATRICES

Algorithm 3 accelerates repeated odd-polynomial iterations on rectangular matrices by avoiding repeated rectangular multiplications, with benefits that grow for large aspect ratios. The section also examines numerical safeguards and a transformer application, where these computational changes had little effect at the tested scale.

  • Fast polynomial iteration: Algorithm 3 applies any sequence of odd-polynomial iterations while using only two rectangular matrix multiplications across all T iterations.This replaces the naive method’s two rectangular multiplications per iteration.
  • Runtime: Algorithm 3 can be selected at runtime when α > 1.5 T/(T − 1), because its estimated cost is then lower than the baseline.The criterion compares Algorithm 3 with the naive method for the relevant dimensions and iteration count.
  • Numerical stability: Low-precision computation can become unstable from ill-conditioned X and spurious negative eigenvalues in Y.The proposed remedies are restarting after three iterations and adding 10^-3I during the first application of Algorithm 3.
  • Runtime: 6 degree-5 iterations achieved almost a 2x GPU speedup at α = 4 and a 5x speedup at α = 32.Restarting every three iterations preserved the trend but reduced the runtime savings.
  • Application to Muon: In a GPT-Small FineWeb1B experiment, head splitting and Algorithm 3 negligibly affected optimization quality, while all 12 runtimes were nearly identical.Final loss changed from 3.59 to 3.55 with head splitting, and Algorithm 3 changed it to 3.60 without splitting or 3.56 with splitting.
Loading 2505.16932v5…