Adam and Its Variants

Learning Objectives

By the end of this section, you will be able to:

1. Explain why momentum alone is insufficient when different parameters need different learning rates
2. Derive the RMSprop update and explain how dividing by the running RMS gradient creates per-parameter adaptive learning rates
3. Derive the full Adam algorithm as the combination of momentum (first moment) and RMSprop (second moment) with bias correction
4. Implement Adam from scratch in NumPy and trace the m, v, $\hat{m}$, $\hat{v}$ buffers step by step
5. Use PyTorch's Adam optimizer to train a real network and compare its convergence against SGD with momentum
6. Explain AdamW and why decoupled weight decay is the standard in modern transformer training

Where Momentum Falls Short

In Section 1, we saw that momentum gives every parameter the same treatment: accumulate past gradients with $\beta = 0.9$, then step with learning rate $\eta$. This works brilliantly when the loss surface is smooth and all directions have similar curvature. But real neural networks have a much harder problem.

The Per-Parameter Problem

Consider our diagonal flip network with 31 parameters. Some of these are weights in $\mathbf{W}_1$ (layer 1), some are biases in $\mathbf{b}_1$, and others are in $\mathbf{W}_2$ and $\mathbf{b}_2$. During training:

• Layer 1 weights see gradients that are attenuated by the ReLU and layer 2 weights (chain rule). Typical gradient magnitude: ~0.1
• Layer 2 biases see gradients directly from the loss. Typical gradient magnitude: ~1.0
• Difference: 10× between the smallest and largest gradients among the 31 parameters

With momentum, the learning rate $\eta = 0.05$ applies to all 31 parameters equally. If we choose $\eta$ to work well for the large-gradient parameters (layer 2 biases), it is too small for the small-gradient parameters (layer 1 weights). If we increase $\eta$, the large-gradient parameters overshoot.

The fundamental limitation of momentum: it smooths the gradient direction but cannot adapt the step SIZE per parameter. All parameters share one learning rate. In a network where gradient magnitudes span 10× or 100× across parameters, one learning rate cannot serve them all optimally.

What we need is an optimizer that automatically gives larger steps to parameters with small gradients and smaller steps to parameters with large gradients. This is the idea behind adaptive learning rate methods.

Adaptive Learning Rates: The Key Insight

The core idea is simple: divide each parameter's update by a measure of how large its gradients typically are. If a parameter consistently has large gradients, divide by a large number (take smaller steps). If it has small gradients, divide by a small number (take larger steps).

AdaGrad: The First Attempt (2011)

AdaGrad (Adaptive Gradient) introduced this idea. For each parameter $w_i$, it accumulates the sum of all past squared gradients:

$v_t = v_{t-1} + g_t^2$

Then the update divides by the square root of this sum:

$w_{t+1} = w_t - \frac{\eta}{\sqrt{v_t} + \varepsilon} \cdot g_t$

For a parameter with consistently large gradients (say $g \approx 10$), $v$ grows quickly and $\sqrt{v}$ becomes large, shrinking the effective step. For a parameter with small gradients ($g \approx 0.1$), $v$ stays small and the effective step stays large. Per-parameter adaptation, automatically.

But AdaGrad has a fatal flaw: $v$ only grows, never shrinks. After thousands of steps, $v$ becomes so large that the effective learning rate approaches zero. The optimizer effectively stops learning.

RMSprop: Scaling by Gradient History

RMSprop (Root Mean Square Propagation), proposed by Geoffrey Hinton in 2012, fixes AdaGrad's decay problem with a simple change: replace the running sum of squared gradients with an exponential moving average:

$v_t = \rho \cdot v_{t-1} + (1 - \rho) \cdot g_t^2$

$w_{t+1} = w_t - \frac{\eta}{\sqrt{v_t} + \varepsilon} \cdot g_t$

Here $\rho$ (typically 0.9 or 0.999) is the decay rate. Unlike AdaGrad's sum, $v$ now tracks the recent gradient variance. Old gradient information fades away exponentially, so $v$ stays bounded and the learning rate never dies.

