科学空间|Scientific Spaces

众所周知，我们很早就开始尝试将Muon用于大规模LLM的训练。特别地，在《Muon续集：为什么我们选择尝试Muon？》中，我们提出了“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训练，参数是$\beta_1=0.9,\beta_2=0.95$。由于共性很明显，所以这大概率不是巧合，因此笔者尝试分析背后的原理。

点击阅读全文...

在之前的文章《当Batch Size增大时，学习率该如何随之变化？》和《Adam的epsilon如何影响学习率的Scaling Law？》中，我们从理论上讨论了学习率随Batch Size的变化规律，其中比较经典的部分是由OpenAI提出的展开到二阶的分析。然而，当我们要处理非SGD优化器时，这套分析方法的计算过程往往会相当复杂，有种无从下手的感觉。

接下来的几篇文章，笔者将重新整理和思考上述文章中的相关细节，尝试简化其中的一些推导步骤，给出一条更通用、更轻盈的推导路径，并且探讨推广到Muon优化器的可能性。

方法大意

首先回顾一下之前的分析方法。在《当Batch Size增大时，学习率该如何随之变化？》中，我们介绍了多种分析学习率与Batch Size规律的思路，其中OpenAI在《An Empirical Model of Large-Batch Training》提出的二阶近似分析占了主要篇幅，本文也是沿用同样的思路。

点击阅读全文...

看完了前三篇的读者，想必已经熟悉本系列的“套路”——先提出更新量的约束，寻找最速下降方向，接着再给参数也加上约束，寻找新的最速下降方向。在求解参数约束问题时，我们采用的是“一阶近似够用”原则来简化约束形式，这在几何上对应于“切空间”。然后，我们用待定系数法转化无约束形式来写出解析解，最后再数值求解待定系数。

这篇文章我们再来求解一个新例子——谱球面约束下的Muon——它是第一篇文章《流形上的最速下降：1. SGD + 超球面》的类比推广，当我们希望参数的谱范数始终不变时可以考虑它。当然，也可以单纯作为一道练习题来练手。

问题描述

在《流形上的最速下降：2. Muon + 正交》和《流形上的最速下降：3. Muon + Stiefel》中，我们已经详细讨论了Muon与正交约束的碰撞，所以相关背景我们就不展开了，直接给出问题形式：
\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\max_{\boldsymbol{\Phi}} \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{\Phi}\Vert_2 = 1,\,\, \Vert\boldsymbol{W}\Vert_2 = 1,\,\,\Vert\boldsymbol{W} - \eta \boldsymbol{\Phi}\Vert_2=1\end{equation}

点击阅读全文...

今天水一点轻松的内容，它基于笔者这两天意识到的一个恒等式。这个恒等式实际上很简单，但初看之下会有点意料之外的感觉，所以来记录一下。

基本结果

我们知道$\newcommand{relu}{\mathop{\text{relu}}}\relu(x) = \max(x, 0)$，容易证明如下恒等式
\begin{equation}x = \relu(x) - \relu(-x)\end{equation}
如果$x$是一个向量，那么上式就更直观了，$\relu(x)$是提取出$x$的正分量，$- \relu(-x)$是提取出$x$的负分量，两者相加就得到原本的向量。

点击阅读全文...

上回说到，当我们把优化对象从向量参数转移到矩阵参数，并选用更适合矩阵的谱范数约束后，Muon优化器便自然而然地出现了。进一步地，我们考虑了给参数加上正交约束后的最速下降方向，这其中又分方阵和非方阵两部分讨论，其中方阵的求解我们在上一篇文章已经完成，但非方阵部分依然悬而未决。

本文的目标，则是把非方阵部分的求解补上，使得正交约束下的优化得以完全解决。

任务信息

先简单回顾一下上文《流形上的最速下降：2. Muon + 正交》的结果。我们要求解的目标是
\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\max_{\boldsymbol{\Phi}} \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{\Phi}\Vert_2 = 1,\,\, \boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I},\,\,(\boldsymbol{W} - \eta \boldsymbol{\Phi})^{\top}(\boldsymbol{W} - \eta \boldsymbol{\Phi})=\boldsymbol{I}\end{equation}

