Learning-rate schedules, warmup and gradient clipping in practice

Before this lesson: lesson 05 (the noise ball of SGD, the linear scaling rule) and lesson 06 (Adam and AdamW). The stability limit $\eta < 2/\lambda_{\max}$ comes from lesson 03, the descent lemma from lesson 02. Time: about 80 minutes, plus the notebook.

Why this matters #

Almost every large model is trained with AdamW, a warmup, a decay and gradient clipping. GPT-3 used a linear warmup over the first 375 million tokens, a cosine decay to 10% of the peak learning rate, and clipped the global gradient norm at 1.0 (Brown et al. 2020, appendix B). LLaMA used 2,000 warmup steps, a cosine decay to 10% of the peak and clipping at 1.0 (Touvron et al. 2023). GraphCast, a weather model, used a linear warmup from 0 to $10^{-3}$ over 1,000 steps and a half-cosine decay to 0 (Lam et al. 2023). Newer recipes replace the cosine by warmup–stable–decay (WSD), because it does not fix the training length in advance. When a run diverges or stalls, the schedule, the peak learning rate and the clipping threshold are the first things to check.

1. Predict first #

2. The picture #

A schedule is a function $\eta(t)$. In this lesson and in lesson 06 the step counter is $t = 1, 2, \dots, T$, as in the Adam paper and the PyTorch documentation; $T$ is the total number of steps. Here are the common shapes. All of them start with the same linear warmup over $W$ steps.

Why decay at the end. Lesson 05 showed that SGD with a constant step does not converge to the minimizer. It reaches a noise ball: a cloud around the minimizer whose size is proportional to $\eta$. A smaller $\eta$ gives a smaller cloud and a lower loss. But a small $\eta$ from the start would make slow progress. So we keep $\eta$ large while we are far from the minimum, and decay it at the end.

Why warm up at the start. The first steps are the least stable: the weights are random, the curvature can be high, and Adam's estimate of the gradient scale comes from only a few gradients. A short linear warmup lets the model reach a region where a large learning rate is safe. We look at the evidence in section 3.

The next figure runs SGD on a noisy quadratic with five schedules. It plots the exact expected loss, so the curves are smooth.

Clipping limits the length of a step. Some losses have regions where the gradient and the curvature are both very large. A step of fixed size $\eta$ then jumps far past the minimum. For $f(x) = x^4/4$ with $\eta = 0.1$, gradient descent from $x_0 = 5$ goes to $-7.5$, then $34.7$, then $-4139$: it diverges. Clipping caps the step length. With steps of length at most 1, the same run goes $5 \to 4 \to 3 \to 2 \to 1.2 \to 1.027 \to \dots$ and converges.

3. Definitions and main results #

Schedules #

Notation: $T$ total steps, $W$ warmup steps, $D$ cooldown steps, $\eta_{\max}$ the peak and $\eta_{\min}$ the final value. Every schedule below starts with the linear warmup $\eta(t) = \eta_{\max}\,t/W$ for $t < W$.

A warmup–stable–decay (WSD) schedule has three phases: warmup, a long stable phase at $\eta_{\max}$, and a short cooldown (decay) phase. The name comes from MiniCPM (Hu et al. 2024). The 1−sqrt cooldown drops fast at first and then spends more of the cooldown at a small learning rate: at the middle of the cooldown ($p = 0.5$) it is at $1 - \sqrt{0.5} \approx 0.29$ of the peak, the linear one at $0.5$.

Why decay: the noise ball is proportional to $\eta$ #

Proposition (noise ball of SGD in one direction). Let $f(x) = \tfrac{\lambda}{2}x^2$ with $\lambda > 0$, and let SGD use the stochastic gradient $g_t = \lambda x_t + \xi_t$, where the noise $\xi_t$ has mean 0 and variance $\sigma^2$ and is independent of $x_t$. With a constant step $0 < \eta < 2/\lambda$,

Proof.

