把 Muon 的每一步都算到最优：Polar Express 如何让 GPT-2 训练更快更稳

TL;DR

Muon 优化器之所以快，是因为它每一步不沿着动量方向走，而是沿着动量矩阵的“最接近半正交矩阵”走——也就是求一次极分解。但在 GPU 上精确求极分解太贵，于是大家一直用 Newton-Schulz 这类多项式近似：只用矩阵乘法，精度粗糙但够用。ICLR 2026 荣誉提名论文 Polar Express 问得更彻底：既然每步只能用一个低次多项式，能不能让每一步的多项式都是当时条件下的最优解？答案是可以，而且可以贪心地一步步求出来：每轮解一个 minimax 问题，把当前奇异值区间尽可能压向 1。结果是理论上每步最优、渐近超指数收敛，实践中 5 到 6 步就够用；塞进 Muon 之后，774M 参数的 GPT-2-Large 在 10 亿 FineWeb token 上，最优学习率下的验证损失从前任方法的 3.398 降到 3.340，而且在全部学习率下都更好。这不是换模型或数据，而是把一个被当成“差不多就行”的数值子程序真正算到了最优。

论文信息卡

标题：The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm
作者：Noah Amsel、David Persson、Christopher Musco、Robert M. Gower
单位：纽约大学、Flatiron Institute
会议：ICLR 2026 荣誉提名（Honorable Mention）
原文链接：https://arxiv.org/abs/2505.16932
代码：https://github.com/NoahAmsel/PolarExpress

背景：Muon 火了，但它最核心的一步一直是“凑合”

要理解这篇论文，先要理解 Muon 到底在算什么。标准带动量的 SGD 先把梯度做滑动平均得到动量 $M_t$ ，再沿着 $-M_t$ 更新权重。Muon 只改了一处：不直接走 $-M_t$ ，而是走 $-\mathrm{polar}(M_t)$ 。这里的 $\mathrm{polar}(M)$ 是极分解中的正交因子：若 $M = U\Sigma V^T$ ，则 $\mathrm{polar}(M) = UV^T$ ，也就是在谱范数意义下离 $M$ 最近的半正交矩阵。直观上说，Muon 把动量矩阵的每个奇异值都“压平”成 1：方向保留，幅度归一。这等价于谱范数下的最速下降方向。

问题出在“压平”这一步怎么算。直接做 SVD 是 $O(\min(mn^2, nm^2))$ 的量级，而且经典 SVD 算法很难吃满 GPU 并行。数值分析领域几十年来的解法五花八门：Newton 法、Padé 迭代、QDWH、Zolo-pd 等等，但它们大多要算矩阵求逆或 QR 分解——在 GPU 上不好并行，在半精度下也不稳定。深度学习要的恰恰相反：只要矩阵乘法，只要半精度，只要几步就能出一个“够用”的近似，16 位小数精度纯属浪费。于是 Muon 最初选了最老的办法 Newton-Schulz：3 次版本 $X_{t+1} = \frac{3}{2}X_t - \frac{1}{2}X_tX_t^TX_t$ ，相当于对每个奇异值反复套用 $p(x) = \frac{3}{2}x - \frac{1}{2}x^3$ ；5 次版本 $(15x - 10x^3 + 3x^5)/8$ 更快。但它有个结构性毛病：前期慢，奇异值离 1 远时几乎原地踏步，而 Muon 每步只给 5 步左右预算，前期慢就等于全程慢。

后来社区出了两个启发式补丁：Jordan 的 5 次多项式起步极快，但误差 plateau 在约 0.3、根本不收敛；You 连用 6 个不同的多项式，精度更好，但同样不收敛。两者系数均为启发式搜索所得，没有最优性保证。

方法拆解：把“求极分解”翻译成“每步压区间”

第一步：矩阵问题退化成标量问题

关键观察是，多项式作用在矩阵上，等价于逐个作用在它的奇异值上。若 $M = U\Sigma V^T$ ，奇次多项式 $p$ 作用后就是 $p(M) = Up(\Sigma)V^T$ ，其中 $p(\Sigma)$ 是对角线上每个奇异值套用 $p$ 。而 $\mathrm{polar}(M) = UV^T$ ，相当于把每个奇异值都变成 1。于是矩阵近似误差（谱范数）就等于标量误差 $\max_{x \in [\ell, u]} |1 - p(x)|$ ，其中 $[\ell, u]$ 是奇异值所在的区间。求极分解的问题，就这样退化成：在给定区间上，用若干个奇次多项式的复合去一致逼近常数函数 1。

第二步：贪心居然就是全局最优

