Why AdamW is the King of AI Training

Topic Overview

What is Adam?

Adam (Adaptive Moment Estimation) is the undisputed default optimizer for modern deep learning. It is a mathematically rigorous combination of two previous algorithms: it calculates the first moment (mean) of the gradients to build directional inertia (like Momentum), and the second moment (uncentered variance) of the gradients to scale the step size per parameter (like RMSProp).

Why does it exist?

Before Adam, engineers had to make a difficult choice. If your loss landscape had steep cliffs, you used RMSProp to prevent divergence. If your loss landscape had long, flat plateaus, you used Momentum to avoid stalling. In complex models like Transformers, the loss landscape has both simultaneously, and different layers exhibit different behaviors. Manually tuning between these extremes was unmaintainable. Adam exists to combine both mechanics into a single, robust algorithm that “just works” out of the box.

What problem does it solve?

It solves the problem of highly non-stationary, sparse, and noisy gradients in massive parameter spaces. It ensures that a 7-billion parameter language model can update its embedding layers (which see sparse, massive gradients) and its deep attention layers (which see dense, small gradients) using the exact same optimizer configuration, without either layer diverging or stalling.

Why should engineers care?

If you train a Large Language Model (LLM), a Vision Transformer (ViT), or a Diffusion Model, you are using AdamW (a fixed variant of Adam). Understanding Adam means understanding exactly why your GPU VRAM is maxed out (the 3x memory multiplier), why you must use weight decay correctly, and why the “cold start” bias correction is the secret sauce that makes the first few hundred steps of training actually work.

Prerequisites

1. Momentum

Why it is needed: Adam’s \(\mathbf{m}_t\) equation is literally the Momentum equation. You must understand how an Exponentially Weighted Moving Average (EWMA) builds directional speed and dampens oscillations to understand Adam’s first moment.
Where it appears: The \(\mathbf{m}_t = \beta_1 \mathbf{m}_{t-1} + (1-\beta_1)\mathbf{g}_t\) calculation.

2. RMSProp

Why it is needed: Adam’s \(\mathbf{v}_t\) equation is literally the RMSProp equation. You must understand how accumulating squared gradients normalizes the learning rate per dimension to understand Adam’s second moment.
Where it appears: The \(\mathbf{v}_t = \beta_2 \mathbf{v}_{t-1} + (1-\beta_2)\mathbf{g}_t^2\) calculation.

Historical Motivation

In 2014, the deep learning community was fractured. Computer Vision researchers swore by SGD + Momentum. Natural Language Processing researchers swore by RMSProp (or Adagrad) because word embeddings created massively sparse gradients that destroyed Momentum.

The Engineering Failure: Researchers tried manually combining them (e.g., dividing the Momentum update by the RMSProp denominator). This led to chaotic training. The core mathematical issue was the “cold start” problem: both Momentum and RMSProp initialize their state buffers to zero. In the first few steps, these buffers are heavily biased towards zero, causing the optimizer to take artificially tiny steps, regardless of the actual gradient magnitude. Manual hacks to fix this were brittle.

The Emergence: Diederik Kingma and Jimmy Ba published “Adam: A Method for Stochastic Optimization.” Their engineering contribution wasn’t just adding Momentum to RMSProp; it was deriving a mathematically exact bias correction term that instantly unbias those zero-initialized buffers at step \(t=1\). This made Adam stable, self-contained, and immediately effective across all domains.

Intuitive Explanation

Imagine you are driving a heavy manual-transmission car down a mountain (Vanilla SGD). * Momentum adds a heavy flywheel to the engine. It accelerates smoothly on straightaways but struggles to shift gears on steep inclines. * RMSProp adds an automatic transmission. It shifts gears perfectly for the incline, but the car has no mass, so a slight breeze (noisy gradient) stops it entirely.

Adam is the combination: A heavy car with an automatic transmission. It has the mass to power through flat plateaus (Momentum), and the transmission to crawl safely down steep cliffs without crashing (RMSProp).

