Adam Optimizer: Adaptive Learning Rates Explained

Training a deep neural network is fundamentally an optimization problem: you have a loss function and you want to find the weights that minimize it. As we saw in the Stochastic Gradient Descent chapter, SGD moves weights in the direction of steepest descent, and momentum adds a velocity term to smooth out oscillations and speed past shallow gradients. Both methods use a single global learning rate for every parameter, which creates a fundamental tension. Some parameters receive dense gradient signals and should move cautiously; others are updated rarely and need larger steps to converge at all. A single learning rate cannot serve both needs well.

Adam (Adaptive Moment Estimation), introduced by Diederik Kingma and Jimmy Ba in 2014, resolves this tension by maintaining an independent learning rate for every parameter, scaled automatically based on the history of gradients that parameter has seen. It does this by tracking two running statistics: the mean of past gradients (the first moment) and the uncentered variance of past gradients (the second moment). These two statistics together give Adam enough information to set an appropriate step size for each parameter individually. The result is an optimizer that requires almost no per-problem tuning and converges reliably across a wide range of architectures, from convolutional networks to transformers.

Understanding Adam deeply means understanding why the two moment estimates are needed, where the bias correction terms come from, why the epsilon term in the denominator matters more than its small value suggests, and what the practical failure modes look like in production training runs. This chapter covers all of that, from first principles through a full PyTorch implementation, a survey of Adam's most important variants, and a discussion of the memory implications that make Adam challenging at large scale.

To see why a single global learning rate is limiting, consider a sparse input problem like word embeddings. In language modeling, some words appear thousands of times per batch; others might appear once per epoch. When you compute gradients and average them over a minibatch, frequently occurring words generate large, consistent gradient signals. Rare words generate tiny gradients most of the time, with occasional large spikes. If you set the learning rate high enough for rare words to learn, frequent words overshoot. If you set it low enough for frequent words, rare words barely move.

The same tension appears in deeper networks. Parameters close to the output receive large, stable gradients; parameters in early layers receive small, noisy gradients that have been attenuated through many multiplications during backpropagation. We saw in the backpropagation chapter how the chain rule multiplies gradients across layers, which can leave early-layer parameters chronically under-updated or chronically over-updated depending on the global learning rate.

The problem with a single learning rate becomes especially vivid when you think about the geometry of the loss surface. In a well-conditioned problem, the loss function is roughly spherical: the curvature is similar in every direction. SGD with a fixed learning rate moves at roughly the same speed in all directions. In an ill-conditioned problem, the loss surface looks more like a narrow canyon: very steep walls on one side, a very flat floor in another. The optimal path to the minimum runs along the canyon floor, but SGD with a fixed learning rate has to choose between two bad options. Use a large step size and you bounce back and forth off the steep walls, oscillating violently. Use a small step size and you crawl along the flat floor at a glacial pace.

This ill-conditioning is not a pathological edge case. It arises naturally in deep networks from weight initialization asymmetries, from the scale differences between input features, and from the structure of the weight matrices at different depths. A typical transformer training run will have gradient norms that vary by orders of magnitude across different parameter groups. The first embedding layer, the output projection, and the attention weight matrices all occupy different regions of the loss landscape with vastly different curvatures.

Adaptive learning rates address the ill-conditioning problem by measuring, for each parameter independently, how much curvature has been encountered. Parameters in steep dimensions accumulate large second-moment estimates and receive smaller effective learning rates. Parameters in flat dimensions accumulate small second-moment estimates and receive larger effective learning rates. The optimizer self-calibrates its step size to match the local geometry of each dimension simultaneously, without any manual tuning.

The idea of parameter-specific step sizes had been explored before Adam. AdaGrad (Duchi, Hazan, and Singer, 2011) was the first major adaptive optimizer. It accumulated the sum of all squared gradients seen so far for each parameter, and used that sum in the denominator:

where $G_t = \sum_{i=1}^{t} g_i^2$ is the cumulative sum of squared gradients. AdaGrad was a significant step forward for sparse problems: parameters that receive occasional large gradients quickly build up a large $G_t$, causing their learning rate to shrink. Parameters that receive small gradients maintain a small $G_t$ and continue to learn quickly. This worked well for training word vectors and other models with sparse inputs.

However, AdaGrad has a critical flaw for long training runs: the accumulator $G_t$ only ever grows. It grows every time a gradient is non-zero, and it never forgets old gradient information. After enough training steps, $G_t$ becomes so large that the effective learning rate $\alpha / \sqrt{G_t}$ shrinks to near zero for all parameters. Training effectively stalls. In practice, AdaGrad works well for shallow models trained for a small number of epochs, but fails to converge on deep networks trained for hundreds of epochs.

RMSProp, proposed by Geoffrey Hinton in an unpublished lecture (2012), addressed the monotonic accumulation problem by replacing the cumulative sum with an exponential moving average:

