A2 · Hyperparameter Scaling：从一维扫描到联合 scaling

做什么

题面来自官方材料（本地原件在 _local/，官方链接见「材料」）。 编号是原件的编号（notebook 是 cell 序号、PDF 是页码），方便跟官方题面对号。 本项目不提供答案（D53）。

第 1 页

Assignment 2: Hyperparameter Scaling（作业 2：Hyperparameter Scaling）

CS 312

说明（Instructions）. 本作业不计分，但完成作业会显著帮助你在这个课程的 quiz 中获得部分分数。请按课程日程中的截止时间，向 Gradescope 提交一份 PDF writeup。

请将 PDF 提交限制在 20 页以内，bonus question 不计入页数。

对于每个问题，你的交付物包含两部分：对所提供的实验以及每个问题中要求的新实验进行分析；以及一段简短（约一段话）的综述，说明你对这些实验的解读。请用这两部分来检查并发展你的理解。请记住，你可以在 quiz 中引用你的作业来为自己的答案提供依据；因此，为了让 quiz 获得部分分数，准备一套详细、组织良好且你完全理解的实验对你最有利。

在最后一个可选问题中，你可以设计一个与我们的 quiz 类似、有趣的结果预测问题：你需要提交一段来自我们代码库的简短 code diff（少于 100 行），以及一个预测任务，形式为 single-answer 或 multiple-answer 的 multiple-choice question。如果我们选择以任何方式（包括以某种修改后的形式）在实际 quiz 中使用你提交的 quiz question，你将获得相当于该 quiz question 分值的 bonus points。示例问题的答案在本讲义末尾。

概览（Overview）

训练神经网络要求我们推理 scale：我们会训练更大的模型、使用更多数据，并增大 batch sizes。在某个 scale 下有效的 hyperparameters，在另一个 scale 下可能并不有效。在本作业中，你将研究 hyperparameter scaling：当模型和训练 scale 变化时，如何调整 hyperparameters。

本作业涉及以下问题：

• 当我们用更多数据训练时，learning rate 和 weight decay 这类 hyperparameters 应如何缩放？
• initialization 和 learning-rate scaling 能否让我们在扩展时复用同一个 optimal learning rate？
• 缩放 batch size 会如何影响 hyperparameters？

目标是了解：一个 scale 下的实验何时能够指导另一个 scale 下的训练，以及如何检验这些 scaling rules 的极限。

作为这一过程的一部分，我们会仔细研究 toy examples：一些我们有可靠理论和预测的简单 optimization tasks。Toy examples 很重要，因为它们构成了研究者许多基于理论的直觉的基础；理解这些模型的威力和局限，对于理解我们应当在多大程度上信任 deep learning optimization 中的理论至关重要。

第 2 页

我们的模型（Our model）

我们使用 Llama 风格的架构，基线维度如下：8 层、hidden size 512、8 个 attention heads、4,096-token 的 vocabulary，以及约 35M parameters。模型使用 query/key (QK) normalization，即在每个 attention head 内对 queries 和 keys 应用 root mean square normalization。基线 context length 为 1,024 tokens。RoPE 不缩放，训练使用完整 source dataset 的一个 prefix，并以 seed 42 打乱。我们在 fresh data 上训练，并评估 validation loss。在基线模型大小和 batch size 为 64 sequences 时，在单张 H100 上训练 614.4M tokens 大约需要 15 分钟。

提供的实验（Provided experiments）

对于 Problems 1 和 2，使用提供的 source sweeps。在 starter repository 根目录下，用以下命令列出某个 subpart 的 run links：

uv run python -m experiments.a2.provided_sweeps --part P1a --links

对于其他 supplied runs，将 P1a 替换为 P1b、P1d、P1e 或 P2a。该命令会列出每个 run 的 token budget、optimizer、schedule、learning rate、weight decay 和 W&B link。使用不带 --links 的命令可离线读取 measurements。

预计实验（Estimated experiments）

建议的 grids 总共约有 166 个新 runs，其中包括 55 个 five-step tests；supplied runs 和可选探索（包括 Muon）不计入。跨 parts 复用匹配的 runs。请保持在 48 GPU-hour budget 之内：计划最多用 36 小时做实验，并预留 12 小时用于 debugging 和 reruns。在启动 sweeps 前，先为每个 model size 计时一个代表性 run；按需缩小 grids，并报告任何未测试的比较。

默认训练设置（Default training settings）

除非某个问题另有说明，否则 language-model experiments 使用以下默认设置。用 $\eta$ 表示 peak learning rate (LR)，用 $\lambda$ 表示 weight-decay coefficient (WD)。每个问题都会指定其 token budget。默认模型是 d8，训练设置如下。

第 3 页

AdamW：带 decoupled weight decay 的 Adam

在本作业中，我们的许多问题会用到 AdamW。为完整起见，我们在此定义这里所说的 AdamW optimizer 是什么。AdamW 将 adaptive gradient update 与单独的 weights shrinkage 结合起来。在 update $t \ge 0$ 时，令 $g_t$ 为提供给 Adam 的、关于 weight matrix $W_t$ 的 gradient。给定 moment coefficients $\beta_1, \beta_2$ 和数值稳定项 $\epsilon > 0$，定义 gradient 与 squared-gradient moments $m_t, v_t$、它们的 bias-corrected 版本 $\hat{m}_t, \hat{v}_t$，以及下面的 Adam direction $U_t$。从 $m_0 = v_0 = 0$ 开始；所有运算均为 elementwise。

$$ \begin{aligned} m_{t+1} &= \beta_1 m_t + (1-\beta_1) g_t, \\ v_{t+1} &= \beta_2 v_t + (1-\beta_2) g_t^2, \\ \hat{m}_{t+1} &= \frac{m_{t+1}}{1-\beta_1^{t+1}}, \\ \hat{v}_{t+1} &= \frac{v_{t+1}}{1-\beta_2^{t+1}}, \\ U_t &= \frac{\hat{m}_{t+1}}{\sqrt{\hat{v}_{t+1}}+\epsilon}. \end{aligned} $$

给定当前 LR $\eta_t$，AdamW 按如下方式更新该矩阵：

$$ W_{t+1} = W_t - \eta_t U_t - \eta_t \lambda W_t. $$

我们遵循常见做法，将 embeddings、normalization parameters 和 biases 排除在 weight decay 之外。

使用 AdamW 的 learning-rate scaling

Optimal learning rates 取决于其他 hyperparameters，例如 batch size。当我们扩展数据或模型时，它们应如何变化？

我们可能会设想，这个问题有许多可能的答案。由于模型相同，且 optimal learning rate 通常相当稳定，我们可能会设想 optimal learning rate 保持相近。或者，我们可能会查看 convex optimization 的理论，并根据收敛定理得出结论：正确的做法是在 $T$ 步内以 $O(1/\sqrt{T})$ 或 $O(1/T)$ 的速率降低 learning rates。

这是一个简单，但出奇复杂且影响重大的主题。即便针对这个问题进行了大规模、严肃的实证努力，也会导致分歧，正如你从 StepFun learning rate scaling paper 中这个例子所看到的那样。

文献中识别出的各种 hyperparameter scaling rules 的示例。注意：可信的论文与机构之间存在广泛分歧。

现在我们将研究 learning rate scaling 在我们的 d8 系统中如何运作，并看看能否对哪些类型的 scaling rules 合理建立直觉。通篇用 $D$ 表示 training-token budget。

第 4 页

Problem 1: When can training on less data guide large scale runs?（9 个新 runs）

