Adam from Scratch — Two Moving Averages and the Correction Nobody Explains

Adam is the default optimizer in almost every deep-learning codebase, and almost nobody can say what its four lines do. This article assembles it from the two problems it solves — the zigzag of plain SGD and the mismatched gradient scales between layers — then takes the bias correction apart with numbers: without it, Adam's very first step is 3.2× too large and the error peaks at 6.6×. Ends on the traps: weight decay belongs in AdamW, not Adam, and the epsilon sits outside the square root for a reason.

Adam: A Method for Stochastic Optimization

Where plain SGD runs out of road

Gradient descent is one line. You stand at parameters $w$, you compute the gradient $g_t = \nabla L(w_{t-1})$ — the direction in which the loss climbs fastest — and you step the other way, scaled by a learning rate $\eta$:

$$w_t = w_{t-1} - \eta\, g_t$$

In words: look at which way is uphill, and step downhill by $\eta$ times the steepness. That is the whole of SGD, and if you want it built up from the ground the place to start is Gradient Descent from Scratch. Here we only need the two ways it fails, because Adam is nothing more than one fix for each.

Failure one: the zigzag. Real loss surfaces have ravines — directions in which the surface is steeply walled, crossed with directions in which it is nearly flat. On a canyon wall the gradient points mostly across the canyon, toward the near wall, and barely at all along it, toward the exit. So SGD bounces wall to wall, crawling down the gentle slope it actually needs to descend. The steep direction demands a small $\eta$ to avoid overshooting; the flat direction wants a large one to make any progress at all. One number cannot serve both.

Failure two: mismatched scales. In a deep network, gradients in one layer routinely arrive a hundred or a thousand times larger than gradients in another — different layer widths, different activations, different distances from the loss. A single global $\eta$ that is sane for the large-gradient layer is invisible to the small-gradient one, and vice versa. What you want is a per-parameter step size, set automatically.

Adam (Adaptive Moment estimation, Kingma & Ba 2014) solves failure one by remembering the direction of recent gradients, and failure two by remembering their size. Both memories are the same cheap device: an exponential moving average.

The device: an exponential moving average

Every piece of Adam is built from one recurrence. Given a stream of numbers $x_1, x_2, \dots$ and a decay rate $\beta$ between 0 and 1, the exponential moving average (EMA) is

$$a_t = \beta\, a_{t-1} + (1 - \beta)\, x_t, \qquad a_0 = 0.$$

Read it as a sentence: keep $\beta$ of what you already believed, and mix in $1-\beta$ of what you just saw. Unrolling it shows what it really is — a weighted average of the whole history, with weights that shrink geometrically into the past:

$$a_t = (1-\beta)\left(x_t + \beta x_{t-1} + \beta^2 x_{t-2} + \dots\right).$$

The useful intuition is that an EMA with decay $\beta$ has an effective memory of roughly $1/(1-\beta)$ recent values. So $\beta = 0.9$ remembers about the last 10 steps, and $\beta = 0.999$ about the last 1000. That single fact will explain both of Adam's default decay rates and, later, why one of them needs a correction that lasts for hundreds of steps.

Note the initialisation: $a_0 = 0$. We start with no history, so we start from zero. That choice is the entire origin of the bias problem two sections from now — hold on to it.

First moment: remembering the direction

Apply the EMA to the gradients themselves. This is the first moment estimate, $m_t$ — an estimate of the mean gradient:

In words: the direction we trust is mostly the direction we trusted a moment ago, nudged 10% toward the gradient we just saw. Piece by piece: $m_{t-1}$ is what we believed the typical gradient direction was before this step; $g_t$ is the gradient we just measured on this mini-batch; $\beta_1$ decides how much of the old belief survives. Stepping along $m_t$ instead of along $g_t$ is momentum, and it fixes the zigzag directly. The side-to-side components of the ravine bounce keep reversing sign, so in the running average they cancel. The down-the-canyon component has the same sign every step, so it accumulates. The optimizer behaves like a heavy ball: it stops rattling between the walls and keeps its downhill roll.

Momentum also smooths the mini-batch noise. Each $g_t$ is a noisy estimate of the true gradient, computed from a handful of examples; averaging ten of them gives a steadier direction than trusting any one.