By using a decaying average rather than a cumulative sum, RMSProp's denominator reflects recent gradient magnitudes rather than the entire training history. Parameters that were frequently updated early in training but are now quiet will see their $v_t$ decay, allowing their effective learning rate to recover. This makes RMSProp viable for long training runs where gradient statistics change substantially over time.

Adam combines RMSProp's adaptive second moment with momentum's first moment, wrapping both in a bias correction that makes the estimates accurate from the very first update. It is, in a sense, the synthesis of a decade of adaptive optimizer research into a single coherent algorithm.

Adam names its running statistics after statistical moments. Recall that for a random variable $g$:

• The first moment is the mean: $\mathbb{E}[g]$
• The second moment is the mean of the square: $\mathbb{E}[g^2]$ (uncentered variance)

Adam does not compute these exactly. Instead it maintains exponential moving averages that approximate them. Let $g_t$ be the gradient of the loss with respect to a single parameter at step $t$. Adam maintains:

where:

• $m_t$: the first moment estimate (exponential moving average of gradients)
• $v_t$: the second moment estimate (exponential moving average of squared gradients)
• $\beta_1$: the decay rate for the first moment, typically 0.9
• $\beta_2$: the decay rate for the second moment, typically 0.999
• $g_t$: the gradient at time step $t$

The first moment $m_t$ is recognizable as the velocity term from momentum. When gradients consistently point in one direction, $m_t$ builds up in that direction. When gradients oscillate, contributions cancel out. This is exactly the momentum behavior we saw in Chapter 7.

The second moment $v_t$ tracks the recent average of squared gradients. If a parameter has been receiving large gradients (either positive or negative), $v_t$ will be large. If it has received small or near-zero gradients, $v_t$ will be small. Squaring keeps $v_t$ positive and captures magnitude regardless of sign.

The Adam update will divide the first moment by the square root of the second moment. To see why this achieves adaptive learning rates, consider two parameters:

Parameter A (dense gradients): receives gradients like $[0.8, 0.7, 0.9, 0.8, ...]$ repeatedly. The second moment $v_t \approx 0.64$. The effective learning rate is scaled by $1/\sqrt{0.64} \approx 1.25$. Not much amplification because this parameter is already getting strong updates.

Parameter B (sparse gradients): receives gradients like $[0.01, 0.0, 0.0, 0.0, 0.5, 0.0, ...]$ rarely. The second moment $v_t \approx 0.002$ (most steps contribute zero). The effective learning rate is scaled by $1/\sqrt{0.002} \approx 22$. This rare-but-important parameter gets a much larger effective step.

The division by $\sqrt{v_t}$ naturally equalizes the learning pace: parameters with large historical gradients get a small effective learning rate, and parameters with small historical gradients get a large effective learning rate. The optimizer finds a reasonable pace for each parameter without any manual tuning.

There is a useful geometric interpretation of what Adam is doing. The second moment estimate $\sqrt{v_t}$ approximates the root-mean-square gradient for each parameter. Dividing by this quantity normalizes the update so that each parameter receives a step whose size is roughly proportional to the inverse of its typical gradient magnitude. In a sense, Adam is performing a diagonal approximation to the natural gradient: it rescales the gradient to account for the curvature of the loss surface along each parameter axis independently. Full natural gradient methods (like K-FAC) use the full Fisher information matrix to perform this rescaling, which is computationally prohibitive at large scale. Adam's diagonal approximation is cheap to compute and surprisingly effective in practice, capturing the most important per-axis curvature information while staying computationally tractable.

There is a subtle but important flaw in the raw moment estimates. Both $m_t$ and $v_t$ are initialized to zero at step $t=0$. This means that early in training, the estimates are biased toward zero. At step $t=1$:

With $\beta_1 = 0.9$, we have $m_1 = 0.1 \cdot g_1$ even though our best estimate of the mean gradient is $g_1$ itself. The true mean should be $g_1$, but our estimate is only 10% of it.

This zero-initialization bias decays over time. After $t$ steps, the first moment is a weighted sum of past gradients:

Taking the expectation (assuming $g_i$ are drawn from the same distribution with mean $\mu$):

where:

• $\mu$: the true mean gradient
• $\beta_1^t$: the decay factor raised to the $t$-th power
• $(1 - \beta_1^t)$: the bias factor, which starts near 0 and approaches 1 as $t$ grows

The bias factor $(1 - \beta_1^t)$ starts at $1 - 0.9 = 0.1$ at step $t=1$ (meaning the estimate is only 10% of the truth) and approaches 1 as $t \to \infty$ (where it becomes accurate). To correct for this initialization bias, Adam divides each moment estimate by its bias correction factor:

where:

• $\hat{m}_t$: the bias-corrected first moment estimate
• $\hat{v}_t$: the bias-corrected second moment estimate
• $1 - \beta_1^t$: correction factor for the first moment at step $t$
• $1 - \beta_2^t$: correction factor for the second moment at step $t$

