Muon优化器

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

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

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

初步介绍

其中NewtonSchulz5被定义为

# 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优化器赏析：从向量到矩阵的本质跨越