科学空间|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，跟优化器是否有效，没有任何关系^[原文]”。

点击阅读全文...

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

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

基本概念

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

点击阅读全文...

在《高阶MuP：更简明但更高明的谱条件缩放》的“近似估计”一节中，我们曾“预支”了一个结论：“一个服从标准正态分布的$n\times m$大小的随机矩阵，它的谱范数大致是$\sqrt{n}+\sqrt{m}$。”

这篇文章我们来补充讨论这个结论，给出随机矩阵谱范数的快速估计方法。

随机矩阵论

设有随机矩阵$\boldsymbol{W}\in\mathbb{R}^{n\times m}$，每个元素都是从标准正态分布$\mathcal{N}(0,1)$中独立重复地采样出来的，要求估计$\boldsymbol{W}$的谱范数，也就是最大奇异值，我们以$\mathbb{E}[\Vert\boldsymbol{W}\Vert_2]$为最终的估计结果。

点击阅读全文...

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

这篇文章我们再来求解一个新例子——谱球面约束下的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}

点击阅读全文...

上回说到，当我们把优化对象从向量参数转移到矩阵参数，并选用更适合矩阵的谱范数约束后，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}

点击阅读全文...

上一篇文章《矩阵平方根和逆平方根的高效计算》中，笔者从$\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}

点击阅读全文...

设$\boldsymbol{P}\in\mathbb{R}^{n\times n}$是一个特征值都是非负实数的$n$阶方阵，本文来讨论它的平方根$\boldsymbol{P}^{1/2}$和逆平方根$\boldsymbol{P}^{-1/2}$的计算。

基本概念

矩阵$\boldsymbol{P}$的平方根，指的是满足$\boldsymbol{X}^2=\boldsymbol{P}$的矩阵$\boldsymbol{X}$。我们知道正数都有两个平方根，因此不难想象矩阵平方根一般也不唯一。不过，“算术平方根”是唯一的，一个正数的算术平方根是正的那个平方根，类似地，我们将$\boldsymbol{P}$的特征值全是非负数的那个平方根称为算术平方根。本文要求的矩阵平方根，默认都是指算术平方根。

点击阅读全文...