The “Cold Start” / Bias Correction Intuition: There is one final problem. When you first get into this car, the automatic transmission computer is completely blank (initialized to zero). It thinks you are going 0 mph, so it refuses to shift out of 1st gear. Bias Correction is a mathematical trick that acts like a sensor calibrator. It looks at the fact that the computer has only been running for 1 second, realizes the 0 mph reading is a lie caused by the blank slate, and artificially amplifies the reading to what it should be. By step 100, the computer has enough history that the calibrator fades away entirely.

Visual Understanding

Draw a 2D contour plot of an extreme loss landscape: a long, flat plateau that suddenly drops off into a narrow, steep canyon.

Vanilla SGD: Crawls across the plateau. Hits the canyon wall and violently explodes (diverges). Momentum: Accelerates nicely across the plateau. Hits the canyon and bounces violently left and right, slowly inching down. RMSProp: Crawls slowly across the plateau (no mass to carry it). Hits the canyon and safely steps down the steep wall, but takes forever because steps are tiny. Adam: 1. On the plateau: The second moment (\(\mathbf{v}\)) sees tiny gradients and scales the learning rate up. The first moment (\(\mathbf{m}\)) accumulates this scaled gradient. The car accelerates smoothly and rapidly across the plateau. 2. At the cliff: The gradient suddenly spikes. The second moment (\(\mathbf{v}\)) accumulates this spike and scales the learning rate down heavily. The first moment (\(\mathbf{m}\)) is immediately dampened by this scaling. The car safely brakes and steps smoothly down the cliff wall without bouncing.

Mathematical Foundations

Let’s build Adam piece by piece.

Step 1: The First Moment (Momentum) \[\mathbf{m}_t = \beta_1 \mathbf{m}_{t-1} + (1-\beta_1)\mathbf{g}_t\] * \(\mathbf{m}_t\): The running mean (direction + scaled magnitude) of the gradient. * \(\beta_1\): Usually \(0.9\). Looks back roughly 10 steps.

Step 2: The Second Moment (RMSProp) \[\mathbf{v}_t = \beta_2 \mathbf{v}_{t-1} + (1-\beta_2)\mathbf{g}_t^2\] * \(\mathbf{v}_t\): The running uncentered variance (magnitude only) of the gradient. * \(\beta_2\): Usually \(0.999\). Looks back roughly 1000 steps. (Why so long? We want the step-size scaling to be incredibly stable and immune to sudden gradient spikes).

Step 3: The Bias Correction (The Genius of Adam) Because \(\mathbf{m}_0 = 0\) and \(\mathbf{v}_0 = 0\), these equations are heavily biased towards zero at the beginning of training. For example, at \(t=1\): \(\mathbf{m}_1 = (1-\beta_1)\mathbf{g}_1\). This is \(0.1 \times \mathbf{g}_1\), meaning our estimate of the gradient’s mean is 10x too small.

We derive the correction by taking the expected value of the EWMA. To get an unbiased estimate, we divide by the sum of the geometric series weights: \[\hat{\mathbf{m}}_t = \frac{\mathbf{m}_t}{1 - \beta_1^t}\] \[\hat{\mathbf{v}}_t = \frac{\mathbf{v}_t}{1 - \beta_2^t}\] * At \(t=1\), denominator is \(1 - 0.9^1 = 0.1\). We divide \(\mathbf{m}_1\) by \(0.1\), perfectly restoring the true gradient magnitude. * At \(t \to \infty\), \(\beta^t \to 0\), so the denominator \(\to 1\), and the correction vanishes.

Step 4: The Update \[\theta_{t+1} = \theta_t - \eta \frac{\hat{\mathbf{m}}_t}{\sqrt{\hat{\mathbf{v}}_t} + \epsilon}\] * We step using the bias-corrected direction (\(\hat{\mathbf{m}}\)), scaled by the bias-corrected RMSProp denominator (\(\sqrt{\hat{\mathbf{v}}}\)).

Worked Numerical Example

