LM5-T2 — Backpropagation and Gradient-Based Optimization

Learning objectives

After this lesson, you should be able to:

• Apply the chain rule to a tiny network and compare with PyTorch autograd.
• Write the gradient-descent update and explain the role of the learning rate.
• Compute one Adam step by hand and explain AdamW’s decoupled weight decay.
• Describe warmup plus cosine decay and gradient-norm clipping, and compute a schedule value.

The idea: walk downhill on the loss surface

Imagine the loss as the height of a landscape whose coordinates are all the parameters of the model. At the current point, the gradient is the direction of steepest ascent. Moving a small step the opposite way lowers the loss. Repeat, and you walk downhill. This is gradient descent. A model with 124 million parameters is a point in a space of 124 million dimensions, but the procedure is the same.

Backpropagation

The loss is a composition of many simple functions: embedding lookup, linear layers, attention, normalisation, softmax, the logarithm. The chain rule says that the derivative of a composition is the product of the derivatives of its parts. Backpropagation applies it efficiently: it runs the network forward and stores intermediate values, then walks backward from the loss, multiplying local derivatives, and obtains the gradient with respect to every parameter in a single pass costing about as much as two forward passes. Frameworks do this automatically (autograd); you write the forward pass and call loss.backward().

The learning rate

The gradient gives a direction, not a distance. The learning rate scales the step. Too small and training crawls; too large and the steps overshoot, the loss bounces or explodes. It is the most important hyperparameter, and it should change during training.

Adam and AdamW

Different parameters have gradients of very different sizes, and the noise from random batches is large. Adam keeps a running average of each parameter’s gradient (the momentum, which smooths noise) and of its squared gradient (which estimates its typical scale). Each update is the smoothed gradient divided by the square root of the typical squared gradient, so every parameter moves by about the learning rate per step, whatever the scale of its gradient. AdamW additionally shrinks the weights slightly at each step (weight decay) in a way that is decoupled from the gradient, which regularises the model and is the standard in language-model training.

Schedule and clipping

Early in training, Adam’s statistics are unreliable and the model is fragile, so the learning rate starts near zero and rises linearly during a warmup. After the peak it decays smoothly, usually along a half cosine, to a small floor, so that late updates are fine adjustments. Gradient clipping rescales the whole gradient when its norm exceeds a limit, so that one bad batch cannot throw the weights far away.

The mathematics: gradient, Adam, schedule, clipping

Gradient descent. For parameters w and loss L, with learning rate η:

Example: L(w) = (w − 3)², gradient 2(w − 3), η = 0.1, start at w = 0. Step 1: 0 − 0.1 · (−6) = 0.6. Step 2: 0.6 − 0.1 · (−4.8) = 1.08. Step 3: 1.08 − 0.1 · (−3.84) = 1.464. Each step covers 20% of the remaining distance to 3.

Chain rule. For a chain L = ℓ(a), a = g(b), b = h(w): dL/dw = ℓ′(a) · g′(b) · h′(w). For a linear layer y = Wx with upstream gradient ∂L/∂y = δ, the gradient of the weights is the outer product δx^T (summed over the positions in the batch), and the gradient passed back to the input is W^Tδ.

Adam. With gradient g_t at step t, decay rates β₁ = 0.9, β₂ = 0.999 (common values), ε = 10⁻⁸:

m_t

running average of the gradient (first moment); smooths the noise.

v_t

running average of the squared gradient (second moment); the typical gradient scale.

m̂_t, v̂_t

bias-corrected values; m and v start at 0, so early averages are biased toward 0.

η

the learning rate; ε avoids division by zero.

First step with g = 2: m = 0.1 · 2 = 0.2, v = 0.001 · 4 = 0.004; corrections divide by 0.1 and 0.001, giving m̂ = 2 and v̂ = 4; the update is η · 2 / (2 + 10⁻⁸) ≈ η. The first step has size η whatever the gradient was: for g = 200 it is the same. That is the scale invariance of Adam.

AdamW adds decoupled weight decay λ (for example 0.1):

The decay acts on the weights directly, not through the gradient, so it is not rescaled by Adam’s denominator.

Schedule. With peak η_max, floor η_min, warmup T_w steps and T total steps:

For T = 1,000, T_w = 100, η_max = 6·10⁻⁴, η_min = 6·10⁻⁵: η(50) = 3·10⁻⁴; η(100) = 6·10⁻⁴; η(550) = 6·10⁻⁵ + ½ · 5.4·10⁻⁴ · (1 + cos(π/2)) = 3.3·10⁻⁴; η(1000) = 6·10⁻⁵.

Gradient clipping. With maximum norm c: if ‖g‖ > c then g ← g · c / ‖g‖. For g = [3, 4] (norm 5) and c = 1: [0.6, 0.8]. The direction is kept; only the length changes. The norm is taken over all parameters together.

A worked example: checking gradients and one Adam step

The code checks a hand-derived gradient against autograd, repeats the gradient descent example, performs the first Adam step by hand and with torch.optim.Adam, evaluates the schedule and clips a gradient.

import math
import torch

