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

Polar Express：最优矩阵符号方法及其在 Muon 算法中的应用

Noah Amsel Thanks:  New York University. noah.amsel@nyu.edu    David Persson Thanks:  New York University and Flatiron Institute. dup210@nyu.edu , dpersson@flatironinstitute.org    Christopher Musco Thanks:  New York University. cmusco@nyu.edu    Robert M. Gower Thanks:  Flatiron Institute. rgower@flatironinstitute.org

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.

计算极分解以及与之相关的矩阵符号函数，是数值分析里研究了数十年的老问题。最近它又成了训练深度网络的 Muon 算法里的关键子程序。不过，这一应用提出的要求与经典场景差得很远：深度学习需要照顾 GPU 的算法，更看重吞吐，而不是极高的精度。我们提出 Polar Express，一种新的极分解算法。它和 Newton–Schulz 以及其他经典多项式方法一样，只用矩阵-矩阵乘法，因而在 GPU 上很高效。受 Chen & Chow 以及 Nakatsukasa & Freund 较早工作的启发，Polar Express 在每一步都通过求解一个 minimax 优化问题来调整更新规则。我们证明这一策略能在最坏情形下把误差压到最小，从而让 Polar Express 无论在前期迭代还是在渐近阶段都能尽快收敛。我们还处理了有限精度下的问题，使其能切实用于 bfloat16。把该方法嵌进 Muon 训练框架后，在 FineWeb 数据集的 10 亿 token 上训练 GPT-2 时，验证损失得到了稳定的改善，并在一系列学习率下都优于近期的替代方法。

1  Introduction

1 引言

Contents

Advanced linear algebra is making its way into deep learning. Efficient algorithms for computing matrix functions have found exciting new applications in training neural networks. In particular, approximations to the matrix-inverse are used in the full Adagrad method [13], the matrix square-root and quarter-root appear as subroutines in the Shampoo and Soap optimizers [17, 43, 47], and most recently, the matrix sign function has become a key ingredient of the Muon optimizer [5, 4, 22].

目录

高等线性代数正在进入深度学习。计算矩阵函数的高效算法，已经在神经网络训练里找到了令人兴奋的新用途。具体来说，矩阵求逆的近似出现在完整的 Adagrad 方法中 [13]，矩阵平方根和四分之一次根是 Shampoo 与 Soap 优化器的子程序 [17, 43, 47]，而最近，矩阵符号函数又成了 Muon 优化器的关键构件 [5, 4, 22]。

While the problem of computing these matrix functions has been studied by numerical analysts for decades, applications in deep learning come with different requirements than those in computational science. For deep learning, it is critical to take maximum advantage of GPU-friendly operations like matrix-matrix products and to avoid less parallel operations. Moreover, memory overhead must be small to handle large models. On the other hand, high accuracy is typically less important; the gold standard of sixteen digits of accuracy is overkill in deep learning.

尽管计算这些矩阵函数的问题已被数值分析学家研究了数十年，深度学习应用提出的要求与计算科学并不相同。对深度学习来说，关键是尽量用满 GPU 友好的操作（例如矩阵-矩阵乘法），并避开并行度较差的运算。内存开销也必须很小，才能撑住大模型。另一方面，高精度通常并不那么重要；科学计算里追求的十六位有效数字，对深度学习来说过头了。

Given these considerations, there is a need to develop new matrix function methods that are tailor-made for deep learning applications. We take on this challenge by designing a state-of-the-art, GPU-friendly algorithm for computing the matrix sign function, or more generally, for computing the polar decomposition of a rectangular matrix. We apply our new Polar Express method (Algorithm 1) to compute the descent direction in the increasingly popular Muon optimizer. In Figure 1, we show that using Polar Express within Muon consistently results in lower validation loss across all learning rates when training a GPT-2 model, as compared to other matrix sign methods [9, 44, 22].

基于这些考虑，有必要专门为深度学习设计新的矩阵函数方法。我们接受这一挑战，设计了一种面向 GPU、用于计算矩阵符号函数——更一般地说，用于计算矩形矩阵极分解——的先进算法。我们把新的 Polar Express 方法（算法 1）用到日益流行的 Muon 优化器里，去计算下降方向。图 1 表明，在训练 GPT-2 时，把 Polar Express 放进 Muon，在所有学习率下都能得到更低的验证损失，优于其他矩阵符号方法 [9, 44, 22]。

from itertools import repeat
import torch

coeffs_list = [
(8.28721201814563, -23.595886519098837, 17.300387312530933),
(4.107059111542203, -2.9478499167379106, 0.5448431082926601),
(3.9486908534822946, -2.908902115962949, 0.5518191394370137),
(3.3184196573706015, -2.488488024314874, 0.51004894012372),
(2.300652019954817, -1.6689039845747493, 0.4188073119525673),
(1.891301407787398, -1.2679958271945868, 0.37680408948524835),
(1.8750014808534479, -1.2500016453999487, 0.3750001645474248),
(1.875, -1.25, 0.375), # subsequent coeffs equal this numerically
]

safety factor for numerical stability (but exclude last polynomial)

coeffs_list = [(a / 1.01, b / 1.013, c / 1.015)
for (a, b, c) in coeffs_list[:-1]] + [coeffs_list[-1]]

@torch.compile
def PolarExpress(G: torch.Tensor, steps: int) -> torch.Tensor:
assert G.ndim >= 2
X = G.bfloat16() # for speed
if G.size(-2) > G.size(-1): X = X.mT # this reduces FLOPs
X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01)
hs = coeffs_list[:steps] + list(
repeat(coeffs_list[-1], steps - len(coeffs_list)))
for a, b, c in hs:
A = X @ X.mT
B = b * A + c * A @ A
X = a * X + B @ X # X <- aX + bX^3 + cX^5
if G.size(-2) > G.size(-1): X = X.mT
return X

Algorithm 1 Python code for the Polar Express of degree = 5.

from itertools import repeat
import torch

coeffs_list = [
(8.28721201814563, -23.595886519098837, 17.300387312530933),
(4.107059111542203, -2.9478499167379106, 0.5448431082926601),
(3.9486908534822946, -2.908902115962949, 0.5518191394370137),
(3.3184196573706015, -2.488488024314874, 0.51004894012372),
(2.300652019954817, -1.6689039845747493, 0.4188073119525673),
(1.891301407787398, -1.2679958271945868, 0.37680408948524835),
(1.8750014808534479, -1.2500016453999487, 0.3750001645474248),
(1.875, -1.25, 0.375), # subsequent coeffs equal this numerically
]

safety factor for numerical stability (but exclude last polynomial)

coeffs_list = [(a / 1.01, b / 1.013, c / 1.015)
for (a, b, c) in coeffs_list[:-1]] + [coeffs_list[-1]]

@torch.compile
def PolarExpress(G: torch.Tensor, steps: int) -> torch.Tensor:
assert G.ndim >= 2
X = G.bfloat16() # for speed
if G.size(-2) > G.size(-1): X = X.mT # this reduces FLOPs
X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01)
hs = coeffs_list[:steps] + list(
repeat(coeffs_list[-1], steps - len(coeffs_list)))
for a, b, c in hs:
A = X @ X.mT
B = b * A + c * A @ A
X = a * X + B @ X # X <- aX + bX^3 + cX^5
if G.size(-2) > G.size(-1): X = X.mT
return X

算法 1 次数 $d=5$ 的 Polar Express 的 Python 代码。

Figure 1: Training a GPT-2-Large model (774M params) on 1 billion tokens from the FineWeb dataset [2]. The label muon- refers to implementing Muon using 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$.

图 1：在 FineWeb 数据集 [2] 的 10 亿 token 上训练 GPT-2-Large 模型（774M 参数）。标签 muon- 表示用 来计算极因子，以实现 Muon。左：各学习率下的最终验证损失。右：采用各自最优学习率时，验证损失随 epoch 的变化。各方法的最优学习率（$lr$）和最终验证损失分别为：adamw $(lr=0.0001)$：$4.172$，muon-You $(lr=0.02)$：$3.400$，muon-Jordan $(lr=0.02)$：$3.398$，muon-PolarExp $(lr=0.02)$：$3.340$。

1.1  The Muon Method

1.1 Muon 方法

The Muon optimizer has recently gained popularity for training large language models, often outperforming state-of-the-art adaptive gradient methods like Adam and AdamW [26, 31]. Muon has been used to set records for the NanoGPT speedrun [22] and to expand the Pareto frontier of performance versus training FLOPs for large language models [30, 42].

Muon 优化器最近在训练大语言模型时颇受欢迎，常常超过 Adam 和 AdamW 这类先进的自适应梯度方法 [26, 31]。Muon 曾用来刷新 NanoGPT speedrun 纪录 [22]，也把大语言模型在性能与训练 FLOPs 之间的帕累托前沿往外推了一截 [30, 42]。