论文的核心定理（定理 4.1）说：贪心就是最优。做法是：

已知当前奇异值区间 $[\ell_t, u_t]$ ，选一个 $d$ 次奇多项式 $p_t$ ，让它在该区间上离 1 的最大偏差最小，即解 minimax 问题 $\arg\min_{p} \max_{x \in [\ell_t, u_t]} |1 - p(x)|$ ；
把区间推进一格：新的下界 $\ell_{t+1} = p_t(\ell_t)$ ，新的上界 $u_{t+1} = 2 - \ell_{t+1}$ ，这一步的误差恰好是 $1 - \ell_{t+1}$ ；
对新区间重复上述过程。

定理保证，这样贪心拼出来的复合多项式 $p^\star = p_T \circ \cdots \circ p_1$ ，就是同类方法中最优的。更进一步，它超指数收敛：对 $d = 2q + 1$ 次多项式，误差上界是 $|1 - \ell^2|^{(q+1)^T}$ ， $d = 3$ 时二次收敛， $d = 5$ 时三次收敛；实际大约比同次 Newton-Schulz 快一倍。

为什么贪心能成立？直觉是：每一步的最优多项式只依赖当前区间，而当前区间又完全由上一步的最优结果决定——前一步省下的误差，不会在后一步被浪费。

第三步：怎么算出每步的最优多项式

Minimax 多项式有经典的刻画工具：等振荡定理。最优的 $d = 2q + 1$ 次奇多项式，一定存在 $q + 2$ 个等振荡点，在这些点上误差取到 $\pm E$ 且正负交替。 $d = 3$ 时有闭式解：最优多项式是对 Newton-Schulz 多项式的缩放，恰与 Chen 与 Chow 早年的结果一致。

对实际主推的 $d = 5$ ，论文给了一个比通用 Remez 算法简单得多的专用算法：在区间内猜四个等振荡点 $\ell < q < r < u$ ，解一个 $4 \times 4$ 线性方程组得到系数与误差 $E$ ，再由多项式导数的极值点更新 $q, r$ ，交替迭代直到收敛。整个离线阶段用 float64 算一次，系数存下来之后在线直接复用——这就是论文算法 1 里那张 coeffs_list 的来历：前 7 组系数各不相同，第 8 组起收敛到固定的 $(1.875, -1.25, 0.375)$ ，也就是 5 次 Newton-Schulz 多项式本身。区间压近 1 之后，最优多项式自然退化为经典方法。

第四步：让理论在 bfloat16 里活下来

纯理论的最优多项式搬到半精度会出三种毛病，补丁均可在离线阶段改系数完成，不增加在线开销：

超界发散：舍入可能让某个奇异值略大于当前上界 $u_t$ ，最优多项式在界外可能把它越推越大。办法是把每个多项式换成 $p_t(x/1.01)$ ，留 1% 安全裕度，代价只是收敛目标从 1 变成 0.999998，最后一轮去掉该因子即可。
非单调掉精度：最优多项式在区间内来回振荡，可能把靠近 $u_t$ 的奇异值一下子打回 $\ell_{t+1}$ ，极端情况下甚至让 $p_t(\sigma_i)$ 变负、符号翻转。办法是“垫底”：当 $\ell_t < u_t/10$ 时，就按 $\ell_t = u_t/10$ 来选多项式，保证 $p_t(x)/x \geq 0.236$ ，收敛只慢一点点。
归一化偏移：沿用 Muon 原实现，用 $\|M\|_F + 10^{-2}$ 而不是 $\|M\|_F$ 做归一化。

最终的在线算法干净得可以当 drop-in 替换：矩阵转 bfloat16，必要时转置以省 FLOPs，除以 Frobenius 范数，每步算 $X \leftarrow aX + bX^3 + cX^5$ ，5 到 6 步结束。作者推荐 $d = 5$ 、 $T = 5$ 或 $6$ 、 $\ell_1 = 10^{-3}$ ——其中 $10^{-3}$ 也有讲究：bfloat16 的机器精度是 $2^{-7} \approx 0.0078$ ，奇异值一般不会比它小太多，猜错一两个数量级也只拖慢几步。

Polar Express 的复合多项式与误差收敛对比

图：复合多项式对符号函数的逼近与收敛对比，Newton-Schulz 前期慢的问题一目了然（来源：论文原文）

前三个最优多项式与区间演化

图：前三个最优多项式 p1、p2、p3 的演化及奇异值上下界的收窄过程（来源：论文原文）

实验与数据：不是只赢一点，而是全学习率都赢

收敛实验：前期慢的病治好了，平台期的病也没有

