小模型上调好的学习率,换到更宽的网络上往往不再合适。宽度改变了每层参与求和的项数,也改变了梯度对特征的作用。即使初始化时的激活方差相近,训练中的特征更新仍可能随着宽度增大而爆炸或消失。

$\mu\mathrm{P}$(Maximal Update Parametrization,最大更新参数化)通过配套的初始化和更新规则,在稳定性约束下保留不随宽度消失的特征更新。理解这套规则,可以从初始化、神经网络高斯过程(NNGP)、神经正切核(NTK)与均值场的关系入手1。

一、$\mu\mathrm{P}$ 的前身

这些理论在历史上有不少并行发展,讨论的对象也不同:有的分析随机初始化,有的描述训练过程。以下固定输入维度与网络深度,让隐藏层宽度 $n$ 增长。$O(1)$ 表示量级不随宽度增长,而非数值恰好等于一2。

标准参数化与参数坐标

常见的标准参数化(Standard Parametrization,SP)直接存储各层权重,按 fan-in 一类规则初始化,再用指定的学习率训练。比较不同宽度时,学习率是否随宽度变化也是参数化的一部分。

完整的参数化包括实际存储的参数、前向计算中的倍率、初始化分布和优化器的学习率。权重的初始化标准差相同,训练过程仍可能不同。

对于固定的重参数化 $w=s\theta$,普通 SGD 在实际权重坐标下的有效学习率为

$$ \eta_w=\eta_\theta s^2. $$

因此,要让两种参数坐标对应相同的训练轨迹,前向倍率、初始化与学习率都需要匹配。

重参数化下的 SGD 学习率

设真正参与计算的权重为 $w$,程序存储的参数为 $\theta$,二者满足

$$ w=s\theta, $$

其中 $s$ 可以依赖模型宽度,但在一次训练中固定。这里的 $\mathcal L$ 仅指数据损失,不含另外施加的权重衰减。使用普通 SGD,根据链式法则,

$$ \nabla_\theta\mathcal L=s\nabla_w\mathcal L, \qquad \Delta w=s\Delta\theta =-\eta_\theta s^2\nabla_w\mathcal L. $$

因此,以 $\theta$ 为坐标训练时,在 $w$ 坐标下对应的有效学习率为

$$ \eta_w=\eta_\theta s^2. $$

例如,把 $s$ 改成 $1/\sqrt n$,同时保持 $\eta_\theta$ 不变,会让 $w$ 坐标下的有效学习率多出 $1/n$。即使通过调整初始化使两个实现拥有相同的初始函数分布,它们后续的训练轨迹仍可能不同。

上述 $s^2$ 关系限定于普通 SGD。Adam 会利用梯度历史归一化更新,在相同的参数重缩放下,需要另外计算有效学习率。

Xavier 与 He 初始化

考虑一层线性变换

$$ z_j=\sum_{i=1}^{n_{\rm in}}W_{ji}h_i. $$

在零均值权重、适当独立性等初始化近似下,它的二阶矩满足

$$ \mathbb E[z_j^2] \approx n_{\rm in}\operatorname{Var}(W_{ji})\mathbb E[h_i^2]. $$

要保持 $z_j$ 的二阶矩,输入项变多时,权重方差需要相应缩小。反向传播也有类似计算,但求和项数与输出维度有关。

Xavier 初始化综合考虑前向与反向传播,在接近线性的分析下给出常见的方差选择[1] :

$$ \operatorname{Var}(W_{ji}) =\frac{2}{n_{\rm in}+n_{\rm out}}. $$

He 初始化进一步考虑 ReLU:对称输入经过 ReLU 后,二阶矩减半,因此保持前向二阶矩的常见选择为[2]

$$ \operatorname{Var}(W_{ji})=\frac{2}{n_{\rm in}}. $$

这些规则约束随机初始化时的激活与梯度尺度。输入、输出维度不同时,前向和反向的稳定性通常还需要取舍。训练中的特征变化还取决于学习率和梯度作用于输入的方式;仅有初始化的方差条件,不能保证特征更新在不同宽度下保持相近尺度。

初始化的高斯过程极限

把隐藏层宽度记为 $n$,考虑标量输出

$$ f_n(x)=\frac{1}{\sqrt n} \sum_{i=1}^{n}a_i\phi(w_i^\top x). $$

