A3 · Optimization 2:损失景观与尖锐度
| 年份 | 2026(Fall 2026)—— 课程进行中,后续版本可能调整 |
| 分值 | 不计分(Ungraded) |
| 截止 | 10/20(周二)13:00 · 报告 ≤ 20 页 |
| 算力 | 与 A2 同一套基线(35M 级小模型,单张 H100 一次实验十几分钟);要跑多组对照 |
| 配哪几讲 | 候选(机器按术语匹配,供参考):第5讲 §八、河谷模型:理论学习率与不可消除的稳态风险;第5讲 §六、Hessian:为曲率提供一套语言;第6讲 §开场:课程事务与本讲路线;第6讲 §完整训练曲线与 Edge of Stability 的提出 |
| 官方页面 | https://deep-learning-alchemy.github.io/#schedule |
做什么
题面来自官方材料(本地原件在
_local/,官方链接见「材料」)。 编号是原件的编号(notebook 是 cell 序号、PDF 是页码),方便跟官方题面对号。 本项目不提供答案(D53)。
第 1 页
作业 3:Optimization 2
CS 312 · Due date(截止日期): Oct 20 by 1 PM (Quiz time) · 不计分(Ungraded)
基本设置。 除非某道题另有说明,实验都使用课程的 language-modeling 设置:在提供的 DCLM 风格数据上训练模型做 next-token prediction,用 validation loss 评估,并在相同的数据、token budget 和评估协议下比较各次运行。Validation loss 越低越好。
说明。 提交一份简短的实验报告,页数最多 20 页。对每道题,说明你做了哪项比较、扫描了哪些超参数、你找到的最佳 validation loss 和诊断信息,以及你从结果中得出的解释。充分调整 learning rate 和其他核心超参数,使结论不是糟糕运行造成的假象。当题目要求画图时,请包含清晰的坐标轴、标签,以及足够上下文,使趋势可解释。
代码。 这些题目的 starter launchers 位于 experiments/a3_optimization。每个 launcher 会列出该题的 baseline runs;TODO 标出了你需要自己构建的部分。
第 2 页
背景
Optimization 是寻找某个 loss $\ell$ 的 $\arg\min_\theta \ell(\theta)$ 的过程。到目前为止,我们对 $\ell(\theta)$ 的结构和本质了解得很少,只是把它当作一个难以理解的黑盒。但结果表明,深入理解 $\ell$ 对于更好地建立对 neural network training dynamics 的直觉相当重要。Deep networks 为什么能可靠地训练?Training curves 为什么长成那样?我们为什么需要 learning rate schedules?如果不对 loss landscape 有更深的理解,就很难回答这些问题;本作业将尝试建立这种理解。
Loss landscape $\ell(\theta)$ 高度复杂且具有随机性。为了聚焦研究,我们要思考两个主要问题:
- 有多少个 minima,它们长什么样?
- 在走向某个 minimum 的过程中,loss 的形状是什么?
这些问题分别对应 minima 和 sharpness 的研究。
首先,我们会采取更宏观的视角,观察 loss 对简单干预(例如对模型做 averaging,或改变 learning rate schedules)会如何反应。在我们对 loss 的行为建立了一些直觉之后,我们会通过仔细检查 loss 的 Hessian,来研究一种更局部、更细粒度的 loss 结构。这个 Hessian 会揭示 loss landscape 中一些有趣且令人惊讶的事实,并印证我们的宏观研究。
Mode Connectivity
我们先从深度学习中的 minima 研究开始。深度学习中的目标函数是非凸的,而且几乎可以肯定,对任意给定的深度学习目标,都存在许多许多 minima。在深度学习早期,这个 local minimum 问题曾是 optimization 的核心问题之一。然而,由于多种因素共同作用,这些问题今天看起来远没有那么突出,包括 deep networks 中使用 residual connections、overparametrization,以及优化器中的高随机性。
(a) 没有 skip connections
(b) 有 skip connections
有/无 skip connections 的 ResNet-56 的 loss surfaces。Li, Xu, Taylor, Studer, and Goldstein, “Visualizing the Loss Landscape of Neural Nets,” NeurIPS 2018 的 Figure 1。
现在我们要研究:这些干预是否已经完全消除了 local minima,或者在一个 deep neural network 中是否在某种意义上“所有 minima 都是好的 minima”。为了看清这一点,我们会转向 convex optimization 中的一个基本事实。
第 3 页
对于凸函数 $F$,由 Jensen's inequality 可得
$$F\left(\frac{x+y}{2}\right)\le \frac{F(x)+F(y)}{2}.$$
我们可以利用这个事实,为 deep neural network 的 minima 结构构建一个简单的探针。我们可以多次运行 stochastic optimization 过程,记录最终迭代点 $(\theta_1, \theta_2, \ldots)$,并研究 loss 如何随这些参数的凸插值($\hat{\theta}=\sum_i \alpha_i \theta_i$,其中 $\alpha$ 位于 simplex 上)而变化。
Problem 1: Mode Connectivity
两个训练好的网络可能有相近的 validation loss,却处在参数空间中非常不同的位置。在这道题中,你将检验 parameter-space averaging 何时能保持 validation loss。
给定两个 checkpoints $\theta_1$ 和 $\theta_2$,定义 midpoint barrier
$$B(\theta_1,\theta_2)=L\left(\frac{\theta_1+\theta_2}{2}\right)-\frac{L(\theta_1)+L(\theta_2)}{2}.$$
这里 $L$ 是 validation loss。以 $B\le 0.01$ 作为 weight averaging 对这两个 checkpoints 有效的判据。
使用标准的 depth-8、614M-token transformer 设置,采用默认 optimizer,但把默认 schedule 替换为 warmup 之后的 constant learning rate。
Starter 用 TrainConfig 字段设置这个 sweep。一次 prefix run 会在每个 $x$ 处保留其 model 和 optimizer state。每个 branch 从该 state 分叉(fork_from_run、fork_from_step),使用自己的 data order(batch_order_seed),并在再经过 500 次 updates 后停止(stop_at_step),同时在每个 $x+y$ 处保留一个 checkpoint。因为 learning rate 是常数,一个 500-update branch 就覆盖了所有 $y$。measure 阶段会评估每个 cell 的 interpolation loss 和 ensemble loss。
experiments/a3_optimization/p1_mode_connectivity.py
BASE = TrainConfig(lr_schedule="wsd0.0")
# constant after warmup
def branch(x, order_seed, base=BASE):
fork = {}
if x > 0:
fork = dict(fork_from_run=training_run_name(prefix_for(base)),
fork_from_step=x)
return replace(base, batch_order_seed=order_seed,
stop_at_step=x + max(BRANCH_STEPS),
keep_checkpoint_steps=tuple(x + y for y in BRANCH_STEPS),
**fork)
(a)
设置 forked-run sweep。对每一对 $(x,y)$,先训练一个共享 prefix 走 $x$ 步,复制 model weights 和 optimizer state,然后用相同的 constant learning rate、相同的数据分布、不同的 data order,继续训练两个 branch 各 $y$ 个额外 steps。扫描 $x\in\{0,50,100,250,500,1000\}$,$y\in\{25,50,100,250,500\}$。
对每个 cell,评估两个 endpoint losses,以及 0.5/0.5 prediction ensemble 的 loss。对于 next-token probabilities $p_{\theta_1}$ 和 $p_{\theta_2}$,使用
$$p_{\text{ens}}(\cdot\mid c)=\frac{1}{2}p_{\theta_1}(\cdot\mid c)+\frac{1}{2}p_{\theta_2}(\cdot\mid c).$$
第 4 页
在 $(x,y)$ grid 上绘制 endpoint loss 和 ensemble loss。
(b)
在同一个 $(x,y)$ grid 上,评估 0.5/0.5 weight averaging:
$$\theta_{\text{avg}}=\frac{1}{2}\theta_1+\frac{1}{2}\theta_2.$$
对每个 cell,报告两个 endpoint losses、weight-averaged loss $L(\theta_{\text{avg}})$,以及 barrier $B(x,y)$。把 $B(x,y)$ 画成 heatmap,并标出 $B(x,y)\le 0.01$ 的 cells。根据这张 heatmap,描述哪些地方的 runs 可以 averaging、哪些地方不可以。
(c)
从 $(x,y)$ grid 中选择有代表性的 cells,评估完整的 interpolation curve
$$\theta(\alpha)=(1-\alpha)\theta_1+\alpha\theta_2,\quad \alpha\in\{0,0.25,0.5,0.75,1\}.$$
绘制 validation loss 对 $\alpha$ 的图。把 interpolation curves 与 midpoint barrier 以及 0.5/0.5 prediction ensemble 进行比较。
(d)
探索哪些训练选择会控制模型最终是否落入同一个 low-loss basin,判断依据是上面的 midpoint barrier 和 interpolation curves。尝试改变 learning rate、batch size 和 learning-rate schedule 等 knobs。你发现了哪些模式,哪些进一步的实验有助于检验你的解释?
第 5 页
Assignment 3: Optimization 2
截止日期:10 月 20 日下午 1 点(Quiz 时间)
River-Valley Dynamics
现在让我们把注意力从极小值转向动力学。在实际优化中,迭代点会如何变化?我们过去见过一些动力学模型(例如带噪二次模型,出现在我们之前的作业和讲座中),它们对超参数有一些宏观上的预测价值。
我们现在想更深入一些,观察优化景观的细粒度结构。我们可能好奇的一个现象是:为什么 learning rate schedules 对深度学习中的优化如此重要。
即使在深度学习早期,人们也早已知道,像 cosine schedule(Loshchilov and Hutter, 2017)这样的调度器会显著改善优化,即使在 batch optimization 这类我们并不预期强烈需要 learning rate decay 的情形中也是如此。此外,当时也不清楚应该如何推理 learning rate decay 的最优形状;理论告诉我们,它也许应该看起来像 $O(1/\sqrt{t})$ 或 $O(1/t)$。
一条新兴的研究路线(Wen et al., 2024; Dremov, Hägele, Kosson, and Jaggi, 2025)一直主张:深度学习优化器通常并不是可预测地沿着平滑景观下降(如 NQM 所预测的那样带一些噪声),而是具有复杂、混沌的动力学。$\theta_t$ 的某些分量可能剧烈地来回振荡,在优化景观中的“墙壁”之间弹跳,而其他分量可能随时间更缓慢、更可靠地移动。这个过程的一个生动意象是“river-valley”景观:优化器沿着山谷下降,但在山谷两侧来回弹跳。
River-valley 图景。在较大的 learning rate 下,迭代点在山谷陡峭的两侧之间弹跳,同时沿着河流漂移;衰减 learning rate 会把它们带到河流中。图来自 Wen, Li, Wang, Hall, Liang, and Ma, “Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective,” ICLR 2025 的 Figure 1(a)。
我们将研究这一现象,首先通过 learning rate schedules 和平滑的视角,然后通过 Hessian 的视角。
Problem 2: River-Valley Dynamics
使用默认的 depth-8、614M-token 设置,只改变 schedule 和 EMA 设置。在相同 token 数下比较 validation losses。相对于某个 reference run 定义 regret 为
$$R_M(t) = L(\theta_t^M) - L(\theta_t^{\text{ref}}).$$
Starter 会为 (a) 和 (b) 部分启动四个 schedules,并将仅用于评估的 EMAs 记录为
第 6 页
val_ema{decay}_loss:
experiments/a3_optimization/p2_river_valley.py
EMA_DECAYS = (0.99, 0.999)
SCHEDULE_RUNS = [
TrainConfig(lr_schedule=lr_schedule, ema_decays=EMA_DECAYS)
for lr_schedule in ("wsd0.0", "linear", "cos", "wsd0.2")
]
这里 wsd0.0 表示 warmup 之后保持常数,wsd0.2 表示在训练最后 20% 内衰减。
对于 (c) 部分,lr_dip=(start, bottom, end, scale) 会把任意 schedule 乘以一个临时的 dip,其中各时间以训练进度的比例给出。
对于 (d) 部分,lr_schedule="invsqrt" 和 "sgdrT"(配合 min_lr_ratio)实现下面两个 schedules,而 warmup_steps 固定 updates 中的 warmup 长度。
(a) 比较 constant-after-warmup、linear、cosine 和 plateau-then-decay schedules。绘制训练过程中的 learning rate、training loss 和 validation loss。描述 loss 形状。为确保结果稳健,你可能需要在备选值(例如 $1.5 \times 10^{-3}$ 和 $6 \times 10^{-3}$)上调节 peak learning rate。使用最终 validation loss 最低的 run 作为 (b) 部分的 reference。
(b) 在每个 run 中跟踪模型参数的 exponential moving average (EMA):
$$\bar{\theta}_0 = \theta_0,\qquad \bar{\theta}_t = \beta_{\text{EMA}}\bar{\theta}_{t-1} + (1-\beta_{\text{EMA}})\theta_t.$$
EMA 仅用于评估。在所有 schedules 上尝试几个 averaging timescales,并绘制 raw 和 EMA validation loss 以及 regret。平均之后,schedules 之间的差距还剩多少?
(c) 在训练中途衰减 learning rate,然后在不重置 model 或 optimizer state 的情况下把它 warm 回之前的水平。改变这一 intervention 的深度和持续时间。使用一个不做 intervention 的 control,并让两个 runs 具有相同的 final cooldown。绘制相对于这个 control 的 regret。loss 改善在 rewarmup 之后是否仍然存在,final regret 是多少?是否可能根据 learning rate schedule 等超参数预测 regret 的持续时间和幅度?
(d) 测试下面两个 horizon-free learning-rate schedules,保持 AdamW 不变:
- Inverse-square-root decay(Vaswani et al., 2017, Section 5.3):对于 update $t \ge 1$,使用
$$\alpha_t = \alpha_{\max}\min\{t/w, \sqrt{w/t}\},$$
其中 $w$ 是固定的 warmup 长度(以 updates 计)。 - Cosine warm restarts(Loshchilov and Hutter, 2017, SGDR):在固定 warmup 之后,重复
$$\alpha_t = \alpha_{\min} + \frac{1}{2}(\alpha_{\max}-\alpha_{\min})\left[1+\cos(\pi s/T)\right],$$
其中 $s$ 是固定长度 $T$ 的 cycle 内的 update index。只重启 learning rate,不重启 model 或 optimizer state。
独立于最终 budget 选择 $w$、$T$ 和 learning-rate bounds。让每个 schedule 在多个 stopping budgets 上保持不变,不添加针对 budget 的 cooldown,并将 validation loss 和 regret 与针对每个 budget 调过的 schedules 比较。当训练 budget 未知时,这些是有用的替代方案吗?
第 7 页
Hessian and Sharpness
我们现在更详细地研究 river valley 景观结构。损失景观中的哪些结构性因素可能导致这些行为?为了理解这一点,我们会稍微绕道,看看一个非常基本的凸优化现象,以及它与深度学习训练动力学之间的联系。
考虑步长为 $\eta$ 的 gradient descent,作用于一个二次函数,其 Hessian 的最大特征值为 $\lambda_{\max}$。如果有 $\eta\lambda_{\max} < 2$,迭代点会收敛;而如果有 $\eta\lambda_{\max} > 2$,迭代点会振荡并最终爆炸。noisy quadratic model 会暗示,即使是在深度神经网络中,$\eta < 2/\lambda_{\max}$ 也是收敛所必需的。
然而,Cohen et al. (2021) 发现训练并不是这样运作的。在真实网络上的 full-batch gradient descent 中,sharpness $\lambda_{\max}$ 往往会在训练期间上升(这一现象称为 progressive sharpening),直到达到 $2/\eta$,然后停止上升并停留在那里。这通常被称为“edge of stability”,因为优化过程似乎会在景观中找到动力学勉强稳定的位置,并停留在那里。
在 Adam 这类我们采取 preconditioned gradient step 的情形中,有一个类似的量称为 preconditioned sharpness(Cohen et al., 2022),它也会逐渐增加,并稳定在一个依赖于 $\beta$ 的特定值上(对于 $\beta_1 = 0.9$,该值为 $38/\eta$)。
这些是来自小型 full-batch 实验的干净结果。在这个问题中,你将检查它们是否对我们的 transformer 成立,然后看看在我们实际进行的 minibatch training 中还剩下什么。
更正式地说,对于参数 $\theta$,sharpness $\lambda_{\max}$ 是损失在一个固定 probe set of sequences 上的 Hessian $H$ 的最大特征值。对于 AdamW,我们有 preconditioned sharpness
$$\lambda_{\max}^{\text{pre}} = \lambda_{\max}\left(D^{-1/2} H D^{-1/2}\right), \qquad D = \sqrt{\hat{v} + \epsilon},$$
其中 $\hat{v}$ 是 AdamW 的偏差校正二阶矩估计。对于 SGD 和 Adam,我们预期 $\eta\lambda_{\max}$ 和 $\eta\lambda_{\max}^{\text{pre}}$ 分别 sharpen 并稳定到一个常数。在经典的 edge-of-stability 情形中,这个稳定点发生在二次函数的 edge of stability 处;对 SGD 而言,这发生在 $\eta\lambda_{\max} < 2$ 时。对于 $\beta_1 = 0.9$ 且 preconditioner 冻结的 Adam,这变成 $\eta\lambda_{\max}^{\text{pre}} < 2(1+\beta_1)/(1-\beta_1) = 38$。
我们将在整个训练过程中计算所有这些量。计算 Hessians 的计算代价相当高,所以我们将不得不缩小一些实验规模,并利用巧妙的计算技巧,例如在 Hessian-vector products 上使用 Lanczos iteration 来找到最大特征值。Starter code(见下)已经提供了所有这些可用工具。
Problem 3: The Edge of Stability (22 runs)
这个问题大量使用 Hessian eigenvalues,helper code 位于 experiments/a3_optimization/p3_edge_of_stability.py。它有两个 helpers,会返回一个附带了 sharpness logger 的 TrainConfig:用于 (a) 部分的 full_batch(N, updates, ...),以及用于 (b) 和 (c) 部分的 minibatch(...)。提供了 sharpness.py 来读取存储在 weights and biases 中的 Hessian statistics。
experiments/a3_optimization/p3_edge_of_stability.py
FULL_BATCH_RUNS = [
full_batch(32, 1500, learning_rate=0.01),
# (a) i: the clean case
]
MINIBATCH_RUNS = [
minibatch(),
# (b): the base recipe
]
第 8 页
Part (a) 使用全批量训练,在少量序列上进行,即 Cohen et al. 的设置。Part (b) 和 (c) 使用我们实际采用的小批量训练。整个过程里,你会希望把 sharpness 画在对数轴上,并在稳定性阈值处画一条虚线。
(a) Replicating Cohen et al (Full Batch).(复现 Cohen et al.,全批量)
在前 $N$ 条训练序列上,使用 16 倍 embedding 初始化,运行梯度下降,共 1,500 次更新(full_batch(N, 1500, learning_rate=...))。(3 runs)
i. The clean case.(干净情形)
在 $N=32$、$\eta=0.01$ 下训练。画出 $\lambda_{\max}$ 随 update 的变化,并在第二个面板中画出它随 training loss 的变化。描述发生了什么。这是否符合 edge-of-stability 的预测?与 Cohen et al. 中的图作比较。
ii. Changing $\eta$ mid-run.(在运行中途改变 $\eta$)
重跑 clean case,但在 update 750 时将 $\eta$ 减半(lr_schedule="step750x0.5")。然后改为将 $\eta$ 加倍再跑一次。先预测每种情况下 $\lambda_{\max}$ 会到哪里,再检查。这种响应是对称的吗?
(b) Minibatch training.(小批量训练)
现在研究小批量情形下的 edge of stability,类似我们标准的 d8 设置。minibatch() 给出一个简单配方(但使用常数 LR 和短得多的 horizon,以节省计算)。在 held-out probe 上画出 $\eta_t \lambda_{\max}^{\text{pre}}$ 随 training tokens 的变化,去掉 warmup。minibatch AdamW 会在 edge 上训练吗?(1 run)
(c) Open ended hyperparameter exploration.(开放式超参数探索)
研究超参数对 edge of stability 的影响。哪些因素会影响模型进入并停留在 edge of stability 的倾向?(18 runs)
第 9 页
Inside the Hessian(Hessian 内部)
在 edge-of-stability 问题中,我们看到 Hessian 的最大特征值在深度学习优化中表现出非常典型的行为。事实证明,Hessian(包括最大特征值以及其特征值谱)一直是开发新的优化和学习算法的关键角色,所以我们要再进一步,以获得对 Hessian 的直觉。
Our objects of study(我们的研究对象)
直接构造 Hessian 在计算上不可行。对于一个有 $m$ 个参数的模型,Hessian 有 $m^2$ 个元素。因此,我们会把研究集中在 Hessian 的几个关键方面。对于特征值,我们关心最大特征值的大小(sharpness)、它相对于迹的大小(归一化 sharpness)、负特征值(saddles)以及整体分布。
在特征值之外,我们还关心 Hessian 的块结构。Hessian 中的非对角元素可以看作变量之间的“依赖关系”(即,改变一个变量要求我们改变另一个变量,以保持 loss 相近)。为了度量这一点,我们会查看 Hessian $H$ 中对应于参数族的子矩阵,例如:input embedding、attention projections、QK-norm gains、MLP weights、其他 RMSNorm gains 以及 output head。如果这些 Hessian 子矩阵的 Frobenius 范数接近整体 Hessian 的 Frobenius 范数,这就意味着变量之间的相互作用大多是“within block”的。
Tools(工具)
Hessian 是计算昂贵、难以驾驭的对象,所以我们会给你一些标准工具,以便进行这项分析。$H$ 是在固定 probe 上平均 next-token loss 的 Hessian,该 probe 由 edge-of-stability 问题 part (b) 中使用的 64 条 held-out training-source 序列组成。为了避免我们的算法中可能出现的数值问题,我们会使用 FP32 并关闭 TF32。
- 对 Hessian-vector products 做 Lanczos iterations,可以计算 $H$ 的顶部和底部特征对,以及残差 $\|Au - \lambda u\| / |\lambda|$。这正是我们在 edge-of-stability 问题中用来获得最大特征值的方法。我们会更完整地利用 Lanczos iteration 的信息,包括用它的 Ritz values 估计特征值密度。
- Hutchinson。我们经常需要估计 Hessian 矩阵的迹和平方迹。这可以用 Hutchinson estimator 轻松估计。对于随机符号向量 $z$,$\mathbb{E}[z^\top A z] = \operatorname{tr} A$,且 $\mathbb{E}[z \odot A z] = \operatorname{diag} A$。重要的是,这可以用来估计子矩阵中捕获的 Hessian 比例,因为如果 $z$ 在块 $b$ 上是随机符号向量、在其他位置为零,那么对每个块 $c$,有 $\mathbb{E}\|(Hz)_c\|^2 = \|H_{cb}\|_F^2$。
Problem 4: Inside the Hessian (8–20 runs)(问题 4:Hessian 内部)
starter 是 experiments/a3_optimization/p4_inside_the_hessian.py,每个部分(a、b-rescale、b-train、b-measure、c)各有一个 stage,外加一个 fetch stage。工具在同一个文件夹的 hessian.py 中(ProbeHessian(model, probe) 给出 H.hvp(v);partition 和 restrict 给出 blocks),hessian_measure.py 中有 measure() 和 part (c) 的 continue_projected(),hessian_checkpoints.py 加载提供的 runs 和你自己的 runs,hessian_jobs.load_result(name) 读取已 fetch 的结果。
experiments/a3_optimization/p4_inside_the_hessian.py
LOOK_JOBS = [
# part (a)
{"name": f"a-adamw-{step}",
"measure": {"run": "adamw", "step": step, "extremes": True, "families": True,
"energy": ("family", "layer"),
"density_runs": 4 if step in DENSITY_AT else 0}}
for step in CHECKPOINTS
]
SUBSPACE_JOBS = [
# part (c)
{"name": f"c-{arm}", "continue": {**START, "transform": arm}}
for arm in ("full", "top", "removed")
]
第 10 页
(接上页代码块)
"measure": {"run": "adamw", "step": step, "extremes": True, "families": True,
"energy": ("family", "layer"),
"density_runs": 4 if step in DENSITY_AT else 0}}
for step in CHECKPOINTS
]
SUBSPACE_JOBS = [
# part (c)
{"name": f"c-{arm}", "continue": {**START, "transform": arm}}
for arm in ("full", "top", "removed")
]
Checkpoints。默认的 d8 AdamW run 和一个带有 momentum 的 SGD run(峰值学习率 1、momentum 0.9、无 weight decay)已提供,带有从 0 到 9,375 的 13 个 step 的 optimizer state(614M tokens;step 61 是 4M tokens,153 是 10M,305 是 20M,1,526 是 100M)。
在 H100 上,probe 上的一个 Hessian-vector product 大约需要 2 秒,所以要通过计算 products 来规划测量:使用这些选项时,一个 part (a) checkpoint 大约需要 13 分钟(一个 density 会再增加 13 分钟),而 part (c) 每 50 步大约需要 20 分钟。始终使用全部 64 个 probe rows。
(a) What the Hessian looks like (0 runs).(Hessian 长什么样)
使用 AdamW checkpoints,描述 Hessian 在训练过程中如何变化。
- 哪个 family 最 sharp?这如何变化?画出每个 family block 的最大特征值、$\lambda_{\max}$ 和 $\operatorname{tr}(H^2)/d$ 随 tokens 的变化。
- 谱看起来是什么样?估计 initialization、10M、100M 和 614M tokens 处的特征值密度。有离群值吗?有负特征值吗?这对训练有什么含义?
- 它的块对角程度如何?在训练过程中,按 family 和按 layer 将 $\operatorname{tr}(H^2)=\|H\|_F^2$ 拆分为块内部分和块间部分。与把参数随机划分成相同大小的块作比较。曲率会随深度增长吗?
(b) Move $\lambda_{\max}$ without moving the loss (about 4–10 runs).(在不改变 loss 的情况下移动 $\lambda_{\max}$)
在三个 seeds 上,默认配方最终得到的 $\lambda_{\max}$ 在 4.5 到 6.0 之间,validation loss 在 2.925 到 2.931 之间。找到对训练配置的改变,使最终 $\lambda_{\max}$ 至少改变 $2\times$,同时让最终 validation loss 保持在该范围内。解释你采取这些干预的理由。
- Starting point。有没有办法改变一些权重,使模型做出相同的预测,同时其 Hessian 发生变化?(Hint:回想我们在课上关于 loss invariance 的讨论)。在我们的架构中找到这样一种变换,手工把它应用到最终 checkpoint 上,并测量受影响 blocks 在变换前后的 loss 和曲率。
- Sharpening in training。找到一种对训练配方的改变,能在没有人工干预的情况下诱导出这种 sharpening 效应。仅靠超参数干预可能让你取得不错的效果,但最强的干预可能还涉及一些代码改动。如果我们想跨 runs 比较 $\lambda_{\max}$,这意味着什么?
(c) Training in and out of the top subspace (about 3–10 short runs).(在顶部子空间内外训练)
Gur-Ari et al. (2018) 发现梯度主要位于前几个 Hessian 特征向量的张成空间中。测试这在我们的 d8 系统中使用带动量的 SGD 是否成立。(我们从 Adam 切换到 SGD,因为这样更容易推理 Hessian 和我们的 updates 之间的关系)。
- Starting point。从 100M tokens 的 SGDM checkpoint 开始,将 momentum 置零,并继续 300 步,使用三种 updates:完整 update、投影到前 10 个特征向量上的 update,以及去掉该分量的 update。重新计算特征向量
第 11 页
每 25 步。记录 loss、gradient 在 subspace 中的占比,以及 $\lambda_{\max}$ 相对于 heavy-ball threshold $\frac{2(1+\beta)}{\eta_t}$ 的取值。
- 调查这一现象。其中一些结果可能出乎意料。改变一系列你选择的 hyperparameters(例如 $k$、refresh interval、learning rate、starting point),看看你的结论是否会改变。是否有改动让训练更快?你对 Hessian 中的 eigendirections(特征方向)如何关联到 optimizer trajectory 有什么心智模型?
可选探索(Optional exploration)
为了更广泛地探索本作业中的主题,构造你自己的 prediction problem(预测问题)。这个最后的练习与 quiz question 无关,但可能带来 bonus points。
Problem 5:你自己的预测问题(5–10 次运行)
你现在已经充分探索了本作业中的现象,可能发现了出乎意料或令人惊讶的现象。我们希望基于你的洞察,帮助制作更好、更有趣的 quizzes。
在你的探索中包括以下内容:
- (a) 一份题目说明(question writeup),类似于示例问题或之前的 quizzes。答案应当可以基于某段 code 的执行结果进行客观检查(最好是 multiple choice 或 numerical prediction)。
- (b) 一个 code diff,以及可执行 code,能够得出题目中的答案(例如打印 loss)。code diff 应当紧凑($<100$ 行代码)。
目标是提出一个公平、新颖的问题:其结果应当能够利用本作业进行推理,同时测试超出所给示例的洞察。课程工作人员可能会在未来的 quiz 中使用有趣的问题。
材料
自查
- 交付物:每题写清对照是什么、扫了哪些超参、最好的 val loss 与诊断量、你的解读。
- 报告 ≤ 20 页;曲线要能“自解释”(读者不看正文也大致看懂对比)。
这份作业和书的关系
配第 5–6 讲《损失景观(上/下):几何、尖锐度与稳定边缘》。 书里讲的“盆地/河谷/临界点”“edge of stability”,这里要用 loss 曲面的实验把它测出来。