初探优化器（一）：从SGD到SOAP-Muon

最近一段时间在看优化器相关的工作。一开始的目的是想跟一下 Modded NanoGPT，研究一下训练加速 oriented 的模型架构设计，发现还有个专门做加速器的 Track3。看的过程中感觉自己对优化器的理解太零散，Adam 大概知道在做什么，Muon 大概知道是把奇异值拉平，Shampoo 和 SOAP 就只记得是在梯度两边各乘一个矩阵，至于这些方法之间是什么关系、为什么有的在 Transformer 上好用有的不好用，基本说不太清楚。所以干脆把这段时间读的论文和做的实验整理成一个系列，算是给自己补课。

这一篇先把已有的优化器过一遍，从最基础的 SGD 一直讲到 Track 3 现在纪录里在用的 SOAP-Muon，整理一下逻辑的思路。

按时间排下来大概是这样（2010 年之前的部分压缩了），我把它们分成了三条线：逐元素（elementwise）的方法，从 SGD、动量一直到 Adam 和 AdamW；利用曲率（curvature）的方法，从 Newton 法到 Natural Gradient 和 K-FAC；以及把权重矩阵当成矩阵来处理的方法，Shampoo、Muon、SOAP 和 SOAP-Muon。最后这条线 2024 年之后变得非常快，Track 3 上现在排在前面的基本都是这一类。

0. Preliminary

0.1. 优化器在做什么

考虑一组参数 $\theta$ 和一个训练要最小化的 loss $L(\theta)$，每一步从训练集里取一个 batch $B_t$，算出这个 batch 上的梯度

然后更新参数

$\eta_t$ 是学习率，$\Delta_t$ 是这一步的更新方向。这篇里讲到的所有优化器，说到底都在回答同一个问题：手上有到目前为止的所有梯度 $g_1, \dots, g_t$，再加上自己存的一些状态（比如动量），这一步的 $\Delta_t$ 应该怎么取。最简单的取法是 $\Delta_t = g_t$，也就是 SGD。

另一个会一直出现的对象是线性层。$y = Wx$，权重 $W \in \mathbb{R}^{m \times n}$ 把 $n$ 维的输入映射成 $m$ 维的输出，Transformer 里绝大部分参数都在这种矩阵里（Attention 的 $W_Q, W_K, W_V, W_O$，MLP 的两个矩阵）。一个 batch 里有 $N$ 个 token，第 $i$ 个 token 的输入记为 $x_i$，loss 对它输出的梯度记为 $\delta_i = \partial L / \partial y_i$，那么 $W$ 的梯度是

也就是 $N$ 个秩为 1 的外积加在一起。讲逐元素方法的时候用不太到这个式子，把所有参数拉平成一个长向量就行；但到了 K-FAC、Muon、Shampoo 这些方法，基本都是在 ”梯度是一堆外积之和” 这件事上做文章。

0.2. 两个麻烦：噪声和形状

如果每一步都能拿到整个训练集上的精确梯度，最朴素的梯度下降其实已经能用了，麻烦在于实际训练里有两件事会捣乱。

一个是噪声。batch 只是训练集里很小的一部分，$g_t$ 等于真实梯度 $\nabla L(\theta_t)$ 加上一个随机误差，这个误差期望是 0，方差大致和 batch size 成反比，batch 越小，每一步的方向就越不靠谱。

另一个是形状。loss 曲面在不同方向上的弯曲程度可以差很多，一个简单的例子是：

一个沿 $x$ 方向很平、沿 $y$ 方向很陡的山谷。梯度下降的学习率受最陡的方向限制，超过 $2/\kappa$ 就会在 $y$ 方向发散，于是在平坦的 $x$ 方向上每一步只能挪一点点。曲率最大值和最小值之比 $\kappa$ 叫条件数，梯度下降需要的步数大致和它成正比。

