Muon优化器
近日,Moonshot开源了改进版 Muon 优化算法及用 Muon 训练的SOTA级的MoE小模型。开启了Muon在大模型应用的局面。也许新的优化器时代即将到来!
PS:像谷歌23年提出的Lion(EvoLved Sign Momentum)优化器也号称比AdamW好,但是缺乏在大模型上的成功实验,大多数人还是选择Adam/AdamW。
苏神和官方作者Keller Jordan都写了一些博客来介绍这个内容。
初步介绍
其中NewtonSchulz5被定义为
1 |
|
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)。
作者也尝试用了使用三阶和七阶多项式(注意现在是五阶),但并没有达到更好的效果。