假设各神经元独立同分布初始化,$a_i$ 与 $w_i$ 独立,且 $a_i$ 零均值、相关二阶矩有限。对固定输入,输出是许多独立随机贡献的归一化求和。对一组固定输入同时应用中心极限定理,它们的输出可以收敛到联合高斯分布。

在这个设定下,输出方差为

$$ \operatorname{Var}[f_n(x)] =\operatorname{Var}[a\phi(w^\top x)]. $$

若右侧是有限非零常数,初始输出的波动便保持 $O(1)$。这里使用了独立、零均值贡献的求和及其归一化,单个神经元的激活本身不需要服从高斯分布。

由此得到的高斯过程视角,把随机网络看成一个函数先验。其协方差由初始化和激活函数决定;Lee 等人的工作将这一对应关系系统地扩展到深层网络[3] 3。

这个函数先验给出了初始化时 $f_0$ 的随机分布,可用于相应的贝叶斯推断。梯度下降中的隐藏特征是否变化,还取决于函数对参数的敏感度和更新尺度。

NTK 极限中的训练

神经正切核将参数梯度组合成一个输入之间的相似度:

$$ K_\theta(x,x') =\left\langle\nabla_\theta f(x), \nabla_\theta f(x')\right\rangle. $$

仍使用上一节的 $1/\sqrt n$ 两层网络,并令存储参数 $a_i,w_i$ 的初始化量级为 $O(1)$。若损失导数、输入与相关激活导数保持有界,对单个隐藏权重有

$$ \nabla_{w_i}f_n(x) =\frac{a_i}{\sqrt n}\phi'(w_i^\top x)x =O(n^{-1/2}). $$

因此,在学习率 $\eta=O(1)$ 的普通 SGD 下,固定索引的典型单元有单步 $\Delta w_i=O(n^{-1/2})$。在不随宽度增长的有限训练时间内,并满足相应正则性与稳定性条件时,这类单元的特征变化可以随宽度增大而消失4。

大量微小的参数更新共同作用于输出,仍然可以使网络函数发生 $O(1)$ 变化。NTK 理论刻画的一类极限中,初始化核收敛到确定性极限,并在训练期间保持不变,函数更新由固定核控制[4] 。

例如,在梯度流下,把训练集上的输出记为向量 $f_t$,便得到

$$ \frac{\mathrm d f_t}{\mathrm dt} =-K_0\nabla_f\mathcal L(f_t). $$

固定核使训练分析大幅简化,同时限制了模型在这一极限中的表示适应能力。有限宽网络会偏离这个极限,而该近似是否适用,还取决于参数化、学习率和训练时间。

均值场训练动力学

另一种两层网络尺度是

$$ f_n(x)=\frac{1}{n} \sum_{i=1}^{n}a_i\phi(w_i^\top x). $$

现在单个参数对输出的梯度带有 $1/n$。若在上述存储坐标中配合学习率

$$ \eta=n\eta_0, $$

那么宽度因子会在更新中抵消。例如,对单个样本的损失 $\ell(f_n(x),y)$,

$$ \Delta w_i =-\eta_0\,\partial_f\ell\, a_i\phi'(w_i^\top x)x. $$

当相关量保持有界、损失梯度非退化且 $\eta_0$ 适当时,单个神经元可以保持 $O(1)$ 的更新尺度,特征变化不会被宽度强制压到零。具体梯度仍受数据、激活和对称性影响,部分神经元的更新可能为零。这里的 $\eta\propto n$ 对应当前存储坐标;换一种参数表示后,需要重新换算有效学习率。

均值场方法把神经元群体写成经验分布

$$ \rho_n=\frac1n\sum_{i=1}^{n}\delta_{(a_i,w_i)}, \qquad f_n(x)=\int a\phi(w^\top x)\,\mathrm d\rho_n(a,w), $$

再研究这个分布随训练的演化。Mei 等人的工作说明,在适当缩放极限下,两层网络的 SGD 可以用分布动力学描述[5] 。

若初始化时各项独立且零均值,$1/n$ 平均的输出波动通常为 $O(n^{-1/2})$,所以初始输出可以趋近零5。初始输出方差消失,与训练中产生非退化的输出和特征变化是相容的。

这个均值场例子与前面的 NTK 例子,在初始函数分布和训练尺度上都有差异。前向倍率与学习率共同决定了各自的极限,单独替换其中一项通常会改变结论。

信号传播中的均值场

前面的均值场描述神经元分布随训练时间的演化。初始化理论中也有一类均值场分析,研究信号沿网络深度的传播。

