Optimization: gradient descent & SGD

By the end of this page you can derive the gradient-descent update from nothing but the sign of a derivative, predict for any learning rate whether it converges, oscillates or explodes, take the gradient of $\|y-X\theta\|_2^2$ by hand, and explain exactly why sampling one training point at a time still works.

ERM is an optimization problem

Every chapter so far has ended at the same door — problem formulation, decision trees, linear regression. The door is empirical risk minimization:

$\H$ is the hypothesis class, $h$ one candidate from it, $(x^{(i)},y^{(i)})$ the $i$-th training pair (the parenthesised superscript is an index, not a power), $\ell$ the loss, and $\argmin$ returns the argument achieving the smallest value, not the value.

"Efficient" is load-bearing. The previous chapter's closed form $\theta^\star=(X^\top X)^{-1}X^\top y$ costs $O(Nd^2+d^3)$; at $d=10^6$ that cubic term is $10^{18}$ operations, which is not "slow" but "never". Worse, for almost any $\H$ other than a linear model the closed form does not exist at all. So: take the simplest instance of the problem, solve it completely, and generalise the algorithm that falls out.

The simplest possible case: one sample, one parameter

Strip linear regression as far as it goes: $\X=\R$, intercept $b=0$ so $h(x)=xw$ with one knob, and $N=1$ — a single training point $(x^{(1)},y^{(1)})$. ERM becomes

$J$ eats a candidate weight and returns the squared error it incurs; the data are fixed constants, the parameters-are-the-variables view from the previous chapter. (The lecture notes write the weight as $\omega$; same object.) Expanding,

a constant plus a linear term plus a positive multiple of $w^2$: a parabola opening upward, with exactly one bottom and no surprises. It stays a parabola for general $N$, because a sum of upward parabolas in $w$ is another upward parabola in $w$.

The one observation the whole algorithm is built on

Stand at any point $w$ on the parabola and compute the slope $\frac{dJ(w)}{dw}$.

• Everywhere to the right of the minimiser the derivative is positive — the curve climbs as you walk right.
• Everywhere to the left of it the derivative is negative.

That is not a coincidence about this parabola; it is what "the minimum" means, plus the fact that a parabola has only one. And it is actionable, because you can compute a derivative at a single point without knowing where the minimum is.

From the observation to an algorithm

The notes turn it into a procedure: start anywhere; if the derivative is positive move left, if negative move right; repeat until a stopping criterion fires. That answers "which way" and leaves open how far — which is where all the interesting behaviour lives.

A fixed step of $0.1$ would be terrible: start a mile away and you need forever, get within a hair of the bottom and you step straight past it, forever. We want big strides when far away and tiny ones when close — but "how far away am I" is exactly what we do not know. The picture supplies the missing information: out at the edges the tangent is nearly vertical, near the bottom nearly flat.

"Move opposite to the sign" and "move proportionally to the magnitude" are both accomplished by a single subtraction:

A derivative of $+8$ moves us $8\alpha$ left; $-0.2$ moves us $0.02\alpha$ right. The sign handles direction, the magnitude handles step length, and nothing else in the update was designed.

The gradient descent algorithm

Every symbol. The superscript $(t)$ is the iteration index, not a power. $w^{(0)}$ is wherever you start, and for a parabola "any point" genuinely means any point. The vertical bar means "evaluate that derivative function at the current iterate". And $\alpha$ is a positive number we choose.

GradientDescent(J, alpha, w0):
    w <- w0
    for t = 1, 2, 3, ...:
        g <- dJ/dw evaluated at w      # one number
        w <- w - alpha * g             # the whole algorithm
        if stopping_criterion(): break
    return w

Stopping criteria

The loop says "for $t=1,2,\dots$" with no upper limit, so something has to stop it. The notes list four, and they are genuinely different tests:

In practice you combine them: run for at most $T$ iterations, stopping early if the objective has improved by less than $10^{-6}$ for ten iterations in a row.

$\alpha$: the learning rate

$\alpha$ is the learning rate or step size, and the notes give three families of choice:

