D3.3 · LESSON 3 OF 6

SGD and Adam

Diagnose an unstable training run using losses, gradient magnitudes, initialization, normalization and regularization.

THE IDEA

From a gradient to a step

The gradient says which way is uphill. An optimizer decides how far to move. SGD moves by the learning rate times the gradient. Momentum keeps a velocity (a running sum of recent gradients), so consistent directions build speed and zig-zags cancel. Adam keeps two running averages, one of the gradient and one of its square, and divides one by the square root of the other, so each weight gets a step of about the learning rate in size no matter how large its gradient is.

Adam, one step (t starts at 1)

m = b1*m + (1-b1)*g                  # running average of the gradient
v = b2*v + (1-b2)*g*g                # running average of its square
m_hat = m / (1 - b1**t)              # bias correction: m and v start at 0
v_hat = v / (1 - b2**t)
w = w - lr * m_hat / (sqrt(v_hat) + eps)

sgd (lr 0.05)momentum (lr 0.05)adam (lr 0.05)

Same network, same learning rate 0.05. After 200 steps SGD has reached 0.37, Adam 0.001. Their sensible learning rates differ by orders of magnitude, so they should not share one.

Optimizer state is part of the training

Adam's two averages are state that is saved separately from the weights. If you restart training from saved weights but with fresh (zeroed) averages, the next steps behave differently from an uninterrupted run, even though the model is identical. The demo clears the state at step 100.

clear Adam's moment estimates at step 100 (the weights stay exactly the same)

no resetreset at step 100

loss at step 101: no reset 0.0035 this run 0.0290

The green line is the uninterrupted run. The red dashed one clears Adam's moments at step 100.

Predict first. You resume training by loading only the model weights (not the optimizer state). Is that exactly equivalent to continuing the original Adam run?

Chapter context

A differentiable model can still be difficult to train. Learning rate, parameter scale, data scale and optimizer state interact. Inspect the run before changing many settings at once.

TRY IT YOURSELF

Part 1. Do it hands-on

Choose an optimizer, a learning rate and (optionally) when to clear its state. Every run is a real 200-step training on the same data.

learning rate0.05clear optimizer state at stepnever

training loss

2-8-8-1 tanh network, 200 steps final training loss 0.3749 accuracy 80% slow: falling, but far from done

To do: Reach a final training loss below 0.05 with plain SGD (200 steps)
To do: Reach below 0.05 with Adam using a learning rate of 0.02 or smaller
To do: Break Adam: find a learning rate of 0.5 or more where it fails to get below 0.05
To do: Pause Adam at step 100 and clear its moment estimates (reset at 100). The loss at step 101 should be at least 0.01 higher than the same run without the reset

Part 2. Do it in code

momentum_step(w, v, g, lr, beta): v_new = beta * v + g, then w_new = w - lr * v_new. Return (w_new, v_new). Do not modify the inputs.
adam_step(w, m, v, g, t, lr=0.001, b1=0.9, b2=0.999, eps=1e-8): t counts steps starting at 1 (raise ValueError if t < 1). m_new = b1 m + (1 - b1) g; v_new = b2 v + (1 - b2) g^2; bias-correct m_hat = m_new / (1 - b1^t) and v_hat = v_new / (1 - b2^t); w_new = w - lr * m_hat / (sqrt(v_hat) + eps). Return (w_new, m_new, v_new).

import numpy as np

def momentum_step(w, v, g, lr, beta):
    pass

def adam_step(w, m, v, g, t, lr=0.001, b1=0.9, b2=0.999, eps=1e-8):
    pass

Ctrl/⌘ + Enter runs it

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

CHECK YOUR UNDERSTANDING

Is resuming only model weights exactly equivalent to resuming Adam training?

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

Norm clipping uses

g'=g\min\left(1,\frac{c}{\|g\|_2}\right).

For inverted dropout with keep probability q, $\tilde h=m h/q$ where $m\sim\operatorname{Bernoulli}(q)$ , so $E[\tilde h]=h$ . This expectation is over masks, not a claim that each individual masked output is unchanged.

Worked example

The gradient [3,4] has norm 5. Clipping its norm to 2 yields [1.2,1.6], preserving direction. For a scalar activation 4 and keep probability 0.5, dropout gives either 8 or 0 with equal probability, averaging to 4.

Reference code

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

from math import sqrt, isclose
g = [3,4]
scale = min(1, 2/sqrt(sum(v*v for v in g)))
clipped = [v*scale for v in g]
assert all(isclose(a,b) for a,b in zip(clipped,[1.2,1.6]))
assert isclose(sqrt(sum(v*v for v in clipped)),2)
assert .5*8+.5*0 == 4

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

Common mistakes

Clipping is not the same as normalizing every gradient to the limit. Lower training loss can accompany worse validation. A numerical fix should not conceal invalid data or labels.