原创 苏剑林 2025-09-12 20:37 北京
训练一稳定，Adam就变0.2？
©PaperWeekly 原创 · 作者 | 苏剑林
单位 | 科学空间
研究方向 | NLP、神经网络
众所周知，我们很早就开始尝试将 Muon 用于大规模 LLM 的训练。
特别地，在QK-Clip巧解MaxLogit爆炸难题：让Muon在Scaleup之路上更进一步中，我们提出了 “Match Adam Update RMS” 的技巧，以便快速从 Adam 迁移到 Muon 上，这个技巧同样用到了 Kimi K2 的训练中。
该技巧是指将 Muon 的 Update RMS 统一成 0.2，这使得我们复用 Adam 的学习率和权重衰减率。
这一技巧的背后，是我们观察到 Adam 的 Update RMS 约等于 0.2，并且这一现象是稳定且可复现的。这便引发了一个有趣的问题：为什么 Adam 的 Update RMS 是 0.2？我们可以从理论上解释它吗？
问题引入
首先描述一下现象：从实验中我们观察到，大致上在 Warmup 结束、模型进入正式训练后，Adam 的 Update RMS 几乎都保持在 0.2～0.3 之间，并且不同尺寸的模型也呈现出相似的规律。
这些模型的共同点是都用 Adam 训练，参数是 。由于共性很明显，所以这大概率不是巧合，因此笔者尝试分析背后的原理。
然后我们回顾一下 Adam 优化器的形式：
注意：本文所有向量的乘除法，包括平方，默认都是指 Hadamard 积/商，即 Element-wise 的乘/除。
我们要做的事情，就是证明 ，至少在 这组设置下如此。
由于我们关心的是稳定训练后的情形，因此可以认为 t 足够大，以至于 都足够接近于 0，那么就不用区分 和 、 和 。同时，我们假设 足够小，也可以忽略，于是有 。
对于 ，我们可以得到展开式
数值模拟
如果我们假设 都是从同一个分布采样出来的，那么我们就可以直接用数值模拟的方法估计 。事不宜迟，让我们从最简单的标准正态分布 进行尝试，参考代码如下：
import numpy as npN, T = 10000, 2000beta1, beta2 = 0.9, 0.95m, v = 0, 0for i in range(T): g = np.random.randn(N) m = beta1 * m + (1 - beta1) * g v = beta2 * v + (1 - beta2) * g**2u = m / v**0.5rms = (u**2).mean()**0.5print(rms)大家猜猜结果是多少？答案大概是 0.225，居然跟实验结果惊人地相似！这反过来表明我们的模拟假设跟实际情况还是很吻合的。
可能有读者觉得不对， 不是纯噪声了吗，这还能吻合？实际训练当然不可能是纯噪声，只能说单次梯度的信噪比小得可怜，因此可以用纯噪声来模拟。
读者可以自行折腾一下上述参考代码，观察 Update RMS 的影响变量，大体结论是：Update RMS 跟 正相关，跟 似乎关系不大，如果 的分布具有非零均值（相当于增大梯度的信噪比），那么 Update RMS 也会变大。
平均近似
这一节笔者尝试从理论方面推导上述模拟结果的一个近似解析解。首先，我们从 RMS 的定义可知，要求 ，需要先求 。笔者的想法是，用 的期望作为它的近似，并进一步转化为平均场近似：
可能会有读者质疑最后一步近似的合理性。笔者的建议是，先不管这些细枝末节，就好比上一节假设 一样，先算了再说，如果结果合理那么过程必然一定程度上也是合理的。
现在我们分别算分子、分母，这次我们一般地设 ，其中分母比较简单
至于分子，可以直接展开平方计算，也可以稍微偷懒一下：我们要求的是 的二阶矩 ，它又等于 ，由于 是 的加权平均，所以必然有 ；至于方差，它具有平方可加性，因此
所以
结果分析
由于 已经是平方后的向量，所以为了估计 ，我们只需要对各个分量求平均然后开平方。求平均这一步，我们不妨再来一次平均场近似（分子分母分别求平均），最终将得到
它有两个影响因子：一是 ，这可以看成是梯度的信噪比（SNR）；二是 ，这 是Adam 的超参数之一。特别地，结果不依赖于 ，这跟前面的模拟结果吻合。那么这个式子究竟近似得好不好呢？我们不妨考虑最简单的特例 ，此时
代入 ，结果是 ，跟模拟结果和实践表现居然都很吻合！进一步地，它跟模拟结果的多个对比如下：
▲ 模拟结果与平均场近似（不同 beta1、beta2）
应该说，近似程度还是不错的，特别是 之后，结果几乎跟平均场近似重合了。至于考虑 SNR 的比较结果如下：
▲ 模拟结果与平均场近似（不同 beta1、SNR）
当信噪比增大时，平均场近似的误差开始变大，不过仍旧能预测一个整体趋势。事实上，实际训练中梯度的信噪比很少机会能有接近 1 这么大，因此依然可以认为平均场是一个良好近似。
反向预测
如果我们已经接受平均场近似（7），那么可以反过来用它估算梯度的信噪比：
在实际训练中， 是给定的，（也就是 Adam 的 Update RMS）也是可以直接估算的，所以上式是可计算的。当然，这个式子只对 Adam 适用，有没有更一般的估计思路呢？还真有！别忘了前面我们估计得到
那么对它的分量求和然后开平方，我们认为它会是 的一个近似：
至于二阶矩是 ，而像 Muon 之类的优化器并没有二阶矩可用，但是我们留意到二阶矩的结果是跟 无关的，所以我们不妨考虑一个最简单的特例—— ——此时 。
当然这可能有点勉强，但估算嘛肯定是怎么方便怎么来。这个“近似”意味着成立 ，于是我们有
右端的形式跟式（7）如出一辙，所以我们可以写出
也就是用 替代 ，这就给出了一种带动量优化器通用的估计 的思路。
可能还有读者想问动量都没有咋办？这就真没有办法了，因为这里的 属于跨优化轨迹的统计量，我们总得有些跨轨迹的统计信息，才有可能去估计它。
文章小结
本文主要从模拟实验和理论近似两个角度探讨了 Adam 的 Update RMS，它可以作为我们在 Muon 优化器中将 Update RMS 对齐到 0.2 的理论依据之一。
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：hr@paperweekly.site
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
·