Let’s track a single weight through the first 2 iterations to see the bias correction in action. * \(\theta_0 = 2.0\), \(\eta = 0.1\), \(\beta_1 = 0.9\), \(\beta_2 = 0.999\), \(\epsilon = 10^{-8}\) * \(\mathbf{m}_0 = 0\), \(\mathbf{v}_0 = 0\) * Assume constant gradient: \(\mathbf{g}_1 = \mathbf{g}_2 = 4.0\)

Iteration 1 (\(t=1\)): 1. \(\mathbf{m}_1 = 0.9(0) + 0.1(4.0) = \mathbf{0.4}\) 2. \(\mathbf{v}_1 = 0.999(0) + 0.001(16.0) = \mathbf{0.016}\) 3. Bias Correction: * \(\hat{\mathbf{m}}_1 = 0.4 / (1 - 0.9^1) = 0.4 / 0.1 = \mathbf{4.0}\) (Corrected back to true gradient!) * \(\hat{\mathbf{v}}_1 = 0.016 / (1 - 0.999^1) = 0.016 / 0.001 = \mathbf{16.0}\) (Corrected back to true squared gradient!) 4. Update: * Denominator = \(\sqrt{16.0} + 0 = 4.0\) * Step = \(0.1 \cdot (4.0 / 4.0) = 0.1\) * \(\theta_1 = 2.0 - 0.1 = \mathbf{1.9}\)

Notice: Without bias correction, the step would have been \(0.1 \cdot (0.4 / \sqrt{0.016}) \approx 0.316\). The bias correction perfectly normalized the first step.

Iteration 2 (\(t=2\)): 1. \(\mathbf{m}_2 = 0.9(0.4) + 0.1(4.0) = 0.36 + 0.4 = \mathbf{0.76}\) 2. \(\mathbf{v}_2 = 0.999(0.016) + 0.001(16.0) = 0.015984 + 0.016 = \mathbf{0.031984}\) 3. Bias Correction: * \(1 - \beta_1^2 = 1 - 0.81 = 0.19\) * \(\hat{\mathbf{m}}_2 = 0.76 / 0.19 \approx \mathbf{4.0}\) * \(1 - \beta_2^2 = 1 - 0.998001 = 0.001999\) * \(\hat{\mathbf{v}}_2 = 0.031984 / 0.001999 \approx \mathbf{16.0}\) 4. Update: Step remains \(\approx 0.1\). \(\theta_2 \approx \mathbf{1.8}\)

Computational Interpretation

What is the exact hardware cost of this “just works” optimizer?

Memory Representation: To train a model, the GPU must store three distinct states per parameter: 1. \(\theta\): The weights themselves (FP16 or FP32). 2. \(\mathbf{m}\): First moment buffer (Must be FP32 to prevent precision loss in the EWMA). 3. \(\mathbf{v}\): Second moment buffer (Must be FP32).

The 3x Memory Wall: If you have a 7 Billion parameter model (like LLaMA-7B): * Weights in FP16: \(\sim 14\) GB * Gradients in FP16: \(\sim 14\) GB (temporary) * Adam \(\mathbf{m}\) buffer in FP32: \(\sim 28\) GB * Adam \(\mathbf{v}\) buffer in FP32: \(\sim 28\) GB * Total VRAM just for the optimizer: \(\mathbf{56}\) GB. * This is the primary reason why training LLMs requires massive clusters of 80GB A100 GPUs. Adam is incredibly memory-hungry.

Compute Operations: It requires 4 multiplies, 2 adds, 1 square root, 1 division per parameter. The square root and division are computationally expensive compared to SGD’s simple addition, but modern GPUs handle them efficiently via vectorized math libraries (cuBLAS).

Implementation Perspective

1. NumPy (From Scratch)

import numpy as np

def adam_step(theta, grad, m, v, t, lr=0.001, b1=0.9, b2=0.999, eps=1e-8):
    """
    theta: Parameters (N-dim)
    grad: Gradients (N-dim)
    m: First moment buffer (N-dim)
    v: Second moment buffer (N-dim)
    t: Current timestep integer (CRITICAL for bias correction)
    """
    # 1. Update biased first moment estimate
    m = b1 * m + (1.0 - b1) * grad
    
    # 2. Update biased second raw moment estimate
    v = b2 * v + (1.0 - b2) * (grad ** 2)
    
    # 3. Compute bias-corrected first moment estimate
    m_hat = m / (1.0 - b1 ** t)
    
    # 4. Compute bias-corrected second raw moment estimate
    v_hat = v / (1.0 - b2 ** t)
    
    # 5. Update parameters
    theta = theta - lr * (m_hat / (np.sqrt(v_hat) + eps))
    
    return theta, m, v

