Reading

Muon


Fanyi Pu · GPT-6 Astra

· Updated

如果第一次听说 Muon,可以先把它理解成一种专门处理神经网络二维权重矩阵的 optimizer (Jordan et al., 2024)。它不改变模型结构、loss 或 backpropagation;forward 和 backward 仍然照常进行。改变的是 backward 之后的 optimizer step:算出的梯度最终怎样变成权重更新。

但梯度本身还没有完整回答“应该走哪一步”。它能用一阶近似比较不同更新谁更可能降低 loss,却没有规定多大的更新才算同样大。对向量参数,我们常把 Euclidean norm 当作默认尺度;对 linear layer 的矩阵 W,还可以直接关心它作为映射 y=Wx 时,一步更新最多会让输出特征改变多少。选择不同的尺度,就会设计出不同的 optimizer。

因此本文不从背诵 Muon 的流程开始,而是假设我们正在自己设计它,依次回答三个问题:

  1. mini-batch 梯度有噪声,先怎样得到比较稳定的更新意图?
  2. 如果限制一次矩阵更新对特征的最大影响,怎样在这个限制内尽量降低 loss?
  3. 理想答案需要 SVD,训练时怎样只用 GPU 擅长的矩阵乘法逼近它?

回答完这三个问题,momentum、polar factor 和 Newton–Schulz 就不再是突然出现的名词,而是分别解决一个具体困难。Muon 的名字 MomentUm Orthogonalized by Newton–Schulz 正好概括了最终得到的主要步骤。

下面先看它在 Transformer 训练中的位置。Muon 通常负责隐藏层里的 attention projection 和 MLP 权重,embedding、LM head、bias 与 normalization 参数仍交给 AdamW;θt 表示整个模型在第 t 步的参数。

Muon 在一次 Transformer 训练 step 中的位置一个 mini-batch 经过参数为 θ_t 的模型 forward,得到 loss;backward 算出各参数的梯度。Optimizer step 按参数用途分组:隐藏层的 attention projection 和 MLP 矩阵由 Muon 更新,embedding、LM head、bias 和 normalization 参数由 AdamW 更新。两组更新共同组成下一步的模型参数 θ_{t+1},再用于下一个 mini-batch 的 forward。两条优化器分支属于同一个模型。Mini-batch输入与训练目标模型 · forward
θt
Loss
Lt
Backward
∇θLt
按参数用途分组Muon隐藏层 attention projection与 MLP 权重矩阵AdamWEmbedding、LM headBias、normalization 参数更新后的同一个模型
θt+1
下一步,使用新 batch
左右滑动查看完整示意图
实线跟随当前 batch 的计算,虚线把更新后的模型送到下一步。Muon 和 AdamW 分别更新同一模型中的不同参数组;下面会放大左侧的 Muon 分支。

先决定怎样衡量“走了一步”

设 W∈Rm×n,当前 mini-batch 经过 backward 后得到梯度 Gt=∇WLt。如果打算加入更新 ΔW,loss 在当前位置附近满足

L(W+ΔW)≈L(W)+⟨Gt,ΔW⟩F,⟨A,B⟩F=tr⁡(A⊤B).

于是 ⟨Gt,ΔW⟩F 越负,一阶近似下这一步越有利。但如果不给 ΔW 限制,把它无限放大就能让这个线性式无限减小;所以设计更新方向之前,必须先定义“一步最多走多远”。

对 y=Wx,权重改变 ΔW 会带来 Δy=ΔWx,而

maxx≠0‖Δy‖2‖x‖2=‖ΔW‖2.

右边是 ΔW 的 spectral norm,也就是最大奇异值。因此,约束 ‖ΔW‖2≤η 有一个直接含义:无论输入朝哪个方向,一步更新造成的输出变化都不超过输入长度的 η 倍。 这不是衡量矩阵步长的唯一方式,但它和 linear layer 对特征的作用直接相连。

