M5.4 · LESSON 4 OF 6

Stochastic gradients

Perform gradient descent updates and explain learning rate, stochastic gradients, momentum and Adam.

THE IDEA

Why estimate the gradient

Training losses are usually averages over a dataset: L(θ) = (1/N) Σ L_i(θ). The full gradient averages N per-example gradients, so one update costs N gradient evaluations. A minibatch averages only k randomly chosen examples. It is a cheaper estimate of the same gradient: on average it is exactly right, but any single minibatch can point somewhat elsewhere.

Our toy dataset has 8 points in the plane, and each example's loss is ½·|θ − c|². The full-data minimum is the mean of the points, [1.5, 1.25]. At θ = [-4, -3] the full gradient is [-5.5, -4.25], while the single example #3 alone says [-7, -5].

Watch the noise

Batch size1Learning rate0.4

decay: η / (1 + 0.1·(t − 1))Seed

iteration 0

full batch (all 8)minibatch of 1true optimum [1.5, 1.25]

Press Step. Rings mark the examples used in the current minibatch. excess loss: minibatch 24.156 full batch 24.156 noise floor over steps 41 to 60 (minibatch): 1.154

Same seed, same settings, same path: re-run it and nothing changes. Change the seed and the minibatch path changes, while the full-batch path (blue) never does.

Blue: full batch, always the same. Orange: minibatches drawn with a seeded random generator, so the run is repeatable. Rings mark the examples in the current batch. Press Step, then change the batch size or the seed.

With a constant learning rate the minibatch path reaches the neighbourhood of the optimum quickly and then keeps jittering around it, never settling. The size of that jitter, the noise floor, grows with the learning rate and shrinks with the batch size. Two standard remedies are a larger batch and a learning rate that decays over time.

Predict first. Averaged over every possible minibatch of size 2, how does the minibatch gradient compare with the full gradient?

Common mistakes

Treating one minibatch as the truth. The minibatch gradient and the minibatch loss are estimates. They bounce from step to step even when training is healthy.
Comparing two runs that differ in luck. Different seeds or shuffling give different paths. When comparing batch sizes or learning rates, fix the seed, and repeat with several seeds before believing a small difference.
Mixing up the cost. Halving the batch size halves the cost of a step, but the estimate gets noisier. There is no free choice.

Chapter context

The gradient points toward the steepest local increase under Euclidean geometry. To lower a differentiable loss, a basic optimizer steps in the opposite direction. A local direction says little about how large a safe step can be.

TRY IT YOURSELF

Part 1. Trade noise against speed

Seed 7, 60 steps, start (−4, −3). Tasks use the real seeded runs. The seed box in the widget is for exploring, the tasks always use seed 7.

Batch size8 (all)Learning rate0.4

decay: η / (1 + 0.1·(t − 1))Seed

iteration 0

full batch (all 8)minibatch of 8true optimum [1.5, 1.25]

Press Step. Rings mark the examples used in the current minibatch. excess loss: minibatch 24.156 full batch 24.156 noise floor over steps 41 to 60 (minibatch): 0

Same seed, same settings, same path: re-run it and nothing changes. Change the seed and the minibatch path changes, while the full-batch path (blue) never does.

To do: Make the noise visible: choose a batch size below 8 and a learning rate for which the noise floor (mean excess loss over the last 20 steps) is at least 1
To do: With batch size 1, tame the noise: get the noise floor down to 0.1 or lower while still starting far from the optimum
To do: At constant learning rate 0.4, find a batch size (below 8) whose noise floor is 0.35 or lower

Part 2. Write minibatch SGD

The loss is the mean of 0.5·(θ − c)² over the data c (one number per example). Write full_grad(data, theta): the gradient over all examples.
Write batch_grad(data, theta, idx): the same gradient computed only from the examples whose positions are listed in idx.
Write sgd(data, theta0, lr, batches): batches is a list of index lists, one per step. Return the list of thetas, starting with theta0.

def full_grad(data, theta):
    # mean over all examples of (theta - c)
    pass

def batch_grad(data, theta, idx):
    # mean over the examples data[i] for i in idx of (theta - c)
    pass

def sgd(data, theta0, lr, batches):
    # one update per batch; return [theta0, theta1, ...]
    pass

print(full_grad([1, 2, 3, 6], 0.0))

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Is one minibatch gradient necessarily equal to the full-data gradient?

AYesBNo

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

\theta_{t+1}=\theta_t-\eta\nabla L(\theta_t).

Adam uses $m_t=\beta_1m_{t-1}+(1-\beta_1)g_t$ and $v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2$ , coordinatewise. Divide by $1-\beta_1^t$ and $1-\beta_2^t$ before using $\hat m_t/(\sqrt{\hat v_t}+\epsilon)$ in the update.

For the special quadratic $L(\theta)=(\theta-c)^2$ , the residual follows $\theta_{t+1}-c=(1-2\eta)(\theta_t-c)$ . Starting away from c, convergence requires $|1-2\eta|<1$ , or $0<\eta<1$ . This interval depends on this particular curvature and loss normalization.

Worked example

For L=(theta−3)^2, start at 0 and use eta=0.25. The gradients are −6, −3 and −1.5. Three updates produce theta values 1.5, 2.25 and 2.625. Their losses are 2.25, 0.5625 and 0.140625. A fixed learning rate works here because the example is a simple quadratic.

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import isclose
theta = 0.0
positions = []
for _ in range(3):
    theta -= 0.25 * 2*(theta-3)
    positions.append(theta)
assert positions == [1.5, 2.25, 2.625]
assert isclose((theta-3)**2, 0.140625)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Subtracting loss instead of its gradient is not gradient descent. Squared gradients in Adam are coordinatewise. A zero gradient may identify a saddle point or maximum, not just a minimum.