The Muon update rule [5] is defined as follows. Let $\lambda,\beta>0$ be the learning rate and momentum coefficient hyperparameters. (By default, $\beta=0.9$.) Let $\bm{W}_{t}\in\mathbb{R}^{m\times n}$ be the weight matrix of a given neural network layer at iteration $t$, and let $\bm{G}_{t}\in\mathbb{R}^{m\times n}$ be its (stochastic) gradient. Let $\bm{M}_{t}\in\mathbb{R}^{m\times n}$ be the running momentum estimate of the gradient, where $\bm{M}_{0}=\bm{0}$. The Muon update is given by

Muon 的更新规则 [5] 如下。设 $\lambda,\beta>0$ 为学习率和动量系数超参数（默认 $\beta=0.9$）。设 $\bm{W}_{t}\in\mathbb{R}^{m\times n}$ 为某层网络在第 $t$ 步的权重矩阵，$\bm{G}_{t}\in\mathbb{R}^{m\times n}$ 为其（随机）梯度。设 $\bm{M}_{t}\in\mathbb{R}^{m\times n}$ 为梯度的滑动动量估计，且 $\bm{M}_{0}=\bm{0}$。Muon 更新由下式给出

$$\displaystyle\bm{M}_{t} \\ \displaystyle=\beta\bm{M}_{t-1}+(1-\beta)\bm{G}_{t} \\ \displaystyle\bm{W}_{t+1} \\ \displaystyle=\bm{W}_{t}-\lambda\polar(\bm{M}_{t}).$$

Whereas standard stochastic gradient descent (SGD) with momentum updates the weight matrix by taking a step in the direction $-\bm{M}_{t}$, the Muon method steps in the direction $-\polar(\bm{M}_{t})$, where $\polar(\bm{M})$ denotes the closest semi-orthogonal matrix to $\bm{M}$ [19, Chapter 8]. Concretely, if $\bm{M}=\bm{U}\bm{\Sigma}\bm{V}^{\mathsf{T}}$ is the singular value decomposition (SVD) of $\bm{M}$, then

标准的带动量随机梯度下降（SGD）沿 $-\bm{M}_{t}$ 方向更新权重；Muon 则沿 $-\polar(\bm{M}_{t})$ 迈步，其中 $\polar(\bm{M})$ 表示离 $\bm{M}$ 最近的半正交矩阵 [19, 第 8 章]。具体来说，若 $\bm{M}=\bm{U}\bm{\Sigma}\bm{V}^{\mathsf{T}}$ 是 $\bm{M}$ 的奇异值分解（SVD），则

$$\polar(\bm{M}):=\bm{U}\bm{V}^{\mathsf{T}}. \tag{1}$$

The matrix $\polar(\bm{M})$ can be seen as a generalization of the matrix sign function to rectangular matrices [3]. Indeed, when $\bm{M}$ is square symmetric with eigendecomposition $\bm{M}=\bm{V}\bm{\Lambda}\bm{V}^{\mathsf{T}}$, $\polar(\bm{M})$ exactly coincides with the matrix sign function $\sign(\bm{M})=\bm{V}\sign(\bm{\Lambda})\bm{V}^{\mathsf{T}}$ [19, Chapter 5]. Equivalently, $\polar(\bm{M})$ is the left orthogonal factor of the polar decomposition of $\bm{M}$ [19, Chapter 8]. The motivation for Muon is that $-\polar(\bm{M})$ gives the steepest-descent direction with respect to the spectral norm (instead of the Frobenius norm, as in standard SGD).

矩阵 $\polar(\bm{M})$ 可以看成矩阵符号函数向矩形矩阵的推广 [3]。确实，当 $\bm{M}$ 为方阵且对称、特征分解为 $\bm{M}=\bm{V}\bm{\Lambda}\bm{V}^{\mathsf{T}}$ 时，$\polar(\bm{M})$ 恰好等于矩阵符号函数 $\sign(\bm{M})=\bm{V}\sign(\bm{\Lambda})\bm{V}^{\mathsf{T}}$ [19, 第 5 章]。等价地说，$\polar(\bm{M})$ 就是 $\bm{M}$ 的极分解中左侧的正交因子 [19, 第 8 章]。Muon 的动机是：$-\polar(\bm{M})$ 给出的是关于谱范数的最速下降方向（而标准 SGD 用的是 Frobenius 范数）。

Recent work [40] shows that Muon can be viewed as a conditional gradient (Frank-Wolfe) method with a trust region defined by the spectral norm. In the same work, the authors also provide a convergence theory for the smooth and non-convex setting, as well as for the stochastic non-convex case. The analysis of Muon was further refined in [41], which proves convergence under a layerwise $(L_{0},L_{1})$-smoothness assumption, in both the stochastic non-convex and stochastic Polyak–Łojasiewicz settings. We refer the reader to [22] and [5] for further background. In this paper, we take the Muon update rule as given and focus on the problem of efficiently computing the polar decomposition $\polar(\bm{M})$.

近期工作 [40] 表明，Muon 可以看成一种条件梯度（Frank-Wolfe）方法，其信赖域由谱范数划定。同一工作还给出了光滑非凸情形以及随机非凸情形下的收敛理论。[41] 进一步细化了分析，在逐层 $(L_{0},L_{1})$-光滑假设下，证明了随机非凸和随机 Polyak–Łojasiewicz 两种设定中的收敛性。更多背景见 [22] 和 [5]。本文把 Muon 的更新规则当作给定，集中讨论如何高效计算极分解 $\polar(\bm{M})$。

1.2  Computing the Polar Factor

1.2 计算极因子

Although $\polar(\bm{M})$ can be computed directly via an SVD in $O(\min(mn^{2},nm^{2}))$ time, doing so is prohibitively expensive in deep learning applications, especially as standard SVD algorithms fail to take full advantage of the parallelism available on GPUs. There has been significant work on highly-parallel methods for the SVD, but the most common approaches actually require computing the matrix-sign function as a subroutine [33, 35]. Numerical analysts have spent decades developing iterative methods for computing $\polar(\bm{M})$. This rich line of work includes Newton–Schulz [19, Chapter 8], Padé iteration [25, 18], the Newton and scaled Newton iterations [19, Chapter 8], the QWHD iteration [32, 35], and Zolo-pd (Zolotarev polar decomposition) [33]. Unfortunately, as discussed in Section 2, most of these methods are based on rational approximations to the function $\sign(x)$ and require computing matrix inverses or QR decompositions. Such methods are ill-suited to GPU acceleration and deep learning applications. In contrast, the older Newton-Schulz method is based on polynomial approximation of $\sign(x)$ and uses only matrix-matrix products. Thus, Muon initially used Newton-Schulz [4]. Indeed, Muon stands for “MomentUm Orthogonalized by Newton-Schulz” [22].

虽然可以直接用 SVD 在 $O(\min(mn^{2},nm^{2}))$ 时间内算出 $\polar(\bm{M})$，这在深度学习里贵得难以承受，尤其是标准 SVD 算法没法充分利用 GPU 上的并行能力。高度并行的 SVD 方法也有不少研究，但最常见的路线其实还是把计算矩阵符号函数当作子程序 [33, 35]。数值分析学家花了数十年发展计算 $\polar(\bm{M})$ 的迭代法。这条丰富的脉络包括 Newton–Schulz [19, 第 8 章]、Padé 迭代 [25, 18]、Newton 迭代及其缩放版本 [19, 第 8 章]、QWHD 迭代 [32, 35]，以及 Zolo-pd（Zolotarev 极分解）[33]。可惜，如第 2 节所讨论的，这些方法大多基于对函数 $\sign(x)$ 的有理逼近，需要计算矩阵逆或 QR 分解。这类方法既不适合 GPU 加速，也不适合深度学习。相比之下，更早的 Newton-Schulz 方法基于对 $\sign(x)$ 的多项式逼近，只用矩阵-矩阵乘法。因此 Muon 起初采用的就是 Newton-Schulz [4]。Muon 这个名字正是 “MomentUm Orthogonalized by Newton-Schulz” 的缩写 [22]。

The Newton-Schulz methods.

Newton-Schulz constructs a sequence of approximations $\bm{X}_{t}\approx\polar(\bm{M})$ as follows:

Newton-Schulz 方法。

Newton-Schulz 按如下方式构造一列近似 $\bm{X}_{t}\approx\polar(\bm{M})$：

$$\displaystyle\bm{X}_{0}=\bm{M}/\|\bm{M}\|_{\text{F}} \\ \displaystyle\bm{X}_{t+1}=\frac{3}{2}\bm{X}_{t}-\frac{1}{2}\bm{X}_{t}\bm{X}_{t}^{\top}\bm{X}_{t} \tag{2}$$