还有一个实际问题:单个 mini-batch 的 Gt 很嘈杂,而后面要做的变换会主动抬高一些较弱方向。与其先把每个 noisy gradient 的强弱抹平再平均,不如先汇总几步的证据。Muon 因此先使用普通 momentum:

Mt=βMt−1+(1−β)Gt,M−1=0.

Mt 仍然和 W 同形状。它在时间上平滑梯度,接下来再把这个较稳定的矩阵当作“希望下降的方向”,求解

min‖ΔW‖2≤η⟨Mt,ΔW⟩F.

先用一个刻意简化的 Mt 看这个约束会把我们带到哪里。暂时省略时间下标,令

M=diag⁡(8,2,0.5).

把三个坐标轴上的单位输入依次记作 e1,e2,e3,就有 Me1=8e1、Me2=2e2、Me3=0.5e3。这三个方向对下降的贡献强弱不同,但 spectral-norm budget 允许我们在每个方向上最多走 η。为了让内积尽可能负,最自然的选择不是继续保留 8:2:0.5 的比例,而是在三个方向都用满预算:

ΔW=−ηdiag⁡(1,1,1).

如果只把 M 除以一个整体 norm,三个数会一起缩小,16:4:1 的比例不变,较弱方向分到的步长预算仍然很少。把三个值都变成 1 并非凭空规定的一条 heuristic;在这个例子里,它正是 spectral-norm 约束下的一阶最优解。这些“方向强度”的正式名称是 singular values(奇异值)。

从约束推出 polar factor

在刚才的对角例子中,矩阵方向恰好就是三个坐标轴。一般矩阵会混合坐标,但 compact SVD 仍能找出它真正作用的方向。暂时省略时间下标:

M=UΣV⊤=∑i=1rσiuivi⊤,r=rank⁡(M),σi>0.

r 是矩阵的 rank。vi 是输入空间里的单位方向,ui 是对应的输出方向,两组向量各自正交。等式 Mvi=σiui 说的是:输入沿着 vi 变化时,M 会把它转到 ui,强度为 σi。所以 SVD 不是为了给算法增加一层抽象,而是在一般坐标中重新找到对角例子里的三根轴。

现在回到前面的约束问题。把更新写成 ΔW=−ηD,那么 ‖D‖2≤1,我们需要让 ⟨M,D⟩F 尽可能大。对任意这样的 D,

⟨M,D⟩F=∑i=1rσiui⊤Dvi≤∑i=1rσi,

因为 D 在任何单位输入方向上的输出长度都不超过 1。怎样同时取到每一项的上界?让 D 把每个 vi 都映到对应的 ui 即可:

D=UV⊤=∑i=1ruivi⊤.

这个 D 就叫 M 的 polar factor:

Polar⁡(M)=UV⊤,ΔW⋆=−ηPolar⁡(M).

于是“把每个非零奇异值变成 1”不是先规定一条 heuristic、再寻找理由;它是 spectral-norm budget 下把每个有下降信号的奇异方向都用满后自然得到的结果 (Bernstein & Newhouse, 2024; Su, 2025)。如果改用 Frobenius-norm budget,最优方向会是 −M/‖M‖F,所有奇异值只按同一比例缩放。选择怎样衡量一步,决定了怎样处理梯度。

对于前面的 diag⁡(8,2,0.5),U,V 都是单位矩阵,所以 polar factor 正是 I3。一般矩阵的 UV⊤ 也会保留每一对 vi→ui,只是去掉原来的 σi。

这个变换也可以直接从几何上看。矩阵 M 把输入空间的单位球映成椭球:vi 指向输入主轴,ui 指向对应的输出主轴,σi 就是这条输出半轴的长度。UV⊤ 保留同一组 vi→ui 对应关系,但去掉 Σ 带来的不均匀伸缩,因此把非零子空间里的单位球仍映成单位球。