使用上面的 baseline model 和默认训练设置。改变 LR 和 training-token 预算，其他设置保持不变。对于 (a) 和 (b)，分析提供的 W&B sweeps 或 starter 的 CSV。自行拟合曲线并作图。

(a) 从小训练预算中学习（provided runs）

在 153.6M、307.2M 和 614.4M tokens 下，使用提供的在 $\{.0015, .003, .006\}$ 处的测量值，研究最终 validation loss 如何随 peak LR 变化。对每个 scale，拟合一条将 peak LR 映射到 loss 的曲线，并据此估计最优 LR。用这些最优值提出一个关于 training-token 预算的 scaling rule。画出拟合出的最优值和你的 rule。你观察到什么趋势？

(b) 在增加预算之前先预测（provided runs）

在查看 1.2288B、1.8432B 和 2.4576B tokens 处提供的测量值之前，先记录你预测的这些预算下的最优 LR。然后拟合它们的 loss–LR 曲线，并将你的预测与拟合出的最优值比较。来自小预算的趋势是否延续？

(c) 比较并测试两种 scaling rules（5 个新 runs）

在来自 (a) 和 (b) 的全部六个 source budgets 上拟合一个 learning-rate power law，并仅在 1.2288B、1.8432B 和 2.4576B tokens 上拟合另一个。用每个规则预测 4.9152B 处的最优 LR。你预期哪个预测会给出更低的 loss，为什么？

在两个预测的 LR 以及 target-grid LRs $\{.0015, .003, .006\}$ 处训练。将你的预测与拟合出的 target optimum 比较，并报告它们相对于 target grid 中最佳采样 LR 的 measured loss gaps。

使用较小预算的拟合是否带来了收益？思考拟合 scaling laws 时的 bias-variance tradeoffs，以及你能将 scaling law 拟合外推多远。

(d) Hyperball 与 AdamW（推荐：4 个新 runs）

learning rate scaling 的形状和可预测性可能因 optimizer 不同而有显著差异，而新的、非传统优化方法的一个动机是它们可能比标准方法 scaling 得更可预测。为此，我们将探索 Hyperball，它是对已有充分研究的 Adam optimizer 的一个简单修改。Hyperball 使用 Adam 方向，但约束每个被选中的 weight matrix 保持其初始 norm。对于线性层 weight matrix $W_t$，令 $U_t$ 为 Adam update direction。用 $\|\cdot\|_F$ 表示 Frobenius norm（各元素平方和的平方根）。对于非零 norms，Hyperball update 为

$$W_{t+1} = \frac{\|W_t\|_F}{\|\tilde{W}_{t+1}\|_F}\tilde{W}_{t+1}$$

$$\tilde{W}_{t+1} = W_t - \eta_t \frac{\|W_t\|_F}{\|U_t\|_F} U_t.$$

我们通过相对于 weight norm 缩放标准 Adam update 来提出一个 update $\tilde{W}_{t+1}$，并强制在 step 之后 norm 保持不变。由于 hyperball 已经控制 norms，Hyperball 不使用 weight decay。Embeddings、normalization parameters 和 biases 使用普通 Adam。

所提供的 runs 对 Adam 参数使用 peak LR $(.000656/.00630\,\eta \approx .1041\eta)$，其中 $\eta$ 是所列出的 Hyperball peak LR。对两组应用相同的 warmup 和 decay multiplier，并在你的 target runs 中保持这个 LR ratio 固定。

使用 153.6M、307.2M 和 614.4M tokens 下提供的 24 个 Hyperball runs（P1d），每个预算有八个 LRs。在这三个预算处分别拟合一个 optimum，并对这些 optima 拟合一个 power law。将其趋势与在同样三个预算上拟合的 AdamW 比较；没有提供更大预算的 Hyperball source budgets。在运行之前，记录一个 1.2288B tokens 处的预测。

第 5 页

（接上页）该 target。测试预测的 LR 以及你选择的三个附近 LRs（最多四个新 runs，可复用重复项）。将预测与最佳采样的 target loss 比较。最优 LR 是否随预算变化得更规律？讨论仅从三个 source budgets 外推的不确定性。

(e) 探索 schedule variants（可选）

到目前为止，我们研究的是单一 schedule 下的最优 LR scaling。我们的发现能否迁移到不同 learning rate schedules？我们能否直接把最优 LR 从一个 schedule 迁移到另一个？使用相同方法探索不同的 learning-rate schedules，例如 cosine decay 或 constant learning rate。

第 6 页

Joint learning-rate and weight-decay scaling

到这里，你应该已经对 learning rates 如何作为 scale 的函数表现有了总体感觉，并且我们在 hyperball 实验中看到了一些有趣的迹象：learning rates 和参数的 norm 可能相关（至少在我们的 transformer model 中）。在这个问题中，我们将深入这一现象，并尝试理解 learning rates 和 weight decay 如何相互作用。

Problem 2: How should learning rate and weight decay scale together?（4 个新 runs）

(a) 让 LR 和 WD 一起变化（provided runs）

分析提供的以下 LR–WD grids 的 W&B sweeps。我们的 runs 包含下表中给出的 peak LR 和 WD 的所有组合，

由于我们有两个 hyperparameters，我们将不得不扩展先前的 scaling rules，并考虑 loss 的 bivariate functional form。令 $x = \log \eta$ 和 $y = \log \lambda$。对每个 scale（token budget），使用系数 $a, b, c, d, e, f$ 拟合一个 local quadratic $\hat{L}(x, y)$ 来近似最终 val loss：

$$\hat{L}(x, y) = a + bx + cy + dx^2 + exy + fy^2,$$

交叉项 $exy$ 允许 LR 和 WD 相互作用。

对每个预算，画出其 contours，并标出每个预算处拟合出的 optimum。如果我们做两张图，一张针对 LR、一张针对 WD，分别画出 optimal LR 或 WD 作为 scale 的函数，我们会看到什么？LR 趋势与 Problem 1 中的趋势相比如何？任一坐标能否用单个 power law 很好地描述？

将 joint tuning 得到的最佳 measured validation losses 与 Problem 1 中相同预算下的结果比较。把两者都画成 tokens 的函数，并画出它们的差。调 WD 的收益如何依赖于 training budget？

(b) 从 coupling 到 joint scaling law（无新 runs）

在每个预算处，令 $(\eta^*, \lambda^*)$ 为 (a) 中拟合出的 optimal pair。在其 contour plot 上，画出你的 quadratic fit 和 observed data，两个轴分别为 peak LR 和 WD（使用 log-log scale）。增加一个 hyperparameter 能否补偿另一个 hyperparameter 的减少？在该图上画出 $\eta\lambda = \eta^*\lambda^*$ 的线，并看看它与 contour plot 的关系。

现在画出 $\eta^*\lambda^*$ 随 tokens 的变化，并与单独的 LR 和 WD 趋势比较。product 是否比任一坐标变化得更规律？测量这些不同 scaling fits 各自的 $R^2$，并找出这些变量中哪个具有最紧的 power law scaling。

第 7 页

(c) 测试 joint rule（4 个新 runs）

考虑一个在每个预算处都使用 peak LR .003 的策略（注意，这不是要求 constant LR schedule，我们只是跨 scales 固定相同的 peak LR）。我们的目标是看看能否通过调整 weight decay 来“挽救”这个 LR policy。

1. Fit. 使用在 $\{.1536, .3072, .6144, 1.2288\}$B tokens 上拟合的来自 (b) 的 product law。

2. Test. 在 2.4576B tokens 处验证，该预算被排除在拟合之外。在 peak LR .003 下，测试由预测的 product 除以 .003 得到的 WD。所有 runs 使用 linear decay。