At each iteration, this rule effectively applies the cubic polynomial $p(x)=\frac{3}{2}x-\frac{1}{2}x^{3}$ to each singular value of $\bm{X}_{t}$. The scalar fixed-point iteration $x_{t+1}=p(x_{t})$ converges to $\sign(x_{0})$ as $t\to\infty$, provided $|x_{0}|\leq 1$. As a result, the matrix iteration satisfies $\lim\limits_{t\to\infty}\bm{X}_{t}=\bm{U}\bm{V}^{\top}=\polar(\bm{X}_{0})$. Higher-degree versions of Newton-Schulz follow the same principle. For example, the degree-5 polynomial $p(x)=(15x-10x^{3}+3x^{5})/8$ converges even faster. The Newton-Schulz iterations converge super-exponentially when $\bm{X}_{t}$ is sufficiently close to $\polar(\bm{M})$, but they suffer from slow initial convergence; when $\bm{X}_{0}$ is far from $\polar(\bm{M})$, the approximation improves slowly over the first few iterations.

每一步迭代，这条规则实际上是把三次多项式 $p(x)=\frac{3}{2}x-\frac{1}{2}x^{3}$ 作用到 $\bm{X}_{t}$ 的每一个奇异值上。标量不动点迭代 $x_{t+1}=p(x_{t})$ 在 $|x_{0}|\leq 1$ 时会收敛到 $\sign(x_{0})$（当 $t\to\infty$）。因此矩阵迭代满足 $\lim\limits_{t\to\infty}\bm{X}_{t}=\bm{U}\bm{V}^{\top}=\polar(\bm{X}_{0})$。更高次数的 Newton-Schulz 遵循同一原理。例如五次多项式 $p(x)=(15x-10x^{3}+3x^{5})/8$ 收敛得更快。当 $\bm{X}_{t}$ 已经足够接近 $\polar(\bm{M})$ 时，Newton-Schulz 迭代超指数收敛；但起步很慢：当 $\bm{X}_{0}$ 离 $\polar(\bm{M})$ 还远时，前几步近似改善得很有限。

The Jordan and You methods.

In Muon, high accuracy approximations to $\polar(\bm{M})$ are usually not necessary. The primary goal is instead to compute a coarse approximation in as few iterations as possible. To accelerate convergence in the low-accuracy regime, Jordan recently proposed a fixed-point iteration based on the polynomial $p(x)=3.4445x-4.7750x^{3}+2.0315x^{5}$ [22], which was found using a heuristic numerical search. Unlike Newton-Schulz, the scheme that Jordan proposed does not converge to $\polar(\bm{M})$. Instead, it plateaus at an error of $\approx 0.3$. However, it reaches this level of accuracy rapidly. As a result, when the number of iterations is smaller than about $10$, Jordan’s method outperforms the Newton-Schulz iteration. Building on this idea, You [9] proposed a method that applies six different polynomial updates in succession, which were again found by heuristic search. This method achieves better accuracy than Jordan’s but still fails to converge.

Jordan 方法与 You 方法。

在 Muon 里，通常并不需要高精度地逼近 $\polar(\bm{M})$。首要目标是用尽量少的迭代算出一个粗糙近似。为了加快低精度区间里的收敛，Jordan 最近提出了一种基于多项式 $p(x)=3.4445x-4.7750x^{3}+2.0315x^{5}$ 的不动点迭代 [22]，系数来自启发式数值搜索。与 Newton-Schulz 不同，Jordan 提出的格式并不会收敛到 $\polar(\bm{M})$，而是在误差约 $0.3$ 处停下。不过它达到这一精度的速度很快。因此当迭代次数大约小于 $10$ 时，Jordan 的方法会超过 Newton-Schulz 迭代。顺着这一思路，You [9] 提出连续使用六种不同的多项式更新，同样由启发式搜索得到。该方法精度优于 Jordan，但仍然不会真正收敛。

We introduce a new method. In particular, we derive polynomial update rules that are optimal at every iteration, outperforming all previous polynomial methods in our setting.

我们提出一种新方法。具体来说，我们导出在每一步都最优的多项式更新规则，在我们所关心的设定下超过此前所有多项式方法。

1.3  Contributions

1.3 贡献

We present Polar Express, an iterative method for approximating $\polar(\bm{M})$. Our method dynamically adapts the polynomial update rule at each iteration, prioritizing rapid progress in the initial stage and high accuracy in the later stage. Polar Express constructs polynomials $p_{1},\ldots,p_{T}$ so that the resulting composition is the optimal approximation to the sign function with respect to the supremum ($L^{\infty}$) norm (Theorem 4.1). By iteratively applying these polynomials to $\bm{M}$, Polar Express computes an approximation to $\polar(\bm{M})$ that is optimal in the worst-case at every iteration. Our method converges to $\polar(\bm{M})$ super-exponentially (Theorem 4.3), and it quickly reaches a good approximation within just five to ten iterations. This early-stage acceleration is especially valuable in deep learning applications, where runtime efficiency takes precedence over high accuracy. In contrast, classical methods like Newton-Schulz suffer from a slow initial convergence, while recent heuristic proposals [22, 9] fail to converge. Our method is efficient to run on GPUs, using only a few matrix-matrix products per iteration.[^1]

我们提出 Polar Express，一种逼近 $\polar(\bm{M})$ 的迭代法。该方法会动态调整每一步的多项式更新规则，前期优先快速推进，后期再追求高精度。Polar Express 构造多项式 $p_{1},\ldots,p_{T}$，使得所得复合是关于上确界（$L^{\infty}$）范数对符号函数的最优逼近（定理 4.1）。把这些多项式迭代作用到 $\bm{M}$ 上，Polar Express 就能在每一步都得到最坏情形下最优的 $\polar(\bm{M})$ 近似。我们的方法超指数收敛到 $\polar(\bm{M})$（定理 4.3），并且只需五到十步就能得到不错的近似。这种前期加速对深度学习尤其有价值，因为那里更看重运行效率而不是高精度。相比之下，Newton-Schulz 这类经典方法起步很慢，而近期的启发式方案 [22, 9] 又无法真正收敛。我们的方法在 GPU 上很高效，每步只用几次矩阵-矩阵乘法。[^1]

We give an explicit instantiation of Polar Express in Algorithm 1, which incorporates minor modifications to make it compatible with half-precision arithmetic (see Section 4.4). Algorithm 1 can be used as a drop-in replacement for previous methods. In numerical experiments, Polar Express outperforms previous methods on synthetic matrices and gradient matrices from a GPT-2 transformer (Figure 4). We demonstrate the effectiveness of using Polar Express within the Muon optimizer in Figure 1, showing that it consistently improves the training of GPT-2 language models on 1 billion tokens of the FineWeb dataset [2].

我们在算法 1 中给出 Polar Express 的一个显式实例，其中加入了少许改动，使之能兼容半精度算术（见第 4.4 节）。算法 1 可以当作此前方法的即插即用替换。数值实验中，Polar Express 在合成矩阵以及来自 GPT-2 Transformer 的梯度矩阵上都超过了以往方法（图 4）。图 1 展示了把 Polar Express 放进 Muon 优化器后的效果：在 FineWeb 数据集 [2] 的 10 亿 token 上训练 GPT-2 语言模型时，它稳定地改善了训练表现。

Notation.

We let $\|\bm{M}\|_{\text{F}}$ and $\|\bm{M}\|_{2}$ denote the Frobenius norm and spectral norm (largest singular value) of a matrix $\bm{M}$, respectively. We denote the spectrum (set of singular values) by $\sigma(\bm{M})$.

记号。

我们分别用 $\|\bm{M}\|_{\text{F}}$ 和 $\|\bm{M}\|_{2}$ 表示矩阵 $\bm{M}$ 的 Frobenius 范数和谱范数（最大奇异值）。用 $\sigma(\bm{M})$ 表示谱（奇异值集合）。

Let $\mathbb{P}_{d}$ be the set of polynomials of degree at most $d$. For odd $d$, $\mathbb{P}_{d}^{\odd}$ denotes the set of polynomials of degree at most $d$ containing only odd-degree monomials. For a polynomial $p$, $\deg(p)$ is its degree. Let $\sign(x)$ be the scalar sign function, which satisfies $\sign(0)=0$, $\sign(x)=1$ if $x>0$ and $\sign(x)=-1$ if $x<0$.

设 $\mathbb{P}_{d}$ 为次数至多为 $d$ 的多项式集合。对奇数 $d$，$\mathbb{P}_{d}^{\odd}$ 表示次数至多为 $d$、且只含奇次单项式的多项式集合。对多项式 $p$，$\deg(p)$ 为其次数。设 $\sign(x)$ 为标量符号函数，满足 $\sign(0)=0$，$x>0$ 时 $\sign(x)=1$，$x<0$ 时 $\sign(x)=-1$。

For a polynomial $p\in\mathbb{P}_{d}^{\odd}$ and a matrix $\bm{M}$ with rank reduced SVD given by $\bm{M}=\bm{U}\bm{\Sigma}\bm{V}^{\mathsf{T}}$ and positive singular values $\sigma_{1}\geq\cdots\geq\sigma_{\rank(\bm{M})}>0$, we define $p(\bm{M}):=\bm{U}p(\bm{\Sigma})\bm{V}^{\mathsf{T}}$, where $p(\bm{\Sigma})$ is the diagonal matrix with diagonal entries $p(\sigma_{i})$ for $i=1,\ldots,\rank(\bm{M})$.