What the Denominator Really Does

The quantity $\sqrt{v_t}$ is an estimate of the root mean square (RMS) of recent gradients. Dividing by it normalizes each parameter's step by its gradient's typical magnitude:

All three parameters end up with an effective step of approximately $\eta$, regardless of their gradient magnitude. RMSprop equalizes the step size across parameters automatically.

RMSprop's Limitation

RMSprop has no momentum. It adapts the magnitude of each step but doesn't smooth the direction. On a noisy loss surface, RMSprop still zig-zags — it just zig-zags with adaptive step sizes. What if we could combine the directional smoothing of momentum with the magnitude adaptation of RMSprop?

Adam: The Best of Both Worlds

Adam (Adaptive Moment Estimation), published by Kingma & Ba in 2015, combines momentum and RMSprop into a single optimizer. It maintains two exponential moving averages:

• First moment $m_t$ — the mean of gradients (like momentum). Smooths direction.
• Second moment $v_t$ — the mean of squared gradients (like RMSprop). Scales magnitude.

The Full Algorithm

Given gradient $g_t$ at step $t$:

Step 1: Update the first moment (momentum)

$m_t = \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t$

Step 2: Update the second moment (RMSprop)

$v_t = \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2$

Step 3: Bias correction

$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}$

Step 4: Update the weights

$w_{t+1} = w_t - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \varepsilon}$

Default hyperparameters: $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\varepsilon = 10^{-8}$, $\eta = 0.001$.

Reading the Update Rule

The update $\eta \cdot \hat{m}_t / (\sqrt{\hat{v}_t} + \varepsilon)$ can be read as:

• $\hat{m}_t$ says: "the average gradient direction is this way"
• $\sqrt{\hat{v}_t}$ says: "the typical gradient magnitude is this big"
• Their ratio $\hat{m}_t / \sqrt{\hat{v}_t}$ is approximately $\pm 1$ when the gradient is consistent, giving a step of approximately $\eta$
• When gradients are noisy (alternating sign), $|\hat{m}_t|$ is small but $\hat{v}_t$ is still large, so the step shrinks — automatic noise dampening

Understanding Bias Correction

Both $m$ and $v$ are initialized to zero. This creates a problem: in the first few steps, they are biased toward zero. Bias correction fixes this.

Why Zero Initialization Causes Bias

At step 1 with $\beta_1 = 0.9$:

$m_1 = 0.9 \times 0 + 0.1 \times g_1 = 0.1 \cdot g_1$

The expected value of $m_1$ is $0.1 \cdot E[g]$, not $E[g]$. It underestimates the true gradient mean by a factor of 10. In general, after $t$ steps:

$E[m_t] = (1 - \beta_1^t) \cdot E[g]$

Dividing by $(1 - \beta_1^t)$ gives $E[\hat{m}_t] = E[g]$ — an unbiased estimate.

The Second Moment Bias Is Much Worse

For $v$ with $\beta_2 = 0.999$, the bias at step 1 is:

$v_1 = 0.001 \cdot g_1^2 \quad \Rightarrow \quad E[v_1] = 0.001 \cdot E[g^2]$

That is 1000× too small. Without correction, $\sqrt{v_1}$ would be $\sqrt{0.001} \approx 0.032$ times the true RMS, making the step ~30× too large. The bias correction divides by $(1 - 0.999^1) = 0.001$, multiplying $v$ by 1000 to recover the correct scale.

Notice: the first moment correction vanishes by step 50, but the second moment correction is still significant at step 1000. This is because $\beta_2 = 0.999$ is so close to 1 that it takes ~1000 steps for $v$ to reach steady state.

Interactive: Adam Step by Step

The visualization below lets you step through Adam on $f(w) = w^2$ starting at $w = 5$. Watch how the raw moments (m, v) slowly ramp up while the bias-corrected moments ($\hat{m}$, $\hat{v}$) stay stable from step 1. The purple (Adam) dot moves in constant steps of ~lr, while the green (Momentum) dot accelerates.

