Moving averages and the Adam optimizer
Raw gradients are noisy. Smoothing them with a moving average makes training steadier. Adam adds a per-weight step size.
Words on this page (9)
- gradient
- : The direction and rate at which the loss changes when each parameter changes. Training steps against it.
- training
- : Adjusting a model's parameters on data so its predictions improve.
- Adam
- : The most common optimizer. It smooths gradients and adapts the step size for each parameter.
- weights
- : The learned numbers that define a trained model. A model whose weights are published is called open-weight.
- learning rate
- : How big each gradient descent step is.
- optimizer
- : The algorithm that turns gradients into parameter updates, such as SGD or Adam.
- gradient descent
- : Repeatedly nudging parameters opposite to the gradient to reduce the loss.
- loss
- : A number measuring how wrong the model's output is. Training tries to reduce it.
- logarithm
- : The power you raise a base to in order to get a number. log base 2 of 8 is 3 because 2 to the power 3 is 8.
Exponential moving average
An EMA blends the old average with the new value: m ← β·m + (1 - β)·g, where β is near 1, such as 0.9. Recent values weigh more, old values fade. With β = 0.9 and 1, 1, 1, starting from m = 0: m = 0.1, 0.19, 0.271. It starts low, a bias that corrects by dividing by (1 - βᵗ). At step 1 the average is 0.1 and the correction divides by 1 - 0.9 = 0.1, giving 1.0, the true gradient. At step 2 it is 0.19 / 0.19 = 1.0 again.
Adam
Adam keeps two EMAs per : the average of gradients (momentum, direction) and the average of squared gradients (scale). The step is × momentum / (square root of scale + a tiny constant). Weights with consistently large gradients take smaller steps and weights with small gradients take relatively bigger ones. It is the default for transformers, usually with weight decay (AdamW), which gently shrinks weights toward zero.
Learning rate schedules
A common recipe is warmup (rise from near zero over the first steps) then decay (fall, often along a cosine curve). Large models are sensitive to this, since a rate that is too high early can wreck .
Worked example: smoothing
Daily gradients for one weight: 1.0, -0.8, 1.2, 0.9. A raw update would zigzag. With beta 0.9, momentum builds as 0.1×1.0 = 0.10, then 0.9×0.10 + 0.1×(-0.8) = 0.01, then 0.9×0.01 + 0.1×1.2 = 0.129, then 0.9×0.129 + 0.1×0.9 = 0.206. The smoothed value changes gently. It starts low because it begins at zero, so real optimizers divide by a correction factor (here 1 − 0.9^4 = 0.344, giving 0.206 / 0.344 = 0.60, close to the plain average of 0.575).
Why Adam is the default
Different weights need different step sizes. Adam divides each step by how large that weight's recent gradients have been, so noisy weights move carefully and quiet ones move more. It works well out of the box, which is why most transformers use it.
How one Adam step is built
Adam tracks two averages per weight, m for gradients and v for squared gradients. A typical choice is β₁ = 0.9 for m and β₂ = 0.999 for v. After bias correction on step 1, m̂ equals the gradient g and v̂ equals g², so the step is the learning rate × g / (|g| + a tiny number). That is about plus or minus the learning rate, whatever g was. Check: with g = 0.5, 50 and 0.001 the ratio g / |g| is 1 in all three cases. The first steps all have about the same size, even when gradient scales differ by a factor of 50,000. That is the "per-weight step size" idea in action.
A second example: a weight that zigzags
Gradients +1, -1, +1, -1 for one weight, with β = 0.9. The raw average m runs 0.1, -0.01, 0.091, -0.0181. After dividing by 1 - 0.9ᵗ (0.1, 0.19, 0.271, 0.3439) you get 1.0, -0.053, 0.336, -0.053. A plain gradient step would swing by ±1 every time. The smoothed signal says: mostly cancelling, barely move. Compare a steady gradient of 1, where the corrected value stays at exactly 1 and the weight moves at full speed.
Trade-offs
Adam stores two extra numbers for every weight. Invented size: a model with 1 billion weights stored in 32 bits (4 bytes each) needs 4 GB for the weights, 4 GB for gradients and 8 GB for the two Adam averages, 16 GB in total before any activations. Plain with no memory needs only 8 GB for weights plus gradients. This extra cost is one reason people use lower-precision storage and other tricks. Adam is also not always best: it needs less tuning, but a well-tuned simpler method can generalize as well on some tasks.
What goes wrong, and how to tell
Symptom: the spikes up and recovers, or turns to nan, usually early. Likely cause: learning rate too high for the model size, so try a longer warmup. Another symptom: loss falls fast then plateaus above where it should be, often a rate that decayed too early. To confirm, the learning rate next to the loss and see whether the spike lines up with the rate peak.
Common mistake
Using a high learning rate from step one on a large model. Warmup, starting low and rising, prevents early blow-ups.
When you check your weight every morning, the number jumps around by a pound or two from water and meals. A smarter reading is the average of the last week, which shows the real trend. Training does the same with its noisy gradient signals: Adam keeps a moving average so one odd batch does not send the model off course.
Check yourself
In an exponential moving average with beta 0.9, which values count the most?
Not quite. Older values fade geometrically, so recent values weigh more.
Why do large models use learning-rate warmup?
Not quite. The rate rises slowly until training is stable.
Gradients for one weight bounce between +1 and -1 every step. What does Adam's moving average do with this?
Not quite. Averaging noisy gradients cancels the swings, so the update follows the overall trend.
You got 0 of 3 right.
Go deeper
- ArticleWhy Momentum Really WorksArticle
Gabriel Goh · An interactive explanation of momentum.
- Paper
Kingma and Ba · The original Adam paper.
- Video
DeepBean · Each optimizer explained with animations.
Spotted a mistake or something unclear? Report it on GitHub.