论文先在“已知谱”的理想场景里验收：构造奇异值在 $10^{-6}$ 到 1 之间对数均匀分布的随机矩阵，比较 5 次 Newton-Schulz、Jordan、You、Chen & Chow 与 Polar Express。结果很有戏剧性：Newton-Schulz 前 17 步几乎不动，之后才加速；Jordan 11 步就冲到约 0.3 的误差，但此后彻底 plateau；You 只定义了 6 步，走势与 Jordan 类似。调好 $\ell = \sigma_{\min}$ 的 Polar Express 则在每一步都压过所有对手，11 步就达到很好的精度，到达任一给定误差的速度约为 Newton-Schulz 的两倍。即使 $\ell$ 猜错两个数量级，它仍有竞争力。

在真实 GPT-2 梯度矩阵上（第 4 个 Transformer 块的权重梯度，Tiny Shakespeare 批次），结论基本复现，只有一个提醒：把 $\ell$ 猜得过小很多个数量级时，收敛会被明显拖慢，甚至不如 Jordan。

GPT-2 训练：大模型与小模型，验证损失双双更低

真正的考验是放进 Muon 训语言模型。实验基于 Jordan 的 nanoGPT 改版，两个模型、同一套协议：GPT-Small（124M，12 层）与 GPT-Large（774M，36 层），词表 50257、上下文 1024，在 FineWeb 的 10 亿 token 上训 1 个 epoch、batch 32、4 张 H100、bfloat16 混合精度；学习率先恒定 40% 步数再线性衰减；Muon 只用于二维以上权重矩阵（其余用 AdamW）；矩阵符号计算统一用 5 步 5 次多项式——三个方法每步计算量完全相同，拼的纯粹是多项式质量。

GPT-Large 的结果是图 1 的主角：最优学习率 $lr = 0.02$ 下，AdamW 验证损失 4.172，muon-You 3.400，muon-Jordan 3.398，muon-PolarExp 3.340。关键不在 0.058 的绝对值，而在“在全部学习率下都更好”：左图里 PolarExp 的曲线整条压在另外两条下面，右图里用最优学习率训全程，它的验证损失曲线也始终最低。GPT-Small 的图 6 复现了同一模式：最优验证损失 AdamW 4.197、muon-Jordan 3.639、muon-You 3.629、muon-PolarExp 3.588；训练损失同样更低，而且因为每步开销相同，按 epoch 赢就等于按 wall-clock 时间赢。两个量级、全部学习率——一致性本身就值得信任。

局限与开放问题

小模型、小 token 量的验证：正文实验停留在 GPT-2 量级（124M/774M）与 10 亿 token，能否平移到十亿参数以上、万亿 token 的训练，论文未披露。
$\ell$ 猜测仍是启发式：理论允许猜错，但实验也显示猜得过小会明显拖慢收敛；一个全局 $\ell = 10^{-3}$ 是否对所有层都合适，仍待检验。
只动了 Muon 的一个子程序：Muon 本身的动量、学习率调度、与 AdamW 的分工都沿用前人设定，本文不回答“Muon 为什么好”。
半精度之外的场景未验证：bfloat16 下的三处稳定补丁是针对深度学习定制的；在 float32/float64 科学计算里，它与 QDWH、Zolo-pd 的性价比如何，论文未做正面比较。
实现细节仍有取舍：安全因子 1.01、垫底阈值 $u_t/10$ 都是“够用就好”的常数。

一点思考

这篇论文最值得带走的是两条判断。第一，优化器里的数值子程序值得像架构一样认真对待：当大家默认“近似一下就行”时，一个有证明的最优近似能在不增加任何计算量的前提下换来全学习率的损失下降，这种“免费午餐”在训练栈里并不多见。第二，别把启发式搜索的结果当成终点：Jordan 与 You 的方法在实践中已经很好用，但它们 plateau 在 0.3、永不收敛的毛病，恰恰是“目标函数定错了”的症状——把目标改成区间 minimax，系数就变成了可证明的最优解。我的判断是：如果你已经在用 Muon，Polar Express 值得直接替换现有的 Newton-Schulz/Jordan 实现（作者已开源、系数即插即用）；还没用 Muon 的话，这篇不是入坑理由——它优化的是 Muon 内部一步，而非 Muon 本身的适用性。

原文链接

论文：https://arxiv.org/abs/2505.16932
代码：https://github.com/NoahAmsel/PolarExpress
会议荣誉：ICLR 2026 Outstanding Papers 公告（Honorable Mention）https://blog.iclr.cc/2026/04/23/announcing-the-iclr-2026-outstanding-papers/