Gradient Descent & Optimization

A learning algorithm is only as good as the optimiser that drives it. Given a cost function that measures how wrong the model is on the training set, the optimiser's job is to find the parameters that minimise that cost. This lesson is the engine room of that search: how the gradient tells the optimiser which way is downhill, why three flavours of "how much data per step" give different costs and different noise, how a learning rate that is too high blows up and a rate that is too low stalls, why momentum, RMSProp and Adam are the default choices in modern deep learning, and what convex and non-convex loss surfaces look like and why high-dimensional landscapes are usually benign for stochastic optimisers.

Learning Objectives

1. Write the batch, stochastic and mini-batch gradient descent update rules, explain the cost-per-update and noise trade-off between them, and identify the regime where mini-batch is the practical default.
2. Diagnose the symptoms of a learning rate that is too high (divergence, oscillation, NaN loss) and one that is too low (stalling, plateau, wasted compute), with concrete numeric traces.
3. Derive momentum as a velocity term with the ball-rolling analogy, write its update rule, and explain why it dampens oscillation in ravines and accelerates along consistent gradients.
4. Derive RMSProp as an exponential moving average of squared gradients and explain how its per-parameter scaling adapts the effective step size to each parameter's gradient magnitude.
5. State the Adam update rule as momentum on the first moment combined with RMSProp on the second moment, write the bias-correction terms, and explain why Adam is the most common default optimiser.
6. Distinguish convex and non-convex loss surfaces, sketch the local-minima intuition, explain saddle points, and argue why high-dimensional non-convex landscapes are mostly benign for stochastic optimisers.
7. Apply a learning-rate schedule (step decay, exponential decay, cosine annealing, warmup) and read a loss curve to diagnose whether the optimiser is converging, stalling, or oscillating.

1. The Optimisation Problem

Given a dataset with $n$ examples and a parameterised model $f_\theta(x)$ with parameters $\theta \in \mathbb{R}^d$, training means solving

where $J(\theta)$ is a loss function that measures how badly $f_\theta$ fits the data. The most common loss in regression is mean-squared error,

and in classification the cross-entropy that pairs with logistic regression or a softmax output.

For many models — linear and logistic regression with convex losses — there is a closed-form solution. For everything else (kernel SVMs, neural networks, anything non-linear) the cost has no algebraic solution and you must search for the minimum with an iterative optimiser.

The workhorse of that search is gradient descent: at every step, compute the gradient of the loss with respect to the parameters, and move the parameters a small distance in the negative-gradient direction. The intuition is that the gradient points uphill, so its negation points downhill, so a small step in that direction reduces the loss.

where $\eta > 0$ is the learning rate (step size). This single rule, with embellishments, is what every optimiser in the lesson does.

2. Batch, Stochastic, and Mini-Batch Gradient Descent

The rule above has a hidden choice: which examples do you average over to form $\nabla_\theta J$? Three answers give three different optimisers.

2.1 Batch gradient descent

Use every training example to form the gradient at each step:

This is the textbook definition. The gradient is exact, so each step reliably decreases the loss (with a small enough $\eta$) and the trajectory is smooth. The cost is one full pass over the data per update — if $n = 1{,}200{,}000$ ImageNet images, a single step requires $1.2$ million forward and backward passes.

Batch GD works for small datasets (a few thousand rows on a single CPU) but becomes impractical for anything modern. It also cannot escape the I/O wall: each update is gated by reading every example from disk.

2.2 Stochastic gradient descent (SGD)

Use one example at a time:

where $i_t$ is drawn uniformly at random from $\{1, \ldots, n\}$. The gradient is now a noisy estimate of the true direction, but each update is cheap — one example in, one example out, one gradient step. One epoch (one pass over the data) produces $n$ updates.

The noise has two faces. On the bad side, the trajectory zig-zags and never settles exactly at the minimum — the parameters jitter around it. On the good side, the noise lets the optimiser escape shallow local minima and saddle points, and the cheap updates mean you can take vastly more of them.

2.3 Mini-batch gradient descent