For $\beta_2 = 0.999$, the second moment bias correction is even more severe early on: $1 - 0.999^1 = 0.001$, so $\hat{v}_1 = v_1 / 0.001 = 1000 \cdot v_1$. Without this correction, the second moment estimate is far too small for the first thousand or so steps, which would cause Adam to take enormous steps with a tiny denominator. The bias correction stabilizes training in these early steps, making Adam's behavior in the first epoch much more predictable.

As training progresses, $\beta_1^t \to 0$ and $\beta_2^t \to 0$, so $\hat{m}_t \to m_t$ and $\hat{v}_t \to v_t$. The correction becomes negligible after enough steps. For $\beta_1 = 0.9$, the first moment correction becomes less than 1% after about $t = 46$ steps. For $\beta_2 = 0.999$, the second moment correction becomes less than 1% after about $t = 4,600$ steps.

Without bias correction, you would observe a characteristic artifact at the start of training: the learning rate would appear enormous for the first few hundred steps, because the second moment estimate $v_t$ would be near zero, making $\alpha / \sqrt{v_t}$ very large. This would cause gradient explosions or wildly oscillating loss values before eventually stabilizing. With bias correction, both moments are immediately calibrated to their true statistical values, and Adam's step sizes remain reasonable from step one. This is especially important for transformer training with learning rate warmup: the warmup schedule deliberately starts with a small learning rate and increases it over the first several thousand steps. During this warmup period, the bias correction is doing exactly its intended job. Both the scheduled learning rate and the moment estimates are warming up together, creating doubly conservative early steps that prevent gradient explosions. The two mechanisms are complementary, not redundant.

With bias-corrected moment estimates in hand, the Adam parameter update is:

where:

• $\theta_t$: the parameter vector at step $t$
• $\alpha$: the global learning rate (typically 0.001)
• $\hat{m}_t$: bias-corrected first moment (direction and momentum)
• $\hat{v}_t$: bias-corrected second moment (per-parameter scale)
• $\epsilon$: a small constant for numerical stability, typically $10^{-8}$

The full Adam algorithm, written out step by step, is:

where:

• $\mathcal{L}$: the loss function
• All other variables are as defined above

The $\epsilon$ term in the denominator prevents division by zero when $\hat{v}_t \approx 0$, which happens for parameters that have received no gradients. But $\epsilon$ has an additional effect on Adam's behavior that is often overlooked.

Consider a parameter where $\hat{v}_t$ is very small, say $10^{-12}$. Then $\sqrt{\hat{v}_t} + \epsilon \approx \epsilon = 10^{-8}$, and the effective learning rate is $\alpha / \epsilon = 0.001 / 10^{-8} = 100,000$. This is enormous. In practice, if a parameter truly receives near-zero gradients consistently, its $\hat{m}_t$ is also near zero, so the actual update $(\alpha / \epsilon) \cdot \hat{m}_t$ remains small. The two small quantities cancel out.

However, if $\epsilon$ is set too small (e.g., $10^{-16}$), numerical precision issues can cause $\hat{v}_t$ to underflow to zero in floating-point arithmetic, resulting in infs or nans in the updates. If $\epsilon$ is set too large (e.g., 1.0), Adam degenerates toward vanilla SGD because the denominator becomes dominated by $\epsilon$ rather than $\sqrt{\hat{v}_t}$, eliminating the adaptive behavior. The default $\epsilon = 10^{-8}$ is a pragmatic choice that avoids both failure modes in typical float32 computation.

PyTorch's torch.optim.Adam uses $\epsilon = 10^{-8}$ by default but exposes it as a tunable parameter. Some practitioners increase it to $10^{-7}$ or even $10^{-5}$ when training with mixed precision (float16), since smaller values can cause numerical instability at reduced precision.

Let's trace a single parameter through five Adam update steps to make the mechanics concrete. Suppose the parameter starts at $\theta_0 = 2.0$, and receives gradients $g = [0.5, 0.4, 0.6, 0.3, 0.5]$ over five steps. We'll use the standard hyperparameters: $\alpha = 0.01$, $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$.

Initialize: $m_0 = 0$, $v_0 = 0$.

Step 1 ($t=1$, $g_1 = 0.5$):

Notice that after bias correction, $\hat{m}_1 = 0.5 = g_1$ and $\hat{v}_1 = 0.25 = g_1^2$. The bias correction brings the estimates in line with the actual gradient magnitude at step 1, giving an effective step of exactly $\alpha / \sqrt{g_1^2} \cdot g_1 = \alpha \cdot \text{sign}(g_1)$. In other words, early in training the Adam update approximates a step of constant size $\alpha$ in the gradient direction, regardless of gradient magnitude. This is the trust region interpretation of Adam.

Step 2 ($t=2$, $g_2 = 0.4$):

The parameter continues to decrease, with each step sized close to $\alpha = 0.01$. This is characteristic of Adam in early training: the per-parameter normalization makes step sizes roughly equal to the global learning rate $\alpha$, giving predictable training dynamics regardless of gradient scale.