点击阅读全文...

本文继续我们的约束优化系列。在上文《流形上的最速下降：1. SGD + 超球面》中，我们重温了优化器的“最小作用量”原理，提出不同优化器的核心差异在于给更新量施加的不同约束，如果这个约束是欧几里得范数，那么对应的最速下降便是SGD。进一步地，我们还讨论了同时给参数增加模长约束后的结果，这构成了超球面流形上的最速下降。

不过，上文只能算是“热身”，因为它处理的是相对简单的向量参数优化。本文正式进入更具挑战性的部分——优化参数从向量变成矩阵，并且增量约束改为谱范数，由此衍生出Muon优化器；接着，我们再给参数添加正交约束，这将得到正交流形下的Muon优化器。

命题描述

设待优化参数具有矩阵形式$\boldsymbol{W}\in\mathbb{R}^{n\times m}$，不失一般性，设$n\geq m$。根据上一篇文章的“最小作用量”原理，我们得出最速下降的增量$\Delta\boldsymbol{W}$应该满足
\begin{equation}\min_{\Delta \boldsymbol{W}} \mathcal{L}(\boldsymbol{W} +\Delta\boldsymbol{W}) \qquad \text{s.t.}\qquad \rho(\Delta\boldsymbol{W})\leq \eta\end{equation}

点击阅读全文...

类似“梯度的反方向是下降最快的方向”的描述，经常用于介绍梯度下降（SGD）的原理。然而，这句话是有条件的，比如“方向”在数学上是单位向量，它依赖于“范数（模长）”的定义，不同范数的结论也不同，Muon实际上就是给矩阵参数换了个谱范数，从而得到了新的下降方向。又比如，当我们从无约束优化转移到约束优化时，下降最快的方向也未必是梯度的反方向。

为此，在这篇文章中，我们将新开一个系列，以“约束”为主线，重新审视“最速下降”这一命题，探查不同条件下的“下降最快的方向”指向何方。

优化原理

作为第一篇文章，我们先从SGD出发，理解“梯度的反方向是下降最快的方向”这句话背后的数学意义，然后应用于超球面上的优化。不过在此之前，笔者还想带大家重温一下《Muon续集：为什么我们选择尝试Muon？》所提的关于优化器的“最小作用量原理（Least Action Principle）”。

点击阅读全文...

上一篇文章《矩阵平方根和逆平方根的高效计算》中，笔者从$\newcommand{mcsgn}{\mathop{\text{mcsgn}}}\mcsgn$算子出发，提出了一种很漂亮的矩阵平方根和逆平方根的计算方法。比较神奇的是，该方案经过化简之后，最终公式已经看不到最初$\mcsgn$形式的样子。这不禁引发了更深层的思考：该方案更本质的工作原理是什么？是否有推广到任意$r$次方根的可能性？

沿着这个角度进行分析后，笔者惊喜地发现，我们可以从一个更简单的角度去理解之前的迭代算法，并且在新角度下可以很轻松推广到任意$r$次方根和逆$r$次方根的计算。接下来我们将分享这一过程。

前情回顾

设$\boldsymbol{G}\in\mathbb{R}^{m\times n}$是任意矩阵，$\boldsymbol{P}\in\mathbb{R}^{n\times n}$是任意特征值都在$[0,1]$内的矩阵，上一篇文章给出：
\begin{gather}
\boldsymbol{G}_0 = \boldsymbol{G}, \quad \boldsymbol{P}_0 = \boldsymbol{P} \notag\\[6pt]
\boldsymbol{G}_{t+1} = \boldsymbol{G}_t(a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2) \label{eq:r2-rsqrt}\\[6pt]
\boldsymbol{P}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2)^2\boldsymbol{P}_t \label{eq:r3-rsqrt}\\[6pt]
\lim_{t\to\infty} \boldsymbol{G}_t = \boldsymbol{G}\boldsymbol{P}^{-1/2}\notag
\end{gather}

点击阅读全文...