3. Compare. 使用两个 baselines，二者都训练 2.4576B tokens：(i) 复用 153.6M sweep 中 validation loss 最低的 LR–WD pair，不重新调参；(ii) 将 peak LR 固定为 .003，测试 WD $\{.05, .1, .2\}$，并取最佳 validation loss。报告预测相对两个 baselines 的 loss gap。product rule 是否给出了有效的 training recipe？

第 8 页

Example question

Difficulty: Medium

我们能否推迟 weight decay？使用 153.6M tokens 下的默认训练设置，peak LR 为 .0015，WD 为 $\lambda = .4$，初始化 seeds 为 42–44，并固定 data order。我们修改标准训练过程，使 weight decay 不在每一步发生，而是每 $K$ 步施加一个更强的 weight decay。

令 $\eta_t$ 为 update $t$ 处实际使用的 LR（即调用 optimizer.step 之前 optimizer 的 LR，而不是推进其 scheduler 之后的 LR）。在正 warmup 下，第一次 update 使用零 LR，而 schedule 只在最后一次 update 之后才达到零。普通 AdamW 在应用 gradient update 之前，按 $r_t = 1 - \eta_t\lambda$ 缩小 weights。比较三种 schedules：

• 像通常一样，每次 update 都施加 weight decay。
• 仅在每两个 updates 组成的 block 结束时施加。
• 仅在每四个 updates 组成的 block 结束时施加。

对于一个在 update $t$ 结束、长度为 $k$ 的 block，延迟 schedule 每 $k$ 步施加一个更强的 weight decay

$$\tilde{r}_t = \prod_{j=t-k+1}^{t}(1 - \eta_j\lambda).$$

通常的 schedule 在三个 seeds 上给出的 mean final validation loss 约为 3.209。预测每个延迟 schedule 的 mean final validation loss。将每个 mean 四舍五入到最接近的 0.01 的倍数。

• 每两次 updates 施加 weight decay：____。
• 每四次 updates 施加 weight decay：____。

第 9 页

batch size 如何影响超参数？

我们目前已经对双变量超参数缩放有了些体会（在 learning rate 与 weight decay 调参的那个问题里）。我们可能会想，是否存在更好的办法，可以让我们从数学上推导出超参数之间正确的缩放关系。

为此，我们将研究 batch size、learning rate 与 scale 三者如何相互作用。batch size 是一个理想的研究变量，因为 batch size、learning rate 与收敛速率之间的关系已被充分研究并理解，属于 online optimization 理论的一部分。

我们将利用一个著名的玩具模型——noisy quadratic model，看看它如何给出（可能的）可以迁移到真实语言模型训练中的 learning-rate 规则。

背景：noisy quadratic model

noisy quadratic model（NQM）用一个二次型损失加带噪声的梯度来替代神经网络训练。noisy quadratic 有两个我们需要变化并理解的关键性质。二次型的 curvature 控制问题的优化难度，也控制不同方向能被优化的快慢；gradient noise 则代表训练样本之间的差异，并告诉我们更大的 batch 是否有帮助。我们将跟随 Guodong Zhang et al. (2019), Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model，该文用这个简单模型解释 learning rate、momentum 和优化器如何影响 batch-size scaling。

我们将考察两个结论，以及它们成立的条件：

1. 最优 LR 可以随 batch size 呈幂律关系。设 $B$ 为 batch size，$\eta^*(B)$ 为在固定总训练样本数下使期望最终损失最小的 LR。把候选规则写成 $\eta^*(B) \approx cB^p$，其中 $c$ 与 $p$ 为拟合常数。对于小步长区间的 SGD，NQM 预测近似线性缩放，$p = 1$：batch size 翻倍，LR 也应翻倍。当然，总会到达某个点，这个预测会失效，因为任意大的 learning rate 会导致发散。我们将检验这条缩放规则，并找出它有用的参数范围。
2. Momentum 在大 batch 下可能比在小 batch 下更有帮助。经典 NQM 的预测是：在噪声主导的小 batch 区间，SGD 的 momentum 主要只是重新缩放有效 LR，在调好 LR 之后几乎不带来额外收益。在更大的 batch 下，噪声减小可以让 momentum 加速收敛。我们将在有限训练预算下检验这一预测，并与 Adam 做比较。

我们的第一个问题将使用一个简单的二维二次型模型供你上手。

2D 二次型玩具模型

令 $w = (w_1, w_2)^\top$ 为参数，并固定 curvature 矩阵 $H = \mathrm{diag}(1, 10)$。损失及其精确梯度为

$$f(w) = \frac{1}{2}w^\top H w = \frac{1}{2}(w_1^2 + 10 w_2^2), \qquad \nabla f(w) = Hw.$$

为了定义这个二次型的噪声部分，我们考虑梯度是真实梯度的带噪版本。对于 batch size 为 $B$ 的一次更新，我们定义相应的 minibatch 梯度为

$$g_t = Hw_t + \epsilon_t, \qquad \epsilon_t \sim \mathcal{N}(0, I_2/B)$$

在各次更新之间独立，其中 $I_2$ 是二维单位矩阵。这个噪声模型来自一个朴素假设：梯度是 i.i.d. 样本的平均，因此我们预期方差按 $O(1/B)$ 缩放。

第 10 页

优化器

在这个问题中，我们将使用一个简化的优化器模型：常数 LR $\eta$，且没有 weight decay。

至于优化器本身，我们将研究 Stochastic gradient descent (SGD)（定义为 $w_{t+1} = w_t - \eta g_t$）、Adam（$\epsilon = 10^{-8}$，$\beta_1 = 0.9$，$\beta_2 = 0.95$），以及 RMSProp —— 即 $\beta_1 = 0$ 的 Adam。

问题 3.1：这个模型对 batch size 给出了什么预测？（在本机 CPU 上运行）

考虑上面我们的 2D 玩具二次型模型。你将使用带噪声的梯度在这个玩具模型上做随机优化，并比较超参数如何改变最终损失。

具体地，令 $N = 8192$ 为数据集总大小，$B$ 为 batch size（在下面指定），并取 $w_{0,1} \sim \mathcal{N}(0, 1)$、$w_{0,2} \sim \mathcal{N}(0, .1)$ 作为初始化。

比较期望最终损失 $L_N(B, \eta) = \mathbb{E}[f(w_{N/B})]$，其估计方式是对独立的初始化与噪声取平均（至少 512 个样本）。

(a) SGD. 在 batch $B \in \{1, 2, 4, 8, 16, 32, 64\}$ 下扫描 LR，并对估计出的最优 LR 拟合幂律。预测 $B = 256$ 时的最优 LR，然后用该 batch 下的 LR 扫描检验该预测。把 batch 范围扩展到 512，并画出调完 LR 后的最优 LR 与最佳最终损失。这条规则在什么范围内有效？想一想梯度噪声减小与固定 $N$ 下更新次数变少之间如何竞争。

(b) Adaptive updates. 用 RMSProp 和 Adam 重复该实验，分别调 LR。比较 source 指数，以及在 batch 256 处预测 LR 与调优 LR 对应的损失。这些优化器是否遵循同一条规则？相似的指数是否意味着相似的迁移准确率？说出你对语言模型设定的预测。最优 LR 应如何随 batch size 变化，你预期该预测在什么范围内成立？

(c) Explore the assumptions. 把梯度噪声协方差 $I_2$ 换成 $\sigma^2 I_2$，$\sigma \in \{1, 10, 100, 300\}$。重复 RMSProp 和 Adam 的 source 拟合与 batch-256 检验。画出拟合指数随 $\sigma$ 的变化。改变梯度信号与噪声的比例是否会改变缩放规则？把你的观察与二阶矩的各项贡献联系起来。