1. The update is $x_{t+1} = x_t - \eta(\lambda x_t + \xi_t) = (1 - \eta\lambda)x_t - \eta\xi_t$.
2. Square and take expectations. The cross term $\mathbb E[x_t\xi_t] = \mathbb E[x_t]\,\mathbb E[\xi_t] = 0$ because $\xi_t$ is independent of $x_t$. So $\mathbb E[x_{t+1}^2] = (1 - \eta\lambda)^2\,\mathbb E[x_t^2] + \eta^2\sigma^2$.
3. Write $a = (1 - \eta\lambda)^2$. Since $0 < \eta\lambda < 2$, $a < 1$. The fixed point solves $v = av + \eta^2\sigma^2$, so $v^\star = \eta^2\sigma^2/(1 - a)$.
4. $1 - a = 1 - (1 - \eta\lambda)^2 = \eta\lambda(2 - \eta\lambda)$, so $v^\star = \eta\sigma^2/(\lambda(2 - \eta\lambda))$.
5. Subtract the fixed-point equation from step 2: $\mathbb E[x_{t+1}^2] - v^\star = a\,(\mathbb E[x_t^2] - v^\star)$. The gap shrinks by the factor $a$ per step, so $\mathbb E[x_t^2] \to v^\star$. Multiply by $\lambda/2$ for the loss. $\blacksquare$

Two consequences explain the shape of the loss curves.

• Halving $\eta$ about halves the plateau. The plateau is proportional to $\eta$ (exactly, up to the factor $1/(2 - \eta\lambda)$).
• The drop after a decay is fast in sharp directions. By step 5, the gap to the new plateau shrinks by $(1 - \eta\lambda)^2$ per step, so it takes about $1/(2\eta\lambda)$ steps. Directions with large curvature $\lambda$ adjust in a few steps. This is why the loss falls sharply at the start of a cooldown. It is the noise ball shrinking, not a new minimum.

Without noise ($\sigma = 0$) the plateau is 0, and a decay only slows the progress. Decay is a tool against noise.

Why a cosine cannot be stopped early, and why WSD can #

The cosine value at step $t$ depends on $T$. At $t = T/2$ (with $W \ll T$) we have $p \approx 1/2$ and $\eta \approx (\eta_{\max} + \eta_{\min})/2$. So if you stop a cosine run early, the model never had a phase with a small learning rate: it is still inside a large noise ball. If you want to train longer than planned, the learning rate is already low, and you need a new warmup. Every training length needs its own run.

WSD removes this problem (Hägele et al. 2024, section 3):

1. The stable phase does not depend on $T$. You can save checkpoints during it and start a short cooldown from any of them. One long run gives models for many training lengths. Studies of scaling laws, which need models trained for many lengths, therefore need far fewer runs.
2. A cooldown of about 10–20% of the steps matched a tuned cosine schedule in their language-model experiments; the gains stopped growing near 20%. The 1−sqrt shape gave a lower loss than the linear shape.
3. Averaging the weights of recent checkpoints (see the deeper block) also helped, at no extra training cost.

Bergsma et al. (2025) compared decay shapes for language models trained at compute-optimal budgets. A linear decay to zero beat a cosine decay to 10% of the peak. In one of their experiments, a 610M-parameter model trained on 80 tokens per parameter with decay to zero beat the 10× decay trained on 200 tokens per parameter, a compute saving of about 60%.

Why warmup helps: the evidence #

1. Sharpness. Kalra and Barkeshli (2024) studied warmup in detail. Its main effect is that the network can tolerate a larger target learning rate. During warmup, training moves to regions with lower sharpness $\lambda_{\max}(\nabla^2 L)$, the top eigenvalue of the Hessian of the loss. So the large learning rate that comes later stays below the stability limit $2/\lambda_{\max}$ of lesson 03. Whether the sharpness first falls or first rises depends on the initialization and the parametrization. Lesson 11 studies sharpness and the edge of stability.
2. Adam's second moment. Early in training, Adam's $\hat v_t$ is estimated from very few gradients, so the first steps are badly scaled. The same paper shows that a better initial value of $v$ gives some of the benefit of warmup.
3. Large batches. Goyal et al. (2017) trained ResNet-50 with batches of 8,192 images. With the linear scaling rule (learning rate proportional to the batch size, lesson 05), the scaled learning rate was too large for the first epochs. A gradual warmup over 5 epochs fixed this, and the large-batch model matched the small-batch accuracy.