Compromise: use a batch of $B$ examples per step,

where $\mathcal{B}_t$ is a randomly drawn batch of size $B$ (commonly $32$, $64$, $128$, $256$). The gradient is a less-noisy estimator than the one-example SGD version, but each step is still cheap. One epoch produces $n / B$ updates.

Mini-batch is the practical default. It maps cleanly onto GPU parallelism (a batch of $64$–$512$ examples saturates a GPU's matrix-multiply units), and the noise from random sampling acts as a regulariser that improves generalisation. Modern deep-learning frameworks call this "SGD" by default; the textbook one-example SGD is rarely used.

2.4 Comparison

For $n = 1{,}200{,}000$ and $B = 128$, one epoch produces $9{,}375$ updates — each using only about $0.01\%$ of the data. That is roughly $9{,}000\times$ more updates than batch GD per epoch, at slightly higher per-update noise. In practice mini-batch converges in a tiny fraction of the wall-clock time batch GD needs.

2.5 Implementation

import numpy as np

def mini_batch_sgd(X, y, lr=0.01, batch_size=32, epochs=10):
    """Plain mini-batch gradient descent on MSE loss with a linear model.

    X: (n, d) feature matrix; y: (n,) target vector; theta starts at zeros.
    Returns the parameter trajectory and the per-epoch average loss.
    """
    n, d = X.shape
    theta = np.zeros(d)
    losses = []
    for epoch in range(epochs):
        # shuffle indices once per epoch so each example is seen once
        perm = np.random.permutation(n)
        epoch_loss = 0.0
        for start in range(0, n, batch_size):
            batch = perm[start:start + batch_size]
            Xb, yb = X[batch], y[batch]
            # MSE gradient for linear model y_hat = Xb @ theta
            grad = (2.0 / len(batch)) * Xb.T @ (Xb @ theta - yb)
            theta -= lr * grad            # the update rule
            epoch_loss += np.mean((Xb @ theta - yb) ** 2) * len(batch)
        losses.append(epoch_loss / n)
    return theta, losses

The body of the inner loop is the literal equation $\theta_{t+1} = \theta_t - \eta \nabla J$. Everything else — shuffling, batching, averaging — is plumbing.

3. The Learning Rate: Too High, Too Low

The single most important hyperparameter in optimisation is the learning rate $\eta$. Choose it badly and nothing else matters.

3.1 Too high: divergence

If $\eta$ is larger than twice the curvature of the loss along the gradient direction, each step overshoots the minimum and the next gradient points back the other way. The sequence of parameters oscillates with growing amplitude, and the loss curve explodes.

Concretely, suppose the loss is locally quadratic near some $\theta$ with curvature $L$ (the Lipschitz constant of the gradient, or the largest eigenvalue of the Hessian). The step $\eta \nabla J$ followed by re-evaluation of the gradient has magnitude $\eta L$ in the steepest direction. The step is contractive when

equivalently $\eta < 2 / L$. Anything larger and the iteration is unstable.

A worked trace on a one-dimensional quadratic $J(\theta) = \tfrac{1}{2} \theta^2$ (so $L = 1$) with $\theta_0 = 1$:

With $\eta = 0.1$ the loss decays cleanly; with $\eta = 1.0$ the iteration snaps to the minimum in one step (the boundary case); with $\eta = 1.5$ the iteration oscillates and decays slowly. With $\eta = 2.1$ it diverges.

In practice the visible symptoms of "too high" are:

• The training loss grows epoch by epoch instead of shrinking.
• The loss becomes NaN (overflow in the forward pass, often within the first few hundred steps).
• Validation loss bounces wildly from epoch to epoch.
• Gradient norms grow monotonically across the first few epochs.

3.2 Too low: stalling

If $\eta$ is too small, the optimiser moves, but glacially. Symptoms:

• Training loss decreases by $0.1\%$ per epoch instead of $10\%$ per epoch.
• After $k$ hours of GPU time the model is barely better than random.
• Gradient norms look healthy — the optimiser is just taking tiny steps relative to where it needs to go.
• Convergence is eventually reached, but you used $10\times$ the wall-clock time you could have.

3.3 The right ballpark

A practical recipe is to start with $\eta = 0.01$ for SGD on a normalised dataset, or follow the optimiser's default (Adam typically wants $10^{-3}$, RMSProp $10^{-3}$, SGD with momentum $10^{-2}$ to $10^{-1}$). Run for $5$–$10$ epochs and watch the loss; if it explodes, divide $\eta$ by $10$ and restart; if it crawls, multiply by $3$–$10$ and restart.

A learning-rate finder (cyclically sweep $\eta$ from $10^{-7}$ to $10^{0}$ for one epoch and plot loss vs $\eta$) gives a direct read of where the loss starts climbing — set $\eta$ a factor of $3$–$10$ below that cliff.

4. Momentum: The Ball-Rolling Analogy

Plain gradient descent wastes motion on inconsistent gradients. Imagine the loss surface has a long shallow valley: the gradient always points slightly across the valley (steepest direction is perpendicular to the valley axis), so plain GD zig-zags across the valley floor and crawls along it.

Momentum borrows the physics of a heavy ball rolling downhill. The ball accumulates velocity, not just position, so consistent downhill pushes build up speed, while sideways pushes cancel out.

4.1 The update rule

where $v_t$ is the velocity vector, $\beta \in [0, 1)$ is the momentum coefficient (typically $0.9$), and the gradient is added (not subtracted) into the velocity. Some texts flip the sign convention by defining the velocity as a moving average of the negative gradient; the algebra is the same up to a sign.

When $\beta = 0$ the rule collapses to plain GD. When $\beta = 0.9$ the optimiser keeps roughly the last $1 / (1 - \beta) = 10$ gradients in memory, exponentially weighted — recent steps count more, older steps count less, but the influence decays smoothly rather than cutting off.

4.2 Why it helps

Two effects:

1. Acceleration along consistent gradients. If the gradient points roughly the same direction for many steps, the velocity accumulates and the effective step size becomes $\eta / (1 - \beta)$ — about $10\times$ larger than plain GD with the same $\eta$ at $\beta = 0.9$.
2. Dampening of oscillation. In the ravine example, the gradients perpendicular to the valley axis cancel from step to step, while the gradients along the valley axis reinforce. The velocity along the valley axis grows; the velocity perpendicular stays near zero.

The ball-rolling analogy is precise: in the absence of friction, a ball on a frictionless slope accelerates without bound. In optimisation, the $\eta$ in front of $v_{t+1}$ acts as the "time step", and $\beta$ as a friction coefficient (smaller $\beta$ = more friction = faster decay of velocity).

4.3 Nesterov momentum

A small but useful variant: instead of computing the gradient at $\theta_t$, compute it at $\theta_t - \eta \beta v_t$ — the position you are about to step to. This "look-ahead" gradient gives an earlier correction when the optimiser is about to overshoot. Nesterov momentum converges faster than classical momentum on convex losses.

5. RMSProp: Adaptive Per-Parameter Scaling

Momentum handles the case where the direction of the gradient is noisy. RMSProp handles the case where the magnitude of the gradient varies wildly across parameters.

Consider a neural network where one layer's weights need updates $1000\times$ smaller than another layer's because the gradient magnitudes are $1000\times$ different. A single global $\eta$ either blows up the first layer or starves the second.

5.1 The update rule

Maintain an exponential moving average of the squared gradient, then divide the gradient by its root:

where $\gamma \in [0, 1)$ is the decay rate (typically $0.99$), the squaring is element-wise, and $\epsilon \approx 10^{-8}$ prevents division by zero.

The squared-gradient average $s_t$ estimates the recent magnitude of $\nabla J$ per parameter. Dividing by $\sqrt{s_t}$ rescales the step so that parameters with consistently large gradients take smaller relative steps, and parameters with consistently small gradients take larger relative steps. The effective per-parameter step is roughly $\eta$ regardless of the gradient scale.

5.2 Why it helps

If parameter $j$ has gradient magnitude $g$ consistently, then $s_t \to g^2$ and the update is $\eta g / \sqrt{g^2 + \epsilon} \approx \eta$ — a unit-sized step. If parameter $j$ has gradient magnitude $10 g$ consistently, the update is also $\approx \eta$. The optimiser self-calibrates per parameter.

RMSProp also dampens the zig-zag pattern in ravines. The "downhill" direction has high gradient magnitude and takes a normal step; the "sideways" directions have oscillating gradients and the squared average smooths them out.

5.3 Implementation

def rmsprop_step(theta, grad, s, lr=1e-3, gamma=0.99, eps=1e-8):
    """One RMSProp update; returns new theta and new squared-gradient state.

    theta, grad, s are 1-D arrays of identical shape.
    """
    s_new = gamma * s + (1 - gamma) * (grad ** 2)   # EMA of squared gradient
    theta_new = theta - lr * grad / (np.sqrt(s_new) + eps)
    return theta_new, s_new

The signature change from plain SGD is the extra state s that the optimiser carries between steps.

6. Adam: Combining Momentum and RMSProp

Adam (Kingma and Ba, 2014) is "momentum on the first moment + RMSProp on the second moment", plus bias correction for the early steps where the moving averages are still warming up.

6.1 The update rule

where $\beta_1$ controls the momentum (default $0.9$) and $\beta_2$ controls the RMSProp decay (default $0.999$). The $\hat{m}$ and $\hat{s}$ are the bias-corrected versions of $m$ and $s$.

6.2 Why bias correction matters

The moving averages are initialised to zero, so at step $t = 1$:

This is much smaller than the true gradient — the moving average has not "warmed up". Without correction, Adam's first few updates would be roughly $10\times$ smaller than intended, and the optimiser would stall at the start.

The bias-correction terms $\hat{m} = m / (1 - \beta_1^t)$ and $\hat{s} = s / (1 - \beta_2^t)$ undo exactly this: at $t = 1$ they scale the estimates up by $1 / (1 - \beta_1)$ and $1 / (1 - \beta_2)$; as $t \to \infty$ the correction factor $\to 1$ and the bias vanishes.

6.3 Why Adam is the default

Adam combines three independent wins:

1. Momentum ($m_t$) smooths noisy gradients and accelerates consistent directions.
2. Adaptive scaling ($s_t$) normalises step size per parameter.
3. Bias correction ($\hat{m}$, $\hat{s}$) gives well-behaved early steps.

The result is an optimiser that converges quickly on a wide range of learning rates, works on noisy gradients (mini-batch), works on sparse gradients (NLP embeddings), and rarely needs careful tuning. The default hyperparameters $\eta = 10^{-3}$, $\beta_1 = 0.9$, $\beta_2 = 0.999$ work on most problems.

The trade-off is that Adam's generalisation can be slightly worse than well-tuned SGD-with-momentum on some computer-vision benchmarks, and recent work (e.g. AdamW, decoupled weight decay) refines the weight-decay interaction. For most practitioners, Adam is still the right first choice.

6.4 Update-rule summary

Reading this table left-to-right is a one-sentence history of optimisation: plain GD sets the template, momentum adds inertia, RMSProp adds per-parameter adaptation, Adam combines both with bias correction.

7. Convex vs Non-Convex Loss Surfaces

The geometry of the loss function determines what optimisation can and cannot guarantee.

7.1 Convex losses

A loss $J(\theta)$ is convex if for any two points $\theta_a, \theta_b$ and any $\alpha \in [0, 1]$:

Geometrically, the chord between any two points on the loss surface lies above the surface. Equivalently, the surface has a single bowl shape with one minimum, and any local minimum is the global minimum.

Linear regression with MSE loss is convex. Logistic regression with cross-entropy loss is convex. SVM hinge loss is convex. For these models, plain gradient descent from any starting point converges to the global optimum (with a small enough $\eta$, or with the right step-size schedule).

7.2 Non-convex losses

A loss is non-convex when the surface has multiple basins. Neural networks with one or more hidden layers and most non-linear models are non-convex. The surface can have:

• Multiple local minima — basins of different depth.
• Multiple saddle points — flat regions where the gradient is zero but the curvature has both positive and negative eigenvalues.
• Plateaus — wide regions of near-constant loss where the gradient is tiny.

For non-convex losses, gradient descent is not guaranteed to find the global minimum. It will converge to some local minimum (or a saddle point, in pathological cases), and which one depends on the initialisation and the noise.

7.3 Why this matters in practice

Two observations balance each other:

1. In principle, non-convex optimisation is hard. There is no guarantee that the local minimum the optimiser finds is a good one.
2. In practice, deep networks are trained successfully every day. The reason is that the local minima that stochastic optimisation finds in modern architectures are usually good enough — often nearly as good as the global minimum.

The "local minima are not a problem" claim is empirical and architecture-dependent. The more recent and more precise statement is that saddle points are the dominant obstacle in high dimensions, not local minima. Section 8 explains why.

8. Saddle Points and Why High Dimensions Are Benign

A saddle point is a stationary point of $J$ where the Hessian has both positive and negative eigenvalues. The surface curves up in some directions and down in others — like a horse saddle, or a Pringle chip. At a saddle, the gradient is zero, so plain gradient descent stalls.

8.1 Local minima vs saddle points

In one dimension, every stationary point is either a local min or a local max. In two dimensions, a saddle point is one additional case. In $d$ dimensions, the number of possible curvature signatures grows combinatorially: the Hessian has $d$ eigenvalues, and each can be positive, negative, or zero. A stationary point is a local minimum only when all eigenvalues are non-negative; it is a saddle point whenever the eigenvalues have mixed signs.

For a random point on a random smooth surface in $d$ dimensions, the probability that all eigenvalues have the same sign drops exponentially with $d$. Most stationary points of a high-dimensional loss are saddles, not local minima.

8.2 Why this is good news for SGD

SGD has a property that plain GD does not: gradient noise. A noisy gradient at a saddle point typically has a non-zero component along a direction of negative curvature, which pushes the parameters off the saddle. Empirically, SGD escapes saddle points within a few hundred steps in typical deep-learning settings.

The intuition is that in a $10{,}000{,}000$-parameter network, the loss surface is almost everywhere a saddle in some directions and a downhill in others. The optimiser, slightly noisy, naturally follows the downhill directions and ignores the saddle ones. It does not need to "find" a global minimum — it needs to keep descending, and there are always directions to descend in.

8.3 Practical implications

• Don't panic at non-monotonic loss curves. A spike in training loss (followed by continued decrease) is often the optimiser escaping a saddle or a sharp region — not a sign of failure.
• Don't trust a single training run. Different random seeds explore different basins; train $3$–$5$ times with different seeds and report the spread.
• Use the noise of mini-batch SGD. Full-batch GD on a non-convex loss can sit at a saddle indefinitely; mini-batch SGD's gradient noise is the mechanism that escapes it.

9. Learning-Rate Schedules and Convergence Diagnosis

A constant learning rate is rarely optimal. The most common refinement is to anneal $\eta$ over training: large steps early to make progress, smaller steps late to settle into a good minimum.

9.1 Common schedules

• Step decay. Multiply $\eta$ by a factor (e.g. $0.1$) at fixed epochs. ResNet-style: $\eta_0 = 0.1$, divide by $10$ at epochs $30$ and $60$.
• Exponential decay. $\eta_t = \eta_0 e^{-\lambda t}$ for some decay rate $\lambda > 0$.
• Cosine annealing. $\eta_t = \eta_{\min} + \tfrac{1}{2}(\eta_{\max} - \eta_{\min})(1 + \cos(\pi t / T))$, which smoothly drops $\eta$ from $\eta_{\max}$ to $\eta_{\min}$ over $T$ steps.
• Warmup. Start $\eta$ near zero and ramp it up over the first $1{,}000$–$10{,}000$ steps. Used with large-batch training and Transformers because the early Adam updates are dominated by the bias-correction terms and benefit from a cautious start.
• Cyclical learning rates. Swing $\eta$ between $\eta_{\min}$ and $\eta_{\max}$ periodically; sometimes helps escape saddle points but rarely the default.

9.2 Reading the loss curve

A loss curve is the diagnostic. Three canonical shapes:

1. Decreasing then flat at a low value — the optimiser converged. Stop training and report test metrics.
2. Decreasing then flat at a high value — the optimiser is stuck. Likely causes: learning rate too low, model too small, or the loss has plateaued at a poor local minimum. Try lowering $\eta$, switching to Adam, or adding capacity.
3. Decreasing then increasing — overfitting (train loss down, val loss up) or learning rate too high. Reduce $\eta$, add regularisation, or stop earlier.

A loss curve that never goes down from the start is a learning-rate-too-high symptom. A loss curve that goes down for a few epochs and then becomes perfectly flat is a learning-rate-too-low symptom or a saddle-point stall.

9.3 Implementation

def cosine_anneal(step, total_steps, eta_max=1e-3, eta_min=1e-5):
    """Cosine learning-rate schedule; call once per step."""
    return eta_min + 0.5 * (eta_max - eta_min) * (1 + np.cos(np.pi * step / total_steps))

Pair the schedule with the optimiser:

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer,
    lr_lambda=lambda step: cosine_anneal(step, total_steps=10000) / 1e-3,
)
for step in range(10000):
    # ... forward, backward, optimizer.step() ...
    scheduler.step()    # advance the schedule

