这篇文章我们来推导$\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来计算mclip(奇异值裁剪)
By 苏剑林 | 2025-06-07 | 3004位读者 | 引用前面我们用了两篇文章《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迭代(下)
By 苏剑林 | 2025-06-05 | 4158位读者 | 引用在上文《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}
msign算子的Newton-Schulz迭代(上)
By 苏剑林 | 2025-05-11 | 13933位读者 | 引用在之前的《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:更简明但更高明的谱条件缩放
By 苏剑林 | 2025-03-24 | 19223位读者 | 引用在文章《初探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续集:为什么我们选择尝试Muon?
By 苏剑林 | 2025-02-27 | 39167位读者 | 引用本文解读一下我们最新的技术报告《Muon is Scalable for LLM Training》,里边分享了我们之前在《Muon优化器赏析:从向量到矩阵的本质跨越》介绍过的Muon优化器的一次较大规模的实践,并开源了相应的模型(我们称之为“Moonlight”,目前是一个3B/16B的MoE模型)。我们发现了一个比较惊人的结论:在我们的实验设置下,Muon相比Adam能够达到将近2倍的训练效率。
优化器的工作说多不多,但说少也不少,为什么我们会选择Muon来作为新的尝试方向呢?已经调好超参的Adam优化器,怎么快速切换到Muon上进行尝试呢?模型Scale上去之后,Muon与Adam的性能效果差异如何?接下来将分享我们的思考过程。
Muon优化器赏析:从向量到矩阵的本质跨越
By 苏剑林 | 2024-12-10 | 60012位读者 | 引用随着LLM时代的到来,学术界对于优化器的研究热情似乎有所减退。这主要是因为目前主流的AdamW已经能够满足大多数需求,而如果对优化器“大动干戈”,那么需要巨大的验证成本。因此,当前优化器的变化,多数都只是工业界根据自己的训练经验来对AdamW打的一些小补丁。
不过,最近推特上一个名为“Muon”的优化器颇为热闹,它声称比AdamW更为高效,且并不只是在Adam基础上的“小打小闹”,而是体现了关于向量与矩阵差异的一些值得深思的原理。本文让我们一起赏析一番。
最近评论