Choosing the peak learning rate #

The peak $\eta_{\max}$ is the most important number of the schedule. Too small: slow progress. Too large: the run diverges or keeps a large noise ball.

The learning-rate range test (Smith 2017; the "LR finder" of the fastai library):

1. Start from the initialization, as in a real run.
2. Train for a few hundred steps while $\eta$ grows, for example exponentially from $10^{-7}$ to $10$. (Smith increased it linearly over a few epochs; the exponential version is common today.)
3. Record a smoothed training loss (an exponential moving average of the minibatch losses).
4. Stop when the smoothed loss exceeds a few times its best value (fastai uses 4 times).
5. Pick a learning rate clearly below the blow-up point. A common rule of thumb is one tenth of the $\eta$ with the lowest smoothed loss. This is a heuristic, not a theorem: it gives a safe start, and a short sweep around it (for example ×0.3 to ×10) often finds a faster value. The notebook shows such a case.

Link to the stability limit. On a quadratic with Hessian $H$, gradient descent diverges exactly when $\eta > 2/\lambda_{\max}(H)$ (lesson 03). The range test finds this limit by experiment, for the real loss, at the start of training. Two cautions: the sharpness of a network changes during training (warmup lowers it, lesson 11 shows that it can also rise), and for Adam the relevant limit is for the preconditioned Hessian. So the range test gives a good starting value, not a guarantee.

Changing the batch size. For SGD, the linear scaling rule multiplies $\eta$ by $k$ when the batch size is multiplied by $k$ (Goyal et al. 2017; lesson 05). It stops working above the critical batch size. For Adam, Malladi et al. (2022) derived a square-root scaling rule ($\eta \propto \sqrt B$, together with changes of $\beta_1$, $\beta_2$ and $\varepsilon$). Treat both as starting points and check them with a short run.

Gradient clipping by global norm #

Let $g = (g^{(1)}, \dots, g^{(m)})$ be the gradients of all $m$ parameter tensors of a model, viewed as one long vector. Its global norm is $\|g\| = \big(\sum_i \|g^{(i)}\|^2\big)^{1/2}$. Clipping by global norm with threshold $c > 0$ replaces $g$ by

Properties.

1. The direction is unchanged: the result is a positive multiple of $g$.
2. The length is $\min(\|g\|, c)$.
3. If $\|g\| \le c$, nothing changes. So clipping acts only on rare, large gradients.
4. With step $\eta$, a clipped step is $x_{k+1} = x_k - h_k\nabla f(x_k)$ with $h_k = \min(\eta,\ \eta c/\|\nabla f(x_k)\|)$. So the step length is $\min(\eta\|\nabla f\|,\ \eta c)$: at most $\eta c$.

In PyTorch, torch.nn.utils.clip_grad_norm_(params, max_norm) computes the global norm, multiplies every gradient by $c/(\|g\| + 10^{-6})$ when this factor is below 1, and returns the norm before clipping. Log this number: a rising gradient norm often comes before a loss spike.

Other kinds of clipping change the direction. Clipping each tensor separately (per-tensor norm) or each entry separately (clip_grad_value_, which clamps every entry to $[-c, c]$) rescales different parts by different factors.

Why clipping helps: $(L_0, L_1)$-smoothness #

The convergence theory of lesson 02 assumes $L$-smoothness, $\|\nabla^2 f(x)\| \le L$ everywhere. Then the safe fixed step is about $1/L$. But for $f(x) = x^4/4$ the curvature $f''(x) = 3x^2$ has no upper bound, so no single $L$ exists. Zhang, He, Sra and Jadbabaie (2020) proposed a weaker condition.

Definition. A twice differentiable $f$ is $(L_0, L_1)$-smooth if