By step 5, the moment estimates have accumulated enough history that the bias corrections become smaller. The first moment $m_t$ is now a smooth average of all five gradients, pointing consistently in the positive direction (since all gradients are positive). The second moment $v_t$ has accumulated a sense of the typical gradient magnitude squared, and the effective step size is settling toward a stable value. This accumulation phase is what makes Adam's convergence so reliable: within a few dozen steps, it has calibrated per-parameter step sizes that remain appropriate throughout the rest of training. If you were to plot the effective learning rate $\alpha / (\sqrt{\hat{v}_t} + \epsilon)$ over time, you would see it start near $\alpha$ during the bias-dominated early steps, then gradually settle as $\hat{v}_t$ stabilizes around the true mean squared gradient. This settling process takes longer for parameters with smaller gradients, because their $v_t$ grows more slowly and the bias correction factor remains significant for more steps.

Adam's design draws from two predecessors, and it helps to understand what each contributes.

Vanilla SGD uses a fixed learning rate:

SGD is simple and works well when gradients are well-scaled, but it has no memory of past gradients and no ability to adapt its step size per parameter.

Momentum adds a velocity term:

Momentum smooths oscillations and builds up speed in consistent gradient directions, which helps escape flat regions and ravines. But it still uses a single global learning rate.

RMSProp adapts the learning rate per parameter using the second moment:

RMSProp gives each parameter its own adaptive step size but applies the raw gradient $g_t$ rather than a momentum-smoothed version. This means the update direction can be noisy when individual gradients are noisy.

Adam combines both improvements:

The numerator $\hat{m}_t$ provides the momentum benefit (smoothed direction), and the denominator $\sqrt{\hat{v}_t} + \epsilon$ provides the adaptive scale benefit (per-parameter step size). The bias correction makes both estimates reliable from the start.

The practical implication of this comparison is that Adam's 2x memory overhead over SGD (two tensors per parameter instead of one) is the price you pay for the two improvements. For most modern deep learning workloads where GPU memory is the binding constraint, this tradeoff is well worth it. The improvement in convergence speed typically more than compensates for the smaller effective batch size you can fit on a GPU. However, at very large model scales (billions of parameters), the memory overhead becomes a substantial engineering challenge, as discussed later in the chapter.

The original Adam paper provides a convergence proof for the online learning setting (non-stationary objectives) under the assumption of convexity. In practice, neural network training is non-convex, and the theoretical guarantees do not directly apply. Nevertheless, several empirical and theoretical observations are useful.

For convex problems, Adam achieves $O(\sqrt{T})$ regret, where $T$ is the number of steps. This is the same asymptotic rate as AdaGrad, though Adam's per-step constants are generally better because the exponential moving average avoids the excessive accumulation that slows AdaGrad in long runs.

For non-convex problems, Adam can fail to converge to a stationary point in theory. A paper by Reddi, Kale, and Kumar (2018) demonstrated specific synthetic examples where Adam oscillates between two solutions rather than converging. This led to the AMSGrad variant (discussed below), though in practice Adam converges reliably on most machine learning tasks.

Adam outperforms SGD most clearly in situations with:

• Sparse features: embedding tables, one-hot inputs, high-dimensional vocabularies. The adaptive per-parameter learning rate is decisive here. Embedding rows that are rarely activated can receive large effective learning rates because their second moment accumulates slowly. This is one reason Adam became dominant for NLP before it became dominant for vision.
• Non-stationary objectives: Adam's exponential weighting means it adapts to changing gradient distributions, which matters in curriculum learning or when the data distribution shifts during training.
• Noisy gradients: the momentum term smooths noise, and the adaptive denominator prevents large oscillations from dominating the update.
• Deep networks: parameters at different depths receive gradients of very different magnitude. Adam compensates for this discrepancy without requiring layer-wise learning rate tuning.
• Fast initial convergence requirements: because Adam normalizes step sizes to be approximately $\alpha$ regardless of gradient scale, it tends to make rapid progress in early training. When compute is limited and you cannot afford to train to convergence, Adam often finds a better solution faster.

Careful research has shown that carefully tuned SGD with momentum can match Adam's generalization performance on image classification benchmarks, and sometimes exceeds it. The intuition is that Adam's adaptive learning rates may allow it to find sharp, narrow minima that generalize less well than the broader minima that SGD tends to find. The sharp minimum hypothesis (from Keskar et al., 2017) suggests that minima discovered through large, consistently-directed updates tend to be broader and generalize better. SGD with a large, slowly decaying learning rate naturally finds such broad minima. Adam's adaptive normalization tends to use more cautious, consistent step sizes that can lock in on narrow minima earlier in training.

However, this tradeoff is highly architecture- and task-dependent. For language modeling and transformer training, Adam or its variants are essentially universal. The datasets are large enough, and the loss landscapes well-behaved enough, that the generalization gap between Adam and SGD rarely manifests.

The practical advice from researchers at Google, OpenAI, and DeepMind converges on: use Adam (or AdamW) by default; switch to carefully-tuned SGD with momentum only if you have the compute budget for extensive hyperparameter search and there is evidence from prior work that SGD generalizes better for your specific task.