下图截取例子中最大、最小奇异方向张成的二维平面。为了在同一尺度下比较形状,上路先把 M 除以最大奇异值;这不会改变 16:1 的半轴比例,也不会改变它的 polar factor。

同一个单位圆经过原矩阵与理想 polar factor 后的几何对比取最大和最小奇异方向张成的二维截面。上路中,归一化后的原矩阵把输入单位圆映成横纵半轴比例为 16 比 1 的细长椭圆;下路中,polar factor 把相同的输入单位圆映成单位圆。两条映射都把输入主轴 v_1 和 v_3 对应到输出主轴 u_1 和 u_3,区别只在输出半轴长度分别为 1 与十六分之一,或 1 与 1。
M/‖M‖2
Q=Polar⁡(M)
u1
u3
u1
u3
(1,1/16)
(1,1)
input unit circleinput unit circle
左右滑动查看完整示意图
为比较形状,上路先除以最大奇异值;正的整体缩放不改变 polar factor。原矩阵与 polar factor 保留同一组输入—输出主轴对应关系,但前者产生不均匀伸缩,后者把所有非零半轴长度设为 1。

这也解释了 Muon 名字中的 orthogonalized:它描述的是 polar 处理后的更新矩阵。加上这个更新之后,权重 W 本身不必正交,也不是在要求相邻两个训练 step 的更新互相垂直。对于满秩矩形矩阵,只能有 O⊤O=In(m≥n)或 OO⊤=Im(m≤n),所以更准确地说是 semi-orthogonal。

如果 M 秩亏,上面的 compact SVD 只保留非零奇异值,UV⊤ 是一个 partial isometry,不能声称它的所有行或所有列都正交。多项式迭代也不会凭空生成缺失的方向:在精确算术下,零奇异值始终为零。

和 sign、向量归一化有什么关系?

苏剑林在《Muon 优化器赏析》(Su, 2024) 中将这个 polar 变换记作 msign。对非零标量 z,sign⁡(z)=z/z2;矩阵也有类似的表达式:

Polar⁡(M)=M(M⊤M)−1/2.

这里的逆平方根作用于矩阵的特征值;遇到零特征值时采用伪逆约定,将对应值保留为零。因此它一般不同于逐元素取 sign。两个特例能说明矩阵形状的意义:

Polar⁡(diag⁡(z1,…,zn))=diag⁡(sign⁡(z1),…,sign⁡(zn)),Polar⁡(v)=v‖v‖2,v∈Rn×1, v≠0.

同一组数值,作为对角矩阵中的独立缩放系数,结果是逐元素取符号;作为一个列向量,结果是整体归一化。参数在计算中扮演什么角色,会影响合适的更新方式。这里的 msign 指基于奇异值的变换,不是一般数值线性代数中基于特征值定义的 matrix sign function。

这个推导给出了 Muon 的局部设计原则,不是效果保证。如果较小奇异值对应的方向也包含可靠的学习信号,用满这些方向的预算可能比让最大方向主导更新更有效;但它们也可能只是噪声。Momentum 只能帮助平滑信号,不能保证 polar 处理后一定泛化更好。整个网络的特征还会随其他层一起改变,实际 Muon 又只做有限步近似,因此不能把单层的一阶问题当成完整训练动力学。

与此相比,AdamW (Loshchilov & Hutter, 2019) 保存一阶、二阶矩估计,并按参数坐标逐元素缩放。它选择的是另一种更新几何,不会显式把同一个权重矩阵分解成输入—输出方向。Muon 调整的是由整个矩阵共同决定的奇异方向,不是逐元素取 sign,也不是把所有矩阵元素改成一样大。

用 Newton–Schulz 避免 SVD

到这里,理想更新已经设计出来了,但直接实现会卡在 SVD:训练的每个 step 都要对许多权重矩阵分解,代价太高。我们需要一种不显式求 U,Σ,V,却能保留奇异向量并把奇异值推向 1 的运算。