作为 curvature 的对照，用标量损失 $f(w) = w^2/2$、初始 $w_0 \sim \mathcal{N}(0, 2)$、梯度噪声方差 $\sigma^2/B$ 重复这些实验。哪些结论对 curvature 模型敏感？

(d) Momentum at small and large batches. 回到 $H = \mathrm{diag}(1, 10)$ 与噪声协方差 $I_2$。对带 momentum 的 SGD，初始化 $s_0 = 0$，并使用

$$s_{t+1} = \mu s_t + g_t, \qquad w_{t+1} = w_t - \eta s_{t+1},$$

其中 $\mu$ 是 momentum 系数。在小 batch $B = 16$ 与大 batch $B = 256$ 下，比较 $\mu = 0$ 与 $.9$，分别调 LR。对 Adam，保持 $\beta_2 = .95$，在每个 batch 下联合扫描 $\beta_1$ 与 LR（提示：你应该包含 $\beta_1 = 0$、$0.9$ 以及更接近 1 的其他取值）。画出调完 LR 后每个 $\beta_1$ 对应的最佳损失。momentum 的收益如何依赖于 batch size？把最佳 $(\beta, \mathrm{LR})$ 组合的学习曲线与最佳无 momentum 基线比较。你预期这些观察会迁移到语言模型训练中吗？

第 11 页

语言模型中的 batch size 缩放

NQM 中的缩放规则与 momentum 效应会迁移到语言模型训练吗？使用 614.4M tokens 上的默认模型与训练设置，变化 batch size 与指定的超参数。

这里 $B$ 是每次 optimizer update 的序列数。使用 $\min(B, 64)$ 条序列的 microbatches，并报告实际处理的 token 数。使用线性 learning-rate schedule，在前 1% 的 optimizer updates 上做 warmup。下面的 batch-switch 示例以处理过的 tokens 为单位指定自己的 schedule。

问题 3.2：当 batch size 变化时，应该缩放什么？（建议：48 次新运行）

(a) Tune LR with weight decay (9 new runs). 把 WD 设为 .1，在 $B \in \{8, 16, 32, 64\}$ 下扫描峰值 LR。画出 loss–LR 曲线、拟合出的最优 LR，以及最佳实测损失随 batch size 的变化。最优 LR 与性能如何随 batch size 改变？

(b) Compare two scaling hypotheses (suggested: 27 new runs). NQM 也可以给出 weight decay 与 batch 如何联合缩放的预测。在 $B \in \{8, 16, 32, 64\}$ 下设计实验，检验两个假设：

i. 跨 batch 固定 WD 为 .1，缩放最优峰值 LR。

ii. 跨 batch 固定峰值 LR 为 .0015，缩放最优 WD。

用 source 结果预测 $B = 128$ 与 $B = 256$ 的配置，然后检验哪个假设迁移得更好。讨论你的发现是否与 NQM 一致。

(c) Ablate momentum (suggested: 12 new runs). 在 $B = 8$ 与 $B = 256$ 下，固定每个 batch 实测最佳的 LR–WD 组合与 $\beta_2 = .95$。设计包含 $\beta_1 = 0$ 对照的实验，以检验 momentum 的效果如何随 batch size 变化。与我们之前重新调 LR 的那个问题做比较。

(d) What does the NQM explain? (no new runs). 你二次型实验中的哪些预测与你的语言模型实验一致，哪些不一致？

给出一个简短的预测–实测表，涵盖 LR–batch-size 缩放、weight decay 和 momentum。附上支持你结论的实测目标损失。区分「失败的预测」与「NQM 未建模的效应」。你会如何修改 NQM 来解释其中一处不一致？

延伸阅读

Power Lines: WD–LR–batch coupling；Step Law: optimal LR and batch-size scaling；Hyperball: control of weight and update norms。

第 12 页

示例问题

难度： Hard

在训练期间将 batch size 加倍

用 AdamW 训练 d8 language model 614.4M tokens，初始 batch size 为 64，peak LR 为 $0.0015$，WD 为 $1.6$。令 $\eta_{\text{base}}(D)$ 表示处理 $D$ 个 tokens 后的 learning rate：在前 $1\%$ 的 tokens 上 warm up，然后在 614.4M tokens 时线性衰减到零。

在处理恰好 307.2M tokens 后，将 batch size 加倍到 128。用 4,687 个包含 64 个 sequences 的 batches，随后再用一个包含 32 个 sequences 的 batch 到达中点；这在切换前给出 4,688 次 updates。然后使用 2,343 个包含 128 个 sequences 的 batches 和一个包含 96 个 sequences 的最终 batch，在切换后给出 2,344 次 updates。用每个 update 实际的 batch size 计算其 mean loss。保持 Adam 的累积 moments 和两个 $\beta$ 值不变。对于训练的后半段，考虑五种 policies：

所有 runs 都处理相同顺序的 tokens。在每个 update 之前所消耗的 tokens 数量处评估 $\eta_{\text{base}}(D)$，包括那些 partial batches。改变 batch size 不得重新开始 warmup。

你预测哪个 policy 会在 3 个 seeds（42,43,44）上给出最低的 mean final validation loss？选择一个，并解释当切换后每个 token 的 optimizer updates 大约只有一半时，改变 LR 或 WD 会如何影响训练。

第 13 页

Part 4: Initialization, feature learning, and transfer

到目前为止，我们依赖于一种基于 scaling-law 的方法来在大规模下寻找好的超参数。然而，即便是这种理论指导的方法，也可能显得不令人满意。为什么最优超参数非得发生变化？如果我们把架构和优化器设计正确，也许只需调一次超参数，然后在所有地方都使用同一个最优超参数。这正是各种 hyperparameter transfer 方案所承诺的，而我们将深入讨论其中一种流行的迁移方案，称为 µP。

目标和 abc 参数化

我们想在小模型上调超参数，并在扩大 width 和 depth 时复用它们。事实证明，我们可以通过调整初始化、输出 scale 和学习率参数化方式中的一些小细节，来修补现有架构和优化器，使其具有合理的 hyperparameter transfer。

沿用 Everett et al. 的做法，考虑一个线性层，其 fan-in 为 $n$，有 $p$ 个输出，输入特征为 $x\in\mathbb{R}^n$。我们用三个指数 $(a,b,c)$ 调整该层的标准形式，它们分别修改 forward multiplier、initialization scale 和 learning rate。令 $\sigma$ 与 base LR $\eta$ 为与 scale 无关的常数（即我们的、随 scale 保持稳定的超参数），

$$ z=n^{-a}Wx,\qquad W\in\mathbb{R}^{p\times n},\qquad (W_0)_{ij}\overset{\text{iid}}{\sim}\mathcal{N}(0,\sigma^2 n^{-2b}),\qquad \eta_W=\eta n^{-c}. $$

两个原则：稳定性与特征学习

我们用 RMS norm 比较逐坐标的 scale：

$$ \|v\|_{\mathrm{RMS}}^2=\frac{1}{n}\sum_{j=1}^{n}v_j^2,\qquad \|A\|_{\mathrm{RMS}}^2=\frac{1}{pn}\sum_{i,j}A_{ij}^2 \quad (A\in\mathbb{R}^{p\times n}). $$

