Gradient Descent & Variants
Batch, stochastic & mini-batch descent, convergence rates, and the learning rate.
On this page
Beginner: gradient descent is "repeatedly take a small step downhill." The gradient of a function at a point is the direction of steepest increase — so if you want to go downhill as fast as possible right now, you step in exactly the opposite direction: the negative gradient. Take a small step that way, recompute the direction from your new spot (the downhill direction usually changes a little as you move), take another small step, and repeat. Stop once the steps stop helping — you've reached a point where every direction looks flat or uphill, which is (at least locally) the bottom.
Intermediate: the interesting question is what you compute the gradient from at each step, and the three classical variants answer it differently. Imagine training on one million labeled examples. Batch gradient descent computes the gradient using every one of those million examples before taking a single step — the direction it moves in is as accurate as the data allows, but each step is extremely expensive, since it requires a full pass over the entire dataset just to move once. Stochastic gradient descent (SGD) goes to the opposite extreme: pick one random example, compute the gradient from that example alone, and step immediately. Any single step's direction is a noisy, often quite wrong, estimate of the true downhill direction — but steps are nearly free, so you can take a million of them in the time batch GD takes to take one. Mini-batch gradient descent is the practical middle ground: compute the gradient from a small random subset — commonly somewhere between 32 and 512 examples — and step. This is what essentially every real training pipeline actually runs; "batch" and "pure SGD" are the two theoretical extremes that mini-batch sits between.
Advanced: the size of the step, the learning rate (written α), has to be tuned regardless of which variant you use, and it's a genuine trade-off in both directions. Too large a learning rate and each step overshoots the bottom — the iterate can oscillate back and forth across the minimum, or in the worst case diverge outright, with the loss growing instead of shrinking. Too small a learning rate and every individual step is safe, but convergence crawls: you're leaving performance on the table by taking tiny, needlessly cautious steps when larger ones would still have been safe. The classical fix is a learning-rate schedule — start with a larger α for fast early progress, then decay it over time (step decay, exponential decay, cosine decay) so later steps get more cautious as you approach the optimum. There's also a subtler point worth sitting with: SGD's noise, which is a pure cost in the convex setting this lesson is built around — noisier steps just mean slower, less reliable convergence to the one global minimum — can actually become a genuine benefit once the loss surface is non-convex, as it is for essentially all deep neural networks. A noisy step can kick the iterate out of a shallow local minimum or off a saddle point that a noise-free batch step would have gotten stuck at or slowed to a crawl near. The previous lesson's whole point about convexity — that it's what makes "reached a flat point" mean "reached the global optimum" — is exactly why this trade-off flips: in the convex world noise has nothing useful to escape from, so it's only ever a tax on convergence speed.
Every variant uses this exact same update rule — what changes is what R̂, the empirical risk being minimized at step k, is computed over:
Batch sums over all n examples every step; SGD sums over exactly 1, freshly and uniformly sampled each step; mini-batch sums over a freshly sampled subset B with |B| typically in the dozens to low hundreds — strictly between the other two. As |B| grows from 1 toward n, mini-batch GD interpolates continuously between pure SGD and full batch GD; it is not a fourth, separate algorithm.
Take gradient descent on a single deterministic, convex function f (this is the batch-GD case with no sampling noise; the previous lesson on convex optimization is the reason a "global minimum" is even a meaningful, reachable target here). Assume f is L-smooth: its gradient doesn't change arbitrarily fast, formally
A standard consequence of L-smoothness (stated here rather than re-derived, since proving it is a short calculus exercise orthogonal to the point of this derivation) is a quadratic upper bound on f around any point:
Now substitute in the actual gradient descent update, y = x_{k+1} = x_k - α ∇f(x_k) and x = x_k, so y - x = -α∇f(x_k):
The first extra term is −α‖∇f(x_k)‖² exactly (a dot product of a vector with itself, times −α), and the last term is (Lα²/2)‖∇f(x_k)‖² (the squared norm pulls α² out and leaves the same ‖∇f(x_k)‖² factor). Both terms share that factor, so they combine directly:
That last line is the classical descent lemma. Look at the coefficient (1 − Lα/2): as long as the step size satisfies α ≤ 1/L, that coefficient is at least 1/2, which is positive — and since α and ‖∇f(x_k)‖² are both non-negative, the entire subtracted term is non-negative. That means:
— the function value is guaranteed non-increasing every single step. This is exactly the mechanism behind the diagram below: its toy surface has L = 5 (the curvature of its steep direction), and monotonic descent along that direction is guaranteed only up to α ≤ 1/5 = 0.2; push past α = 2/L = 0.4 and the coefficient (1 − Lα/2) turns negative, the bound stops guaranteeing descent, and the diagram visibly diverges.
Chaining this per-step inequality over many iterations (cited here rather than re-derived term by term, since it's a standard telescoping-sum argument once the descent lemma is established) is what produces the two classical headline rates. For f merely convex and L-smooth, running GD with α = 1/L gives
— error shrinks like O(1/k): sublinear, and to halve the remaining error you roughly have to double the number of iterations. That's the entire guarantee plain convexity buys you: eventual convergence to the global minimum, at this specific, fairly slow rate. Under the strictly stronger assumption of μ-strong convexity — precisely the extra condition the previous lesson flagged as the thing that upgrades "eventually converges" into an actual provable rate — the same algorithm instead gets
which shrinks geometrically in k. Solving for how many iterations are needed to reach a target error ε gives O(log(1/ε)) iterations — a linear (in the optimization sense) rate, exponentially faster than the merely-convex O(1/k) case. The ratio κ = L/μ, the condition number, controls exactly how fast: a well-conditioned bowl (κ close to 1) converges quickly, while an ill-conditioned, elongated bowl like the one in the diagram below (large κ) converges slowly even though the rate is still technically geometric.
Where this is used: the α ≤ 1/L bound derived above is precisely why learning-rate tuning is not a minor implementation detail — get α wrong relative to the (usually unknown) smoothness constant L of your actual loss landscape, and you lose either the descent guarantee (too large) or most of your convergence speed (too small), exactly as the diagram below demonstrates directly. It's also exactly why adaptive optimizers — the subject of the next lesson, 2.2.3 — exist at all: methods like AdaGrad, RMSProp, and Adam adjust per-parameter step sizes on the fly using observed gradient statistics, so that in practice you never have to hand-estimate a global L and pick one fixed α ≤ 1/L for an entire, possibly highly ill-conditioned, high-dimensional loss surface.
Drag the learning-rate slider for an instant preview of where each path ends up, then hit Replay to watch the animated race — batch (blue) takes the exact gradient, SGD (amber) visibly wanders on a single noisy example per step, and mini-batch (green) tracks close to batch with only light jitter.
All three implementations run the identical experiment — the same synthetic data, the same hand-derived squared-error gradient — differing only in how many examples each gradient estimate is averaged over per update. The printed epoch counts and final losses make the speed/noise trade-off concrete rather than just asserted: batch GD needs the fewest epochs but only ever performs one parameter update per epoch, SGD performs n noisy updates per epoch, and mini-batch sits in between on both counts.
import random
random.seed(0)
# ---- synthetic linear regression data: y = 3x + 5 + noise ----
n = 200
xs = [random.uniform(-5, 5) for _ in range(n)]
true_w, true_b = 3.0, 5.0
ys = [true_w * x + true_b + random.gauss(0, 0.6) for x in xs]
def mse(w, b, xs, ys):
total = 0.0
for x, y in zip(xs, ys):
e = w * x + b - y
total += e * e
return total / len(xs)
def grad_on_indices(w, b, xs, ys, idxs):
# Gradient of mean squared error, averaged over exactly the given indices --
# pass all indices for batch GD, one index for SGD, a handful for mini-batch.
dw = 0.0
db = 0.0
for i in idxs:
err = w * xs[i] + b - ys[i]
dw += 2 * err * xs[i]
db += 2 * err
m = len(idxs)
return dw / m, db / m
TOL = 0.40 # stop once full-dataset MSE drops below this
MAX_EPOCHS = 400
def batch_gd(alpha):
w, b = 0.0, 0.0
all_idx = list(range(n))
for epoch in range(1, MAX_EPOCHS + 1):
dw, db = grad_on_indices(w, b, xs, ys, all_idx) # every example, one update
w -= alpha * dw
b -= alpha * db
loss = mse(w, b, xs, ys)
if loss < TOL:
return epoch, loss
return MAX_EPOCHS, mse(w, b, xs, ys)
def stochastic_gd(alpha):
w, b = 0.0, 0.0
order = list(range(n))
for epoch in range(1, MAX_EPOCHS + 1):
random.shuffle(order) # reshuffle every epoch -- see the Pitfalls section
for i in order:
dw, db = grad_on_indices(w, b, xs, ys, [i]) # one example, one update
w -= alpha * dw
b -= alpha * db
loss = mse(w, b, xs, ys)
if loss < TOL:
return epoch, loss
return MAX_EPOCHS, mse(w, b, xs, ys)
def minibatch_gd(alpha, batch_size):
w, b = 0.0, 0.0
order = list(range(n))
for epoch in range(1, MAX_EPOCHS + 1):
random.shuffle(order)
for start in range(0, n, batch_size):
idxs = order[start:start + batch_size] # ~32 examples, one update
dw, db = grad_on_indices(w, b, xs, ys, idxs)
w -= alpha * dw
b -= alpha * db
loss = mse(w, b, xs, ys)
if loss < TOL:
return epoch, loss
return MAX_EPOCHS, mse(w, b, xs, ys)
b_epochs, b_loss = batch_gd(alpha=0.05)
s_epochs, s_loss = stochastic_gd(alpha=0.01)
m_epochs, m_loss = minibatch_gd(alpha=0.02, batch_size=32)
print(f"Batch GD: {b_epochs:4d} epochs (1 update/epoch), final MSE = {b_loss:.4f}")
print(f"SGD: {s_epochs:4d} epochs ({n} updates/epoch), final MSE = {s_loss:.4f}")
print(f"Mini-batch GD: {m_epochs:4d} epochs ({-(-n // 32)} updates/epoch), final MSE = {m_loss:.4f}")
# Same tolerance, three very different numbers of parameter updates to get there --
# that gap is exactly the batch-size / update-count trade-off this lesson is about.- Virtually all deep learning training uses mini-batch SGD, not either pure extreme. Modern GPUs are throughput machines — computing a gradient over 128 examples costs barely more wall-clock time than computing it over 1, because the examples are processed in parallel — so mini-batches land in a genuine sweet spot: cheap enough per step to take many steps, averaged enough to keep the gradient noise from dominating, and large enough to actually saturate the hardware's parallelism.
- Learning-rate warmup and decay schedules are standard in large language model training: start with a small learning rate and ramp it up over the first few thousand steps (warmup, to avoid destabilizing randomly-initialized weights with large early updates), hold or slowly decay it through the bulk of training, then decay it further (often following a cosine curve) toward the end — the exact "classical fix" the plain-English section above describes, at a scale where getting the schedule wrong can waste weeks of compute.
- Online learning systems — ad click-through-rate prediction being the canonical example — are a natural fit for pure SGD rather than mini-batch. There is no fixed dataset to batch over: each new impression-and-click (or non-click) arrives as a single fresh example, the model takes one SGD step on it, and moves on. Waiting to accumulate a "batch" would mean either delaying model updates or artificially buffering a stream that's naturally one-example-at-a-time.
- Batch gradient descent is still fine, even preferred, for small datasets that fit comfortably in memory — a few thousand rows of a classical regression or classification problem, say. When computing the exact gradient over the whole dataset costs almost nothing, there's no speed advantage to sampling, and batch GD's smooth, noise-free convergence is simply the better-behaved choice.
- An "epoch" is one full pass of mini-batch (or pure SGD) updates over the training set — split the data into mini-batches, run through all of them once, and that's epoch 1; reshuffle and repeat for epoch 2, and so on. Training for "50 epochs" is shorthand for "50 full passes," each made up of many individual mini-batch steps — exactly the loop structure the code examples above implement directly.
- Large-scale distributed training pushes batch size up into the thousands or tens of thousands (spreading one giant batch across many GPUs simultaneously), and empirically also needs the learning rate scaled up to match — the "linear scaling rule" used at large tech companies training on big GPU clusters is a direct, practical acknowledgment that batch size and learning rate are coupled knobs, not independent ones.
- Picking a learning rate by intuition or leaving a framework default untouched, then being surprised when training diverges (rate too large) or crawls for hours with barely moving loss (rate too small). The Derivation above isn't decorative — the stable range for α genuinely depends on the specific loss surface's smoothness, and "the default worked on a different problem" is not evidence it will work here.
- Assuming a bigger batch size is unconditionally better because it means a less noisy gradient estimate. It's a real trade-off, not a free lunch: a larger batch does reduce per-step noise, but each step also costs more compute, and past a certain size the reduction in noise stops meaningfully improving either convergence speed or the final model's quality — you're just paying more per step for a gradient estimate that was already accurate enough.
- Forgetting to reshuffle the data between epochs in SGD or mini-batch training. If the data has any ordering structure — sorted by label, grouped by time, grouped by source — training in the same fixed order every epoch means every early step of every epoch sees a skewed slice of the data, quietly biasing the whole training trajectory in a way that's easy to miss and annoying to diagnose after the fact.
Going deeper
SGD's gradient noise, which the plain-English section already flagged as sometimes helpful in non-convex landscapes, has a more specific informal story behind it: it appears to act as an implicit regularizer. Rather than just randomly kicking the iterate around, the noise seems to bias SGD's trajectory away from sharp, narrow minima (where a small perturbation in the parameters causes a large jump in loss) and toward flatter, wider ones — and flat minima have been empirically linked to better generalization than sharp ones reaching the same training loss. This connects directly back to the double-descent and generalization discussion from Module 1: there, the surprising finding was that heavily over-parameterized models often generalize far better than the classical bias-variance picture predicts, and part of the informal explanation offered there was that gradient-based optimizers tend to land on comparatively simple, well-behaved solutions among the many that fit the training data — SGD's noise is one of the candidate mechanisms proposed for exactly why that happens, not a separate phenomenon.
Treat this as an active, evolving research area rather than settled fact. The flat-minima-generalize-better story has real empirical support and some theoretical backing in simplified settings, but there are also known counterexamples and open debates about how universally it holds across architectures and tasks. What is settled is the more basic point this lesson leans on throughout: a property that looks purely like a weakness through a convex-optimization lens (noisier steps, worse per-step guarantees) is not automatically a weakness once the assumptions — convexity, in this case — no longer hold.
You're training on 5 million examples. Switching from batch GD to mini-batch GD with a batch size of 256 changes the number of parameter updates per epoch from 1 to roughly 19,500. Given that each individual mini-batch step is noisier than the one exact batch step, why does training with mini-batches typically still reach a good model faster in wall-clock time?
Because 'noisier per step' and 'slower overall' are not the same thing. Batch GD's single update per epoch is exact but you only get to move once per full pass over 5 million examples — an enormous amount of computation spent on one step. Mini-batch GD spends roughly the same total amount of computation per epoch, but spreads it across ~19,500 much cheaper steps, each computed from a batch small enough for a GPU to process in parallel almost as fast as a single example. Even though every individual mini-batch step is a noisier estimate of the true gradient than the one batch step, taking thousands of noisy-but-cheap steps per epoch converges faster in wall-clock time than taking one expensive-but-exact step, because the model gets to update its parameters far more often for roughly the same total compute.
Batch, stochastic, and mini-batch gradient descent are the same update rule — step opposite the gradient of the empirical risk — differing only in how many examples that risk is averaged over per step, trading exactness for speed as you move from batch toward SGD. Mini-batch, sitting between the extremes, is what nearly all practical training actually runs. The learning rate governs a separate but equally important trade-off, provably bounded by a step-size threshold set by the loss surface's smoothness (α ≤ 1/L) for guaranteed descent, with strong convexity the extra ingredient that upgrades a merely-eventual convergence guarantee into a genuinely fast one — and it's precisely the difficulty of knowing that threshold in practice that motivates the adaptive optimizers coming up next.