先把行数较多的矩阵转置,使下面的 X 满足行数不多于列数;最后再转置回来。暂且假设它满行秩,那么目标 polar factor 应满足

XX⊤=I.

这个等式也告诉我们当前还差多少。令

Rk=I−XkXk⊤.

如果某个奇异值是 σ,Rk 在对应左奇异方向上的值就是 1−σ2:轴太短时为正,太长时为负。把 Rk 左乘回 Xk,就能在不旋转奇异方向的前提下,对短轴做正修正、对长轴做负修正。因此可以先尝试

Xk+1=Xk+αRkXk.

沿一个奇异方向看,这等价于

fα(σ)=σ+α(1−σ2)σ.

σ=1 自动是固定点。我们还希望固定点附近的一阶误差项直接消失,因此令 fα′(1)=1−2α=0,得到 α=1/2,也就推导出经典的三次 Newton–Schulz 迭代:

Xk+1=32Xk−12XkXk⊤Xk.

现在才需要选择起点。理想 polar factor 不受正的整体缩放影响,所以先取

X0=M‖M‖F+ϵ.

因为最大奇异值不超过 Frobenius norm,X0 的奇异值都不超过 1,正好落在上述修正会把非零短轴向 1 推动的区间。将 Xk=UΣkV⊤ 代入矩阵式,可以核对

Xk+1=U(32Σk−12Σk3)V⊤.

这说明矩阵乘法确实在保持 U,V 的同时,对每个奇异值分别应用同一个标量函数。SVD 只用来解释发生了什么,实现并没有算它。

经典三次式在原点的斜率只有 f′(0)=1.5,很小的奇异值经过少数几步仍可能很小。如果训练时只愿意支付固定的几次矩阵乘法,我们自然会问:能不能用更高次的多项式,换取零附近更大的放大倍率?把同类修正扩展到五次,可以用 Gram matrix 整理成

Ak=XkXk⊤,Xk+1=aXk+(bAk+cAk2)Xk,f(σ)=aσ+bσ3+cσ5.

这里 Ak 是 Gram matrix:它保留左奇异方向,并把每个奇异值变成 σ2。多项式 aI+bAk+cAk2 因而能根据 σ2 为每个方向算出不同增益,最后乘回 Xk。这正好满足“只改轴长、不显式寻找轴”的要求。

Muon 的参考实现选择

(a,b,c)=(3.4445,−4.7750,2.0315),

通常重复 5 步。原作者用 ad-hoc gradient-based search 选择这组系数,目标是在有限步预算下尽量增大零附近的斜率 (Jordan et al., 2024)。这是一个有意的工程取舍:f′(0)=a=3.4445,小奇异值起初增长更快;但 f(1)=0.701,所以 1 甚至不是固定点。它并不让所有奇异值严格收敛到 1,而是快速送进相近的范围。实际输出更像 UΣ~V⊤,不能把“多做几步”简单理解为“更接近精确 polar factor”。

一次 Newton–Schulz 迭代怎样用 Gram matrix 计算方向增益上半部分展示实际矩阵计算:从 X_k 计算 A_k 等于 X_k X_k 转置,再计算多项式增益矩阵 P_k 等于 aI 加 bA_k 加 cA_k 平方,最后左乘 X_k 得到 X_{k+1}。下半部分沿单个奇异方向展开同一运算:先平方奇异值,代入增益多项式,再乘回原奇异值。示例第一步中,三个奇异值和增益分别为 0.968 乘 0.753 得 0.729;0.242 乘 3.172 得 0.768;0.061 乘 3.427 得 0.207。Matrix computation · no SVD requirediterate
Xk
Gram matrix
Ak=XkXk⊤
polynomial gain
Pk=aI+bAk+cAk2
next iterate
Xk+1=PkXk
The same step along singular direction iaxis lengthsquare itcompute gainmultiply back
σi(k)
zi=(σi(k))2
gi=a+bzi+czi2
σi(k+1)=giσi(k)
First iteration in the running example
σ1:0.968×0.753⟶0.729
gain < 1 · shrink
σ2:0.242×3.172⟶0.768
gain > 1 · expand
σ3:0.061×3.427⟶0.207
gain > 1 · expand
左右滑动查看完整示意图
上半部分是实现真正执行的矩阵乘法,不需要求出 SVD;下半部分只是用 SVD 坐标解释这些乘法。Gram matrix 将奇异值变成平方,矩阵多项式据此给每个方向算出增益,再乘回原矩阵:第一步压低最强方向,同时快速放大两个较弱方向。