对于 $\mathbb{R}^p$ 中的输出向量，把向量 RMS 定义中的 $n$ 换成 $p$。令 width $n\to\infty$，训练步 $t$ 固定。对于满足 $\|x\|_{\mathrm{RMS}}=\Theta(1)$ 的固定 probe 输入 $x$，令 $z_t=n^{-a}W_t x$。这里 $O(1)$ 表示随 width 有界，$\Theta(1)$ 表示有界且不消失。稳定性要求初始 hidden features 可用，且变化有界：

$$ \|z_0\|_{\mathrm{RMS}}=\Theta(1),\qquad \|z_t-z_0\|_{\mathrm{RMS}}=O(1). $$

初始 readout logits 则可能消失，因此其初始 scale 只需为 $O(1)$。上面 $\Theta(1)$ 的初始化要求适用于 hidden features。特征学习要求在某固定步有非消失的变化：

$$ \|z_t-z_0\|_{\mathrm{RMS}} =\|n^{-a}(W_t-W_0)x\|_{\mathrm{RMS}} =\Theta(1). $$

如果该变化按 $n^{-r}$ 缩放，那么 $r<0$ 不稳定，$r=0$ 给出特征学习，$r>0$ 给出消失的变化。

从稳定性推导 Kaiming initialization

我们先用初始化时的稳定性推导 Kaiming initialization。由于 $(W_0)_{ij}\overset{\text{iid}}{\sim}\mathcal{N}(0,\sigma^2 n^{-2b})$，且 $W_0$ 与输入 $x$ 独立，交叉项消失：

$$ \mathbb{E}[z_{0,i}^2\mid x] =n^{-2a}\sum_{j=1}^{n}\mathbb{E}[(W_0)_{ij}^2]\,x_j^2 =\sigma^2 n^{-2(a+b)}\|x\|_2^2. $$

第 14 页

为了稳定性，我们希望所有 hidden features（包括 $x$ 和 $z$）在初始化时都具有 $\Theta(1)$ 的 RMS norm。这促使我们选择 $a+b=1/2$：

$$ a+b=\frac12 \quad\Longrightarrow\quad \mathbb{E}\big[\|z_0\|_{\mathrm{RMS}}^2\mid x\big] =\sigma^2\|x\|_{\mathrm{RMS}}^2. $$

用于 ReLU 的 Kaiming initialization 选择 $a=0$、$b=1/2$ 和 $\sigma^2=2$。常数因子 $2$ 的推导见 He et al. (2015), Section 2.2。

Kaiming initialization 没有处理训练过程中的特征学习，这促使了 µP 理论。为了理解它，我们引入用于描述权重更新与输入特征之间 alignment 的记号。

权重更新与输入特征之间的 alignment

令 $\Delta W := W_t-W_0$。我们关心的训练对是 $(\Delta W,x)$。按定义，一次权重更新给出 $\Delta z=n^{-a}\Delta W x$。对于非零的 $\Delta W$ 和 $x$，定义

$$ R(\Delta W,x) =\frac{\|\Delta W x\|_{\mathrm{RMS}}}{\|\Delta W\|_{\mathrm{RMS}}\|x\|_{\mathrm{RMS}}}, \qquad \|\Delta z\|_{\mathrm{RMS}} =n^{-a}R(\Delta W,x)\|\Delta W\|_{\mathrm{RMS}}\|x\|_{\mathrm{RMS}}. $$

比值 $R$ 表示参数更新与输入之间的 alignment。注意，alignment 如何让我们用其组成部分（参数变化与输入变化）来简化特征的变化。现在让我们通过两个特例来获得关于特征更新数量级 scale 的一些直觉。假设 $\Delta W\in\mathbb{R}^{p\times n}$，其中 $n$ 是 fan-in。

无 alignment： $\Delta W$ 具有 iid 中心化 Gaussian 条目，且与 $x$ 独立，因此 $R(\Delta W,x)\approx\sqrt{n}$。假设更新条目 $\Delta W_{ij}$ 是相互独立的 $\mathcal{N}(0,\tau^2)$ 变量，且与 $x$ 独立。对于 $j\ne r$，有 $\mathbb{E}[\Delta W_{ij}\Delta W_{ir}]=0$，因此

$$ \mathbb{E}[(\Delta W x)_i^2\mid x] =\sum_{j,r}x_j x_r\mathbb{E}[\Delta W_{ij}\Delta W_{ir}] =\tau^2\sum_j x_j^2 =n\tau^2\|x\|_{\mathrm{RMS}}^2. $$

于是 $\Delta W x$ 的 RMS scale 为 $\sqrt{n}\,\tau\|x\|_{\mathrm{RMS}}$，而 $\Delta W$ 的 RMS scale 为 $\tau$。Gaussian 旋转对称性给出该比值精确的 RMS scale：

$$ \big(\mathbb{E}[R(\Delta W,x)^2\mid x]\big)^{1/2}=\sqrt{n}. $$

该等式在随机更新上的 RMS 意义下成立；单个更新不必恰好具有这个比值。

Full alignment： $\Delta W$ 是 rank one，且 $x$ 沿其右奇异向量方向，因此 $R(\Delta W,x)=n$。由输入构造出的梯度具有这种形式。对于 $t=1$，从 $W_0$ 出发，在单样本损失 $\ell(z)$ 上走一步 SGD，其中 $z=n^{-a}W_0 x$，$\delta=\partial\ell/\partial z$ 是损失关于输出的梯度。学习率为 $\eta_W$ 时，更新为

$$ \Delta W=-\eta_W n^{-a}\delta x^\top,\qquad \Delta W x=-\eta_W n^{-a}\delta\|x\|_2^2. $$

对于非零更新，$\Delta W$ 是 rank one，且 $x/\|x\|_2$ 是其具有非零奇异值的右奇异向量。每一行都与 $x$ 成比例，因此

$$ \|\Delta W\|_{\mathrm{RMS}} =\eta_W n^{-a}\|\delta\|_{\mathrm{RMS}}\|x\|_{\mathrm{RMS}}, \qquad R(\Delta W,x) =\frac{\|x\|_2^2}{\|x\|_{\mathrm{RMS}}^2}=n. $$

这达到了 Cauchy–Schwarz 上界 $R(\Delta W,x)\le n$。我们看到 $\sqrt{n}$ 和 $n$ 都是合理的 scaling。稍后我们会看到二者中哪一个在经验上拟合得更好。

第 15 页

readout head 与其输入特征变化之间的 alignment

令 $V_0\in\mathbb{R}^{n\times q}$ 为转置约定下的初始 readout 矩阵，其中类别数 $q$ 不随 $n$ 缩放。对于最后一层的特征变化 $\Delta x\in\mathbb{R}^n$，初始 readout 贡献

$$ \Delta z_{\text{feature}}=n^{-a}V_0^\top \Delta x. $$

对于非零的 $V_0$ 和 $\Delta x$，定义

$$ S(V_0,\Delta x) =\frac{\|V_0^\top\Delta x\|_{\mathrm{RMS}}} {\|V_0\|_{\mathrm{RMS}}\|\Delta x\|_{\mathrm{RMS}}}, \qquad \|\Delta z_{\text{feature}}\|_{\mathrm{RMS}} =n^{-a}S(V_0,\Delta x)\|V_0\|_{\mathrm{RMS}}\|\Delta x\|_{\mathrm{RMS}}. $$

有两个参考情形： - 无 alignment： iid 中心化 Gaussian $V_0$ 与 $\Delta x$ 独立，给出

$$ \big(\mathbb{E}[S(V_0,\Delta x)^2\mid \Delta x]\big)^{1/2}=\sqrt{n}. $$

• Full alignment： $V_0$ 的各列与 $\Delta x$ 成比例，给出 $S(V_0,\Delta x)=n$。

