M1.4 — Optimization and Gradient Descent

Learning objectives

After this lesson, you should be able to:

• Write the gradient descent update and explain the learning rate.
• Show when gradient descent converges, oscillates or diverges.
• Explain momentum and Adam at a high level.
• Explain why mini-batch gradients are unbiased but noisy, and how batch size affects the noise.
• Explain local minima and why the starting point matters.
• Describe learning-rate warmup and decay and gradient clipping.

The one-sentence idea

Training repeatedly steps the parameters against the gradient, and getting the step size and direction handling right decides whether it converges, crawls or blows up.

Downhill is easy to find. How far to step is the hard part.

Training is an optimization problem

Optimization means finding the parameters θ that make a function, here the loss L(θ), as small as possible. Everything you trained in M2.4 to M2.5 was one: the loss was low for good parameters, and training searched for them. Calculus (M1.2) says which direction is downhill, and gradient descent repeatedly steps that way:

The step size η (eta) is the learning rate. It is the most important setting in training, as the first experiment shows. Take the simplest bowl-shaped loss L(x) = x², whose gradient is 2x and whose minimum is at 0, and start at x = 5 with three learning rates:

import numpy as np

def gd(eta, steps=8, x=5.0):
    path = [x]
    for _ in range(steps):
        x = x - eta * 2 * x                  # the gradient of x^2 is 2x
        path.append(round(x, 3))
    return path

print(gd(0.1))
print(gd(1.0))
print(gd(1.1))

[5.0, 4.0, 3.2, 2.56, 2.048, 1.638, 1.311, 1.049, 0.839]
[5.0, -5.0, 5.0, -5.0, 5.0, -5.0, 5.0, -5.0, 5.0]
[5.0, -6.0, 7.2, -8.64, 10.368, -12.442, 14.93, -17.916, 21.499]

Each step multiplies x by 1 − 2η. With η = 0.1 that factor is 0.8, so x shrinks smoothly toward the minimum, slowly. With η = 1.0 the factor is −1: each step jumps exactly to the other side and bounces forever. With η = 1.1 the factor is −1.2, so the jumps grow and the process diverges. Too small wastes time, and too large breaks training. It is the same stability condition as the proportional control of M5.6, because both are feedback loops.

Valleys, momentum and adaptive methods

Real loss surfaces are rarely round bowls. Often some directions are steep and others almost flat, a long narrow valley. The learning rate then has to stay small enough for the steep direction, which makes progress along the flat direction painfully slow. Take L = ½(x² + 25 y²), steep in y and gentle in x, and count the steps needed to get the loss below 10⁻⁶:

grad = lambda p: np.array([p[0], 25.0 * p[1]])          # gradient of 0.5 * (x^2 + 25 y^2)
loss = lambda p: 0.5 * (p[0] ** 2 + 25.0 * p[1] ** 2)

def steps_to_converge(update, p0=(10.0, 1.0), limit=20000):
    p, state = np.array(p0), {}
    for t in range(1, limit + 1):
        p = update(p, grad(p), state, t)
        if loss(p) < 1e-6:
            return t
    return None

def plain(p, g, s, t, lr=0.07):
    return p - lr * g

def momentum(p, g, s, t, lr=0.07, beta=0.9):
    s["v"] = beta * s.get("v", 0) + g                  # a running sum of past gradients
    return p - lr * s["v"] * (1 - beta)

def adam(p, g, s, t, lr=0.3, b1=0.9, b2=0.999, eps=1e-8):
    s["m"] = b1 * s.get("m", 0) + (1 - b1) * g         # smoothed gradient
    s["v"] = b2 * s.get("v", 0) + (1 - b2) * g * g     # smoothed squared gradient
    m_hat, v_hat = s["m"] / (1 - b1 ** t), s["v"] / (1 - b2 ** t)
    return p - lr * m_hat / (np.sqrt(v_hat) + eps)

print(steps_to_converge(plain), steps_to_converge(momentum), steps_to_converge(adam))

123 163 157

The counts are 123 steps for plain gradient descent, 163 for momentum and 157 for Adam. Plain descent is not worse here, which is the honest result: on a simple quadratic with a well-chosen learning rate there is little for the fancier methods to fix (their rates were set by hand), and their advantages show up on harder, noisier problems and when good rates are not known in advance. What each method does is the point. Plain gradient descent is limited by the steepest direction. Momentum keeps a running average of past gradients, which damps back-and-forth zig-zags and builds speed along consistent directions. Adam also rescales each parameter’s step by the size of its recent gradients, so parameters with tiny gradients are not left behind. That adaptivity is why Adam is the common default for large models and for the training loops of M2.5, where the right learning rate differs enormously between layers.