上图把一次实际的五次迭代拆成了“测量轴长—计算增益—乘回原矩阵”。把相同操作重复五次,三个奇异值便沿着下面的曲线变化。曲线可能越过 1 再折返,而不是把椭球单调地磨成单位球。

五次多项式如何改变每个奇异值沿用同一个对角矩阵,横轴是一次更新内部的迭代次数 0 到 5,纵轴是各个原始奇异方向上的值。小值被迅速放大,之后围绕 1 上下变化,并非单调收敛。每步数值见图后的折叠表格。One update · five polynomial iterations0.00.51.01.5012345
k
σ1
σ2
σ3
左右滑动查看完整示意图
实线、长虚线和点线分别跟踪三个原始奇异方向;圆点表示实际迭代值,连线仅用于跟踪。虚线水平线是理想目标 1。这是矩阵内部的数值迭代,不是训练 loss 曲线。
示例每步的奇异值

从对角矩阵的奇异值 8、2、0.5 开始,先除以 Frobenius norm 加 1e-7,再使用正文的五次多项式。表格保留三位小数,各列跟踪原始方向。

迭代次数σ1σ2σ3
00.9680.2420.061
10.7290.7680.207
21.0791.0250.673
30.6890.6871.143
41.1271.1290.771
50.7400.7441.021

极小的非零奇异值在有限步后仍可能很小。苏剑林在 Newton–Schulz 专文 (Su, 2025b) 中讨论了下一步可以怎样继续设计:每一步采用不同的 (ak,bk,ck),根据前一步剩余的奇异值区间安排后续变换。同样的迭代步数不必共享系数;固定系数五步法只是一个具体的速度—精度选择,而近似 polar 更准确也不自动意味着训练效果更好。

前面的转置还有计算上的理由:它让 Ak=XkXk⊤ 的边长为 min(m,n),临时方阵更小。矩阵乘法可以在 BF16 下进行;上图用 FP64 演示标量递推,实际低精度结果会有差异。如果输入秩亏,σ=0 在这些奇函数下始终为零;迭代不会凭空创造缺失方向。

被 polar 抹掉的尺度怎样补回来

Polar factor 决定了更新的形状,也刻意抹掉了 M 原有的整体大小。若直接令 ΔW=−ηO,learning rate η 可以解释成 spectral norm 意义下的步长;但换一个形状的权重矩阵,同样的 O 会有不同的逐元素尺度。把 Muon 用到宽窄各异的层、再与 AdamW 参数组共用训练配置时,我们还必须决定想让哪个量保持一致。

先算出未经额外缩放的基准。对满秩矩阵的理想 polar 结果,令 r=min(m,n),因为共有 r 个值为 1 的奇异值,

‖O‖F2=r,RMS⁡(O)=‖O‖Fmn=1max(m,n).

如果直接使用 O,矩阵越大,平均每个元素的更新反而越小。假如设计目标是让理想更新的元素 RMS 固定为某个值 ρ,上式直接给出

sm,n=ρmax(m,n).

Moonlight (Liu et al., 2025) 取经验值 ρ=0.2,也就是 sm,n=0.2max(m,n),以便和 AdamW 的更新尺度对齐。这里的 0.2 不是 AdamW update RMS 的普适定理;有限步近似和秩亏也会使实际 RMS 偏离目标。