对多项式 $p\in\mathbb{P}_{d}^{\odd}$，以及具有降秩 SVD $\bm{M}=\bm{U}\bm{\Sigma}\bm{V}^{\mathsf{T}}$、正奇异值为 $\sigma_{1}\geq\cdots\geq\sigma_{\rank(\bm{M})}>0$ 的矩阵 $\bm{M}$，我们定义 $p(\bm{M}):=\bm{U}p(\bm{\Sigma})\bm{V}^{\mathsf{T}}$，其中 $p(\bm{\Sigma})$ 是对角元为 $p(\sigma_{i})$（$i=1,\ldots,\rank(\bm{M})$）的对角矩阵。

2  Related Work

2 相关工作

Computing $\polar(\bm{M})$ is an important and longstanding problem in numerical linear algebra, with applications spanning electronic structure calculations, lattice quantum chromodynamics, orthogonal Procrustes analysis, parallel algorithms for computing the SVD, and beyond; see e.g. [18, 23, 8, 16, 36, 45].

计算 $\polar(\bm{M})$ 是数值线性代数里一个重要而悠久的问题，应用涵盖电子结构计算、格点量子色动力学、正交 Procrustes 分析、并行 SVD 算法等等；参见 [18, 23, 8, 16, 36, 45]。

Newton-Schulz and polynomial Padé methods.

The earliest methods in the literature are polynomial iterations like (2). Several nearly simultaneous papers introduced the family of polynomial Padé iterations, comprising Newton-Schulz and its higher-degree analogues [27, 6, 18, 29]. These higher-degree methods are also sometimes called “Newton-Schulz”; when doing so, we will specify the degree for clarity. In these methods, each iteration refines the current approximation $\bm{X}_{t}$ by applying a low-degree odd matrix polynomial, where any odd monomial $x\mapsto x^{2q+1}$ is defined for rectangular matrices by the formula $\bm{X}_{t}\mapsto\bm{X}_{t}\left(\bm{X}_{t}^{\top}\bm{X}_{t}\right)^{q}$. Our Polar Express method also takes this form, though unlike Newton-Schulz, it changes the polynomial at each iteration.

Newton-Schulz 与多项式 Padé 方法。

文献中最早的方法是形如 (2) 的多项式迭代。几篇几乎同时出现的论文引入了多项式 Padé 迭代族，其中包括 Newton-Schulz 及其更高次数的同类方法 [27, 6, 18, 29]。这些更高次数的方法有时也被称作 “Newton-Schulz”；这样称呼时我们会标明次数以免混淆。在这些方法里，每一步都用一个低次奇次矩阵多项式来改进当前近似 $\bm{X}_{t}$，其中任意奇次单项式 $x\mapsto x^{2q+1}$ 对矩形矩阵由公式 $\bm{X}_{t}\mapsto\bm{X}_{t}\left(\bm{X}_{t}^{\top}\bm{X}_{t}\right)^{q}$ 定义。我们的 Polar Express 也取这一形式，不过与 Newton-Schulz 不同，它在每一步都会更换多项式。

The polynomials used in Padé methods are chosen to match the value and first few derivatives of $\sign(x)$ at the points $x=\pm 1$. For instance, the update rule of the third method in this family is defined by $p(x)=\frac{1}{16}\left(35x-35x^{3}+21x^{5}-5x^{7}\right)$, which is the unique degree-7 polynomial satisfying $p(\pm 1)=\pm 1$ and $p^{\prime}(\pm 1)=p^{\prime\prime}(\pm 1)=p^{\prime\prime\prime}(\pm 1)=0$. These methods converge so long as all singular values of $\bm{X}_{0}$ lie in $(0,1]$, a condition guaranteed by the initialization of (2). Furthermore, the order of convergence of the degree $2q+1$ method is $q+1$ [6]. In particular, the Newton-Schulz method ($q=1$) converges quadratically.

Padé 方法所用的多项式，是为了在 $x=\pm 1$ 处匹配 $\sign(x)$ 的函数值及其前若干阶导数。例如该族中第三种方法的更新规则由 $p(x)=\frac{1}{16}\left(35x-35x^{3}+21x^{5}-5x^{7}\right)$ 给出，这是满足 $p(\pm 1)=\pm 1$ 且 $p^{\prime}(\pm 1)=p^{\prime\prime}(\pm 1)=p^{\prime\prime\prime}(\pm 1)=0$ 的唯一七次多项式。只要 $\bm{X}_{0}$ 的全部奇异值落在 $(0,1]$ 内——(2) 的初始化就保证了这一点——这些方法就会收敛。此外，次数为 $2q+1$ 的方法收敛阶为 $q+1$ [6]。特别地，Newton-Schulz 方法（$q=1$）二次收敛。

Newton’s method and rational Padé.

In the numerical analysis literature, polynomial methods were succeeded by rational iterations like Newton’s method [18], defined as follows[^2]:

Newton 方法与有理 Padé。

在数值分析文献中，多项式方法之后出现的是像 Newton 方法这样的有理迭代 [18]，定义如下[^2]：

$$\displaystyle\bm{X}_{0}=\bm{M} \\ \displaystyle\bm{X}_{t+1}=\frac{1}{2}\left(\bm{X}_{t}+\bm{X}_{t}^{-\top}\right) \tag{3}$$

Newton’s method also converges quadratically. Like Newton-Schulz, it works because the rational function $r(x)=\frac{1}{2}(x+x^{-1})$ has a stable fixed point at $1$; unlike for Newton-Schulz, this point is a global attractor for the whole positive real line. At first glance, Newton’s method has nothing to do with the Padé iterations discussed above. However, after a change of variables $\bm{Y}_{t}=\bm{X}_{t}^{-1}$, it can be reinterpreted as $\bm{Y}_{t+1}=2\bm{Y}_{t}(\bm{I}+\bm{Y}_{t}^{\top}\bm{Y}_{t})^{-1}$, which is sometimes called inverse Newton. Observing that $r(x)=\frac{2x}{1+x^{2}}$ satisfies $r(\pm 1)=\pm 1$ and $r^{\prime}(\pm 1)=0$, we see that (inverse) Newton is also a Padé method, though a rational rather than polynomial one. In fact, given a odd degree $2q_{n}+1$ for the numerator and an even degree $2q_{d}$ for the denominator, there is a unique rational function that matches the value and first $q_{n}+q_{d}$ derivatives of $\sign(x)$ at $x=\pm 1$. This directly yields a Padé method for computing $\polar(\bm{M})$ whose order of convergence is $q_{n}+q_{d}+1$. For instance, $r(x)=\frac{3x+x^{3}}{1+3x^{2}}$ is called Halley’s method, which converges cubically. When $q_{d}=0$, we recover the polynomial Padé methods.

Newton 方法同样二次收敛。和 Newton-Schulz 一样，它之所以有效，是因为有理函数 $r(x)=\frac{1}{2}(x+x^{-1})$ 在 $1$ 处有一个稳定不动点；与 Newton-Schulz 不同的是，这个点是整条正实轴的全局吸引子。乍看之下，Newton 方法与上面讨论的 Padé 迭代毫无关系。不过做变量替换 $\bm{Y}_{t}=\bm{X}_{t}^{-1}$ 之后，它可以改写成 $\bm{Y}_{t+1}=2\bm{Y}_{t}(\bm{I}+\bm{Y}_{t}^{\top}\bm{Y}_{t})^{-1}$，有时称作逆 Newton。注意到 $r(x)=\frac{2x}{1+x^{2}}$ 满足 $r(\pm 1)=\pm 1$ 且 $r^{\prime}(\pm 1)=0$，可见（逆）Newton 其实也是一种 Padé 方法，只不过是有理的而非多项式的。事实上，给定分子的奇次数 $2q_{n}+1$ 和分母的偶次数 $2q_{d}$，存在唯一的有理函数能在 $x=\pm 1$ 处匹配 $\sign(x)$ 的函数值及其前 $q_{n}+q_{d}$ 阶导数。这直接给出一种计算 $\polar(\bm{M})$ 的 Padé 方法，收敛阶为 $q_{n}+q_{d}+1$。例如 $r(x)=\frac{3x+x^{3}}{1+3x^{2}}$ 称为 Halley 方法，三次收敛。当 $q_{d}=0$ 时，就回到多项式 Padé 方法。

There are two main weakness of Newton’s method and the Padé iterations: slow convergence in the initial phase and the need to compute explicit inverses. To accelerate initial convergence, Higham popularized the technique of rescaling the matrix after every Newton iteration [18]. Intuitively, rescaling $\bm{X}_{t}$ so that $\sigma_{\max}=1/\sigma_{\min}$ centers the spectrum around $1$, where convergence is fastest. Several easily-computable choices of scaling factor exist to accomplish this approximately. Note that this rescaling scheme would fail for Newton-Schulz, which likewise suffers from slow initial convergence but which would diverge if $\sigma_{\max}\gg 1$.