Poole 等人的工作[6] 和 Schoenholz 等人的 Deep Information Propagation[7] 追踪随机初始化时的信号尺度与输入间相关性,分析它们经过每一层网络后的变化。

在基本的宽全连接网络中,把激活前数值的二阶矩记为 $q_\ell$,可以得到这样的递推:

$$ q_{\ell+1} =\sigma_w^2\, \mathbb E_{Z\sim\mathcal N(0,1)} [\phi(\sqrt{q_\ell}Z)^2] +\sigma_b^2. $$

这里假设各层独立高斯初始化,权重方差为 $\sigma_w^2/n_\ell$,偏置方差为 $\sigma_b^2$,并在固定深度下取大宽度。$q_\ell$ 的稳定性约束单个输入的信号尺度;不同输入之间的信息能否被保留,还需要同时追踪它们的协方差。

Tensor Programs 用更通用的计算语言分析大宽度极限。Tensor Programs I 将初始化时的高斯过程对应扩展到更广泛的网络计算[8] ;后续工作进一步处理训练中的参数与特征相关性。

二、$\mu\mathrm{P}$ 的尺度规则

两个隐藏层的 MLP

初始化信号稳定、输出能够学习、隐藏特征持续变化,是三个不同的条件。NTK 极限允许输出学习,隐藏特征却在领先阶冻结。

Yang 与 Hu 在 Tensor Programs IV 中研究了宽度相关参数化及其训练极限,提出最大更新参数化 $\mu\mathrm{P}$[9] 。这里的“最大”指在稳定性约束下,让各层特征保留不随宽度消失的更新尺度;它不要求把数值学习率或参数更新范数设得尽可能大。

考虑两个隐藏层、标量输出的 ReLU MLP:

$$ \begin{aligned} z_1&=W_1x, & h_1&=\operatorname{ReLU}(z_1),\\ z_2&=W_2h_1, & h_2&=\operatorname{ReLU}(z_2),\\ f(x)&=W_3h_2. \end{aligned} $$

其中 $W_1\in\mathbb R^{n\times d}$、$W_2\in\mathbb R^{n\times n}$、$W_3\in\mathbb R^{1\times n}$。固定输入维度 $d$、深度、批量和数据,只增加宽度 $n$。程序直接存储这三个有效权重,不额外添加前向倍率;暂不引入偏置、归一化或权重衰减。表中的学习率直接作用于公式里的 $W$。

梯度与输入的相关性

初始化方差的计算依赖独立随机项的抵消。训练时,同一个输入同时出现在梯度和前向计算中,这一独立性不再成立。

对单个样本的损失 $\ell(f(x),y)$,令 $\delta_2=\partial\ell/\partial z_2$。隐藏层梯度与 SGD 更新满足

$$ \nabla_{W_2}\ell=\delta_2h_1^\top, \qquad (\Delta W_2)h_1 =-\eta_2\delta_2\|h_1\|^2. $$

右侧的平方和 $\|h_1\|^2$ 通常具有 $n$ 的量级,而不是随机符号求和的 $\sqrt n$ 量级。梯度与输入的对齐,使许多很小的参数更新仍能对特征产生有限作用。

三类权重因此需要不同的缩放:输入层只对固定的 $d$ 个输入求和,后两层对 $n$ 个特征求和,读出层还决定反传到隐藏层的梯度尺度。

三层 SGD 学习率的尺度推导

取输入与隐藏层的初始化标准差分别为 $\sigma_1/\sqrt d$、$\sigma_2/\sqrt n$,读出层取 $\sigma_3/n$,其中 $\sigma_i$ 不随宽度变化。前两个隐藏表示的典型坐标保持常数量级。设损失对输出的导数 $\ell'$ 有界且非退化。

固定 $h_2$,读出层更新对输出的作用为

$$ \Delta f_{\rm readout} =-\eta_{\rm out}\ell'\|h_2\|^2 =O(\eta_{\rm out}n). $$

因此应取 $\eta_{\rm out}\propto 1/n$,让这部分输出变化保持常数量级。

在上述随机初始化下,第二层反传信号

$$ \delta_2=\ell'W_3^\top\odot\operatorname{ReLU}'(z_2) $$

的典型坐标为 $O(1/n)$。在更新恒等式中,$1/n$ 与 $\|h_1\|^2=O(n)$ 抵消,因此 $\eta_{\rm hidden}\propto 1$。