• Constant. Pick a number and keep it, e.g. $\alpha=0.01$. The default starting point.
• Decaying with the iteration number, for example $\alpha_t = \frac{1}{\sqrt{t}}$: confident strides early, cautious nudges later. The standard choice when gradients are noisy (see SGD), because a shrinking step averages the noise out instead of bouncing forever.
• Adaptive methods that tune the rate per parameter from the gradient history. beyond syllabus AdaGrad, RMSProp, Adam — what people actually run, and out of scope here.

The four learning rates, computed in full

The notes' exercise, left for the poll: let $J(w)=w^2$ and $w^{(0)}=1$. Run three iterations. What is $w^{(3)}$ for (i) $\alpha=0.25$, (ii) $\alpha=0.5$, (iii) $\alpha=1$, (iv) $\alpha=2$?

Set up once. Since $\frac{dJ}{dw} = 2w$, the update $w^{(t)} = w^{(t-1)} - \alpha\cdot 2w^{(t-1)}$ factors:

Each iteration multiplies the iterate by the same contraction factor $1-2\alpha$, so $w^{(t)} = (1-2\alpha)^t w^{(0)}$ and $w^{(3)} = (1-2\alpha)^3$. Everything about the behaviour is in that one number — but grind through the arithmetic anyway, because the shape of the numbers is the lesson.

Read the "step" columns as $w^{(t-1)} - \alpha\cdot 2w^{(t-1)}$, with the derivative $2w^{(t-1)}$ shown in the parentheses. Row four, step 2: the iterate is $-3$, the derivative there is $2(-3)=-6$, so $-3 - 2(-6) = 9$.

What each case is telling you

(i) $\alpha=0.25$ — healthy convergence. Factor $0.5$, so every step halves the distance: $1, \tfrac12, \tfrac14, \tfrac18, \dots$. It never exactly arrives — convergence is geometric — but a stopping criterion fires long before that matters.