# Init
params = np.array([2.0])
m = np.zeros_like(params)
v = np.zeros_like(params)

# Step 1
params, m, v = adam_step(params, np.array([4.0]), m, v, t=1)
# Step 2
params, m, v = adam_step(params, np.array([4.0]), m, v, t=2)

2. PyTorch (Production Standard - AdamW)

CRITICAL ENGINEERING NOTE: You will almost never use torch.optim.Adam in modern deep learning. You will use torch.optim.AdamW.

Original Adam coupled L2 regularization (weight decay) into the gradient update. Because Adam divides by \(\sqrt{\mathbf{v}}\), the weight decay was scaled inversely by the gradient magnitude, effectively destroying the decay for parameters with large gradients. AdamW decouples this, applying weight decay directly to the weights.

import torch
import torch.nn as nn
import torch.optim as optim

layer = nn.Linear(10, 2)

# AdamW applies weight decay directly: theta = theta - lr * wd * theta
# This is mathematically correct for Transformers.
optimizer = optim.AdamW(
    layer.parameters(), 
    lr=1e-4,              # Standard for fine-tuning
    betas=(0.9, 0.999),   # Defaults
    eps=1e-8,              # Default. DO NOT change to 1.0
    weight_decay=0.01      # Standard regularization strength
)

# Advanced: Per-layer learning rates (Common in NLP)
# Often, we want the classification head to learn faster than the pretrained backbone
optimizer = optim.AdamW([
    {'params': model.transformer.parameters(), 'lr': 1e-5}, # Slow
    {'params': model.classifier.parameters(), 'lr': 1e-3}   # Fast
], weight_decay=0.01)

Where It Appears In AI

AI Component	Exact Usage
Transformers (Self-Attention)	The Q, K, V projection matrices have highly unstable gradient spectra. Adam’s second moment (\(\mathbf{v}\)) normalizes these spectra, preventing the attention logits from exploding. SGD rarely works here.
LLM Pre-training (GPT, LLaMA)	AdamW is the sole optimizer used. The massive embedding layers (sparse gradients) and deep MLPs (dense gradients) require the dual Momentum + RMSProp mechanics to train stably.
Diffusion Models (U-Net, DiT)	AdamW is used to predict noise. The variance of the loss shifts drastically as the diffusion timestep changes; Adam’s fast \(\beta_2\) adaptation handles this non-stationarity.
Fine-tuning (LoRA)	When training Low-Rank Adaptations, AdamW is used on the tiny LoRA matrices (\(A\) and \(B\)), providing rapid convergence with minimal hyperparameter tuning.
Reinforcement Learning (PPO)	Used to update the policy and value networks in RLHF. The high-variance policy gradients are smoothed by Adam’s first moment.

Deep Dive Into Real Models

In a Transformer (LLM Pre-training)

Consider the training step of a 70B parameter model. The batch consists of sequences of text.

Embedding Layer: The gradient is sparse. Only the specific token IDs present in the batch get updates.
- Adam’s role: The \(\mathbf{v}\) buffer for unused tokens remains frozen (decays slightly due to \(\beta_2\)). The \(\mathbf{v}\) buffer for used tokens spikes. Adam automatically scales down the step for frequently seen tokens and scales it up for rare tokens.
LayerNorm / RMSNorm: Gradients here can be massive and sudden.
- Adam’s role: The \(\mathbf{v}\) buffer acts as a shock absorber. The denominator \(\sqrt{\mathbf{v}}\) instantly grows, throttling the learning rate and preventing the layer norm weights from destroying the activation distribution.
The AdamW Update: After the Adam step calculates the new weights based on the gradient, the AdamW algorithm executes: \(\theta = \theta - \eta \cdot \lambda \cdot \theta\). This direct decay prevents the massive weight matrices from growing unboundedly, independent of the gradient statistics.