Adam has four hyperparameters: $\alpha$, $\beta_1$, $\beta_2$, and $\epsilon$. Understanding what each controls helps you diagnose training problems.

The learning rate remains the most important hyperparameter. Adam's default of $\alpha = 10^{-3}$ works for a surprisingly wide range of tasks, but you should expect to tune it. Rules of thumb:

• Start with $\alpha = 10^{-3}$ for new problems
• If training loss explodes or diverges immediately: reduce by 10x
• If training is too slow after 10% of planned epochs: increase by 3-10x
• Use a learning rate finder (warm-up scan) to identify the range of valid learning rates

One important difference from SGD: Adam's effective per-parameter learning rate is $\alpha / (\sqrt{\hat{v}_t} + \epsilon)$. When you change the global $\alpha$, you are scaling all per-parameter learning rates proportionally. The adaptive normalization means Adam is generally less sensitive to the precise choice of $\alpha$ than SGD is, but the relative magnitudes still matter.

$\beta_1$ controls how much past gradients influence the current update direction. With $\beta_1 = 0.9$ (default), the effective momentum window is approximately $1 / (1 - 0.9) = 10$ steps. This means Adam averages gradient direction over the last ~10 steps.

Lowering $\beta_1$ makes Adam respond more quickly to recent gradient changes, which can help in non-stationary settings but increases noise. The default 0.9 is almost never worth changing. Some practitioners use $\beta_1 = 0$ (effectively no momentum) for specific research experiments to isolate the adaptive learning rate effect, but this is rarely beneficial in practice.

$\beta_2$ controls how long Adam remembers past gradient magnitudes. With $\beta_2 = 0.999$ (default), the effective window is $1 / (1 - 0.999) = 1000$ steps. The second moment tracks a much longer history than the first moment.

This asymmetry is intentional: you want the denominator to be a stable estimate of the gradient's typical magnitude, not one that jumps around with each new batch. However, $\beta_2 = 0.999$ can be slow to adapt if the gradient scale changes dramatically during training (for example, when learning rate warm-up causes a large jump in gradient norms). In such cases, $\beta_2 = 0.99$ or even $0.98$ can respond faster.

In transformer training, $\beta_2 = 0.98$ has been found to work better than the default 0.999 in some settings (the original "Attention Is All You Need" paper used $\beta_2 = 0.98$).

As discussed above, $\epsilon = 10^{-8}$ is the default. For float16 training, $\epsilon = 10^{-6}$ or $10^{-5}$ is often more stable.

Adam has inspired a family of related optimizers, each addressing a specific limitation. Understanding the variants helps you choose the right optimizer for your use case.

AMSGrad (Reddi, Kale, Kumar, 2018) was proposed to fix the theoretical non-convergence problem in Adam. Recall that Adam's second moment estimate is an exponential moving average, which can "forget" past large gradients. AMSGrad uses the maximum of all past second moment estimates instead:

By using $\hat{v}_t^{\max}$ in the denominator instead of $\hat{v}_t$, AMSGrad ensures that the effective learning rate never increases. Once a large gradient has been seen, the corresponding dimension's effective learning rate stays compressed. This monotonically non-increasing property is what enables the convergence guarantee.

In practice, AMSGrad often converges more slowly than Adam because the non-increasing effective learning rate can make it too conservative in the latter part of training. Most practitioners find that Adam with a well-tuned schedule outperforms AMSGrad, which is why AMSGrad has seen limited adoption despite its theoretical appeal. You can enable it in PyTorch with torch.optim.Adam(..., amsgrad=True).

Nadam (Dozat, 2016) combines Adam with Nesterov momentum. Recall that Nesterov momentum modifies the standard momentum update by computing the gradient at the "lookahead" position rather than the current position. Nesterov momentum typically converges slightly faster than standard momentum because the gradient is evaluated at a more future-looking position, meaning the update step already accounts for where you will be after applying the current velocity.

Nadam incorporates this improvement into Adam by substituting the Nesterov first moment estimate into the update rule. The difference from Adam is subtle but can matter in settings where the gradient changes direction quickly. PyTorch provides Nadam as torch.optim.NAdam. For most applications, the improvement over Adam is marginal, but it is worth knowing about as the theoretically cleaner variant.

AdamW (Loshchilov and Hutter, 2017) is arguably the most important Adam variant and deserves a full chapter of its own (which follows this one). The key insight is that Adam's weight_decay parameter does not implement true weight decay: it adds $\lambda \theta$ to the gradient before the Adam update, which means the regularization term is divided by $\sqrt{\hat{v}_t}$ along with the gradient. This weakens the regularization for frequently updated parameters and strengthens it for rarely updated ones, which is the opposite of what you want.

AdamW implements decoupled weight decay: it applies the weight decay step directly to the parameters after the Adam update, independently of the gradient:

This simple decoupling produces substantially better generalization in language modeling and has become the standard optimizer for training transformers. All modern large language models use AdamW rather than vanilla Adam.