With $L_1 = 0$ this is $L$-smoothness. The new condition allows the curvature to grow where the gradient is large.

Example. $f(x) = x^4/4$ is $(3, 3)$-smooth: we need $3x^2 \le 3 + 3|x|^3$. For $|x| \le 1$, $3x^2 \le 3$. For $|x| \ge 1$, $3x^2 \le 3|x|^3$. Zhang et al. also measured the local smoothness along real training runs and found that it grows with the gradient norm, most clearly for an LSTM language model.

Why a fixed step is slow and a clipped step is fast. A fixed step must be safe at the worst point the run can visit, where the curvature can be $L_0 + L_1 M$ ($M$ the largest gradient norm on the way). So it must be about $1/(L_1 M)$, far too small elsewhere. The clipped step $\min(\eta, \eta c/\|\nabla f\|)$ behaves like $1/(L_0 + L_1\|\nabla f(x)\|)$: the safe step for the local curvature.

Reading loss curves #

A loss curve is the main tool to debug a schedule. These patterns are common; the first thing to try is listed.

Large language-model runs sometimes show loss spikes that do not recover. The PaLM team restarted such runs from a checkpoint about 100 steps before the spike and skipped about 200–500 batches of data (Chowdhery et al. 2022).

Answers to section 1. (1) No. At step 50,000 the cosine planned for 100,000 steps is about halfway, so its learning rate is still about $(\eta_{\max} + \eta_{\min})/2$: the model never had a low-learning-rate phase and sits in a large noise ball. The cosine planned for 50,000 steps has decayed to $\eta_{\min}$ and reaches a lower loss. (2) The gradient of norm 50 is scaled by $1/50$: same direction, length 1. The gradient of norm 0.5 is not changed. (3) The loss drops within a short time to a new plateau about half as high (the noise ball is proportional to $\eta$), and then progress continues more slowly.

4. Worked example #

(a) Schedule values. $T = 20{,}000$, $W = 2{,}000$, $\eta_{\max} = 10^{-3}$.

Warmup + cosine with $\eta_{\min} = 10^{-4}$:

• $t = 1{,}000$ (warmup): $\eta = 10^{-3}\cdot 1{,}000/2{,}000 = 5\cdot10^{-4}$.
• $t = 2{,}000$: the warmup ends and $p = 0$, so $\eta = 10^{-4} + \tfrac12\cdot 9\cdot10^{-4}\cdot 2 = 10^{-3}$.
• $t = 6{,}500$: $p = 4{,}500/18{,}000 = 0.25$, $\cos(\pi/4) = 0.7071$, so $\eta = 10^{-4} + 4.5\cdot10^{-4}\cdot 1.7071 = 8.68\cdot10^{-4}$.
• $t = 11{,}000$: $p = 0.5$, $\cos(\pi/2) = 0$, so $\eta = 10^{-4} + 4.5\cdot10^{-4} = 5.5\cdot10^{-4}$.
• $t = 20{,}000$: $p = 1$, $\cos\pi = -1$, so $\eta = 10^{-4}$.

WSD with $D = 4{,}000$ (20% of $T$) and $\eta_{\min} = 0$: $\eta = 10^{-3}$ until $t = 16{,}000$.

• $t = 18{,}000$: $p = 2{,}000/4{,}000 = 0.5$. Linear: $10^{-3}(1 - 0.5) = 5\cdot10^{-4}$. 1−sqrt: $10^{-3}(1 - 0.7071) = 2.93\cdot10^{-4}$.
• $t = 19{,}000$: $p = 0.75$. Linear: $2.5\cdot10^{-4}$. 1−sqrt: $10^{-3}(1 - 0.8660) = 1.34\cdot10^{-4}$.

(b) Clipping. Two tensors have gradients $g^{(1)} = (2, -1)$ and $g^{(2)} = (2)$; the threshold is $c = 1.5$.

• Global norm: $\|g\| = \sqrt{4 + 1 + 4} = 3$.
• Factor: $1.5/(3 + 10^{-6}) \approx 0.5$.
• Clipped gradients: $(1, -0.5)$ and $(1)$. The direction is the same; the norm is 1.5.

