为什么官方版Muon比MuP版多出一个max(1, ⋅)？

在文章《Muon优化器指南：快速上手与关键细节》中，我们罗列了Muon的几个版本，它们的区别是学习率的矩阵形状相关的缩放因子不同，其中“官方版（KellerJordan版）”只比“MuP版”多出了一个 max(1,⋅)max(1,⋅)的截断操作。本文就来专门讨论一下，这个截断操作是怎么来的呢？

几个版本[#](https://spaces.ac.cn/archives/11772#%E5%87%A0%E4%B8%AA%E7%89%88%E6%9C%AC)

Muon的更新规则，可以统一地写成

M t=β M t−1+G t W t=W t−1−η t(α msign(M t)+λ W t−1)

M t=W t=β M t−1+G t W t−1−η t(α msign(M t)+λ W t−1)(1)

几个不同的版本区别在于 α α，分别是：

α={1(朴素版)√max(1,d o u t/d i n)(KellerJordan版)√d o u t/d i n(MuP版)0.2×√max(d o u t,d i n)(Moonlight版) α=⎧⎩⎨⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪⎪1 max(1,d o u t/d i n)−−−−−−−−−−−−−√d o u t/d i n−−−−−−−√0.2×max(d o u t,d i n)−−−−−−−−−−−√(朴 素 版)(KellerJordan 版)(MuP 版)(Moonlight 版)

其中矩阵 W∈R d i n×d o u t W∈R d i n×d o u t，代表线性层 y=x W y=x W 的训练参数，其中输入 x∈R d i n x∈R d i n 是行向量。

本文主要关心“KellerJordan版”和“MuP版”，前者在后者的基础上多加了个 max(1,)max(1,)。按照我们在《高阶MuP：更简明但更高明的谱条件缩放》、《MuP之上：2. 线性层与最速下降》等文章的分析，在MuP相关的谱条件约束下，最速下降应该就是MuP版Muon才对，多出来的 max(1,⋅)max(1,⋅)该如何解释呢？

特征增量[#](https://spaces.ac.cn/archives/11772#%E7%89%B9%E5%BE%81%E5%A2%9E%E9%87%8F)

简单起见，下面的讨论都省掉下标 t t。不失一般性，我们假设动量 M M 是满秩的，那么 Φ=msign(M)Φ=msign(M)的奇异值全为1，那么当 d i n≤d o u t d i n≤d o u t 时，Φ Φ⊤=I d i n Φ Φ⊤=I d i n，当 d i n>d o u t d i n>d o u t 时，Φ⊤Φ=I d o u t Φ⊤Φ=I d o u t。

记 Δ W=η α Φ Δ W=η α Φ，我们要做的事情，就是寻找 α α 与 d i n,d o u t d i n,d o u t 的关系。从《为什么我们偏爱各向同性？基于最速下降的理解》中我们知道，参数实际上只是模型的副产物，特征层面的变化可能会更加本质。将 Δ W Δ W 转换到特征层面是 Δ y=x Δ W=η α x Φ Δ y=x Δ W=η α x Φ，那么‖Δ y‖R M S=α‖x Φ‖R M S∥Δ y∥R M S=α∥x Φ∥R M S。

接下来要分情况讨论。首先是当 d i n≤d o u t d i n≤d o u t 时，Φ Φ 可以写成 U[I d i n,0 d i n×(d o u t−d i n)]V⊤U[I d i n,0 d i n×(d o u t−d i n)]V⊤的形式，其中 U∈R d i n×d i n,V∈R d o u t×d o u t U∈R d i n×d i n,V∈R d o u t×d o u t 都是正交矩阵，那么

‖Δ y‖R M S=η α‖x U[I d i n,0 d i n×(d o u t−d i n)]V⊤‖R M S=η α‖x U[I d i n,0 d i n×(d o u t−d i n)]‖R M S=η α‖[x U,0 d o u t−d i n]‖R M S=η α√d i n d o u t‖x U‖R M S=η α√d i n d o u t‖x‖R M S

∥Δ y∥R M S=====η α∥∥x U[I d i n,0 d i n×(d o u t−d i n)]V⊤∥∥R M S η α∥∥x U[I d i n,0 d i n×(d o u t−d i n)]∥∥R M S η α∥∥[x U,0 d o u t−d i n]∥∥R M S η α d i n d o u t−−−−√∥x U∥R M S η α d i n d o u t−−−−√∥x∥R M S(2)

注意全程都是等号，所以我们只需要设 α=√d o u t/d i n α=d o u t/d i n−−−−−−−√，就可以让“每一个”Δ y Δ y 的RMS都是 η‖x‖R M S η∥x∥R M S，即所有Token的相对更新幅度一致。

各向同性[#](https://spaces.ac.cn/archives/11772#%E5%90%84%E5%90%91%E5%90%8C%E6%80%A7)

很遗憾，对于第二种情况 d i n>d o u t d i n>d o u t，上述“完全一致”的目标并无法实现。具体来说，此时 Φ Φ 的SVD要写成 U[I d o u t 0(d i n−d o u t)×d o u t]V⊤，于是

‖Δ y‖R M S=η α‖x U[I d o u t 0(d i n−d o u t)×d o u t]V⊤‖R M S=η α‖x U[I d o u t 0(d i n−d o u t)×d o u t]‖R M S=η α‖(x U)[:d o u t]‖R M S

x U 是一个 d i n 维向量，d i n>d o u t，所以(x U)[:d o u t]只是截取了 x U 的前 d o u t 个维度去算RMS，这时候它的RMS是不确定的，最大可以去到√d i n/d o u t‖x‖R M S（Worst Case），最小可以是0。

我们知道正交矩阵不改变RMS，所以‖x U‖R M S=‖x‖R M S，当 x 的分布足够各向同性时，我们可以认为这代表着 x U 每个分量的平均尺度都是‖x‖R M S，于是取前 d o u t 个分量去RMS，平均来说也约等于‖x‖R M S，即‖Δ y‖R M S≈α‖x‖R M S，因此只需要取 α=1，就可以实现跟上一节类似的效果。

各向异性[#](https://spaces.ac.cn/archives/11772#%E5%90%84%E5%90%91%E5%BC%82%E6%80%A7)

将前两节的结果综合起来，我们就得到

α=√max(1,d o u t d i n)

这正是KellerJordan版Muon出现的 max(1,⋅)。

然而，上一节的结论依赖于输入 x 足够各向同性的假设，这在训练早期或许近似成立，但随着训练的推进，特征的分布会逐渐各向异性，集中在让‖Δ y‖R M S 最大化的“Worst Case”上面，此时平均近似‖Δ y‖R M S≈η α‖x‖R M S 就不够准确了，反而是最大值 η α√d i n/d o u t‖x‖R M S 会更加准确。

这种情况下，让‖Δ y‖R M S≈η‖x‖R M S 的 α 是√d o u t/d i n，跟 d i n≤d o u t 时的结论一致，重新得到MuP版结果。也就是说，对于训练的中后期，MuP版Muon更加科学。对于这种不一致性，我们有两种策略：一是始终用MuP版Muon，这会稍微降低一点初期的收敛速度，但毕竟中后期才是训练的“重中之重”；二是将缩放因子改为

α=√max(τ t,d o u t d i n)

其中 τ t 从1到0单调衰减，这样实现了从KellerJordan版到MuP版的渐变，代价是多了一个Schedule要调。

文章小结[#](https://spaces.ac.cn/archives/11772#%E6%96%87%E7%AB%A0%E5%B0%8F%E7%BB%93)

本文主要从“特征增量”均匀性的角度解释了KellerJordan版中 max(1,⋅)的来源。

_转载到请包括本文地址：[https://spaces.ac.cn/archives/11772](https://spaces.ac.cn/archives/11772 "为什么官方版Muon比MuP版多出一个max(1, ⋅)？")_

_更详细的转载事宜请参考：_[《科学空间FAQ》](https://spaces.ac.cn/archives/6508#%E6%96%87%E7%AB%A0%E5%A6%82%E4%BD%95%E8%BD%AC%E8%BD%BD/%E5%BC%95%E7%94%A8 "《科学空间FAQ》")