Hexo
0%

Muon优化器原理

Posted on Edited on

参考文章

https://spaces.ac.cn/archives/10592

recap: AdamW

AdamW（Adam with Weight Decay）是目前深度学习中最常用的优化器之一，它是对经典 Adam 算法的一个重要改进，解决了权重衰减（Weight Decay）在自适应梯度算法中失效的问题。

1. AdamW 的更新公式

AdamW 的核心思想是将权重衰减从梯度更新中分离出来。其更新步骤如下：

第一步：计算一阶和二阶矩（动量）

其中 $g_t$ 是当前步的梯度，$m_t$ 是梯度的指数移动平均（一阶矩），$v_t$ 是梯度平方的指数移动平均（二阶矩）。

第二步：偏差纠正（Bias Correction）

第三步：参数更新（关键差异点）

注意：$\lambda \theta_{t-1}$ 这一项是直接作用在参数更新上的，而不是加在梯度 $g_t$ 里。

2. 直觉解释

你可以把 AdamW 想象成一个带刹车的智能小球：

  • $m_t$ (一阶矩)：就像惯性。球不会因为路面一个小坑就立刻停下，而是会根据之前的速度继续滑动，这有助于跳出局部最优。
  • $v_t$ (二阶矩)：就像路面的粗糙度自适应。如果某个方向坡度一直很大（梯度大），它会增加“摩擦力”让更新变稳；如果坡度很小，它会放大更新，确保能走得动。
  • Weight Decay ($\lambda$)：就像向心力。它不断地把参数往原点（0）拉，防止某些参数数值爆炸，从而起到正则化、减少过拟合的作用。

3. 为什么一定要进行偏差纠正？

这是许多初学者容易忽略的地方。我们需要证明：如果不纠正，初始阶段的估算是有偏的。

数学证明

以一阶矩 $m_t$ 为例，其展开式为：

我们希望 $m_t$ 是梯度期望 $E[g_t]$ 的无偏估计。对等式两边取期望：

假设真实梯度分布相对稳定，$E[g_i] \approx E[g_t]$，则：

利用等比数列求和公式 $\sum_{i=0}^{n-1} r^i = \frac{1-r^n}{1-r}$，我们可以得到：

结论

显而易见，$E[m_t] \neq E[g_t]$，而是多了一个因子 $(1-\beta_1^t)$。

  • 在训练初期（$t$ 较小）：由于 $m_0$ 通常初始化为 0，$\beta_1$ 又接近 1（如 0.9），导致 $m_t$ 会严重向 0 偏移，取值远小于真实期望。
  • 随着训练进行（$t \to \infty$）：$\beta_1^t$ 趋于 0，偏移逐渐消失。

偏差纠正的目的就是通过除以 $(1-\beta_t^t)$，抵消掉这个初始阶段的缩小效果，使优化器从第一步起就能准确反映梯度的真实规模。

4. 常用超参数设置

在处理 LLM、Transformer 或视觉模型时，以下是业界公认的“黄金标准”设置：

参数 常用值 说明
Learning Rate ($\eta$) $10^{-3}$ ~ $10^{-4}$ 核心参数，通常配合 Linear Warmup 使用。
$\beta_1$ $0.9$ 控制一阶动量，几乎不需要改动。
$\beta_2$ $0.999$ 或 $0.95$ 控制二阶动量。在大模型训练中，有时调低至 0.95 可增加稳定性。
$\epsilon$ $10^{-8}$ 防止除零的小数，通常保持默认。
Weight Decay ($\lambda$) $0.01$ ~ $0.1$ AdamW 的精髓，负责正则化。

使用建议

  • 搭配调度器：AdamW 必须搭配 Learning Rate Scheduler（如 Cosine Decay）。
  • 解耦学习率：如果你调大了 $\lambda$，通常不需要剧烈改动 $\eta$，这是 AdamW 优于 Adam（加 L2 正则）的地方。

recap: 正交化

正交矩阵的定义以及精确对矩阵正交化的方法(gram-schmidz)

Gemini小解释

这是一个非常精妙的对比视角。你的理解触及了“局部几何”与“全局几何”的差异，但我们可以把这个直觉再推敲得更准一点：

1. 你的直觉对了一半：Adam 的“逐元素独立性”

实际上，Adam 比你想象的更“分裂”。它并没有把所有参数视为一个整体空间中的大向量，而是把每一个标量参数（每一个权重数字）都视为一个独立的维度。

  • Adam 的视角：如果你的模型有 10 亿个参数，Adam 就像是在一个 10 亿维 的空间里。它给每一个维度（每一颗螺丝钉）都装了一个独立的自适应减震器（即 $v_t$ 项，二阶动量）。
  • 这种做法的问题：它完全无视了参数之间的结构关系。它不知道某 100 万个数字其实构成了一个“线性层矩阵”，它只把它们当成 100 万个互不相干的点。

