Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

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

TL;DR

The paper introduces Polar Express, a GPU-friendly method for computing the polar decomposition via adaptive compositions of odd polynomials, optimized to minimize worst-case error in the spectral norm. By solving a sequence of per-iteration minimax problems, Polar Express achieves rapid early convergence and strong asymptotic behavior without inverses or QR factorizations, making it well-suited for low-precision GPU workflows. The method is integrated into the Muon optimizer for large-scale GPT-2 training, where it improves validation loss across learning rates compared to prior polynomial-based approaches. Comprehensive analysis shows quadratic to cubic convergence depending on degree, finite-precision stabilization techniques, and practical efficacy demonstrated on GPT-2 models trained on 1B tokens from the FineWeb dataset. Overall, Polar Express provides a principled, high-throughput alternative to traditional polar-decomposition methods in deep learning contexts.

Abstract

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 algorithm 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 the Muon training framework, our method leads to consistent improvements in validation loss when training a GPT-2 model on one billion tokens from the FineWeb dataset, outperforming recent alternatives across a range of learning rates.

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

TL;DR

The paper introduces Polar Express, a GPU-friendly method for computing the polar decomposition via adaptive compositions of odd polynomials, optimized to minimize worst-case error in the spectral norm. By solving a sequence of per-iteration minimax problems, Polar Express achieves rapid early convergence and strong asymptotic behavior without inverses or QR factorizations, making it well-suited for low-precision GPU workflows. The method is integrated into the Muon optimizer for large-scale GPT-2 training, where it improves validation loss across learning rates compared to prior polynomial-based approaches. Comprehensive analysis shows quadratic to cubic convergence depending on degree, finite-precision stabilization techniques, and practical efficacy demonstrated on GPT-2 models trained on 1B tokens from the FineWeb dataset. Overall, Polar Express provides a principled, high-throughput alternative to traditional polar-decomposition methods in deep learning contexts.

Abstract

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 algorithm 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 the Muon training framework, our method leads to consistent improvements in validation loss when training a GPT-2 model on one billion tokens from the FineWeb dataset, outperforming recent alternatives across a range of learning rates.

Paper Structure

This paper contains 31 sections, 6 theorems, 53 equations, 9 figures, 4 algorithms.

Table of Contents

  1. Introduction
  2. The Muon Method
  3. Computing the Polar Factor
  4. The Newton-Schulz methods.
  5. The Jordan and You methods.
  6. Contributions
  7. Notation.
  8. Related Work
  9. Newton-Schulz and polynomial Padé methods.
  10. Newton's method and rational Padé.
  11. Adaptive rational methods from optimal approximations.
  12. Adaptive polynomial methods.
  13. Application within Muon.
  14. Approximations by Compositions of Polynomials
  15. The Polar Express
  16. ...and 16 more sections

Key Result

Theorem 0

Let $d$ be odd and define $\ell_1 = \ell$ and $u_1 = u$. For $t = 1,\ldots,T$ define The resulting composition $p^{\star}:=p_T \circ p_{T-1} \circ \cdots \circ p_1$ is optimal and the error is given by: Furthermore the new error, lower and upper bounds can be computed through

Figures (9)

  • Figure 1: Training a GPT-2-Large model (774M params) on 1 billion tokens from the FineWeb dataset aroca2023fineweb. The label muon-<name> refers to implementing Muon using <name> to compute the polar factor. Left: final validation loss across learning rates. Right: validation loss across epochs using the best learning rate. The best learning rate ($lr$) and final validation loss for each method was adamw$(lr= 0.0001)$: $4.172$, muon-You$(lr= 0.02)$: $3.400$, muon-Jordan$(lr= 0.02)$: $3.398$ and muon-PolarExp$(lr= 0.02)$: $3.340$.
  • Figure 2:
  • Figure 3: Effects of stabilizing the update rules with a safety factor and cushioning, as described in \ref{['sec:finite_prec']}. The blue curve is the optimal degree-5 polynomial for the interval $[0.005, 1]$. It is has numerical issues because it maps singular values near $0.8$ down to almost zero and maps $1+\epsilon$ to $\approx u_{t+1}+25\epsilon$. The stabilized version is better because it ensures $\frac{p_t(x)}{x} \geq 0.236$ and maps all $x \leq 1.01$ to at most $u_{t+1}$.
  • Figure 4: Convergence of various degree-5 polynomial methods in the spectral norm. When tuned properly, Polar Express attains outperforms the other methods at every iteration. Left panel: synthetic matrix with $\sigma_{\max} = 1$, $\sigma_{\min} = 10^{-6}$. Right panel: gradient of a certain weight matrix of a randomly-initialized GPT-2 architecture on a batch of language modeling data, normalized by the Frobenius norm.
  • Figure 5: Convergence of polynomial methods in the Frobenius norm on GPT-2 gradient matrices. The number of matrix-matrix products is $T(d+1)/2$, where $d$ is the degree ($3$ for Chen & Chow; $5$ for all others) and $T$ is the number of iterations.
  • ...and 4 more figures

Theorems & Definitions (13)

  • Theorem 0
  • proof
  • remark 1: Why a fixed degree?
  • Theorem 1
  • proof
  • lemma 2
  • proof
  • lemma 3
  • proof
  • Theorem 3
  • ...and 3 more