Second moment: remembering the size

Now apply the same EMA to the squared gradients. This is the second moment estimate, $v_t$:

In words: keep a slow-moving record of how big this parameter's gradients usually are, ignoring which way they point. The squaring is elementwise — each parameter keeps its own $v$, tracking how large that parameter's gradients have recently been, regardless of sign. So $\sqrt{v_t}$ is a running estimate of the typical magnitude of that coordinate's gradient (its root-mean-square), which is where the name RMSProp comes from: this line is exactly RMSProp, the method Adam inherits it from.

The plan is to divide the step by $\sqrt{v_t}$. Think about what that does. A coordinate whose gradients are consistently huge gets a large $\sqrt{v_t}$ and therefore a small step; a coordinate whose gradients are consistently tiny gets a small $\sqrt{v_t}$ and a large step. The division cancels the scale of the gradient and leaves only its sign and consistency — which is precisely the per-parameter learning rate that failure two was asking for. Layers with mismatched gradient magnitudes end up taking comparable steps without anyone tuning them separately.

Bias correction: the first steps are lying to you

We initialised both averages at zero. That was an honest choice — we had no history — but it makes the early estimates wrong in a known direction, and the fix is the one part of Adam that gets skipped in every casual explanation.

Watch the very first step. Suppose the gradient is $g_1 = 1$. Then

$$m_1 = \beta_1 \cdot 0 + (1 - \beta_1) \cdot 1 = 1 - \beta_1 = 0.1.$$

We measured a gradient of 1 and our estimate of the typical gradient is 0.1. It is not that the estimate is noisy; it is that it has been dragged 90% of the way toward the fictitious zero we seeded it with. The same thing happens to the second moment, only far worse, because $\beta_2$ is closer to 1: $v_1 = (1 - \beta_2) g_1^2 = 0.001$.

The size of the lie is exactly computable. Take the unrolled EMA with every $x_t$ equal to the same value $x$. The weights are $(1-\beta)(1 + \beta + \beta^2 + \dots + \beta^{t-1})$, a geometric series that sums to $1 - \beta^t$ rather than to 1. So

$$\mathbb{E}[m_t] \approx (1 - \beta_1^t)\, \mathbb{E}[g_t], \qquad \mathbb{E}[v_t] \approx (1 - \beta_2^t)\, \mathbb{E}[g_t^2].$$

In words: after $t$ steps, each average holds only a $1-\beta^t$ fraction of the true value; the rest is still the zero we started from. The estimate is shrunk by a factor we know. Divide it back out, and the bias is gone. Those two divisions are the bias-corrected moments:

In words: scale each average back up by exactly the amount the zero start shrank it. Check the first step again: $\hat{m}_1 = 0.1 / (1 - 0.9) = 0.1/0.1 = 1$, which is the gradient we actually measured. And $\hat{v}_1 = 0.001/(1 - 0.999) = 0.001/0.001 = 1 = g_1^2$. The correction is exact at $t=1$ and fades as $\beta^t \to 0$, so it matters early and disappears later — which is exactly the shape of the problem.

Now, does it matter in practice? It is tempting to argue that both moments are shrunk, so the ratio $m_t/\sqrt{v_t}$ cancels the error out. It does not, because the two are shrunk by different amounts: the numerator carries $1-\beta_1^t$ and the denominator carries $\sqrt{1-\beta_2^t}$. The ratio of the uncorrected step to the correct one is

$$\frac{1 - \beta_1^{\,t}}{\sqrt{1 - \beta_2^{\,t}}}.$$

In words: the numerator recovers from its zero start within about 10 steps, the denominator takes about 1000, so for a long stretch you are dividing a nearly-honest number by a still-shrunken one. Feed a constant gradient $g = 1$ into Adam with the standard $\beta_1 = 0.9$, $\beta_2 = 0.999$ and the corrected update takes a step of exactly $\eta$ from step 1 onward — the right answer. The uncorrected update takes 3.16× that on step 1, rises to a peak of 6.57× around step 12, is still 3.2× too large at step 100 and 1.4× too large at step 693, and only settles within 1% of correct after roughly 5000 steps. Skipping bias correction does not cost you a couple of shaky steps; it hands you thousands of steps of inflated learning rate, at exactly the moment when the parameters are furthest from any minimum. That is why so much early Adam folklore insisted on a warmup — a good part of what warmup was patching is this.