# 1. a chain rule by hand versus autograd: L = (sigmoid(w*x) - y)^2
w = torch.tensor(0.5, requires_grad=True)
x, y = 2.0, 1.0
s = torch.sigmoid(w * x)
L = (s - y) ** 2
L.backward()
manual = 2 * (s - y) * s * (1 - s) * x                 # dL/ds * ds/d(wx) * d(wx)/dw
print(round(w.grad.item(), 6), round(manual.item(), 6))   # equal

# 2. gradient descent on (w - 3)^2
v = 0.0
for _ in range(3):
    v -= 0.1 * 2 * (v - 3)
print(round(v, 3))                                     # 1.464

# 3. the first Adam step is about lr, whatever the gradient
for g in (2.0, 200.0):
    p = torch.nn.Parameter(torch.tensor([1.0]))
    opt = torch.optim.Adam([p], lr=0.01)
    p.grad = torch.tensor([g])
    opt.step()
    print(g, round(1.0 - p.item(), 5))                 # 0.01 both times

# 4. warmup and cosine schedule
def lr_at(t, T=1000, Tw=100, hi=6e-4, lo=6e-5):
    if t < Tw:
        return hi * t / Tw
    return lo + 0.5 * (hi - lo) * (1 + math.cos(math.pi * (t - Tw) / (T - Tw)))
print([round(lr_at(t), 6) for t in (0, 50, 100, 550, 1000)])   # [0.0, 0.0003, 0.0006, 0.00033, 6e-05]

# 5. clipping a gradient of norm 5 to norm 1
q = torch.nn.Parameter(torch.zeros(2))
q.grad = torch.tensor([3.0, 4.0])
norm_before = torch.nn.utils.clip_grad_norm_([q], max_norm=1.0)
print(round(norm_before.item(), 3), [round(v, 3) for v in q.grad.tolist()])   # 5.0 [0.6, 0.8]

Part 3 shows what Adam does to gradients of very different size: both parameters move by 0.01 in the first step. Plain gradient descent would have moved the second parameter 100 times further.

Engineering implications and limitations

Typical values. AdamW with β₁ = 0.9, β₂ = 0.95 to 0.999, weight decay 0.1, peak learning rate of order 10⁻⁴ to 10⁻³ for small models (smaller for larger ones), warmup of 0.5 to 2% of the steps, cosine decay to a tenth of the peak, clipping at norm 1.0. The tiny notebook model tolerates a higher rate (3·10⁻³).

Order of calls. optimizer.zero_grad(), loss.backward(), optional clipping, optimizer.step(), then the scheduler step. Gradients accumulate across backward() calls, so forgetting zero_grad() silently sums gradients over steps.

Exclude norms and biases from weight decay. Decay is normally applied to weight matrices only. The book’s simple loop applies it to all; large-scale recipes use parameter groups.

Learning-rate finder. If a run is unstable, the first thing to try is a lower peak learning rate or a longer warmup. A short sweep on a few hundred steps finds the stable range quickly.

Watch the gradient norm. Log the norm returned by clip_grad_norm_. A sudden spike precedes many divergences, and a norm that is always at the clip limit means the clipping is doing the optimizer’s job.

Optimizer state is memory. Adam stores two extra numbers per parameter (Lesson LM4-T5). Memory-saving variants (8-bit states, Adafactor) exist and are discussed in Module 8.

Check gradients when you write a custom operation. Compare with finite differences or with a reference implementation before trusting it in a long run.

Common misconceptions

Engineering takeaways

1. Gradient descent: w ← w − η∇L. Backpropagation computes ∇L with the chain rule in one backward pass; autograd does it for you.
2. Adam keeps moving averages of the gradient and its square, so each parameter moves by about η per step; its first step is ≈ η for any gradient size.
3. AdamW applies weight decay directly to the weights: w ← w − η(m̂/(√v̂ + ε) + λw).
4. Schedule: linear warmup to a peak, then cosine decay to a floor; e.g. 3.3e-4 at the midpoint of a 6e-4 → 6e-5 decay.
5. Clip the global gradient norm to protect against bad batches; clipping keeps the direction.

References and further reading

See Chapter 5 of the book (training an LLM) and its Appendix D (adding bells and whistles to the training loop: warmup, cosine decay, gradient clipping).

• Kingma, D. P. and Ba, J. (2015). Adam: A Method for Stochastic Optimization. arXiv:1412.6980.
• Loshchilov, I. and Hutter, F. (2019). Decoupled Weight Decay Regularization. arXiv:1711.05101.
• Build a Large Language Model (From Scratch), Chapter 5: Pretraining on Unlabeled Data (primary book reference)
• Build a Large Language Model (From Scratch), Appendix D: Adding Bells and Whistles to the Training Loop (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: the training loop itself, with validation, logging and checkpoints.

Related content

Nothing below is required to finish this lesson.

• Before
  LM5-T1: The Next-Token Objective and Cross-Entropy

  Derive the training loss from the probability of the correct next token and relate it to perplexity.

• Lab
  LM5-L3: Pretrain a Tiny GPT and Diagnose the Run

  Train a small GPT on a small corpus, plot training and validation loss, sample at different stages and diagnose a bad run.