科学空间|Scientific Spaces

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

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

写在前面

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

点击阅读全文...

我们在《重新思考学习率与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}

点击阅读全文...

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

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

方法大意

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

点击阅读全文...

在文章《初探MuP：超参数的跨模型尺度迁移规律》中，我们基于前向传播、反向传播、损失增量和特征变化的尺度不变性推导了MuP（Maximal Update Parametrization）。可能对于部分读者来说，这一过程还是显得有些繁琐，但实际上它比原始论文已经明显简化。要知道，我们是在单篇文章内相对完整地介绍的MuP，而MuP的论文实际上是作者Tensor Programs系列论文的第5篇！

不过好消息是，作者在后续的研究《A Spectral Condition for Feature Learning》中，发现了一种新的理解方式（下称“谱条件”），它比MuP的原始推导和笔者的推导都更加直观和简洁，但却能得到比MuP更丰富的结果，可谓MuP的高阶版本，简明且不失高明的代表作。

准备工作

顾名思义，谱条件（Spectral Condition）跟谱范数（Spectral Norm）相关，它的出发点是谱范数的一个基本不等式：
\begin{equation}\Vert\boldsymbol{x}\boldsymbol{W}\Vert_2\leq \Vert\boldsymbol{x}\Vert_2 \Vert\boldsymbol{W}\Vert_2\label{neq:spec-2}\end{equation}

点击阅读全文...

众所周知，完整训练一次大型LLM的成本是昂贵的，这就决定了我们不可能直接在大型LLM上反复测试超参数。一个很自然的想法是希望可以在同结构的小模型上仔细搜索超参数，找到最优组合后直接迁移到大模型上。尽管这个想法很朴素，但要实现它并不平凡，它需要我们了解常见的超参数与模型尺度之间的缩放规律，而MuP正是这个想法的一个实践。

MuP，有时也写$\mu P$，全名是Maximal Update Parametrization，出自论文《Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer》，随着LLM训练的普及，它逐渐已经成为了科学炼丹的事实标配之一。

方法大意

在接入主题之前，必须先吐槽一下MuP原论文写得实在太过晦涩，并且结论的表达也不够清晰，平白增加了不少理解难度，所以接下来笔者尽量以一种（自认为）简明扼要的方式来复现MuP的结论。

点击阅读全文...

上一篇文章《当Batch Size增大时，学习率该如何随之变化？》我们从多个角度讨论了学习率与Batch Size之间的缩放规律，其中对于Adam优化器我们采用了SignSGD近似，这是分析Adam优化器常用的手段。那么一个很自然的问题就是：用SignSGD来近似Adam究竟有多科学呢？

我们知道，Adam优化器的更新量分母会带有一个$\epsilon$，初衷是预防除零错误，所以其值通常很接近于零，以至于我们做理论分析的时候通常选择忽略掉它。然而，当前LLM的训练尤其是低精度训练，我们往往会选择偏大的$\epsilon$，这导致在训练的中、后期$\epsilon$往往已经超过梯度平方大小，所以$\epsilon$的存在事实上已经不可忽略。

因此，这篇文章我们试图探索$\epsilon$如何影响Adam的学习率与Batch Size的Scaling Law，为相关问题提供一个参考的计算方案。

点击阅读全文...