为什么我们偏爱各向同性？基于最速下降的理解

从机器学习时代的数据白化预处理，到深度学习时代的BatchNorm、InstanceNorm、LayerNorm、RMSNorm等花样百出的Normalization方法，本质上都体现了我们对“各向同性（Isotropy）”的偏爱。为什么我们会倾向于各向同性的特征呢？它有什么实际上的好处呢？这个问题能找到很多答案，比如对齐尺度、减少冗余、去相关性等等，但多是流于表面的感觉。

近日，笔者在读论文《The Affine Divergence: Aligning Activation Updates Beyond Normalisation》时，悟到了该问题在优化视角下的一个新理解，个人认为它相对来说还是比较贴近本质的，所以写出来跟大家分享和讨论一下。

最速下降 #

我们从最简单的线性层出发
\begin{equation}\boldsymbol{Y} = \boldsymbol{X}\boldsymbol{W}\end{equation}
其中$\boldsymbol{X}\in\mathbb{R}^{b\times d_{in}}$是当前层输入，$\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}$是权重，$\boldsymbol{Y}\in\mathbb{R}^{b\times d_{out}}$是输出。损失函数记为$\mathcal{L}(\boldsymbol{Y}) = \mathcal{L}(\boldsymbol{X}\boldsymbol{W})$，那么有
\begin{equation}\frac{\partial \mathcal{L}}{\partial\boldsymbol{W}} = \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\end{equation}
我们以梯度下降为例，此时更新规则为
\begin{equation}\boldsymbol{W}\quad\leftarrow\quad \boldsymbol{W} - \eta\frac{\partial \mathcal{L}}{\partial\boldsymbol{W}} = \boldsymbol{W} - \eta \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\end{equation}
梯度下降的基本原理是广为人知的“梯度反方向是损失下降最快的方向”，但这结论是有前提的。最关键的前提是它选取的度量是欧氏范数，如果换一个范数，那么最速方向也就变了，对此我们在《Muon续集：为什么我们选择尝试Muon？》、《流形上的最速下降：1. SGD + 超球面》等文章已有讨论。

切换视角 #

本文主要关注另一个不那么容易留意到的前提：视角，或者说立场。

假设我们认同“梯度反方向是损失下降最快的方向”，那么问题是：“谁”的梯度呢？可能有读者说那自然是参数的梯度了，这确实是标准答案，但未必是最佳答案。参数本质上是模型的副产品，我们实际上不关心参数如何，我们关心的是模型输入输出是否构成了期望的映射。

所以，输入输出的特征变化，是我们更关心的对象。如果我们从特征的视角出发，那么结论就会不一样。具体来说，参数从$\boldsymbol{W}$变成$\boldsymbol{W} - \eta \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}$，输出特征$\boldsymbol{Y}$的变化量是
\begin{equation}\Delta \boldsymbol{Y} = \boldsymbol{X}\left(\boldsymbol{W} - \eta \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right) - \boldsymbol{X}\boldsymbol{W} = - \eta \boldsymbol{X} \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\end{equation}
根据“梯度反方向是损失下降最快的方向”，那么站在$\boldsymbol{Y}$的立场看，我们要想$\boldsymbol{Y}$的变化量让损失下降得最快，那么应该满足$\Delta \boldsymbol{Y}\propto -\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}$才对，现在却多出了一个Gram矩阵$\boldsymbol{X} \boldsymbol{X}^{\top}$，这意味着它并没有沿着下降最快的方向走。

各向同性 #

一个朴素的想法是：如果$\boldsymbol{X} \boldsymbol{X}^{\top}$直觉等于单位阵（的若干倍）就好了。我们知道$\boldsymbol{X}\in\mathbb{R}^{b\times d_{in}}$，若$b\leq d_{in}$，那么$\boldsymbol{X} \boldsymbol{X}^{\top} = \boldsymbol{I}$意味着这$b$个向量构成了一组标准正交基；但实际通常都有$b > d_{in}$，所以$\boldsymbol{X} \boldsymbol{X}^{\top} = \boldsymbol{I}$在数学上无法严格成立。

在这种情况下，我们只能期望这$b$个向量能尽可能均匀地分布在单位超球面上，以使得$\boldsymbol{X} \boldsymbol{X}^{\top} = \boldsymbol{I}$近似成立，这便是各向同性的缘由。也就是说：

如果输入特征能满足各向同性，那么可以让参数上的最速下降同时约等于特征上的最速下降，“双管齐下”，充分提高模型的学习效率。