(ii) $\alpha=0.5$ — the bullseye. Factor exactly $0$: we land on the minimiser in one step and stay (the gradient at $0$ is $0$, so later steps do nothing). This is possible because for a quadratic the "proportional to the derivative" heuristic is exact: for $J(w)=cw^2$ the derivative at $w$ is $2cw$, so the distance to the minimum is exactly $\frac{1}{2c}$ times the derivative, and $\alpha=\frac{1}{2c}$ makes the step precisely that distance. (Dividing the gradient by the second derivative instead of guessing $\alpha$ is Newton's method beyond syllabus; it hits the same $d^3$ wall as the closed form, because the second derivative is a $d\times d$ matrix.)

(iii) $\alpha=1$ — perpetual oscillation. Factor $-1$: every step flips the sign and preserves the magnitude, $1 \to -1 \to 1 \to -1$ forever. The direction is right every time, but we overshoot by exactly twice the distance to the minimum, landing at the mirror point with the mirror gradient. $J$ never decreases, and neither the improvement criterion nor the $|w^{(t)}-w^{(t-1)}|$ criterion will ever fire.

(iv) $\alpha=2$ — divergence. Factor $-3$: each step flips the sign and triples the magnitude, $1\to-3\to9\to-27\to81$, with $J$ going $1\to 9\to 81\to 729$. In floating point this becomes inf and then NaN within a few dozen iterations; a training loss printing nan is almost always this, in a thousand dimensions.

The convergence condition, derived

Generalise to $J(w)=cw^2$ with $c>0$. Now $\frac{dJ}{dw}=2cw$, the update is $w^{(t)}=w^{(t-1)}(1-2\alpha c)$, and the same argument gives $|1-2\alpha c| < 1$, i.e.

with the midpoint $\alpha=\frac{1}{2c}$ converging in one step. Sharper curvature forces smaller steps — a narrow steep valley must be walked carefully.

Writing the curvature as the second derivative $L = J'' = 2c$, the condition $\alpha < 1/c$ reads

and this is the form that generalises: for any $L$-smooth function (one whose gradient satisfies $\|\nabla J(u)-\nabla J(v)\|\le L\|u-v\|$), gradient descent with $\alpha < 2/L$ is stable and $\alpha = 1/L$ is the standard safe choice. $L$ is the largest curvature anywhere in the landscape — in $d$ dimensions the largest eigenvalue of the Hessian, which for linear regression is the largest eigenvalue of $\frac{2}{N}X^\top X$.

Why this already works for all $N$

Everything so far was derived on a single sample. With $N$ samples (still $\X=\R$, still $b=0$) the objective is the mean squared error

Expand the square inside the sum and collect powers of $w$:

That is $J(w) = cw^2 - 2mw + k$: a parabola again, opening upward because $c$ is a sum of squares. So the single-sample analysis transfers verbatim — same update, same sign argument, same threshold $\alpha < 1/c$ — only now $c$ is the mean squared input.

Two flags. $c$ depends on the scale of your features: multiply every $x^{(i)}$ by $1000$ and $c$ grows by $10^6$, so the largest safe $\alpha$ shrinks by $10^6$ — which is why people standardise features before running gradient descent (see feature engineering). And the previous chapter's corner case reappears: if $\sum_i (x^{(i)})^2 = 0$ then $c=0$, $J$ is flat, the gradient is $0$ everywhere, and gradient descent returns whatever $w^{(0)}$ you typed. The closed form reported that by dividing by zero; both are saying "the data does not determine $w$".

Multivariate calculus: what a gradient is

Real problems have many parameters. Extend to $\X=\R^d$ and a general objective $J:\R^d\to\R$ consuming a parameter vector $\theta$ and returning one number. Before running the algorithm we need to know what replaces $dJ/dw$.

In the scalar case $\frac{\partial J(\theta)}{\partial\theta} \in \R$: one input, one derivative. When $J:\R^d\to\R$ there are $d$ inputs, so $d$ partial derivatives, stacked into a column:

What the gradient means geometrically

Picture $J$ as a landscape, with $\theta=(\theta_1,\theta_2)$ a map coordinate and $J(\theta)$ the altitude, and draw its level sets (contours of constant altitude) as on a topographic map.

• $\nabla J(\theta)$ points in the direction of steepest ascent: of all directions you could step, that one increases $J$ fastest, and its length $\|\nabla J\|$ is how fast.
• Therefore $-\nabla J(\theta)$ is the direction of steepest descent. That is the direction gradient descent walks.
• The gradient is perpendicular to the level set through that point. Walking along a contour does not change altitude, so the rate of change in that direction is zero, so the gradient has no component along it.

The two identities you need — both proved

These two are the only gradient formulas linear regression needs, and both are proved the same way: write out a single coordinate as an ordinary scalar partial derivative, differentiate with high-school rules, then recognise the resulting vector. That technique works for every matrix-calculus identity you will meet.

Shapes, and which layout a formula is written in

More generally you may need $\frac{\partial J(\theta)}{\partial\theta}$ where $\theta$ and the output of $J$ are each a scalar, a vector or a matrix. Do not memorise a table of cases; remember what the notes say:

The content is never in doubt, only the arrangement, and there are two conventions:

• Denominator layout (this course, and the more popular convention in ML): the result is shaped like the denominator, i.e. like $\theta$ — a $d\times1$ column for $J:\R^d\to\R$.
• Numerator layout: shaped like the numerator, i.e. like $J$ — the same object as a $1\times d$ row.

With $d=3$, $J(\theta)=\theta^\top v$ and $v=(4,5,6)^\top$, the three scalar partials are $4,5,6$ in both conventions; only the stacking differs:

How to tell at a glance: check a known identity — a source writing $\nabla_\theta(\theta^\top A\theta) = \theta^\top(A+A^\top)$ (a row) is in numerator layout; our $(A+A^\top)\theta$ (a column) is denominator layout. They differ by a transpose and nothing else, so nobody is wrong — but mixing them mid-derivation produces maddening shape errors.

Back to optimization: linear regression in matrix form

Take $\X=\R^d$, $\Y=\R$, and ERM for linear regression as the previous chapter left it, with samples as rows of $X\in\R^{N\times d}$ and labels in $y\in\R^N$:

The middle step is the definition $\|v\|_2^2 = v^\top v$; the last multiplies out, with the two cross terms $-y^\top X\theta$ and $-\theta^\top X^\top y$ each $1\times1$ — scalars equal to their own transposes, hence the same number — combining into $-2\theta^\top X^\top y$.

$J$ is a bowl, and we want $\nabla_\theta J(\theta)=0$. Differentiate term by term:

• $y^\top y$ contains no $\theta$. Gradient: $\mathbf{0}$.
• $-2\theta^\top X^\top y$ has the form $\theta^\top v$ with $v = -2X^\top y$. By identity (i), gradient: $-2X^\top y$.
• $\theta^\top X^\top X\theta$ has the form $\theta^\top A\theta$ with $A = X^\top X$, symmetric. By identity (ii), gradient: $2X^\top X\theta$.

Setting it to zero:

assuming $X^\top X$ is invertible, which requires $X$ to have full column rank.

Why this is a lucky accident, not the general situation

Setting the gradient to zero produced an equation we could solve exactly because the gradient was linear in $\theta$ — and it was linear because $J$ was quadratic. Most problems are not like this: set the gradient of a neural network's loss to zero and you get millions of coupled nonlinear equations with no solution formula. Gradient descent works there because it never needs to solve $\nabla J = 0$; it only needs to evaluate $\nabla J$ where it stands, and evaluation is always possible.

Local minima, global minima, and getting stuck

Gradient descent may not find the optimal solution unless the function takes some special form (like our bowl). It may get stuck in a local minimum. The notes' definitions, exactly:

Note the asymmetry: "local" quantifies over a neighbourhood, "global" over all of $\R^d$. Every global minimum is local; the converse fails. And the gradient cannot tell the difference — at a local minimum $\nabla J = 0$ and the surface curves upward in every direction, exactly what a global minimum looks like from up close. In the fog, a small dip and the deepest valley feel identical.

So where you start determines where you end. Two starting points on opposite sides of a ridge slide into different basins with different final losses. That is not randomness — plain gradient descent is deterministic — it is sensitivity to the initialisation. Try it In Figure 2, press Function for the multi-minimum landscape and move the start slider: $w^{(0)}=-2.5$ ends in the deep left valley, $w^{(0)}=2.0$ in the shallower right one.

Saddle points and convexity, briefly

Saddle points. beyond syllabus The gradient vanishes, but the point is a minimum along some directions and a maximum along others. Gradient descent crawls near one without being trapped, and eventually escapes downhill. In high dimensions saddles vastly outnumber local minima — a local minimum needs all $d$ curvature directions pointing up, a demanding coincidence when $d$ is a million — so the practical enemy in deep learning is saddles and flat plateaus, not bad local minima.

Convex problems. beyond syllabus A function is convex if the chord between any two points on its graph lies on or above the graph; bowls, parabolas and $\|y-X\theta\|_2^2$ all qualify. A convex function has no bad local minima: every local minimum is global. That is why gradient descent on linear regression converges to the same answer from every start. Logistic regression is also convex; neural networks emphatically are not.

Batch gradient descent, and the problem with it

ERM once more, with the average and a general hypothesis:

where $h_\theta(x) = \theta^\top x$ for linear regression, but $h_\theta$ could be a neural network — all we need is that $\ell$ is differentiable in $\theta$. Because the objective is a sum and differentiation is linear, the gradient is a sum too:

So to take one step the algorithm must compute a gradient for every single training example and average them. Gradient descent applied to ERM is therefore also called batch gradient descent: all $N$ samples per update.

Put numbers on it: at ImageNet scale, $N = 1{,}280{,}000$ images at about $0.5$ ms per forward-and-backward pass, one full gradient costs $1{,}280{,}000 \times 0.5\,\text{ms} \approx 640$ s $\approx 10.7$ minutes — and that buys one update. Ten thousand steps would take 74 days. It is also wasteful: the first thousand images already said roughly which way is downhill, and the remaining 1.279 million only refine that into the exact average, which you then multiply by $\alpha=0.001$ and immediately invalidate by moving. We compute a very precise number and barely use it.

Stochastic gradient descent

The cost per update drops by a factor of $N$. On the ImageNet numbers, an SGD update costs $0.5$ ms instead of $10.7$ minutes, so in the time batch GD makes one step SGD makes $1{,}280{,}000$. As the notes put it: a very noisy version of the full gradient, but fast to compute.

Why does this work at all?

A single random training point is a terrible summary of a dataset, so why doesn't following it just wander? Because the single-sample gradient is an unbiased estimate of the full gradient: wrong on any given step, but right on average, with no systematic lean.

The training-error curves

A note on these two figures. The lecture notes contain two plots — "Figure 1: Training error of gradient descent" and "Figure 2: Training error of (batch) GD and SGD" — whose images did not survive extraction from the PDF. The curves below are reconstructions: real simulations run live in your browser, fitting a linear model to 200 synthetic points with both algorithms, so the shapes are genuine rather than drawn by hand. Press Re-run to draw a fresh random dataset and watch the qualitative story stay the same.

Mini-batch gradient descent

Batch GD uses all $N$ samples per update; SGD uses exactly $1$. The intermediate version is the obvious one, and it is what everyone actually runs.

The tradeoff, quantified

Averaging $K$ independent unbiased estimates gives an estimate that is still unbiased (the average of things each averaging to $\nabla J$ averages to $\nabla J$) but less noisy. Specifically, for independent draws the variance of the mean is

where $\sigma^2$ is the per-sample gradient variance. So:

• Noise falls like $1/K$. The standard deviation — the actual scale of the wobble — falls like $1/\sqrt{K}$.
• Cost rises like $K$. Each iteration does $K$ times the work.

Multiply those together: noise reduction per unit of compute is roughly constant. Going from $K=1$ to $K=100$ makes each step $100\times$ more expensive to make the noise $10\times$ smaller, so in pure flop-counting terms small $K$ wins. Everyone still uses $K=32$ or $K=256$, because the flop count is not the wall clock:

• Hardware parallelism. A GPU computing $256$ sample gradients at once takes barely longer than computing one — the arithmetic units are idle otherwise, so mini-batching is nearly free in real time up to saturation.
• Fixed per-step overhead. Kernel launches, memory transfers and the optimizer update are paid per iteration, not per sample, and dominate at $K=1$.
• Stability. A less noisy gradient tolerates a larger $\alpha$, recovering some of the lost steps.

beyond syllabus $K$ is conventionally a power of two for hardware reasons (warp size, memory coalescing, tensor-core dimensions), not statistical ones. Also, implementations rarely sample with replacement as the definition says: they shuffle once and cut the data into consecutive blocks of $K$, so every sample is used once per pass — one such pass is an epoch.

Recap

Where the rule came from: derivative positive $\Rightarrow$ step left, negative $\Rightarrow$ step right, magnitude large when far away $\Rightarrow$ step proportionally — all three at once in $\theta \leftarrow \theta - \alpha\nabla_\theta J$. And the dial: on $J(w)=cw^2$ it converges iff $0<\alpha<1/c$, is exact in one step at $\alpha = 1/(2c)$, oscillates on the boundary and diverges beyond; in general $\alpha < 2/L$ with $L$ the largest curvature.

• Everything was derived on a parabola, because a one-sample squared loss is a parabola and so is the $N$-sample average — but the algorithm only ever uses a derivative, which is why it transfers to landscapes of any shape.
• Two identities do all the work: $\nabla_\theta(\theta^\top v) = v$ and $\nabla_\theta(\theta^\top A\theta) = (A+A^\top)\theta$, which is $2A\theta$ for symmetric $A$ — and $X^\top X$ is always symmetric. Together they give $\nabla_\theta\|y-X\theta\|_2^2 = -2X^\top y + 2X^\top X\theta$ and hence the normal equations, agreeing with the completing-the-square derivation in the previous chapter.
• Linear regression is the lucky case. Almost nothing else has a closed form; gradient descent is what you use when $\nabla J = 0$ cannot be solved.
• Local minima are real, and gradient descent can get stuck in one depending on where it starts. It remains overwhelmingly popular anyway, because a good solution found efficiently beats a perfect one found never.
• Judge optimizers by error versus compute, never versus iterations.

Next: feature engineering, where the shape of the bowl you are descending turns out to be something you get to choose.