last time i took a two-layer network and computed all four gradients by hand, and then ran the same thing in pytorch and autograd landed on the same four numbers, 30, 30, 20, 10. and then i stopped there, which is exactly one line too early, because in a real training loop the next thing that happens is optimizer.step() and that line is where the numbers actually turn into learning. so this time i want to do the same trick with the optimizer, take those four gradients and apply one update by hand, first with plain SGD, then with momentum, then with Adam, and see what each of them does to the same starting point. the whole thing fits in a few lines of arithmetic and it removes away a lot of vague feeling about learning rates, because you get to watch the same network survive one step and die on another.
Where we are
same network as last time, two layers, one hidden ReLU, squared error:
x ──▶ h1 = w1·x + b1 ──▶ a1 = ReLU(h1) ──▶ h2 = w2·a1 + b2 ──▶ L = (h2 − y)²
x = 1.0 y = 2.0
w1 = 2.0 b1 = 0.0
w2 = 3.0 b2 = 1.0
forward gives h1 = 2, a1 = 2, h2 = 7, L = 25, and backward gives
∂L/∂w1 = 30 ∂L/∂b1 = 30 ∂L/∂w2 = 20 ∂L/∂b2 = 10
an optimizer's entire job is to take that list of four numbers and decide how far, and in what direction, each of the four parameters moves. that is all it is. it does not see the network, it does not know there is a ReLU in there, it gets a parameter and the gradient sitting on it and nothing else.
SGD, which is one line of arithmetic
the rule is param = param − lr × grad, and there is nothing hidden in it. with lr = 0.1:
w1 = 2.0 − 0.1·30 = −1.0
b1 = 0.0 − 0.1·30 = −3.0
w2 = 3.0 − 0.1·20 = 1.0
b2 = 1.0 − 0.1·10 = 0.0
import torch
def net(w1, b1, w2, b2, x=torch.tensor(1.0), y=torch.tensor(2.0)):
h1 = w1 * x + b1
a1 = torch.relu(h1)
h2 = w2 * a1 + b2
return (h2 - y) ** 2, h1
p =[torch.tensor(v, requires_grad=True) for v in (2.0, 0.0, 3.0, 1.0)]
opt = torch.optim.SGD(p, lr=0.1)
loss, h1 = net(*p)
opt.zero_grad()
loss.backward()
opt.step()
print([round(t.item(), 4) for t in p]) # [-1.0, -3.0, 1.0, 0.0]
print(net(*p)[0].item(), net(*p)[1].item()) # 4.0 -4.0
now look at what that did. the loss went from 25 to 4, which reads like a great first step. but the new h1 is −4, and a ReLU with a negative input outputs zero and has derivative zero, so the hidden
neuron is dead, and from the previous article we know exactly what that means for the gradients,w1, b1 and w2 all get zero from here on and only b2 keeps learning. one step. the network
lost its hidden layer on the very first update and the loss number went down while it happened, so the loss will not tell you.
same thing at lr = 0.01:
w1 = 1.7 b1 = −0.3 w2 = 2.8 b2 = 0.9
h1 = 1.4 loss = 7.9524
still alive, loss 25 → 7.95, and it keeps going down from there and reaches basically zero after about thirty steps with h1 still positive at 0.48. so the difference between a network that trains
and a network that quietly stopped having a hidden layer, on this problem, is one number in the constructor. that is why the learning rate is the hyperparameter people actually tune, and why "the
loss went down" is not evidence that anything is healthy.
Momentum, and the part where it does not do what you expect
pytorch writes momentum like this, and it is worth knowing the exact form because a lot of textbooks put the learning rate inside the buffer but pytorch does not:
buf = momentum · buf + grad (on the first step, buf = grad)
param = param − lr · buf
so the buffer is a running sum of gradients with the old ones decayed by 0.9 each step, and the update uses that instead of the raw gradient. run it on our network with lr = 0.01, momentum = 0.9
and do the first two steps by hand.
step 1. the buffer starts empty, so buf = grad = [30, 30, 20, 10] and the update is identical to plain SGD. momentum does literally nothing on the first step, it has no history yet.
w1 = 1.7 b1 = −0.3 w2 = 2.8 b2 = 0.9 loss = 7.9524
step 2. gradients at the new point are [15.792, 15.792, 7.896, 5.64]. the buffer is
buf = 0.9·[30, 30, 20, 10] + [15.792, 15.792, 7.896, 5.64]
= [42.792, 42.792, 25.896, 14.64]
42.792 where the gradient alone is 15.792, so the step is nearly three times bigger than the
gradient asked for, because the previous gradient pointed the same way and momentum is adding them
up. and it pays off immediately:
w1 = 1.27208 b1 = −0.72792 w2 = 2.54104 b2 = 0.7536
h1 = 0.54416 loss = 0.018587
plain SGD at the same step 2 is at 3.217154. momentum is at 0.0186. that is the case for it in one line, on a slope where the gradients keep pointing the same direction it accumulates and covers
ground much faster than the raw gradient would.
step 3, where it goes wrong. we overshot. h1 was 0.544 and the buffer is still carrying 39.2 of leftover velocity, so the next update pushes h1 to −0.239953 and the ReLU shuts. loss back up to 1.906848. and here is the thing that makes momentum worth understanding rather than just switching on, at step 4 the gradients on w1 and b1 are exactly 0, the neuron is dead, and the
buffer is still 35.2851, so the parameters keep moving anyway. that is what momentum means, taken literally. no gradient, still moving. h1 goes to −0.9457, then −1.58, and it is being
driven further into the dead zone by the memory of gradients from three steps ago.
run it out to 60 steps and the ending is the bit i would not have guessed:
momentum, 60 steps: loss = 0.005930 w1 = −2.6398 b1 = −4.6398 w2 = 0.2008 b2 = 1.9230
plain SGD, 30 steps: loss ≈ 0.000000 w1 = 1.2415 b1 = −0.7585 w2 = 2.6270 b2 = 0.7313
the momentum run reports a loss of 0.0059 and looks converged. but h1 = −2.6398 − 4.6398 = −7.28,
the hidden neuron has been dead for fifty-odd steps, w2 decayed to 0.2 because it is reading from a constant zero, and the whole thing is now the output bias b2 = 1.923 predicting a target of 2.
it solved the problem by throwing the network away and keeping the bias. a low loss on a tiny problem can mean the model found the answer or it can mean the model routed around its own dead parts, and
the only way to see which one you have is to print something other than the loss.
i am not saying do not use momentum, on real problems with mini-batch noise it is almost always better and it costs one keyword. i am saying the failure mode is specific and worth having seen once,
because "momentum keeps moving when the gradient is zero" is the same property that gets you through a flat region and the same property that drives a neuron off a cliff.
Adam, and what its learning rate actually is
Adam keeps two running averages per parameter, one of the gradient and one of the squared gradient, then divides one by the square root of the other:
m = β1·m + (1−β1)·g β1 = 0.9 (average gradient)
v = β2·v + (1−β2)·g² β2 = 0.999 (average squared gradient)
m̂ = m / (1 − β1^t) bias correction, t = step number
v̂ = v / (1 − β2^t)
param = param − lr · m̂ / (√v̂ + ε)
the bias correction exists because m and v start at zero, so without it the first few steps would be far too small. now do step 1 by hand and watch something nice happen. m and v are both zero, so
m = 0.1·g m̂ = m / (1 − 0.9) = g
v = 0.001·g² v̂ = v / (1 − 0.999) = g²
so m̂ / √v̂ = g / |g| = 1, for every parameter, whatever its gradient was. the update is just lr.
for our four gradients 30, 30, 20, 10 at lr = 1e-3:
w1 = 1.999 b1 = −0.001 w2 = 2.999 b2 = 0.999 loss = 24.910101
all four moved by exactly 0.001. the one with a gradient of 30 and the one with a gradient of 10 moved the same distance. i checked this to twelve decimals out of suspicion and it is 0.001000000000
for three of them and 0.000999999999 for the fourth, the difference being the ε = 1e-8 in the denominator.
that is the single most useful thing to know about Adam. its learning rate is a step size, not a multiplier on the gradient. in SGD, lr = 1e-3 means "move a thousandth of whatever the gradient
says", so the actual distance depends on gradient scale, which depends on your loss and your data and your initialization. in Adam, lr = 1e-3 means "move about 0.001", full stop. that is why 1e-3 is a sane default across wildly different models and why the same value in SGD would be useless on our problem here.
it also means Adam is not automatically faster:
Adam lr=1e-3, 60 steps: loss = 20.04 (still crawling, it moves 0.001 per step)
Adam lr=0.1, 60 steps: loss = 0.0146
plain SGD lr=0.01, 30 steps: loss ≈ 0
on four parameters whose gradients are all roughly the same size, adapting per parameter buys you nothing and the fixed step size just makes it slow. so where does it earn its place.
The experiment where Adam is obviously right
change one thing, set x = 0.01 instead of 1.0. nothing else. now the gradients are
∂L/∂w1 = −0.0564 ∂L/∂b1 = −5.64 ∂L/∂w2 = −0.0376 ∂L/∂b2 = −1.88
w1 and b1 sit in the same layer and their gradients are a hundred times apart, because ∂L/∂w1 = ∂L/∂h1 · x and ∂L/∂b1 = ∂L/∂h1 · 1, and x is 0.01. this is not a contrived setup, it is what a badly scaled input feature does to you, and in a real network the same spread happens between embedding weights, biases and normalization scales.
train it 200 steps both ways and look at how far each parameter actually travelled:
SGD lr=0.01 : w1 moved 0.0028 b1 moved 0.2808 (ratio 100)
Adam lr=0.01 : w1 moved 0.2207 b1 moved 0.2207 (ratio 1)
SGD's single learning rate gets divided among the parameters by their gradient magnitudes, so w1 barely moves at all and effectively is not being trained, while b1 does all the work. Adam moved
both by the identical amount, because dividing by √v̂ cancels the scale of each parameter's own gradient. that is what "per-parameter learning rates" means concretely, and it is the reason it is
the default for anything with heterogeneous parameter groups in it.
both runs reach a near-zero loss here, this network is too small for it to matter. the point is not which one won, the point is what the movement numbers say about where the learning rate went.
The two lines, and the order
whichever one you pick, the loop is the same three calls and swapping optimizer changes nothing else:
for epoch in range(200):
pred = model(x)
loss = loss_fn(pred, y)
optimizer.zero_grad() # clear last step's gradients
loss.backward() # compute this step's gradients
optimizer.step() # apply the update rule
the order is load-bearing. step() before backward() updates the parameters with last step's gradients, and the loop still runs and the loss still prints, it is just always one step stale.
skipping zero_grad() is worse, gradients accumulate rather than overwrite, which i wrote about in the autograd piece, and the effective step grows every iteration until the thing diverges. neither of those raises an error. they just train badly, which is the pattern with most of this.
Try it before you close the tab
take the network at the top and run one SGD step at lr = 0.05 instead of 0.1 or 0.01, and work out h1 before you run it. predict whether the hidden neuron survives. there is a specific learning rate where it dies and you can solve for it in one line, h1 = (2 − 30·lr) + (0 − 30·lr), so it goes
negative at lr = 1/30. anything above 0.0333 kills the neuron on the first step. finding that number yourself is worth more than the rest of this article.
then set momentum = 0.9 and print h1 at every step for ten steps, and watch it cross zero and keep going while the gradient reads0.0.
what actually made you change optimizer the last time you changed it? for me it has mostly been switching to Adam because something would not train and then never switching back, which is not
really a reason, and the honest answer is i did not look at the gradient scales until i started writing this chapter.
This is one chapter's worth of an idea from my book, PyTorch From Ground Up, which builds everything from tensors upward so nothing stays vague. If it helped:
8 chapters are free, no email required,
there's a free one-page tensor cheat-sheet here, every example runs in
the companion notebooks on GitHub, and the
full book is on Leanpub or in
paperback and Kindle on Amazon.
More in this series
How Training Actually Works, the part where the training loop stops being magic:
-
PyTorch Autograd Explained: What
.backward()Actually Does - Backpropagation by Hand: Two Layers, a Pen, and Then Autograd Agrees
The shape mechanics underneath all of it, worth having solid first:
Coming next in How Training Actually Works: what the loss function is actually choosing for you.