Failure Modes

1. The 3x Memory Wall

Cause: Storing \(\mathbf{m}\) and \(\mathbf{v}\) in FP32 for billions of parameters.
Symptom: torch.cuda.OutOfMemoryError during optimizer.step(), even though the forward/backward pass fit in memory.
Solution: Use bitsandbytes library to implement 8-bit Adam. It quantizes the \(\mathbf{m}\) and \(\mathbf{v}\) states to 8-bit integers, dropping optimizer memory from 56GB to ~14GB for a 7B model, with zero degradation in performance.

2. Generalization Gap

Cause: Adam’s aggressive per-parameter adaptation tends to find extremely sharp, narrow minima in the loss landscape.
Symptoms: Training loss drops beautifully, but validation loss is noticeably higher than a model trained with SGD+Momentum. The model is brittle to out-of-distribution data.
Solution: For Computer Vision (ResNet, ViT image classification), SGD+Momentum is still often preferred for final accuracy. For NLP/Generative models, the gap is mitigated by using strong weight decay (AdamW) and data augmentation.

3. The \(\epsilon\) Mismatch

Cause: Copying an optimizer from TensorFlow (which sometimes default \(\epsilon=1.0\)) to PyTorch (which defaults \(\epsilon=10^{-8}\)).
Symptoms: Training completely stalls. Loss does not decrease.
Solution: Never change \(\epsilon\) from \(10^{-8}\) in PyTorch unless explicitly tuning for a specific mathematical reason (like stabilizing extremely small gradients in RL).

Engineering Insights

1. FP16 Mixed Precision Interactions When using PyTorch’s automatic mixed precision (AMP), the gradients are computed in FP16. If you pass these FP16 gradients directly into Adam’s \(\mathbf{v} = \mathbf{v} + g^2\) calculation, the \(g^2\) operation will underflow to zero for any gradient \(< 0.00006\). * Insight: PyTorch’s AdamW implementation automatically casts the incoming gradients to FP32 before updating the \(\mathbf{m}\) and \(\mathbf{v}\) buffers. Never write a custom optimizer that skips this cast.

2. Fused Kernels A standard Adam update requires reading \(\theta\), reading \(\mathbf{m}\), reading \(\mathbf{v}\), reading \(g\), and writing them all back. This causes massive memory bandwidth bottlenecks. NVIDIA provides FusedAdam (used natively in PyTorch via torch.optim.AdamW(fused=True)). It launches a single GPU kernel that does all the math in registers/cache, reducing memory read/writes by roughly 50% and speeding up the optimizer step significantly.

Interview Questions

Beginner

Q: What are the two “moments” that Adam estimates? A: The first moment is the mean of the gradients (which provides directional inertia, like Momentum). The second moment is the uncentered variance of the gradients (which provides per-parameter step-size scaling, like RMSProp).

Intermediate

Q: Why does Adam require bias correction, whereas vanilla Momentum and RMSProp do not? A: Vanilla Momentum and RMSProp are often implemented without strict mathematical guarantees on the exact magnitude of the first step; they just “warm up.” Adam, however, applies the update as a fraction (\(\mathbf{m} / \sqrt{\mathbf{v}}\)). If both are initialized to zero, this fraction is \(0/0\), and for the first few steps, it is heavily skewed towards zero, causing the optimizer to take artificially tiny steps. Bias correction mathematically inflates the buffers early on to counteract the zero-initialization, ensuring the step size is accurate from step 1.

Advanced

Q: Explain exactly why the original Adam’s implementation of L2 regularization (weight decay) was mathematically flawed, and how AdamW fixes it. A: In standard L2, we add \(\lambda \theta\) to the gradient: \(g_{new} = g + \lambda \theta\). In Adam, this new gradient is divided by \(\sqrt{\mathbf{v}_t}\). If a weight \(\theta\) is very large, the weight decay term is large, but if the historical gradients \(\mathbf{v}_t\) for that weight are also large, the denominator scales the decay down, preventing the large weight from shrinking properly. AdamW decouples this: it applies the Adam update using only the data gradient \(g\), and then separately applies the decay directly to the weights: \(\theta = \theta - \eta \lambda \theta\). This ensures weight decay is proportional to the weight’s magnitude, independent of the gradient history.