按层类型推导 µP

在此推导中，我们假设两个 full-alignment 条件同时成立： 1. 权重更新与输入特征之间 full alignment。 2. 初始 readout head 与输入特征变化之间 full alignment。 Problem 4.0 要求你推导这些 alignment 假设的所有四种组合下的 scaling rules。Hidden 和 readout 共享权重更新约束。 Maximal-update parameterization（µP）保持更新不消失、信号有界。将 Adam 更新写成 $\Delta W=-\eta n^{-c}U$，其中

$$ U=\frac{\hat m}{\sqrt{\hat v}+\epsilon} $$

是归一化方向。对于 Adam，我们可以假设 $\|U\|_{\mathrm{RMS}}=\Theta(1)$。[注 1] 定义 $\alpha$ 使得 $R(U,x)=\Theta(n^\alpha)$。对于 order-one 输入，

$$ \|n^{-a}\Delta W x\|_{\mathrm{RMS}} =\eta n^{-a-c}R(U,x)\|U\|_{\mathrm{RMS}}\|x\|_{\mathrm{RMS}} =\Theta(\eta n^{\alpha-a-c}). $$

为了特征学习，我们希望当 $n\to\infty$ 时，$\|n^{-a}\Delta W x\|_{\mathrm{RMS}}$ 保持在常数、非零的 scale。在 full alignment（$\alpha=1$）下，这给出我们的第一个方程：

$$ a+c=1. $$

特征变化约束取决于该层的形状。从 $(W_0,x_0)$ 出发经过一次更新后，变化 $\Delta W$ 和 $\Delta x$ 给出

$$ \Delta z=n^{-a}\Big( \underbrace{\Delta W x_0}_{\text{weight update}} +\underbrace{W_0\Delta x}_{\text{feature change}} +\underbrace{\Delta W\Delta x}_{\text{interaction}} \Big). $$

我们已经证明，直接的 weight-update 项满足 $\|n^{-a}\Delta W x_0\|_{\mathrm{RMS}}=\Theta(1)$。对于 $\|\Delta x\|_{\mathrm{RMS}}=\Theta(1)$，interaction 项满足

$$ \|n^{-a}\Delta W\Delta x\|_{\mathrm{RMS}} \le n^{1-a}\|\Delta W\|_{\mathrm{RMS}}\|\Delta x\|_{\mathrm{RMS}} =O(\eta n^{1-a-c})=O(\eta). $$

剩下的就是控制 $W_0\Delta x$。假设 $W_0$ 有 $p$ 个输出通道。为了表达这一项对 $n$ 的依赖，我们使用 Gaussian matrix norm bound。令

$$ \|A\|_{\mathrm{op}}=\sup_{u\ne0}\frac{\|Au\|_2}{\|u\|_2} $$

为 operator norm。Gaussian matrix norm bound 给出 $\|W_0\|_{\mathrm{op}} =$

[注 1] Adam 可以看作带 momentum 的 signSGD，使用 $\operatorname{sign}(\hat m)$，并结合 variance adaptation；见 Balles and Hennig (2018), Dissecting Adam。

第 16 页

以高概率为 $O(n^{-b}(\sqrt{p}+\sqrt{n}))$。

转到 RMS 范数后，每个 RMS 为 $\Theta(1)$ 的 $\Delta x$ 都被约束住，包括依赖于 $W_0$ 的变化：

$$\|n^{-a}W_0\Delta x\|_{\mathrm{RMS}} \le n^{-a}\sqrt{\frac{n}{p}}\|W_0\|_{\mathrm{op}}\|\Delta x\|_{\mathrm{RMS}} = O\left(n^{-a-b}\left(\sqrt{n}+\frac{n}{\sqrt{p}}\right)\right).$$

1. Hidden 矩阵：两个维度一起缩放，$p\propto n$。 当 $p\propto n$ 时，特征变化上界为 $O(n^{1/2-a-b})$。初始化要求 $a+b=1/2$ 已经使它为 $O(1)$，即使 $\Delta x$ 依赖于 $W_0$ 也是如此。不需要更强的初始化缩放。

2. Readout：readout 与 feature-change 完全对齐给出 $a+b=1$。 对 readout 而言，$p=q$ 固定。 用转置约定把初始 readout 写成 $V_0=W_0^\top\in\mathbb{R}^{n\times q}$。与学习到的特征变化完全对齐意味着 $S(V_0,\Delta x)=\Theta(n)$。在 $\|V_0\|_{\mathrm{RMS}}=\Theta(n^{-b})$ 且 $\|\Delta x\|_{\mathrm{RMS}}=\Theta(1)$ 时，得到

$$\|n^{-a}V_0^\top\Delta x\|_{\mathrm{RMS}} = n^{-a}S(V_0,\Delta x)\|V_0\|_{\mathrm{RMS}}\|\Delta x\|_{\mathrm{RMS}} = \Theta(n^{1-a-b}).$$

稳定性要求 $a+b\ge 1$。要求对特征变化有非零的、量级为 1 的响应，就得到 $a+b=1$。初始化时，readout 与其输入特征相互独立，因此初始 logits 的 RMS 尺度为 $\Theta(n^{1/2-a-b})=\Theta(n^{-1/2})$。

1. Input 矩阵：fan-in 保持固定。 在 fan-in $d$ 固定且外部输入不变时，初始化和直接更新的尺度分别为 $n^{-a-b}$ 与 $\eta n^{-a-c}$，这里假设 update 对齐非零。让二者都保持量级 1，得到 $a+b=a+c=0$。

我们的推导表明，不同层需要不同的 $(a,b,c)$ 选择，才能得到稳定且非零的特征更新。下面总结其中一种选择。这里 $m=n/n_0$，$n_0$ 是参考宽度，$k$ 是每个矩阵实际的 fan-in；这些缩放在 $n_0$ 处与基线一致。

第 17 页

Problem 4.0：推导 update 与 feature-change 对齐的缩放规则

考虑两个假设的所有四种组合，每个假设取“完全对齐”或“不对齐”：

• 权重更新方向 $U$ 与更新前输入 $x_0$ 对齐。
• 初始 readout head $V_0\in\mathbb{R}^{n\times q}$ 与其输入特征变化 $\Delta x=x_1-x_0$ 对齐。

这里 $V_0$ 采用转置约定，因此它对特征变化的响应为 $n^{-a}V_0^\top\Delta x$。记 $\Delta W=-\eta n^{-c}U$，其中 $\eta>0$ 固定，且 $\|U\|_{\mathrm{RMS}}=\Theta(1)$。对隐藏矩阵和 readout 矩阵，假设 $x_0$ 与 $\Delta x$ 的 RMS 为量级 1，并记

$$R(U,x_0)=\Theta(n^\alpha),\qquad S(V_0,\Delta x)=\Theta(n^{\omega_{\text{move}}}).$$

对于交互项，再假设 $R(U,\Delta x)=O(n^\alpha)$。对隐藏矩阵和 input 矩阵取 $a=0$，对 readout 取 $c=0$。

(a) 对每种组合，确定 $\alpha$ 与 $\omega_{\text{move}}$，并推导对 $(a,b,c)$ 的约束。要求初始化时隐藏特征为量级 1、直接更新为量级 1，以及初始 readout 给出量级为 1 的响应 $n^{-a}V_0^\top\Delta x$。用一张表总结隐藏、readout 和固定 fan-in 的 input 矩阵的 $(a,b,c)$。