2. Muon 的视角：尊重“结构”

Muon 的“聪明”之处在于它意识到：神经网络不是数字的堆砌，而是变换的组合。

  • 矩阵作为整体：Muon 认为，一个线性层的 $d \times d$ 矩阵不应该是 $d^2$ 个独立的维度，而是一个整体的变换算子。
  • 空间独立：Muon 确实为每个矩阵实施正交化。它在每个层的“局部空间”内，确保该层的变换方向（奇异向量）是均匀且正交的。

3. “整体”与“局部”的降维打击

我们可以用一个形象的比喻来总结 Adam 和 Muon 的空间观：

优化器 空间观 行为比喻
Adam 极其微观：每个参数都是一个独立的轴。 像是一个巨大的合唱团，指挥要求每个人根据自己的音量（梯度历史）独立调整声音大小，不管邻居唱什么。
Muon 结构宏观：每个层（矩阵）是一个独立的空间。 指挥将合唱团按声部（矩阵）分组。要求每个声部内部音域必须横跨高低音（正交化），且组与组之间独立推进。

4. 为什么 Muon 的“局部正交”比 Adam 的“全局自适应”快？

因为 Adam 在每个坐标轴上做的缩放，实际上可能破坏矩阵的性质。
例如，Adam 可能会放大梯度中原本就很小的噪声方向，缩小重要的主成分方向，因为它试图让每个轴的更新步长都变成一样的。

而 Muon 在局部矩阵空间里做 $U V^T$，它保证了：

  1. 不改变能量：正交变换不改变向量的范数，梯度的能量在通过层时不会被扭曲。
  2. 最大的探索率：因为它强迫每个奇异值都为 1，这相当于强迫模型在每一个可能的特征方向上都用相同的力度去“试探”，而不是只盯着那一两个已经很明显的特征（大奇异值方向）。

总结

  • Adam：是坐标轴级别的自适应（Coordinate-wise）。它把所有参数拉平，无视矩阵结构。
  • Muon：是算子级别的自适应（Operator-wise）。它承认矩阵的几何结构，并在每个矩阵内部寻求最优的正交骨架。

既然你提到了“整体空间”，你是否想了解在分布式训练中，Muon 是如何利用这种矩阵独立性来做 Distributed Data Parallel (DDP) 下的通信优化的？（这正是 Muon 能在大型集群上跑得飞快的原因之一）。

GPT解释(超级无敌详细)

可以。下面我按这篇文章的主线，把核心公式一步一步推导出来，并在每一步都说明“为什么要这么做”。我会集中讲三条最关键的线：

  1. Muon 到底在算什么：
  2. 为什么它等于
  3. 为什么可以用 Newton–Schulz 迭代逼近它
    以及这个迭代本质上是在“把奇异值推向 1”。 (科学空间)

0. 先把目标说清楚

文章里 Muon 的更新是

所以真正关键是：给定一个矩阵 $M$，什么叫 $\mathrm{msign}(M)$。文章给出的定义是：若

是 $M$ 的 SVD，那么

其中 $r=\mathrm{rank}(M)$。直觉上，这一步是保留左右奇异向量定义的“方向结构”，但把奇异值统一处理掉，不再让不同方向的尺度差异主导更新。(科学空间)

1. 先复习 SVD，因为后面全靠它

任意矩阵 $M\in\mathbb{R}^{n\times m}$ 都可以写成

这里：

  • $U$ 的列是左奇异向量；
  • $V$ 的列是右奇异向量；
  • $\Sigma$ 是“对角形”的，主对角线上是非负奇异值

更具体地，若只看秩 $r$ 的有效部分，可以写成

其中

  • $U_r\in\mathbb{R}^{n\times r}$,
  • $\Sigma_r=\mathrm{diag}(\sigma_1,\dots,\sigma_r)$,
  • $V_r\in\mathbb{R}^{m\times r}$。

这个写法最重要的直觉是：

也就是说，$M$ 是若干个 rank-1 方向 $u_i v_i^\top$ 的叠加，而 $\sigma_i$ 控制每个方向的强弱。Muon 的核心想法就是：我想保留这些方向 $u_i,v_i$，但不想让 $\sigma_i$ 的不均衡破坏更新。 (科学空间)

2. 为什么 $\mathrm{msign}(M)=UV^\top$ 很自然

先看标量的 sign：

它做了什么？
它保留符号方向，但去掉绝对值大小。

矩阵版想做类似的事：
保留矩阵的“方向信息”，但去掉奇异值大小。

因为标量的“大小”是 $|x|$，矩阵最自然的“大小分解”就是 SVD 里的奇异值 $\sigma_i$。所以如果

