Adam: Two Memories and Bias Correction

September 2026 · 5 min read + interactive lab

Build Adam from its first and second moments, derive why bias correction divides by $1-\beta^t$, and see why AdamW decouples weight decay.

Adam combines Momentum (an exponential moving average of signed gradients $m_t$) with RMSprop (an exponential moving average of squared gradients $v_t$), plus a bias correction that removes the startup shrinkage caused by initializing both buffers at zero.

1. Two memories with different jobs

Starting from $m_0=0$ and $v_0=0$, each step $t$ updates two coordinate-wise exponential moving averages of the gradient $g_t$ (typically with $\beta_1=0.9$, $\beta_2=0.999$):

Here $m_t$ tracks signed direction (opposite signs cancel), while $v_t$ tracks squared magnitude (always nonnegative). Adam first divides each buffer by $1-\beta^t$ to correct for zero initialization, then takes a normalized step:

2. Why zero initialization needs bias correction

Unrolling $m_t$ from $m_0=0$ writes $m_t$ as a weighted sum of past gradients $g_1,\dots,g_t$:

Why don’t the $g_k$ terms block the geometric series? In $\sum_{k=1}^t\beta_1^{t-k}g_k$, each $g_k$ varies from step to step, so we first take the expected value $\mathbb E[m_t]$ (assuming the gradients have stationary mean $\mathbb E[g_k]=\mu$, or equivalently a steady gradient $g_k=\mu$). Because $\mathbb E[g_k]=\mu$ is the same at every step $k$, $\mu$ factors out in front of the sum, leaving only the powers of $\beta_1$ inside the geometric series $1+\beta_1+\cdots+\beta_1^{t-1}=\frac{1-\beta_1^t}{1-\beta_1}$:

By the exact same step for $v_t=(1-\beta_2)\sum_{k=1}^t\beta_2^{t-k}g_k^2$ with $\mathbb E[g_k^2]=\sigma^2$, $\sigma^2$ factors out to give $\mathbb E[v_t]=(1-\beta_2^t)\sigma^2$. Because $m_0=0$ and $v_0=0$ contributed the missing $\beta_1^t$ and $\beta_2^t$ weight, both raw buffers are biased toward zero; dividing by $1-\beta_1^t$ and $1-\beta_2^t$ cancels that factor:

As $t\to\infty$, $\beta_1^t\to 0$ and $\beta_2^t\to 0$, so $1-\beta_1^t\to 1$ and $1-\beta_2^t\to 1$ and the bias correction naturally fades away.

3. What the ratio $\widehat m_t/\sqrt{\widehat v_t}$ actually measures

Since $\widehat m_t\approx\mathbb E[g]$ estimates the mean gradient (signal) and $\widehat v_t\approx\mathbb E[g^2]=(\mathbb E[g])^2+\operatorname{Var}(g)$ estimates the uncentered second moment, the Adam ratio acts as a coordinate-wise signal-to-noise normalizer ($\epsilon\approx 0$):

• Steady gradient ($\operatorname{Var}(g)\ll(\mathbb E[g])^2$): $\widehat m_t/\sqrt{\widehat v_t}\approx\operatorname{sign}(\mathbb E[g])=\pm 1$, so the coordinate moves by $\approx\pm\eta$ whether its gradient magnitude is $10^{-3}$ or $10^3$.
• Noisy or oscillating gradient ($\operatorname{Var}(g)\gg(\mathbb E[g])^2$): opposite signs cancel in $\widehat m_t$ while squares add in $\widehat v_t$, so $|\widehat m_t|\ll\sqrt{\widehat v_t}$ and the step automatically shrinks.

4. Lab: inspect what the update actually divides

On the shared quadratic dataset, toggle bias correction on and off to compare the raw moments $(m_t,v_t)$ against the corrected values $(\widehat m_t,\widehat v_t)$ used in the step.

5. Why AdamW separates weight decay from $\ell_2$ regularization

Adding an $\ell_2$ penalty $\tfrac\lambda2\|w\|^2$ to the loss changes the gradient to $g_t^{\text{reg}}=g_t+\lambda w_{t-1}$. In standard Adam, that penalty enters $m_t$ and gets divided by $\sqrt{\widehat v_t}+\epsilon$, whereas AdamW applies weight decay directly to $w_{t-1}$ outside the adaptive ratio:

Notice that the adaptive multiplier $\frac{\eta}{\sqrt{\widehat v_t}+\epsilon}$ behaves the exact same way as in Adagrad and RMSprop—it boosts sparse/tiny-gradient coordinates ($\sqrt{\widehat v_t}\ll 1$) and reduces dense/large-gradient coordinates ($\sqrt{\widehat v_t}\gg 1$). That is helpful for the task gradient $\widehat m_t$, but harmful when it also multiplies the $\ell_2$ penalty $\lambda w_{t-1}$:

• Sparse / tiny-gradient weights ($\sqrt{\widehat v_t}\ll 1$, multiplier is boosted): the $\ell_2$ penalty term $\left(\frac{\eta}{\sqrt{\widehat v_t}+\epsilon}\right)\lambda w_{t-1}$ gets amplified, over-decaying rare-feature weights toward $0$.
• Dense / large-gradient weights ($\sqrt{\widehat v_t}\gg 1$, multiplier is reduced): the $\ell_2$ penalty term $\left(\frac{\eta}{\sqrt{\widehat v_t}+\epsilon}\right)\lambda w_{t-1}$ gets dampened, leaving high-gradient weights almost unregularized.

AdamW fixes this by decoupling weight decay from $\sqrt{\widehat v_t}$: every weight decays uniformly by $\eta\lambda w_{t-1}$, while the adaptive multiplier $\frac{\eta}{\sqrt{\widehat v_t}+\epsilon}$ is used purely for the loss gradient $\widehat m_t$.

6. Implementation and state memory

# One AdamW step (moments use only the task gradient)
t += 1
m = beta1 * m + (1 - beta1) * grad
v = beta2 * v + (1 - beta2) * grad.square()
m_hat = m / (1 - beta1 ** t)
v_hat = v / (1 - beta2 ** t)
parameter *= 1 - lr * weight_decay
parameter -= lr * m_hat / (v_hat.sqrt() + eps)

optimizer = torch.optim.AdamW(
    model.parameters(), lr=3e-4, betas=(0.9, 0.999),
    eps=1e-8, weight_decay=0.01
)

Optimizer state memory: For every model parameter $w_i$, Adam/AdamW stores two state values—the first moment $m_{t,i}$ and second moment $v_{t,i}$—plus the step count $t$. Every weight tensor W of shape $(d_{\text{out}}, d_{\text{in}})$ requires two matching tensors m and v of shape $(d_{\text{out}}, d_{\text{in}})$, adding $+8\text{ bytes/parameter}$ in FP32 ($2\times$ the state memory of SGD with Momentum).

7. Questions worth answering

Sources

Kingma and Ba, Adam: A Method for Stochastic Optimization; Loshchilov and Hutter, Decoupled Weight Decay Regularization; PyTorch AdamW.