在这一初始化和固定深度设定下,继续反传得到的 $\delta_1$ 的典型坐标也为 $O(1/n)$。输入层的直接变化为

$$ (\Delta W_1)x=-\eta_{\rm in}\delta_1\|x\|^2. $$

$d$ 固定,$\|x\|^2$ 不随 $n$ 增长,因此需要 $\eta_{\rm in}\propto n$。

这些计算给出了初始化附近的尺度选择。训练还会改变各层输入,并形成参数与特征之间的相关性。例如,联合更新隐藏层时有

$$ \Delta z_2=(\Delta W_2)h_1+W_2\Delta h_1 +(\Delta W_2)\Delta h_1. $$

只分析第一项不足以证明整个训练轨迹的性质;Tensor Programs 系统地处理了这些相关性与训练极限[9][10] 。

有效权重坐标下的缩放规则

在上述有效权重坐标中,三类参数的宽度依赖关系如下[10] 。$\sigma_i$ 与 $\eta_i$ 是不随宽度改变的常数,可以逐层选择;缩放规则确定了宽度的幂次,这些常数仍需调优。

参数 初始化标准差 SGD 学习率 Adam 学习率
输入层 $W_1$ $\sigma_1/\sqrt d$ $\eta_1 n$ $\eta_1$
隐藏层 $W_2$ $\sigma_2/\sqrt n$ $\eta_2$ $\eta_2/n$
读出层 $W_3$ $\sigma_3/n$ $\eta_3/n$ $\eta_3/n$

例如,宽度从 64 增加到 256:初始化标准差分别乘 $1,1/2,1/4$;SGD 学习率分别乘 $4,1,1/4$;Adam 学习率分别乘 $1,1/4,1/4$。这里比较的是同一套参数坐标下的两个模型。

Adam 按梯度的二阶矩归一化更新。在 $\epsilon$ 可忽略、梯度非退化的尺度分析下,归一化后的典型坐标更新为常数量级,实际步幅由学习率决定6。SGD 中补偿小梯度的输入层 $n$ 倍学习率,因此不能原样用于 Adam。

读出层采用独立零均值初始化、标准差按 $1/n$ 缩放时,初始输出通常随宽度缩小,训练中的输出更新仍可保持常数量级。若使用全零读出初始化,第一步的隐藏层梯度会为零。

mup 的存储坐标与有效权重

缩放因子可以放在存储参数或前向计算中。原作者公开的 mup 实现使用 MuReadout、set_base_shapes 和 MuAdam 配合工作[35] 。

设相对基准模型的宽度倍率为 $r=n/n_0$。默认输出倍率下,MuReadout 存储参数 $V$,前向计算使用的有效权重是

$$ W_3=V/r. $$

初始化经过 set_base_shapes 调整后,$V$ 的标准差相对基准模型保持不变,因此 $W_3$ 的标准差随宽度按 $1/r$ 缩小。MuAdam 对 $V$ 使用不随宽度变化的学习率,但 $\Delta W_3=\Delta V/r$,有效更新已经缩小了 $1/r$。隐藏矩阵有两个随宽度增长的维度,MuAdam 则直接将其参数组学习率除以 $r$。

上表与 mup 采用了不同的存储坐标。使用 MuReadout 与 MuAdam 时,手工将读出层学习率再除一次 $r$ 会重复缩放。普通 PyTorch 初始化之后,调用 set_base_shapes 完成尺度调整,再创建优化器;若在标记后重新初始化,应使用 mup.init 或显式遵守同一套尺度。注意力头数、参数共享方式或网络结构改变后,哪些维度随宽度增长也需要重新确认。

单步分层更新实验

实验用 NumPy 实现上述 MLP 并手算梯度。固定一批训练数据和一批探测输入,宽度取 32、64、128、256、512、1024,每个宽度运行 16 个随机种子。计算一次 SGD 更新后,固定该层输入特征,仅替换这一层权重,测量预激活的变化:

$$ (\Delta W_1)x,\qquad (\Delta W_2)h_1,\qquad (\Delta W_3)h_2. $$

纵轴是这些变化在探测批量上的 RMS,即所有坐标平方均值的平方根,统计值按坐标数量归一化。SP 基线对三层使用同一个、不随宽度变化的学习率;$\mu\mathrm{P}$ 则按上表逐层缩放。

固定数据与 SGD 基准步长时,SP 的输入层直接更新随宽度缩小,隐藏层与读出层更新增大;μP 的三个面板大致水平。

