Muon优化器指南：快速上手与关键细节
原文链接 | 2025-11-19 | 苏剑林
这段时间，相信很多读者已经刷到过Muon优化器的相关消息。实际上，Muon的提出时间大致是去年的10月份，由 Keller Jordan 在推特上提出，距今也不过一年多一点。然而，就在这一年里，Muon已经经历了百亿、千亿乃至万亿参数模型的训练考验，足以表明它是一个相当有竞争力的优化器。
如今，Muon已经内置在Torch、Keras等训练框架中，就连Megatron这样的大型框架也逐渐开始支持，这意味它已经获得了业界的普遍认可。不过，对于仅熟悉Adam的读者来说，如何快速有效地切换到Muon，可能依然是一件让人困惑的事情。所以，本文试图给出一个快速上手教程。
简要介绍#
Muon的正式提出者是 Keller Jordan ，目前任职于OpenAI。开头说了，Muon最早发表在推特上，而直到现在，作者也只是多写了篇博客《Muon: An optimizer for hidden layers in neural networks》而不是一篇Paper，作者的观点是“是否写成Paper，跟优化器是否有效，没有任何关系[原文]”。
Muon是一个专门为矩阵参数定制的优化器，也有一些相关工作具有类似的特点，比如Shampoo，还有更早一些的Stochastic Spectral Descent，等等。很多工作或多或少都能关联上Muon，但没有一个是能够完全覆盖Muon的，所以在笔者看来Muon算是一个全新的工作。
在国内，最早向大家科普Muon的文章，应该是笔者的博客《Muon优化器赏析：从向量到矩阵的本质跨越》，而首次在较大规模的模型上验证Muon，应该是我们二月份发布的Moonlight，其所提的Moonlight版Muon，用到了后来的万亿参数的K2中。K2之后，GLM-4.5同样用到了这个Muon变体。
跟Muon的作者之一Jeremy Bernstein在他的博客《Deriving Muon》所说的一样，对笔者而言，Muon的独特之处在于它可以基于更本质的优化原理推导而来，并在实践中有效，相比之下，虽然Adam也很有效，但它更像是一种启发式方案。
四个版本#
本文不打算介绍Muon的数学细节，也不打算介绍Muon的实现，而是主要介绍从Adam切换到Muon的一些技术细节和注意事项。刚才说了，Muon是专用于矩阵参数优化的，并且是非Element-wise的更新规则，这可能导致新用户上手起来会比较困惑。
还有，据笔者所知，Muon目前至少有四个略微不同的版本，这种多版本现象也加剧了这种困惑。如果用户不了解其中的细节，可能会因为调错超参数（特别是学习率）而导致不好的效果，下面会着重理清楚这些内容。首先，对于一个矩阵\(\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}\)，\(\boldsymbol{G}\)是它的梯度，四个Muon变体分别是：
\[\begin{aligned}\newcommand{msign}{\mathop{\text{msign}}} &\quad\boldsymbol{M}_t \quad=\quad \beta \boldsymbol{M}_{t-1} + \boldsymbol{G}_t \\[10pt] &\quad\boldsymbol{W}_t = \left\{ \begin{aligned} &\boldsymbol{W}_{t-1} - \eta_t \left(\msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}\right) &\color{skyblue}{(\text{朴素版})} \\[5pt] & \boldsymbol{W}_{t-1} - \eta_t \left(\sqrt{\max(1, d_{out}/d_{in})}\msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}\right) &\color{skyblue}{(\text{KellerJordan版})} \\[5pt] & \boldsymbol{W}_{t-1} - \eta_t \left(\sqrt{ d_{out}/d_{in}}\msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}\right) &\color{skyblue}{(\text{MuP版})} \\[5pt] & \boldsymbol{W}_{t-1} - \eta_t \left(0.2\times\sqrt{\max(d_{out},d_{in})}\msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}\right) &\color{skyblue}{(\text{Moonlight版})} \end{aligned}\right. \end{aligned}\]
如果要开Nesterov动量则将\(\msign(\boldsymbol{M}_t)\)换成\(\msign(\beta\boldsymbol{M}_t + \boldsymbol{G}_t)\)，其中\(\msign\)在实现中通常以
zeropower_via_newtonschulz命名，具体实现细节普通用户可以不管。
四个版本的唯一区别是\(\msign\)前的缩放因子，其中“KellerJordan版”和“MuP版”大同小异，“Moonlight版”则稍微特别一点。Keras只实现了“KellerJordan版”，而Torch则实现了“KellerJordan版”和“Moonlight版”，朴素版目前似乎比较少见，对笔者来说常用的是自己写的“MuP版”。
两个维度#
这里我们要注意一个重要细节，“KellerJordan版”和“MuP版”对\(d_{in},d_{out}\)的顺序是敏感的，所以第一件事情就是要搞清楚\(d_{in},d_{out}\)的含义，并不是说矩阵的第一个维度就一定是\(d_{in}\)、第二个维度就一定是\(d_{out}\)。
\(d_{in}\)、\(d_{out}\)的含义分别是线性层的输入、输出维度，所以哪个是\(d_{in}\)哪个是\(d_{out}\)，要看线性层的具体实现。比如Keras的Dense层实现是\(\boldsymbol{x}\boldsymbol{W}\)，那么矩阵\(\boldsymbol{W}\)的第一个维度是\(d_{in}\)、第二个维度是\(d_{out}\)；然而，Torch的Linear层实现的是\(\boldsymbol{x}\boldsymbol{W}^{\top}\)，所以矩阵\(\boldsymbol{W}\)的第二个维度是\(d_{in}\)，第一个维度才是\(d_{out}\)。
所以，如果要实现“KellerJordan版”的Muon，对于Torch的Linear层，缩放因子应该是
max(1, W.shape[0]/W.shape[1])**0.5，而对于Keras则应该是
max(1, W.shape[1]/W.shape[0])**0.5。所以，当前Keras（3.12版）的Muon实现其实是不正确的，因为它照搬了Torch的缩放因子实现[源码]。
如果是自己写的模型，则需要根据自己的写法谨慎判断了，比如不排除在Torch中混用了自带的Linear层和手写的
x @ W，这样就不能一概而论是
W.shape[0]/W.shape[1]还是
W.shape[1]/W.shape[0]了。当然，如果你嫌搞清楚这些比较麻烦，那就可以考虑用“Moonlight版”，它的缩放因子关于\(d_{in},d_{out}\)是对称的。
超参设置#
搞清楚\(d_{in},d_{out}\)后，剩下就是学习率\(\eta_t\)和权重衰减系数\(\lambda\)怎么设置了。这里的假设是用户已经有Adam的调参经验，在Adam下已经得到了不错的效果，想要快速迁移到Muon体验一波。
我们先来看“Moonlight版”，它的缩放因子是通过对齐Adam的Update RMS得到的，如果想要了解细节，可以参考《Muon续集：为什么我们选择尝试Muon？》，至于\(0.2\)这个“Magic Number”，可以参考《为什么Adam的Update RMS是0.2？》。简单来说，“Moonlight版”Muon对齐了Adam的更新幅度，所以从Adam迁移过来的最简单的做法是：啥也不用改，用回Adam的\(\eta_t\)和\(\lambda\)就行。
然后看剩余三个版本。我们知道，主流模型通常有个hidden_size（记为\(d\)），模型的矩阵形状多数不会明显偏离\(d\times d\)，因此我们以\(d_{in}=d_{out}=d\)来近似处理，此时这三个版本都是一样的，相比“Moonlight版”则少了个\(0.2\sqrt{d}\)。既然“Moonlight版”对齐了Adam更新幅度可以不改变超参，那么这三个版本的学习率应该要放大\(0.2\sqrt{d}\)倍，才能对齐Adam的更新幅度，相应地，\(\lambda\)则要除以\(0.2\sqrt{d}\)。
代入\(d=1024,2048,4096\)，结果分别是\(6.4, 9, 12.8\)，如果记不住\(0.2\sqrt{d}\)，那么可以简单记住，如果我们使用另外三个版本的Muon，那么直接将Adam的学习率放大10倍来作为Muon的学习率。如果直接代入Adam的学习率到Muon中，就会因为欠拟合得到Muon远不如Adam的结论，据笔者所知，Muon的一些差评便来源于此。
这样看还是“Moonlight版”更好用？“Moonlight版”确实有不错的实践效果，但如果就此说它更好用，其实是站在Adam的角度下评价了。“MuP版”或“KellerJordan版”的好处是学习率可迁移，即在小模型调好学习率后，直接用到大模型往往也有不错的效果，这部分可以参考Jeremy Bernstein的博客《Deriving Muon》或笔者的博客《高阶MuP：更简明但更高明的谱条件缩放》
其他参数#
如果说Muon只管矩阵参数，那么其余参数怎么办呢？比如线性层的Bias项、RMSNorm的gamma项，这些是1维的参数；又比如卷积层可能出现3维、4维数组的参数。
这里先要更正一下，Muon并不是只管矩阵参数，Muon是只管“稠密输入的线性层的矩阵参数”，如果读者觉得这个比较费解，那么只需要记住Embedding层和最后的分类层（包括GPT的LM Head）的矩阵参数都不能用Muon，否则效果会明显差。这些不能用Muon的矩阵参数，还有1维、3维及更高维的参数，如果读者不想费太多心思，那么直接用Adam就行，基本上Muon实现都是混合了Adam的，用户可选某些层用Adam。
如果读者愿意捣鼓，那么像卷积层的3、4维参数，也可以用上Muon。以Conv2D为例，卷积核形状通常是\((w, h, d_{in}, d_{out})\)，它的等效实现其实是把\((w, h, d_{in})\)的Patch输入展平为\(w \times h \times d_{in}\)的向量，然后卷积核也Reshape为\((w\times h \times d_{in}, d_{out})\)，最后做矩阵乘法，所以它想要用Muon，那就要先将动量Reshape为\((w\times h \times d_{in}, d_{out})\)，计算\(\msign\)后再Reshape回去更新。
类似地，还有RMSNorm的gamma参数，可以视为跟对角矩阵做乘法，于是将它的动量视为对角矩阵，也可以算\(\msign\)，结果等效于SignSGDM；Embedding层可以视为多个\((1,d)\)矩阵去计算\(\msign\)，结果是Normalized SGDM（参考《Muon优化器赏析：从向量到矩阵的本质跨越》）。如果还想折腾，比如Multi-Head的Attention，每个Head的投影矩阵是不是可以考虑单独拿出来分别做\(\msign\)…
生命不息，折腾不止～
期望结果#
最后，如果用户按照上述说明正确设置并跑起来了，那么就可以开始祈祷幸运之神的降临了。
我们应该期待一个什么样的结果呢？如果没有出现梯度爆炸等异常情况，那么多数情况下Muon会比Adam略好一些，当然也不排除某些情况Muon会略差，但不论如何，它们的差距不会非常大。如果出现一方比另一方好非常多的现象，那么可能得反思一下哪边的设置出现问题了。
不过，这都不是绝对的，比如某些极端的设置下，确实也会出现Muon比Adam好很多，Adam怎么调也不行的现象。总之，祝你幸运。如果出现有意思的现象，欢迎一起交流分析。