上图是让几种方法在这个山谷上各走 50 步。左边是沿坐标轴放的山谷：GD 先在陡的方向上来回震荡，之后沿着谷底一点点往前挪，50 步之后离终点还剩起点距离的四成多；Adam 给每个坐标单独定步长，刚好能抵消沿坐标轴的形状，很快就到了终点附近（因为带了动量，在终点附近还冲过头绕了一下）；Newton 法一步到位。右边是把同一个山谷整体转 45°，起点也跟着转。GD 和 Newton 的轨迹和左边完全一样，只是整体转了一下，它们本来就不关心坐标轴朝哪；Adam 就不行了，逐坐标的缩放这时候和山谷的方向对不上，在起点附近来回折腾，50 步之后还剩六成多的距离，比 GD 还差。

换个坐标系，Adam 就从最好的变成了最差的。神经网络里的权重矩阵其实天然就有这个问题：对输入或者输出空间做一个旋转，网络表示的函数可以完全不变，逐坐标的优化器看到的东西却全变了。

0.3. NanoGPT speedrun Track 3

1. SGD 与动量

1.1. SGD

$g_t$ 是真实梯度的无偏估计，所以只要学习率按合适的方式衰减，SGD 在凸问题上能收敛到最优解。不过深度学习里基本没人按那些理论条件去衰减学习率。Track 3 的 baseline 是前 30% 恒定、后 70% 线性衰减到 0。schedule 对最终结果的影响很大，Track 3 上不少纪录都有一部分收益来自改 schedule，但它和优化器算是两件独立的事，这篇就不展开了。

1.2. 动量

$\beta \in [0, 1)$ 一般取 0.9 或 0.95。展开之后 $m_t = \sum_{s \le t} \beta^{t-s} g_s$，是过去所有梯度的指数加权和。改动很小，但同时干了三件事：

• 平均掉噪声。如果各步的噪声相互独立，把 $m_t$ 归一化成加权平均之后，它的噪声方差大约是单步的 $(1-\beta)/(1+\beta)$，$\beta = 0.95$ 时差不多相当于把 39 步的梯度平均在一起；
• 在一致的方向上加速。梯度一直指向同一个方向的话，速度会累积到单步的 $1/(1-\beta)$ 倍；
• 在来回震荡的方向上减速。上面山谷里 $y$ 方向的梯度每步变号，在 $m_t$ 里会互相抵消。

对二次函数，Polyak 证明了选好 $\eta$ 和 $\beta$ 之后，收敛需要的步数可以从和 $\kappa$ 成正比降到和 $\sqrt{\kappa}$ 成正比。实现上还有一种 EMA（指数滑动平均）的写法 $m_t = \beta m_{t-1} + (1-\beta) g_t$，两者只差一个常数倍，可以吸收进学习率里，Adam 用的是 EMA 形式，PyTorch 的 SGD 和 Muon 用的是前一种，同样叫 lr，两种写法在 $\beta = 0.95$ 时能差出 20 倍，读代码的时候还是要注意一下这个问题。

1.3. Nesterov

和普通动量相比，它在更新里多给了当前梯度一点权重，反应更快一些，Muon 默认用的就是这种。

1.4. Gradient noise scale

batch size 为 $B$ 时，梯度里噪声和信号的能量之比大约是 $B_{\mathrm{noise}}/B$。$B$ 远小于 $B_{\mathrm{noise}}$ 时，每一步的方向主要由噪声决定，batch 加倍差不多能让需要的步数减半；$B$ 远大于 $B_{\mathrm{noise}}$ 时，梯度已经很准了，再加大 batch 几乎省不了步数，$B_{\mathrm{noise}}$ 附近就是所谓的临界 batch size。训练中 $B_{\mathrm{noise}}$ 会慢慢变大，因为 loss 越低，真实梯度越小，而单个样本梯度的波动并不会同比例下降。

这个量对读 Track 3 的结果很重要。Track 3 的 batch 是固定的，所有优化器面对的噪声水平都一样；但自己做实验的时候不可能每个想法都跑一遍完整的 Track 3，经常要先在更小的模型、更小的 batch 上试，这时候噪声在梯度里占的比例会大很多，一些对付噪声的技巧在小 batch 上效果明显，换到大 batch 上就不一定了。

2. 逐坐标自适应：从 Adagrad 到 AdamW

2.1. Adagrad 与 RMSProp

这里的平方、开方和除法都是逐元素的。累积梯度一直很大的坐标步长变小，很少有梯度的坐标步长相对变大。

