Adam 如何為每個參數調整步長？

One step size for every dial

An AI learns by turning dials. A network has millions of them, its weights, and training nudges every one a little at a time. To know which way to turn a dial, training measures how the mistake changes when that dial moves. That measurement is the dial's push, the gradient. It points downhill, toward a smaller mistake.

Plain gradient descent moves every dial with the same rule:

w ← w − step × push

There is one step size, and every dial shares it. The trouble is that hills differ. Take two dials. Both start at w = −2, and both make the smallest mistake at w = 0.

• The top dial sits on a steep, narrow hill: mistake = 25 · w². Its push at the start is 100.
• The bottom dial sits on a wide, nearly flat hill: mistake = ½ · w². Its push is 2.

Same position, but on the steep hill the push is fifty times bigger.

Two dials, two hills: the same start, but the steep hill's push is fifty times bigger, and one slider sets the step size for both.

Try a big step size, 0.045. The steep dial jumps right past the bottom: from −2 to 2.5, then −3.125, then 3.906. Every jump is bigger than the last, and it flies off. Try a small step size, 0.015, which suits the steep hill: it settles in three hops (−0.5, −0.125, −0.031). Now the flat dial crawls. It needs 245 steps to get within 0.05 of the bottom. No single step size fits both.

With a big step size the steep dial jumps past the bottom, higher each time, and flies off.

Each dial's own step size

The fix is to stop sharing. Each dial remembers how big its pushes usually are, and divides its push by that. A dial with big pushes takes careful steps. A dial with tiny pushes takes bolder ones.

On the very first step, each push divided by its usual size is exactly 1 for both dials (100 out of 100, and 2 out of 2). So both dials move exactly the same distance: 0.1. After that their paths stay identical. They roll down together, overshoot a little to +0.26, swing back, and both stay within 0.05 of the bottom from step 71.

Both dials take the same steps: a line joining them stays vertical as they roll down together.

What Adam computes

This method is Adam, from a paper by Diederik Kingma and Jimmy Ba, published at ICLR 2015. The name comes from "adaptive moment estimation". It combines two ideas: momentum, from the last episode, and its own step size for every dial.

For each dial, g is its push at this step and t counts the steps (1, 2, 3, …). Adam keeps two running averages and turns them into a step:

• m = 0.9 · m + 0.1 · g (the average push: this is momentum)
• v = 0.999 · v + 0.001 · g² (the average squared push: how big the pushes usually are)
• m̂ = m / (1 − 0.9ᵗ) and v̂ = v / (1 − 0.999ᵗ) (m and v start at zero, so early on they come out too small; this corrects them)
• step = lr · m̂ / (√v̂ + 1e-8)

The 1e-8 is only there to avoid dividing by zero. The key part is the division by √v̂, which has the same size as the push. If a dial's pushes are all fifty times bigger, its m̂ and its √v̂ are both fifty times bigger, and they cancel. That is why the steep dial and the flat dial move the same.

The steeper-or-flatter test

Now make both hills a hundred times steeper, then a hundred times flatter, and keep every setting the same.

• A hundred times steeper: the shared step size of 0.015 blows up on the steep hill. After one step the dial is at 148, after two at −10,952.
• A hundred times flatter: the same shared step size crawls. The flat dial needs 24,591 steps.
• Adam (step size 0.1) takes 71 steps on the original hills, 71 on the steeper ones and 71 on the flatter ones, on exactly the same path every time.

That is a big part of why Adam is so popular: when a problem's scale changes, you don't have to hunt for a new step size.

The old way blows up or crawls for 24,591 steps; Adam takes the same path every time: 71 · 71 · 71.

Try it in PyTorch

In PyTorch, Adam is one line. A learning rate of 1e-3 is the usual starting point:

opt = torch.optim.Adam(model.parameters(), lr=1e-3)

Here are the two dials from this article, for one step:

import torch as t
w=t.nn.Parameter(-2*t.ones(2))
o=t.optim.Adam([w],lr=.1)
L=25*w[0]**2+w[1]**2/2
L.backward();o.step()
print(w.data)

It prints tensor([-1.9000, -1.9000]). Both dials moved exactly 0.1, even though one push was fifty times bigger. Change 25 to 2500 and run it again: still 0.1.

References

• Kingma, D. P., & Ba, J. (2015). Adam: A Method for Stochastic Optimization. ICLR 2015. https://arxiv.org/abs/1412.6980
• CS878 Deep Learning, SEECS NUST, lectures 07–08 (course slides).

This article is part of Deep Learning, Visually by Logit Solutions AI: short animated explainers that cover one idea at a time, each with a written version like this one.