The long tail comes from $\beta_2 = 0.999$: the factor $1 - \beta_2^t$ only reaches 0.5 at $t \approx 693$. The second moment has a memory of a thousand steps, so its start-up transient lasts a thousand steps too.

The update rule, assembled

Four lines per step, in order:

Read the last line as one sentence. $\hat{m}_t$ is the direction, smoothed over recent history. $\sqrt{\hat{v}_t}$ is the typical size of this coordinate's gradient. Dividing the first by the second produces a signal-to-noise ratio: large when recent gradients agree in sign (so the mean is nearly as big as the RMS, and the ratio approaches 1), small when they thrash and cancel (mean near zero, RMS still large, ratio near 0). Adam steps confidently where the gradient is consistent and timidly where it is contradictory, one coordinate at a time. And $\epsilon$ is a tiny constant that stops the division from exploding when a coordinate's gradients are all ~0.

The bound is worth noticing: because the ratio $|\hat{m}_t| / \sqrt{\hat{v}_t}$ is at most about 1 in the well-behaved case, a single Adam step moves each parameter by at most about $\eta$, whatever the gradient's units. That is why an Adam learning rate transfers between models far better than an SGD one — $\eta$ is a step size in parameter space, not a scale factor on an unknown gradient.

The defaults, and what each one is for

The paper's defaults are $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$, and they are the values in your framework right now. Each has a reading.

$\beta_1 = 0.9$ gives the direction a memory of about $1/(1-0.9) = 10$ steps. Long enough to cancel the ravine wobble and average out mini-batch noise, short enough to turn when the surface genuinely turns. Raising it to 0.95 or 0.98 (common for large language models, where batches are huge and each gradient is already low-noise) buys a smoother trajectory at the cost of slower response.

$\beta_2 = 0.999$ gives the scale a memory of about 1000 steps. Scale should be estimated over a much longer window than direction, because it is a statistic about the distribution of gradients rather than a heading, and a jumpy denominator would make the step size jitter. The price is the thousand-step start-up transient from the previous section. Practitioners drop it to 0.95–0.98 when training is unstable and gradient magnitudes are shifting fast.

$\epsilon = 10^{-8}$ is numerical insurance for the division. It matters more than its size suggests, because it also silently caps the step: when a coordinate's $\sqrt{\hat{v}_t}$ falls well below $\epsilon$, the ratio becomes $\hat{m}_t/\epsilon$-limited rather than scale-free, and that coordinate stops getting its adaptive boost. Very small $\epsilon$ makes flat coordinates take enormous steps; larger values ($10^{-6}$, and $10^{-4}$ or higher in some mixed-precision and RL setups) make Adam behave more like SGD with momentum. It is not a dummy constant.

And $\eta$ is still the dial that matters most. $10^{-3}$ is the classic default; fine-tuning a transformer usually wants $10^{-5}$ to $5 \times 10^{-5}$. Adam adapts the relative step between parameters; it does not choose the absolute size for you.

Seeing it in code

The whole optimizer is a dozen lines. Here it is minimizing the quadratic bowl $L(w) = w^2$, whose gradient is $2w$, starting from $w = 5$:

import numpy as np

def grad(w):  return 2 * w          # dL/dw for L = w^2

w = 5.0
m = v = 0.0                          # both moments start at zero
beta1, beta2, eta, eps = 0.9, 0.999, 0.1, 1e-8

for t in range(1, 201):              # t starts at 1 — the correction needs it
    g = grad(w)
    m = beta1 * m + (1 - beta1) * g          # 1st moment: direction
    v = beta2 * v + (1 - beta2) * g * g      # 2nd moment: size
    m_hat = m / (1 - beta1 ** t)             # bias correction
    v_hat = v / (1 - beta2 ** t)
    w -= eta * m_hat / (np.sqrt(v_hat) + eps)
    if t <= 5 or t == 200:
        print(f"step {t:3d}:  w = {w:+.4f}")

