Plain SGD updates:
Adam adds memory.
For each parameter it tracks:
These answer two questions: 1. Has the gradient consistently pointed this way? 2. How large are gradients for this parameter normally?
Follow one parameter¶
Suppose:
and gradients are:
Use easy-to-follow values:
Start:
At step 1:
Bias correction:
Adam update:
Ignoring \(\epsilon\), the normalized term is:
What does v do?¶
Suppose: - parameter A usually has gradients ~0.001 - parameter B usually has gradients ~10
Plain SGD naturally moves B vastly more.
Adam divides by each parameter’s historical gradient scale, making updates more comparable.
What does m do?¶
Consistent gradients:
build persistent momentum.
Oscillating gradients:
tend to cancel in the running average.
Memory cost¶
For 70B parameters, two FP32 moment tensors cost:
each, or:
together.
This is why optimizer sharding matters so much in large-scale training.
Why Adam is so popular for Transformers¶
Transformer parameters can live in very different numerical regimes. One projection matrix may routinely see large gradients while another sees tiny ones.
A single global SGD learning rate has to compromise across all of them.
Adam’s per-parameter normalization makes optimization substantially more forgiving.
Bias correction matters early¶
Because \(m_0=v_0=0\), the raw moving averages are biased toward zero during the first few steps.
The corrections:
compensate for that startup bias.
This connects to learning-rate warmup: the optimizer’s statistical state is itself settling during early training.
AdamW and weight decay¶
Modern Transformer training often uses AdamW, which separates weight decay from the adaptive gradient update.
The conceptual goal of weight decay is to discourage weights from drifting unnecessarily large, providing a regularizing pressure.
The distinction to retain:
- loss tells us what prediction was bad,
- backprop computes sensitivities,
- optimizer decides how those sensitivities translate into parameter changes.
Adam's epsilon and numerical stability¶
The denominator:
contains a small \(\epsilon\) to avoid division by zero and stabilize updates when squared-gradient estimates are tiny.
That constant looks insignificant, but in low-gradient regimes it can materially affect the effective step size.
Optimizer state is persistent training memory¶
Activations live only for one forward/backward pass.
Adam's \(m\) and \(v\) persist across the whole training run.
So the optimizer carries a form of long-term numerical memory about each parameter's gradient history.
This is why resuming training from a checkpoint often means saving more than model weights. To resume the same optimization trajectory faithfully, systems may also save: - optimizer moments, - learning-rate scheduler state, - gradient-scaling state, - random-number generator state.
A “model checkpoint” for inference and a “training checkpoint” are therefore very different artifacts.
A second Adam step¶
Continue the toy parameter from the first step. Suppose:
With \(\beta_1=0.5\):
With \(\beta_2=0.9\):
Bias correction gives:
and:
So Adam's next update depends not only on \(g_2\), but on the history encoded in both moments. That is the essential difference from plain SGD.
Why this makes checkpoint files huge¶
For a model with \(N\) parameters, an inference checkpoint may need only the weights. A resumable Adam training checkpoint can additionally contain two optimizer tensors and sometimes master-precision weights. The on-disk training state can therefore be several times larger than the deployable model.
This is another place where “the model” and “the state required to train the model” are very different objects.