图中实线为 $\mu\mathrm{P}$,虚线为标准参数化 SP;线表示 16 个种子的中位数,阴影表示四分位区间。宽度从 32 增至 1024 时,SP 输入层的直接更新约缩至原来的 0.17 倍,隐藏层增至 4.30 倍,读出层增至 24.59 倍;$\mu\mathrm{P}$ 三者分别约为 1.06、0.99、0.88 倍。这里比较各层随宽度变化的趋势,不同层的绝对高度可以不同。

降低 SP 的全局学习率虽然能压低读出层更新,也会进一步缩小输入层已经很小的更新。$\mu\mathrm{P}$ 通过逐层缩放,同时控制这几种宽度依赖。

实验设置、代码与数据

实验使用平方损失、输入维度 16、训练批量 32、探测批量 64,基准步长为 $10^{-3}$。输入为固定的高斯样本,目标由固定的非线性教师函数生成,并包含常数偏置,使初始批量的输出梯度不会因强制去均值而抵消。SP 与 $\mu\mathrm{P}$ 共享基础随机权重,前两层初始化相同;读出标准差分别为 $1/\sqrt n$ 与 $1/n$。二者的初始输出分布并不相同,所以这张图也不是性能优劣的比较。

实验将各层初始化常数设为 1,SGD 表里的三个 $\eta_i$ 都取同一个 eta,以观察宽度幂次的作用。$\mu\mathrm{P}$ 的更新为:

# W1: (n, d), W2: (n, n), W3: (1, n)
# g1, g2, g3 are gradients of the same mean squared loss.
W1 -= eta * n * g1
W2 -= eta * g2
W3 -= eta / n * g3

SP 则对三层都使用 eta。完整脚本还包含中心有限差分检查,用来核对手写反向传播;本次检查的相对误差低于 $3\times10^{-10}$。

下载完整 Python 脚本 · 下载逐种子实验数据 。在装有 NumPy 与 Matplotlib 的环境中,进入下载目录运行:

python3 mup-coordinate-check.py --csv coordinate-check.csv --figure coordinate-check.svg

所有随机种子、数据生成方式与统计定义都保存在脚本中。

这个实验只检查单步、分层的直接更新。完整的 coordinate check 还需追踪多个训练步的激活,以及同一探测输入上联合更新前后的特征变化。数据导致的梯度抵消、有限宽度和小学习率都会影响曲线,曲线平坦本身不足以判定实现正确。

超参数迁移的验证

Tensor Programs V 将这一思路用于 $\mu$Transfer[10] :在较小的模型上搜索超参数,再按照 $\mu\mathrm{P}$ 约定迁移到更宽的目标模型。

验证实现时,固定数据、深度、模块结构、优化器、批量和训练步数,只改变宽度,检查逐层激活、单层直接更新与联合更新的尺度。随后在小宽度上搜索基准学习率,并在较大宽度上复查:每个候选值都按同一套 $\mu\mathrm{P}$ 规则生成参数组学习率,其余训练与评估协议保持一致。

不同宽度的“基准学习率—验证损失”曲线及多个随机种子,可以检验小模型选出的学习率在大模型上是否仍接近最优;不同宽度的损失曲线不必完全重合。

本节的 MLP 实验没有扫描学习率,仅提供部分尺度检查的证据。

三、后续发展

$\mu\mathrm{P}$ 的超参数迁移首先在宽度方向展开。深度、模块结构、优化器、专家配置和训练预算的变化,也会影响初始化与更新的尺度。

谱条件

Tensor Programs V 用配套的初始化与更新规则实现宽度方向的超参数迁移[10] 。A Spectral Condition for Feature Learning 进一步用权重矩阵和更新矩阵的谱范数组织这些规则[14] ,把参数尺度与特征学习联系起来。

参数更新对特征的直接作用是 $(\Delta W)h$。对于 $n_{\rm out}\times n_{\rm in}$ 的矩阵,欧氏范数不等式给出

$$ \operatorname{RMS}((\Delta W)h) \leq \|\Delta W\|_{\rm op} \sqrt{\frac{n_{\rm in}}{n_{\rm out}}} \operatorname{RMS}(h). $$

谱范数 $\|\cdot\|_{\rm op}$ 衡量矩阵对输入方向可能产生的最大放大,形状因子则把向量总范数换成坐标 RMS。这个范数不等式只有上界信息,不能保证特征更新随宽度增加仍不消失,也不能推出有限宽模型的最优学习率一致。实际更新还取决于输入与更新方向的关系。