那去掉大小、只保留方向，最自然就是把 $\Sigma_r$ 里的每个正奇异值都换成 1：

于是得到

这一步的 intuition 很重要：

  • Adam 类方法是按坐标缩放；
  • Muon是按矩阵的奇异方向来规整更新；
  • 它不是“逐元素地把大数变小、小数变大”，而是“把不同奇异方向的尺度统一”。 (科学空间)

3. 推导文章中的恒等式

$\mathrm{msign}(M)=M(M^\top M)^{-1/2}$

这是全文最关键的恒等式之一。文章直接给了结果，但我们可以完整推出来。(科学空间)

第一步：从 SVD 出发

设

那么

因为 $U_r^\top U_r=I_r$，所以

这里 $\Sigma_r^2=\mathrm{diag}(\sigma_1^2,\dots,\sigma_r^2)$。

第二步：开平方再取逆

因为

所以它的平方根是

其逆平方根（若有零奇异值则理解为伪逆）就是

这里

第三步：右乘回去

现在计算

代入上面的表达式：

因为 $V_r^\top V_r=I_r$，所以

而这正是定义中的

所以我们得到

4. 同理推导

$\mathrm{msign}(M)=(MM^\top)^{-1/2}M$

同样从

出发，有

所以

然后左乘 $M$：

因此

这条公式的直觉是：
$\mathrm{msign}(M)$ 就像对 $M$ 做一种“矩阵归一化”。
标量里是

矩阵里对应地变成

文章也明确指出这是它作为 sign 的矩阵推广的关键理解。(科学空间)

5. 为什么对向量时它会退化成 $l_2$ 归一化

文章说，如果把向量 $m\in\mathbb{R}^n$ 看成 $n\times 1$ 矩阵，那么

我们现在直接推。(科学空间)

因为 $m$ 是列向量，所以

是一个 $1\times 1$ 矩阵，也就是标量：

于是

代回恒等式

就得到

这特别值得记住，因为它说明：

  • 标量情形：sign 是“除以绝对值”；
  • 向量情形：msign 是“除以 $l_2$ 范数”；
  • 矩阵情形：msign 是“除以矩阵意义下的尺度”，即奇异值。 (科学空间)

6. 为什么对对角矩阵时它退化成逐元素 sign

设

则

所以

因此

也就是

这正是逐元素 sign。文章借此说明：Muon 可以看成把 SignSGD/Tiger 的逐元素规整化，推广成了矩阵级规整化。(科学空间)

7. 推导“最优正交近似”公式

$\mathrm{msign}(M)=\arg\min_{O^\top O=I}|M-O|_F^2$

这也是文章的核心结论之一，而且非常有解释力：
Muon 不是随便把奇异值变成 1，而是在 Frobenius 范数下，找离 $M$ 最近的正交矩阵。 (科学空间)

我们一步一步来。

第一步：展开平方

目标函数是

利用 Frobenius 范数恒等式：

所以

Frobenius 内积满足

于是

第二步：利用 $O$ 是正交矩阵

若 $O^\top O=I$ 且 $O$ 是 $n\times n$ 正交矩阵，那么

所以目标变成

注意这里前两项与 $O$ 无关，所以最小化距离等价于最大化

这一步 intuition 很关键：
“离 $M$ 最近”这件事，最后变成了“让 $O$ 与 $M$ 的对齐程度最大”。(科学空间)

第三步：代入 SVD

令

则

利用 trace 的循环不变性：

可写成

记

因为 $U,V,O$ 都是正交矩阵，$Q$ 也是正交矩阵。于是

第四步：把 trace 展开

由于 $\Sigma$ 是对角的，

而正交矩阵每个对角元都满足

又因为 $\sigma_i\ge 0$，所以要让

最大，就应该让每个 $Q_{ii}$ 尽量大，也就是取

最理想的情况就是

于是

所以最优解是

这条结论给了一个非常漂亮的几何解释：

  • $M$ 可能“歪歪扭扭”、不同方向尺度不同；
  • $\mathrm{msign}(M)$ 是最接近它的“纯正交版本”；
  • Muon 其实是在拿这个“最接近的正交更新”来替代原始矩阵更新。(科学空间)

8. 现在解释它为什么像“自适应学习率”

文章指出，Muon 与 Adam 类似，也有“尺度不敏感”和“更各向同性”的性质。(科学空间)

8.1 为什么损失整体缩放不影响方向

若把损失乘上常数 $\lambda$，梯度矩阵也乘上 $\lambda$，于是动量矩阵 $M$ 也会乘上 $\lambda$。

设

那么

于是

因为奇异值全被“置一”了，所以尺度因子 $\lambda$ 被消掉了。

