A Proof of Learning Rate Transfer under $μ$P

1. 研究背景与问题 (Problem)

在深度学习的扩展（Scaling）中，随着神经网络宽度（hidden dimension）的增加，超参数（Hyperparameters, HPs）通常需要重新调整以避免数值溢出或训练不稳定。

• 核心问题：是否存在一种参数化方法，使得在宽度趋于无穷大时，最优学习率（Optimal Learning Rate, LR）收敛到一个非零常数？如果存在，则意味着可以在小模型上微调学习率，然后直接迁移到大模型上（即学习率迁移，Learning Rate Transfer），从而大幅降低调参成本。
• 现有局限：
  • 标准参数化 (Standard Parametrization, SP)：随着宽度增加，最优学习率通常趋向于 0，导致无法直接迁移。
  • 神经切线核参数化 (NTP)：导致“懒惰训练”（Lazy Training），特征学习几乎停止，最优学习率随宽度发散或行为异常。
  • 最大更新参数化 (µP)：由 Yang 和 Hu 提出，旨在最大化无限宽度下的特征学习。虽然经验上观察到 µP 具有学习率迁移特性，但此前缺乏严格的数学证明，特别是针对多层网络（Multi-layer）的情况。

2. 方法论 (Methodology)

作者通过理论推导和实证实验，在线性多层感知机 (Linear MLP) 的框架下证明了 µP 下的学习率迁移性质。

2.1 模型设定

• 模型：深度为 $L$ 的线性 MLP，$f(x) = V^\top W_L \dots W_1 W_0 x$。
• 训练：使用梯度下降 (GD) 最小化二次损失函数。
• 参数化定义：
  • µP：初始化方差 $W_\ell \sim \mathcal{N}(0, n^{-1})$，$V \sim \mathcal{N}(0, n^{-2})$，学习率 $\eta$ 为常数（不随 $n$ 缩放）。
  • SP：$V \sim \mathcal{N}(0, n^{-1})$，其他同 µP，学习率 $\eta$ 为常数。

2.2 核心分析思路

作者的关键洞察是：在固定训练步数 $t$ 下，模型输出 $f^{(t)}(x)$ 可以表示为学习率 $\eta$ 的多项式函数。
$$f^{(t)}(x) = \sum_{k} c_k \eta^k$$
其中系数 $c_k$ 取决于初始化权重。

• 分析策略：
  1. 研究当宽度 $n \to \infty$ 时，该多项式各项系数的渐近行为（收敛性）。
  2. 分析损失函数 $L(\eta)$ 作为 $\eta$ 的函数的极限形式。
  3. 证明在 µP 下，最优学习率 $\eta_n^*$ 收敛到一个非零确定性常数 $\eta_\infty^*$；而在 SP 下，$\eta_n^* \to 0$。

3. 主要贡献与理论结果 (Key Contributions & Results)

3.1 单步训练 ($t=1$) 的完整刻画

• µP 下的收敛性：
  • 证明了在 µP 下，损失函数关于 $\eta$ 的多项式中，只有一次项（线性项）的系数收敛到一个非零常数，而所有高阶项（$\eta^k, k \ge 2$）的系数在 $L^2$ 范数下以 $O(n^{-(k-1)/2})$ 的速度收敛到 0。
  • 因此，当 $n \to \infty$ 时，损失函数渐近地表现为关于 $\eta$ 的二次函数。
  • 定理 1：最优学习率 $\eta_n^{(1)}$ 以概率收敛到一个非零常数 $\eta_\infty^{(1)}$，且收敛速度为 $O(n^{-1/2})$。
  • 极限值公式：$\eta_\infty^{(1)} = \frac{m}{L} \frac{y^\top K y}{\|Ky\|^2}$，其中 $K$ 是归一化的输入 Gram 矩阵。
• SP 下的失效：
  • 定理 2：在 SP 下，由于 $V$ 的初始化方差不同（$n^{-1}$ 而非 $n^{-2}$），导致所有多项式系数被放大了 $\sqrt{n}$ 倍。为了补偿这种增长，最优学习率必须趋向于 0（$\eta_n^{(1)} \xrightarrow{P} 0$）。因此，SP 不具备学习率迁移能力。

3.2 多步训练 ($t \ge 2$) 的推广

• 复杂性增加：在 $t \ge 2$ 时，模型输出关于 $\eta$ 的多项式阶数显著增加（例如 $t=2$ 时最高阶可达 $3L-1$），且系数依赖于之前的梯度更新，分析更为复杂。
• 定理 3 (一般步数)：
  • 证明了在 µP 下，对于任意训练步数 $t$，模型输出 $f^{(t)}(x)$ 几乎必然收敛到一个确定性系数的多项式 $f_\infty^{(t)}(\eta)$。
  • 损失函数 $L_n^{(t)}(\eta)$ 一致收敛到极限损失 $L_\infty^{(t)}(\eta)$。
  • 在假设极限损失存在唯一最小值的前提下，证明了最优学习率 $\eta_n^{(t)}$ 几乎必然收敛到非零常数 $\eta_\infty^{(t)}$。
  • 结论：µP 下的学习率迁移性质对任意深度 $L$ 和任意训练步数 $t$ 均成立。

3.3 实证验证

• 线性网络：在不同宽度（$n=128$ 到 $8192$）下，µP 的最优学习率稳定收敛，而 SP 的最优学习率随宽度增加迅速衰减至 0。
• 扩展实验：
  • 深度影响：在 $L=3, 9, 27$ 的不同深度下，µP 均表现出学习率迁移。
  • 非线性与优化器：在引入 ReLU 激活函数和使用 Adam 优化器的情况下，经验结果依然支持学习率迁移，尽管理论证明目前仅限于线性网络和 GD。

4. 意义与影响 (Significance)

1. 理论突破：这是首个关于 µP 下学习率迁移的严格数学证明。它填补了从经验观察到理论保证之间的空白，解释了为什么 µP 能够实现“零样本超参数迁移”（Zero-shot Hyperparameter Transfer）。
2. 机制解释：证明了 µP 通过特定的初始化缩放（特别是输出层 $V$ 的 $n^{-2}$ 缩放），使得在无限宽度极限下，特征更新（Feature Learning）保持为 $\Theta(1)$ 量级，从而避免了特征更新的消失或爆炸，保证了最优学习率的稳定性。
3. 实践指导：
   • 为大规模模型训练提供了理论依据：研究者可以在较小的模型上高效地搜索最优学习率，然后直接应用于更大规模的模型，无需昂贵的重新调参。
   • 明确了 SP 和 NTP 在超参数迁移上的局限性，强调了 µP 在扩展性（Scalability）方面的优势。
4. 未来方向：虽然当前证明基于线性网络，但作者指出其结论很可能适用于非线性网络，未来的工作将致力于将证明扩展到更复杂的架构和优化器。

总结

该论文通过深入分析线性 MLP 在无限宽度极限下的多项式结构，严格证明了 µP 参数化 能够保证最优学习率随宽度增加而收敛至非零常数，从而实现了学习率迁移。相比之下，标准参数化 (SP) 会导致最优学习率趋于零。这一结果为 µP 在大规模深度学习中的广泛应用提供了坚实的理论基础。