Try these experiments:

1. Press play at default settings (lr=0.01, β&sub1;=0.9, β&sub2;=0.999). Watch momentum race ahead while Adam takes tiny constant steps. On this 1D problem, momentum wins.
2. Increase lr to 0.05: Adam takes bigger steps but still constant-sized. Momentum still wins because velocity builds up.
3. Set β&sub1;=0.0: disables the momentum component. Adam becomes pure RMSprop-with-bias-correction. The steps are still approximately lr because the second moment normalizes the gradient.
4. Set β&sub2;=0.9 (instead of 0.999): the second moment tracks recent gradients more aggressively. Watch how the raw v ramps up faster, while $\hat{v}$ stays approximately the same due to bias correction.

Why momentum wins here: on a 1D problem, there is only one parameter. Adam's per-parameter adaptation has nothing to adapt between. The ratio $\hat{m}/\sqrt{\hat{v}} \approx 1$ always, so the step is just lr. Momentum accelerates because velocity compounds in a consistent direction. Adam's advantage appears when parameters have different gradient scales — which requires at least two dimensions.

Interactive: Optimizer Comparison

The visualization below shows all four optimizers — SGD (red), Momentum (green), RMSprop (blue), and Adam (purple) — racing toward the minimum of the Rosenbrock function: $f(x, y) = (1 - x)^2 + 100(y - x^2)^2$. This function has a famous curved valley (banana shape) where the minimum sits at $(1, 1)$. It is a classic test because the curvature differs enormously along the valley floor versus across the walls.

What to look for:

• SGD (red) barely moves — the learning rate must be tiny to avoid divergence in the steep direction, so progress in the gentle direction is glacial
• Momentum (green) follows the valley but oscillates across the walls, especially at higher learning rates
• RMSprop (blue) makes equal-sized steps in both directions but has no momentum to smooth the path
• Adam (purple) combines the best of both — it follows the valley floor smoothly with adaptive steps in both directions

NumPy: Adam from Scratch

Let's implement Adam from scratch and trace every internal variable. We use the same $f(w) = w^2$ problem from Section 1 so we can directly compare Adam against vanilla SGD (w=4.25 after 8 steps) and momentum (w=2.44 after 8 steps).

1import numpy as np
2
3# ── Minimize f(w) = w² starting at w=5 ──
4# Gradient: df/dw = 2w
5# Minimum: w* = 0
6
7# === Adam Optimizer ===
8w = 5.0
9lr = 0.01
10beta1 = 0.9
11beta2 = 0.999
12eps = 1e-8
13m = 0.0
14v = 0.0
15
16print("=== Adam (lr=0.01, β₁=0.9, β₂=0.999) ===")
17for t in range(1, 9):
18    grad = 2 * w
19    m = beta1 * m + (1 - beta1) * grad
20    v = beta2 * v + (1 - beta2) * grad ** 2
21    m_hat = m / (1 - beta1 ** t)
22    v_hat = v / (1 - beta2 ** t)
23    w = w - lr * m_hat / (np.sqrt(v_hat) + eps)
24    print(f"t={t}: g={grad:.3f}  m̂={m_hat:.3f}  w={w:.4f}")
25
26print(f"\nAfter 8 steps:")
27print(f"  Adam:        w={w:.4f}  loss={w**2:.2f}")
28print(f"  Momentum:    w=2.4421  loss=5.96")
29print(f"  Vanilla SGD: w=4.2539  loss=18.10")

The results reveal Adam's behavior clearly:

Adam is the slowest on this 1D problem. This is not a bug — it is an important property to understand.