我们还可以证明，如果随机变量服从$d_{in}$维的标准正态分布，那么它在满足各向同性的同时，其模长也会高度集中在$\sqrt{d_{in}}$附近，即近似满足在半径为$\sqrt{d_{in}}$的超球面上。反过来，如果能将输入$\boldsymbol{X}$标准化为零均值、单位协方差的矩阵，我们认为$\boldsymbol{X} \boldsymbol{X}^{\top} = d_{in}\boldsymbol{I}$也近似成立，而这操作便是白化。

标准化层 #

除了事前对数据进行白化外，我们还可以在模型中间引入一些标准化运算，使得中间特征也能近似满足期望性质。更具体地说，我们设法寻找一个$d_{in}\times d_{in}$的变换矩阵$\boldsymbol{A}$，使得变换后的特征$\boldsymbol{X}\boldsymbol{A}$尽量满足$(\boldsymbol{X}\boldsymbol{A})(\boldsymbol{X}\boldsymbol{A})^{\top}=\boldsymbol{I}$，即
\begin{equation}\min_{\boldsymbol{A}} \Vert \boldsymbol{X}\boldsymbol{A}\boldsymbol{A}^{\top}\boldsymbol{X}^{\top} - \boldsymbol{I}\Vert_F\end{equation}
这个问题的解可以用伪逆表示：
\begin{equation}\boldsymbol{A}\boldsymbol{A}^{\top} = \boldsymbol{X}^{\dagger}(\boldsymbol{X}^{\top})^{\dagger} = (\boldsymbol{X}^{\top}\boldsymbol{X})^{-1}\end{equation}
这里假设了$\boldsymbol{X}^{\top}\boldsymbol{X}$可逆。根据上式，我们得到一个可行解为$\boldsymbol{A} = (\boldsymbol{X}^{\top}\boldsymbol{X})^{-1/2}$，那么对应的变换是
\begin{equation}\boldsymbol{X}(\boldsymbol{X}^{\top}\boldsymbol{X})^{-1/2}\end{equation}
这正是没有中心化的白化。考虑算$\boldsymbol{X}^{\top}\boldsymbol{X}^{-1/2}$比较昂贵，如果改用对角线近似，那么就是将每一维分别标准化，视颗粒度的不同，它分别对应着BatchNorm、InstanceNorm等运算；如果我们更关心“超球面”，则可以考虑分别标准化每个样本的模长，这对应于LayerNorm、RMSNorm等运算。

SGD之外 #

有趣的是，“如果输入特征能满足各向同性，那么参数上的最速下降同时约等于特征上的最速下降”这一结论不仅适用于SGD，还适用于Muon。考虑没有动量的Muon，更新规则是
\begin{equation}\newcommand{msign}{\mathop{\text{msign}}} \boldsymbol{W}\quad\leftarrow\quad \boldsymbol{W} - \eta\msign\left(\frac{\partial \mathcal{L}}{\partial\boldsymbol{W}}\right) = \boldsymbol{W} - \eta \msign\left(\boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right)\end{equation}
根据$\msign(\boldsymbol{M}) = \boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^{-1/2}$，我们有
\begin{equation}\Delta \boldsymbol{Y} = - \eta \boldsymbol{X} \msign\left(\boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right) = - \eta \boldsymbol{X} \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}} \left(\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\vphantom{\bigg|}^{\top}\boldsymbol{X}\boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right)^{-1/2}\end{equation}
由此可见，如果$\boldsymbol{X}\boldsymbol{X}^{\top}\approx \boldsymbol{I}$，那么
\begin{equation}\Delta \boldsymbol{Y} \approx - \eta \frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}} \left(\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\vphantom{\bigg|}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right)^{-1/2} = -\eta\msign\left(\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right)\end{equation}
我们知道，Muon是谱范数下的最速下降，上式则意味着，如果输入特征能满足各向同性，那么参数上的谱范数最速下降同时约等于特征上的谱范数最速下降，非常精妙！然而，如果是其他优化器比如SignSGD，则无法复现类似性质，这一区别也许是Muon优异性的背后原因之一。

文章小结 #

这篇文章我们讨论了一个问题：参数层面上的最速下降，何时才刚好是特征层面上的最速下降？答案正是标题所述的“各向同性”，由此我们得出为什么偏爱各向同性的一个解释——它能同步两个层面上的最速下降，提高训练效率。

转载到请包括本文地址：https://spaces.ac.cn/archives/11549

更详细的转载事宜请参考：《科学空间FAQ》