Adagrad 来自在线凸优化（online convex optimization），在那个框架下它的有一个很漂亮的 regret 界。论文里其实还给了一个 full-matrix 版本，用整个梯度外积的累积 $\sum_s g_s g_s^\top$ 而不只是它的对角线，更新为 $(\sum_s g_s g_s^\top)^{-1/2} g_t$。这个版本对大模型完全没法算，但 Shampoo 就是从这里出发的。

这就是 RMSProp。它一直没有正式发表，但用得非常广，引用的时候只能引一页课程 slides，也挺有意思的。

2.2. Adam 与 AdamW

$m_t$ 和 $v_t$ 分别是梯度一阶矩和（未中心化的）二阶矩的滑动估计。第三行是偏差修正：$m_0 = v_0 = 0$，训练刚开始时两个估计都偏小，除以 $1-\beta^t$ 把它们校正回来。常用超参是 $\beta_1 = 0.9$、$\beta_2 = 0.999$，大模型训练里 $\beta_2$ 一般取 0.95。每个参数要额外存 $m$ 和 $v$ 两个状态，fp32 下优化器状态是参数量的两倍。

Adam 为什么好用，我看到的解释大概有三种，它们之间并不完全一致：

范数取 $\ell_2$，解是 $g / \|g\|_2$，就是归一化的 SGD；取 $\ell_\infty$（每个坐标的绝对值都不超过 1），解是 $\mathrm{sign}(g)$，就是符号下降。换句话说，选什么范数，就是在选怎么衡量 ”一步走了多远”，更新的方向也就跟着定了。

大模型里用 Adam 还有两个实践上的细节。一个是 $\beta_2$：0.999 意味着 $v$ 大概记得最近一千步的梯度大小，梯度的尺度变化很快的时候（比如训练刚开始，或者 loss spike 之后），$v$ 跟不上，$\hat m / \sqrt{\hat v}$ 会突然变得很大，一般认为这是大模型训练里把 $\beta_2$ 调到 0.95 的主要原因。另一个是 $\epsilon$：对梯度本来就很小的坐标，分母里的 $\epsilon$ 可能比 $\sqrt{\hat v}$ 还大，这时 Adam 在这些坐标上就退化成了 SGD，所以 $\epsilon$ 的取值是会改变优化器行为的。Track 3 baseline 里 aux AdamW 的 $\epsilon$ 是 $10^{-10}$，比 PyTorch 默认的 $10^{-8}$ 小两个数量级，One-sided Shampoo 那个提交里 grafting 用的 Adam 更是设到了 $10^{-15}$，基本就是不想让 $\epsilon$ 起任何作用的意思。至于 $10^{-10}$ 和 $10^{-8}$ 在 Track 3 上到底能差多少，我还没有专门测过。

改动看起来很小，效果却明显更好，之后 AdamW 就成了大模型训练的默认选择，常见配置是 $\beta = (0.9, 0.95)$、weight decay 0.1，再加上 warmup 和 cosine 衰减。后面的 Muon 和 SOAP 基本都沿用了这种解耦的权重衰减。

2.3. Adafactor

状态从 $mn$ 降到 $m + n$。论文里还有几个配套的改动：去掉一阶动量进一步省显存、对更新做裁剪（update clipping）、按参数自身的大小决定相对步长、让 $\beta_2$ 随训练步数增加。T5 和 PaLM 都用过它（PaLM 用的是不做分解的版本）。

我觉得 Adafactor 在这条线上的位置有一些特别之处：它是第一个用上了” 参数是一个矩阵，有行有列” 这个结构的主流优化器，虽然出发点只是压缩二阶矩，和矩阵的几何没什么关系。有意思的是，几年之后 SOAP 的作者回过头发现，Shampoo 的一个变体其实就是在某个旋转过的坐标系里跑 Adafactor。

3. 曲率：Newton、Natural Gradient 与 K-FAC

3.1. Newton 法

回到 0.2 节的山谷，GD 慢是因为它对所有方向一视同仁，最自然的想法就是把曲率也考虑进来。在当前点附近把 loss 展开到二阶，记位移为 $d$，