(b) 使用 $m=n/n_0$，给出每种层类型的前向乘子、初始化方差和 LR，并在 $n_0$ 处与前面的表一致。各对齐假设会怎样改变这些缩放规则？检查交互项保持有界，并指出哪一种组合恢复了前文的 µP 规则。

第 18 页

4.1: Toy Example：五步 Transformer 压力测试

测试 µP 缩放如何在极端宽度和深度之间迁移超参数。

训练协议

从第一次更新起就用恒定 LR 训练五次更新，不使用 warmup、weight decay 和 gradient clipping。使用未缩放的 rotary position embeddings（RoPE）。在不同 parameterization 之间使用同样的五个训练 batch 和 64 条固定验证序列；按 microbatches 评估，并用第五次更新后的 validation loss 调整 LR。

宽度设置

令 $m=n/n_0$，$k$ 为每个矩阵的 fan-in。用下面给出的方差，以独立零均值高斯条目初始化权重矩阵。

深度设置

把宽度和参考宽度都固定为 64，使用一个 attention head，参考深度为 2。参数、梯度、Adam 状态和残差加法保持为 FP32；分支计算使用 BF16 autocast。令 $L$ 表示 Transformer 层数，$r=L/2$，并用 $c_L$ 在残差加法前乘到每个 attention 和 feedforward 分支上。把这些深度乘子叠加在 µP 宽度方案之上：

Block 乘子作用于隐藏矩阵和 block 内归一化。Embedding、readout 和最终归一化的 LR 与 stabilizer 保持不变。

第 19 页

Problem 4.1：在五次更新中测试宽度和深度缩放（约 55 次运行）

测量特征移动、update 对齐以及 readout–feature-change 对齐。

令 $h_t$ 是宽度为 $n$ 的隐藏特征向量，在一组 $N$ 个固定 probe $x_1,\ldots,x_N$ 上求值。对 fan-in 为 $k$ 的矩阵 $A_t$，令 $\Delta A_t=A_t-A_{t-1}$ 为其实际更新，$u_{t-1}$ 为其更新前输入。用下式测量相对初始化的特征移动和 update 对齐：

$$M_t=\left[\frac{1}{Nn}\sum_{i=1}^{N}\|h_t(x_i)-h_0(x_i)\|_2^2\right]^{1/2},\qquad \alpha_{\text{upd}}^{(t)}=\log_k R(\Delta A_t,u_{t-1}).$$

把特征 RMS 在序列、token 位置和坐标上汇总；权重 RMS 在条目上计算。分别检查归一化后的特征和 residual stream。

对 readout probe，取 $h_t$ 为 readout 乘子之前的最终 RMSNorm 输出，并令 $V_0\in\mathbb{R}^{n\times q}$ 为转置约定下的初始 readout 权重。测量

$$\omega_{\text{move}}^{(t)}=\log_n S(V_0,h_t-h_0).$$

这探测推导中的 initial-readout/feature-change 项。使用与上面相同的固定 probes，并以相同方式汇总 RMS；当特征移动为零时省略该比值。

(a) 跨宽度迁移（约 30 次五步运行）。预测 Kaiming 和标准 µP 在宽度 640、2,560 和 5,120 之间迁移其最佳 base LR 的效果。为每个宽度和 parameterization 选择一组 LR sweep。用合适的模型拟合每条 loss–LR 曲线。画出拟合得到的最佳 base LR 和拟合得到的最小 validation loss 随宽度变化的图。哪一种 parameterization 能更好地迁移其调好的 base LR？它是否也达到更低的 loss？

(b) 探测宽度缩放（复用 (a) 的配置）。在每个宽度和 parameterization 的最佳采样 LR 下，追踪固定 probes 上 logit RMS 以及最终 RMSNorm 输出的 RMS 和移动 $M_t$，贯穿五次更新。画出每个 block 的 query、key、value、attention-output 以及 feedforward gate/up/down 矩阵和最终 readout 的 update 对齐 $\alpha_{\text{upd}}^{(t)}$。使用每个矩阵的更新前输入及其实际 fan-in $k$。再画出 $\omega_{\text{move}}^{(t)}$，并与 no-alignment（$1/2$）和 full-alignment（$1$）尺度比较。比较不同层和不同宽度下的特征移动与对齐。前文推导中的哪些假设和预测与你的测量一致？

(c) 跨深度迁移（约 25 次五步运行）。预测哪一种深度方案（depth prescription）应能最可靠地迁移其 base LR。比较标准 µP、Depth-µP 和 CompleteP 在深度 2、100 和 1,000 下的表现，为每个深度和 prescription 选择一组 LR sweep。画出拟合得到的最佳 base LR 和拟合得到的最小 validation loss 随深度变化的图。更好的 LR 迁移是否意味着五次更新后的 loss 更低？

(d) 探测深度缩放（复用 (c) 的配置）。在每个深度和 prescription 的最佳采样 LR 下，追踪归一化前的 residual-stream RMS，以及乘以 $c_L$ 之前的分支输出 RMS。在相匹配的相对位置（例如中点层和最终层）测量归一化后 residual-stream 特征的移动。还要在全部五次更新中追踪 $\omega_{\text{move}}^{(t)}$，并在不同深度和 prescription 之间比较。相似的归一化激活是否会掩盖一个不断增大的 residual stream？把这些 probe 与 (c) 中的 LR 迁移结果联系起来，区分有界信号与非零特征变化。

第 20 页

4.2：缩放规则能预测更长训练吗？

Problem 4.1 中的 scaling ideas 在更长训练中是否有帮助？使用默认模型和训练设置，训练 153.6M tokens。

课程基线使用默认 architecture 和 initialization。对于 µP，只改动下表中的四个设置，使用 reference width $n_0 = 512$。令 $n$ 为 hidden width，$m = n/n_0$，$\eta$ 为 peak base LR。令 $G$ 为一个共享矩阵，其元素为标准 Gaussian 抽样并在 $\pm 3$ 处截断，用于 embedding 和 readout 两者的初始化，并在各 prescription 之间共享。

保持所有其他训练设置固定，包括 RMSNorm $\epsilon = 10^{-5}$ 和 gradient clipping 为 1。Normalization 和 clipping 可能对不同 widths 产生不同影响：baseline embedding scale $1/n$ 可能与 $\sqrt{\epsilon}$ 相当。使用 feature logs 和 gradients.jsonl 报告 pre-clipping gradient norms、embedding 占平方梯度范数的比例，以及 clipping coefficient。在解释 transfer 时考虑这些效应。复用提供的 width-512 diagnostics；可选的 normalization 或 clipping 消融必须在你 compute budget 之内完成。

Problem 4.2：从五次更新到更长训练预算（建议：50 次新 runs）

(a) 迁移 reference LR（20 次初始 runs）

预测 Problem 1 在 153.6M tokens 下的 best sampled LR 是否能在两种 prescription 下迁移到 widths 128 和 256。为每个 configuration 选择一个 LR sweep，其中包含 transferred source LR。在新 runs 之间保持 data 和其他设置固定；与 Problem 1 提供的 width-512 curve 比较。

在 source LR 处，对两种 prescription 和所有三个 widths 绘制训练全程的 $\alpha_{\text{upd}}(t)$。对每个 block 的 attention 和 feedforward 矩阵以及 readout，使用固定输入和每个矩阵的 fan-in。也绘制来自 Problem 4.1 的 $\omega_{\text{move}}(t)$。在最初五次更新之后，update alignment 和 readout–feature-change alignment 如何演化？

(b) 比较 transfer 和 tuned performance（不跑新 runs）

