在深度学习中最常用的优化方法是梯度下降方法及其变体。在过去很长一段时间中，Adam优化器是NLP社区的默认选择，在ViT出现之后，CV方面的工作也逐渐开始使用Adam和Adam的变体（在ViT之前，一种常见的观点是Adam不适用于Vision任务）。

最近Muon优化器在Kimi的新工作带动下又火了一把，相较于Adam优化器需要同时维护一阶、二阶矩估计量，Muon只需要维护一份梯度的动量估计，因此在大规模训练中有很大优势。最近笔者顺着Muon的reference看了Jeremy Bernstein在优化的一些文章，觉得很有意思，因此写这篇文章梳理一下这一系列工作的部分精要。本文的核心论点是：使用诱导范数来约束梯度更新，可以推导出最近的一些新出现的优化方法，这也可能是未来深度学习优化的一个有潜力的探索方向。

梯度下降（Recap）

当前深度学习优化算法的基石是梯度下降。之前笔者写过一篇拙文（自然梯度（二）：黎曼距离下的最速下降）整理过梯度下降的推导，核心的结论是：当我们假设参数空间是一个欧几里得空间、参数的距离可以用欧几里得距离来衡量时，我们在某个点约束 $\|\Delta\theta\|_2\le\epsilon(\epsilon>0)$ 时， $\Delta\theta$ 取梯度的反方向时可以让目标函数下降最多（具体的证明请参阅上述引文）。

使用梯度下降最大的问题是，它实际上忽略了模型的结构。换句话说，梯度下降相当于将模型所有参数展平为1维向量，并且用向量2范数来衡量每次更新的「步长」。这种抽象是实用的，但是也存在一定的问题。两组参数有可能在欧几里得空间中距离很近，但是诱导的模型输出空间距离很远。造成的结果就是更新的方向实际上不是目标函数下降最快的方向。

这个问题要如何解决呢？在自然梯度（二）：黎曼距离下的最速下降中，我们介绍了自然梯度方法，即使用Fisher信息矩阵的逆作为梯度的pre-conditioner来矫正梯度下降的方向，从原理上是使用参数更新前后引导的概率分布的KL散度作为每次更新的步长约束。但是对于常见的深度神经网络来说，这样做仍然是不切实际的，因为FIM是一个 $N\times N$ 的大矩阵（其中 $N$ 是参数量），对于这么大的矩阵存储或求逆都是很难做到的。

诱导范数作为步长约束

是否有一种更「廉价」的方法，可以考虑模型的参数结构，同时将参数的变化对于输出的影响作为约束呢？

幸运的是，对于当下最流行的神经网络（e.g., Transformer）而言，模型往往可以拆解为很多小模块，其中最常见的是Linear模块（线性映射，这里忽略bias term）

$$ f(\boldsymbol{x};\boldsymbol{W})=\boldsymbol{Wx},\ \boldsymbol{W}\in\mathbb{R}^{n\times m},\boldsymbol{x}\in \mathbb{R}^{m}\\ $$

在标准的Transformer中，Attention、FFN、LM分类器都是由Linear模块组成的，Embedding从数学原理上也是输入为one-hot encoding的线性映射。假设现在对于某个Linear模块的参数 $\boldsymbol{W}$ 做 $\Delta\boldsymbol{W}$ 的更新（ $\boldsymbol{W}'\leftarrow \boldsymbol{W}+\Delta\boldsymbol{W}$ ），我们需要衡量这个更新对于最终输出的影响是多少（从而可以约束这个影响）。由于神经网络比较复杂，衡量 $\Delta\boldsymbol{W}$ 对于最终目标函数的影响是相对繁琐的，但我们可以退而求其次，衡量 $\Delta\boldsymbol{W}$ 对于这个Linear模块的输出 $\boldsymbol{Wx}$ 的影响。考虑线性模块的输入与输出空间的距离都使用欧几里得范数 $\|\cdot\|_{\ell_2}$ 衡量，那么这个约束可以通过如下不等式实现

$$ \|\Delta\boldsymbol{W}\boldsymbol{x}\|_{\ell_2} \le {\color[rgb]{0, 0.5, 0.8}{\|\Delta\boldsymbol{W}\|_{\ell_2\to\ell_2}}}\|\boldsymbol{x}\|_{\ell_2} $$ 这里的 ${\color[rgb]{0, 0.5, 0.8}{\|\Delta\boldsymbol{W}\|_{\ell_2\to\ell_2}}}$ 是矩阵2范数。这个不等式告诉我们，如果约束了参数更新量的谱范数（不等式右侧），也就约束了更新前后这个线性模块输出的变化量。