$H$ 是 Hessian 矩阵。这个二次函数的最小值点在 $d = -H^{-1} g$，按 0.1 节 $\theta_{t+1} = \theta_t - \eta_t \Delta_t$ 的写法，就是取 $\Delta = H^{-1} g$、$\eta = 1$，这就是 Newton 法。对山谷来说 $H = \mathrm{diag}(1, \kappa)$，Newton 法一步就到最优点，而且不管山谷怎么转都一样，因为它对参数的线性变换是不变的（affine invariant）。

问题在于规模。一个 $10^8$ 参数的模型，Hessian 有 $10^{16}$ 个元素，存都存不下，更别说求逆。神经网络的 loss 也不是凸的，$H$ 可能有负特征值，直接用 $H^{-1}$ 会往上坡走；从一个 batch 估出来的 $H$ 噪声也很大。所以实际能用的二阶方法，基本都在做两件事：找一个半正定的曲率矩阵代替 $H$，再找一种结构化的近似让它算得起。

3.2. Gauss-Newton、Fisher 与 Natural Gradient

loss 是负对数似然的时候（比如语言模型的交叉熵），有一个天然的半正定替代品，Fisher 信息矩阵：

把 $\hat y$ 换成真实标签，得到的是经验 Fisher $\mathbb{E}[g g^\top]$，算起来方便很多，Adam 的 $v$ 就是它的对角线。但经验 Fisher 和真正的 Fisher 在很多情况下差别很大，拿它做曲率近似要小心¹³¹³。

3.3. K-FAC

约等号来自假设输入 $x$ 和输出梯度 $\delta$ 相互独立。$A \in \mathbb{R}^{n\times n}$ 是输入的二阶矩，$B \in \mathbb{R}^{m\times m}$ 是输出梯度的二阶矩。Kronecker 积的逆等于逆的 Kronecker 积，于是 Natural Gradient 的更新可以写回矩阵形式：

原本要对一个 $mn \times mn$ 的矩阵求逆，现在只需要对一个 $m \times m$ 和一个 $n \times n$ 的矩阵求逆，对 $4096 \times 4096$ 的层，就是从一个约 $1.7\times 10^7$ 维的问题变成两个 4096 维的问题。实现里 $A$ 和 $B$ 用滑动平均估计，每隔若干步重新求一次逆，再加上 damping 保证数值稳定。

K-FAC 在研究里影响很大，但据我所知在 LLM 预训练里用得不多，求逆的开销、对 damping 的敏感，以及实现的复杂度大概都是原因。

4. 把矩阵当作矩阵：Muon

4.1. 动机

前面的逐元素方法有一个共同点：它们把一个 $m\times n$ 的权重矩阵当成 $mn$ 个互不相关的数，可网络用 $W$ 的方式，是把它当成一个线性映射作用在输入上。0.2 节的右图其实已经把这个问题摆出来了，坐标系一转，逐坐标缩放的效果就全变了。

另一个观察来自梯度本身的形状。$G = \sum_i \delta_i x_i^\top$ 是一堆外积之和，实际训练里这种矩阵的奇异值往往衰减得很快，少数几个方向占了绝大部分能量。直接拿 $G$（或者它的动量）去更新，步长几乎全花在那几个主导方向上，其他方向基本学不动。

4.2. 算法

第一行是 Nesterov 动量，$U \Sigma V^\top$ 是 $\tilde M_t$ 的瘦 SVD。$U V^\top$ 在线性代数里叫极分解（polar decomposition）的正交因子，后面记作 $\mathrm{polar}(\cdot)$，它保留了原矩阵的奇异向量，也就是 ”往哪些方向走”，但把每个方向上的步长都拉平成 1。$s$ 是一个和矩阵形状有关的缩放系数，Jordan 的实现里取 $s = \sqrt{\max(1, m/n)}$，$m$ 是输出维度，$n$ 是输入维度。

Muon 只管网络中间的二维权重矩阵，embedding、最后的 LM head、RMSNorm 的 gain 和各种 bias 还是交给 AdamW，Track 3 里管这部分叫 aux AdamW。embedding 的输入是 one-hot 向量，每一步只有被用到的那几行有梯度，“整个矩阵对任意输入的作用” 这种看法对它不太适用；LM head 在 Jordan 的实验里用 Muon 也不如用 AdamW；一维的 gain 和 bias 则根本没有奇异值可言。所以实际跑起来的 ”Muon” 一直是 Muon 和 AdamW 的混合体，aux AdamW 这部分的超参调起来也挺费劲的，Track 3 #36 那个调好的 baseline，改的超参里一大半都是 aux AdamW 的。

