在 Muon 优化器:矩阵正交化驱动的梯度更新 中,我们建立了 msign 算子的数学骨架:把梯度矩阵 投影到最近的正交矩阵 ,并用 Newton-Schulz 迭代避免完整的 SVD。这套方案已经在 Kimi K2 上实现了 2× 训练加速,但它有一个隐疾:Newton-Schulz 在矩阵奇异值跨度大时,前十几步几乎不动。
ICLR 2026 Honorable Mention 论文 The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm (Amsel et al.) 把这个问题钉到了棺材里。他们证明:对任意奇异值区间 ,存在唯一的区间最优 odd quintic 多项式——而 Newton-Schulz 的 不是这个最优解。把每一步的多项式都换成对应区间的最优解,复合下来在 步 deg-5 多项式空间中全局最优。
实测上,Polar Express 比 Newton-Schulz 快约 2 倍达到机器精度;在 GPT-2 Large + FineWeb 1B tokens 的训练中,val loss 在所有 learning rate 上一致优于 Newton-Schulz 和此前流行的 Jordan / You 启发式。本文从 Newton-Schulz 的痛点出发,把 Jordan 启发式、You 六步法、Polar Express 三家解法放在同一个数学框架下对比,详解其背后的 Chebyshev 一致逼近 + Remez 算法,并逐行解读 80 行开源实现。
痛点:Newton-Schulz 的"前 17 步几乎不动"
Newton-Schulz 迭代
在标量上等价于 。它是把所有奇异值"推向 1"的一个标量映射,作用于每个奇异值独立。
固定点分析:。在 内 ,所以奇异值单调上升趋于 1。问题出在收敛速度上。 在 附近的 Taylor 展开
二次收敛——靠近 1 时极快。但在 时
只有 1.5 倍的线性增长。如果某个奇异值是 (接近实际深度学习梯度矩阵的最小奇异值),要把它推到 需要 步线性增长——而每步矩阵乘法都很贵。
论文 Figure 4 用合成矩阵( 均匀对数分布)展示:Newton-Schulz 前 17 步几乎不动,第 18 步之后才进入二次收敛区域,加起来需要 25-30 步。
这就是问题:经典数值分析时代追求"双精度收敛到机器精度",迭代次数本来就多;深度学习时代每步都要走 GPU 矩阵乘,前 17 步浪费的 FLOPs 就是真金白银。
修补尝试:Jordan 启发式与 You 六步法
社区早就意识到这个痛点。Jordan(Muon 原作者)在博客里给出一个度-5 多项式
代替 Newton-Schulz 的 。这个系数是用近似梯度下降在某个奇异值分布上优化出来的——本质是启发式调参。它在前几步比 Newton-Schulz 快得多,但代价致命:收敛性丢了。论文 Figure 4 显示 Jordan 迭代到第 11 步左右就卡死在误差 不动,无论再迭代多少次都过不去。
You 进一步把每一步用不同的 deg-5 系数,称作"六步法"。这是 Jordan 思路的扩展,但仍然不收敛——只是把 plateau 的高度从 0.3 压低一点,最终上限和 Jordan 同级。
两家的共同问题:用启发式调参代替严格的最优逼近。他们直觉对了(区间最优多项式应该自适应),但缺一个数学工具——能严格回答"给定区间 ,对常数 1 的最优 odd quintic 逼近是什么"。
这个工具是 Chebyshev 一致逼近理论里的 Remez 算法。
区间最优多项式:Chebyshev 与 Remez
一致逼近问题
设 是目标函数(这里 ), 是次数 的奇多项式空间(即 , 为奇数)。 意义下的最佳逼近是
为什么是 而非 ?因为我们要的是任意奇异值都收敛——最差情况上界比平均误差更重要。
等振荡定理(Equioscillation Theorem)
Chebyshev 1854 年的关键定理: 是 的 最佳次数 多项式逼近,当且仅当误差函数 在 上至少 个点处取得 且符号交替。
直观图像:最优逼近误差曲线像一条"波浪",在区间内的极值点处误差正负交替触顶。如果某段误差比其他段小,那一定不是最优——我们可以重新分配多项式系数让所有极值更均衡。
对 odd quintic ,,所以等振荡要求 个交替点。设这 4 个点为 (边界 + 两个内部极值),则有:
写成矩阵形式(把 视为未知数):
只要知道极值点位置 ,解一次 4×4 线性方程就拿到 。但极值点位置本身依赖于 ——是个鸡生蛋问题。
Remez 算法:交换迭代
Remez (1934) 给出了解法:
初始化 一组候选极值点(论文用四分位 ) 循环:
- 用当前 解线性系统 (1) 得
- 计算误差函数 的真实极值点位置:,令 解二次方程得新的
- 收敛判据: 则停
实际中通常 3-5 轮就收敛。optimal_quintic 的 Python 实现只有 20 行:
| |
完美对应上面三个步骤:LHS 是等振荡线性系统,np.linalg.solve 是步骤 1,sqrt(...) 一行是步骤 2 的二次方程显式解( 化为 的二次方程,闭式根公式)。
近似收敛的 shortcut(if 1 - 5e-6 <= l / u)值得一说:当区间已经收缩到 ,多项式逼近的数值条件数会爆炸,Remez 迭代失效。论文证明此时极限多项式恰好是 (即标准 deg-5 Newton-Schulz)经 缩放后的形式,直接返回。
区间复合:从单步最优到全局最优
单步最优只解决"给定区间 怎么找最好的多项式",但 Muon 迭代是要复合 个多项式。一个自然的问题:贪心地每步取当前区间的最优多项式,复合下来还是不是全局最优?
直觉上不一定——贪心算法在很多问题上得不到全局最优。但 Polar Express 论文的 Theorem 4.1 给出了肯定回答:
Theorem 4.1(区间最优的递归结构)。设 。定义
并令 ,。则
在所有 个 deg- 奇多项式复合中全局最优。
证明的核心引理是:在区间 上最优的 ,其最小值恰在 处,最大值在某个内部点 处。由等振荡对称性 给出 。这意味着:
计算下一轮区间不需要矩阵运算,只需 1 次标量求值 。
收敛速率(Theorem 4.3):对 ,
deg-3()二次收敛,deg-5()三次收敛——这是 Newton-Schulz 也具有的渐近阶。Polar Express 的优势不在渐近,而在早期收敛:从 出发不再"前 17 步几乎不动"。
把 Remez 求解和区间递推串起来就是 Algorithm 1 的离线阶段(optimal_composition):
| |
l = a*l + b*l**3 + c*l**5 一行直接对应 ,纯标量运算,毫秒级。整个离线阶段用 float64 跑一次得到 个系数三元组 ,烤进 coeffs_list 全局变量,线上推理直接查表用。
三家解法横向对照
把 Jordan 启发式、You 六步法、Polar Express 放在同一框架下比较:
| 方法 | 每步多项式 | 收敛性 | 早期表现 | 渐近阶 | bfloat16 友好性 |
|---|---|---|---|---|---|
| Newton-Schulz (deg-5) | 固定 | √ 全局收敛 | 极慢(前 17 步几乎不动) | 三次 | 需注意 |
| Jordan 启发式 | 固定 | × 卡在误差 0.3 | 快 | — | 未验证 |
| You 六步法 | 6 个固定 deg-5 多项式 | × 仍不收敛 | 快(前 6 步) | — | 未验证 |
| Polar Express (deg-5, ) | 8 个区间最优多项式(Remez 求解) | √ Theorem 4.1 保证最优 | 快(5-6 步到机器精度) | 三次 | √ 显式工程设计 |
三家的本质分歧:
- Newton-Schulz:用"无脑通用"多项式——对任何区间都用同一个。代价是初期慢,因为没有针对当前区间优化。
- Jordan / You:意识到要"自适应",但用启发式调参——快但失去收敛保证。卡在 0.3 不动是因为某些奇异值落在多项式的"死区"。
- Polar Express:用严格的最优逼近理论。每一步多项式都是当前区间 上的 最佳逼近,且组合后全局最优。
这里有个有趣的辩证:Newton-Schulz 的"固定多项式"是数值分析时代的产物——那时假设输入分布未知,必须用对任意输入都好的多项式。深度学习时代输入分布是已知的(梯度矩阵的奇异值有典型分布),可以预先针对这个分布优化。Polar Express 的本质是把输入分布的先验知识编进多项式系数。
苏剑林在 msign 算子的 Newton-Schulz 迭代(下) 中也指出了同一方向:
Newton-Schulz 迭代的优化已经形成了一个小的研究方向,主要思路是利用输入分布的先验信息来调整多项式系数。
Polar Express 给出了这个方向上的"理论最优"答案。Jordan 和 You 是这个方向上的早期尝试,Polar Express 是定论。
bfloat16 工程取舍:safety factor 与 cushion
理论最优在 float64 下完美,搬到 bfloat16(深度学习训练的实际精度)会出意外。论文 Section 4.4 给出三个工程干预:
Safety factor
理论上每步要求奇异值始终留在 内。但 bfloat16 的 rounding 可能让某个 一点点,落到 区间。在那里,多项式 超出了它的"设计区域",会把误差放大而非缩小。Nakatsukasa & Higham (2012) 给出的常数:每次 rounding 误差 经过 后被放大到约 ,多步累积就发散。
干预方法:把每个 替换为 ,等价于把"安全区间"略微往下搬一点。代码实现极简:
| |
为什么不同次幂的除数不同?因为 。最后一步(iter == num_iters - 1)不应用 safety factor——它不需要为下一步预留误差余量。
Cushion(自由下界保护)
等振荡多项式有一个数值病态行为:当 极小(比如 )时,多项式在 端的"压缩比" 会非常接近 0。这意味着大奇异值会被压得极小,下一轮的 就变得非常接近 ——区间在 1 附近过度收缩,加上 bfloat16 的精度损失,后续迭代基本失效。
cushion 是个软限制:当 (论文用 )时,Remez 求解用区间 而非 ,然后重缩放系数让 (恢复等振荡对称):
| |
代价是稍微减慢收敛——但论文说不显著。
预生成系数表
最后一个工程取舍:所有 个多项式系数在程序加载时一次性算好,固化进 coeffs_list 全局变量。线上推理只跑矩阵乘法循环:
| |
这意味着 Polar Express 在线推理的每步开销与 Newton-Schulz 完全相同——3 次小矩阵乘 + 1 次大矩阵乘。唯一新增的是查表读 ,可以忽略。
Horner 形式的等价性值得一验。设 为 SVD,则
所以 在 SVD 意义下就是把 的奇异值立方(同时保留奇异向量)。同理 是奇异值五次方。因此
恰是把多项式 作用到每个奇异值上。
FLOPs 计数:设 且 (转置后保证), 是 小矩阵。每步 3 次 矩阵乘 + 1 次 矩阵乘,约 FLOPs,与 Newton-Schulz 持平。
实验:GPT-2 上的实战
论文 Section 5 给出三组实验:合成矩阵收敛、真实梯度收敛、GPT-2 端到端训练。
合成矩阵(论文 Figure 4): 在 对数均匀分布,左右奇异向量随机。结果:
| 迭代步数 | Newton-Schulz | Jordan | You (6 步) | Polar Express () |
|---|---|---|---|---|
| 5 步 | 误差 | |||
| 10 步 | 卡在 0.3 | |||
| 15 步 | 卡在 0.3 | 卡在 0.3 | 机器精度 | |
| 25 步 | 机器精度 | 卡在 0.3 | 卡在 0.3 | 机器精度 |
最反直觉的发现:Jordan 启发式即使迭代 100 步也卡在误差 0.3——它从一开始就不收敛。Polar Express 第 15 步已经到机器精度。
GPT-2 Large + FineWeb 1B tokens(论文 Figure 1):
| 优化器 | 最佳 val loss | 学习率 |
|---|---|---|
| AdamW | 4.172 | 1e-4 |
| Muon (Newton-Schulz) | — | — |
| Muon (Jordan) | 3.398 | 2e-2 |
| Muon (You) | 3.400 | 2e-2 |
| Muon (Polar Express) | 3.340 | 2e-2 |
差距看似不大(3.340 vs 3.398 = 0.058),但这是在所有 learning rate 上稳定的优势。论文 Figure 6 把 lr 横扫一遍,Polar Express 在每个 lr 点都领先。
实际工程意义:
- 减少 Muon 内 msign 子例程的迭代步数(论文实验用 5 步,业界常规是 5-7 步)
- 在 lr 容忍度上有更宽的甜区——重要,因为 lr 是大模型训练最不希望调的超参
- 训练曲线更稳定,方差更小
后续工作 / 与其他方向的接口
Polar Express 把"Muon 子例程优化"这条线推到了理论最优。但它不是 Muon 工程化的终点——还有几条平行的方向:
流式幂迭代(苏剑林在 基于流式幂迭代的 Muon 实现 中讨论的方案):不是在每步训练中跑完整的 Newton-Schulz / Polar Express,而是把 SVD 估计平摊到训练循环里——本质是把多项式迭代用幂迭代替代。两条路线互补:Polar Express 优化"单次正交化",流式幂迭代优化"训练循环里的均摊成本"。
MuP(Maximal Update Parametrization)框架:Muon 在 MuP 下天然满足缩放稳定性,使得 hyperparameter 可以从小模型迁移到大模型。Polar Express 不改变 Muon 的 MuP 性质,因为它只优化 msign 子例程,不动外层。
流形上的最速下降:苏剑林在 流形上的最速下降系列 中把 Muon 视为 Stiefel 流形上的优化器——Polar Express 是其投影算子的更高效实现,没有改变流形结构本身。
Halley 迭代 / QDWH 等经典数值分析方法:它们渐近收敛阶更高(Halley 是三次,QDWH 是六次),但依赖矩阵求逆或 QR 分解——GPU 不友好。Polar Express 在 GPU 时代的胜利是保留了"只用矩阵乘"的特性,同时获得了区间最优收敛速度。
小结
Polar Express 的核心贡献用一句话概括:用 Chebyshev 一致逼近 + Remez 算法,把 Newton-Schulz 的"固定多项式"换成"区间最优多项式",并证明贪心区间复合在 步 deg-5 多项式空间中全局最优。这个改动只增加了离线一次性的多项式求解成本(毫秒级),线上推理 FLOPs 与 Newton-Schulz 完全相同,但前期收敛速度快约 2 倍。
从 Newton-Schulz → Jordan / You → Polar Express 的演进,展示了一条非常清晰的研究路径:
- 观察痛点(NS 前 17 步几乎不动)
- 启发式尝试(Jordan 调参快但不收敛)
- 发现根本工具(Chebyshev / Remez 是这个问题的数学语言)
- 严格最优化(Polar Express)
每一步都是站在前人肩膀上的小步前进,但累积下来是从"启发式 + 失败"到"理论最优 + GPT-2 端到端胜利"的质变。这也是数学工具进入工程的典型方式:以经典理论(Chebyshev 1854)的翻译为载体,进入新的工程语境(bfloat16 GPU)。
对深度学习优化器设计的启示:矩阵正交化子例程的最优化已经接近天花板(理论最优 + 工程实现完备)。下一个 2× 加速大概率来自其他方向——流式幂迭代均摊、动量与正交化的联合设计、或者跳出 Muon 框架的全新优化器范式。
相关概念
- SVD 与极分解 — 极分解 的几何,msign 就是其中的 ,详见 奇异值分解与低秩近似
- 谱范数与条件数 — Newton-Schulz 收敛性的核心工具,Polar Express 区间最优本质是谱范数意义下的最佳逼近,详见 谱范数、条件数与优化景观
- Muon 优化器 — 本文的直接前置,msign 算子 / Newton-Schulz 迭代推导 / 与 Adam/Shampoo 的对比,详见 Muon 优化器:矩阵正交化驱动的梯度更新
参考文献
- Amsel, N., Persson, D., Musco, C., & Gower, R. M. (2025). The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm. ICLR 2026 (Honorable Mention). arXiv:2505.16932
- 代码: github.com/NoahAmsel/PolarExpress
- 苏剑林. msign 算子的 Newton-Schulz 迭代(上). https://kexue.fm/archives/10922
- 苏剑林. msign 算子的 Newton-Schulz 迭代(下). https://kexue.fm/archives/10996
- 苏剑林. 流形上的最速下降:3. Muon + Stiefel. https://kexue.fm/archives/11221
- 苏剑林. 基于流式幂迭代的 Muon 实现:1. 初识. https://kexue.fm/archives/11654
- Jordan, K., et al. (2024). Muon: An optimizer for hidden layers in neural networks. arXiv preprint.
- Chen, J., & Chow, E. (2014). A stable scaling of Newton-Schulz for improving the sign function computation of a Hermitian matrix. Argonne National Laboratory Tech Report ANL/MCS-P5059-0114.
- Nakatsukasa, Y., & Freund, R. W. (2016). Computing fundamental matrix decompositions accurately via the matrix sign function in two iterations. SIAM Review, 58(3).
- Nakatsukasa, Y., & Higham, N. J. (2012). Backward stability of iterations for computing the polar decomposition. SIAM J. Matrix Anal. Appl., 33(2).
- Higham, N. J. (2008). Functions of Matrices: Theory and Computation. SIAM.
- Achieser, N. I. (1956). Theory of Approximation. Frederick Ungar Publishing.
- Trefethen, L. N. (2013). Approximation Theory and Approximation Practice. SIAM.