Adafactor (Shazeer and Stern, 2018) addresses the memory overhead problem. For a weight matrix of shape $[r, c]$, Adam stores two additional matrices of the same shape ($m_t$ and $v_t$), tripling the memory for that parameter. Adafactor factors the second moment matrix into two rank-one vectors of shape $[r]$ and $[c]$, reducing the optimizer state from $O(rc)$ to $O(r + c)$.

For large embedding tables or weight matrices, this can be a massive memory saving. A 50,000-word vocabulary with 1024-dimensional embeddings requires a matrix of shape $[50000, 1024]$. Adam stores two copies of this, totaling $2 \times 50000 \times 1024 = 102$ million floats for the optimizer state alone. Adafactor reduces this to $(50000 + 1024) \times 2 \approx 102,000$ floats. The tradeoff is that the factored approximation is less accurate than the true per-element second moment, which can slow convergence and require more careful hyperparameter tuning. Adafactor is used in training T5 and some other very large models where memory is the binding constraint.

Adam is robust but not failure-proof. These are the most common issues and their solutions.

Symptom: Loss is NaN or increases rapidly in the first few hundred steps.

Cause: Gradients are very large at initialization, causing step sizes that blow up the parameters. With poorly initialized weights, the first gradient can be orders of magnitude larger than what Adam stabilizes to after warm-up.

Fix: Use learning rate warm-up. Start with $\alpha$ at 1-10% of the target value and linearly increase it over the first $T_{\text{warmup}}$ steps (typically 1,000-10,000 steps). This gives Adam time to accumulate reliable moment estimates before taking large steps. Warm-up is essentially standard practice for transformer training.

Symptom: Training loss stops decreasing after a few thousand steps and stays flat.

Cause: Either the learning rate is too low, the moment estimates have accumulated high curvature noise, or the model is stuck in a flat region.

Fix: Check the learning rate first. Then consider learning rate decay: cosine annealing or step decay after a plateau. A learning rate restart (cyclic warm-up) can also help escape flat regions.

Symptom: Training with gradient accumulation over $K$ steps gives different results than training with batch size $K$ times larger.

Cause: When you accumulate gradients over $K$ steps and then take one Adam step, you are mixing per-step gradients differently than if you averaged them into one batch gradient. The moment estimates see $K$ small updates instead of one averaged update, leading to subtly different dynamics.

Fix: When using gradient accumulation for large effective batch sizes, scale $\alpha$ by $\sqrt{K}$ (not $K$) and keep the moment estimates consistent by only calling the optimizer step once per accumulation cycle.

Symptom: Adam achieves lower training loss than SGD but higher test error.

Cause: The adaptive learning rate may cause Adam to focus on dimensions with low variance, which can correspond to spurious correlations in the training data rather than generalizable features.

Fix: Switch to SGD with momentum for tasks where this generalization gap matters (primarily image classification with standard architectures). For NLP tasks, this problem is much less pronounced.

PyTorch provides torch.optim.Adam as a built-in optimizer. Understanding how to use it correctly, including how to set per-layer learning rates and integrate it with learning rate scheduling, is essential for practical deep learning.

First, install the required libraries.

A complete training loop using Adam over a synthetic binary classification dataset.

The model converges within 200 epochs using Adam's default learning rate of 0.001. This shows that the default configuration works well without any tuning on this type of task. Notice the standard pattern: optimizer.zero_grad() clears gradients from the previous step, loss.backward() computes new gradients, and optimizer.step() applies the Adam update. Forgetting zero_grad() is one of the most common bugs in PyTorch training: gradients accumulate across steps, effectively multiplying the gradient by the number of steps since the last reset, which causes wildly incorrect updates.

You can inspect Adam's internal state to understand what the optimizer has learned about each parameter.

The exp_avg and exp_avg_sq tensors correspond to $m_t$ and $v_t$ in Adam's equations. After training, $v_t$ tells you how large the gradients have been for each parameter. Parameters with large exp_avg_sq values have been actively updated throughout training. This state information can be very useful for debugging: if exp_avg_sq is near zero for many parameters, those parameters have barely been updated, which often signals a problem in the gradient flow (dead neurons, disconnected components, or missing gradients through a non-differentiable operation).

When you save and restore a model checkpoint, you should save the optimizer state alongside the model weights. Resuming training from a checkpoint without the optimizer state means Adam loses all its accumulated moment estimates and must restart the calibration process from zero, causing a temporary dip in training efficiency and sometimes a brief spike in loss.

Adam supports different learning rates for different parameter groups, which is useful when fine-tuning pretrained models where you want the base model to update slowly and the new head to update quickly.

This is a common pattern for transfer learning, where the pretrained backbone should update at a small fraction of the head's learning rate to preserve learned representations. The intuition is that the backbone already encodes useful features; updating it too aggressively will destroy that knowledge before the head has a chance to use it. A common heuristic is to set backbone learning rates to 1/10th or 1/100th of the head learning rate, though optimal values vary by task and model architecture.