4.3. Newton-Schulz 迭代

每个矩阵每一步都做一次 SVD，在 GPU 上太慢，也不适合 bf16。Muon 用 Newton-Schulz 迭代来近似 $U V^\top$：先把矩阵除以它的 Frobenius 范数，保证所有奇异值都不超过 1，然后反复做

这个迭代只用矩阵乘法，而且不改变奇异向量，每个奇异值 $\sigma$ 各自经历同一个标量映射 $\sigma \mapsto a\sigma + b\sigma^3 + c\sigma^5$。所以选系数的问题就变成了：选一个奇多项式，让它迭代几次之后，把 $(0, 1]$ 里的数都推到 1 附近。

左图是三组系数下的标量映射。经典的三次迭代 $(1.5, -0.5, 0)$ 在 1 附近收敛很快，但对很小的奇异值每次只能放大 1.5 倍，迭代 12 步之后，$10^{-2}$ 以下的奇异值还明显没到 1。Jordan 调出来的五次多项式 $(3.4445, -4.7750, 2.0315)$ 在 0 附近斜率很大，5 步就能把 $10^{-3}$ 量级的奇异值推到 0.5 左右，$3\times 10^{-3}$ 以上的基本都到了 1 附近；代价是它并不收敛到 1，在 $[10^{-2}, 1]$ 的输入上，5 步之后的输出在 0.68 到 1.13 之间来回摆。Jordan 的判断是这点误差对训练基本没有影响，换来的是更少的迭代次数，speedrun 比的是墙钟时间，这个取舍很合理。Track 3 比的是步数，算得慢一点无所谓，所以 baseline 换成了另一组系数 $(2, -1.5, 0.5)$，迭代 12 步，它在 0 附近斜率为 2，在 1 处精确收敛，$10^{-3}$ 以上的奇异值最后都是 1。

右图是一个合成的动量矩阵（$256 \times 1024$，奇异值按幂律衰减再加上噪声）经过两种迭代之后的奇异值，原始的奇异值跨了将近两个数量级，迭代之后基本都被拉到了 1 附近。按 Jordan 的估算，Newton-Schulz 带来的额外计算在典型的语言模型训练里不到 1%²¹²¹。

系数和缩放这些东西看起来只是实现细节，我在小模型上就被它坑过一次。一开始用的是一个常见的 Muon 库实现（Jordan 的五次系数、不加 weight decay、库自带的步长缩放），在这个设置下看到某个改动比 SOAP 还好 0.01 左右；后来把这些全部换成 Track 3 的约定重跑，收益只剩下五分之一，置信区间也跨过了 0。也就是说之前看到的 ”收益”，有一大半是在弥补库实现本身的问题。从那之后我的小模型实验都尽量和 Track 3 的约定对齐。

4.4. 为什么有效

Muon 为什么有效，我目前看到的主要有三种说法，彼此相容，只是角度不同。

第一种是谱范数下的最速下降¹⁵¹⁵。把 2.2 节的范数换成矩阵的谱范数（spectral norm，也就是最大奇异值 $\|\Delta\|_{\mathrm{op}}$），

输入 $x$ 的每个分量大小在 1 左右，也就是 RMS 为 1、$\|x\|_2 = \sqrt{n}$，那么一个所有奇异值都是 $\sigma$ 的更新 $\Delta W$ 最多能让输出变化 $\sigma\sqrt{n}$，换算成输出的 RMS 是 $\sigma\sqrt{n/m}$。想让输出的 RMS 变化和输入在同一个量级，就要取 $\sigma = \sqrt{m/n}$。Jordan 的 $\sqrt{\max(1, m/n)}$ 只在输出维度大于输入维度时放大，算是这个式子的一个变体。

第二种是把 Muon 看成没有记忆的 Shampoo，下一节讲 Shampoo 的时候再说。