深度缩放

深度增长时,更多残差分支的贡献会在表示中累积。宽度方向的缩放规则没有规定这些分支的尺度及其相互作用。

Depth-$\mu\mathrm{P}$ 在 Tensor Programs VI 中研究先取无限宽、再取无限深的残差网络极限[11] 。除了特征更新的大小,它还关注不同深度位置产生的特征是否足够多样。

其中,每个残差块只有一层计算时的结果有明确的结构条件。论文也讨论了所研究参数化在更复杂残差块中的限制,因而不能将单层分支的结论直接用于任意 Transformer。

CompleteP 研究深 Transformer 的计算效率[13] 。即使基准学习率能够跨深度迁移,某些参数化的残差块对块内参数的响应仍可能在大深度极限下趋于线性化。CompleteP 将避免这种惰性作为额外设计条件,结合深度方向的超参数迁移与更充分的特征学习。这里的块内线性化与整个网络进入固定 NTK 极限不同。

Zheng 等人将谱条件扩展到宽度与深度的联合缩放,在残差块内变换数固定时,分别分析单变换和多变换分支[15] 。CompleteP for RL 在强化学习的控制任务中考察这种参数化[16] 。与固定数据集上的监督学习不同,策略更新会改变采样分布,特征学习也随这一过程演化。

模块与注意力

即使总宽度一样,改变注意力头数、KV 共享方式或模块倍率,也会改变参数对特征的作用。一个 Transformer 包含注意力、前馈层、归一化和残差分支,它们未必适合共用同一组倍率。

Mlodozeniec 等人将逐模块调参纳入迁移框架[17] ,研究学习率、AdamW 参数、权重衰减、初始化和残差乘子在宽度、深度、批量与训练时长变化时如何配合。逐模块迁移由文中的 Transformer 实验支持,批量与时长的规则还依赖随机微分方程(SDE)近似及所用调度。

GQA-$\mu\mathrm{P}$ 针对分组查询注意力分析投影结构与参数缩放[18] ,引入期望算子范数(expected operator norm),处理 KV 头重复形成的秩亏有效投影,并研究不同重复系数之间的学习率迁移。KV 头共享方式带来的差异,无法仅由模型的隐藏宽度描述。

优化器与有限宽度修正

Muon is Scalable for LLM Training 通过权重衰减和更新尺度调整,将 Muon 用于大规模语言模型训练[19] 。经过矩阵预条件或正交化以后,更新随矩阵形状变化的规律也会改变。

Extending $\mu\mathrm{P}$ 用谱条件为多种优化器推导宽度缩放规则[20] ;Xu 等人从矩阵算子范数下的优化几何出发,研究行列归一化与宽度迁移[21] 。这些分析把优化器的更新方向、矩阵形状和学习率尺度放在一起考察,不能直接沿用普通 SGD 的坐标缩放公式。

即使渐近尺度匹配,有限宽度下的最优学习率仍可能漂移。Qiu 等人研究 Shampoo、SOAP 和 Muon 的学习率及权重衰减迁移,并分析分块与显式谱归一化对有限宽偏差的修正[22] 。修正效果依赖优化器和块大小,叠加这些操作并不一定带来收益。

训练中的尺度控制也有不同做法。Muon++ 研究在学习率与谱间隙等条件下维持 $\mu\mathrm{P}$ 的谱条件[23] ,Spectral Sphere Optimizer(SSO)则对隐藏层二维权重及其更新施加模块级谱范数约束[24] 。Learnable Multipliers 让矩阵、行或列的倍率参与学习,以调整特征尺度[25] ;这减少了前向及权重衰减乘子的手工调节需求,逐模块学习率乘子仍需调优。

HyperP 围绕 Frobenius 球面上的 MuonH 优化,研究宽度、深度、token 预算与 MoE 粒度之间的迁移[26] 。其宽深分析包含近似各向同性、小步长等条件;随 token 预算调整学习率的关系则来自特定模型与预算范围的经验拟合。

MoE 与路由学习

MoE 增加了专家数、专家隐藏维度和每个 token 激活的专家数 top-$K$ 等缩放变量。早期的 $\mu$-Parametrization for Mixture of Experts 在固定专家数与 top-$K$、专家隐藏维度随模型宽度同比增长的设定下研究宽度迁移[27] 。当专家粒度也变化时,这套结果不能直接照搬。