假设现在需要优化的神经网络是由一系列的线性模块堆叠组成（e.g., MLP），我们可以参照梯度下降的推导构造如下的更新¹

$$ \underset{\Delta\boldsymbol{W}_1,\ldots,\Delta\boldsymbol{W}_L}{\text{arg min}}\left[ \sum_{l=1}^L {\langle \boldsymbol{G}_l, \Delta\boldsymbol{W}_l \rangle}_F + \frac{\lambda}{2}\max_{l=1}^L{\color[rgb]{0, 0.5, 0.8}{\|\Delta\boldsymbol{W}_l\|^2_{\ell_2\to\ell_2}}} \right]\\ $$

这里 $\boldsymbol{G}_l$ 表示 $\boldsymbol{W}_l$ 对应的梯度矩阵（布局与原参数矩阵相同）， ${\langle \cdot, \cdot \rangle}_F$ 表示Frobenius内积（对矩阵而言，逐元素相乘求和）。这里之所以使用 $\max_{l=1}^L$ （而不是直接求和），是因为我们引入这个约束时希望目标函数在 $\Delta\boldsymbol{W}_l$ 变化下，能够保持平滑的性质²，因此需要bound所有参数矩阵更新量的谱范数的最大值。

我们来逐步推导这个最小值成立时的 $\Delta\boldsymbol{W}_1,\ldots,\Delta\boldsymbol{W}_L$ 取值³。为了方便，把每个 $\Delta\boldsymbol{W}_l$ 拆解成大小和方向两部分： $\Delta\boldsymbol{W}_l=c_l\boldsymbol{T}_l(c_l\triangleq\|\Delta\boldsymbol{W}_l\|_{\ell_2\to\ell_2})$

（为了可读性，下面的 $\|\cdots\|$ 均表示谱范数 $\|\cdot\|_{\ell_2\to\ell_2}$ ）