第三种是 Jordan 原始博客里的直觉：SGD 和 Adam 产生的更新矩阵往往接近低秩，被少数方向主导，正交化之后，那些原本很小、但可能很重要的方向拿到了和主导方向一样的步长。这个说法很直观，但 ”可能很重要” 基本算纯猜。奇异值很小的那些方向，到底是被压制住的有用信号，还是就是噪声？正交化对两者是一视同仁地放大的。后面的不少工作，都是在追问这个问题。

4.5. 规模化与变体

5. 带记忆的两侧统计量：Shampoo

5.1. 算法

这两个矩阵的含义挺直观的。$G G^\top$ 的第 $(i, j)$ 个元素是 $G$ 第 $i$ 行和第 $j$ 行的内积，也就是第 $i$ 个和第 $j$ 个输出神经元收到的梯度有多相似，累积起来，$L$ 记录的就是梯度在输出空间里各个方向上的能量和相关性；$G^\top G$ 对输入空间做同样的事。左右各乘一个负幂次，可以理解成在输出和输入两个空间里分别做白化（whitening）：历史上梯度能量一直很大的方向被压下去，一直很小的方向被放大。至于为什么是 $-1/4$：full-matrix Adagrad 要的是那个大矩阵的 $-1/2$ 次幂，论文证明了它可以被 $L_t^{1/2} \otimes R_t^{1/2}$（乘上一个和梯度秩有关的常数）从上方控制住，代进去再取 $-1/2$ 次幂，每一侧就分到了 $-1/4$。

5.2. 和 Muon 的关系

如果去掉累积，只用当前这一步的梯度，那么 $L = G G^\top$，$R = G^\top G$。把 $G = U\Sigma V^\top$ 代进去：

这里要稍微小心一点：$G$ 不是方阵的时候，$L$ 和 $R$ 至少有一个是奇异的，上面的 $-1/4$ 次幂要在伪逆、或者加上 $\epsilon I$ 再让 $\epsilon \to 0$ 的意义下理解。所以没有记忆的 Shampoo 就是 Muon 的正交化，Muon 只是额外在动量上做这件事。Track 3 的 #22 干脆把 SpectralDescent、不带动量的 Muon 和 betas 取 (0, 0) 的 Shampoo 写成了同一行，8225 步。

5.3. 工程上的麻烦

Shampoo 一直没有像 Adam 那样普及，主要是工程上的原因：

• 求 $-1/4$ 次幂需要特征分解，要在 fp32 甚至 fp64 下做，数值上很敏感，代价是 $O(m^3 + n^3)$，所以一般每隔几十到上百步才更新一次预条件矩阵；
• 需要 $\epsilon$ 一类的 damping，结果对它比较敏感；
• 实际实现里常用 grafting，用 Adam 或 Adagrad 的更新大小配上 Shampoo 的更新方向，来稳定每一层的步长；
• 很大的矩阵需要分块，分布式训练需要专门的实现。

grafting 这一条我第一次看到的时候觉得有点奇怪，等于是承认 Shampoo 自己定不好每一层的步长，方向用 Shampoo 的，大小还得借 Adam 的。Track 3 上那个 Shampoo 的提交，打开配置一看也是开着 Adam grafting 的。我自己在小模型上跑 SOAP 一类的方法，fp32 的 torch.linalg.eigh 有时（可能 5%？）会报不收敛直接退出，最后只好在失败的时候退回 fp64 重算。

把这几类方法的开销放在一起（对一个 $m \times n$ 的矩阵）：

6. SOAP：在 Shampoo 的特征基里跑 Adam

6.1. 出发点

推导很短。设 $L = Q_L \Lambda_L Q_L^\top$，$R = Q_R \Lambda_R Q_R^\top$，旋转后的梯度为 $G' = Q_L^\top G Q_R$，那么

也就是在旋转后的坐标系里，第 $(i, j)$ 个元素被除以 $\sqrt{\lambda_{L,i}\,\lambda_{R,j}}$。而 $L$ 的第 $i$ 个特征值 $\lambda_{L,i}$ 正好等于（累积意义下的）$\sum_j G'^2_{ij}$，$\lambda_{R,j}$ 等于 $\sum_i G'^2_{ij}$，分别是旋转后梯度平方的行和与列和，也就是 Adafactor 用来重建二阶矩的 $r$ 和 $c$。两者只差一个标量：Adafactor 的 $\hat V$ 分母里还有一个 $\mathbf{1}^\top r = \operatorname{tr}(L)$，所以 Adafactor 的更新是这个 Shampoo 的 $\sqrt{\operatorname{tr}(L)}$ 倍。这个因子对同一个矩阵里的所有元素都一样，但会随层和训练步数变化，相当于每一层每一步换了一个学习率，方向上两者是一致的。