Newton 方法与 Padé 迭代有两个主要弱点：初期收敛慢，以及需要显式求逆。为加快初期收敛，Higham 推广了每步 Newton 迭代后重新缩放矩阵的技巧 [18]。直观上，把 $\bm{X}_{t}$ 缩放到 $\sigma_{\max}=1/\sigma_{\min}$，就能把谱居中到 $1$ 附近，而那里收敛最快。有若干易于计算的缩放因子可以近似做到这一点。需要注意，这种缩放方案对 Newton-Schulz 行不通：后者同样受初期收敛慢的困扰，但如果 $\sigma_{\max}\gg 1$ 它还会发散。

Computing matrix inverses is difficult to parallelize and to implement stably in low precision arithmetic. However, a trick was developed for stably computing many rational methods without explicit inverses; QR decompositions can be used instead [32, 49]. Applying this trick to Halley’s method and combining with a special rescaling scheme yields the QDWH (QR-based dynamically weighted Halley) method, which converges in just six iterations for any reasonably conditioned matrix [32].

计算矩阵逆既难并行，也难在低精度算术里稳定实现。不过人们找到了一个窍门，可以在不显式求逆的情况下稳定地计算许多有理方法：改用 QR 分解 [32, 49]。把这一窍门用到 Halley 方法上，再配上一种特殊的缩放方案，就得到 QDWH（基于 QR、动态加权的 Halley）方法，对任何条件数不太离谱的矩阵都能在六步内收敛 [32]。

Adaptive rational methods from optimal approximations.

A landmark 2016 paper introduced a new paradigm to design iterative methods for computing $\polar(\bm{M})$ [33]. We describe this paradigm in more detail in Section 4, but the main insight is as follows. Padé methods choose the update rule to be an approximation to $\sign(x)$ of a given degree that is optimally accurate in the neighborhood of $x=1$. Instead, we should choose the approximation to $\sign(x)$ that is optimal over an interval $[\ell,1]\subset\mathbb{R}_{\geq 0}$ that contains the singular values. Moreover, after each step of the algorithm, the range of the singular values changes; therefore, we adapt the update rule at each iteration to match the new interval. When the range of the singular values is large, this approach ensures that the update rule shrinks it as quickly as possible. As the algorithm proceeds and the interval shrinks to a small neighborhood of $1$, the update rule approaches that of a Padé method, maintaining the same high order of convergence as it has.

由最优逼近得到的自适应有理方法。

2016 年一篇标志性论文引入了一种设计 $\polar(\bm{M})$ 迭代法的新范式 [33]。我们在第 4 节更详细地描述这一范式，其核心洞见如下。Padé 方法把更新规则选成在 $x=1$ 邻域内、给定次数下对 $\sign(x)$ 最优精确的逼近。相反，我们应当选择在包含奇异值的区间 $[\ell,1]\subset\mathbb{R}_{\geq 0}$ 上最优的 $\sign(x)$ 逼近。而且每走一步，奇异值的取值范围都会变；因此我们在每一步都按新区间调整更新规则。当奇异值范围很大时，这种方法能尽快把它缩小。随着算法推进、区间收缩到 $1$ 的一个小邻域，更新规则会逼近 Padé 方法，从而保持同样高的收敛阶。

Within the class of odd rational functions whose numerators and denominators have degree $2q+1$ and $2q$, respectively, an explicit formula for this optimal approximation to $\sign(x)$ on any interval $[\ell,1]$ was found by Zolotarev. It was shown that these rationals have remarkable convergence properties for any $q$ [33]. For $q=1$, this optimal approximation coincides exactly with the dynamically weighted Halley’s method (QDWH) referenced above. For even faster convergence than QDWH, [33] proposed the zolo-pd method, which uses $q=17$. Finally, these methods all admit the same QR-based implementation trick as QDWH.

在分子、分母次数分别为 $2q+1$ 和 $2q$ 的奇次有理函数类中，Zolotarev 给出了任意区间 $[\ell,1]$ 上对 $\sign(x)$ 最优逼近的显式公式。已经证明，这些有理函数对任意 $q$ 都有出色的收敛性质 [33]。当 $q=1$ 时，这一最优逼近恰好就是上面提到的动态加权 Halley 方法（QDWH）。为了比 QDWH 收敛得更快，[33] 提出了取 $q=17$ 的 zolo-pd 方法。这些方法都可以沿用 QDWH 那套基于 QR 的实现技巧。

Adaptive polynomial methods.

In this paper, we adopt the paradigm of zolo-pd [33] but with polynomials rather than rationals of degree $(2q+1,2q)$. This choice avoids the need for QR factorizations, relying solely on GPU-friendly matrix-matrix multiplications in low-precision arithmetic. While this class of methods has not been fully developed in the numerical analysis literature, similar ideas have been rediscovered in different guises. In an unpublished manuscript that predates zolo-pd, Chen and Chow [11] describe a rescaling strategy for Newton-Schulz. Though motivated differently, their method is equivalent to ours for degree-3 polynomials. They also observe numerical instability that prevents the method from converging to all the way to machine precision. Using the insights of [34], they propose a simple mitigation for this issue that we adopt in Section 4.4. Our work gives this method a theoretical foundation that connects it to the paradigm of zolo-pd, and we prove its optimality in the sense of (6). In addition, we study odd polynomials of arbitrary degree and focus particularly on the degree-5 case. Independently, a group of cryptographers developed a similar method for approximating the scalar function $\sign(x)$ in the context of homomorphic encryption schemes [28]. Their focus is mainly on tuning the analogues in their setting of the polynomial degree and number of iterations, whereas we focus on demonstrating optimality and efficiently constructing the update polynomials for degree $3$ and $5$. In addition, we consider matrix-valued inputs in low-precision arithmetic—not scalars in exact arithmetic—and we demonstrate our method’s effectiveness within the Muon algorithm for training deep neural networks.

自适应多项式方法。

本文沿用 zolo-pd [33] 的范式，但把次数为 $(2q+1,2q)$ 的有理函数换成多项式。这一选择避免了 QR 分解，只依赖在低精度算术下对 GPU 友好的矩阵-矩阵乘法。这类方法在数值分析文献中尚未得到充分发展，但类似想法以不同面貌被重新发现过。在早于 zolo-pd 的一份未发表手稿中，Chen 和 Chow [11] 描述了一种针对 Newton-Schulz 的缩放策略。尽管动机不同，他们的方法在三次多项式时与我们等价。他们还观察到数值不稳定性会阻止方法一直收敛到机器精度。借助 [34] 的见解，他们提出了一种简单的缓解办法，我们在第 4.4 节加以采用。我们的工作为这一方法提供了理论根基，把它连接到 zolo-pd 的范式，并在 (6) 的意义下证明了最优性。此外，我们研究任意次数的奇次多项式，并特别关注五次情形。独立地，一组密码学家在同态加密方案的语境下，为逼近标量函数 $\sign(x)$ 发展了类似方法 [28]。他们主要调节各自设定中与多项式次数和迭代步数对应的量，而我们侧重于证明最优性，并高效构造三次和五次的更新多项式。此外，我们考虑低精度算术下的矩阵值输入——而不是精确算术下的标量——并展示该方法在用于训练深度网络的 Muon 算法中的效果。

Application within Muon.

The designers of Muon realized that, due to the extreme efficiency requirements and lax accuracy requirements of their setting, rational-based methods from the numerical analysis literature are inapplicable. However, polynomial-based iteration schemes can take full advantage of GPUs because they use only matrix-matrix products in half-precision arithmetic, not inverses or QR decompositions. The preference for speed over accuracy motivates methods that aim to quickly produce coarse approximations, even at the cost of asymptotic convergence. Examples include the proposals of Jordan [22] and You [44, 9], as discussed in Section 1.2. Like Chen and Chow [11], Jordan found that convergence in the initial phase can be accelerated by choosing update rules that have a large derivative near zero, so as to increase the small singular values as much as possible at each iteration. You furthermore chose to use different update rules at each iteration, allowing extra flexibility to tune the trade-off between speed and accuracy. Both used degree-5 polynomials that were found through gradient descent on heuristic objective functions. These proposals were previously compared to Newton-Schultz[^3], but never to Chen and Chow’s method. We find that our method outperforms them all.

在 Muon 中的应用。