Stochastic gradient descent: noisy but cheap

The gradient of the loss over the whole dataset is expensive to compute at each step. Stochastic gradient descent (SGD) estimates it from a random mini-batch (M2.5). The estimate is unbiased: averaged over many batches it equals the true gradient, but any single one is noisy, and, by the √n law of M1.3, the noise shrinks as the batch grows:

rng = np.random.default_rng(0)
x = rng.normal(size=10000)
y = 3 * x + rng.normal(0, 0.5, 10000)
w = 0.0                                                   # evaluate the gradient of the mean squared error at w = 0

per_example = -2 * x * (y - w * x)                        # gradient of (y - w x)^2 for each example
print(round(float(per_example.mean()), 2))                # the full-data gradient
for B in (1, 10, 100, 1000):
    est = [per_example[rng.integers(0, 10000, B)].mean() for _ in range(2000)]
    print(B, round(float(np.mean(est)), 2), round(float(np.std(est)), 2))

-5.96
1 -5.75 7.99
10 -6.07 2.66
100 -5.96 0.84
1000 -5.96 0.26

The full gradient is −5.96. Batches of every size average out to roughly that (unbiased), but the spread of a single-example estimate is 7.99, larger than the quantity itself, shrinking about as 1/√B to 0.26 at B = 1000. Some noise is useful, because it can shake training out of poor regions, and large batches allow larger learning rates, but too much slows convergence. Batch size is another setting chosen on the validation set (M6.4).

Non-convex losses, schedules and clipping

A bowl has one minimum. The loss of a neural network does not: it has many local minima and flat saddle regions, and where you end up depends on where you start. A one-dimensional picture with two valleys shows the principle:

f_grad = lambda x: 4 * x ** 3 - 6 * x + 1           # derivative of f(x) = x^4 - 3x^2 + x
for start in (-2.0, 2.0):
    x = start
    for _ in range(200):
        x -= 0.01 * f_grad(x)
    print(start, round(x, 3), round(x ** 4 - 3 * x ** 2 + x, 3))

-2.0 -1.301 -3.514
2.0 1.131 -1.07

Starting from −2 the descent settles at x = −1.301, where the function is −3.514, the global minimum. Starting from +2 it settles in a different valley, x = 1.131, with a higher value of −1.07. In very high dimensions, practice has shown that the many minima of large networks tend to be good enough, but initialisation, learning rate and data order still change the result. That is why training is repeated with different random seeds when comparing methods.

Two further practical tools appear in nearly every large training recipe:

• Learning-rate schedules. The rate is not held constant: it is often raised from near zero over a short warmup, then decayed, for example along a cosine curve, so that early steps are cautious and late steps are fine.
• Gradient clipping. If a gradient’s norm exceeds a threshold, it is scaled back so that one bad batch cannot throw the parameters far away.

def lr_at(step, total, peak=1e-3, warmup=100):
    if step < warmup:
        return peak * step / warmup                                   # linear warmup
    progress = (step - warmup) / (total - warmup)
    return peak * 0.5 * (1 + np.cos(np.pi * progress))                # cosine decay to zero

print([round(float(lr_at(s, 1000) * 1e3), 3) for s in (0, 50, 100, 550, 1000)])

def clip_norm(g, max_norm):
    n = np.linalg.norm(g)
    return g * (max_norm / n) if n > max_norm else g

print(clip_norm(np.array([30.0, 40.0]), 5.0))                         # norm 50 scaled down to norm 5

[0.0, 0.5, 1.0, 0.5, 0.0]
[3. 4.]

The schedule, shown in units of 10⁻³, climbs from 0 to 1.0 over the first 100 steps, falls to half at the midpoint, and reaches 0 at the end. The clipped gradient [3, 4] has the same direction as the original [30, 40] but a norm of 5. You will see exactly these knobs, optimizer, learning rate, schedule and clipping, in the training configuration of any VLA fine-tuning recipe.

Common misconceptions

Engineering takeaways

1. Gradient descent updates θ ← θ − η ∇L.
2. Too small a learning rate crawls, a larger one oscillates, and a too-large one diverges.
3. Narrow valleys limit the learning rate and slow progress along flat directions.
4. Momentum smooths and accelerates, and Adam adapts the step per parameter.
5. Mini-batch gradients are unbiased with noise shrinking as 1/√B.
6. Non-convex losses have several minima, and the start point matters.
7. Warmup and decay schedules and gradient clipping are standard in large training runs.
8. Run several seeds when comparing methods.

Optimization is the quiet engine under every model: a loop of small, well-judged steps downhill.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

That completes Module 1, the mathematics. Module 2 turns to the tools: first, the Python needed to read and write model and data code.

Back to the VLA roadmap prerequisites