这和 Adam 用归一化来减少尺度敏感性是同一种精神，只不过 Adam 在坐标维度上做，Muon 在奇异方向上做。(科学空间)

8.2 为什么说它让更新更“各向同性”

如果

那 $\Sigma$ 里的不同奇异值，代表不同奇异方向上的拉伸强弱不同。
奇异值差异越大，矩阵越“各向异性”。

Muon 做的是把

变成

也就是把所有奇异方向的尺度统一。于是更新不再偏向某几个奇异值特别大的方向，而是在有效子空间里更均匀。这就是文章说的“更各向同性”。(科学空间)

9. 为什么可以不用 SVD，改用 Newton–Schulz 迭代

SVD 太贵，所以文章接下来从恒等式

出发，想逼近矩阵函数 $X\mapsto X^{-1/2}$。(科学空间)

核心想法是：先把标量函数

在 $t=1$ 附近做泰勒展开，然后把这个多项式“搬到矩阵上”。

文章给出：

保留到二阶后可整理成

(科学空间)

我们把这一步完整展开一下。

10. 从泰勒展开推到多项式近似

从

开始，逐项展开。

先算

再算

三项加起来：

把常数项合并：

把 $t$ 项合并：

保留 $t^2$ 项：

所以

于是

再左乘 $M$：

这就是文章里的近似公式。(科学空间)

11. 为什么这会导出迭代格式

文章接着说：如果 $X_t$ 已经是 $\mathrm{msign}(M)$ 的一个近似，那把同样的多项式作用到 $X_t$ 上，可能得到更好的近似，于是写成

(科学空间)

这一步的 intuition 是：

  • 真正想要的结果 $Y=\mathrm{msign}(M)$ 满足（若满列秩且 $n\ge m$，则有 $Y^\top Y=I$）；
  • 所以如果当前 $X_t$ 还没做到 $X_t^\top X_t=I$，我们就用一个多项式，把它朝这个条件“推过去”。

换句话说，这个迭代不是在直接逼近矩阵元素，而是在逼近“正交性条件”。(科学空间)

12. 文章最关键的 insight：

这个迭代本质上是在迭代奇异值

这是全文最值得吃透的一步。文章把一般迭代写成

然后说明它等价于对每个奇异值单独做一个标量迭代。(科学空间)

我们完整推一次。

第一步：假设 $X_t$ 的 SVD

设

那么

进一步

第二步：分别代入各项

第一项：

第二项：

第三项：

第三步：合并

所以

因为 $\Sigma_t$ 是对角矩阵，对角矩阵的幂就是把每个对角元单独取幂。
如果某个奇异值是 $\sigma_t$，那么它的更新就是

这就是文章那句“问题简化成单个奇异值的迭代”的完整原因。(科学空间)

13. 为什么目标是把奇异值迭代到 1

因为如果

最后我们希望 $X_t$ 收敛到

那就意味着希望

也就是说，每个奇异值都要满足

所以整个矩阵迭代的本质，就是设计一个标量函数 $g$，让所有初始奇异值都尽快被送到 1 附近。这就是文章后面把系数选择写成优化问题的原因。(科学空间)

14. 为什么初始化要除以 Frobenius 范数

文章取

理由是：缩放不会改变 $U,V$，只会缩放奇异值；而除以 $|M|_F$ 后，初始奇异值会被压到更标准的范围内，文章表述为落在 $[0,1]$ 之间，从而更适合迭代。(科学空间)

更直观地说：

  • 我们的多项式近似是围绕 $t=1$ 展开的；
  • 所以希望 $X_t^\top X_t$ 的谱不要太离谱；
  • 先把 $M$ 缩到一个温和尺度，迭代更稳。

15. 把整件事串成一句真正“懂了”的话

Muon 的核心不是“神秘地做正交化”，而是：

  1. 把梯度/动量矩阵 $M$ 做 SVD：
  2. 只保留方向 $U,V$，把尺度 $\Sigma$ 统一成 $I$：
  3. 这等价于求离 $M$ 最近的正交矩阵；
  4. 实现时不直接做 SVD，而是通过迭代不断把奇异值 $\sigma_i$ 推向 1。

所以 Muon 的“正交化”本质上不是你在线代课上学的 Gram–Schmidt 那种“把一组向量一个个变垂直”，而是：

把整个梯度矩阵投到“正交矩阵流形”附近，得到一个最接近原矩阵、但奇异值全被标准化的更新方向。 这正是文章强调的“从向量到矩阵的本质跨越”。(科学空间)

你要是愿意，我下一条可以继续做两件事里的任意一个：

  1. 用一个 $(2\times 2)$ 的具体数值例子，把上面每一步真的算出来；
  2. 解释 这和你熟悉的 Gram–Schmidt / QR 正交化到底是什么关系、又有什么本质区别。