Muon 的设计者意识到，由于该设定对效率要求极高、对精度要求却很宽松，数值分析文献中的有理方法并不适用。但基于多项式的迭代格式可以充分利用 GPU，因为它们在半精度算术下只用矩阵-矩阵乘法，不要求逆或 QR 分解。对速度优先、精度其次的偏好，促使人们去设计能尽快给出粗糙近似的方法，哪怕牺牲渐近收敛。例子包括第 1.2 节讨论的 Jordan [22] 和 You [44, 9] 的方案。与 Chen 和 Chow [11] 类似，Jordan 发现：若选择在零点附近导数很大的更新规则，以便每一步尽可能抬升较小的奇异值，就能加快初期收敛。You 更进一步，在每一步使用不同的更新规则，从而有更大的灵活性来权衡速度与精度。两者都使用通过启发式目标函数上的梯度下降找到的五次多项式。这些方案此前曾与 Newton-Schultz[^3] 比较过，但从未与 Chen 和 Chow 的方法比较。我们发现我们的方法超过了它们全部。

3  Approximations by Compositions of Polynomials

3 用多项式复合做逼近

To design a GPU-friendly method for computing $\polar(\bm{M})$, we limit ourselves to the following GPU-friendly operations:

为了设计一种对 GPU 友好的 $\polar(\bm{M})$ 计算方法，我们把自己限制在下列对 GPU 友好的操作上：

i) Linear combinations: given scalars $\beta,\gamma\in\mathbb{R}$ and matrices $\bm{B}$ and $\bm{C}$, compute $\beta\bm{B}+\gamma\bm{C}$,

i) 线性组合：给定标量 $\beta,\gamma\in\mathbb{R}$ 以及矩阵 $\bm{B}$ 和 $\bm{C}$，计算 $\beta\bm{B}+\gamma\bm{C}$，

ii) Matrix-matrix products: compute $\bm{B}\bm{C}$.

ii) 矩阵-矩阵乘法：计算 $\bm{B}\bm{C}$。

While both these computational primitives are well-suited for parallel computing environments, matrix-matrix products come at a higher computational cost than linear combinations. Therefore, our method attempts to minimize the number of matrix-matrix products. A key observation is that we can compute odd monomials of $\bm{M}=\bm{U}\bm{\Sigma}\bm{V}^{\mathsf{T}}$ using the following formula:

这两种计算原语都适合并行计算环境，但矩阵-矩阵乘法的计算代价高于线性组合。因此我们的方法尽量减少矩阵-矩阵乘法的次数。一个关键观察是，对 $\bm{M}=\bm{U}\bm{\Sigma}\bm{V}^{\mathsf{T}}$ 的奇次单项式可以用下面的公式计算：

$$\bm{M}^{2q+1}:=\bm{U}\bm{\Sigma}^{2q+1}\bm{V}^{\mathsf{T}}=\bm{M}(\bm{M}^{\mathsf{T}}\bm{M})^{q}.$$

Hence, for an odd polynomial $p(x)=a_{0}x+a_{1}x^{3}+\cdots+a_{q}x^{2q+1}$ we can compute

因此，对奇次多项式 $p(x)=a_{0}x+a_{1}x^{3}+\cdots+a_{q}x^{2q+1}$，我们可以计算

$$p(\bm{M}):=a_{0}\bm{M}+a_{1}\bm{M}(\bm{M}^{\mathsf{T}}\bm{M})+\cdots+a_{q}\bm{M}(\bm{M}^{\mathsf{T}}\bm{M})^{q}.$$

It has been shown that for an arbitrary polynomial $p$, one requires $\Theta(\deg(p)^{1/2})$ products to compute $p(\bm{M})$ [39]; see also [20] for related work. This compares favorably to the naive approach that forms all monomials in $p$ and then sums them together, which requires $\Omega(\deg(p))$ products. However, if $p$ can be expressed as a composition of $T$ polynomials, each of degree $d$

已经证明，对任意多项式 $p$，计算 $p(\bm{M})$ 需要 $\Theta(\deg(p)^{1/2})$ 次乘法 [39]；相关工作另见 [20]。这比先形成 $p$ 中所有单项式再相加的朴素做法要好，后者需要 $\Omega(\deg(p))$ 次乘法。不过，如果 $p$ 能写成 $T$ 个次数均为 $d$ 的多项式的复合

$$p=p_{T}\circ p_{T-1}\circ\cdots\circ p_{1}, \tag{4}$$

then the degree of $p$ is $d^{T}$, and $p(\bm{M})$ can be efficiently computed recursively by

那么 $p$ 的次数就是 $d^{T}$，并且 $p(\bm{M})$ 可以高效地递归计算：

$$\bm{X}_{0}=\bm{M},\quad\bm{X}_{t}=p_{t}(\bm{X}_{t-1})\text{ for }t=1,2,\ldots,T. \tag{5}$$

The final iterate is $\bm{X}_{T}=p(\bm{M})$, which we compute with just $O(Td)$ matrix-matrix products.

最后一次迭代是 $\bm{X}_{T}=p(\bm{M})$，我们只用 $O(Td)$ 次矩阵-矩阵乘法就能算出它。

Iterative methods for $\polar(\bm{M})$ can be seen in this light. For instance, the degree-5 Newton-Schulz method uses the polynomial update $p_{t}(x)=\frac{15}{8}x-\frac{10}{8}x^{3}+\frac{3}{8}x^{5}$ for each $t=1,\ldots,T$. The composition $p=p_{T}\circ\cdots\circ p_{1}$ approximates $\sign(x)$, and the approximation error goes to $0$ as $T$ grows. In this paper, we ask the following question: what choice of $p_{T}\circ\cdots\circ p_{1}$ gives the best approximation to $\sign(x)$?

计算 $\polar(\bm{M})$ 的迭代法都可以从这个角度看。例如，五次 Newton-Schulz 方法对每个 $t=1,\ldots,T$ 都使用多项式更新 $p_{t}(x)=\frac{15}{8}x-\frac{10}{8}x^{3}+\frac{3}{8}x^{5}$。复合 $p=p_{T}\circ\cdots\circ p_{1}$ 逼近 $\sign(x)$，逼近误差随 $T$ 增大而趋于 $0$。本文要问的是：怎样选择 $p_{T}\circ\cdots\circ p_{1}$，才能给出对 $\sign(x)$ 的最好逼近？

The method we will present is optimal in the following sense: given lower and upper bounds $\ell$ and $u$ on the singular values of $\bm{M}$, an odd degree $d\in\mathbb{N}$, and the number of iterations $T\in\mathbb{N}$, our method computes the composition $p^{\star}(\bm{M})$ that minimizes the worst-case error in the spectral norm. That is,

我们将给出的方法在下述意义上是最优的：给定 $\bm{M}$ 的奇异值的下界 $\ell$ 和上界 $u$、一个奇次数 $d\in\mathbb{N}$，以及迭代次数 $T\in\mathbb{N}$，该方法算出的复合 $p^{\star}(\bm{M})$ 能把谱范数下的最坏情形误差降到最小。也就是说，

$$p^{\star}=\argmin_{\begin{subarray}{c}p=p_{T}\circ p_{T-1}\circ\cdots\circ p_{1}\
p_{t}\in\mathbb{P}{d}^{\odd}\end{subarray}}\max{\begin{subarray}{c}\bm{M}\in\mathbb{R}^{m\times n}\
\sigma(\bm{M})\subset[\ell,u]\end{subarray}}\left|\polar(\bm{M})-p(\bm{M})\right|_{2}. \tag{6}$$

Given that $\polar(\bm{M})-p(\bm{M})=\bm{U}(\bm{I}-p(\bm{\Sigma}))\bm{V}^{\mathsf{T}}$, and by the unitary invariance of the spectral norm, we have that (6) is equivalent to

鉴于 $\polar(\bm{M})-p(\bm{M})=\bm{U}(\bm{I}-p(\bm{\Sigma}))\bm{V}^{\mathsf{T}}$，再由谱范数的酉不变性，(6) 等价于

$$p^{\star};=;\argmin_{\begin{subarray}{c}p=p_{T}\circ p_{T-1}\circ\cdots\circ p_{1}\
p_{t}\in\mathbb{P}{d}^{\odd}\end{subarray}},\max{x\in[\ell,u]}\left|1-p(x)\right|. \tag{7}$$

For completeness, the equivalence between (6) and (7) is proven in Appendix C.

为完整起见，(6) 与 (7) 的等价性在附录 C 中证明。

(a) The left figure compares the composition (for $T=6$ and $d=5$) of polynomials given by Polar Express ($\ell=0.001$), You’s method (which is defined up to 6 iterations), Newton-Schulz, and Jordan’s method for approximating $\sign(x)$. The right figure demonstrates the convergence of the methods on $[0.001,1]$. Note the slow initial convergence of Newton-Schulz.

(a) 左图比较 Polar Express（$\ell=0.001$）、You 的方法（定义到 6 步）、Newton-Schulz 以及 Jordan 的方法在逼近 $\sign(x)$ 时所用多项式的复合（$T=6$，$d=5$）。右图展示这些方法在 $[0.001,1]$ 上的收敛情况。注意 Newton-Schulz 初期收敛很慢。