Learning rate warmup is nearly universal in transformer training. The most common schedule is linear warmup followed by cosine or inverse-square-root decay.

The warmup ramps the learning rate from zero to the target value $\alpha$. During this phase, Adam's moment estimates are also near zero (from their zero initialization), so the combined effect of small $\alpha$ and small moment magnitudes means the first warmup steps are very conservative. As both the schedule and the moment estimates warm up together, training becomes stable.

The key parameters for PyTorch's Adam optimizer are:

• lr: Global learning rate $\alpha$. Default 0.001. This is the most impactful hyperparameter.
• betas: Tuple $(\beta_1, \beta_2)$ for first and second moment decay. Default (0.9, 0.999). Rarely need changing.
• eps: Numerical stability constant $\epsilon$. Default 1e-8. Increase to 1e-6 for float16 training.
• weight_decay: L2 regularization coefficient. Default 0. For true decoupled weight decay, use AdamW instead (covered in the next chapter).
• amsgrad: Whether to use AMSGrad variant. Default False. Rarely needed in practice.

To understand how Adam adapts to different gradient patterns, let's build visualizations that trace the moment estimates and effective learning rates for parameters with different gradient histories.

One reason Adam dominates practical deep learning is its robustness to gradient noise. When training with small minibatches, the gradient estimate is noisy: each batch samples only a fraction of the training data, and the resulting gradient is a stochastic approximation of the true gradient. This noise has two distinct effects on optimization.

The first effect is noise in gradient direction. On any given step, the gradient might point in a slightly wrong direction due to the random batch composition. Momentum (the first moment $m_t$) mitigates this by averaging over recent steps. If the true gradient points consistently in one direction but is corrupted by zero-mean noise, the average over $1/(1-\beta_1)$ steps will be much closer to the true direction than any individual step.

The second effect is noise in gradient magnitude. Stochastic gradients have variable magnitude even when the direction is correct. A batch that happens to contain many hard examples will produce a large gradient; a batch of easy examples produces a small one. The second moment $v_t$ captures this variability. By dividing by $\sqrt{v_t}$, Adam normalizes away the magnitude variation, making step sizes more consistent even when individual batch gradients vary widely.

This dual smoothing is why Adam with small batches often converges faster in terms of wall-clock time than SGD with large batches. With large batches, SGD has low gradient variance and converges well; with small batches (which process more data per unit time on a single GPU), Adam's noise robustness is decisive. In language model training, where batch sizes are often constrained by memory and gradients are computed from short context windows, this property is particularly valuable.

At large scale, Adam's memory requirements become a significant engineering constraint. Consider a language model with $N$ parameters. A standard training setup in float32 requires:

• $N$ floats for the model parameters themselves
• $N$ floats for the parameter gradients (computed during backpropagation)
• $2N$ floats for the Adam first and second moments

The total is $4N$ floats, four times the model's parameter count. For float32 training, each float is 4 bytes, so a 7 billion parameter model requires $4 \times 7 \times 10^9 \times 4 = 112$ GB of memory just for the training state. This far exceeds the capacity of any single GPU (which typically has 40-80 GB of HBM memory), requiring distributed training across many devices.

Several strategies reduce this memory footprint:

ZeRO (Zero Redundancy Optimizer) from Microsoft DeepSpeed shards the optimizer state, gradients, and parameters across all GPUs in a distributed training job. In the most aggressive mode (ZeRO-3), each GPU holds only $1/N_{\text{GPU}}$ of the model weights, gradients, and Adam moments. The memory per GPU scales inversely with the number of GPUs, making it possible to train models far larger than any single GPU's capacity.

8-bit Adam (Dettmers et al., 2022, available via the bitsandbytes library) quantizes the optimizer states to 8-bit integers, reducing the 8 bytes per parameter for Adam state (two float32 tensors) to 2 bytes per parameter. The quantization is done with a dynamic block-wise scheme that preserves the optimizer's behavior almost exactly, with negligible accuracy loss. This is perhaps the most practical approach for single-GPU training of large models.

Mixed precision training uses float16 for the forward and backward passes (saving memory on activations and intermediate computations) while keeping the optimizer state and master copy of the weights in float32. The Adam step itself is computed in float32 and then cast back to float16 for the next forward pass. This is the standard approach for training transformers on modern hardware and roughly halves memory usage compared to full float32 training.

Adam is a powerful and broadly applicable optimizer, but it has several limitations that practitioners encounter regularly.

Adam stores two additional tensors per parameter: $m_t$ and $v_t$. For large models, this triples the memory required for optimizer state compared to SGD (which stores one tensor per parameter for momentum, or zero for vanilla SGD). A 7 billion parameter model requires roughly 7B floats for the parameters themselves. Adam requires another 14B floats (two moments), for a total of 21B floats of optimizer state. At float32, that is 84 GB, which is prohibitive on a single GPU.

Remedies include:

