M5.6 · LESSON 6 OF 6

Adam and local geometry

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

THE IDEA

Adam gives every coordinate its own step size

On the valley the gradient is [−0.2, 4] at the start: one coordinate is 20 times larger. Plain descent moves 20 times farther in y than in x. Adam keeps two running averages per coordinate:

m ← β₁·m + (1 − β₁)·g the average gradient (the momentum-like part),
v ← β₂·v + (1 − β₂)·g² the average squared gradient (how big this coordinate's gradients typically are).

The update is θ ← θ − η · m̂ / (√v̂ + ε), computed for each coordinate. Dividing by √v̂ rescales every coordinate to roughly the same size, so a coordinate with huge gradients and one with tiny gradients both move about η per step. The defaults are β₁ = 0.9, β₂ = 0.999 and ε = 10⁻⁸ (a small number that prevents dividing by zero).

Adam learning rate0.5 bias correction

t	g (x, y)	m (x, y)	v (x, y)	m̂ = m / (1 − β₁ᵗ)	v̂ = v / (1 − β₂ᵗ)	Adam step (x, y)	plain GD step
1	-0.2, 4	-0.02, 0.4	0.00004, 0.016	-0.2, 4	0.04, 16	0.5, -0.5	0.1, -2
2	-0.18, 3	-0.036, 0.66	0.00007236, 0.02498	-0.1895, 3.474	0.0362, 12.5	0.4979, -0.4913	0.098, 0
3	-0.1601, 2.017	-0.04841, 0.7957	0.00009791, 0.02903	-0.1786, 2.936	0.03267, 9.686	0.4941, -0.4717	0.09604, 0
4	-0.1403, 1.074	-0.0576, 0.8236	0.0001175, 0.03015	-0.1675, 2.395	0.02942, 7.55	0.4882, -0.4358	0.09412, 0

The valley gradient at the start is [−0.2, 4]: the y component is 20 times larger. Plain gradient descent steps are proportional to that, Adam's are not.

iteration 0

Plain gradient descent (same lr): position [-5, 2], loss 4.5. At the start. Press Step.

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

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

Table of the first four updates on the valley, with the plain-gradient-descent step for comparison. Move the learning rate and switch bias correction off to see the first-step size change.

Bias correction, in numbers

m and v start at zero. After one step m = 0.1·g and v = 0.001·g²: far too small, only because the averages have not had time to fill. Without a fix the first update would be lr · 0.1 / √0.001 = lr · 3.16 times the right size, and in the early steps the estimates are biased toward zero. Dividing by 1 − βᵗ removes that exactly: for a gradient that never changes, m̂ and v̂ equal g and g² at every step.

Constant gradient g2β₁0.9β₂0.999

t	m_t	1 − β₁ᵗ	m̂_t	v_t	1 − β₂ᵗ	v̂_t
1	0.2	0.1	2	0.004	0.001	4
2	0.38	0.19	2	0.007996	0.001999	4
3	0.542	0.271	2	0.01199	0.002997	4
5	0.819	0.4095	2	0.01996	0.00499	4
10	1.303	0.6513	2	0.03982	0.009955	4
50	1.99	0.9948	2	0.1952	0.04879	4

The gradient never changes, so the right answers are m = g = 2 and v = g² = 4. The raw m and v start near zero and creep up. Dividing by 1 − βᵗ gives exactly g and g² at every step.

A constant gradient. Raw m_t and v_t creep toward g and g². The corrected columns are g and g² from step 1.

Predict first. At the very first Adam step on the valley (lr = 0.5, bias correction on), how far does x move, given its gradient is only −0.2?

Local geometry: Adam is not a guarantee

Adam rescales steps by the history of gradients; it does not know the global shape. On a bumpy surface it still slides into whichever local minimum is near its start. Below, press Run 50 with Adam, then click a different start point and run again.

Learning rate0.3

L = 0.04 (x² + y²) + 0.5 sin(2x) sin(2y). Many local minima. Where you end up depends on where you start and on the optimizer.

iteration 0

Adam: position [-3, 2.5], loss 0.476. At the start. Press Step.

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

Adam on the bumpy surface. Different starting points end in different valleys with different losses.

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 Adam

Adam learning rate0.05 bias correction

t	g (x, y)	m (x, y)	v (x, y)	m̂ = m / (1 − β₁ᵗ)	v̂ = v / (1 − β₂ᵗ)	Adam step (x, y)	plain GD step
1	-0.2, 4	-0.02, 0.4	0.00004, 0.016	-0.2, 4	0.04, 16	0.05, -0.05	0.01, -0.2
2	-0.198, 3.9	-0.0378, 0.75	0.00007916, 0.03119	-0.1989, 3.947	0.0396, 15.6	0.04999, -0.04996	0.00998, -0.18
3	-0.196, 3.8	-0.05362, 1.055	0.0001175, 0.0456	-0.1979, 3.893	0.03921, 15.22	0.04996, -0.0499	0.00996, -0.162
4	-0.194, 3.7	-0.06766, 1.32	0.000155, 0.05925	-0.1967, 3.837	0.03881, 14.83	0.04993, -0.04981	0.00994, -0.1458

The valley gradient at the start is [−0.2, 4]: the y component is 20 times larger. Plain gradient descent steps are proportional to that, Adam's are not.

iteration 0

Plain gradient descent (same lr): position [-5, 2], loss 4.5. At the start. Press Step.

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

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

To do: Reach loss 0.05 or lower within 30 steps using Adam (set the learning rate)
To do: With bias correction on, make the first Adam update move x by 0.5 (within 0.01)
To do: Beat plain gradient descent by a factor of two: reach loss 0.05 or lower within 15 steps (the best plain gradient descent needs about 31)
To do: Switch bias correction off and choose the learning rate so the first update still moves x by 0.5 (within 0.02). Compare with the corrected size

Part 2. Write Adam yourself

Write adam_step(theta, m, v, grad, t, lr, b1, b2, eps) for step number t (counting from 1). Update m and v, divide by 1 - b1**t and 1 - b2**t, then move by lr * m_hat / (sqrt(v_hat) + eps). Return (theta_new, m_new, v_new).
Write adam(grad_fn, theta0, lr, steps, b1=0.9, b2=0.999, eps=1e-8): start with m = v = 0 and return the path of steps + 1 positions.

from math import sqrt

def adam_step(theta, m, v, grad, t, lr, b1, b2, eps):
    # m = b1*m + (1-b1)*g ; v = b2*v + (1-b2)*g*g ; bias-correct ; move coordinatewise
    pass

def adam(grad_fn, theta0, lr, steps, b1=0.9, b2=0.999, eps=1e-8):
    # start with m = v = 0, t = 1 for the first update; return the path
    pass

print(adam_step([-5.0, 2.0], [0.0, 0.0], [0.0, 0.0], [-0.2, 4.0], 1, 0.5, 0.9, 0.999, 1e-8))

Ctrl/⌘ + Enter runs it

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

CHECK YOUR UNDERSTANDING

Why correct moments initialized to zero?

ABecause the learning rate would otherwise change sign in the first steps.BTheir early values are biased toward zero.CBecause the moments would otherwise become larger than the gradients.

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.