In other words, the problem given in (6) reduces to that of finding a “uniform” or “minimax” approximation to the constant function $x\mapsto 1$ over the interval $[\ell,u]$, as given in (7). Uniform approximation on an interval by polynomials or rational functions of a given degree is a central topic in approximation theory; see e.g. [46]. Here, we seek an approximation of a particular form—a composition of odd polynomials of fixed degrees. In the next section, we solve the optimization problem of (7) and use the solution to create Polar Express. Figure 2 (a) shows the resulting $p^{*}$ polynomial labeled as PolarExp, as compared to the Jordan’s method in [22], and the six iterations of You’s method in [9].

换句话说，(6) 中的问题归结为在区间 $[\ell,u]$ 上寻找对常值函数 $x\mapsto 1$ 的“一致”或“minimax”逼近，即 (7)。用给定次数的多项式或有理函数在区间上做一致逼近，是逼近论的中心课题；参见 [46]。这里我们要找的是一种特殊形式的逼近——固定次数的奇次多项式的复合。下一节我们求解 (7) 的优化问题，并用其解来构造 Polar Express。图 2 (a) 给出所得的 $p^{*}$ 多项式（标记为 PolarExp），并与 [22] 中 Jordan 的方法以及 [9] 中 You 方法的六步迭代相比较。

4  The Polar Express

4 Polar Express

4.1  Greedy is optimal

4.1 贪心就是最优

The key observation is that the polynomial used in each iteration can be chosen greedily, given the choice of polynomials from the previous iterations. For the first iteration, we choose $p_{1}$ so as to map the interval $[\ell,u]$ as close to $1$ as possible. That is, it minimizes $\max_{x\in[\ell,u]}|1-p_{1}(x)|$. The image of $p_{1}$ will be a new interval $[\ell_{2},u_{2}]$, where

关键观察是：每一步所用的多项式都可以在给定此前各步选择的前提下贪心地确定。对第一步，我们选择 $p_{1}$，使区间 $[\ell,u]$ 被映射到尽可能靠近 $1$ 的地方。也就是说，它最小化 $\max_{x\in[\ell,u]}|1-p_{1}(x)|$。$p_{1}$ 的像会是一个新区间 $[\ell_{2},u_{2}]$，其中

$$\ell_{2}=\min_{x\in[\ell,u]}p_{1}(x)\qquad\qquad u_{2}=\max_{x\in[\ell,u]}p_{1}(x) \tag{8}$$

We now pick $p_{2}$ to map the interval $[\ell_{2},u_{2}]$ as close to $1$ as possible, obtaining a new interval $[\ell_{3},u_{3}]$ that is the image of $[\ell,u]$ through $p_{2}\circ p_{1}$. We continue this process for as many iterations as desired.

接着我们选择 $p_{2}$，使区间 $[\ell_{2},u_{2}]$ 被映射到尽可能靠近 $1$ 的地方，得到新区间 $[\ell_{3},u_{3}]$，它是 $[\ell,u]$ 经过 $p_{2}\circ p_{1}$ 的像。我们按需要的步数继续这一过程。

The following theorem guarantees that this process finds the solution to (7), and thereby also (6). The scheme is also outlined in Figure 2 (b), which demonstrates the evolution of the lower bounds $\ell_{t}$, the upper bounds $u_{t}$, and the polynomials $p_{t}$ across iterations.

下面的定理保证这一过程能找到 (7) 的解，从而也找到 (6) 的解。该方案也在图 2 (b) 中示意，展示了下界 $\ell_{t}$、上界 $u_{t}$ 以及多项式 $p_{t}$ 随迭代的演变。

Theorem 4.1.

Let $d$ be odd and define $\ell_{1}=\ell$ and $u_{1}=u$. For $t=1,\ldots,T$ define

定理 4.1。

设 $d$ 为奇数，并定义 $\ell_{1}=\ell$、$u_{1}=u$。对 $t=1,\ldots,T$ 定义

$$\displaystyle p_{t} \\ \displaystyle=\;\argmin_{\begin{subarray}{c}p\in\mathbb{P}_{d}^{\odd}\end{subarray}}\,\max_{x\in[\ell_{t},u_{t}]}|1-p(x)| \\ \displaystyle\ell_{t+1} \\ \displaystyle=\;\min_{x\in[\ell_{t},u_{t}]}p_{t}(x) \\ \displaystyle u_{t+1} \\ \displaystyle=\;\max_{x\in[\ell_{t},u_{t}]}p_{t}(x) \tag{9}$$

Then the new error, lower and upper bounds can be computed through

那么新的误差以及新的下界、上界可以通过下式计算

$$\ell_{t+1}=p_{t}(\ell_{t}),\quad u_{t+1}=2-\ell_{t+1},\quad\text{ and }\quad\max\limits_{x\in[\ell_{t},u_{t}]}|1-p_{t}(x)|=1-\ell_{t+1}. \tag{10}$$

Furthermore, the composition $p^{\star}:=p_{T}\circ p_{T-1}\circ\cdots\circ p_{1}$ is optimal and the error is given by:

此外，复合 $p^{\star}:=p_{T}\circ p_{T-1}\circ\cdots\circ p_{1}$ 是最优的，误差由下式给出：

$$\max\limits_{x\in[\ell,u]}|1-p^{\star}(x)|\quad=\quad\min_{\begin{subarray}{c}p=p_{T}\circ p_{T-1}\circ\cdots\circ p_{1}\
p_{t}\in\mathbb{P}{d}^{\odd}\end{subarray}},\max{x\in[\ell,u]}\left|1-p(x)\right|=1-\ell_{T+1}. \tag{11}$$

Proof.

See Appendix A. ∎

证明。

见附录 A。∎

4.2****Why a fixed degree?.

We note that choice of the degree of each $p_{1},p_{2},\ldots,p_{T}$ need not be the same for Theorem 4.1 to hold. More generally, one may specify a sequence of degrees $d_{1},\ldots,d_{T}$ and define each $p_{t}$ as

4.2****为何固定次数？

我们指出，定理 4.1 并不要求每个 $p_{1},p_{2},\ldots,p_{T}$ 的次数相同。更一般地，可以指定一列次数 $d_{1},\ldots,d_{T}$，并把每个 $p_{t}$ 定义为

$$p_{t}=\argmin_{\begin{subarray}{c}p\in\mathbb{P}_{d_{t}}^{\odd}\end{subarray}}\,\max_{x\in[\ell_{t},u_{t}]}|p(x)-1|,\qquad\mbox{for }t=1,\ldots,T.$$

Our theory translates entirely to this more general case. However, for simplicity we assume $d=d_{t}$ for all $t=1,\ldots,T$. Our setting is similar to that of [28], which considers the closely related problem of choosing the depth $T$ and degrees $d_{1},\ldots,d_{T}$ such that $p$ approximates $\sign$ up to a prescribed error tolerance while minimizing the number of scalar multiplications. Interestingly, from [28, Table 2] the optimal choice of degrees is $d_{t}=5$ for almost all iterations. This justifies choosing $d$ to be a constant and our use of $d=5$ in particular.

我们的理论完全适用于这一更一般的情形。不过为简单起见，我们假定对所有 $t=1,\ldots,T$ 都有 $d=d_{t}$。我们的设定与 [28] 相近，后者考虑的相关问题是：选择深度 $T$ 和次数 $d_{1},\ldots,d_{T}$，使 $p$ 在达到指定误差容限的同时尽量减少标量乘法次数。有趣的是，从 [28, 表 2] 来看，几乎每一步的最优次数都是 $d_{t}=5$。这为把 $d$ 取成常数、尤其是取 $d=5$ 提供了理由。

Fortunately, (10) shows that once $p_{t}$ has been found, we can compute the new lower and upper bounds $\ell_{t+1}$ and $u_{t+1}$ and the approximation error simply by evaluating $p_{t}(\ell_{t})$. Hence, for any fixed upper and lower bounds on the singular values of $\bm{M}$, we can precompute the polynomials $p_{1},\ldots,p_{T}$ and the bounds $[\ell_{1},u_{1}],\ldots,[\ell_{T+1},u_{T+1}]$. Then, applying the iterative procedure of (5), the final iterate $\bm{X}_{T}$ will satisfy the following error bound

幸运的是，(10) 表明一旦找到 $p_{t}$，只需计算 $p_{t}(\ell_{t})$ 就能得到新的下界 $\ell_{t+1}$、上界 $u_{t+1}$ 以及逼近误差。因此，对 $\bm{M}$ 的奇异值的任意固定上、下界，我们可以预先计算多项式 $p_{1},\ldots,p_{T}$ 以及区间 $[\ell_{1},u_{1}],\ldots,[\ell_{T+1},u_{T+1}]$。然后执行 (5) 的迭代过程，最终迭代 $\bm{X}_{T}$ 将满足如下误差界

$$\|\polar(\bm{M})-\bm{X}_{T}\|_{2}=\|\polar(\bm{M})-p^{\star}(\bm{M})\|_{2}\leq 1-\ell_{T+1}. \tag{12}$$