Running it gives w = +4.9000, +4.8001, +4.7002, +4.6005, +4.5010 and then w = -0.0001 at step 200. Look at the early steps: each one moves $w$ by almost exactly $0.1 = \eta$, regardless of the fact that the gradient there is 10. That is the scale-free property, visible in five numbers.

Now delete the two correction lines and use m and v directly. The first step becomes w = +4.6838 — a move of 0.3162, which is $3.16\eta$, the exact inflation factor predicted in the bias-correction section. The theory and the print statements agree to four digits.

Pitfalls

Weight decay in Adam is not weight decay — use AdamW. The textbook way to add L2 regularization is to append $\lambda w$ to the gradient. Do that inside Adam and the penalty goes through the $1/\sqrt{\hat{v}_t}$ division along with everything else, so parameters with large gradients get less decay than parameters with small ones — the strength of your regularizer becomes an accident of gradient magnitude. Loshchilov & Hutter's fix (AdamW, 2017) is to decouple it: apply the adaptive step, then subtract $\eta \lambda w$ separately, outside the division. This is not a micro-optimization; the paper reports that it substantially improves Adam's generalization, letting it compete with SGD with momentum on image classification, and it is the reason essentially every modern language model is trained with AdamW. Note also that PyTorch's Adam(weight_decay=...) is the coupled, worse version — reach for AdamW explicitly.

Adam converges faster but can generalize worse. On many vision tasks, well-tuned SGD with momentum still reaches a better test accuracy than Adam, even though Adam gets the training loss down sooner. The usual explanation is that the adaptive denominator suppresses exactly the gradient noise that pushes SGD out of sharp minima and into flatter, better-generalizing ones. Practical reading: for a transformer or anything with wildly uneven gradient scales, Adam/AdamW is the default and SGD often does not train at all; for a convnet on a fixed benchmark where you can afford to tune, SGD+momentum is still worth a run.

The epsilon goes inside, next to the square root — not under it. The correct denominator is $\sqrt{\hat{v}_t} + \epsilon$. Writing $\sqrt{\hat{v}_t + \epsilon}$ is a different algorithm: the second form adds $\epsilon$ in units of squared gradient and so kicks in at a completely different threshold ($\sqrt{10^{-8}} = 10^{-4}$ rather than $10^{-8}$). Both forms appear in real implementations, so if you port an $\epsilon$ value between frameworks, check which convention you are landing in before trusting it.

Adam costs memory. Two extra float buffers per parameter, each the size of the model. For a 7B-parameter model in fp32 that is roughly 56 GB of optimizer state on top of the weights and gradients, which is why 8-bit optimizers, Adafactor, and sharded states exist. SGD with momentum needs one such buffer; plain SGD needs none.

$t$ must be the global step, and it must start at 1. At $t=0$ the correction divides by $1 - \beta^0 = 0$. If you reset $t$ when resuming from a checkpoint, you reintroduce the whole start-up transient — a silent 3–6× learning-rate spike right after every restart.

In practice

Who touches this, and when. An ML engineer fine-tuning a pretrained model, writing the training config. A researcher training a new architecture from scratch, chasing down why the loss spikes or goes NaN early on. In both cases, knowing which of Adam's four lines is doing the work lets you pick settings for a reason instead of by incantation.

The names you actually type. In PyTorch it is torch.optim.AdamW(params, lr=..., betas=(0.9, 0.999), eps=1e-8, weight_decay=...). Hugging Face's TrainingArguments exposes the same knobs as learning_rate, adam_beta1, adam_beta2, adam_epsilon and weight_decay. When memory is tight, the 8-bit variant from bitsandbytes (AdamW8bit) is the usual candidate. The optimizer's state lives in optimizer.state_dict(), where each parameter has a step, an exp_avg ($m_t$) and an exp_avg_sq ($v_t$). The symbols in this article are literally the key names, so while debugging you can open the state and look for, say, a layer whose $v_t$ is essentially zero.