后续工作开始联合分析这些变量。Jiang 等人利用动态均值场理论研究模型宽度、深度、专家数及专家隐藏维度的扩展;其中增加专家数的分析保持激活专家数与总专家数的比例不变[28] 。How to Scale Mixture-of-Experts 区分三种联合缩放区间,提出最大尺度稳定参数化 MSSP,并指出在这些设定中,最大更新条件本身不足以保证稳定的学习率迁移[29] 。Complete-muE 则通过 dense FFN、全专家激活的 dense MoE 到 sparse MoE 的映射,研究从密集模型向稀疏模型的近最优超参数迁移[30] 。

Why Routers Freeze 在固定专家数的无限宽极限下分析路由饱和、对称性与专家专门化[31] ,并讨论趋于零的路由 logits 有限宽波动如何仍然决定 Top-$K$ 选择、打破专家间的对称性。专家激活的尺度稳定,并不直接约束这些路由学习行为。

训练预算

小模型训练多少 token,也会影响最优学习率。Zhou 等人考察模型规模、深度、权重衰减和训练 token 总量(token horizon)对最优学习率的影响,并比较拟合与迁移方法[32] ,报告了所测迁移配方在大规模预训练中的局限。

Scaling with Collapse 研究 $\mu\mathrm{P}$ 下不同规模训练损失曲线的近似重合[33] 。其设置保持学习率日程、每参数 token 数(TPP)和归一化 AdamW 时间尺度一致,将横轴换成训练完成比例,并用最终损失归一化纵轴。这一归一化用于训练诊断与调参早停,曲线重合依赖上述预算和优化设置。

Let’s Scale Step by Step 将模型规模迁移与 token-horizon 校准分步应用于采用 Muon 和多头潜在注意力(MLA)的 MoE[34] :在激活专家数和每个专家的隐藏维度固定时,联合扩大模型宽度和总专家数,再用经验规律外推更长的训练预算。宽度迁移与训练时长校准是两个步骤,前者的结果不能代替对后者及其组合效果的验证。

数值精度

u-$\mu\mathrm{P}$ 将最大更新参数化与 Unit Scaling 结合[12] ,以权重、激活和梯度在初始化时接近单位尺度为目标,协调低精度训练和超参数迁移。这个目标依赖推导缩放因子时的统计假设;论文也观察到,注意力引入的相关性会使部分中间张量在初始化时就偏离单位尺度。因此,它并不保证所有张量从初始化到整个训练过程都严格保持单位方差。


  1. 第一部分对两种均值场路线的区分,参考了 LawrenceC 的 Maybe I was too harsh on deep learning theory (three days ago) 。公式与技术结论引用对应的原始论文。 ↩︎

  2. 文中的量级分析把数据、输入维度和相应训练时间约定固定。若这些量也随宽度变化,需要重新分析。上文重参数化的学习率推导针对数据损失的普通 SGD,额外正则项与 Adam 更新需要单独处理。 ↩︎

  3. Matthews 等人的 Gaussian Process Behaviour in Wide Deep Neural Networks 也研究了深网络的高斯过程极限,各项工作的取极限方式与假设有所不同。 ↩︎

  4. 若训练时间、学习率或其他尺度也随宽度增长,结论可能变化。有限宽网络可以偏离 NTK 极限,并学习新的特征。 ↩︎

  5. 对独立零均值的初始化贡献,均值场的初始输出波动通常随宽度缩小;这与训练中的非退化特征变化相容。 ↩︎

  6. 这里指不含权重衰减的 Adam 尺度分析,假设 $\epsilon$ 相对二阶矩平方根足够小,或随参数坐标一致调整。若 $w=s\theta$ 且 $s>0$,梯度归一化理想情形下有 $\eta_w=s\eta_\theta$,而 SGD 是 $\eta_w=s^2\eta_\theta$。保留 $\epsilon$ 时,还需匹配 $\epsilon_\theta=s\epsilon_w$。实际 MuAdam 不会自动替用户缩放 $\epsilon$;AdamW 的解耦权重衰减也需要另行匹配。 ↩︎

引用本文

bibtex
@misc{jiang2026mup,
  author = {Jiwen Jiang},
  title = {{\ensuremath{\mu}P}:前身、原理与后续发展},
  year = {2026},
  url = {https://jiwenj.github.io/blog/posts/mup/}
}