拟合每条 loss–LR 曲线。对两种 prescription，绘制 fitted optimal base LR 和 fitted minimum validation loss 随 width 的变化。也绘制在 transferred LR 处测得的 loss。LR transfer 更好的 prescription 在 tuning 后是否也能达到更低的 loss？这一比较与 five-step test 有何不同？

(c) 预测一个 held-out width（约 12 次新 runs）

留出 width 1,024。对每种 prescription，用 part (b) 中 widths 128、256 和 512 的 optimal base LRs 拟合 power law，并预测其在 1,024 处的 optimal LR。在预测 LR 和直接迁移的 source LR 处训练。然后在 1,024 处运行 local LR sweep。将两个预测 LR 与拟合出的 local optimum 比较，并报告它们相对于 best sampled target LR 的 measured loss gaps。拟合 width dependence 是否能改进 direct transfer？

(d) 跨 depth 迁移（18 次初始 runs）

固定 width 512，比较 depths 4、8 和 16。使用 Problem 4.1 的规则测试 standard µP、Depth-µP 和 CompleteP，现在 reference depth 为 8：对 layer count $L$，使用 relative depth $r = L/8$。在 depths 4 和 16 处，比较

第 21 页

迁移后的 source LR 与你选择的 LR sweep。绘制 fitted optimal base LR 和 fitted minimum validation loss 随 depth 的变化，并附上在 transferred LR 处测得的 losses。depth correction 改善的是 transfer、tuned performance，还是两者都改善？

(e) 识别 training regime（diagnostics 来自同一批 runs）

对每个 configuration，在 directly transferred LR 和 best sampled LR 处，按 Problem 4.1 的做法，在最初五次更新期间和训练结束时，追踪固定输入上的 logit RMS 和 normalized hidden-feature movement。同时包括 held-out width 和 depth 比较。在相同 checkpoints 处也包括 $\omega_{\text{move}}(t)$。在不同 widths、depths 和训练阶段，哪些 alignment assumptions 得到支持？一个具有更大 early transients 的 configuration 是否仍能取得更好的 final loss？仅凭 optimal LR 的偏移，是否就能表明模型处在 stable µP regime 之外？用你的 diagnostics 解释这些实验确立了什么。使用 Problem 2 的 LR–WD coupling 讨论固定 WD 可能如何影响该比较。

(f) Muon 与 µP（可选）

Muon 是一种基于 momentum 的 optimizer，它在应用更新前，对每个 hidden weight matrix 的 update direction 近似做 orthogonalization。按照 parts (a)–(c) 中相同的方法，探索 Muon 如何与 µP 相互作用。本探索不属于 required workload；只有在完成 required experiments 后仍有 GPU budget 时才尝试。

第 22 页

示例题

难度：Easy

通过复制来加宽 Kaiming 模型

在匹配的 training batches 上，使用 Problem 4.1 的 five-step width-stress-test 设置。假设 Adam 的 numerical stabilizer 在这些更新中可忽略。从作业中的 width-640 Kaiming baseline 开始，包括其 embedding、normalization 和 readout 初始化。其 LR 对每个 parameter group 都是 .001。不要独立初始化一个更宽的模型，而是把这个模型扩大 $k$ 倍，其中 $k = 4$ 或 $8$。

将 embedding coordinates 和 width-dependent normalization gains 复制 $k$ 次。复制完整的 attention heads，保持每个 head 的维度为 64；保持共享的 64-dimensional QK-normalization gains 不变。对每个 attention 或 feedforward 矩阵 $W$（包括 rectangular gate、up 和 down projections），将其替换为 $k\times k$ block matrix，其中每个 block 都是 $W/k$。若 $V$ 是 readout 矩阵，则沿其 input dimension 拼接 $k$ 个副本并除以 $k$：

$$W' = \frac{1}{k}\begin{bmatrix} W & \cdots & W \\ \vdots & \ddots & \vdots \\ W & \cdots & W \end{bmatrix}, \qquad V' = \frac{1}{k}[V \ \cdots \ V].$$

没有额外的 readout multiplier。这些副本是分开的 trainable parameters，不是 tied weights。

哪些 learning rates 应能在五次更新中的每一次之后，最接近地保持 width-640 模型在相同输入上的 logits？每个 width 选择一个 option。下表中的条目是实际的 parameter-group LRs，没有额外的 width multipliers。“Other” 指 embedding 和所有 normalization parameters。

(a) Width 2560

(b) Width 5120

第 23 页

可选探索

为了更广泛地探索本作业中的主题，请构造你自己的 prediction problem。这个可选练习可以按 Instructions 中所述在 quiz 中获得 bonus points。

Problem 5：你自己的 prediction problem（可选）

你可能已经发现了 schedules、LR 与 WD、batch size 或 model size 之间一个意外的 interaction。用这样一个观察来设计一个问题：同学可以基于本作业进行推理，但其答案超出了所给示例。

在你的探索中包含以下内容：

(a) 按示例风格写一个 question writeup

提供 answer options，并说明是选择一个还是选择所有适用项。在测试 target intervention 之前，记录你的 prediction 和 reasoning。使该 outcome 能通过一个指定 model、data、training budget 和 evaluation metric 的实验进行客观检查。

(b) 紧凑的 code diff（少于 100 行）、可执行代码，以及决定答案的 measured evidence

包含 answer key，并解释与你的 prediction 的任何差异。让 diff 聚焦于该 intervention，并复用 shared training interface。如果 runs 之间的 variation 可能改变正确选项，就检查它，并相应修改问题或 option ranges。

目标是提出一个公平、新颖且带有合理替代项的问题。从你已经理解的实验开始，并使用最小的 informative comparison。Course staff 可能会在未来的 quiz 中使用有趣的问题。本探索是在四个 part 的 starting grids 之外的额外内容。

第 24 页

Assignment 2: Hyperparameter Scaling

Example question answer key

请先自己尝试这些示例题，然后再读下面的答案。

推迟 weight decay（第 8 页）

平均 loss：延迟两次更新与延迟四次更新都是 3.21。让累计的 weight retention 相匹配，并不能保持训练轨迹不变，但这两种延迟在每次更新时施加 decay 所得到的平均 loss 都相近。

通过复制来加宽 Kaiming 模型（第 22 页）

(a) 选项 3；(b) 选项 2。对于扩增因子 $k$，hidden matrices 与 readout 使用 LR $0.001/k$，embeddings 与 normalization 参数使用 $0.001$。当 Adam 的 epsilon 可以忽略时，把 gradient history 乘上一个正常数，归一化后的 update 保持不变。hidden 与 readout 的权重必须保持为原权重的 $1/k$，因此它们的 LR 要按 $1/k$ 缩放；embeddings 与 normalization 的 gain 则是不缩放的复制。有限精度与非零的 epsilon 会带来少量数值差异。

训练到一半时把 batch size 加倍（第 12 页）

Policy A：LR 与 WD 都不改。它的平均 loss 是 2.93867，而次优的 policy B 是 2.94538。为「每个 token 施加 weight decay 的次数变少」做补偿，并不一定能改善最终 loss。

材料

• 题面 PDF
• 代码模板（GitHub a2_release）

自查

• 交付物 = 提供的实验的分析 + 你自己新跑的实验 + 每题的短解读（不是只贴 loss 数字）。
• 官方限 20 页（不含 bonus）—— 页面限制本身是“取舍”的训练：哪些图必须留、哪些放附录。

这份作业和书的关系

配第 3 讲《优化器（一）：不变量与超参数》、第 4 讲《参数化与预条件》。 书里讲的“学习率随宽度怎么变”“batch size 与 LR 的联合缩放”，这里要你自己扫出来。