Why Adam Is Slow Here (And Why It Doesn't Matter)

On a 1D problem with a consistent gradient, Adam's ratio $\hat{m}/\sqrt{\hat{v}}$ is always approximately 1. So the step is just $\eta = 0.01$ — no acceleration, no adaptation. Momentum, by contrast, builds up velocity: after 8 steps, its effective step is ~4.5× the raw gradient step.

But on a real neural network with hundreds or thousands of parameters at different scales, Adam's per-parameter adaptation is transformative. A parameter with large gradients ($\sqrt{\hat{v}} = 10$) gets an effective lr of $0.001/10 = 0.0001$. A parameter with tiny gradients ($\sqrt{\hat{v}} = 0.001$) gets an effective lr of $0.001/0.001 = 1.0$. Each parameter gets exactly the learning rate it needs.

PyTorch: Training with Adam

Now let's see Adam on a real network. We train the same diagonal flip network from Section 1, comparing SGD+momentum against Adam. Both models start with identical weights. The only difference is the optimizer.

1import torch
2import torch.nn as nn
3
4torch.manual_seed(42)
5
6class FlipNet(nn.Module):
7    def __init__(self):
8        super().__init__()
9        self.layer1 = nn.Linear(4, 3)
10        self.relu = nn.ReLU()
11        self.layer2 = nn.Linear(3, 4)
12    def forward(self, x):
13        return self.layer2(self.relu(self.layer1(x)))
14
15X = torch.tensor([[a, b, c, d] for a in [0, 1]
16                   for b in [0, 1] for c in [0, 1]
17                   for d in [0, 1]], dtype=torch.float32)
18Y = X[:, [3, 2, 1, 0]]
19
20model_mom = FlipNet()
21model_adam = FlipNet()
22model_adam.load_state_dict(model_mom.state_dict())
23
24opt_mom = torch.optim.SGD(model_mom.parameters(),
25                           lr=0.05, momentum=0.9)
26opt_adam = torch.optim.Adam(model_adam.parameters(),
27                             lr=0.001)
28
29for epoch in range(30):
30    loss_m = loss_a = 0.0
31    for i in range(16):
32        pred = model_mom(X[i])
33        lm = torch.mean((pred - Y[i]) ** 2)
34        opt_mom.zero_grad()
35        lm.backward()
36        opt_mom.step()
37        loss_m += lm.item()
38
39        pred = model_adam(X[i])
40        la = torch.mean((pred - Y[i]) ** 2)
41        opt_adam.zero_grad()
42        la.backward()
43        opt_adam.step()
44        loss_a += la.item()
45
46    if epoch % 5 == 0:
47        print(f"Epoch {epoch:2d}: Mom={loss_m/16:.4f}  Adam={loss_a/16:.4f}")

On the real 31-parameter network, Adam converges significantly faster. The per-parameter adaptation gives each of the 31 parameters an individually tuned learning rate. Layer 1 weights (which receive attenuated gradients through the chain rule) automatically get a larger effective learning rate, while layer 2 biases (which receive direct gradients) get a smaller one.

Why Adam Wins on Real Networks

AdamW: Weight Decay Done Right

In 2019, Loshchilov & Hutter published a seemingly minor paper that turned out to be extremely important. They showed that the standard way of adding L2 regularization to Adam is fundamentally broken, and proposed AdamW as the fix.

The Problem with L2 Regularization in Adam

L2 regularization adds a penalty $\frac{\lambda}{2} \|\mathbf{w}\|^2$ to the loss. The gradient becomes $g_t + \lambda w_t$. In vanilla SGD, this works correctly — the weight decay term $\lambda w_t$ shrinks each weight proportionally.

But in Adam, the gradient goes through the $\hat{m}/\sqrt{\hat{v}}$ normalization. The weight decay term gets scaled differently for each parameter. A parameter with large $\sqrt{\hat{v}}$ gets less weight decay than intended; a parameter with small $\sqrt{\hat{v}}$ gets more. The adaptive scaling, which is perfect for the gradient, distorts the regularization.

The AdamW Fix

AdamW decouples weight decay from the gradient update. Instead of adding $\lambda w$ to the gradient, AdamW applies it directly to the weights after the Adam step:

$w_{t+1} = w_t - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \varepsilon} - \eta \cdot \lambda \cdot w_t$