• 8-bit Adam: The bitsandbytes library implements Adam with 8-bit quantized optimizer states, reducing memory by roughly 4x with minimal accuracy impact.
• ZeRO: Distributed optimizer state sharding across GPUs, so each GPU holds only a fraction of the moments.
• Adafactor: An adaptive optimizer with $O(1)$ optimizer state that factors the second moment matrix, used for very large models.

As noted in the comparison section, Adam can find sharper minima than SGD on some image classification benchmarks, leading to worse test accuracy despite better training accuracy. This problem is less prevalent in NLP, where the datasets are large enough that sharpness does not correlate strongly with generalization. If you observe a training-test gap that seems larger than expected, two practical remedies are AdamW (which decouples weight decay and tends to find broader minima) and data augmentation (which artificially widens the training distribution, making sharp minima less likely to perfectly fit the training set).

The theoretical convergence guarantee for Adam relies on assumptions (convexity, or specific non-convex conditions) that are not always met. AMSGrad modifies the second moment estimate to use the maximum of all past estimates, guaranteeing convergence in the online convex setting. However, empirical comparisons typically show Adam converging as well as or better than AMSGrad on practical tasks, so AMSGrad has seen limited adoption. The theoretical non-convergence of Adam in worst-case settings does not appear to manifest in the loss landscapes of practical neural networks, which have much more benign structure than the adversarial examples used in the theoretical analysis.

A subtle but important issue is that when you add weight_decay to PyTorch's torch.optim.Adam, you are not getting true weight decay. You are getting L2 regularization, which adds $\lambda \theta$ to the gradient before the update. With Adam's adaptive scaling, this L2 penalty is also divided by $\sqrt{\hat{v}_t}$, which means the actual regularization strength varies by parameter and is weaker for frequently updated parameters (which often need regularization most). This interaction defeats much of the purpose of weight decay as a regularizer. This is why AdamW, introduced by Loshchilov and Hutter in 2017, applies weight decay directly to the parameters after the Adam step rather than through the gradient. AdamW is the subject of the next chapter, and it is the recommended optimizer for essentially all modern neural language modeling.

A subtle practical issue is that Adam interacts with learning rate schedules in a non-obvious way. When the learning rate schedule reduces $\alpha$ by a large factor late in training (as cosine annealing does), the moment estimates $m_t$ and $v_t$ were built up under the higher learning rate regime. The second moment $v_t$ reflects the gradient magnitudes that caused the Adam updates when $\alpha$ was large. As $\alpha$ decreases, the adaptive learning rates per parameter shrink proportionally, but the relationship between $v_t$ and the current training dynamics may no longer be well-calibrated. Some practitioners reset the Adam optimizer state when making large learning rate jumps (such as at the beginning of a second training stage) to avoid carrying over stale moment estimates. This is sometimes called "warm restarting" the optimizer, and it can produce a brief spike in loss followed by improved convergence in the new learning rate regime.

Adam brings together two key ideas that address fundamental limitations of fixed-learning-rate optimizers.

The first idea is per-parameter adaptive learning rates. By tracking the exponential moving average of squared gradients ($v_t$), Adam determines how large gradients have been for each parameter. Parameters with historically large gradients get a compressed step size; parameters with historically small gradients get a boosted step size. This equalization means Adam can train models with very different gradient scales across layers and across dense versus sparse parameters without manual per-layer tuning. The mechanism is essentially a diagonal approximation to the natural gradient, rescaling each parameter's update by the inverse of its typical gradient magnitude.

The second idea is bias-corrected moment initialization. Starting both moment estimates at zero creates a cold-start problem where early updates are either too small (first moment) or divide by an incorrectly small estimate (second moment). The bias correction terms $(1 - \beta_1^t)$ and $(1 - \beta_2^t)$ remove this distortion, making Adam's behavior predictable and stable from the very first update. This is especially important because it makes Adam's early training well-behaved without requiring special initialization, which is one reason the optimizer feels so reliable out of the box.

Together with the momentum-smoothed first moment $m_t$, which provides direction stability by averaging out gradient noise, these two mechanisms give Adam its characteristic reliability: it converges quickly on a wide range of architectures with minimal hyperparameter tuning, using the default $\alpha = 0.001$, $\beta_1 = 0.9$, $\beta_2 = 0.999$.

The history from SGD to momentum to AdaGrad to RMSProp to Adam represents a steady accumulation of understanding about what makes gradients difficult to follow. SGD addressed the basic problem of finding a descent direction. Momentum addressed the noise problem by smoothing that direction. AdaGrad addressed the scale problem by normalizing by gradient history. RMSProp fixed AdaGrad's monotonic decay. Adam combined the smoothing and normalization advances with proper statistical initialization. Each step solved the most obvious remaining failure mode of its predecessor.

The key limitations to keep in mind are memory cost (2x more optimizer state than SGD), the distinction between L2 regularization and true weight decay (use AdamW for language models), potential generalization issues on some vision tasks, and the interactions with mixed-precision and distributed training that require careful configuration. The next chapter covers AdamW, which fixes the weight decay issue and has become the standard optimizer for training transformer models from BERT and GPT through to the most recent large language models.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about the Adam optimizer.