The Polar Express Algorithm

In this note we look at how the sign function can be approximated by compositions of polynomials, as in the Polar Express algorithm [1]. This simple idea has a number of applications in modern machine learning and optimization, as it can be used to approximate matrix functions using only matrix-matrix products, which are very efficient to compute on modern hardware such as GPUs. Applied to the singular values of a matrix, the approximation of the sign function gives the polar factor π‘ˆπ‘‰βŠ€. This is, e.g., the basis of the update of the Muon optimizer [2], which would otherwise require an SVD, an expensive computation that is not very GPU friendly. Muon is explained as steepest descent in the spectral norm in Which Way Is Down? It is furthermore possible to approximate the ReLU function in a similar way, which can be used to project matrices onto the cone of positive semidefinite matrices.

1. The Singular Value Transformation

The singular value transformation is defined for any odd function 𝑓, as

𝑓(𝑀)β‰”π‘ˆπ‘“(Ξ£)π‘‰βŠ€,𝑀=π‘ˆΞ£π‘‰βŠ€.

(1)

When 𝑀 is square and symmetric, the singular value transformation of an odd 𝑓 coincides with the matrix function 𝑓(𝑀)=𝑄𝑓(Ξ›)π‘„βŠ€, which applies 𝑓 to the eigenvalues. Writing 𝑀=π‘„Ξ›π‘„βŠ€=(𝑄sign(Ξ›))|Ξ›|π‘„βŠ€ shows that π‘ˆ=𝑄sign(Ξ›), Ξ£=|Ξ›| and 𝑉=𝑄, and since 𝑓 is odd,

𝑓(𝑀)=𝑄𝑓(Ξ›)π‘„βŠ€=𝑄sign(Ξ›)𝑓(|Ξ›|)π‘„βŠ€=π‘ˆπ‘“(Ξ£)π‘‰βŠ€,

(2)

where the sign function is defined as the odd function