The gradient update ($\hat{m}/\sqrt{\hat{v}}$) goes through Adam's normalization. The weight decay ($\lambda w$) is applied directly, without any scaling. Every parameter gets the same proportional decay, regardless of its gradient history.

In PyTorch, switching from Adam to AdamW is a one-line change:

$\texttt{opt = torch.optim.AdamW(params, lr=0.001, weight\_decay=0.01)}$

Connection to Modern Training

Adam and AdamW are the workhorses of modern deep learning. Here is how they appear in practice.

Standard Hyperparameters for Transformers

Learning Rate Warmup

Transformers use a learning rate warmup for the first 1,000–4,000 steps. The learning rate starts at 0 and linearly increases to the target value. This interacts with Adam in two ways:

1. Adam's bias correction already warms up the effective step. At step 1, $\hat{v}$ is corrected by 1000×, preventing oversized steps. Warmup adds additional safety on top of bias correction.
2. The second moment $v$ needs time to converge. During warmup, $v$ builds up a reasonable estimate of the gradient variance. By the time lr reaches its target, $v$ is well-calibrated and the per-parameter scaling is reliable.

Cosine Learning Rate Decay

After warmup, most transformer training schedules follow a cosine decay:

$\eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\left(\frac{t}{T} \cdot \pi\right)\right)$

This smoothly decreases the learning rate from $\eta_{\max}$ to $\eta_{\min}$ over $T$ steps. The combination of AdamW + warmup + cosine decay is the standard training recipe for GPT, LLaMA, and most modern language models.

Memory Cost at Scale

Adam's two buffers per parameter become significant at scale. For a 7B parameter model (like LLaMA-7B):

The optimizer state (m + v) alone is 56 GB for a 7B model — twice the size of the model itself. This is why optimizer state sharding (like ZeRO from DeepSpeed and FSDP from PyTorch) is essential for training large models. These techniques distribute the m and v buffers across multiple GPUs.

Adam Variants in Production

• AdaFactor: reduces memory by storing only row and column statistics instead of the full m and v matrices. Used in T5 and some Google models.
• LAMB (Layer-wise Adaptive Moments for Batch training): scales Adam's update by the ratio of parameter norm to update norm. Enables extremely large batch sizes (up to 64K). Used in BERT pretraining.
• 8-bit Adam: quantizes m and v buffers to 8-bit integers (via the bitsandbytes library), reducing optimizer memory by 4× with minimal quality loss. Increasingly popular for fine-tuning large models on consumer GPUs.
• Lion (EvoLved Sign Momentum): discovered by Google through evolutionary search. Uses only the sign of the momentum update (no second moment), reducing memory to 1 buffer per parameter. Competitive with Adam on many tasks.

Summary

1. Momentum alone cannot adapt per-parameter. All parameters share one learning rate, which fails when gradient magnitudes span orders of magnitude across parameters.
2. RMSprop divides by the RMS gradient to normalize step sizes across parameters. But it has no momentum (no directional smoothing) and no bias correction (misbehaves early in training).
3. Adam combines momentum and RMSprop with bias correction. First moment $m$ smooths direction; second moment $v$ adapts magnitude. The effective step is approximately $\eta$ for every parameter.
4. Bias correction is essential because both $m$ and $v$ start at zero. Without it, the second moment is underestimated by up to 1000× in the first step, causing oversized updates.
5. Adam is not always fastest — on simple 1D problems, momentum wins because it builds velocity. Adam's advantage is on multi-dimensional problems where parameters need different step sizes.
6. AdamW decouples weight decay from the gradient normalization. This ensures consistent regularization across all parameters and is the standard for modern transformer training.
7. The standard transformer recipe: AdamW + linear warmup + cosine decay, with $\beta_1 = 0.9$, $\beta_2 = 0.95$–$0.999$, weight_decay = 0.01–0.1, gradient clipping at max_norm = 1.0.

Looking ahead: Adam and AdamW give us an excellent optimizer, but they still require choosing a learning rate and a training duration. Learning rate scheduling — how to adjust $\eta$ during training — can further improve convergence. We explore this in Section 3.