For comparison, clipping each tensor to norm 1.5 gives $(2, -1)\cdot 1.5/\sqrt5 = (1.342, -0.671)$ and $(1.5)$: the cosine with the original direction is 0.998. Clipping each value to $[-1, 1]$ gives $(1, -1, 1)$, with cosine 0.962. Both change the direction.

(c) Clipped gradient descent on $x^4/4$. $\eta = 0.1$, $c = 10$, $x_0 = 5$. Then $f'(x) = x^3$ and the step length is $\min(0.1|x^3|,\ 1)$.

• $x_0 = 5$: $|f'| = 125 > 10$, so the step has length 1: $x_1 = 4$.
• $x_1 = 4$: $|f'| = 64 > 10$: $x_2 = 3$.
• $x_2 = 3$: $|f'| = 27 > 10$: $x_3 = 2$.
• $x_3 = 2$: $|f'| = 8 \le 10$, no clipping: $x_4 = 2 - 0.1\cdot 8 = 1.2$.
• $x_4 = 1.2$: $x_5 = 1.2 - 0.1\cdot 1.728 = 1.0272$, and plain gradient descent continues to 0.

Plain gradient descent from the same start: $x_1 = 5 - 12.5 = -7.5$, $x_2 = -7.5 + 42.19 = 34.69$, $x_3 \approx -4139$. The local curvature at $x_0$ is $3\cdot 25 = 75$, far above $2/\eta = 20$.

(d) The noise ball. For $f(x) = x^2/2$ ($\lambda = 1$) with noise $\sigma = 1$, the proposition gives the plateau $\eta/(2(2 - \eta))$: $0.1/3.8 = 0.0263$ for $\eta = 0.1$, and $0.01/3.98 = 0.0025$ for $\eta = 0.01$. Dividing $\eta$ by 10 divides the plateau by about 10.

5. In code #

NumPy checks the numbers of section 4. The notebook has these cells and the exercises.

import numpy as np

def warmup_cosine(t, T, W, lr_max, lr_min=0.0):
    if t < W:
        return lr_max * t / W
    p = min(1.0, (t - W) / (T - W))
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + np.cos(np.pi * p))

def wsd(t, T, W, D, lr_max, lr_min=0.0, shape="linear"):
    if t < W:
        return lr_max * t / W
    if t < T - D:
        return lr_max
    p = min(1.0, (t - (T - D)) / D)
    s = 1 - p if shape == "linear" else 1 - np.sqrt(p)
    return lr_min + (lr_max - lr_min) * s

# (a) the schedule values of section 4
T, W = 20_000, 2_000
print([f"{warmup_cosine(t, T, W, 1e-3, 1e-4):.3e}" for t in (1000, 2000, 6500, 11000, 20000)])
# ['5.000e-04', '1.000e-03', '8.682e-04', '5.500e-04', '1.000e-04']
print([f"{wsd(t, T, W, 4000, 1e-3, shape=s):.3e}" for s in ("linear", "1-sqrt") for t in (18000, 19000)])
# ['5.000e-04', '2.500e-04', '2.929e-04', '1.340e-04']

# (b) clipping by global norm keeps the direction
def clip_by_global_norm(grads, max_norm, eps=1e-6):
    total = np.sqrt(sum(np.sum(g**2) for g in grads))
    factor = min(1.0, max_norm / (total + eps))
    return [g * factor for g in grads], total

clipped, total = clip_by_global_norm([np.array([2.0, -1.0]), np.array([2.0])], max_norm=1.5)
print(total, [c.round(4).tolist() for c in clipped])   # 3.0 [[1.0, -0.5], [1.0]]

# (c) GD and clipped GD on f(x) = x^4/4, eta = 0.1, threshold c = 10
x, y, eta, c = 5.0, 5.0, 0.1, 10.0
for k in range(4):
    x = x - eta * x**3
    g = y**3
    y = y - min(eta, eta * c / abs(g)) * g
    print(k + 1, f"{x:.5g}", f"{y:.4g}")