sign(π‘₯):={1,π‘₯>0,0,π‘₯=0,βˆ’1,π‘₯<0,

(3)

2. Approximating the Polar Normalization

Polar normalization corresponds to the singular value transformation when 𝑓 itself is the sign function, i.e.

polar(𝑀)=π‘ˆdiag(sign(𝜎1),…,sign(πœŽπ‘˜))π‘‰βŠ€=π‘ˆπ‘‰βŠ€,

(4)

where the last equality takes π‘ˆ and 𝑉 from the compact SVD, which keeps only the singular vectors of the nonzero singular values. Zero singular values stay zero, as sign(0)=0.

Computed directly, the polar factor requires an SVD, which is computationally expensive and not very GPU friendly (it is primarily memory bound rather than compute bound). The Polar Express algorithm instead approximates the matrix sign function using just matrix-matrix products [1]. The key idea is rather simple: Approximate the function 𝑓 using a composition of polynomials, e.g.,

𝑓(π‘₯)β‰ˆπ‘π‘˜(π‘π‘˜βˆ’1(…𝑝2(𝑝1(π‘₯))…)).

(5)

Since the above is a composition of polynomials, it can be evaluated using just additions and multiplications.

ExampleΒ 1: Newton–Schulz

Example 1 (Newton–Schulz)Β Β  The idea of approximating functions using repeated polynomial application is not new. An example is the Newton–Schulz algorithm which is based on repeatedly applying the same polynomial, i.e. 𝑝1=𝑝2=…=π‘π‘˜, in order to approximate the sign function. For order 3 and 5 the polynomials look as follows

𝑝(π‘₯)=32π‘₯βˆ’12π‘₯3,𝑝(π‘₯)=(15π‘₯βˆ’10π‘₯3+3π‘₯5)/8.

(6)

While the Newton–Schulz method converges for all π‘₯∈[βˆ’1,1], the convergence is only fast for values close to Β±1. Conversely, for values that are close to zero the convergence is slow. The slow convergence is exactly what the Polar Express method explained shortly aims to fix. The graph shows the cubic and quintic polynomials composed π‘˜ times, together with the sign function (dotted). The shaded interval [βˆ’1,1] is where they approximate it.

A note here is that for non-square matrices we cannot simply do 𝑀𝑛. Instead, what is done is that 𝑀2𝑛+1=𝑀(π‘€βŠ€π‘€)𝑛=(π‘€π‘€βŠ€)𝑛𝑀.

3. The Polar Express Algorithm

The basic question that the Polar Express algorithm aims to solve is as follows: Given a set of 𝑇 polynomials {𝑝𝑖(π‘₯)=π‘Ž1𝑖π‘₯+π‘Ž3𝑖π‘₯3+π‘Ž5𝑖π‘₯5}𝑖=1𝑇, can we find an optimal set of coefficients that minimizes the maximum error over a given interval [β„“,𝑒]? Define their composition as 𝑝=π‘π‘‡βš¬π‘π‘‡βˆ’1βš¬β‹―βš¬π‘1. We then solve

(𝑝1⋆,…,𝑝𝑇⋆)=arg min𝑝1,…,𝑝𝑇maxπ‘€βˆˆβ„π‘šΓ—π‘›πœŽ(𝑀)βŠ‚[β„“,𝑒]β€–polar(𝑀)βˆ’π‘(𝑀)β€–2subject toπ‘π‘‘βˆˆβ„5odd[π‘₯],𝑑=1,…,𝑇.

(7)

Here ℝ𝑑odd[π‘₯] denotes the space of real odd polynomials of degree at most 𝑑. Since 𝑝 is odd, 𝑝(𝑀)=π‘ˆπ‘(Ξ£)π‘‰βŠ€, so the error is max𝑖|1βˆ’π‘(πœŽπ‘–)|, and the problem is the scalar one of approximating 1 on [β„“,𝑒], or by oddness the sign function on [βˆ’π‘’,βˆ’β„“]βˆͺ[β„“,𝑒].

In [1] they show that the above can be solved by a greedy approach, with a caveat that the error bound only holds if the smallest singular value is at least β„“, which we do not know in advance. Smaller singular values still converge, but more slowly. A comment here is that the main focus in the paper is on bfloat16. Here the spacing between representable numbers at 1 is 2βˆ’7=0.0078125, so singular values much smaller than this are at the level of rounding errors. They therefore suggest setting β„“=10βˆ’3, and 𝑒=1 after scaling 𝑀 by its Frobenius norm so that 𝜎1≀1. In addition, they provide a few modifications in order to make the method numerically robust. These additions are skipped here, but can be found in the original paper.

Interactive figureΒ 1: Polar Express against Newton–Schulz

Interactive figure 1 (Polar Express against Newton–Schulz)Β Β  The graph compares the first π‘˜ Polar Express polynomials for [β„“,𝑒]=[10βˆ’3,1], composed, with the Newton–Schulz quintic composed π‘˜ times, drawn as in Example 1. Newton–Schulz approaches Β±1 slowly from below, while Polar Express first lifts the small values and then oscillates around Β±1, with every step shrinking the oscillation. The oscillation is the price for the smallest worst-case error. By the equioscillation theorem, the polynomial with the smallest maximum error on an interval reaches that error several times with alternating sign [3], so each optimal step has to oscillate around Β±1. Newton–Schulz instead matches the sign function closely at Β±1, where 𝑝′(1)=0, which makes it accurate close to Β±1 but slow close to zero. After a few steps the oscillations become too dense to be told apart, and soon after that too small to be seen. The readout gives both compositions at β„“, which is too close to zero to be seen in the graph.

Values at 𝜎=β„“, Newton–Schulz and Polar Express:

Interactive figureΒ 2: Polar Express on a matrix

Interactive figure 2 (Polar Express on a matrix)Β Β  The same polynomials applied to a 2Γ—2 matrix 𝑀. A matrix maps the unit circle (dotted) to an ellipse whose semi-axes are its singular values. The dashed ellipse belongs to 𝑀, scaled so that 𝜎1≀1,11 The scaling is a requirement: the polynomials are designed for singular values in [β„“,1], the positive half of [βˆ’1,1] in the previous figure, and larger singular values need not converge. Dividing 𝑀 by its Frobenius norm guarantees 𝜎1≀1, as the Frobenius norm is at least the largest singular value. and its columns (the two points) can be dragged. The solid ellipse and arrows show the matrix after π‘˜ steps. Each step applies the next polynomial to the singular values, so the ellipse turns into the unit circle, which is the ellipse of π‘ˆπ‘‰βŠ€.

After π‘˜ steps:

4. Projection onto the semidefinite cone

The positive semidefinite cone is the set of all real symmetric 𝑛×𝑛 positive semidefinite matrices. That is

π•Š+𝑛:={π‘‹βˆˆβ„π‘›Γ—π‘›|𝑋=π‘‹βŠ€,𝑋βͺ°0}.

(8)

Semidefinite programming (SDP) problems are optimization problems for which matrix variables are constrained to be positive semidefinite, a special case of the conic programs in Conic Optimization. As such, algorithms that aim to solve SDPs need to ensure that the solution satisfies the positive semidefinite constraint. In interior point methods this is done by only taking steps for which the next iterate remains inside the cone, while first-order methods most often rely on projecting back onto the set after taking a step, i.e., for a symmetric 𝑋=π‘ˆΞ›π‘ˆβŠ€,

Ξ π•Š+𝑛(𝑋)=arg minπ‘Œβˆˆπ•Š+π‘›β€–π‘Œβˆ’π‘‹β€–πΉ=π‘ˆdiag(ReLU(πœ†1),…,ReLU(πœ†π‘›))π‘ˆβŠ€.

(9)

While the idea of projecting back onto the semidefinite cone using just matrix-matrix products is not new (a similar fixed-point approach was suggested in [4]), the main takeaway is that the ideas of Polar Express can be used to find a faster converging method [5]. An initial idea would be to directly find polynomials that approximate the ReLU function, i.e., solving

inf𝑓1,…,𝑓𝑇maxπ‘₯∈[βˆ’1,1]|π‘“π‘‡βš¬π‘“π‘‡βˆ’1βš¬β‹―βš¬π‘“1(π‘₯)βˆ’ReLU(π‘₯)|subject toπ‘“π‘‘βˆˆβ„π‘‘π‘‘[π‘₯],𝑑=1,…,𝑇,

(10)

where 𝑇 is the total number of polynomials in the composition, 𝑑𝑑 is the degree of the 𝑑’th polynomial, and the degrees are fixed in advance as part of a β€œmatrix-matrix” multiplication budget. Unfortunately, directly solving the above is challenging, and instead a two-stage approach is proposed. The idea is to realize that we can write the ReLU function as

ReLU(π‘₯)=12π‘₯(1+sign(π‘₯)).

(11)

The first step is to compute optimal coefficients for approximating the sign function using the Polar Express approach. The second step is then to refine the coefficients of the polynomials found in the first step by minimizing the error to the ReLU function on the entire interval [βˆ’1,1], i.e., minimizing the following loss function

𝐿(𝑓𝑇,…,𝑓1):=maxπ‘₯∈[βˆ’1,1]|12π‘₯(1+π‘“π‘‡βš¬π‘“π‘‡βˆ’1βš¬β‹―βš¬π‘“1(π‘₯))βˆ’ReLU(π‘₯)|.

(12)

Interactive figureΒ 3: From the sign function to the ReLU

Interactive figure 3 (From the sign function to the ReLU)Β Β  The graph shows 12π‘₯(1+𝑝(π‘₯)), with 𝑝 the Polar Express composition from the previous figure, together with the ReLU (dotted). The shaded interval [βˆ’1,1] is where it approximates the ReLU. The oscillation is inherited from 𝑝, and the factor π‘₯ damps the error of 𝑝 where it is largest, close to zero.

Interactive figureΒ 4: Projecting a matrix onto the semidefinite cone

Interactive figure 4 (Projecting a matrix onto the semidefinite cone)Β Β  The same idea for a symmetric 2Γ—2 matrix 𝐴 with eigenvalues πœ†1,πœ†2 and an eigenvector at angle πœ‘. For each unit vector 𝑒 the graph draws a point at distance |π‘’βŠ€π΄π‘’| in the direction of 𝑒, solid where π‘’βŠ€π΄π‘’>0 and dashed where it is negative, so the lobes point along the eigenvectors and have lengths |πœ†1| and |πœ†2|. 𝐴 is positive semidefinite exactly when nothing is dashed.22 If 𝐴 is already positive semidefinite, the curve first moves away from 𝐴 and then returns. The first polynomials are tuned for singular values near β„“=10βˆ’3 and temporarily push values close to 1 far off. The faint lobes are 𝐴, the dotted curve its exact projection, and the bold curve the approximation 12𝐴(𝐼+𝑝(𝐴)) after π‘˜ steps. As π‘˜ grows, the dashed lobes vanish and the solid ones grow slightly, since the negative eigenvalue no longer subtracts from them.

After π‘˜ steps:True projection:

Cite this post

bibtex
@misc{paltorp2026polarexpress,
author = {Paltorp, Mikkel},
title = {{The Polar Express Algorithm}},
year = {2026},
month = oct,
url = {https://mipals.github.io/pdf/2026/01-03-polar-express/polar-express.pdf},
note = {Blog post}
}

Bibliography

  • [1] N. Amsel, D. Persson, C. Musco, and R. M. Gower, β€œThe Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm.” [Online]. Available: https://arxiv.org/abs/2505.16932
  • [2] K. Jordan et al., β€œMuon: An optimizer for hidden layers in neural networks.” [Online]. Available: https://kellerjordan.github.io/posts/muon/
  • [3] L. N. Trefethen, Approximation Theory and Approximation Practice, Extended Edition. Society for Industrial, Applied Mathematics, 2019. doi: 10.1137/1.9781611975949.
  • [4] J. B. Francisco and D. S. GonΓ§alves, β€œA fixed-point method for approximate projection onto the positive semidefinite cone,” Linear Algebra and its Applications, vol. 523, pp. 59–78, 2017, doi: 10.1016/j.laa.2017.02.014.
  • [5] S. Kang, H. Han, A. Groudiev, and H. Yang, β€œFactorization-free Orthogonal Projection onto the Positive Semidefinite Cone with Composite Polynomial Filtering.” [Online]. Available: https://arxiv.org/abs/2507.09165