Adam Optimizer

Why This Matters

Adam is the default optimizer for training deep neural networks. It combines momentum (first moment estimation) with adaptive learning rates (second moment estimation) to handle the challenges of non-convex, high-dimensional optimization. Understanding Adam deeply means understanding why bias correction matters, how AdamW differs from Adam+L2 (they are not the same), and when Adam actually hurts generalization compared to plain SGD.

Mental Model

SGD with momentum keeps a running average of gradients to smooth out noise. RMSprop keeps a running average of squared gradients to adapt the learning rate per-parameter (parameters with large gradients get smaller steps). Adam combines both: it maintains a momentum vector (first moment) and an adaptive scaling vector (second moment), with bias correction to handle initialization.

The Algorithm

Definition

Adam Update Rule

Given parameters $\theta$, learning rate $\eta$, decay rates $\beta_1, \beta_2$, and small constant $\epsilon$:

At step $t$, with gradient $g_t = \nabla_\theta \mathcal{L}(\theta_{t-1})$:

$m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t \quad \text{(first moment estimate)}$ $v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \quad \text{(second moment estimate)}$ $\hat{m}_t = \frac{m_t}{1 - \beta_1^t} \quad \text{(bias-corrected first moment)}$ $\hat{v}_t = \frac{v_t}{1 - \beta_2^t} \quad \text{(bias-corrected second moment)}$ $\theta_t = \theta_{t-1} - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$

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

Components Explained

Definition

First Moment (Momentum)

$m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t$ is an exponential moving average of gradients. With $\beta_1 = 0.9$, this averages roughly the last 10 gradients. It smooths out gradient noise and accumulates direction, like a ball rolling downhill with momentum.

Expanding: $m_t = (1-\beta_1)\sum_{i=1}^t \beta_1^{t-i} g_i$.

Definition

Second Moment (Adaptive Scaling)

$v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2$ is an exponential moving average of squared gradients (elementwise). With $\beta_2 = 0.999$, this averages roughly the last 1000 squared gradients. It estimates the variance of each gradient component.

Dividing by $\sqrt{v_t}$ gives each parameter its own effective learning rate: parameters with consistently large gradients get smaller steps, and parameters with small gradients get larger steps.

Bias Correction

Theorem

Bias Correction for Exponential Moving Averages

Statement

If $m_0 = 0$ and $g_t$ are drawn from a stationary distribution with mean $\mathbb{E}[g_t] = g$, then the raw exponential moving average is biased:

$\mathbb{E}[m_t] = (1 - \beta^t) \cdot g$

The bias-corrected estimate $\hat{m}_t = m_t / (1 - \beta^t)$ satisfies $\mathbb{E}[\hat{m}_t] = g$.

Intuition

When you initialize $m_0 = 0$ and start averaging, the early estimates are biased toward zero. After one step with $\beta = 0.9$, you have $m_1 = 0.1 \cdot g_1$, which underestimates the true gradient by a factor of 10. Dividing by $(1 - 0.9^1) = 0.1$ corrects this. The correction matters most in the first few iterations and becomes negligible as $t$ grows (since $\beta^t \to 0$).

Proof Sketch

$m_t = (1-\beta)\sum_{i=1}^t \beta^{t-i} g_i$.

$\mathbb{E}[m_t] = (1-\beta) g \sum_{i=1}^t \beta^{t-i} = (1-\beta) g \cdot \frac{1-\beta^t}{1-\beta} = (1-\beta^t) g$.

So $\mathbb{E}[m_t/(1-\beta^t)] = g$. For the second moment, the same argument applies with $g^2$ replaced by $\mathbb{E}[g_t^2]$.

Why It Matters

Without bias correction, the first few Adam steps take very small updates because $m_t$ and $v_t$ are close to zero. With $\beta_2 = 0.999$, the second moment estimate is severely biased for the first ~1000 steps. Bias correction prevents this slow start and is essential for Adam to work well in practice.

Failure Mode

Bias correction assumes a stationary gradient distribution. In the early phase of training when the loss landscape changes rapidly, the stationarity assumption is violated. This is one motivation for learning rate warmup.

AdamW: Decoupled Weight Decay

Theorem

AdamW Decouples Weight Decay from Gradient Adaptation

Statement

Adam + L2 regularization adds the L2 gradient to the gradient before moment estimation:

$g_t = \nabla \mathcal{L}(\theta_{t-1}) + \lambda \theta_{t-1}$

then runs standard Adam. The effective weight decay for parameter $\theta_j$ is $\lambda \cdot \eta / \sqrt{\hat{v}_j}$, which varies per parameter.

AdamW applies weight decay directly to the parameters, after the adaptive step:

$\theta_t = (1 - \eta\lambda)\theta_{t-1} - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$

In AdamW, the weight decay $\lambda$ is the same for all parameters, regardless of gradient magnitude.

Intuition

In SGD, L2 regularization and weight decay are equivalent. In Adam, they are not. When Adam divides the gradient by $\sqrt{v_t}$, it also divides the L2 gradient, weakening the regularization for parameters with large gradient history. AdamW avoids this by applying decay separately. This means all parameters decay at the same rate, as intended.

Proof Sketch

With Adam + L2: the effective update for parameter $j$ includes $\eta \cdot \lambda\theta_j / \sqrt{\hat{v}_j}$. Parameters with large $\hat{v}_j$ (historically large gradients) receive weaker decay.

