Every neural network on Earth is trained by gradient descent. It's just rolling a ball downhill — repeatedly. GPT-4 cost $100M+ to train. Stripped down, that training is one Python loop: predict, measure error, nudge weights downhill, repeat. By the end of this lesson, you'll know exactly what "training" means, in 20 lines of code.
Learning Objectives
After this lesson, you will be able to:
Explain gradient descent as repeatedly taking small steps downhill to find the lowest point on an error landscape
Predict what happens when your step size is too big (you overshoot and bounce around) or too small (you crawl and never finish), and learn two tricks (momentum and Adam) that help
See that gradient descent is the single algorithm behind how every AI model -- from the simplest to ChatGPT -- actually learns from data
Distinguish between batch, stochastic, and mini-batch gradient descent, and explain why the noise in SGD is a feature rather than a bug
This is THE algorithm. If you only learn one thing from this entire math track, make it gradient descent. It is the beating heart of every AI model ever trained. And the best part? The core idea is something a 10-year-old could understand: feel the slope, step downhill, repeat.
From a simple model with two settings to GPT-4 with hundreds of billions, they all learn by gradient descent. Understanding this algorithm is understanding how machines learn.
Try it! Open the Python REPL (bottom-right of the screen: click Quick Actions, then Python) and type these lines yourself.
Every ML model has a loss function -- a number that measures "how wrong is the model right now?" The goal of training is to make this number as small as possible.
The loss depends on the model's parameters (weights and biases). If you could plot loss as a function of all parameters, you would get a loss landscape -- a surface with mountains (bad parameters), valleys (good parameters), saddle points (tricky flat regions), and plateaus (frustrating flat regions).
For a model with two parameters, the loss landscape is a 3D surface you can visualize. For GPT-4 with ~1.8 trillion parameters, it is a 1.8-trillion-dimensional surface. You cannot visualize it, but the math works the same way.
A function f is convex if a straight line between any two points on its graph lies on or above the graph itself. Formally:
f(λx+(1−λ)y)≤λf(x)+(1−λ)f(y)∀x,y,λ∈[0,1]
A convex function has at most one local minimum, and that local minimum is automatically the global minimum. Gradient descent on a convex loss is provably convergent under mild assumptions. Linear regression, ridge regression, logistic regression, and SVMs all have convex losses -- their training procedures have rigorous convergence guarantees.
Neural network losses are non-convex. The same architecture can have many local minima, saddle points, and broad plateaus. Gradient descent has no guarantee of finding the global minimum; we hope (and empirically observe) that the local minimum it finds is "good enough." The remarkable success of deep learning is partly the discovery that for highly over-parameterized networks, most local minima generalize well -- so the lack of theoretical guarantees does not bite us in practice. But it is why "GD always converges" is a half-truth: convergence to some stationary point is provable under L-smoothness; convergence to a useful solution is empirical.
Toggle between a single-bowl convex landscape and a multi-valley non-convex one, and watch how gradient descent behaves on each. The non-convex case is what every neural network actually faces.
Loading visualization...
Quick check
You are training a logistic regression model and a deep CNN. Both use cross-entropy loss. Which optimization problem is convex?
The gradient is a vector pointing in the direction of steepest uphill. To go downhill, walk in the opposite direction.
∇L=∂w1∂L∂w2∂L⋮∂wn∂L
The gradient field plots a tiny arrow at every point on the loss surface — each arrow shows which way is "uphill" at that location. Gradient descent walks in the opposite direction of whichever arrow it lands on. Drag the start point around and watch the descent trace stream toward the minimum.
Loading visualization...
What Do You Think?
On a circular bowl-shaped loss, you start at point A and take one gradient step. Which direction does the step point?
#The Algorithm: The Most Important Equation in Modern ML
θt+1=θt−η∇L(θt)
Where:
theta_t are the current parameters
eta (Greek letter eta) is the learning rate -- step size
nabla L is the gradient of the loss -- direction
theta_ are the updated parameters
That is the entire algorithm. Three ingredients: current position, step size, and direction. Everything else in modern optimization -- Adam, momentum, learning rate schedules, gradient clipping -- is a refinement of this core idea.
Try it: Explore the gradient descent landscapeInteractive
Loading visualization...
Try this: Start with the default settings and watch the optimizer converge. Then crank the learning rate up to 0.5 and watch it overshoot and oscillate. Then drop it to 0.001 and watch it crawl. The sweet spot is in between. Also try different surfaces -- a bowl is easy, but a surface with saddle points is treacherous.
Before reading on, predict the answer below — it will sharpen your intuition for the next section.
What Do You Think?
You are minimizing f(x) = x² with gradient descent starting at x = 4. What is the LARGEST learning rate η that still converges (i.e., x_t → 0)?
#The Learning Rate: The Most Important Hyperparameter
What Do You Think?
What happens if the learning rate is 10x too large?
The learning rate controls how big each step is. This single number has an outsized effect on whether training succeeds or fails.
With a tiny learning rate, each step barely moves the parameters. The model will improve, but agonizingly slowly. Training might take 100x longer than necessary. You might run out of compute budget before reaching a good solution. The blindfolded hiker takes baby steps -- technically going downhill, but it would take a lifetime to reach the valley.
With a huge learning rate, the model takes giant leaps. Instead of gently descending into a valley, it overshoots the minimum entirely, landing on the opposite mountainside. The next step overshoots again. The loss oscillates wildly or explodes to infinity. The blindfolded hiker takes 50-foot jumps and ends up launching off cliffs.
The ideal learning rate is large enough to make meaningful progress but small enough to not overshoot. In practice, finding this sweet spot requires experimentation or learning rate finders. Modern optimizers like Adam adapt the effective learning rate automatically for each parameter, but the initial learning rate still matters.
A constant learning rate is rarely optimal: you usually want a large rate early in training (to make progress) and a small rate late (to settle into a flat minimum). Schedules vary η over time. Two are essentially universal in modern deep learning:
Why warmup AND decay? Warmup avoids the first-50-step instability when Adam's m_t and v_t moments are still noisy and bias-corrected estimates are unreliable. Decay drives the optimizer into a flatter region of the loss surface near the end, which empirically generalizes better. Almost every transformer paper since GPT-2 uses linear warmup followed by cosine decay; LLaMA, GPT-3, PaLM, and most fine-tuning recipes follow this template.
When gradient norms occasionally explode (RNNs, transformers in FP16, early training), clip them back into a safe range:
g←g⋅min(1,∥g∥2τ)
This single line rescues training from NaN cascades. When training a transformer in mixed precision, an occasional FP16 overflow can make a gradient component infinite -- without clipping, the next parameter update sends weights to NaN and the model is dead. Clipping at τ = 1.0 keeps step sizes bounded regardless. Note: clipping is not a free lunch -- if your gradients are exploding consistently, you have a real problem (bad initialization, learning rate too high, missing layer norm) and clipping is just hiding it.
We start with random parameter values. Our position on the loss landscape is arbitrary -- we have no idea where the minimum is. The loss is high because the model's predictions are essentially random noise.
We calculate the gradient at our current position using backpropagation -- the chain rule applied backward through every layer. For a model with 100M parameters, we get a 100M-dimensional gradient vector in one backward pass.
We update our parameters: theta = theta - lr * gradient. We move downhill. The loss decreases. The model's predictions improve slightly. For GPT-4, this single step updates ~1.8 trillion numbers simultaneously.
Compute a new gradient at the new position. The slope might be different now -- steeper, shallower, or pointing a different direction. Step again. Each iteration brings us closer to a minimum.
Computes the gradient using the entire training dataset before each update. Gives an accurate gradient estimate but is very slow -- imagine recalculating the slope using every data point before each step.
Computes the gradient using a single random example. Extremely fast per step, but the gradient estimate is noisy. The blindfolded hiker feels the slope at one random point under their foot. The noise actually helps escape shallow local minima and saddle points.
The practical middle ground: compute the gradient over a small batch (typically 32-512 examples). Balances accuracy and speed. This is what people usually mean when they say "gradient descent" in the context of deep learning.
Pure gradient descent can oscillate in ravines -- narrow valleys where the gradient zigzags between steep walls instead of rolling smoothly toward the bottom. Momentum fixes this by remembering which direction you have been going.
Adam (Adaptive Moment Estimation) is the workhorse optimizer of modern deep learning. It combines momentum with per-parameter learning rate adaptation:
First moment (m): running average of gradients -- like momentum, it tracks direction
Second moment (v): running average of squared gradients -- tracks magnitude/variance
This means parameters with consistently large gradients get smaller learning rates (preventing explosions), and parameters with small gradients get larger learning rates (ensuring progress). Adam adapts automatically -- it is the default optimizer for most deep learning research and is used for training virtually all large language models.
PPO's clipped surrogate L = E[min(r·A, clip(r, 1-ε, 1+ε)·A)] is gradient ascent with a built-in trust region — the clip prevents the policy from moving too far in one update. RLHF on LLMs uses this exact form.
Watch SGD, Momentum, RMSProp, and Adam compete on the same loss surface:
Try it: Watch optimizers race to the minimumInteractive
Loading visualization...
Try this: Watch how Adam and Momentum find the minimum faster than plain SGD. Notice how SGD zigzags while Momentum smoothly curves toward the answer. On surfaces with saddle points, SGD sometimes gets stuck while Adam pushes through.
Quick check
In the race above, which optimizer is most likely to escape a saddle point first — and why?
Try it: Explore the Loss Landscape in 3DInteractive
Loading visualization...
Explore this: Rotate the 3D loss surface and find the saddle points — flat regions that look like a minimum from one direction but slope upward from another. These fool gradient descent into stopping too early. Enable momentum and watch it roll through flat regions that plain SGD gets stuck in. Try the "valley" surface — notice how SGD oscillates across the valley while Adam finds the centre.
⚡ Playground:Gradient Descent Explorer → — adjust learning rate, momentum, and optimizer and watch convergence live on a 2D loss surface.
Tests · Verify lr=0.1 converges to near 0 in 30 steps. Verify lr=1.0 oscillates without converging. Find the maximum stable learning rate analytically.
Time to do it in numpy on the function from the title: f(x, y) = x² + 4y². This bowl is elongated (4x steeper in y than x), which is exactly the setup that makes plain gradient descent zigzag. Run the cell, then crank lr up toward 0.25 and watch what happens.
Gradient descent is "feel the slope, step downhill". The algorithm iteratively updates parameters by moving in the direction opposite to the gradient, reducing the loss at each step
The learning rate is the most important hyperparameter. Too large causes overshooting and divergence, too small causes agonizingly slow convergence; before redesigning your architecture, try changing the learning rate by factors of 10
Mini-batch SGD is what everyone actually uses. Computing gradients on small random batches (32-512 examples) balances speed and accuracy, and the stochastic noise helps escape saddle points and sharp minima
Momentum smooths out oscillations. By accumulating past gradient direction like a rolling ball, momentum converts zigzag paths in narrow valleys into smooth, fast descents
Adam adapts learning rates per-parameter. By tracking both gradient direction (first moment) and magnitude (second moment), Adam automatically adjusts effective learning rates, making it the default optimizer for most deep learning
What does the gradient of a loss function tell you?
With gradient descent in your toolbox, you now understand the core algorithm that makes machines learn. Next up: Information Theory -- why cross-entropy loss works, what entropy really measures, and the mathematical foundation of every classification model's loss function.
After many iterations (thousands to millions), the gradient becomes tiny. The ground is nearly flat. We have reached a valley -- a minimum. The model has learned. In practice, we stop when the loss stops decreasing meaningfully or when we hit a compute budget.
With millions of parameters, there are many valleys (local minima). Gradient descent finds a minimum, not necessarily the global minimum. Empirical work on over-parameterized neural networks (Choromanska 2015; Dauphin et al. 2014) suggests that in high dimensions, most local minima of large networks are roughly equally good, and saddle points — flat regions where the gradient is near zero but the surface curves down in some direction — are the more common trap. SGD's inherent noise helps escape these. Note this is empirical for over-parameterized DL, not a theorem about general non-convex optimization.