# 1 -7.5 4
# 2 34.688 3
# 3 -4139 2
# 4 7.0906e+09 1.2

# (d) the noise ball of SGD on f(x) = x^2/2 with gradient noise sigma = 1
rng = np.random.default_rng(0)
for eta in (0.1, 0.01):
    x = np.zeros(20_000)                        # 20,000 independent runs
    for _ in range(3000):
        x -= eta * (x + rng.standard_normal(x.size))
    print(eta, f"{np.mean(0.5 * x**2):.4f}", f"{eta / (2 * (2 - eta)):.4f}")
# 0.1 0.0260 0.0263
# 0.01 0.0025 0.0025

In PyTorch, the same schedules come from torch.optim.lr_scheduler (LinearLR, CosineAnnealingLR, SequentialLR, or a LambdaLR with your own function), and clipping is one line before optimizer.step(): total = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0). With mixed precision and a gradient scaler, unscale the gradients first (scaler.unscale_(optimizer)), then clip.

6. Where it shows up #

• Language models. AdamW with a linear warmup, a cosine, WSD or linear decay, and clipping at 1.0 is the standard pretraining recipe (GPT-3, LLaMA, MiniCPM). WSD is popular when the final training length is not known, or when one run should give checkpoints for several lengths.
• Vision. Step decay by 10× at fixed epochs was the classic ResNet recipe; Goyal et al. added a 5-epoch warmup for batches of 8,192 images. Vision transformers use warmup + cosine.
• ML for science. GraphCast used AdamW, a 1,000-step linear warmup and a half-cosine decay to 0 (Lam et al. 2023). Physics-informed neural networks are often trained with Adam and a step or exponential decay of the learning rate, and then refined with L-BFGS (lesson 10).
• Monitoring. Log the learning rate and the gradient norm before clipping at every step. Most training problems show up first in these two curves.

7. Common mistakes #

• Stopping a cosine run early, or extending it. The cosine value depends on $T$. Plan $T$, or use WSD.
• Schedule length in the wrong unit. T_max counted in epochs while scheduler.step() is called every batch (or the reverse): the schedule ends far too early or never decays.
• Calling scheduler.step() before optimizer.step(). Since PyTorch 1.1 the scheduler is stepped after the optimizer; the other order skips the first value of the schedule.
• Clipping per tensor or per value when you meant the global norm. Both change the direction of the update.
• Clipping before unscaling with mixed-precision loss scaling: the threshold then applies to the scaled gradients.
• Resuming without the optimizer and scheduler state. Adam's moments and the step counter restart, which is like a start without warmup at the peak learning rate.
• Reading a drop at the cooldown as a better model architecture or data change. It is the noise ball shrinking. Compare runs at the same point of their schedules.
• Taking the range-test minimum as the learning rate. The minimum is close to the blow-up point; pick a value clearly below it.

8. Check yourself #

1. With $\eta_{\min} = 0$, at which fraction $p$ of the cooldown is the learning rate $\eta_{\max}/2$, for the linear and for the 1−sqrt shape?
   Answer

   Linear: $1 - p = 1/2$, so $p = 0.5$. 1−sqrt: $1 - \sqrt p = 1/2$, so $\sqrt p = 1/2$ and $p = 0.25$. The 1−sqrt cooldown reaches half the peak twice as early.
2. You have a 10,000-step WSD run with a 10% cooldown and a 10,000-step cosine run. Now you need a model trained for 20,000 steps. What can you reuse?
   Answer

   From the WSD run: the checkpoint at step 9,000, the end of the stable phase. Continue at $\eta_{\max}$ to step 18,000 and then cool down for 2,000 steps. The cosine run cannot be extended without a new warmup, because its learning rate has already decayed; you start a new run.