From the optimality guarantee of Theorem 4.1, we know that our method converges at least as fast as the Newton-Schulz iteration of the same degree. Combining this fact with an existing analysis of Newton-Schulz, we immediately get the following convergence guarantee showing that our method enjoys faster than exponential convergence.

由定理 4.1 的最优性保证，我们知道该方法的收敛速度至少不慢于同次数的 Newton-Schulz 迭代。把这一事实与已有的 Newton-Schulz 分析结合起来，立刻得到下面的收敛保证，表明我们的方法具有快于指数的收敛。

Theorem 4.3.

Let $\bm{M}$ be a matrix normalized so that $\sigma(\bm{M})\subset[\ell,1]$. Let $\bm{X}_{T}=p^{\star}(\bm{M})$, where $p^{\star}$ is the polynomial from Theorem 4.1 with $d=2q+1$. Then, we have

定理 4.3。

设 $\bm{M}$ 已归一化，使得 $\sigma(\bm{M})\subset[\ell,1]$。设 $\bm{X}_{T}=p^{\star}(\bm{M})$，其中 $p^{\star}$ 是定理 4.1 中取 $d=2q+1$ 时的多项式。则有

$$\|\polar(\bm{M})-\bm{X}_{T}\|_{2}\leq|1-\ell^{2}|^{(q+1)^{T}}. \tag{13}$$

Hence, for $d=3$ the method converges quadratically and for $d=5$ the method converges cubically.

因此，当 $d=3$ 时方法二次收敛，当 $d=5$ 时方法三次收敛。

Proof.

See Appendix B. ∎

证明。

见附录 B。∎

In fact, Theorem 4.3 underestimates how fast our method converges. For degree $d=5$, our method converges about twice as fast as Newton-Schulz (compare with [11, Section 3.1]). Furthermore, the same analysis applies even if $p^{*}$ is constructed using a “lower bound” $\ell$ that was too high. That is, replacing $\ell$ on the right-hand side of (13) by $\sigma_{\min}$, the theorem holds even if $p^{*}$ is constructed to be optimal on the interval $[\ell,1]$ for $\ell>\sigma_{\min}$. Intuitively, when $\ell=u=1$, the polynomial $p^{*}$ coincides exactly with the Newton-Schulz method. Mistakenly setting $\ell>\sigma_{\min}$, we obtain a polynomial that converges slower than the optimal polynomial but faster than Newton-Schulz, so the guarantee of Theorem 4.3 still holds (cf. [11, Theorem 3.3]).

事实上，定理 4.3 低估了我们方法的收敛速度。对次数 $d=5$，我们的方法大约比 Newton-Schulz 快一倍（可与 [11, 第 3.1 节] 比较）。而且，即便构造 $p^{*}$ 时所用的“下界”$\ell$ 偏高，同样的分析仍然成立。也就是说，把 (13) 右边的 $\ell$ 换成 $\sigma_{\min}$ 之后，即使 $p^{*}$ 是针对某个 $\ell>\sigma_{\min}$ 的区间 $[\ell,1]$ 最优构造的，定理依然成立。直观上，当 $\ell=u=1$ 时，多项式 $p^{*}$ 恰好就是 Newton-Schulz 方法。若误把 $\ell$ 设得大于 $\sigma_{\min}$，得到的多项式会比最优多项式慢，但比 Newton-Schulz 快，因此定理 4.3 的保证仍然成立（参见 [11, 定理 3.3]）。

4.2  Finding the optimal polynomial for each iteration

4.2 为每一步寻找最优多项式

Theorem 4.1 shows that we can solve (7) by greedily choosing the optimal approximation $p_{t}\in\mathbb{P}_{d}^{\odd}$ for each interval $[\ell_{t},u_{t}]$ for $t=1,\ldots,T$. In this section, we show how to find each $p_{t}$. Since we are now focused on just one iteration, we drop the subscripts. Given $\ell$ and $u$, we wish to solve the following optimization problem:

定理 4.1 表明，只要对每个区间 $[\ell_{t},u_{t}]$（$t=1,\ldots,T$）贪心地选择最优逼近 $p_{t}\in\mathbb{P}_{d}^{\odd}$，就能求解 (7)。本节说明如何找出每个 $p_{t}$。既然现在只关注单步，我们省去下标。给定 $\ell$ 和 $u$，我们希望求解如下优化问题：

$$\argmin_{\begin{subarray}{c}p\in\mathbb{P}_{d}^{\odd}\end{subarray}}\,\max_{x\in[\ell,u]}|1-p(x)| \tag{14}$$

That is, we seek a minimax or uniform approximation of the function $x\mapsto 1$ on $[\ell,u]$ from the set of odd polynomials. (Equivalently, we seek a minimax optimal approximation to $\sign(x)$ on $[-u,-\ell]\cup[\ell,u]$.)

也就是说，我们要在奇次多项式集合中，寻找对区间 $[\ell,u]$ 上函数 $x\mapsto 1$ 的 minimax 或一致逼近。（等价地，我们要寻找 $[-u,-\ell]\cup[\ell,u]$ 上对 $\sign(x)$ 的 minimax 最优逼近。）

Problems of the form (14) are well-studied in approximation theory and numerical analysis. The key mathematical insight underlying the solution is the Equioscillation Theorem, which we state formally for our setting in Lemma A.1. This theorem gives a surprising characterization of the optimal solution of (14): an odd $p$ is optimal for degree $2q+1$ if and only if there is a set of $q+2$ equioscillating points. This is a set of points at which $p$ achieves its maximum approximation error $\pm E$, and for which the sign of the error alternates. Even if the optimal approximation error $E$ is not known in advance, finding a set of $q+2$ equioscillating points for a given $E$ serves as a certificate that no better approximation error is achievable. The Equioscillation Theorem is the basis of the Remez algorithm [37, 38], a general tool that can be used to find (nearly) optimal polynomial approximations of a given degree to any function on any interval. With very minor modifications to handle the constraint that $p$ be odd, Remez can be used to directly solve (14).

形如 (14) 的问题在逼近论和数值分析中已被充分研究。求解的关键数学洞见是等振荡定理，我们在引理 A.1 中针对本设定给出正式陈述。该定理给出了 (14) 最优解的一个出人意料的刻画：一个次数为 $2q+1$ 的奇次 $p$ 最优，当且仅当存在一组 $q+2$ 个等振荡点。这是一组使 $p$ 达到最大逼近误差 $\pm E$、且误差符号交替的点。即便事先不知道最优逼近误差 $E$，对给定的 $E$ 找到一组 $q+2$ 个等振荡点，也可以作为“不可能再有更好逼近误差”的证书。等振荡定理是 Remez 算法 [37, 38] 的基础；Remez 是一种通用工具，可用于寻找任意函数在任意区间上、给定次数的（近乎）最优多项式逼近。只需稍加修改以处理 $p$ 必须为奇次的约束，就可以直接用 Remez 求解 (14)。

However, the Remez algorithm is opaque, complex, and difficult to implement correctly. Fortunately, we do not need the Remez algorithm in its full generality to solve our problem. We seek only low degree polynomials, and the function we wish to approximate is the constant function $f(x)\equiv 1$. For $d=3$, we can derive an explicit, closed form solution to (14) using the Equioscillation Theorem. Up to rescaling, the optimal polynomial turns out to be the same one derived in Chen and Chow by different means [11]. For degree $d=5$, we present Algorithm 3, a much simpler way of solving (14) that is mathematically equivalent to Remez in our setting. This algorithm is implemented in its entirety in Appendix G.

不过 Remez 算法晦涩、复杂，也很难正确实现。幸运的是，求解我们的问题并不需要完整的一般 Remez 算法。我们只要低次多项式，而要逼近的函数是常值函数 $f(x)\equiv 1$。对 $d=3$，我们可以用等振荡定理导出 (14) 的显式闭式解。除去缩放之外，最优多项式恰好就是 Chen 和 Chow 用不同手段导出的那一个 [11]。对次数 $d=5$，我们给出算法 3，这是一种在本设定下与 Remez 数学等价、但简单得多的求解 (14) 的办法。该算法的完整实现见附录 G。

We briefly describe the solution for $d=3$. We seek a polynomial of the form $p(x)=ax+bx^{3}$. The Equioscillation Theorem stipulates that $p$ must have an equioscillating set of size 3. For $p$ to achieve its maximum error at a point $x$, $x$ must be a local extremum of $p(x)-1$ on the interval $[\ell,u]$. Thus, for $x$ to be eligible for membership in the equioscillating set, it must either be a true local extremum of $p(x)-1$ that happens to lie in $[\ell,u]$, or else one of the endpoints $\ell,u$. However, because $p$ is an odd cubic, it has at most one true local extremum on $\mathbb{R}_{\geq 0}$. Thus, to build an equioscillating set of three points, we must include $p$’s unique positive local extremum and both endpoints. This local extremum of $p$ occurs at $\sqrt{\frac{-a}{3b}}$. Therefore, we seek $a,b$ such that