Muon优化器

近日,Moonshot开源了改进版 Muon 优化算法及用 Muon 训练的SOTA级的MoE小模型。开启了Muon在大模型应用的局面。也许新的优化器时代即将到来!

PS:像谷歌23年提出的Lion(EvoLved Sign Momentum)优化器也号称比AdamW好,但是缺乏在大模型上的成功实验,大多数人还是选择Adam/AdamW。

苏神和官方作者Keller Jordan都写了一些博客来介绍这个内容。

初步介绍

其中NewtonSchulz5被定义为

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# Pytorch code
def newtonschulz5(G, steps=5, eps=1e-7):
assert G.ndim == 2
a, b, c = (3.4445, -4.7750, 2.0315)
X = G.bfloat16()
X /= (X.norm() + eps)
if G.size(0) > G.size(1):
X = X.T
for _ in range(steps):
A = X @ X.T
B = b * A + c * A @ A
X = a * X + B @ X
if G.size(0) > G.size(1):
X = X.T
return X

Muon取得了以下经验结果。

  • 研究者将训练速度记录提高到在 CIFAR-10 上达到 94%准确率的 2.6 A100-秒,从 3.3 秒提升。
  • 将训练速度记录提升至 FineWeb(一个被称为 NanoGPT 速度跑的竞争性任务)上的 3.28 验证损失,提高了 1.35 倍。
  • 持续在扩展到 774M 和 1.5B 参数时展示训练速度提升。
  • 训练了一个 1.5B 参数的 Transformer,在 HellaSwag 上达到 GPT-2 XL 级别的性能,耗时 10 个 8xH100 小时。使用 AdamW 达到相同结果需要 13.3 小时。

为何选择NewtonSchulz5

正交化矩阵有好几种方法,但并不好用或会出现问题。

(1)SVD

易于理解,但太慢了。

(2)耦合牛顿迭代(Coupled Newton iteration)

Shampoo中用于执行逆四次方根,并且可以轻松地适应进行正交化。

但它在至少 float32 精度下运行以避免数值不稳定性,这使得它在现代 GPU 上运行缓慢。

(3)NewtonSchulz5

可以在bfloat16 中稳定运行

NewtonSchulz5介绍

该算法和msign(矩阵符号函数)有关。

它也与SVD有关。
$$
U, \Sigma, V^{\top}=S V D ( M ) \quad\Rightarrow\quad\mathrm{m s i g n} ( M )=U_{[ : ,: r ]} V_{[ :, : r ]}^{\top}
$$
$U\in R^{n\times n},\Sigma\in R^{n\times m},V\in R^{m\times m}$,r是M的秩。

进一步,
$$
\mathrm{m s i g n} ( M )=( M M^{\top} )^{-1 / 2} M=M ( M^{\top} M )^{-1 / 2}
$$
进一步,当m=n=r时,
$$
msign(M)=argmin_{O^TO=I}||M-O||^2_F
$$

证明:

$||M-O||^2_F=||M||^2_F+||O||^2_F-2\left<M,O\right> _ F$

$=||M||^2_F+n-2Tr(MO^T)$
$=||M||^2_F+n-2Tr(U\Sigma V^TO^T)$
$=||M||^2_F+n-2Tr(\Sigma V^TO^TU)$
$=||M||^2_F+n-2\sum_i^n \Sigma_{i.i}(V^TO^TU)_{i,i})$

最小值时$O=UV^T$

Newton-schulz是对msign(M)的近似。

我们对$(M^TM)^{-\frac{1}{2}}$在$M^TM=I$进行展开到二阶,有
$$
msign(M)=M(M^TM)^{-1/2}\approx \frac{15}{8}M-\frac{5}{4}M(M^TM)+\frac{3}{8}M(M^TM)^2
$$
然而泰勒展开有可能会有误差,既然我们得到了这样的形式,我们干脆假设$msign(M)\approx aM+bM(M^TM)+cM(M^TM)^2$

作者进一步优化得到(a,b,c)=(3.4445, -4.7750, 2.0315)。

作者也尝试用了使用三阶和七阶多项式(注意现在是五阶),但并没有达到更好的效果。

参考资料

  1. Muon: An optimizer for hidden layers in neural networks

  2. Muon优化器赏析:从向量到矩阵的本质跨越


Muon优化器
https://lijianxiong.work/2025/20250224/
作者
LJX
发布于
2025年2月24日
许可协议