With AdamW: the decay term is $\eta\lambda\theta_j$ regardless of $\hat{v}_j$. The gradient-based update and the decay are fully decoupled.

Why It Matters

Loshchilov and Hutter (2019) showed that AdamW consistently outperforms Adam+L2 across tasks. The key insight: regularization should not interact with the adaptive learning rate mechanism. AdamW is the default optimizer for training Transformer-based language models. For convolutional architectures, SGD with momentum remains competitive or superior (Wilson et al. 2017).

Failure Mode

The optimal $\lambda$ for AdamW is different from the optimal $\lambda$ for Adam+L2. You cannot simply swap one for the other without retuning. The typical AdamW weight decay for Transformers is 0.01-0.1, much larger than what you would use for L2 regularization in Adam.

Learning Rate Warmup

In practice, Adam is often combined with learning rate warmup: start with a very small learning rate and linearly increase it over the first $T_w$ steps to the target value. Why?

1. Second moment initialization: At step 1, $\hat{v}_1 = g_1^2$, a noisy single-sample estimate. A large learning rate with a noisy denominator produces wild parameter updates. Warmup gives $v_t$ time to stabilize.

2. Loss landscape curvature: Early in training, the loss landscape may have regions of very high curvature. Large steps in these regions can be catastrophic. Warmup allows the model to reach a better-conditioned region before taking large steps.

A common schedule is linear warmup for 1-10% of total training steps, followed by cosine decay.

When Adam Fails

Adam is not universally superior to SGD:

1. Generalization gap (contested): Wilson et al. (2017) reported that SGD with momentum often generalizes better than Adam on image classification. One proposed mechanism is that Adam finds sharper minima while SGD's larger noise finds flatter ones (Keskar et al. 2017), but the flat-minima hypothesis itself has been challenged on reparameterization grounds (Dinh et al. 2017). The generalization gap is real; its causal mechanism is not settled.

2. Non-convergence: Reddi et al. (2018) showed that Adam can diverge on simple convex problems because the adaptive learning rate can increase without bound when $v_t$ shrinks. AMSGrad fixes this by taking $\hat{v}_t = \max(\hat{v}_{t-1}, v_t)$ to ensure the learning rate never increases. In practice AMSGrad is rarely used; the modification has not consistently improved empirical performance and most codebases default to AdamW.

3. Domain dependence: Adam dominates for NLP/Transformers. SGD dominates for CNNs on vision tasks. The optimal optimizer depends on the architecture and data distribution.

Canonical Examples

Example

Why bias correction matters early

With $\beta_2 = 0.999$ and $m_0 = v_0 = 0$, at step $t=1$:

• Raw: $v_1 = 0.001 \cdot g_1^2$ (underestimates true second moment by 1000x)
• Corrected: $\hat{v}_1 = g_1^2$ (correct order of magnitude)

Without correction, $\sqrt{\hat{v}_1}$ would be ~30x too small, making the effective learning rate ~30x too large. The first step would be catastrophically large.

Common Confusions

Watch Out

Adam and AdamW are NOT interchangeable

Adam with L2 regularization and AdamW with weight decay produce different parameter trajectories, even with the same $\lambda$ value. In Adam+L2, the regularization strength varies per parameter (inversely with gradient history). In AdamW, it is uniform. Always use AdamW when you want weight decay with adaptive optimizers.

Watch Out

The epsilon parameter is not negligible

The default $\epsilon = 10^{-8}$ prevents division by zero, but in half-precision training ($\epsilon_{\text{mach}} \approx 10^{-4}$), this default can cause numerical issues. For fp16 training, set $\epsilon = 10^{-5}$ or larger.

Summary

• Adam = momentum (first moment) + adaptive LR (second moment) + bias correction
• First moment: $m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t$ (smooths gradients)
• Second moment: $v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2$ (scales per-parameter)
• Bias correction: divide by $(1-\beta^t)$ to correct for zero initialization
• AdamW decouples weight decay from adaptive scaling. use AdamW, not Adam+L2
• Adam finds sharper minima than SGD; SGD can generalize better for some tasks
• Warmup stabilizes the second moment estimate in early training

Exercises

ExerciseCore

Problem

Derive the bias-corrected first moment estimate. Starting from $m_0 = 0$ and $m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t$, show that $\mathbb{E}[m_t] = (1-\beta_1^t)\mathbb{E}[g]$ when the gradients have constant expectation $\mathbb{E}[g_t] = g$.

Hint

Reveal Solution

ExerciseAdvanced

Problem

Show that for SGD (no adaptive scaling), L2 regularization ($g_t \leftarrow g_t + \lambda\theta$) is equivalent to weight decay ($\theta \leftarrow (1-\eta\lambda)\theta - \eta g_t$). Then explain why this equivalence breaks for Adam.

Hint

Reveal Solution

Related Comparisons

• Adam vs. SGD

References

Canonical:

• Kingma & Ba, "Adam: A Method for Stochastic Optimization" (2015)
• Loshchilov & Hutter, "Decoupled Weight Decay Regularization" (2019). introduces AdamW

Current:

• Reddi et al., "On the Convergence of Adam and Beyond" (2018). AMSGrad
• Wilson et al., "The Marginal Value of Adaptive Gradient Methods in ML" (2017)

Next Topics

Adam connects to broader optimization topics:

• Convex optimization basics: the theoretical foundation for understanding convergence
• Learning rate schedules: warmup, cosine decay, and their interaction with Adam