科学空间|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）”，但笔者这里还是结合前四篇的内容，将其称为“对偶梯度下降”。

点击阅读全文...

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

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

基本记号

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

点击阅读全文...

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

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

点击阅读全文...

四个月前，我们发布了Moonlight，在16B的MoE模型上验证了Muon优化器的有效性。在Moonlight中，我们确认了给Muon添加Weight Decay的必要性，同时提出了通过Update RMS对齐来迁移Adam超参的技巧，这使得Muon可以快速应用于LLM的训练。然而，当我们尝试将Muon进一步拓展到千亿参数以上的模型时，遇到了新的“拦路虎”——MaxLogit爆炸。

为了解决这个问题，我们提出了一种简单但极其有效的新方法，我们称之为“QK-Clip”。该方法从一个非常本质的角度去看待和解决MaxLogit现象，并且无损模型效果，这成为我们最新发布的万亿参数模型“Kimi K2”的关键训练技术之一。

问题描述

我们先来简单介绍一下MaxLogit爆炸现象。回顾Attention的定义
\begin{equation}\boldsymbol{O} = softmax(\boldsymbol{Q}\boldsymbol{K}^{\top})\boldsymbol{V}\end{equation}

点击阅读全文...

前面我们在《通过msign来计算奇异值裁剪mclip（上）》讨论了奇异值裁剪$\newcommand{mclip}{\mathop{\text{mclip}}}\mclip$的数值计算，核心思路来自 @leloykun 的文章《Numerically Stable Spectral Clipping Via Newton-Schulz Iteration》（现已重新修订和改名），通过寻找基于$\newcommand{msign}{\mathop{\text{msign}}}\msign$的表达式来避免另外寻找Newton-Schulz迭代，在文章中笔者提出了一个计算量更低的嵌套$\msign$方案。

不过前两天，@leloykun 在推特上指出笔者的方案实际计算中存在误差偏大的问题。本文来具体分析一下这个问题，并给出一个更高效、误差更低的新方案。

点击阅读全文...