Choosing what gets weight decay. In practice you rarely pass params as one flat list. The standard move is to split it into parameter groups: one-dimensional parameters such as biases and LayerNorm weights go into a group with weight_decay=0.0, and decay is applied only to the weight matrices. The usual reasoning is that these set scales and offsets, so pulling them toward zero does little against overfitting and mostly just shrinks what the layer can express.

Traps that cause real incidents.

• Checkpointing only the weights. On resume, exp_avg, exp_avg_sq and step restart from zero, and you go through the whole bias-correction start-up again. Save optimizer.state_dict() (and the scheduler's state) alongside the model, and load it back when you resume.
• Using $\epsilon = 10^{-8}$ in fp16. The smallest positive fp16 number (a subnormal) is about $6 \times 10^{-8}$, so $10^{-8}$ rounds to zero. In an implementation that keeps optimizer state in fp16, $\epsilon$ vanishes and the division can blow up on near-zero gradients. That is why mixed-precision training normally keeps the optimizer state in fp32 — and why you should check when you write your own or use a low-precision optimizer.
• Fine-tuning with the default lr=1e-3. As the defaults section showed, Adam only equalizes the relative step. To avoid wrecking pretrained weights, lowering $\eta$ is your job.

Questions that come up in interviews and code review.

• "What's the difference between Adam and AdamW?" → Whether the decay term passes through the $1/\sqrt{\hat{v}_t}$ division. A complete answer adds that when it does, parameters with large gradients get weaker decay.
• "What happens if you drop bias correction?" → The early effective learning rate is inflated by $(1-\beta_1^t)/\sqrt{1-\beta_2^t}$; numerator and denominator recover on different timescales, so they do not cancel.
• "The loss suddenly spiked mid-run. What do you suspect on the optimizer side?" → Optimizer state not restored on resume; a $\hat{v}_t$ so small that $\epsilon$ no longer protects the division; or a $\beta_2$ so large that the denominator cannot keep up with a sudden change in gradient size.

Summary

• Adam = momentum (remember the gradient's direction) + RMSProp (remember its size), each an exponential moving average, plus a bias correction
• $m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t$ cancels the ravine zigzag; $v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2$ gives every parameter its own effective learning rate
• Both averages start at 0 and are therefore shrunk by $1-\beta^t$; dividing that out gives $\hat{m}_t, \hat{v}_t$. Skipping it inflates the first step by 3.16×, peaks at 6.57× near step 12, and takes ~5000 steps to fade
• The update $w_t = w_{t-1} - \eta\,\hat{m}_t/(\sqrt{\hat{v}_t}+\epsilon)$ is a signal-to-noise ratio: bold where gradients agree, cautious where they contradict, and bounded by roughly $\eta$ per step
• Defaults $\beta_1=0.9$, $\beta_2=0.999$, $\epsilon=10^{-8}$ correspond to memories of ~10 and ~1000 steps; $\epsilon$ is a real hyperparameter, not decoration
• Use AdamW for weight decay, budget two extra buffers per parameter, and don't assume Adam generalizes as well as tuned SGD

For where the one-line update comes from, see Gradient Descent from Scratch; for the wider family of optimizers Adam sits in, Beyond SGD; and for the schedules that ride on top of $\eta$, Learning-Rate Schedules.

Exercises

1. With $\beta_1 = 0.9$, a gradient stream that is exactly $g_t = 2$ for every step. Compute $m_1$, $m_2$, $m_3$ and the corrected $\hat{m}_1$, $\hat{m}_2$, $\hat{m}_3$. What does the corrected sequence tell you?
2. Two parameters receive constant gradients: $g^{(a)} = 100$ every step and $g^{(b)} = 0.01$ every step. After the transient has passed, how far does each move per Adam step with $\eta = 10^{-3}$? How far would each move under plain SGD with the same $\eta$?
3. A colleague reports that removing bias correction "made no difference because it cancels in the ratio." Give the exact expression for the inflation factor, evaluate it at $t = 1$ and $t = 100$ with the default betas, and say what he would have to change about $\beta_2$ for his claim to be nearly true.
4. You add weight_decay=0.01 to torch.optim.Adam. Which layer ends up more strongly decayed — an embedding layer whose gradients are tiny, or an output head whose gradients are large? Work it out from the update rule, then name the one-word fix.