M5.3 · LESSON 3 OF 6

Learning-rate stability

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

THE IDEA

A step is a bet that the slope stays the same

The gradient is the slope right here. The update walks a whole step of length η·|g| using that slope, as if the ground stayed the same along the way. For a gentle slope that is fine; for a curved one it can carry you over the bottom and up the opposite side, higher than where you began.

For the parabola L = a·x² the maths is exact. The gradient is 2ax, so the update is x ← x − η·2ax = (1 − 2aη)·x. Every step multiplies x by the same number 1 − 2aη. The iteration converges only if that number lies strictly between −1 and 1, which means η < 1/a (equivalently η < 2/(2a)). Negative values mean the iterate jumps across the minimum each step.

Learning rate η0.9

L(x) = (x - 3)² so g = 2·(x - 3) update: x ← x − η·g now x = 0 (press Step)

step	x	gradient	loss
0	0	-6	9

Each update multiplies the distance to the minimum by 1 − 2·1·η = -0.8. Negative: it overshoots and swings across, but the swings shrink. Limit: η below 1.

L = (x − 3)² has a = 1, so the limit is η = 1. Push the slider past 1 and press +10 steps. Try 0.5, 0.9, 1.0 and 1.1.

Predict first. For f(x) = 2.5·x², which learning rate makes gradient descent diverge?

Curvature sets the limit, and the flattest direction sets the pace

Real losses curve differently in different directions. On the valley L = 0.02x² + y² the y direction has a = 1 (limit η < 1) and the x direction has a = 0.02 (limit η < 50). The step must be safe for the steepest direction, so η cannot exceed 1. But at η close to 1 the x factor is 1 − 2·0.02·η ≈ 0.96: x shrinks only about 4 percent per step. That mismatch is called ill-conditioning.

Learning rate 0.9

iteration 0

Gradient descent: position [-5, 2], loss 4.5. At the start. Press Step.

Click the plot to choose a new starting point. Hollow dot: start. Big dot: where the optimizer is now.

Pick a learning rate for the valley (start at the hollow dot). Press Run 50. Compare 0.1, 0.9 and 1.05. Click the plot to change the start.

The common mistake: judging by one step

With L = θ² and η = 2 the first step looks harmless, it moves to −3, but the iterates go 1, -3, 9, -27, 81, ...: multiplied by −3 every time. Pure divergence often looks like oscillation for a few steps. If the loss rises, or the sign of a coordinate flips every step with growing size, lower η.

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. Tune and explain

Valley: set the learning rate below, press Run 50 to see the path. The tasks are checked against the real simulation of this valley.

Valley learning rate 0.1

iteration 0

Gradient descent: position [-5, 2], loss 4.5. At the start. Press Step.

Hollow dot: start. Big dot: where the optimizer is now.

Quadratic tester: f(x) = 2.5·x², start x = 1. Try learning rates around your guess for the limit.

Learning rate η0.3

L(x) = 2.5 (x )² so g = 5·(x) update: x ← x − η·g now x = 1 (press Step)

step	x	gradient	loss
0	1	5	2.5

Each update multiplies the distance to the minimum by 1 − 2·2.5·η = -0.5. Negative: it overshoots and swings across, but the swings shrink. Limit: η below 0.4.

Your largest stable learning rate

The gradient 2·a·x vanishes when the learning rate reaches 1/a, so the minimum itself starts moving. Each update multiplies x by (1 − 2·a·lr), which shrinks x only while that factor lies between −1 and 1, that is lr < 1/a. Losses above 1 always explode, so any learning rate must stay below 1 whatever the curvature is.

To do: On the valley (start [-5, 2]) choose a learning rate that brings the loss to 0.05 or lower within 40 steps
To do: Now break it: choose a learning rate for which the loss after 40 steps is larger than at the start
To do: For f(x) = 2.5 x², type the largest learning rate that still converges (use the simulator to test your guess)
To do: Pick the correct explanation of that limit

Part 2. Compute the limit in code

For f(x) = a·x² write contraction(a, lr): the factor by which one gradient descent update multiplies x.
Write max_stable_lr(a): the largest learning rate for which the iteration still converges.
Write diverges(a, lr, x0=1.0, steps=200): run the real update loop from x0 and return True if |x| after the steps is larger than |x0|.

def contraction(a, lr):
    # x_new = contraction * x for f(x) = a * x**2
    pass

def max_stable_lr(a):
    # largest lr that still converges for f(x) = a * x**2
    pass

def diverges(a, lr, x0=1.0, steps=200):
    # run x = x - lr * f'(x) for 'steps' updates; True if |x| grew
    pass

print(contraction(2.5, 0.3))

Ctrl/⌘ + Enter runs it

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

CHECK YOUR UNDERSTANDING

For L=theta squared, does eta=2 converge from theta=1?

ANoBYes

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.