Research

Q: How do modern optimizers like Sophia or Lion attempt to improve upon Adam? A: Lion (EvoLoved Sign Momentum) removes the second moment (\(\mathbf{v}\)) entirely to save memory and compute, using only the sign of the momentum update. Sophia addresses Adam’s tendency to find sharp minima by replacing the second moment (\(\mathbf{v}\), which tracks gradient variance) with the diagonal of the Hessian matrix (which tracks actual curvature). By dividing by the Hessian diagonal, Sophia takes steps that are preconditioned by the true loss landscape geometry, leading to faster convergence and flatter minima than Adam.

Research Perspective

Current Limitations: Adam’s biggest limitation is the memory wall. Storing two FP32 buffers for every parameter limits the size of models that can be trained on a given GPU. Secondly, Adam’s diagonal assumption (treating each parameter independently) ignores the off-diagonal correlations in the Hessian matrix, meaning it cannot efficiently navigate long, narrow valleys that are not axis-aligned.

Modern Directions: 1. Memory-Efficient Adaptation: Algorithms like Adafactor factorize the \(\mathbf{v}\) matrix (separating row and column statistics) to drastically reduce memory, allowing billions of parameters to be trained on consumer GPUs. 2. Second-Order Approximations: Moving beyond the diagonal \(\mathbf{v}\) matrix to cheaply estimate the diagonal of the Hessian (like Sophia) or use K-FAC (Kronecker-Factored Approximate Curvature) to capture 2D correlations in weight matrices. 3. Schedule-Free Optimization: Recent research removes the need for learning rate warmup and decay schedules entirely by dynamically adjusting the effective step size based on the distance between the current parameters and a moving average of past parameters.

Connections

Prerequisites: * Momentum (First moment mechanics) * RMSProp (Second moment mechanics) * Exponentially Weighted Moving Averages (Mathematical foundation)

Enables: * Stable training of Sparse Gradients (Embeddings) * Stable training of Non-Stationary Objectives (RL, RNNs) * “Out of the box” optimization requiring minimal LR tuning

Used by: * Every major LLM (GPT, LLaMA, Gemini via AdamW) * Diffusion Models (Stable Diffusion, DiT) * Reinforcement Learning from Human Feedback (PPO)

Extended by: * AdamW (Decoupled weight decay) * RAdam (Rectified Adam, variance stabilization at warmup) * 8-bit Adam (Memory optimization via quantization)

Related concepts: * Learning Rate Scheduling (Often simplified in Adam, but warmup is still used to let the \(\mathbf{v}\) buffer stabilize) * Weight Decay / L2 Regularization (Coupling vs Decoupling) * Loss Landscape Geometry (Sharp vs. Flat minima)

Final Mental Model

To permanently remember Adam, visualize the “Smart Cruise Control with a Calibrator.”

Imagine driving a vehicle over terrain that instantly switches between ice (tiny gradients), gravel (noisy gradients), and steep hills (massive gradients).

The Gas Pedal (First Moment, \(\mathbf{m}\)): Smoothly presses the accelerator based on the average direction you’ve been told to go, ignoring potholes (noise).
The Smart Transmission (Second Moment, \(\sqrt{\mathbf{v}}\)): Senses the terrain. On ice, it drops into low gear (amplifies the small gas pedal input). On gravel, it stabilizes. On steep hills, it shifts way up to prevent you from accelerating too fast.
The Calibrator (Bias Correction): When you first turn the car on, the transmission computer thinks you’re on flat ground (initialized to 0). The calibrator overrides the computer for the first 10 seconds, ensuring you don’t stall at the starting line, then fades away once the computer has real data.

This combination of smooth acceleration, terrain-aware gearing, and safe startup is why Adam drives the entirety of modern AI.

Neural Nexus

Search This Blog