3. Two tensors have gradients $(0.3, 0.4)$ and $(1.2)$. What does clip_grad_norm_ with max_norm=1.0 return, and what are the clipped gradients?
   Answer

   The global norm is $\sqrt{0.09 + 0.16 + 1.44} = \sqrt{1.69} = 1.3$, and the function returns 1.3. The factor is $1/(1.3 + 10^{-6}) \approx 0.769$, so the gradients become about $(0.231, 0.308)$ and $(0.923)$: the same direction, norm 1.
4. For SGD on $f(x) = \tfrac{\lambda}{2}x^2$ with $\lambda = 4$, noise $\sigma = 1$ and $\eta = 0.05$, what is the expected loss on the plateau?
   Answer

   $\eta\sigma^2/(2(2 - \eta\lambda)) = 0.05/(2\cdot 1.8) \approx 0.0139$. (The variance is $v^\star = 0.05/(4\cdot 1.8) \approx 0.00694$, and $\tfrac{\lambda}{2}v^\star = 2\cdot 0.00694 \approx 0.0139$.)
5. Your run diverges at step 40 with $\eta_{\max} = 3\cdot10^{-3}$, a warmup of 10 steps and no clipping. List the first two changes you would try, in order.
   Answer

   First a longer warmup (for example 1–5% of the steps) and/or a lower peak learning rate, checked with a range test. Second, turn on clipping by global norm (for example at 1.0) and log the gradient norm. Check also for bad data or a wrong loss scale.
6. Why is $x^4/4$ not $L$-smooth, and which $(L_0, L_1)$ work for it?
   Answer

   $f''(x) = 3x^2$ grows without bound, so no constant $L$ bounds it. $(L_0, L_1) = (3, 3)$ works: $3x^2 \le 3$ for $|x| \le 1$ and $3x^2 \le 3|x|^3$ for $|x| \ge 1$.

9. Interview angle #

• "Why warm up the learning rate?" — Early training is the least stable: sharpness can be high and Adam's $\hat v$ is based on few gradients. Warmup lets the model move to flatter regions so the peak learning rate stays below $2/\lambda_{\max}$; with large batches and the linear scaling rule it is required (Goyal et al.).
• "Cosine or WSD?" — WSD matches cosine with a 10–20% cooldown, does not need $T$ in advance, and one run serves many lengths. Cosine needs $T$ and cannot be stopped early.
• "Why does the loss drop at the end of training?" — The learning-rate decay shrinks the noise ball of SGD; its loss is proportional to $\eta$.
• "How do you choose the learning rate?" — A range test for the order of magnitude, then a small sweep around it; scale with the batch size (linear for SGD, about square root for Adam) and re-check.
• "Why clip by the global norm and not per value?" — The global norm keeps the update direction; per-value clipping changes it. Clipping helps on losses whose curvature grows with the gradient ($(L_0, L_1)$-smoothness).

10. Practice #

The notebook for this lesson runs in Colab. It has four exercises, each with a check cell and a folded solution:

1. Implement warmup + cosine and WSD (linear and 1−sqrt cooldowns) and check the values of section 4.
2. Run SGD on a noisy quadratic with a constant, a cosine and a WSD schedule; compare the final loss with the noise-ball formula.
3. Implement clipping by global norm over a list of arrays; check the norm, the direction and the returned value.
4. Run a learning-rate range test for softmax regression on the scikit-learn digits data and pick a learning rate.

11. Watch and read #

• Hägele et al., Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations (2024): WSD, cooldown length and shape, weight averaging.
• Kalra and Barkeshli, Why Warmup the Learning Rate? Underlying Mechanisms and Improvements (2024).
• Zhang, He, Sra and Jadbabaie, Why gradient clipping accelerates training (ICLR 2020): $(L_0, L_1)$-smoothness.
• Smith, Cyclical Learning Rates for Training Neural Networks (2017): the range test.
• Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour (2017): linear scaling and warmup.
• Bergsma et al., Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs (2025).
• Defazio et al., The Road Less Scheduled (2024): schedule-free optimization.
• Malladi et al., On the SDEs and Scaling Rules for Adaptive Gradient Algorithms (2022): the square-root scaling rule for Adam.
• PyTorch documentation: how to adjust the learning rate and clip_grad_norm_.