随分な垢

\[\newcommand{\d}{\mathrm d} \newcommand{\t}{\text} \newcommand{\E}{\mathop{\mathbb E}\limits} \newcommand{\b}{\boldsymbol} \newcommand{\p}{\partial} \newcommand{\c}{\mathcal} \newcommand{\ip}[1]{\left<#1\right>} \]

O. Overview

Optimizer：输入梯度，输出参数更新幅度。本质需要的超参只有 \(\eta\) 一个。

• SGD 的更新幅度与梯度大小有关。
• RMSProp / Adam 一系因为有归一化，所以对梯度大小不敏感；同时参考 AdaGrad 所以相当于动态优化一个 Mahalonobis 范数，为不同变量赋以不同权重。
• 动量除了作为记忆，还起到高频滤波器的作用，降低方差。

Scheduler：研究 \(\eta\) 如何 decay。不带 lr decay 的 SGD 只会震荡，带了的才能保证收敛。

• Robbins-Monro 条件保证 \(O(1/t)\) 的 lr decay 在 SGD 上能收敛到最优解。
• cosine schedule 经验上表现好。
• WSD 符合 river-valley 模型。

(可选) Normalizer：带来 scale invariance。

• 作为负反馈动力学，在参数大时传来的梯度较小。虽然在这方面和 Adam 功能重合，不过它也有其它用处，诸如：
• 防止前向传播的激活值消失/爆炸。
• 优化 Hessian 特征值 landscape，减少耦合性，对 Adam 更友好。
• 处理「协变量偏移」也即随着网络加深激活值整体向某个方向进行偏移的现象。
• 减轻层间耦合。

(可选) Weight Decay：带来 regularization。

• SGD 时与 regularization 等价，但是 Adam 时不然。
• 我们真正在意的其实是相对变化量 \(\Delta\b W/\b W\)。Adam 让 \(\Delta\b W\) 规模固定，那么 \(\b W\) 的规模也不能无限膨胀，所以需要 WD。

上面的东西介导单个模型的动力学，而下面的东西研究一系列模型关于某个超参变化时，其它超参应该怎么跟着变。

• μP 本体研究初始参数和初始 lr 与网络宽度的关系。
• critical batchsize 研究 lr 与 batch size 的关系。

I. Formulation

把多层神经网络看做一个对输入分布 \(\mu_0\) 的多次 push-forward，最终得到神经网络输出分布 \(\mu_L\)，和答案分布 \(\rho\) 算某种 loss。

令初始分布为 \(\nu_0,\dots,\nu_L\)，则 μP 的目标可以概述为：

• 在均方根距离 \(\E\|\b x-\b y\|_\t{RMS}:\b x\sim\mu,\b y\sim\nu\) 下，所有分布的改变量为 \(\Theta(1)\)：既不是 \(o(1)\)（NTK regime）也不是 \(\omega(1)\)（unstable）。
• 输出的 loss \(\ell(\mu_L,\rho)\) 的变化量是 \(\Theta(1)\) 的（non-trivial）。

II. μP w/o Tensor Program

考虑 \(\b z=\b W\b x\) 的线性层，使用下标 \(t\) 表示时刻，令 \(\b g_t=\dfrac{\p\c L}{\p\b z}\)，则 \(t\) 时刻时反向传播则有

\[\b W_{t+1}\gets\b W_t+\eta(\b g_t\b x_t^\top) \]

那么 \(t+1\) 时刻的正向传播则是

\[\b z_{t+1}=\b W_{t+1}\b x_{t+1} \\=(\b W_t+\eta\b g_t\b x_t^\top)\b x_{t+1} \\=\b W_t\b x_{t+1}+\eta\b g_t(\b x_t^\top\b x_{t+1}) \\=\b W_t\b x_{t+1}+n\eta\b g_t\cdot\dfrac{\b x_t^\top\b x_{t+1}}n \]

我们一般认为 \(\b x\) 是各向同性的，因此 LLN 和 CLT 适用于 \(\b x_t^\top\b x_{t+1}/n\) 的分析。

不妨认为因为归一化层的存在，有 \(\E\b x=\b0\)。现在考虑 \(\b x_t\) 和 \(\b x_{t+1}\) 的关系：

• 如果它们完全独立，则 \(\E\b x_t^\top\b x_{t+1}=0\)，此时适用于 CLT 的分析，有 \(\b x_t^\top\b x_{t+1}\sim\sqrt n\sigma\c N(0,1)=\Theta(\sqrt n)\)。
• 否则它们不完全独立，在某一维上均值是 \(\mu\)，则由 LLN 可知 \(\b x_t^\top\b x_{t+1}\) 的主项是 \(n\mu=\Theta(n)\)。

初始时网络近似完全随机，因此 \(\b x_t\) 和 \(\b x_{t+1}\) 作为来自不同 batch 的前向传播结果，没有明显的相关性，因此符合 CLT 的场景。但是随着训练进行，前向传播结果的相关性不断增强，最终过渡到 LLN 的场景。所以 LLN 会是更一般的考虑，我们默认 \(\dfrac{\b x_t^\top\b x_{t+1}}n=\Theta(1)\)，现在考虑 \(n\eta\b g_t\) 这一项。

• 裸的 SGD 中有 \(\b g_t=\Theta(1/n)\)，因此固定学习率 \(\eta=\Theta(1)\) 即可保证训练稳定。
• Adam 则因为基于 SignSGD，可以认为 \(\b g_t\) 是 \(\Theta(1)\) 的，那么为了抵消就必须有 \(\eta=\Theta(1/n)\)。

上述分析是 \(n\to n\) 的线性层。然而第一层是 \(d\to n\)（固定→可变），最后一层是 \(n\to d\)（可变→固定），它们的分析有所不同。

• 第一层的 \(\b x_t^\top\b x_{t+1}\) 本就是常数，所以 SGD 取 \(\Theta(n)\)，Adam 取 \(\Theta(1)\)。
• 最后一层因为直接与 loss 对接，不论 SGD 还是 Adam \(\b g\) 都是 \(\Theta(1)\) 的，于是学习率都是 \(\Theta(1/n)\)。

特别地，Adam 对初始化没那么敏感（反正 \(\b g\) 总是 \(\Theta(1)\)），因此就算照抄 SP 的初始化问题也不大，只需要注意把第一层的 lr 取 \(\Theta(1)\) 才能保证 embedding 学习。只改学习率的 μP 被称作 μP(Simple)，也能 work。

而 SGD 的 \(\b g=\Theta(1/n)\) 的推导严格依赖于初始化。我们展开类似分析。令后一层的梯度为 \(\b g\)，前一层的梯度为 \(\b g'\)，则有

\[\b g_t'=\b W_t^\top\b g_t \\=(\b W+\Delta\b W)^\top\b g_t \]

其中 \(\b W^\top\b g_t\) 描述初始化，而 \(\Delta\b W^\top\b g_t\) 描述训练。

• 初始化假设 \(\b W\sim\c N(0,\sigma^2I)\)，则非最后一层可以用 CLT 得到级别为 \(\Theta(\sqrt n\sigma\b g_t)\)，最后一层因为是常数项求和所以直接是 \(\Theta(\sigma\b g_t)\)。
• 训练假设相关性，则分析类似前文，在非最后一层用 LLN 得到 \(\Theta(n\eta\b g_t)\)，最后一层因为是常数项得到 \(\Theta(\eta\b g_t)\)。

我们的目标是初始化信号和训练信号之和是 \(\Theta(1/n)\) 级别的。前文对 \(\eta\) 的设置已经保证了训练信号，而初始化信号就要靠 \(\sigma\) 的正确设置了。

现在使用这个结论确定初始化规模。

• 最后一层的 \(\b g_t\) 是 \(\Theta(1)\)，为了保证 \(\b g'_t=\Theta(1/n)\)，考虑到 \(\eta=\Theta(1/n)\)，所以取 \(\sigma^2=O(1/n^2)\) 即可。注意到这里取的是大 \(O\) 而不是大 \(\Theta\)，因为 \(\eta\) 已经保证了主项。因此实际使用时可以直接设成 \(0\)——虽然这会导致梯度截断，第一步完全没有梯度回传，不过这一步之后 \(\b W_t\) 就用训练数据中的信息更新为非零，整个训练流程也就打通了。另一方面，也不用担心全 \(0\) 带来的维数对称性，因为最后一层直接对接 loss，而我们有理由相信 loss 是非对称的，因此一次反传自动让最后一层对称破缺。
• 中间层希望 \(\sqrt n\sigma=\Theta(1)\)，因此 \(\sigma^2=\Theta(1/n)\)。特别地，这里就要担心对称性了，所以不能设为全 \(0\)。
• 第一层的初始化就不是为了反向传播了（再往前的梯度也没用），而是为了正向信号。因为第一层的输入端固定，所以必须取 \(\sigma=\Theta(1)\)，否则会随着 \(n\to\infty\) 导致正向传播消失。

III. μP is not Universal...

μP 研究在其它条件固定时，学习率/初始化等与模型大小的关联，无法研究其它东西的关联。

• 同一个模型使用不同量的数据训练，则对应的最优 lr 同样不同：更多的数据需要更小的 lr（left shift）。
• 同一个模型使用不同的 batch size 训练相同的 step（实际收到的数据量不同），在不超过 CBS 前，最优 lr 随着 batch size 线性 scale（SGD）或根号 scale（RMSProp, Adam）。
• 进一步，还有对模型层数的 scale。

因此理想的 scaling law 应该涉及 witdth、depth、#token 和 bsz 等多个维数。同时，μP 基于大数定律的一套要在 \(n\to\infty\) 时才生效——而我们的模型是否足够大还要打问号。总之 μP 远称不上完善。

IV. ...But We Can Fix!

不过还是有一些结论可以给出的。

第一个应用是用来调整 AdamW 中的 weight decay。

[2405.13698] How to set AdamW's weight decay as you scale model and dataset size 分析了如何设置 WD 的系数。

首先一个基础的观察是，WD 的系数必须随着 lr schedule 一起 schedule，不然在训练后期就会 WD 太多。因此可以列出一个基础的 WD 式

\[\b\theta_t=\b\theta_{t-1}-\eta_\t{peak}\gamma_t\b g_t-\lambda\gamma_t\b\theta_{t-1} \]

其中 \(\gamma_t\) 是 scheduler 给出的系数，而 \(\eta_\t{peak}\) 则是最大学习率。

现在分析 \(\lambda\) 应当如何设置以保证 WD 的效果可迁移。具体地，把 WD 写成 EMA 的式子

\[\b\theta_t=(1-\underbrace{\lambda\gamma_t}_\alpha)\b\theta_{t-1}-\eta_\t{peak}\gamma_t\b g_t \]

这个 \(\alpha\) 就是 EMA 的特征系数，而 \(\tau_\t{iter}=\alpha^{-1}\) 就是强度衰减到 \(1/e\) 的时刻。

现在 μP 会改变 \(\eta_\t{peak}\)，但是如果没有动 scheduler，则只要保证实际的 \(\lambda\) 不变就可以保证 WD 的动力学相同。

那么如果同步做 data scaling 呢？这是 \(\tau_\t{iter}\) 就不再是最佳的指标了，此时更喜欢 \(\tau_\t{epoch}=\tau_\t{iter}/M\) 其中 \(M\) 是单个 epoch 包含的 step 数，它表示一条数据在再次看到之前的衰减幅度。只需要保证 \(\tau_\t{epoch}\) 在跨数据集迁移中相同即可。

[!TIP]

这篇论文的意义在于，torch 的默认 AdamW 实现是

\[\b\theta_t=\b\theta_{t-1}-\eta_\t{peak}\gamma_t\b g_t-\eta_\t{peak}\lambda\gamma_t\b\theta_{t-1} \]

强行把 WD 与 \(\eta_\t{peak}\) 挂钩。但是现在如果我们根据 μP 来 scale \(\eta_\t{peak}\)，再沿用原本的 \(\lambda\) 就不合适了。本文指出了这一点，认为 \(\lambda\) 和 \(\eta_\t{peak}\) 当是解耦的。

[2505.13738] Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training 则进一步将分析迁移到如 LLaMA 这种 overtrained 的场景。此时所有数据都只看到一次，模型需要更快地遗忘远古数据，因此最优的 \(\tau_\t{epoch}\)（论文中改叫 \(\tau_\t{EMA}\) 不过实际意义一样，都是 \(\tau_\t{iter}\) 除以 step 数）关于 \(\t{TPP}=D/N\) 呈现 power law 式的指数衰减——数据越多就必须忘得越快。

除了这些有数学支持的规则外，还有其它大力出奇迹的 scaling law！

[2401.02954] DeepSeek LLM: Scaling Open-Source Language Models with Longtermism：固定总 FLOPs \(C\) 时的最优 lr 和最优 bsz。

[2503.04715] Predictable Scale: Part I, Step Law -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining：最优 lr 和 bsz 随数据集增长的关系：最优 bsz 与模型大小无关；最优 lr 随模型大小而减少（符合 μP 但指数上没到 \(n^{-1}\)）；稀疏性（如 MoE）不影响最优超参设置。(*不过有人测了后发现这玩意好像不太对得上？)

特别地，这种经验性的 scaling law 都只针对特定架构（保持 FLOPs/token 一致）和特定 scheduler 有效！

---

只要不超过 cbsz，模型表现理论上应与 bsz 无关；但是如果只调 lr 不调其它超参（beta, epsilon, WD）则 optimal batch size 这个概念还是有意义的。

cbsz 与模型大小无关，而与训练的 token 数、训练时长和优化器（Adam/Muon）有关。

V. QKLayernorm

在 CNN 中有很多种 norm，但 LLM 中几乎只用 LayerNorm 一种：对于 \((B\t{-atchsize},S\t{-equence},H\t{-ead},D\t{-im})\) 的 tensor，LN 只对最后一维也即 hidden dimension 做归一化，其它三维独立。大家会在各种地方加归一化，诸如

• QK 后，防止 attention 矩阵出现过大的异常值。
• 混合 pre norm 和 post norm。

特别地，QK-LayerNorm 可以增大 lr basin，减轻对 lr 的敏感度，且与 μP 是兼容的（因为 μP 无法处理 bsz 和 horizon 变化的场景）

但是乱放 QK-LayerNorm 也会有问题，例如在长上下文环境中它可能给关键位置分配的注意力权重不够。

VI. Muon

考虑 RMS 诱导的矩阵范数

\[\|\b W\|_{\t{RMS}\to\t{RMS}}=\max_\b x\dfrac{\|\b W\b x\|_\t{RMS}}{\|\b x\|_\t{RMS}} \]

[!TIP]

为什么使用 RMS？因为 μP 用了它！

众所周知 RMS 范数和 \(\ell_2\) 范数通过 \(\|\b x\|_\t{RMS}=n^{-1/2}\|\b x\|_2\) 挂钩，而 \(\ell_2\) 范数诱导的是谱范数 \(\|\b W\|_2=\sigma_\max\)，故

\[\|\b W\|_{\t{RMS}\to\t{RMS}}=\sqrt{\dfrac{\t{fan-in}}{\t{fan-out}}}\|\b W\|_* \]

于是若有 \(\b y=\b W\b x\) 则有

\[\|\Delta \b y\|_\t{RMS}\leq\|\Delta\b W\|_{\t{RMS}\to\t{RMS}}\|\b x\|_\t{RMS} \]

现在目标是在满足 \(\|\Delta \b y\|_\t{RMS}\leq\eta\) 的前提下，寻找 \(\min_{\Delta\b W}\ip{\nabla_\b W\c L,\Delta\b W}\)。使用 \(\|\Delta\b W\|_{\t{RMS}\to\t{RMS}}\leq\eta\) 来代替。

这个最优解可以显式解得为 \(\Delta\b W=-\eta\sqrt{\dfrac{\t{fan-out}}{\t{fan-in}}}\t{msign}(\nabla_\b W\c L)\)。