6-month learning plan · · session 20 / 13020 / 130

mathbeginner 55m read

S020 · Calculus II — Gradients & Gradient Descent from Scratch

The one algorithm behind all of ML: gradients as the ‘uphill direction’ in n dimensions, SGD as the workhorse, and the learning-rate/momentum knobs that decide whether you converge or explode.

📐MathM02 · Math Foundations· Session 020 of 130 90 min

🎯 Implement gradient descent from scratch on linear regression, feel what learning rate does, and know when to reach for momentum, Adam, or a schedule.

Why this session exists

Every model in machine learning — linear regression, logistic regression, neural nets, transformers with 175 billion parameters — is trained by the same algorithm. Compute the gradient of the loss with respect to the parameters. Step in the opposite direction. Repeat. That's it. If you can implement gradient descent from scratch on a 2-parameter problem, you already understand how GPT is trained — the rest is engineering (parallelism, memory, tricks). This session gets you fluent in the mechanic that underlies everything.

You will be able to

Compute the gradient of a multi-variable function by hand and interpret it as ‘direction of steepest ascent’.
Implement vanilla gradient descent, SGD, and momentum in ≤ 20 lines each — and know when to use which.
Choose a learning rate that neither diverges nor stalls (and diagnose which failure you're seeing).
Explain why deep-learning defaults are Adam + cosine schedule + gradient clipping.
Watch a loss curve and diagnose learning-rate too high / too low / just right in 5 seconds.

Prerequisites

S019 · Derivatives & chain rule — a gradient is just a vector of partial derivatives.
S017/S018 · Linear algebra — you'll do vector arithmetic and matrix updates.

Pre-read · skim before the session (30 min total)

💡 Intuition 22 min

Gradient descent · Chapter 2 · Deep Learning

3Blue1Brown

The visualisation that unlocked gradient descent for a million people. Watch even if you've seen it.

🔬 Deep dive 145 min

The Spelled-out Intro to Neural Networks and Backpropagation

Andrej Karpathy

Watch 0-60 min first: builds micrograd and does gradient descent by hand. Come back for the rest.

🛠️ Hands-on 20 min

Gradient descent, how neural networks learn | Deep Learning Chapter 2

3Blue1Brown

Comparison of the optimisers you'll actually use in PyTorch. Watch before section (b).

📖 Reading 25 min

Deep Learning · Chapter 8 · Optimization

Goodfellow, Bengio, Courville · free online

The reference. Sections 8.1–8.5 cover everything through Adam. Bookmark for whole book.

📖 Reading 20 min

An overview of gradient descent optimization algorithms

Sebastian Ruder

The single most-cited blog post on optimisers. All variants, one page.

📖 Reading 15 min

Why Momentum Really Works

distill.pub

Interactive visualisations. The intuition that makes momentum feel obvious.

(a) Intuition · 5 min

You're on a foggy hillside; find the valley

🌍 Real world

You're on a mountain in dense fog. You want to reach the valley but can only see the ground under your feet. Every few seconds you check which direction is most downhill and take a step that way. If the slope is steep, big step. If the ground is flat, tiny step. Eventually you're at the valley bottom — where the slope in every direction is zero.

That fog is what an ML model sees. It doesn't know the shape of the loss landscape globally. It only knows the local slope. Gradient descent is the ‘check underfoot, step downhill, repeat’ strategy — repeated a million times for a million-parameter landscape.

💻 Code world

In code, the "slope in every direction" is the gradient — a vector where entry i is ∂loss/∂parameter_i. You step against it: params -= lr · gradient. The learning rate lr is your step size. Too big and you jump over the valley. Too small and you never arrive. Getting it right (or getting an adaptive optimiser to get it right for you) is the whole art of training.

Every optimiser you'll ever use — SGD, momentum, RMSprop, Adam, AdamW, Lion — is a variation on this one theme.

The gradient in one line

What the gradient IS

For a scalar function f(x, y, z), the gradient ∇f = [∂f/∂x, ∂f/∂y, ∂f/∂z]. A vector of partial derivatives.
Geometrically: ∇f points in the direction of steepest ASCENT. |∇f| is how steep. −∇f is steepest DESCENT.
At a minimum, ∇f = 0. That's the fixed point gradient descent converges to.
For a loss with 175 billion parameters, ∇f is a vector of 175 billion numbers. Every training step computes it.

A quick history

1847
Cauchy · gradient descent
Augustin-Louis Cauchy publishes the method for solving equations. 175 years old and still the winner.
1951
SGD · Robbins & Monro
Stochastic approximation — descend using noisy per-sample gradient estimates. The seed of mini-batch SGD.
1964
Momentum · Polyak
Add a fraction of the previous step to the current — dampens oscillations, accelerates flat regions.
2012
AdaGrad · Duchi et al.
Per-parameter adaptive learning rates. Then RMSprop, then Adam (2014). Deep learning explodes.
2018+
AdamW · Loshchilov & Hutter
Decoupled weight decay. The default for training modern transformers (BERT, GPT, LLaMA).
2023+
Lion, Sophia, Shampoo
Newer optimisers experimenting with sign-based updates and second-order approximations. Adam still dominates production.

(b) Visual walkthrough · 15 min

The update rule, forever

Batch · mini-batch · stochastic — three flavours

Batch GD

Full dataset per step

True gradient — no noise
One step per epoch
Slow, memory-hungry
OK for classical convex problems, bad for modern deep learning

SGD (stochastic)

One sample per step

Very noisy gradient — but that noise helps escape local minima
Cheap per step
Requires small learning rate
Rarely used pure — mini-batch dominates

Mini-batch SGD

The workhorse

Batch of 32-1024 samples per step
Compromise: gradient noise controlled but still cheap per step
Utilises GPU parallelism efficiently
This is what ‘SGD’ means in modern practice

The optimiser family tree

Optimisers, from simple to state-of-the-art

SGD (vanilla)

θ ← θ − lr · g. Simple, works surprisingly well with the right lr + schedule. Still used for computer-vision training.

1847

SGD + momentum

Track a velocity v = β·v + g, then θ ← θ − lr · v. β ≈ 0.9. Damps oscillations, accelerates flat regions.

1964

AdaGrad / RMSprop

Per-parameter learning rate scaled by 1/sqrt(sum of squared past gradients). Adapts to gradient scale.

2012

Adam

Momentum + adaptive learning rate. β₁≈0.9, β₂≈0.999, lr≈1e-3. Default for 90% of deep learning.

2014

AdamW

Adam with decoupled weight decay. Default for training LLMs, ViTs, and most modern transformers.

2018

Lion (experimental)

Sign-based, memory-efficient. 2-3× fewer optimiser-state bytes. Useful when memory-bound.

2023

Learning-rate diagnostics from a loss curve

Loss goes DOWN smoothly → lr is roughly right

You're descending the loss surface. Consider slightly larger lr to see if you can converge faster.

Loss oscillates or bounces wildly → lr too HIGH

You're overshooting the valley each step. Halve the lr and retry.

Loss barely moves → lr too LOW (or vanishing gradient)

You're taking micro-steps. Double the lr. If still no movement, check gradient magnitudes — may be a vanishing-gradient problem.

Loss explodes to NaN → lr WAY too high or gradient issue

Weights diverged; some intermediate value overflowed. Rewind to a checkpoint, clip gradient norms, use smaller lr.

Loss goes down then plateaus → schedule the lr

Decay the learning rate (cosine, step, or exponential). Modern training almost always uses a schedule.

Common misconception

✗ What most people think

"Gradient descent finds the minimum. If training converged, I'm at the best solution the model can reach — and if it didn't converge, I should train longer."

✓ What is actually true

Gradient descent finds a point where the gradient is approximately zero, reachable from where you started. For non-convex losses that is one of astronomically many stationary points, and which one you land in depends on initialisation, batch order, and learning rate schedule. "Converged" means "stopped moving", not "found the best".

Why the myth is so sticky

Because for the convex problems used to teach it — linear regression, logistic regression — the myth is literally true: there is one minimum, and gradient descent finds it from anywhere. Everyone's first optimisation experience confirms it. Then the same intuition is carried into deep networks where the loss surface is wildly non-convex, and it produces two specific wrong conclusions. First, that a run that plateaued is done, when it may have stalled on a saddle point or in a flat region that momentum or a learning-rate change would escape. Second, that two runs of the same code should agree — they don't, and the spread between seeds is real information about your setup rather than noise to be hidden by fixing the seed and reporting one number.

Prove it to yourself

Same function, same optimiser, different start — different answer:

import numpy as np
f  = lambda x: x**4 - 3*x**3 + 2       # two basins
df = lambda x: 4*x**3 - 9*x**2

def descend(x, lr=0.01, steps=2000):
    for _ in range(steps): x -= lr * df(x)
    return x

print(descend(-1.0), descend(3.0))   # different minima
print(descend(0.0))                  # stuck: gradient is exactly 0 here

From first principles

Start with the question

Why does stochastic gradient descent — using a noisy estimate from a small batch — work better than full-batch descent, not just faster? Using less data per step should be strictly worse information.

1
The true gradient is the average of per-example gradients over the whole dataset, so a random batch's gradient is an unbiased estimator of it.
forced by · the expectation of a uniform random sample's mean equals the population mean
2
The standard error of that estimate shrinks like 1/√B in batch size B. So going from batch 100 to batch 10,000 costs 100× the compute to cut the noise by only 10×.
forced by · estimator variance falls linearly in sample size, so the error falls with its square root
3
Therefore per unit of compute, many noisy small steps make far more progress than one precise large step — the direction only has to be roughly right for the step to reduce the loss.
forced by · descent needs a descent direction, not the optimal one, and any direction with positive inner product with the true gradient works
4
Separately, the noise means the iterate never sits exactly at a stationary point. Random perturbation kicks it off saddle points, where the true gradient is exactly zero and full-batch descent would halt forever.
forced by · a saddle is unstable in at least one direction, so any perturbation with a component along it escapes
5
And the noise scale acts like a temperature: it cannot settle in a minimum whose basin is narrower than the typical step, so it preferentially ends up in wide, flat minima.
forced by · escaping a basin requires a fluctuation larger than the basin's width, which is easy for narrow ones and hard for wide ones

⇒ Therefore

Therefore the noise in SGD is not a tolerated cost of cheap computation — it is doing real work: escaping saddles and biasing the search toward flat regions, which tend to generalise better than sharp ones.

And note what this predicts, all observable: very large batches reduce the noise and are commonly reported to need explicit compensation — learning-rate scaling and warmup — to match small-batch generalisation. It predicts that decaying the learning rate late in training is what finally lets the iterate settle, since it shrinks the effective temperature. And it predicts why batch size and learning rate must be tuned together rather than independently: it is their ratio that sets the noise scale, so changing one alone changes the optimisation regime.

Mental modelBlindfolded on a hillside, feeling the slope

You are on a landscape in the dark. You cannot see the terrain — you can only feel which way is downhill under your feet, and take a step. That is the entire algorithm. The gradient is the local feel of the slope; the learning rate is your stride length; the loss curve is your altitude log.

Everything that goes wrong has an obvious physical reading: too big a stride and you leap across the valley and up the far side (loss diverges); too small and you take a week to descend (slow convergence); a flat plateau and you can't tell which way to go (vanishing gradient); a narrow ravine and you bounce between the walls while barely advancing along it (which is exactly what momentum fixes by accumulating the consistent direction and cancelling the oscillation).

Learning rate is the single most important hyperparameter. Find it first, by sweeping orders of magnitude — 1e-1 to 1e-5 — not by tuning it in fine increments.
Diverging loss or NaN means the step is too large for the local curvature. Lower the learning rate, add warmup, or clip gradients — in that order.
Features on wildly different scales create ravines, because the loss curves far more steeply along one axis than another. Normalising inputs is not cosmetic; it reshapes the landscape into something a single global step size can handle.
Always plot the loss curve. Diverging, plateaued, noisy-but-descending, and overfitting each have a distinct visual signature, and none of them is visible in a final metric.

🔔 Fires when you see

Fire this model the moment you see: loss going to NaN · a training curve that is flat from step one · validation loss rising while training loss falls · results that swing between seeds · a "just train longer" suggestion · any hyperparameter search that didn't start with the learning rate.

The tradeoff

Which optimiser: plain SGD, SGD with momentum, or an adaptive method like Adam?

Plain SGD

+ you gain one hyperparameter, minimal memory (no per-parameter state), and completely predictable behaviour — what you see is the raw gradient

− you pay painfully slow in ravines, oscillating across the steep direction while creeping along the shallow one; extremely sensitive to feature scaling and to the learning rate

pick when convex or well-conditioned problems, or when you want a baseline with no confounding optimiser dynamics

SGD + momentum (+ schedule)

+ you gain accumulates consistent directions and cancels oscillation, so it traverses ravines efficiently; widely reported to reach the best final generalisation in vision, given a good schedule

− you pay two coupled hyperparameters plus a schedule to design; needs more tuning effort and more patience to get right

pick when you have the budget to tune, the training run is long, and final quality matters more than time-to-first-result — the classic choice for large vision training

Adam / AdamW

+ you gain per-parameter adaptive step sizes make it robust to bad scaling and sparse gradients; works acceptably out of the box, which is why it's the default for transformers and anything with embeddings

− you pay two extra state tensors per parameter (roughly triple the optimiser memory), and it can converge to solutions that generalise slightly worse; plain Adam's weight decay is coupled to the adaptive scaling, which is precisely the bug AdamW fixes

pick when gradients are sparse or badly scaled, or you need a good result quickly without a tuning campaign — the sane default for a first run on almost anything

What a senior engineer actually does

Start with AdamW, because time-to-first-signal matters more than the last fraction of a point, and it removes learning-rate sensitivity as a confounder while you're still debugging the data and the model. Move to tuned SGD+momentum only if you have a long training budget and evidence the optimiser is what's limiting you.

And keep the sequence honest: optimiser choice is almost never the binding constraint. Data quality, feature scaling, and the learning rate schedule dominate it by a wide margin. Swapping optimisers to fix a bad training run is a way of avoiding the harder question of whether the loss, the labels, or the inputs are wrong.

(c) Hands-on · 25 min

You'll fit a straight line to noisy data three ways — analytic solution, vanilla GD, and SGD with momentum — and watch the learning-rate knob live. Save as gd_demo.py.

Run it:

uv run --with numpy python gd_demo.py

Expected output:

What each block does

Anatomy of the code

analytic_ols

The exact solution via normal equations: θ = (AᵀA)⁻¹ Aᵀy. Only works for linear regression; for anything nonlinear we NEED gradient descent.

closed-form

mse_gradient

Hand-derived from MSE loss = mean((wx+b-y)²). Chain rule: dLoss/dw = 2·mean((wx+b-y) · x). This is what autograd would compute automatically.

derivation

gradient_descent

Vanilla GD in 6 lines. Full-batch: uses ALL 200 samples per gradient. Great for teaching; too slow for large datasets.

vanilla

sgd_momentum

Two upgrades: (a) mini-batches of 32 samples for cheap noisy gradients; (b) velocity term v that averages recent gradients. Converges in fewer effective steps.

modern

compare_learning_rates

lr=0.001 → too slow (didn't reach optimum in 200 steps). lr=0.05 → just right. lr=0.3 → diverges to NaN. The one-hyperparameter horror show.

sensitivity

diverging_run

Watch weights double-and-flip until they overflow to inf → NaN. This is what a ‘loss went to NaN’ log line looks like from the inside.

explosion

Try itFeel the loss surface as a 3D landscape

Visualise the loss landscape (requires matplotlib):

import matplotlib.pyplot as plt
 
ws = np.linspace(-1, 5, 50)
bs = np.linspace(-2, 4, 50)
W, B = np.meshgrid(ws, bs)
L = np.zeros_like(W)
for i in range(W.shape[0]):
    for j in range(W.shape[1]):
        L[i, j] = mse_loss(np.array([W[i, j], B[i, j]]), X, Y)
 
plt.figure(figsize=(7, 6))
plt.contour(W, B, L, levels=20)
plt.plot(2, 1, "r*", markersize=15, label="true (2, 1)")
plt.xlabel("slope w"); plt.ylabel("intercept b")
plt.title("MSE loss surface for y = wx + b")
plt.legend(); plt.savefig("loss_surface.png", dpi=100)

Now overlay the trajectory of a GD run at three different lrs and see the paths side-by-side. This IS what "training curves" show, just in 2 dimensions instead of a million.

💡 Hint · After you generate the contour, overlay the trajectory of your GD run (`plt.plot([θ[0] for θ, _ in history], [θ[1] for θ, _ in history])`). You'll see the little zig-zag path down to the valley. Try lr=0.5 — the path spirals outward instead of inward. That's divergence in one picture.

(d) Production reality · 15 min

War story Meta AI · training LLaMA· 202365B parameters, 2000+ A100 GPUs

🔥 What broke

The LLaMA paper reports several training runs going wildly divergent — loss spikes to 10× baseline and never recovers. Requires rewinding hundreds of billions of tokens of compute (millions of dollars).

Cause: rare gradient spikes on specific batches (numerical instability + fat-tailed input distribution) that momentarily exceed the ‘safe’ update magnitude.

🧯 The fix

Gradient clipping (clip_grad_norm_(params, max_norm=1.0)) before every optimiser step. If the gradient's total L2 norm exceeds 1.0, rescale it to length 1.0. Doesn't change direction; only magnitude. Also: skip the optimiser step entirely on any batch with NaN loss and rewind. Standard practice for all frontier LLM training.

🎓 Lesson to steal

At scale, ‘just once in a billion’ numerical issues become certain. Gradient clipping + NaN detection + checkpoint-then-skip are cheap insurance against very expensive divergence.

Post-mortem

War story Common failure · not shuffling the datasilent performance loss

🔥 What broke

A junior ML engineer trains an image classifier and gets 62% accuracy. Same code, same hyperparameters, shuffle turned on → 89%. The dataset was sorted by class; mini-batches were always all-cats or all-dogs; the gradient signal was terrible.

🧯 The fix

Always shuffle training data at the start of every epoch. In PyTorch: DataLoader(..., shuffle=True). In fit-loops: rng.permutation(n). The default is ‘on’ for good reason — if you're overriding, know why.

🎓 Lesson to steal

Mini-batch SGD assumes each batch is a rough sample of the whole distribution. Sorted data breaks that assumption invisibly. Shuffle every epoch, no exceptions.

War story Common failure · Adam without weight decayresearch reproducibility

🔥 What broke

A paper reports SOTA using Adam + L2 regularisation. Nobody can reproduce the exact number. Turns out L2 (weight decay implemented as adding λ·w to the gradient) interacts oddly with Adam's per-parameter scaling — the effective regularisation strength depends on gradient magnitudes.

🧯 The fix

Use AdamW, not Adam. AdamW applies weight decay outside the adaptive-lr scaling: θ ← θ − lr · (Adam_update + λ · θ). Loshchilov & Hutter's 2018 paper. It's the default optimiser for modern transformers (BERT, GPT-family, ViT).

🎓 Lesson to steal

Optimiser choice interacts with regularisation in subtle ways. Read the paper. When in doubt, use AdamW for transformers and SGD-momentum for computer vision.

Where this shows up in the rest of the plan

Gradient descent is the algorithm behind all learning

S078 · Linear + logistic regression

Fit models by GD on the log-likelihood or MSE loss.

S083 · Neural networks

Backprop computes gradients; SGD/Adam updates weights. Same loop as this session, 10⁹ params.

S090 · Advanced optimisers

Adam, AdamW, Lion, Sophia — variations on the same theme with better defaults for deep nets.

S099 · Reinforcement learning · policy gradient

Policy = a neural net; loss = negative expected reward; still gradient descent.

S115 · LLM pretraining

Trillions of tokens × billions of parameters × AdamW + cosine schedule. It's this session, scaled.

S120 · Fine-tuning + LoRA

Freeze most weights, train a few low-rank matrices with… gradient descent. Same recipe.

(e) Recall + stretch · 10 min

Recall — click each to reveal · click to reveal

★ = stretch question

Explain-out-loud test

What does gradient descent do, in one paragraph, to a total beginner?
What is the learning rate and what happens if it's too big or too small?
Why does every deep-learning framework compute gradients automatically?

What comes next

📐MathNext →

S021 · Probability — Random Variables, Distributions, Expectation

Hub: The 6-Month Learning Plan

Part of a 130-session evergreen learning series. Session structure: intuition → visual → hands-on → production war stories → recall. Duration: 90 minutes.