科学空间|Scientific Spaces

这段时间，相信很多读者已经刷到过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，跟优化器是否有效，没有任何关系^[原文]”。

点击阅读全文...

在博客《AdamW的Weight RMS的渐近估计（上）》中，我们推导了AdamW训练出来的模型权重的RMS渐近表达式。不过，那会我们假设了Weight Decay和学习率在整个训练过程中是固定的，这跟实际训练并不完全吻合，所以这篇文章我们将之前的结论推广成动态版。

所谓动态版，即允许Weight Decay和学习率都随着训练步数的增加而变化，比如经典的Cosine Decay、WSD（Warmup Stable Decay）等，从而让结论更为通用。

步骤之一

我们的出发点还是AdamW的定义：
\begin{equation}\text{Adam}\color{skyblue}{\text{W}}:=\left\{\begin{aligned}
&\boldsymbol{m}_t = \beta_1 \boldsymbol{m}_{t-1} + \left(1 - \beta_1\right) \boldsymbol{g}_t\\
&\boldsymbol{v}_t = \beta_2 \boldsymbol{v}_{t-1} + \left(1 - \beta_2\right) \boldsymbol{g}_t^2\\
&\hat{\boldsymbol{m}}_t = \boldsymbol{m}_t\left/\left(1 - \beta_1^t\right)\right.\\
&\hat{\boldsymbol{v}}_t = \boldsymbol{v}_t\left/\left(1 - \beta_2^t\right)\right.\\
&\boldsymbol{u}_t =\hat{\boldsymbol{m}}_t\left/\left(\sqrt{\hat{\boldsymbol{v}}_t} + \epsilon\right)\right.\\
&\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t (\boldsymbol{u}_t \color{skyblue}{ + \lambda_t \boldsymbol{\theta}_{t-1}})
\end{aligned}\right.\end{equation}

点击阅读全文...

前四篇文章我们求解了几个具体的给参数加等式约束的最速下降问题，其中第三、四篇的问题没法找到解析解，所以笔者提出了相应的不动点迭代法。其中的其中，第三篇文章《流形上的最速下降：3. Muon + Stiefel》所研究的“Stiefel流形上的Muon”，问题提出自Jeremy Bernstein的《Orthogonal manifold》一文。

对于这个问题，Jeremy Bernstein最后也给出了一个自己的解法，笔者称之为“对偶梯度下降（Dual Gradient Descent）”，也颇为值得学习一番。

基本概念

Jeremy Bernstein的解法，最后发表在Thinking Machines Lab的博客《Modular Manifolds》中，是该实验室的第二篇博客，文章中将它称为“对偶上升（Dual Ascent）”，但笔者这里还是结合前四篇的内容，将其称为“对偶梯度下降”。

点击阅读全文...

不知道大家有没有发现一个有趣的细节，Muon和MuP都是“Mu”开头，但两个“Mu”的原意完全不一样，前者是“MomentUm Orthogonalized by Newton-Schulz”，后者是“Maximal Update Parametrization”，可它们俩之间确实有着非常深刻的联系。也就是说，Muon和MuP有着截然不同的出发点，但最终都走向了相同的方向，甚至无意间取了相似的名字，似乎真应了那句“冥冥中自有安排”。

言归正传。总之，笔者在各种机缘巧合之下，刚好同时学习到了Muon和MuP，这大大加深了笔者对模型优化的理解，同时也让笔者开始思考关于模型优化更本质的原理。经过一段时间的试错，算是有些粗浅的收获，在此跟大家分享一下。

写在前面

按照提出时间的先后顺序，是先有MuP再有Muon，但笔者的学习顺序正好反过来，先学习了Muon然后再学习MuP，事后来看，这也不失为一个不错的学习顺序。

点击阅读全文...

在《为什么Adam的Update RMS是0.2？》中，我们用平均场近似估计了Adam的Update RMS。不久后，读者 @EIFY 指出相同的结果已经出现在论文《Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks》中。阅读后，笔者发现其中不仅包含了Update RMS的估计，还包含了Weight RMS的估计。

也就是说，AdamW训出来的模型，其权重的RMS是可以事先估计出来一个渐近结果的。大家会不会觉得这个结论有点意外？反正笔者第一次看到它是颇为意外的，直觉上权重模长是模型根据训练集自己学出来的，结果它告诉我这已经隐藏在优化器的超参中，可谓很反直觉了。

这篇文章我们还是用平均场近似方法，来复现对Weight RMS的渐近估计。

点击阅读全文...

我们在《重新思考学习率与Batch Size（二）：平均场》中提到，关注SignSGD的原因之一是我们通常将它作为Adam的理论近似，这是Adam做理论分析时常用的简化策略。除了分析学习率的场景外，在《配置不同的学习率，LoRA还能再涨一点？》、《初探MuP：超参数的跨模型尺度迁移规律》等地方我们也用了这个简化。

然而，SignSGD真是Adam的良好近似吗？一个明显差异是SignSGD的Update RMS总是1，而Adam并非如此。笔者发现，导致这一差异的核心原因是动量，它普遍存在于Adam、Lion、Muon等优化器中。所以，本文我们来考察动量——更广义地说是EMA——的影响。

问题分析

从Adam的视角看，SignSGD对应$\beta_1=\beta_2=0$这个特例，或者对应于Adam的第一步更新量（不管$\beta_1,\beta_2$如何）。因此，我们认为它跟Adam肯定有一些共性，能够捕捉到一些通用的规律。

点击阅读全文...

前两篇文章《重新思考学习率与Batch Size（一）：现状》和《重新思考学习率与Batch Size（二）：平均场》中，我们主要是提出了平均场方法，用以简化学习率与Batch Size的相关计算。当时我们分析的优化器是SGD、SignSGD和SoftSignSGD，并且主要目的是简化，本质上没有新的结论。

然而，在如今的优化器盛宴中，怎能少得了Muon的一席之地呢？所以，这篇文章我们就来尝试计算Muon的相关结论，看看它的学习率与Batch Size的关系是否会呈现出新的规律。

基本记号

众所周知，Muon的主要特点就是非Element-wise的更新规则，所以之前在《当Batch Size增大时，学习率该如何随之变化？》和《Adam的epsilon如何影响学习率的Scaling Law？》的Element-wise的计算方法将完全不可用。但幸运的是，上篇文章介绍的平均场依然好使，只需要稍微调整一下细节。

点击阅读全文...

上文《重新思考学习率与Batch Size（一）：现状》末尾我们说到，对于SignSGD、SoftSignSGD等$\tilde{\boldsymbol{\varphi}}_B$非线性依赖于$\tilde{\boldsymbol{g}}_B$的情形，计算过程的心智负担相当沉重，并且面临难以推广的困境。为此，笔者投入了一些精力去尝试简化其中的推导，万幸有些许收获，其中的关键思路便是本文的主题——平均场。

平均场是物理中常见的近似计算方法，它没有固定的形式，但大体思想就是将求平均移到函数之内。事实上，在《为什么Adam的Update RMS是0.2？》中我们就已经窥见过平均场的魅力，而在这篇文章中，我们再来见识它在计算SignSGD/SoftSignSGD的学习率规律上的奇效。

方法大意

沿着上文的记号，对于SignSGD我们有$\newcommand{sign}{\mathop{\text{sign}}}\tilde{\boldsymbol{\varphi}}_B=\sign(\tilde{\boldsymbol{g}}_B)$，我们需要先计算$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]$和$\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$，继而可以算出
\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\eta^* \approx \frac{\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^{\top}\boldsymbol{g}}{\tr(\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]\boldsymbol{H})}\label{eq:eta-opt}\end{equation}

点击阅读全文...