10. Putting It All Together

A practical recipe for "I have a model and a dataset, I want to train it":

1. Normalise the inputs. Zero-mean, unit-variance per feature. Without this, the loss surface is badly scaled and learning rates that work on one feature destroy another.
2. Pick the optimiser. Adam with default $\eta = 10^{-3}$, $\beta_1 = 0.9$, $\beta_2 = 0.999$ is the right first choice for most problems. Switch to SGD-with-momentum if you have a known reason to (e.g. a ResNet on ImageNet where SGD generalises better) and are willing to tune.
3. Pick the batch size. The largest batch that fits in GPU memory, capped by generalisation considerations (very large batches can hurt test accuracy).
4. Run a learning-rate finder. Sweep $\eta$ over $3$–$5$ orders of magnitude for one epoch; pick $\eta$ a factor of $3$–$10$ below the cliff where loss starts climbing.
5. Add a schedule. Cosine annealing or step decay; warmup if using large batches or Transformers.
6. Watch the loss curve. Decreasing-then-flat is good; flat-then-decreasing is suspicious; increasing is bad.
7. Train multiple seeds. Report mean and standard deviation of test metrics across $3$–$5$ runs with different random initialisations.

That recipe gets you a working model on most problems without extensive hyperparameter tuning.

Key Takeaways

• Gradient descent has three variants — batch (full pass per update, exact, slow), stochastic (one example per update, cheap, noisy), and mini-batch (a batch of $B$ examples per update, the practical default).
• The learning rate $\eta$ controls the step size; $\eta$ too high causes divergence and NaN losses, $\eta$ too low causes stalling and wasted compute.
• Momentum adds a velocity term that accelerates consistent gradients and dampens oscillation in ravines, with $\beta \approx 0.9$ as the default.
• RMSProp maintains an EMA of squared gradients per parameter and divides the update by its square root, normalising step size per parameter.
• Adam combines momentum on the first moment and RMSProp on the second moment with bias correction for early steps, and is the most common default optimiser.
• Convex losses have a single global minimum; non-convex losses have multiple local minima and saddle points.
• Saddle points, not local minima, are the dominant obstacle in high-dimensional non-convex optimisation, and SGD's gradient noise naturally escapes them.
• Learning-rate schedules (step decay, exponential, cosine, warmup) anneal $\eta$ over training and improve final performance; the loss curve is the diagnostic.