Newton-Schulz Iteration for the msign Operator (Part 2)
In the previous post Newton-Schulz Iterations for the msign Operator (Part 1), we tried to find a better Newton-Schulz iteration for the $\mathop{\text{msign}}$ operator, aiming to reach as high an approximation quality as possible within a limited number of iteration steps. This problem can be reformulated as finding a polynomial iteration of the same form for the scalar function $\mathop{\text{sign}}(x)$. Back then, our approach was to use the Adam optimizer to find a local optimum end-to-end — effective, but somewhat crude.
A few days ago, a new paper appeared on arXiv, The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm, in which the authors employ a series of elegant mathematical results to give a much more beautiful answer, in a way that is both elegant and rigorous. In this post, let's appreciate and learn from this wonderful paper together.
Problem Statement
We won't repeat the background and the derivation here; we go straight to the problem we want to solve:
\begin{equation}\mathop{\text{argmin}}_f d(f(x),1)\end{equation}more
where $f = f_T \circ \dots \circ f_2 \circ f_1$, $\circ$ denotes function composition, $f_t(x)$ is an odd polynomial in $x$ (containing only odd powers of $x$), and $d(f(x),1)$ is some metric measuring the distance between the function $f(x)$ and $1$. In the previous post, we chose a finite number of uniformly spaced points in $[0,1]$ and used the average of the largest $k$ values of $|f(x)-1|$ as the metric. This time, we directly take the $|f(x)-1|$ maximum over the interval as the metric, namely
\begin{equation}\mathop{\text{argmin}}_f \max_{x\in[l,u]} |f(x) - 1| \label{eq:opt}\end{equation}
where $[l,u]\subset [0,1]$. Note that here $u$ can directly be taken as 1, but $l$ cannot be taken as 0, because $f(0)$ is always 0, which means the above expression is always greater than or equal to 1, and can never converge — so $l$ must be chosen as a value very close to 0. Following the analysis in the previous post, for the sake of generality we should also account for singular values of size $0.001$, so we may consider $l=0.001$.
Before starting the analysis, let's briefly explain the meaning of the word "Polar" in the paper's title — it actually refers to the "Polar Decomposition" of a matrix:
Polar Decomposition For a square matrix $\boldsymbol{M}\in\mathbb{R}^{n\times n}$, its polar decomposition is $\boldsymbol{M}=\boldsymbol{Q}\boldsymbol{S}$, where $\boldsymbol{Q}$ is an orthogonal matrix and $\boldsymbol{S}$ is a positive semi-definite matrix.
If the SVD of $\boldsymbol{M}$ is $\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$, then we exactly have
\begin{equation}\boldsymbol{M} = \boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top} = (\boldsymbol{U}\boldsymbol{V}^{\top})(\boldsymbol{V}\boldsymbol{\Sigma}\boldsymbol{V}^{\top})\end{equation}
meaning $\boldsymbol{Q}=\boldsymbol{U}\boldsymbol{V}^{\top},\boldsymbol{S}=\boldsymbol{V}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$ is precisely one solution to the polar decomposition. And as we know, when $\boldsymbol{M}$ is a full-rank matrix, $\boldsymbol{U}\boldsymbol{V}^{\top}$ is exactly $\mathop{\text{msgin}}(\boldsymbol{M})$. This is exactly why $\mathop{\text{msgin}}$ is associated with "Polar" — once you compute it, you obtain the "Polar Decomposition" of the matrix. In other words, the essential difficulty of the polar decomposition is computing $\mathop{\text{msgin}}$, which is the same challenge Muon faces.
Greedy Is Enough
Back to the main topic. For the problem $\eqref{eq:opt}$, the first — and probably the single most central — conclusion of the paper is: its greedy solution is exactly its global optimal solution! In formula terms, this means that solving problem $\eqref{eq:opt}$ can be transformed into:
\begin{equation}\begin{gathered} f^* = f_T^* \circ \dots \circ f_2^* \circ f_1^* \\[12pt] f_1^* = \mathop{\text{argmin}}_{f_1} \max_{x\in[l_1,u_1]} |f_1(x) - 1| \\ f_2^* = \mathop{\text{argmin}}_{f_2} \max_{x\in[l_2,u_2]} |f_2(x) - 1| \\ \vdots \\ f_T^* = \mathop{\text{argmin}}_{f_T} \max_{x\in[l_T,u_T]} |f_T(x) - 1| \\[24pt] l_1 = l,\quad u_1 = u, \\[8pt] l_{t+1} = \min_{x\in[l_t,u_t]} f_t^*(x),\quad u_{t+1} = \max_{x\in[l_t,u_t]} f_t^*(x) \end{gathered}\end{equation}
I believe this conclusion will surprise many readers — the author, too, was quite astonished on first seeing it, and found it worth applauding. Not only does it greatly reduce the difficulty of the problem, converting the original composite-function optimization over $T$ steps into a step-by-step single-polynomial optimization at $T=1$, it also allows us to extend the solution forward step by step while always retaining optimality (i.e., the optimal solution at step $T+1$ only requires one extra step of computation on top of the optimal solution at step $T$, rather than starting over from scratch).
It's worth pointing out that this conclusion holds even when each $f_t$ has a different degree (here "degree" refers to the highest power of the polynomial) — for example $f_1$ could be degree 3 while $f_2$ is degree 5, and so on — the conclusion that "the greedy solution is exactly the global optimum" still holds. For simplicity, however, below we'll keep all the $f_t$ at the same degree, mainly considering the degree-3 and degree-5 cases.
The full proof of the above conclusion is somewhat involved, so we'll place it at the end and first work through the follow-up steps based on this conclusion.
Equioscillation
Now that we've reduced the original problem to finding the greedy solution, we just need to focus on solving
\begin{equation}\mathop{\text{argmin}}_{f_t} \max_{x\in[l_t,u_t]} |f_t(x) - 1| \label{eq:local}\end{equation}
To solve this, we first need to recall the "Equioscillation Theorem" for odd polynomials, introduced in The Equioscillation Theorem: A Necessary and Sufficient Condition for Optimal Polynomial Approximation:
Equioscillation Theorem (Odd Case) Let $f(x)$ be an odd polynomial of degree at most $2n+1$, and let $g(x)$ be a continuous function on the interval $[a,b]\subset (0,\infty)$. Then
\begin{equation}f^* = \mathop{\text{argmin}}_f \max_{x\in[a,b]} |f(x) - g(x)|\end{equation}
holds if and only if there exist $a\leq x_0 < x_1 < \cdots < x_{n+1} \leq b$ and $\sigma\in\{0,1\}$ such that
\begin{equation}f^*(x_k) - g(x_k) = (-1)^{k+\sigma} \max_{x\in[a,b]} |f^*(x) - g(x)|\end{equation}
Here what we want to solve is $f_t$, and the target $g$ is identically 1. The equioscillation theorem tells us that $|f_t^*(x)-1|$ achieves the maximum error (denoted $\mathcal{E}$) at least $n+2$ times on $[l_t,u_t]$. It's not hard to see that the points where $|f_t^*(x)-1|$ attains its maximum can only be boundary points or extremum points of $f_t^*(x)$, and a degree-$2n+1$ odd polynomial has at most $n$ extremum points on $(0,\infty)$. So in order to "collect" enough — $n+2$ of them — we are "forced" to include the boundary points, which fixes $x_0 = l_t, x_{n+1}=u_t$, while $x_1,\cdots,x_n$ are the zeros of $\frac{d}{dx}f_t^*(x)$.
Furthermore, since the target function is $1$, the slope of $f_t^*(x)$ at $x=0$ must be greater than zero, so $l_t$ can only be the minimum point of $f_t^*(x)$, hence $\sigma=1$. Putting these results together, we actually need to solve the following system of equations:
\begin{equation}f_t(l_t) = 1 - \mathcal{E}, \quad f_t(u_t) = 1 + (-1)^n \mathcal{E},\quad f_t(x_i) = 1 + (-1)^{i+1}\mathcal{E}, \quad f_t'(x_i) = 0\end{equation}
where $i=1,2,3,\cdots,n$. We can see that both the equations and unknowns number $2n+2$, and adding the constraints on $l_t < x_1 < \cdots < x_n < u_t$ and $\mathcal{E} > 0$, in principle we can pin down the solution.
Solving the System
For the degree-3 odd polynomial ($n=1$), the original paper gives an analytical solution, while for the degree-5 odd polynomial ($n=2$), the original paper provides an iterative algorithm: first fix $x_1,x_2$ and solve for $a,b,c$, then fix $f_t(x)$'s $a,b,c$ and solve for $x_1,x_2$, and repeat — essentially a simplified version of the Remez algorithm.
However, the paper's iteration relies on a root-finding formula to compute $x_1,x_2$, which becomes cumbersome for larger $n$. So here I'll change the approach a bit: first parametrize $f_t'(x_i)$ using $x_1,x_2,\cdots,x_n$, i.e., define
\begin{equation}f_t'(x) = k(x^2-x_1^2)(x^2-x_2^2)\cdots (x^2-x_n^2)\end{equation}
Then we have $f_t(x) = \int_0^x f_t'(x) dx$, so that we express $f_t(x)$ in terms of $k$ and $x_1,x_2,\cdots,x_n$, and we only need to solve the system of equations
\begin{equation}f_t(l_t) = 1 - \mathcal{E}, \quad f_t(u_t) = 1 + (-1)^n \mathcal{E},\quad f_t(x_i) = 1 + (-1)^{i+1}\mathcal{E}\end{equation}
thereby avoiding solving equation $f_t'(x) = 0$. When $n=1$, we can solve to get
\begin{equation}x_1 = \sqrt{\frac{l_t^2+l_t u_t + u_t^2}{3}},\quad k = -\frac{6}{l_t^2 u_t + l_t u_t^2 + 2x_1^3}\end{equation}
When $n > 1$, we can just hand it to Mathematica; for instance, when $n=2$:
df[x_] = k*(x^2 - x1^2) (x^2 - x2^2);
f[x_] = Integrate[df[x], {x, 0, x}];
sol = NSolve[{f[l] == 1 - e, f[x1] == 1 + e, f[x2] == 1 - e,
f[u] == 1 + e, l < x1 < x2 < u, e > 0} /. {l -> 0.001,
u -> 1}, {k, x1, x2, e}, Reals]
f[x] /. sol
Finite Precision
At this point, it seems we've completed the solution of the original problem? In theory, yes — but only under infinite precision. In actual computation, precision is finite, and especially since the Muon optimizer uses bfloat16, the precision loss is more severe, which brings up some issues.
The first issue is that each $f_t^*$ is, in theory, only responsible for the interval $[l_t,u_t]$, but under finite precision, singular values may drift outside that interval. When $n$ is even (i.e., $f_t^*$ is a degree-5, 9, ... polynomial), there is a risk of divergence once values exceed $u_t$, because in that regime $f_t^*(x)$ is monotonically increasing toward positive infinity as $x > u_t$ grows, and any slip can cause divergence over iterations. There are two ways to address this: one is to leave a bit more margin for $[l_t,u_t]$ when solving for $f_t^*$, and the other is to keep the interval unchanged but divide the input by a number greater than 1 after obtaining $f_t^*$.
The original paper uses the latter approach, changing $f_t^*(x)$ to $f_t^*(x / 1.01)$. The number 1.01 is roughly the first representable number after 1 under bfloat16 precision (the exact value being 1.00781) — clearly this is meant to guard against numerical error pushing a singular value of 1 up into the next representable value. If computing at higher precision, this value could be reduced accordingly.
The second issue is more subtle, and we'll illustrate it with a concrete example. Let $n=2,l_1=0.001,u_1=1$; we can compute that $f_1^*$ is
\begin{equation}f_1^*(x) = 8.4703 x - 25.1081 x^3 + 18.6293 x^5\end{equation}
where $x_1 = 0.3674, x_2 = 0.8208, \mathcal{E}=0.9915$. What's wrong with this solution? By the equioscillation theorem, we know $f_1^*(x_2) = 1-\mathcal{E} = 0.0085$, meaning it maps $0.8208$ to $0.0085$. However, our ultimate goal is to turn every value in $(0,1]$ into 1, so $f_1^*$ maps a value already very close to the target, $0.8208$, to something far away from the target, $0.0085$. Although $f_2^*,f_3^*,\cdots$ downstream will theoretically pull it back gradually, repeatedly shrinking and then expanding a number under finite precision accumulates significant error.
Of course, by the equioscillation theorem we know this oscillatory behavior is unavoidable — the best we can do is hope that the maximum error $\mathcal{E}$ isn't too close to 1, thereby mitigating this accumulated error. It's not hard to see that the larger the interval $[l_t,u_t]$, the theoretically harder the fit becomes, and the closer the maximum error $\mathcal{E}$ gets to 1. So the paper introduces a hyperparameter $\lambda \in (0, 1)$, changing the optimization interval from $[l_t,u_t]$ to $[\max(l_t, \lambda u_t),u_t]$, limiting the interval size to ensure $\mathcal{E}$ doesn't get too large. (It should be noted that in the main text of the paper, the $\lambda$ used is $0.1$, but the appendix code actually uses $\lambda$ equal to $0.024$.)
But this raises another issue: doesn't the original $l_t$, especially the one we set at the very start, $l$, risk being neglected? To solve this problem, the paper introduces the "Recenter" trick: if the optimization interval is $[l_t,u_t]$, then it will satisfy $f_t^*(l_t) + f_t^*(u_t) = 2$, but after changing the optimization interval to $[\max(l_t, \lambda u_t),u_t]$ this is no longer guaranteed. So we multiply $f_t^*$ by $\gamma$ to make it satisfy this equation:
\begin{equation}\gamma f_t^*(l_t) + \gamma f_t^*(u_t) = 2\qquad \Rightarrow \qquad \gamma = \frac{2}{f_t^*(l_t) + f_t^*(u_t)}\end{equation}
This restores consideration of the original $l_t$.
Reference Code
Here is the complete Mathematica code for $n=2$:
df[x_] = k*(x^2 - x1^2) (x^2 - x2^2);
f[x_] = Integrate[df[x], {x, 0, x}];
sol[l_, u_] :=
NSolve[{f[l] == 1 - e, f[x1] == 1 + e, f[x2] == 1 - e, f[u] == 1 + e,
l < x1 < x2 < u, e > 0, k > 0}, {k, x1, x2, e}]
ff[x_, l_, u_] = f[x]*2/(f[l] + f[u]) // Expand;
lt = 0.001; ut = 1; lambda = 0.02407327424182761;
While[1 - lt > 0.0001,
fff[x_] = ff[x, lt, ut] /. sol[Max[lt, lambda*ut], ut][[1]];
Print[fff[x]];
lt = fff[lt]; ut = 2 - lt]
The results are as follows ($f_t(x) = a_t x + b_t x^3 + c_t x^5$):
$$\begin{array}{c|ccc} \hline t & a\times 1.01 & b\times 1.01^3 & c\times 1.01^5 \\ \hline \quad 1\quad & 8.28721 & -23.5959 & 17.3004 \\ 2 & 4.10706 & -2.94785 & 0.544843 \\ 3 & 3.94869 & -2.9089 & 0.551819 \\ 4 & 3.31842 & -2.48849 & 0.510049 \\ 5 & 2.30065 & -1.6689 & 0.418807 \\ 6 & 1.8913 & -1.268 & 0.376804 \\ 7 & 1.875 & -1.25 & 0.375 \\ 8 & 1.875 & -1.25 & 0.375 \\ \hline \end{array}$$
Note that what's given here is the result before applying the $f_t^*(x / 1.01)$ correction, so the actual $a, b, c$ still needs to be divided further, based on this table, by the $1,3,5$-th power of $1.01$. We don't give the result after dividing by $1.01$ directly, because the convergent value before dividing by $1.01$, i.e. $1.875, -1.25, 0.375$ ($t \geq 7$), is more concise and easier to observe and appreciate. (Exercise: prove that the final convergent value can be derived from $x_1=x_2=1$ and $f(1)=1$.)
The code from the paper's appendix, cleaned up, is as follows:
import numpy as np
def optimal_quintic(l, u):
assert 0 <= l <= u
if 1 - 5e-6 <= l / u:
# Above this threshold, the equoscillating polynomials
# is numerically equal to...
return (15 / 8) / u, (-10 / 8) / (u**3), (3 / 8) / (u**5)
# This initialization becomes exact as l -> u
q = (3 * l + 1) / 4
r = (l + 3) / 4
E, old_E = np.inf, None
while not old_E or abs(old_E - E) > 1e-15:
old_E = E
LHS = np.array([
[l, l**3, l**5, 1],
[q, q**3, q**5, -1],
[r, r**3, r**5, 1],
[u, u**3, u**5, -1],
])
a, b, c, E = np.linalg.solve(LHS, np.ones(4))
q, r = np.sqrt(
(-3 * b + np.array([-1, 1]) * np.sqrt(9 * b**2 - 20 * a * c)) /
(10 * c)
)
return float(a), float(b), float(c)
def optimal_composition(l, num_iters, cushion=0.02407327424182761):
u = 1
coefficients = []
for _ in range(num_iters):
a, b, c = optimal_quintic(max(l, cushion * u), u)
# Due to cushioning , this may be centered around 1 with
# respect to 0.024*u, u. Recenter it around 1 with respect
# to l, u, meaning find c so that 1 - c*p(l) = c*p(u) - 1:
pl = a * l + b * l**3 + c * l**5
pu = a * u + b * u**3 + c * u**5
rescalar = 2 / (pl + pu)
a *= rescalar
b *= rescalar
c *= rescalar
# Optionally incorporate safety factor here :
# a /= 1.01; b /= 1.01**3; c /= 1.01**5
coefficients.append((a, b, c))
l = a * l + b * l**3 + c * l**5
u = 2 - l
return coefficients
print(*optimal_composition(1e-3, 10), sep="\n")
Completing the Proof
In this final section, let's fill in the proof that "the greedy solution is exactly the global optimal solution."
By the equioscillation theorem, we know that the range of $f_t^*$ is $[l_{t+1},u_{t+1}]$, where $l_{t+1}=f_t^*(l_t),u_{t+1}=2-l_{t+1}$, from which we can see that the maximum error of the $T$-step greedy solution is $\mathcal{E}_T = 1 - l_{T+1} = 1 - f_T^*(l_T)$. We only need to show that the maximum error of the $T$-step global optimal solution can also only be reduced to $1 - f_T^*(l_T)$, and we obtain the conclusion that "the greedy solution is exactly the global optimal solution."
The proof proceeds by mathematical induction. Suppose the conclusion holds for $t=1,2,\cdots,T-1$; then $\hat{f} = f_{T-1}^*\circ \cdots \circ f_2^* \circ f_1^*$ is the global optimal solution at step $T-1$, with range $[l_T, u_T]$ and maximum error $\mathcal{E}_{T-1}=1-l_T=u_T-1$. On the other hand, let $\tilde{f} = \tilde{f}_{T-1}\circ \cdots \circ \tilde{f}_2 \circ \tilde{f}_1$ be an arbitrary $T-1$-step solution with range $[a,b]$, and set $c = \frac{2}{a+b}$; then the range of $c\tilde{f}$ is $[ca,cb]$, and clearly $ca\leq 1, cb\geq 1$. By the induction hypothesis, we have
\begin{equation}\begin{aligned} 1 - ca \geq \mathcal{E}_{T-1} \\ cb - 1 \geq \mathcal{E}_{T-1} \end{aligned}\qquad\Rightarrow\qquad \frac{a}{b} \leq \frac{1 - \mathcal{E}_{T-1}}{1 + \mathcal{E}_{T-1}} = \frac{l_T}{u_T} \end{equation}
That is, the relative size of the range of any arbitrary $T-1$-step solution is no smaller than the relative size of the range $[l_T, u_T]$ of the $T-1$-step optimal solution. It follows that
\begin{equation}\begin{aligned} \min_{f_T} \max_{x\in[l,u]} |f_T(\tilde{f}(x)) - 1| =&\, \min_{f_T} \max_{x\in[a,b]} |f_T(x) - 1| \\ =&\, \min_{f_T} \max_{x\in[a/b,1]} |f_T(x) - 1| \\ \geq &\, \min_{f_T} \max_{x\in[l_T/u_T,1]} |f_T(x) - 1| \\ =&\, \min_{f_T} \max_{x\in[l_T,u_T]} |f_T(x) - 1| \\ =&\, \mathcal{E}_T \end{aligned}\end{equation}
In other words, no matter which alternative $T-1$-step solution you pick, its maximum error can at best only match that of the greedy solution — so the maximum error of the greedy solution is already globally optimal, completing the inductive proof. The key step in the above derivation is
\begin{equation}\min_{f_T} \max_{x\in[a,b]} |f_T(x) - 1| = \min_{f_T} \max_{x\in[a/b,1]} |f_T(x) - 1|\end{equation}
This is because we can always set $g_T(y) = f_T(b y)$, and $g_T$ can still represent an arbitrary odd polynomial of the same degree, so $g_T$ and $f_T$ lie in the same function space, and hence the notation can be interchanged, i.e.,
\begin{equation}\min_{f_T}\max_{x\in[a,b]} |f_T(x) - 1| = \min_{g_T}\max_{y\in[a/b,1]} |g_T(y) - 1|= \min_{f_T}\max_{x\in[a/b,1]} |f_T(x) - 1|\end{equation}
Summary
This post introduced the latest progress in finding better Newton-Schulz iterations for the msign operator. By using the equioscillation theorem together with a greedy reformulation, one can directly obtain the theoretically optimal solution — the whole process is quite rigorous and well worth studying.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.