那既然 Shampoo 相当于在特征基里跑 Adafactor，为什么不在这个坐标系里跑效果更好的 Adam？这就是 SOAP。

6.2. 算法

记 $L$ 和 $R$ 的特征向量矩阵为 $Q_L$ 和 $Q_R$，SOAP 每一步做：

1. 用滑动平均更新 $L \leftarrow \beta L + (1-\beta) G G^\top$，$R \leftarrow \beta R + (1-\beta) G^\top G$；
2. 每隔 $f$ 步重新算一次 $Q_L$ 和 $Q_R$，原论文用特征分解，也可以用一步幂迭代加 QR 分解来近似；
3. 把梯度和动量旋转到特征基里：$G' = Q_L^\top G Q_R$，$M' = Q_L^\top M Q_R$；
4. 在旋转后的坐标系里跑 Adam：$V \leftarrow \beta_2 V + (1-\beta_2)\, G' \odot G'$，$N' = M' / (\sqrt{V} + \epsilon)$；
5. 旋转回去并更新：$W \leftarrow W - \eta\, Q_L N' Q_R^\top$。

直观上，旋转到特征基之后，梯度的各个坐标之间近似不相关，Adam 那种逐坐标处理的假设就更接近成立了，相当于先把 0.2 节右边那个转过的山谷转回来，再交给 Adam。另一个好处是稳定性：Shampoo 的预条件矩阵在两次更新之间是冻结的，而 SOAP 在两次更新特征基之间，Adam 的二阶矩还在一直适应，所以特征分解可以做得更稀疏，代价更低。当然在 Track 3 这种只看步数的设置下，这个好处就无所谓了，现在的纪录都是每一步更新特征基。

原论文在 3.6 亿和 6.6 亿参数的语言模型上报告，SOAP 比 AdamW 少用 40% 以上的迭代步数、35% 以上的 WallTime，比 Shampoo 在两项上各好 20% 左右³⁵³⁵。

6.3. 后续的分析

7. SOAP-Muon：先预处理，再正交化

7.1. 想法

SOAP 和 Muon 都好用，那把它们串起来呢？先用 SOAP 对动量做预处理，再把结果送进 Muon 的正交化：

7.2. 在 Track 3 上的演化

SOAP-Muon 是 Track 3 上推动纪录最多的一类方法，大致经历了这几步¹¹：

• #14（2026 年 5 月）：在 MLP 的两个矩阵上，Muon 正交化之前先做 SOAP 预处理，3150 步；
• #16：attention 的矩阵也加上 SOAP，并用一个 trust gate 控制，3125 步（README 里自己也标了，它相对 #14 的 p 值是 0.34）；
• #44（6 月）：所有隐藏层矩阵都用 SOAP-Muon，并且每一步都更新特征基（precondition_frequency = 1），同时调了 aux AdamW 的 $\beta_2$ 和初始化，2750 步。这一步比之前的纪录一下少了 100 步，而且是一边加东西一边删东西，前面好几个纪录里叠上去的 Contra-Muon、Circuit-Muon、Aurora 这些模块，在 #44 里被当作中性的模块删掉了；
• #45：在最后一步把权重往 150 步的 weight EMA 方向混一下，2720 步；
• #46：加上逐行的更新下限（RowFloor）和 Cautious Weight Decay（只在这一步本来就在让权重变小的坐标上做衰减），2690 步，是现在的纪录。

#44 和 #46 里 SOAP-Muon 的实现有几点值得注意：特征基每步只用一次幂迭代加 QR 来更新；Gram 矩阵用的是原始梯度，滑动平均系数 0.9；被预处理的是动量，并且用的是上一步建好的特征基；预处理的输出还会按某种矩阵范数重新缩放，再送进 Newton-Schulz。