这不是唯一合理的设计目标。例如,也可以要求模型宽度改变时,特征扰动的 RMS 尽量保持一致。苏剑林的《Muon 优化器指南》(Su, 2025c) 将常见选择整理为四种约定。沿用本文 y=Wx 的记号,m 是输出维度,n 是输入维度:

约定sm,n缩放方式
朴素版1直接使用近似 polar 结果
KellerJordan 版max(1,m/n)输出维度大于输入维度时放大更新
MuP 版m/n按输出与输入的维度比缩放
Moonlight 版0.2max(m,n)将理想更新的元素 RMS 对齐到 0.2

这些公式不是对同一个目标的四种精度,而是在回答“什么应该跨层保持不变”的不同问题,所以同一个 learning rate 数值不能脱离缩放约定比较。PyTorch 的 Muon 文档 (PyTorch Contributors, n.d.) 将 KellerJordan 版和 Moonlight 版分别称为 original 和 match_rms_adamw。维度比要按 linear layer 的实际输入、输出判断:如果代码写成行向量形式的 x @ W,权重的存储形状就是 [in,out],不能照搬 PyTorch Linear 的轴顺序。为降低 Newton–Schulz 成本而临时转置矩阵,也不应改变缩放因子。

MuP 版与官方版的 max(1, ·) 分别在控制什么?

MuP 关注模型宽度改变时的更新尺度。对 v∈Rd,RMS⁡(v)=‖v‖2/d。用这个量衡量每个特征分量的变化,有

maxx≠0RMS⁡(ΔWx)RMS⁡(x)=nm‖ΔW‖2.

对非零矩阵的理想 polar 更新,ΔW=−ηsm,nO 且 ‖O‖2=1。取 sm,n=m/n,上式就等于 η,不随输入、输出维度改变。这给出了 MuP 版的尺度直觉;它控制特征扰动,Moonlight 版则对齐更新矩阵本身的元素 RMS。

苏剑林在 2026 年的文章 (Su, 2026) 中进一步用各向同性输入解释官方版的截断,下面假设待处理的矩阵满秩。当 m<n 时,理想 O 会丢掉输入在其行空间之外的分量;若固定 O,并假设输入满足 E[xx⊤]=In,就有

E[RMS⁡(Ox)2]=tr⁡(OO⊤)m=1=E[RMS⁡(x)2].

按这种平均尺度,取 sm,n=1 就够了;当 m≥n 时,相同计算则给出 sm,n=m/n,合起来正是官方版的 max(1,⋅)。这是一种依赖输入分布的解释。训练中输入与更新方向可能相关,不能据此认定该缩放在所有阶段都更优。

把设计装回完整的一步

现在各个零件都有了来由:momentum 先汇总跨 mini-batch 的下降信号;polar factor 在 spectral-norm budget 内选择矩阵方向;Newton–Schulz 用有限次矩阵乘法近似它;sm,n 再赋予 learning rate 选定的跨层尺度。

参考实现还在普通 momentum 上使用 Nesterov-style 组合。普通 Mt 更偏向过去的平均方向;在进入非线性的 polar 处理前,再混入当前梯度,可以让本步信号更直接地影响待处理矩阵。采用本文的 EMA 记法,令

Ht=βMt+(1−β)Gt.

不用 Nesterov 时只需取 Ht=Mt。下面放大训练总览中的 Muon 分支,虚线状态和实线计算也区分了两个时间尺度:t 是外层训练 step,k 是同一个 optimizer step 内部的 Newton–Schulz 迭代。

