科学空间|Scientific Spaces

四个月前，我们发布了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 在推特上指出笔者的方案实际计算中存在误差偏大的问题。本文来具体分析一下这个问题，并给出一个更高效、误差更低的新方案。

点击阅读全文...

这篇文章我们来推导$\newcommand{msign}{\mathop{\text{msign}}}\msign$算子的求导公式。如果读者想要像《Test-Time Training Done Right》一样，将TTT和Muon结合起来，那么本文可能会对你有帮助。

两种定义

本文依然假设大家已经对$\msign$有所了解，如果还没有，可以先移步阅读《Muon优化器赏析：从向量到矩阵的本质跨越》和《msign算子的Newton-Schulz迭代（上）》。现设有矩阵$\boldsymbol{M}\in\mathbb{R}^{n\times m}$，那么
\begin{equation}\boldsymbol{U},\boldsymbol{\Sigma},\boldsymbol{V}^{\top} = \text{SVD}(\boldsymbol{M}) \quad\Rightarrow\quad \msign(\boldsymbol{M}) = \boldsymbol{U}_{[:,:r]}\boldsymbol{V}_{[:,:r]}^{\top}\end{equation}
其中$\boldsymbol{U}\in\mathbb{R}^{n\times n},\boldsymbol{\Sigma}\in\mathbb{R}^{n\times m},\boldsymbol{V}\in\mathbb{R}^{m\times m}$，$r$是$\boldsymbol{M}$的秩。简单来说，$\msign$就是把矩阵的所有非零奇异值都变成1后所得的新矩阵。

点击阅读全文...

前面我们用了两篇文章《msign算子的Newton-Schulz迭代（上）》和《msign算子的Newton-Schulz迭代（下）》讨论了矩阵的$\newcommand{msign}{\mathop{\text{msign}}}\newcommand{sign}{\mathop{\text{sign}}}\newcommand{clip}{\mathop{\text{clip}}}\newcommand{mclip}{\mathop{\text{mclip}}}\msign$算子的数值计算，这篇文章我们来关注“奇异值裁剪（Singular Value Clipping）”运算，它最近在 @_arohan_ 的推特上引起了热议，我们此前在《高阶muP：更简明但更高明的谱条件缩放》也提到过，接下来我们简称为$\mclip$。

基本概念

对于标量$x$，$\clip$运算定义为
\begin{equation}\clip(x) = \max(\min(x, 1), -1) = \left\{\begin{aligned}1, &\quad x\geq 1 \\
x, &\quad x\in(-1, 1)\\
-1, &\quad x\leq -1
\end{aligned}\right.\end{equation}

点击阅读全文...

在上文《msign算子的Newton-Schulz迭代（上）》中，我们试图为$\mathop{\text{msign}}$算子寻找更好的Newton-Schulz迭代，以期在有限迭代步数内能达到尽可能高的近似程度，这一过程又可以转化为标量函数$\mathop{\text{sign}}(x)$寻找同样形式的多项式迭代。当时，我们的求解思路是用Adam优化器端到端地求一个局部最优解，虽然有效但稍显粗暴。

而在几天前，arXiv新出了一篇论文《The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm》，作者运用了一系列精妙的数学结论，以优雅且硬核的方式给出了更漂亮的答案。本文让我们一起欣赏和学习一番这篇精彩的论文。

问题描述

相关背景和转化过程我们就不再重复了，直接给出我们要求解的问题是
\begin{equation}\mathop{\text{argmin}}_f d(f(x),1)\end{equation}

点击阅读全文...

在之前的《Muon优化器赏析：从向量到矩阵的本质跨越》、《Muon续集：为什么我们选择尝试Muon？》等文章中，我们介绍了一个极具潜力、有望替代Adam的新兴优化器——“Muon”。随着相关研究的不断深入，Muon优化器受到的关注度也在日益增加。

了解过Muon的读者都知道，Muon的核心运算是$\newcommand{msign}{\mathop{\text{msign}}}\msign$算子，为其寻找更高效的计算方法是学术社区的一个持续目标。本文将总结一下它的最新进展。

写在前面

$\msign$的定义跟SVD密切相关。假设矩阵$\boldsymbol{M}\in\mathbb{R}^{n\times m}$，那么
\begin{equation}\boldsymbol{U},\boldsymbol{\Sigma},\boldsymbol{V}^{\top} = \text{SVD}(\boldsymbol{M}) \quad\Rightarrow\quad \msign(\boldsymbol{M}) = \boldsymbol{U}_{[:,:r]}\boldsymbol{V}_{[:,:r]}^{\top}\end{equation}
其中$\boldsymbol{U}\in\mathbb{R}^{n\times n},\boldsymbol{\Sigma}\in\mathbb{R}^{n\times m},\boldsymbol{V}\in\mathbb{R}^{m\times m}$，$r$是$\boldsymbol{M}$的秩。简单来说，$\msign$就是把矩阵的所有非零奇异值都变成1后所得的新矩阵。

点击阅读全文...

在文章《初探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}

点击阅读全文...

本文解读一下我们最新的技术报告《Muon is Scalable for LLM Training》，里边分享了我们之前在《Muon优化器赏析：从向量到矩阵的本质跨越》介绍过的Muon优化器的一次较大规模的实践，并开源了相应的模型（我们称之为“Moonlight”，目前是一个3B/16B的MoE模型）。我们发现了一个比较惊人的结论：在我们的实验设置下，Muon相比Adam能够达到将近2倍的训练效率。

Muon的Scaling Law及Moonlight的MMLU表现

优化器的工作说多不多，但说少也不少，为什么我们会选择Muon来作为新的尝试方向呢？已经调好超参的Adam优化器，怎么快速切换到Muon上进行尝试呢？模型Scale上去之后，Muon与Adam的性能效果差异如何？接下来将分享我们的思考过程。

点击阅读全文...