Chong Jie 2025-09-15 23:37 北京
©作者 | Chong Jie
单位 | 上海交通大学
研究方向 | 参数高效微调
放在开头:这篇文章会比较零碎,是由我之前做过的一次分享记录而成。我没有特意去梳理文章,直接把讲稿拿来主义,属于知识比较密集,需要读者自行整理。不过总的逻辑是顺的,可以放心阅读。
本文将分为两部分讲清楚 1)什么是 Muon 优化器、他怎么做的、以及一些相关工作,2)我们对 Muon 的探索。
优化算法是现代深度学习的基石。它决定了训练过程的轨迹,并显著影响模型的收敛速度与泛化性能。随着模型规模扩展至数十亿甚至数万亿参数量级,优化器的性能需求也相应提升。
诸如 AdaGrad、RMSProp、Adam 及其变体 AdamW 等自适应方法,已成为训练语言模型的事实标准,其优势在于稳健性高、超参数易于调节,并能在早期训练阶段实现快速收敛。
现有的方法基本上可以总结为以下三步:
1. 计算参数的梯度(向量化或者矩阵化)
2. 对梯度进行变换得到修正后的分量
3. 梯度回传
公式整理为:
1.
2.
3.
所以重点在于如何设计 。
在介绍 Muon 之前,我先来简单介绍两个经典的优化器,并且和 Muon 有很强的关系。
Adam
Adam 是一种常用的自适应学习率优化算法,结合了 RMSProp 和带动量的 SGD 的优点,能够在各种神经网络训练任务中快速收敛。Adam 可以算是至今为止最成功的优化器之一了。
动态学习率:在深度学习训练中,每个参数维度上的梯度大小和变化往往相差巨大,使用单一全局学习率容易导致部分参数更新过快而震荡,或更新过慢而难以收敛。
Adam 通过同时跟踪梯度的一阶矩(即梯度的指数加权平均,用于捕捉梯度方向)和二阶矩(即梯度平方的指数加权平均,用于估计梯度幅度),为每个参数动态调整有效学习率,从而在抑制噪声、加速收敛和兼顾稳定性方面表现出色。
bias correction:由于初始时 和 均为零,这两项在训练早期会产生偏向零的现象,因此需要进行 correction:
为什么用梯度的平方作为二阶矩:牛顿法强调了在优化过程中考虑曲率信息的重要性,通常是通过计算 Hessian 的逆来获得更优的局部下降方向。
虽然在大规模模型中直接计算或存储完整的 Hessian 不现实,但可以用 Fisher 信息矩阵(FIM)来进行近似。Adam 通过对梯度平方的指数滑动平均来估计这一信息,从而实现根据梯度方差动态调整步长。
Shampoo
Shampoo 是由 Google 提出的一个高效二阶优化算法,其核心思想是在不显式计算和存储完整 Hessian 矩阵的情况下,利用梯度的二阶矩信息来为每个参数维度自适应地调整更新方向和步幅。据说 Gemini 就是用它训的。
可以看到,Shampoo 是利用了 Fisher 矩阵的行和列关系。
Muon的背景与具体做法
随着现代大规模模型在参数维度和结构复杂性上的不断增长,其参数空间表现出越来越丰富的依赖关系与矩阵结构几何特性。这一趋势促使研究者越来越关注如何通过挖掘梯度中的全局结构,提升训练过程的条件数和效率。Muon 就是在这么样的背景下提出来的。
我们首先介绍一下 Muon 优化器,这是一种专门为二维参数结构设计的新型优化方法。
给定一个权重矩阵 ,以及当前迭代步 的梯度 ,Muon 首先会计算一个常规的动量缓冲项 ,其中 是动量系数。
不同于常规优化器直接使用动量或其自适应变种,Muon 对动量项进行了几何结构上的变换,从而生成更具结构性的更新方向。
具体来说,它会提取动量矩阵的“正交部分”,即通过极分解计算出 的正交因子 。理论上,这个因子可以看作是对 进行奇异值分解后得到的 。
考虑到完整 SVD 的计算代价较高,Muon 使用了一种更高效的替代方案:Newton–Schulz 迭代,来近似地求解这个正交因子:,
最终,参数更新规则如下:。
这种基于 Newton–Schulz 迭代的矩阵正交化方式,为参数更新带来了明确的结构和几何约束。相比传统的对角类自适应方法,它在高维参数空间中能更稳定地进行训练,并更高效地探索损失函数的形状。
牛顿迭代法的目的是逼近极分解。他的 motivation 是,对于一个需要分解的矩阵 G,我们可以通过矩阵乘法和加法,得到:
那接下来事情就简单多了,我只需要让中间关于 S 的多项式能够趋向于 1 就可以了。这就是牛顿迭代的思想,最终通过设计一个多项式,通过迭代多次,便可以将一个 [0,1] 之间的数映射到 1 或者接近 1。
我个人的一个简单理解:有一个函数,对于一个非常小的数,我将其放大 3 倍(假如),并且让最大值不超过 1。那么经历多次放大(迭代),我便可以将一个非常小的数放大到接近 1 了。
行业现状
Muon 还存在一些问题,针对这些问题,现有的工作提出了不同的技术方案。
4.1 训练稳定性与学习率的对齐问题
Muon 优化器存在以下两个问题:
1. Muon 训练的参数更新量的 RMS norm 不稳定。
2. Muon 在实际使用时是只作用在 2 维参数,1 维参数还是使用 AdamW。这表示需要用到额外一套学习率,增加调参的负担。
为了解决上述两个问题,Moonshot 提出以下解决策略:
1. 引入 weight decay,稳定更新量;
2. 对齐 Muon 和 Adam 更新量的 RMS norm,复用 Adam 学习率。
因此,最终 Muon 的形式为:
4.2 Attn logits 的爆炸问题
Moonshot 还在其实验中观测到 Muon 更容易产生大的 attn logits。
他们也对应的提出了 MuonClip,并用在 Kimi K2 中:
实验
我们自己也做了一些实验,结论是:Muon 可以取得比 Adam 25% 左右的 training efficiency。
上文我们完整回顾了 Muon 优化器的原理、相关工作以及工程挑战。接下来,我们将分享自己在 Muon 上的进一步探索,并介绍改进版本 AdaMuon —— 所有代码均已开源。
https://github.com/Chongjie-Si/AdaMuon
多啰嗦一句,我觉得 nanoGPT 这个代码框是真的好用,改一改模型或者优化器,很快就能出结果了,很适合作为学术界发论文或者工业界的一个小规模验证实验。
另外 Megatron-LM 的代码也一并附上了。
闲聊一下
今年 3 月份,我在 X 上闲逛,偶然间看到一个叫做 Muon 的优化器。当时看到评论区的人有的是一顿狠吹,说是训练效率大幅提升,并会取代 Adam;也有的人唱衰,说不好用,没有任何效果。其实直到今天关于 Muon 也大致分为这两派人,当然这是后话。
我本着吃瓜的态度,搜索了一下 Muon 的论文,结果 google scholar 上没有,arxiv 上也没有,最后才发现是一篇博客。当看完了博客,我的第一反应是质疑,因为 Muon 和 Shampoo 实在是太像了,甚至可以说是 Shampoo 的简化版:
Muon 与 Shampoo 的核心区别在于是否对预条件矩阵(preconditioner)进行历史累计——Shampoo 通过不断累加历史信息(虽不同于典型的 EMA,但本质仍是累积),而 Muon 则没有这一机制。
从动量优化的角度来看,我认为 Shampoo 应当具备更好的效果。之前也听说 Gemini 在用 Shampoo 训练,这也在一定程度上支持了我的判断。
不过,Shampoo 需要计算矩阵逆,计算开销较大;而 Muon 借助牛顿迭代法避免了显式求逆,在一定程度上可能比 Shampoo 更高效。但真正引起我注意的是苏剑林老师团队近期的工作。虽然我没有找到 Muon 的原论文,但发现了一篇刚刚发布的研究:
https://arxiv.org/abs/2502.16982
这一发现让我颇感意外,也隐约看到了新的可能性。既然工业界已经开始关注 Muon,并对其扩展性进行了实证分析,也许 Muon 真的有望成为 Adam 的替代方案。
不过,当时除了 Kimi 之外,我并未看到更多论文对 Muon 展开深入探索,一切尚处于未知阶段。
这也让我有些犹豫:转向一个新方向必然带来一定的成本,而当时我正专注于 model merging 方面的研究,同时 hold 两个方向虽然可行,但前者已是一片红海,后者却仍迷雾重重——鸡蛋固然不应放在一个篮子里,但资源的分配也需谨慎权衡。
但是,在经过短暂的考虑之后,我当下决定切换方向——之前的不做了。学界需要追求新的东西,业界需要的是真正有用的东西。既然现在还没有多少人关注,或许只是因为大家尚在观望。如果日后它真的展现出巨大价值,即便我吃不上肉,说不定也能跟着喝上一口汤。
我到现在还记得,当我告诉导师我打算做优化器时,他脸上那种震惊和不解的神情。也许他也没想明白,为什么在 LLM 如此火热的当下,我会突然对一个小众的领域产生兴趣。
不过谁又能想到,今天 Kimi v2 1T 的模型用 Muon 训练成功了,GLM4.5 也用 Muon 训练成功了呢?
对Muon的分析
好了,闲话不多说。开启一个新方向,第一要务就是打通一套可靠的 baseline。我选择了一个基于 nanoGPT 修改而来的代码仓库作为起点。在实现了 Muon 优化器之后,初步的实验结果却并不理想。Muon 在训练早期往往比 Adam 收敛得更快,但在中后期却容易被反超。
对于一个新兴方法来说,到这个阶段很容易让人怀疑其有效性。但出于一种直觉和信念,我总觉得应该是代码的某个地方出了问题。
经过逐行排查(幸亏代码量不大),我意外地发现了 Hugging Face 在 GPT-2 的注意力机制中的一个设置:
if self.scale_attn_by_inverse_layer_idx: att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)) / float(self.idx_layer + 1))else: att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
简单来说,原实现中对注意力输出额外除以了层数(layer),使得注意力值被人为缩小。这一做法相当奇怪——我从未在论文中见过类似描述,也没有看到任何理论依据,但它却实实在在地出现在了代码中。
在我将这个 hack 掉之后,训练结果立刻回归正常。Muon 的表现显著提升,相比 Adam 取得了 20% 到甚至 50% 的训练效率增益。
通俗来讲,在最优设置下,Muon 仅需训练 10B token 所达到的验证损失,就相当于 Adam 训练 20B token 的结果(所有的实验结果图我会放到最后)。紧接着 Megatron-LM 上的 Muon 也复现成功,在大规模实验上 Muon 也展示出了惊人的效率。
这个过程看似轻松,实际上耗费了我近两个月的时间。此外,它也给了我一个重要的启示:某些模型架构可能对 Muon 并不友好。
由于现有的模型大多基于 Adam 优化器进行训练和调优,许多被广泛采用的设计实际上是在 Adam 的语境下被验证有效的——我们可以认为这些架构是“为 Adam 优化”的。
这一发现也为未来的模型设计提供了新的思路:在探索新的优化方法时,我们或许也需要重新审视和调整模型结构本身,以更好地适配优化器的特性,而不仅仅默认沿用为 Adam 所设计的架构。
既然 Muon 已经验证有效,接下来的目标自然是探索能否构建一个更高效的改进版本。从整体框架来看,Muon 可被视为一种二阶优化器,因为它利用了参数矩阵的行和列信息来构造预条件矩阵。
不过,我更倾向于从矩阵操作的角度来理解:如果说 SGD 是元素级别的一阶优化器,那么 Muon 实际上是一种矩阵级别的一阶优化器——它的核心在于对一阶动量进行矩阵层级的更新。
于是,一个自然的思路是:能否为 Muon 引入矩阵层级的二阶动量,从而将这类 SGD 风格的优化器提升到 Adam 的级别?
这一步目前看来仍有困难。既然如此,我们不妨先退一步思考:能否先为其引入元素级别的二阶动量,以增强其自适应能力?元素级二阶动量的重要性已无需赘言,其有效性早已在大量实验和实践中得到广泛验证。
这就是 AdaMuon 的 motivation。
AdaMuon
论文标题
AdaMuon: Adaptive Muon Optimizer
论文链接
https://www.arxiv.org/abs/2507.11005
8.1 难点
要想实现一个二阶动量加持的 Muon,最直接的想法是像 Adam 那样对 Muon 进行扩展,直接引入一个二阶动量累积项,并利用该方差估计对正交更新矩阵 进行缩放。然而,这种做法引入了两个关键的设计挑战:
1. 首先,应确定对哪个分量进行累积。在 Adam 中,二阶动量自然是在原始梯度 上累积的;但在 Muon 中,更新方向 是通过动量 的极分解得到的。目前尚不清楚方差跟踪应当应用于 、 还是 ,因为每一种选择都会导致不同的稳定性和归一化行为。
2. 其次,如何实现与 Adam 相当的归一化效果。在 Adam 中,一阶矩和二阶矩均基于同一梯度信号计算,因此它们的比值可有效实现逐元素的步长归一化。
然而在 Muon 中, 在训练早期和中期可能波动剧烈,导致生成的 在不同坐标间差异显著。因此,无论二阶动量是在 、 还是 上积,所得的统计量仍不稳定,无法提供可靠的归一化效果,甚至可能放大噪声。
要顺利将方差自适应能力融入 Muon,同时不破坏其正交化优势和固有的尺度不变性,就必须审慎解决上述挑战。
8.2 方法
针对第一个挑战,我们选择在正交矩阵 上累积二阶动量,而非原始梯度 或动量缓冲 。
该设计基于两点考虑:首先,原始梯度 带有病态缩放和方向噪声,而 Muon 的极分解正是为消除此类噪声而设计,因此不适合用于稳定的方差估计;其次,虽然 提供了时间平滑的梯度,但其元素值在训练早中期仍不稳定。
此外,由于实际更新方向来自 ,在 上累积方差会导致与更新所用的缩放基准不匹配。相比之下, 提供了几何归一化且稳定的下降方向,为方差估计提供了更干净的基准。
具体来说,我们像 Adam 一样,对 逐元素平方进行指数移动平均:
其中 表示哈达玛积。 作为正交更新的二阶动量缓冲区,功能类似于 Adam 中的方差累积器。系数 直接沿用 Muon 的动量参数,这样也可以不引入额外超参数。方差归一化后的更新方向为:
其中 表示逐元素除法。之后我们也可以像 Muon 那样,对齐 Adam 更新量的 RMSnorm:
看上去开始像模像样了,但这么做的效果并不算是多好。主要原因是,动量缓冲 在元素层面仍会剧烈变化,导致极分解产生的 在不同坐标间波动显著。直接累积不稳定的 难以实现有意义的归一化,甚至可能放大噪声。
为使 既可用于二阶动量缩放又足够稳定,我们提出先对 施加变换 ,再执行极分解
注意到 具有全局尺度不变性,即 ,但其本身无法稳定元素级波动。因此, 需在保持坐标方向性的同时降低波动。最终我们选择 ,理由如下:
设计这些条件的原因在于:
到这里,AdaMuon 已经具备雏形了。我们将算法整理成表:
8.3 实验
来看看效果。
如果是做预训练的话,我觉得上面三张图已经很能说明问题了。无论是训练效率,还是评测集的跑分,AdaMuon 都能取得很好的结果。我现在还在进行 AdaMuon scaling law 的实验,还有一些其他架构的实验。从目前的初步结果来看,还挺有希望的。
一些讨论
我们讨论一些补充性发现。虽然 AdaMuon 在效果上优于 Muon,但不可否认的是,由于需要额外缓存一组二阶动量矩阵,AdaMuon 对 memory 的需求也相应提高。然而,在消融实验中我们观察到了一个非常有趣的现象:
Muon+Sign(我们暂称其为 SiMuon)的表现显著优于原始 Muon,甚至接近 AdaMuon 的效果。而 SiMuon 相比 Muon 仅增加了一个符号操作,计算开销极小,因此我们认为这是一种极具潜力的 Muon variant。
遗憾的是,受限于计算资源与时间,我未能进行大规模实验验证。然而,从现有实验结果已可清晰看出 SiMuon 所具备的潜力。
最后,我们从学术角度探讨几个可能的改进方向。首先,针对如 ResNet 等广泛使用卷积结构的模型,Muon 目前仅适用于二维矩阵参数,而卷积核通常以高维张量形式存在。因此,能否为高维张量设计相应的正交优化器,是一个值得探索的重要问题。
其次,Muon 在分布式训练中的通信开销仍有优化空间。当前实现需要跨数据并行组 gather 各节点的动量分量,这一操作引入了额外的通信负担。
能否设计一种通信效率更高的 Muon 变体,或寻求一种既保持其正交更新特性、又具备类似 Adam 那样低通信开销的替代方案,也具有重要的研究价值。
最终,一个更为根本的研究目标是:能否设计出一种真正意义上的矩阵级二阶优化器?
此类方法有望在矩阵空间中对曲率进行建模,从而实现对优化方向的更精确调整,从根本上突破现有基于标量或元素级自适应方法的局限。实现这一目标将不仅能够显著提升优化效率,更可能为深度学习的优化理论带来新的突破。
上述三个方向中,解决第一个问题可形成一篇具有扎实学术价值的论文,解决第二个问题将直接推动工业界的大规模应用;而第三个方向——实现真正的矩阵级二阶优化器,我认为堪称优化器领域的一颗明珠。
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。
总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。
📝 稿件基本要求:
• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注
• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题
• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道:
• 投稿邮箱:hr@paperweekly.site
• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿
△长按添加PaperWeekly小编
🔍
现在,在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
·