一个权重矩阵的 Muon 更新:历史状态、内部迭代与 weight decay以默认的 Nesterov-style momentum 为例:上一步保存的 M_{t-1} 和当前梯度 G_t 共同计算 M_t,M_t 保存给下一个训练 step,同时与 G_t 组合成 H_t。H_t 经过 Frobenius 归一化得到 X_0,在当前训练 step 内依次计算 X_1 到 X_5,输出近似 polar 结果 O_t;乘以 learning rate η_t 和形状系数 s_{m,n} 得到要减去的更新。另一条路径从原权重 W_t 直接计算 decoupled weight decay,得到 (1−η_tλ)W_t,再减去缩放后的更新得到 W_{t+1}。保存的 momentum 不经过 Newton–Schulz;weight decay 不经过 Newton–Schulz,也不乘形状系数。历史 momentum
Mt−1
更新 momentum
Mt=βMt−1+(1−β)Gt
保存到下一步
Mt
来自 backward
Gt=∇WLt
Nesterov 组合
Ht=βMt+(1−β)Gt
Newton–Schulz · 本次更新内部的 5 次迭代
X0=Ht‖Ht‖F+ϵ
X0
X1
X2
X3
X4
X5
Ot
当前权重
Wt
Weight decay
(1−ηtλ)Wt
缩放更新
ηtsm,nOt
+
+
−
Wt+1
左右滑动查看完整示意图
实线是本次更新的数据流,虚线表示跨训练 step 保存的 momentum。浅色框里的五次迭代只产生本步更新;下方的 weight decay 路径直接作用于原权重,最后减去右侧的缩放更新。

沿着图读,一个矩阵的完整更新是:

  1. backward 得到 Gt=∇WLt;
  2. 更新并保存 Mt=βMt−1+(1−β)Gt;
  3. 用 Ht=βMt+(1−β)Gt 构造本步要处理的方向;
  4. 归一化 Ht,在本步内部执行 K 次 Newton–Schulz,得到 Ot;
  5. 按矩阵形状缩放 Ot,再与 decoupled weight decay 合并:
Wt+1=(1−ηtλ)Wt−ηtsm,nOt.

Weight decay 的目的在于收缩当前权重,不是在表达这次 mini-batch 的下降方向,因此它直接作用于 Wt,不送进 momentum 或 Newton–Schulz,也不乘 sm,n。常用起点是 β=0.95、K=5。有些实现用不带 (1−β) 的 momentum 累加式;配套的 Nesterov 写法也要一起对应,不能只替换其中一行。正的整体缩放会被 Newton–Schulz 前的归一化消去,所以忽略 ϵ 和舍入时,两套配套写法可以等价。

Weight decay 为什么能限制权重谱范数?

《为什么我们选择尝试 Muon?》(Su, 2025a) 还给出了一个简单的范数界。设 λ>0、0≤ηtλ≤1,矩阵形状和缩放 sm,n 固定。对理想 polar 更新,由三角不等式得到

‖Wt+1‖2≤(1−ηtλ)‖Wt‖2+ηtsm,n≤max(‖Wt‖2,sm,nλ).

逐步使用这个关系,就有 ‖Wt‖2≤max(‖W0‖2,sm,n/λ)。因此,限制每步更新的谱范数,再配合 weight decay,可以限制权重谱范数的累积增长。

实际近似若满足统一的界 ‖Ot‖2≤C,则将上界中的 sm,n 换成 Csm,n。这个界通常很宽松,也没有同时约束整个网络的输入特征,不能把它当作 attention logits 必然稳定的保证。

实际训练中的取舍

Muon 通常用于 Transformer 隐藏层的 attention projection 和 MLP 权重;embedding、LM head、bias 和 normalization 参数交给 AdamW。Embedding 和 LM head 虽然也是二维矩阵,但参数分组还要考虑它们的用途,不能只检查 ndim == 2。作者还报告 (Jordan et al., 2024),把合并存储的 Q、K、V 分开做 Muon 更新更好;矩阵如何分块会改变变换本身。