#44 算是比较扎实的 SOTA，但 #44 之后这两步我是有点怀疑的。#45 和 #46 各自比上一个纪录少 30 步，按 Track 3 自己给的换算大约是 0.0014 的 loss，而按提交的 log 算，seed 之间的标准差在 0.001 左右。拿 README 里两两对比的显著性公式算一下（两个提交的平均 loss 之差，加上步数之差按每 100 步 0.0045 折算出的 loss，再除以 $\sqrt{1/n_1 + 1/n_2}$，同样以 0.004 为线），#45 相对 #44 是 0.0043，刚刚过线；#46 相对 #45 是 0.0035，其实没过（它作为纪录是有效的，因为它自己在 2690 步上显著低于 3.28，只是这些 log 还不足以说明它比 #45 更好）。README 对 #46 的说明也很坦白，好几个细节还没有被单独证明有用，其中最好笑的一个是 Muon 动量的 cooldown 被安排在 2700 到 2900 步，而纪录在 2690 步就验收了，也就是说这段 cooldown 根本没跑到。

自己做实验的时候也有比较多这种奇怪的 setting 导致的问题，反正体感上 #44 之后的几十步，基本都是靠这样一个个开关叠出来的 buff，单看每一个都很难说清楚有没有用。

7.3. 一些有代表性的 Track 3 结果

需要注意的是，不同提交的调参力度差别很大，很多结果还叠加了学习率 schedule、初始化、hyperball 这类和更新方向无关的改动，所以它算不上一个纯算法的排名。另外在 3300 步附近，少 100 步大约相当于 val loss 低 0.0045¹¹，表里几十步的差距对应的 loss 差其实很小。即便如此，大的趋势还是很清楚，加粗的三行就是三个 stage：逐元素的 Adam 明显落后；Muon 这类把矩阵当矩阵的方法是一个大台阶；在 Muon 之上再加 SOAP 的统计量，是第二个 stage。

8. 留给下一篇的问题

把这些方法按时间顺序读下来，一个感受是它们之间的关系比我原来想的紧得多。full-matrix Adagrad 做 Kronecker 近似得到 Shampoo，Shampoo 去掉记忆就是 Muon 的正交化；Adam 的二阶矩按行列分解得到 Adafactor，把 Adafactor 搬到 Shampoo 的特征基里再换回 Adam 就是 SOAP；最后 SOAP 和 Muon 又在 SOAP-Muon 里接到了一起。逐元素和矩阵这两条线，一直在互相借用对方的东西。

读的过程中形成了一些问题：

1. Adam 的 $v$ 到底在估计什么，是噪声的大小，还是某种曲率？full batch 下 Adam 依然更好，似乎说明不只是噪声。
2. Muon 把所有奇异值拉平成 1，奇异值很小的那些方向，是值得放大的信号，还是噪声？
3. K-FAC 用的是单个样本层面的输入和输出二阶矩，Shampoo 和 SOAP 用的是整个 batch 梯度的 Gram 矩阵，这两种统计量包含的信息一样吗？
4. MLP 的权重是 $4d \times d$ 这样很扁的矩阵，左右两侧统计量的维度差了四倍，哪一侧更重要？
5. Muon 最后一步的 $\mathrm{polar}(\cdot)$ 会把所有奇异值强行拉回 1，那 SOAP-Muon 里 SOAP 费了半天劲算出来的逐元素缩放，岂不是全被抹掉了？可实际上先做 SOAP 再正交化，比单独用 Muon 少了好几百步（虽然中间也混了别的改动）。polar 抹掉的只是奇异值，奇异向量还在，而 SOAP 的逐元素缩放恰恰会改变奇异向量，也就是说 $\mathrm{polar}(T(M)) \neq \mathrm{polar}(M)$。那 SOAP 到底把奇异向量往哪边掰了，为什么掰过之后会更好？

第二篇会尝试把这一篇里的方法放进一个统一的描述里，回答其中一部分，剩下的变成可以用实验检验的假说。另外，这一篇引用了不少 Track 3 社区的结果，有的有严格的多 seed 统计，有的只是单次运行，可信程度参差不齐，怎么在小模型上做可靠的对比，也许之后还可以再写一篇单独的文章。