Prologue
如何衡量神经网络参数空间中的距离（i.e., 选取合适的范数）？这个问题我们之前已经在这篇文章中有所涉及，本文是更完整的拓展。
首先应该明确为什么这个问题很重要？因为在梯度下降这样的优化算法框架下，总是依赖目标参数的某种距离的衡量。例如，当目标参数可以对应到欧几里得空间的坐标时，就可以得到局部目标函数下降最快的方向是梯度负方向。
尽管在深度学习优化中，直接使用欧几里得范数（衡量参数空间的距离）非常吸引人，并且实际上也有效。但是，这种做法会丢失结构信息，因为它实质上是将参数看做一个扁平的向量。那么得到的优化方向就可能是相对低效的。
对于常见的神经网络而言，对参数矩阵使用谱范数似乎是个更好的选择[1][2]。使用谱范数的相关约束，可以引出在现在的大规模模型训练中非常重要的两个工作：
$\mu P$(Maximal Update Parameterization)：使用特定的初始化与学习率的参数化，可以做到在小模型上调节部分超参数，迁移到同构的大模型。- Muon优化器的更新规则，即对梯度矩阵做SVD分解，消除奇异值矩阵后将
$-\boldsymbol{U}\boldsymbol{V}^\top$作为对应参数的更新方向。
从而涵盖了给定神经网络结构下的参数初始化、参数更新的步长和方向三个方面（模型成功训练的三个基本要素）。
feature learning的谱条件
为什么需要使用谱范数来衡量参数的大小和参数更新步长呢？在Greg Yang的系列工作中，考虑的主要是feature learning的问题。Yang在他的TP4中发现，在标准参数化或者NTK参数化的设定下，无穷宽的神经网络没有办法学习特征（一次更新后特征与初始化状态无异），这就与传统的linear transfer learning的智慧形成了悖论，因为如果没有feature learning，那么预训练就对后续的transfer没有任何增益。
于是为了达成非平凡的feature learning，就需要对神经网络的特征$\boldsymbol{h}_\ell(\boldsymbol{x})\in\mathbb{R}^{n_\ell}$的范数与一步更新的特征范数变化量$\boldsymbol{h}_\ell(\boldsymbol{x})$做如下的约束[1]：
$$ \|\boldsymbol{h}_\ell\|_2=\Theta(\sqrt{n_\ell})\text{ and } \|\Delta\boldsymbol{h}_\ell\|_2=\Theta(\sqrt{n_\ell}), \text{ for }\ell=1,\ldots,L-1\tag{1} $$
这个条件意味着，对于中间的任何特征向量而言，每个元素平均下来的范数是常数阶的，一次更新后的变化量也是常数阶的（不会随着宽度$n_\ell$趋近无穷而爆炸或者弥散）。
对于常见的由一系列矩阵参数构成的神经网络而言，实现这样的约束，需要对参数矩阵做如下的谱条件约束：
$$ \|\boldsymbol{W}_\ell\|_2=\Theta\left(\sqrt{\frac{n_\ell}{n_{\ell-1}}}\right)\text{ and } \|\Delta\boldsymbol{W}_\ell\|_2=\Theta\left(\sqrt{\frac{n_\ell}{n_{\ell-1}}}\right), \text{ for }\ell=1,\ldots,L-1\tag{2} $$
这里的$\|\cdot\|_2$是矩阵的谱范数，即从向量$\ell_2$空间映射到向量$\ell_2$空间的算子诱导范数
$$ \|\boldsymbol{A}\|_2=\underset{\boldsymbol{x}\in\mathbb{R}^n}{\max} \frac{\|\boldsymbol{A}\boldsymbol{x}\|_2}{\|\boldsymbol{x}\|_2}\text{ for } \boldsymbol{A}\in\mathbb{R}^{m\times n} $$
这个谱条件能够保证公式$(1)$成立，证明思路主要利用到谱范数的性质$\|\boldsymbol{Av}\|_2\le\|\boldsymbol{A}\|_2\|\boldsymbol{v}\|_2$，证明出上述谱条件诱导的$\|\boldsymbol{h}_\ell\|_2,\|\Delta\boldsymbol{h}_\ell\|_2$的上界均为$\Theta(\sqrt{n_\ell})$，接着证明这个上界依概率总是取得。具体可以看这篇论文的第三节（相比Greg的TP系列，还算是很容易理解的）。
公式$(2)$中的第二个条件是很直观的，对于SGD优化器而言，直接按照每个参数矩阵的维度设定对应的学习率即可
$$ \eta_\ell = \Theta\left(\sqrt{\frac{n_\ell}{n_{\ell-1}}}\right) $$
对于第一个条件$\boldsymbol{W}_\ell = \Theta\left(\sqrt{\frac{n_\ell}{n_{\ell-1}}}\right)$，需要修改初始化的方式——我们可以从一个标准正态分布中i.i.d.采样一个权重矩阵$\boldsymbol{W}'\in\mathbb{R}^{n_\ell\times n_{\ell-1}}$，接着做缩放$\boldsymbol{W} = \sigma\boldsymbol{W}'$（这里略去下标$\ell$）得到最终的初始化参数矩阵。
这个$\sigma$怎么确定呢？由于矩阵的stable rank按定义可以将矩阵的谱范数和Frobenius范数联系起来：
$$ \text{stable-rank}(\boldsymbol{W}) = \frac{\|\boldsymbol{W}\|_F^2}{\|\boldsymbol{W}\|_2^2} $$
这里分子是谱范数平方，分母是Frobenius范数平方。对于标准正态分布中i.i.d.采样的矩阵$\boldsymbol{W}$而言，其stable rank在渐进意义下满足
$$ \text{stable-rank}(\boldsymbol{W}) = \Theta(\min\{n_\ell, n_{\ell-1}\}) $$
那么根据RMS范数与F范数、F范数与谱范数的换算关系，只需要让RMS范数满足
$$ \begin{align} \|\boldsymbol{W}\|_{\text{RMS}} &= \frac{1}{\sqrt{n_\ell\cdot n_{\ell-1}}}\|\boldsymbol{W}\|_F\\ &= \frac{1}{\sqrt{n_\ell\cdot n_{\ell-1}}}\cdot \sqrt{\text{stable-rank}(\boldsymbol{W})\cdot \|\boldsymbol{W}\|_2^2}\\ &= \frac{1}{\sqrt{n_\ell\cdot n_{\ell-1}}}\cdot \Theta(\sqrt{\min\{n_\ell, n_{\ell-1}\}})\cdot \Theta\left(\sqrt{\frac{n_\ell}{n_{\ell-1}}}\right)\\ &= \Theta\left( \frac{1}{\sqrt{n_{\ell-1}}}\min\left\{1,\sqrt{\frac{n_\ell}{n_{\ell-1}}}\right\} \right) \end{align} $$
也就是说，我们设定$\sigma=\Theta\left( \frac{1}{\sqrt{n_{\ell-1}}}\min\left\{1,\sqrt{\frac{n_\ell}{n_{\ell-1}}}\right\}\right)$并做缩放$\boldsymbol{W} = \sigma\boldsymbol{W}'$即可满足上述的谱条件。
这个初始化系数与学习率比例就恰好对应了$\mu P$的核心公式！
在这个初始化和学习率设定下，由于可以满足$ \|\Delta\boldsymbol{h}_\ell\|_2=\Theta(\sqrt{n_\ell})$，即$\|\Delta\boldsymbol{h}_\ell\|_{\text{RMS}}=\Theta(1)$，优化过程的特征学习不会随着宽度$n_\ell$增大而爆炸或者弥散，因此可以做到zero-shot hyper-parameter transfer，从而能在较小的模型中调节部分超参数（学习率、初始化方差和放缩系数，在TP4中称为abc-parameterization），满足这组参数化对大的同构模型也是优势超参。
谱条件引导的最速优化方向
除了谱范数约束推导出初始化方差与学习率（的比例）之外，使用谱范数约束推导最速下降方向：
$$ \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{{\|\Delta\boldsymbol{W}_l\|^2_{2}}} \right]\\ $$
可以得出最速下降的解是
$$ \begin{align} \Delta \boldsymbol{W}_l&=-\eta{\boldsymbol{U}\boldsymbol{V}^\top}\\ \end{align} $$
也就是Muon和Shampoo优化器（如果忽略动量或累加项）的优化方向。具体的证明笔者在这篇文章已经介绍过了，读者也可以直接阅读这篇论文。
需要提及的是，运用这种方法的前提是，参数矩阵表达的是某种线性映射，如果是Embeddings矩阵或分类器矩阵，需要参考具体含义使用其他的诱导范数。
结语
本文简要介绍了使用谱范数衡量参数空间的距离，对初始权重矩阵、权重矩阵的更新量做约束，可以推导出零次超参数迁移方法$\mu P$中的初始化、学习率设置，以及SOTA优化器Muon/Shampoo的优化方向。这是我们科学地优化神经网络的一个起点。