比较项AdamWMuon 的矩阵参数部分
更新处理用二阶矩估计逐元素缩放用矩阵多项式调整 momentum 的奇异值
持久优化器状态一阶、二阶矩一个 momentum buffer
每步额外计算主要是逐元素运算多次矩阵乘法和临时方阵
参数切片更新可逐元素独立计算任意切片后分别正交化通常不等价

假设两者的状态都用 FP32,Muon 的矩阵参数每个元素需 4 bytes 的持久 momentum,AdamW 的两个矩估计需 8 bytes。这只比较优化器状态,不包含参数、梯度、master weights、activations、Newton–Schulz workspace,以及仍由 AdamW 管理的参数,不能据此说训练总显存减半。

分布式训练的区别也来自矩阵结构。例如按行分成 M1,M2,完整 Gram matrix 包含 M1M2⊤ 这样的跨分片项;各卡独立处理自己的切片会漏掉它们。结合 FSDP / ZeRO 或 Tensor Parallel 时,需要恢复相关矩阵,或使用数学上对应的分布式矩阵运算。Moonlight 的实现会 gather 相关分片再计算完整更新,并保留本卡需要的部分;这些临时状态和通信都要计入成本。

Moonlight (Liu et al., 2025) 在其 compute-optimal scaling 实验中报告了相对 AdamW 约 2× 的计算效率,并用 Muon 训练了总参数 16B、激活参数约 3B 的 MoE。这个结果描述的是其训练设置下的 compute–loss 关系,不代表一个 optimizer step 快了两倍。Muon 每步增加的矩阵计算能否被达到目标 loss 所需 step 的减少抵消,要在相同硬件和训练设置下测 wall-clock time。

更大规模还会遇到稳定性问题。Kimi K2 的 MuonClip (Kimi Team, 2025) 在 Muon 更新之外加入 QK-Clip:根据每个 attention head 观测到的 logits,在更新后缩放相关 query / key 权重,抑制 logits 过大。它是额外的稳定化机制;仅把更新矩阵的奇异值拉平,并不能限制整个训练过程中权重和 attention logits 的增长。

References

Bernstein, J., & Newhouse, L. (2024). Old Optimizer, New Norm: An Anthology. doi.org
Jordan, K., Jin, Y., Boza, V., You, J., Cesista, F., Newhouse, L., & Bernstein, J. (2024). Muon: An optimizer for hidden layers in neural networks. kellerjordan.github.io
Kimi Team. (2025). Kimi K2: Open Agentic Intelligence. doi.org
Liu, J., Su, J., Yao, X., Jiang, Z., Lai, G., Du, Y., Qin, Y., Xu, W., Lu, E., Yan, J., Chen, Y., Zheng, H., Liu, Y., Liu, S., Yin, B., He, W., Zhu, H., Wang, Y., Wang, J., … Yang, Z. (2025). Muon is Scalable for LLM Training. doi.org
Loshchilov, I., & Hutter, F. (2019). Decoupled Weight Decay Regularization. International Conference on Learning Representations. arxiv.org
PyTorch Contributors. (n.d.). Muon — PyTorch 2.9 documentation. docs.pytorch.org
Su, J. (2024, December). Muon优化器赏析:从向量到矩阵的本质跨越. kexue.fm
Su, J. (2025a, February). Muon续集:为什么我们选择尝试Muon?. kexue.fm
Su, J. (2025b, May). msign算子的Newton-Schulz迭代(上). kexue.fm
Su, J. (2025c, November). Muon优化器指南:快速上手与关键细节. kexue.fm
Su, J. (2026, June). 为什么官方版Muon比MuP版多出一个max(1, ⋅)?. kexue.fm

Cite this post

@misc{pu2026mlmlrevisitoptmmuon,
  author = {Pu, Fanyi and {GPT-6 Astra}},
  title  = {Muon},
  year   = {2026},
  month  = {9},
  url    = {https://pufanyi.com/blog/ml/ml-revisit/optm/muon}
}