$$ \begin{align} &\underset{\Delta\boldsymbol{W}_1,\ldots,\Delta\boldsymbol{W}_L}{\text{min}}\left[ \sum_{l=1}^L {\langle \boldsymbol{G}_l, \Delta\boldsymbol{W}_l \rangle}_F + \frac{\lambda}{2}\max_{l=1}^L\|\Delta\boldsymbol{W}_l\|^2 \right]\\ &=\underset{c_1,\ldots,c_L\ge 0}{\text{min}}\left[ \sum_{l=1}^L c_l\min_{\|\boldsymbol{T}_l\|=1}{\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F + \frac{\lambda}{2}\max_{l=1}^Lc_l^2 \right]\\ &=\underset{c_1,\ldots,c_L\ge 0}{\text{min}}\left[ -\sum_{l=1}^L c_l\max_{\|\boldsymbol{T}_l\|=1}{\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F + \frac{\lambda}{2}\max_{l=1}^Lc_l^2 \right]\\ &=\underset{c_1,\ldots,c_L\ge 0}{\text{min}}\left[ -\sum_{l=1}^L c_l \|\boldsymbol{G}_l\|_* + \frac{\lambda}{2}\max_{l=1}^Lc_l^2 \right]\quad\triangleright\|\cdot\|_*\text{表示核范数}\\ &=\underset{\eta\ge 0}{\text{min}}\left[ -\sum_{l=1}^L \eta\|\boldsymbol{G}_l\|_* + \frac{\lambda}{2}\max_{l=1}^L \eta^2 \right]\tag{1}\\ \end{align} $$

上面的每个等号成立的解释如下：

第一个等号成立，是直接带入 $c_l\triangleq\|\Delta\boldsymbol{W}_l\|_{\ell_2\to\ell_2}$ 的结果，注意到 $\min$ 的条件被拆成了大小和方向两部分。
第二个等号成立，是因为Frobenius内积具有线性性质， $$ \begin{align} \min_{\|\boldsymbol{T}_l\|=1}{\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F&=-\max_{\|\boldsymbol{T}_l\|=1}{-\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F\\ &=-\max_{\|\boldsymbol{T}_l\|=1}{\langle \boldsymbol{G}_l, -\boldsymbol{T}_l \rangle}_F\\ &=-\max_{\|\boldsymbol{T}'_l\|=1}{\langle \boldsymbol{G}_l, \boldsymbol{T}'_l \rangle}_F\\ \end{align} $$
第三个等号成立，是因为Frobenius内积 ${\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F$ 最大值是矩阵 $\boldsymbol{G}_l$ 的核范数（nuclear norm，即奇异值之和）(求解过程放在文末)。
第四个等号成立，即所有系数 $c_l$ 都取相同值 $\eta$ ：假设我们当前有某组 $c_1,\ldots,c_L$ 取到目标最小值，且满足 $\max_i c_i=\eta$ ，此时如果存在 $c_j<\eta$ ，则只需要增加 $c_j$ 的值，即可在第二项 $\lambda/2\max_{l=1}^Lc_l^2$ 不变的情况下，进一步减小第一项 $-\sum_{l=1}^L c_l \|\boldsymbol{G}_l\|_*$ 的值（注意到核范数是非负的）。因此这里令目标函数取最小的必要条件是所有 $c_1,\ldots,c_L$ 均取最大值 $\eta$ 。

式 $(1)$ 看起来可能有一些复杂，但是本质上关于 $\eta$ 是一个简单的二次函数，用初等数学⁴即可求出它的最优值对应的 $\eta$ ： $$ \eta = \frac{1}{\lambda}\sum_l^L \|\boldsymbol{G}_l\|_* $$

现在我们已经求出了 $\Delta \boldsymbol{W}_l$ 的最优范数 $c^*_l=\eta$ 了，还需要知道它的方向 $\boldsymbol{T}^*_l$

$$ \boldsymbol{T}^*_l=\underset{\boldsymbol{T}_l:\|\boldsymbol{T}_l\|=1}{\text{arg min}}{\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F=-\underset{\boldsymbol{T}_l:\|\boldsymbol{T}_l\|=1}{\text{arg max}}{\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F $$

这里直接给出 $\boldsymbol{T}_l$ 的解，求解的过程放在文末。假设 $\boldsymbol{G}_l\in\mathbb{R}^{m\times n}$ 的最简SVD分解为 $\boldsymbol{U}_{m\times r}\boldsymbol{\Sigma}_{r\times r}\boldsymbol{V}_{r\times n}^\top, r=\text{rank}(\boldsymbol{G}_l)$ ，则

$$ \boldsymbol{T}^*_l=-\boldsymbol{U}\boldsymbol{V}^\top $$

也就是说，上述带谱范数约束下的最速下降的解是 $$ \begin{align} \Delta \boldsymbol{W}_l&=-\eta\ {\color[rgb]{0, 0.5, 0.8}{\boldsymbol{U}\boldsymbol{V}^\top}}\\ \end{align} $$

到这一步，如果读者熟悉Muon优化器的话，会发现 $-{\color[rgb]{0, 0.5, 0.8}{\boldsymbol{U}\boldsymbol{V}^\top}}$ 正是Muon优化器的更新方向（如果忽略动量估计的话）。另外，对于Shampoo优化器，如果忽略掉左右conditioner的累加的话，也可以得出同样的更新方向。 $$ \begin{align} &-(\boldsymbol{G}\boldsymbol{G}^\top)^{-\frac{1}{4}}\boldsymbol{G}(\boldsymbol{G}^\top\boldsymbol{G})^{-\frac{1}{4}}\\ =&- (\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^\top\boldsymbol{V}\boldsymbol{\Sigma}\boldsymbol{U}^\top)^{-\frac{1}{4}}\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^\top(\boldsymbol{V}\boldsymbol{\Sigma}\boldsymbol{U}^\top\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^\top)^{-\frac{1}{4}}\\ =&-(\boldsymbol{U}\boldsymbol{\Sigma}^2\boldsymbol{U}^\top)^{-\frac{1}{4}}\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^\top(\boldsymbol{V}\boldsymbol{\Sigma}^2\boldsymbol{V}^\top)^{-\frac{1}{4}}\\ =&-\boldsymbol{U}\boldsymbol{\Sigma}^{-\frac{1}{2}}\boldsymbol{\Sigma}\boldsymbol{\Sigma}^{-\frac{1}{2}}\boldsymbol{V}^\top=-\boldsymbol{U}\boldsymbol{V}^\top\\ \end{align} $$

总结上面的这一小节内容：Muon、Shampoo优化器本质上是约束每个线性模块的算子诱导范数（在上面的例子中，是谱范数 $\|\cdot\|_{\ell_2\to\ell_2}$ ）衡量的参数更新量，从而推导出使用 $\boldsymbol{U}\boldsymbol{V}^\top$ 替代梯度矩阵 $\boldsymbol{G}$ 作为更新方向。这种约束的优势是，考虑了每次参数更新对于线性模块的输出的影响，更新的方向相比于标准梯度下降（使用欧氏距离衡量参数更新量）会更稳定。

在矩阵约束中考虑语义信息

故事到这里还没有结束。在神经网络中，同样是线性模块，根据输入输出语义的不同，表达的含义是不一样的。因此我们可以使用更泛化的诱导范数来约束每个模块参数的更新。

考虑一个矩阵 $\boldsymbol{W}\in\mathbb{R}^{m\times n}$ ，它链接着赋范空间 $(\mathbb{R}^n, \|\cdot\|_\alpha)$ 和 $(\mathbb{R}^m, \|\cdot\|_\beta)$ ，矩阵 $\boldsymbol{W}$ 的算子诱导范数定义如下 $$ \|W\|_{\alpha\to\beta}=\underset{\boldsymbol{x}\in\mathbb{R}^n}{\max} \frac{\|W\boldsymbol{x}\|_\beta}{\|\boldsymbol{x}\|_\alpha} $$

在前面的例子中，我们假设了输入与输出都是使用 $\ell_2$ 范数来衡量的，得到的最优更新方向是 $\boldsymbol{U}\boldsymbol{V}^\top$ 。现在，我们考虑另一个例子——Embedding矩阵，它的输入是one-hot向量(仅token对应的位置是1，其他位置都是0的向量)，此时我们可以使用 $\ell_1\to\text{RMS}$ 的诱导范数（输入空间是用 $\ell_1$ 范数，输出空间用RMS范数） $$ \|W\|_{\ell_1\to\text{RMS}}=\underset{\boldsymbol{x}\in\mathbb{R}^n}{\max} \frac{\|W\boldsymbol{x}\|_\text{RMS}}{\|\boldsymbol{x}\|_{\ell_1}} $$ 考虑到 $\boldsymbol{x}$ 总是one-hot向量，则上述诱导范数退化为Embedding RMS Norm的最大值： $$ \|W\|_{\ell_1\to\text{RMS}}=\underset{1\le j\le n}{\max} \|\boldsymbol{W}_{:,j}\|_\text{RMS} $$ 参照上一小节的推导，假设这个Embedding矩阵的梯度是 $\boldsymbol{G}$ ，则对应的最速更新方向由如下给出： $$ -\underset{\boldsymbol{T}}{\text{arg max}}{\langle \boldsymbol{G}, \boldsymbol{T} \rangle}_F \quad\text{s.t. }\underset{1\le j\le n}{\max} \|\boldsymbol{T}_{:,j}\|_\text{RMS}=1 $$

上述问题的解由如下构造（即对梯度的每列做RMS标准化）： $$ \boldsymbol{T}_{:,j}=\boldsymbol{G}_{:,j}/\|\boldsymbol{G}_{:,j}\|_\text{RMS} $$

在Bernstein & Newhouse 2024.中，作者还针对Conv2D算子给出了更新方向（在Bernstein的框架中，这个更新方向就是给定范数下的对偶映射（duality map））。最近Pethick 2025.把这个框架进一步拓展，给出了 $\ell_1\to\infty$ 诱导范数下的最优更新方向 $\text{sign}(\boldsymbol{G})$ ，他们将这个更新规则用在了分类模型的输出分类器上（最终他们的优化器实现类似一个Muon与SignSGD的混合）。总而言之，根据参数矩阵的输入、输出的语义不同，应当选取不同的诱导范数约束，对应的更新规则也不同。

结语

本篇文章分享了笔者阅读Bernstein一系列文章之后的一些重点摘录。对于深度学习中常见的矩阵参数，可以考虑根据输入与输出的不同语义，对梯度矩阵赋予不同的诱导范数约束，从而得到不同的更新规则。

这一系列的工作还有很多值得挖掘的地方，例如最优的学习率如何分配（对应到上述优化问题中，我们可以为每个诱导范数的约束项增加一个加权系数）？在Large 2024.中作者提到适当地设计模块之间的约束加权，则不同模型宽度（隐藏层维度）下的最优学习率是一致的（learning rate transfer）。如果这点在更多模型上成立的话，那么大规模神经网络的训练就可以有更加第一性的调参路径了。对于这部分相关工作，笔者还没有做深入的阅读，今后有机会可以再写文章分享。

附录

求解 $\underset{\boldsymbol{T}_l:\|\boldsymbol{T}_l\|=1}{\text{max}}{\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F$ 和 $\underset{\boldsymbol{T}_l:\|\boldsymbol{T}_l\|=1}{\text{arg max}}{\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F$ （下面的矩阵范数均为谱范数 $\|\cdot\|_{\ell_2\to\ell_2}$ ）：

假设 $\boldsymbol{G}_l$ 的SVD分解为 $\boldsymbol{U\Sigma V}^\top=\sum_i \sigma_i\boldsymbol{u}_i\boldsymbol{v}_i^\top$ ，则 $$ \begin{align} {\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F &= \text{Tr}\left(\boldsymbol{G}_l^\top \boldsymbol{T}_l\right)\\ &=\text{Tr}\left(\left(\sum_i \sigma_i\boldsymbol{u}_i\boldsymbol{v}_i^\top\right)^\top \boldsymbol{T}_l\right)\\ &=\text{Tr}\left(\sum_i \sigma_i\boldsymbol{v}_i\boldsymbol{u}_i^\top \boldsymbol{T}_l\right)\\ &=\text{Tr}\left(\sum_i \sigma_i\boldsymbol{u}_i^\top \boldsymbol{T}_l\boldsymbol{v}_i\right)\\ &=\sum_i \sigma_i\boldsymbol{u}_i^\top \boldsymbol{T}_l\boldsymbol{v}_i\quad\triangleright\text{括号内是标量}\\ &\le\sum_i \sigma_i\\ \end{align} $$

最后一个不等式成立是由于Cauchy-Schwarz不等式（ $\{\boldsymbol{u}_i\},\{\boldsymbol{u}_i\}$ 是正交规范向量组，范数是1）： $$ |\boldsymbol{u}_i^\top \boldsymbol{T}_l\boldsymbol{v}_i|\le \|\boldsymbol{u}_i\| \|\boldsymbol{T}_l\boldsymbol{v}_i\| \le \|\boldsymbol{u}_i\| \|\boldsymbol{T}_l\|\|\boldsymbol{v}_i\|=1 $$

注意到，如果带入 $\boldsymbol{T}_l=\boldsymbol{U}\boldsymbol{V}^\top=\sum_j \boldsymbol{u}_j\boldsymbol{v}_j^\top$ （可以验证此时 $\|\boldsymbol{T}_l\|=1$ ）

$$ \begin{align} \sum_i \sigma_i\boldsymbol{u}_i^\top \boldsymbol{T}_l\boldsymbol{v}_i &=\sum_i \sigma_i\boldsymbol{u}_i^\top (\sum_j \boldsymbol{u}_j\boldsymbol{v}_j^\top)\boldsymbol{v}_i\\ &=\sum_i \sigma_i \sum_j \boldsymbol{u}_i^\top\boldsymbol{u}_j\boldsymbol{v}_j^\top \boldsymbol{v}_i\\ &=\sum_i \sigma_i \end{align} $$ 因此 $\sum_i\sigma_i$ 这个上界是可取到的，则 $$ \begin{split} \underset{\boldsymbol{T}_l:\|\boldsymbol{T}_l\|=1}{\text{max}}{\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F=\sum_i\sigma_i\|\boldsymbol{G}_l\|_*\\ \boldsymbol{U}\boldsymbol{V}^\top\in \underset{\boldsymbol{T}_l:\|\boldsymbol{T}_l\|=1}{\text{arg max}}{\langle \boldsymbol{G}_l, \boldsymbol{T}_l \rangle}_F\\ \end{split} $$

这里如果 $\boldsymbol{G}_l$ 是满秩的，则argmax问题的解是唯一的（可以用反证法来证明），如果非满秩，则 $\boldsymbol{U}\boldsymbol{V}^\top$ 不是唯一解（不影响主要结论）。

参考阅读

这里直接按照Lagrange乘子展开。 ↩︎
考虑极端情况，如果直接求和作为约束，那么得到的解可能会让某个参数矩阵变化极大，其他矩阵保持不变，那么目标函数的变化可能是非常剧烈的。 ↩︎
参考了Bernstein & Newhouse 2024. Old Optimizer, New Norm: An Anthology ↩︎
$ax^2+bx(a>0)$ 的最小值在 $b/2a$ 取到。 ↩︎

梯度下降（Recap）#

诱导范数作为步长约束#

在矩阵约束中考虑语义信息#

结语#

附录#

参考阅读#

梯度下降（Recap）

诱导范数作为步长约束

在矩阵约束中考虑语义信息

结语

附录

参考阅读