Optimizers & Training Dynamics

Adam was not the first optimizer, nor the last, but somehow it keeps showing up at every training run like a reliable friend who always picks the restaurant. You could do better, but you probably will not.

Scale, Stubbornly Optimizing AI Agent

6.5.1 Stochastic Gradient Descent and Its Limitations

Vanilla SGD updates parameters by subtracting the gradient scaled by a learning rate: $\theta \leftarrow \theta - \eta \nabla L( \theta )$. While SGD with momentum works well for convolutional networks, it performs poorly for transformer training. The core issue is that transformers have parameters operating at vastly different scales: embedding matrices, attention projections, layer norms, and feed-forward layers all have different gradient magnitudes. A single learning rate cannot simultaneously provide appropriate step sizes for all of these.

6.5.2 Adam: The Default Optimizer for Transformers

Adam (Adaptive Moment Estimation) addresses this by maintaining per-parameter adaptive learning rates based on first and second moment estimates of the gradient. The update rules are:

The second moment estimate tracks the squared gradient magnitude, capturing how variable each parameter's gradient has been:

After bias correction, the parameter update divides the momentum by the square root of the variance, providing adaptive per-parameter learning rates:

where $\hat{m}$ and $\hat{v}$ are bias-corrected estimates (dividing by $1 - \beta ^{t}$). The typical hyperparameters are $\beta _{1} = 0.9$, $\beta _{2} = 0.999$, and $\epsilon = 10^{-8}$.

# Adam optimizer memory budget: each parameter needs 12 bytes (FP32 weight
# + FP32 first moment + FP32 second moment), so 7B params requires ~84 GB.
import torch
# Understanding Adam's memory footprint
model_params = 7e9 # 7B parameters
bytes_per_param = 4 # FP32
param_memory = model_params * bytes_per_param
gradient_memory = model_params * bytes_per_param
adam_m_memory = model_params * bytes_per_param # first moment
adam_v_memory = model_params * bytes_per_param # second moment
total = param_memory + gradient_memory + adam_m_memory + adam_v_memory
print(f"Parameters:       {param_memory / 1e9:.1f} GB")
print(f"Gradients:        {gradient_memory / 1e9:.1f} GB")
print(f"Adam m (1st mom): {adam_m_memory / 1e9:.1f} GB")
print(f"Adam v (2nd mom): {adam_v_memory / 1e9:.1f} GB")
print(f"Total:            {total / 1e9:.1f} GB")

from transformers import get_cosine_schedule_with_warmup
import torch
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=2000,
    num_training_steps=100000
    )
# In your training loop:
for step, batch in enumerate(dataloader):
    loss = model(batch).loss
    loss.backward()
    optimizer.step()
    scheduler.step() # updates LR automatically
    optimizer.zero_grad()

6.5.3 AdamW: Decoupled Weight Decay

Loshchilov and Hutter (2019) identified a subtle but important flaw in Adam's handling of weight decay (L2 regularization). In standard Adam, weight decay is applied to the gradient before the adaptive scaling, which means the effective regularization strength varies across parameters based on their second moment estimates. AdamW fixes this by decoupling weight decay from the gradient update:

Here, $\lambda$ is the weight decay coefficient (typically 0.01 to 0.1), and it is applied uniformly to all parameters regardless of their gradient history. This decoupling is particularly important for transformers because different parameter groups (attention weights, embeddings, biases) benefit from consistent regularization strength.

# Minimal AdamW update step in NumPy, matching torch.optim.AdamW.
# The key line is the decoupled decay theta *= (1 - lr * weight_decay)
# applied OUTSIDE the adaptive m / sqrt(v) ratio.
import numpy as np

def adamw_step(theta, grad, m, v, t, lr=3e-4, betas=(0.9, 0.999),
               eps=1e-8, weight_decay=0.01):
    b1, b2 = betas
    m = b1 * m + (1 - b1) * grad
    v = b2 * v + (1 - b2) * grad ** 2
    m_hat = m / (1 - b1 ** t)              # bias-corrected first moment
    v_hat = v / (1 - b2 ** t)              # bias-corrected second moment
    theta = theta * (1.0 - lr * weight_decay)   # decoupled weight decay
    theta = theta - lr * m_hat / (np.sqrt(v_hat) + eps)
    return theta, m, v

theta = np.array([1.0, -2.0, 0.5])
m, v = np.zeros_like(theta), np.zeros_like(theta)
for t in range(1, 4):
    grad = np.array([0.1, -0.2, 0.05])     # toy fixed gradient
    theta, m, v = adamw_step(theta, grad, m, v, t)
    print(f"step {t}: theta = {theta}")

6.5.4 Memory-Efficient Optimizer Alternatives

AdamW wins on accuracy but loses badly on memory: its two moment buffers triple the parameter footprint, and a 70B model needs roughly 840 GB just for optimizer state in FP32. Three families of memory-efficient alternatives (Adafactor, 8-bit Adam, and Lion) sacrifice some of AdamW's robustness in exchange for 50% to 87% less optimizer memory. We work through them in order of increasing aggressiveness.

Adafactor

Adafactor (Shazeer and Stern, 2018) reduces Adam's memory overhead by factoring the second-moment matrix. Instead of storing a full $v$ tensor for each parameter matrix, Adafactor stores only the row and column statistics: two vectors whose outer product approximates the full matrix. For a weight matrix of shape $(m, n)$, this reduces the second-moment storage from $mn$ to $m + n$. Adafactor was used in the T5 model family.

8-bit Adam

Dettmers et al. (2022) showed that Adam's optimizer states (m and v) can be quantized to 8-bit integers with dynamic block-wise quantization, reducing optimizer memory by 75% with negligible impact on training quality. The key insight is that optimizer states do not need full precision: they accumulate slowly changing statistics, and small quantization errors are averaged out over training steps.

LION (Sign-Based Optimizer)

LION (Chen et al., 2023) takes a radically different approach. Instead of using the full gradient magnitude, LION uses only the sign of the momentum update: every parameter is updated by exactly $+ \eta$ or $- \eta$. This eliminates the second moment entirely, cutting optimizer memory in half compared to Adam. LION also uses a different momentum interpolation that mixes past momentum with the current gradient. Despite its simplicity, LION matches or exceeds AdamW on many vision and language tasks.

Muon and Next-Generation Optimizers

While AdamW remains the workhorse of LLM training, a wave of newer optimizers is challenging its dominance. Muon (Jordan et al., 2024), short for "Momentum + Orthogonalization," takes a fundamentally different approach to parameter updates. Instead of maintaining per-parameter adaptive learning rates like Adam, Muon applies orthogonalization to the momentum buffer at each step, projecting the update direction onto the Stiefel manifold. In simpler terms, Muon ensures that each update step moves the weight matrix in a direction that preserves the orthogonality of its rows or columns, which helps maintain well-conditioned representations throughout training.

Muon gained attention when it achieved state-of-the-art results on the nanogpt-speedrun benchmark, training a GPT-2 scale model faster than any AdamW configuration. Its key advantage is that it requires only a single momentum buffer (like LION) rather than two (like Adam), cutting optimizer memory in half (a savings that compounds with the memory reductions from parameter-efficient methods like LoRA and QLoRA). Unlike LION, Muon preserves gradient magnitude information through the orthogonalization step rather than discarding it via sign operations.

Several other optimizers have emerged alongside Muon. SOAP (Vyas et al., 2024) uses Shampoo-style second-order preconditioning but approximates the expensive matrix operations to reduce overhead. Schedule-Free Adam (Defazio and Mishchenko, 2024) eliminates the learning rate schedule entirely by incorporating averaging theory directly into the optimizer update, removing the need to tune warmup steps and decay schedules. Adalayer applies different effective learning rates to different layers automatically, addressing the well-known issue that optimal learning rates vary across a model's depth.

Despite these innovations, AdamW is unlikely to be displaced quickly in production settings. Optimizer changes interact with every other hyperparameter (learning rate, weight decay, batch size, gradient clipping), and the training recipes for major models have been extensively tuned for AdamW. When fine-tuning pretrained models (covered in Chapter 16), these same optimizer choices carry over, though lower learning rates are typically used. New optimizers must demonstrate not just competitive loss curves on small benchmarks but stable, reliable behavior across months-long training runs at trillion-token scale. That said, Muon and Schedule-Free Adam are increasingly appearing in research-scale training runs, and future large-scale training recipes will likely incorporate ideas from these newer designs.

6.5.5 Learning Rate Schedules

Warmup

Transformer training universally begins with a warmup phase where the learning rate increases linearly from near-zero to the peak value over several hundred to several thousand steps. Warmup is necessary because at initialization, the model's loss landscape is poorly conditioned: gradients can be very large and noisy. Starting with a high learning rate would cause catastrophic parameter updates that push the model into a bad region of the loss landscape from which it cannot recover. Warmup allows the optimizer's moment estimates to stabilize before applying large updates.

Cosine Decay

After warmup, the learning rate is typically decayed following a cosine schedule:

where $T$ is the total number of training steps. Cosine decay provides a smooth, gradual reduction that spends most of the training budget at moderate learning rates. The minimum learning rate is typically set to 10% of the peak rate. Some variants use cosine decay with warm restarts, periodically resetting the schedule to escape local minima.

# Warmup + cosine decay LR schedule: ramp linearly for warmup_steps,
# then decay following a cosine curve down to min_lr.
import numpy as np
import matplotlib.pyplot as plt
def lr_schedule(step, total_steps, warmup_steps, peak_lr, min_lr):
    """Warmup + cosine decay schedule."""
    if step < warmup_steps:
        # Linear warmup
        return peak_lr * step / warmup_steps
    else:
        # Cosine decay
        progress = (step - warmup_steps) / (total_steps - warmup_steps)
        return min_lr + 0.5 * (peak_lr - min_lr) * (1 + np.cos(np.pi * progress))
        # Typical LLM training schedule
        total_steps = 100000
        warmup_steps = 2000
        peak_lr = 3e-4
        min_lr = 3e-5
        steps = np.arange(total_steps)
        lrs = [lr_schedule(s, total_steps, warmup_steps, peak_lr, min_lr) for s in steps]
        print(f"Step 0:      LR = {lrs[0]:.2e}")
        print(f"Step 1000:   LR = {lrs[1000]:.2e}")
        print(f"Step 2000:   LR = {lrs[2000]:.2e} (peak)")
        print(f"Step 50000:  LR = {lrs[50000]:.2e}")
        print(f"Step 99999:  LR = {lrs[99999]:.2e}")

PyTorch ships the same schedule as a one-liner via CosineAnnealingLR, which most production training loops use directly rather than rolling the formula by hand:

import torch
from torch.optim.lr_scheduler import CosineAnnealingLR

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scheduler = CosineAnnealingLR(optimizer, T_max=100_000, eta_min=3e-5)

for step in range(100_000):
    train_step(...)
    optimizer.step()
    scheduler.step()         # advances LR along the cosine curve
    optimizer.zero_grad()

# Plug a Hugging Face cosine-with-warmup scheduler into any PyTorch optimizer.
import torch
from transformers import get_cosine_schedule_with_warmup

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=2000,
    num_training_steps=100_000,
)
for step in range(num_training_steps):
    train_step(...)
    optimizer.step()
    scheduler.step()  # advances LR according to the cosine curve.
    optimizer.zero_grad()

The Noam Schedule (Historical Reference)

Before cosine and WSD became default, the Noam schedule, named after Noam Shazeer and introduced alongside the original Transformer in Vaswani et al. (2017, "Attention Is All You Need"), set the historical convention: linear warmup followed by inverse-square-root decay. The learning rate at training step $t$ is

For $t \le T_{\text{warm}}$ the right argument of the $\min$ dominates, giving a linear warmup. For $t > T_{\text{warm}}$ the left argument takes over, decaying the learning rate as $t^{-1/2}$. The peak learning rate is reached exactly at $t = T_{\text{warm}}$ and is scaled by $d_{\text{model}}^{-0.5}$, so wider models automatically use smaller learning rates. The schedule has no notion of a total step budget, which made it convenient when total training horizons were uncertain. Most modern large-scale runs prefer cosine or WSD because the slow square-root decay leaves the model at a non-trivial learning rate for too long, but Noam is still the historical reference and remains a useful baseline for translation-scale transformer experiments.

# Noam learning-rate schedule (Vaswani et al., 2017, "Attention Is All You Need").
# Linear warmup for T_warm steps, then inverse-square-root decay.
import torch

class NoamScheduler(torch.optim.lr_scheduler._LRScheduler):
    def __init__(self, optimizer, d_model: int, warmup_steps: int = 4000, factor: float = 1.0):
        self.d_model = d_model
        self.warmup = warmup_steps
        self.factor = factor
        super().__init__(optimizer)

    def get_lr(self):
        step = max(self._step_count, 1)
        scale = self.factor * (self.d_model ** -0.5) * \
                min(step ** -0.5, step * self.warmup ** -1.5)
        return [scale for _ in self.base_lrs]

# Usage with the canonical Transformer config (d_model=512, warmup=4000):
model = ...  # your Transformer
optimizer = torch.optim.Adam(model.parameters(), lr=1, betas=(0.9, 0.98), eps=1e-9)
scheduler = NoamScheduler(optimizer, d_model=512, warmup_steps=4000)

WSD: Warmup-Stable-Decay

A newer schedule, increasingly preferred for large-scale training, is the Warmup-Stable-Decay (WSD) schedule (also called trapezoidal). Instead of continuously decaying the learning rate after warmup, WSD maintains a constant "stable" phase at the peak learning rate for most of training, then applies a rapid linear or cosine decay only during a short final phase (typically the last 10-20% of steps).

WSD has a crucial practical advantage: because the learning rate stays constant during the stable phase, you can evaluate checkpoints and decide later when to begin the decay phase. With cosine decay, the total training budget must be known in advance. WSD decouples the schedule from the total step count, making it ideal for continued pretraining and for models where the final training duration is not known at the start. Llama-3, DeepSeek V3, and Qwen 2.5 all used WSD schedules (see Section 7.3 for details on these models).

# Warmup-Stable-Decay (WSD) schedule: warmup, hold at peak_lr for a
# stable fraction of training, then decay. Used in recent large runs.
import numpy as np
def wsd_schedule(step, total_steps, warmup_steps, stable_fraction, peak_lr, min_lr):
    """Warmup-Stable-Decay learning rate schedule."""
    decay_start = int(total_steps * stable_fraction)
    if step < warmup_steps:
        return peak_lr * step / warmup_steps
    elif step < decay_start:
        return peak_lr # Constant during stable phase
    else:
        # Linear decay in the final phase
        decay_progress = (step - decay_start) / (total_steps - decay_start)
        return peak_lr + (min_lr - peak_lr) * decay_progress
        total_steps = 100000
        warmup_steps = 2000
        peak_lr, min_lr = 3e-4, 3e-5
        # Stable phase covers 80% of training; decay covers last 20%
        for s in [0, 2000, 40000, 79999, 80000, 90000, 99999]:
            lr = wsd_schedule(s, total_steps, warmup_steps, 0.8, peak_lr, min_lr)
            print(f"Step {s:>6}: LR = {lr:.2e}")

6.5.6 Gradient Accumulation

Large batch sizes improve training stability and efficiency but require more GPU memory. Gradient accumulation simulates large batches without the memory cost: instead of processing the full batch at once, the training loop processes several smaller micro-batches, accumulating their gradients, and only performing the optimizer step after all micro-batches are complete.

import torch
# Gradient accumulation pseudocode
accumulation_steps = 8 # Effective batch = micro_batch * 8
# Reset gradients from previous step
optimizer.zero_grad()
for i, micro_batch in enumerate(dataloader):
    loss = model(micro_batch) / accumulation_steps # Scale loss
    # Compute gradients via backpropagation
    loss.backward() # Accumulate gradients
    if (i + 1) % accumulation_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        # Update weights using computed gradients
        optimizer.step()
        optimizer.zero_grad()

6.5.7 Training Dynamics: The Loss Landscape

The loss landscape of a transformer is a high-dimensional surface with complex geometry. Research by Li et al. (2018) showed that neural network loss landscapes contain wide, flat minima and narrow, sharp minima. Models that converge to flatter minima tend to generalize better because small perturbations to the parameters (as occur with different test inputs) cause smaller changes in loss.

The Grokking Phenomenon

Power et al. (2022) discovered a surprising training dynamic called "grokking": a model can memorize the training set early in training (achieving near-perfect training accuracy) while showing no generalization, and then, after many additional training steps, suddenly transition to perfect generalization. This delayed generalization can occur thousands of steps after the training loss has plateaued.

Grokking challenges the conventional wisdom that training should be stopped when the validation loss plateaus. The phenomenon has been explained through the lens of representation learning: the model first memorizes using inefficient representations, then gradually discovers the underlying algorithm, which requires more training steps to crystallize. Weight decay plays a critical role in enabling grokking by continuously pushing the model away from memorization-based solutions toward simpler, more generalizable representations.

Figure 6.5.4: In grokking, the model achieves perfect training accuracy early but takes many more steps to generalize, with a long validation plateau between memorization and true learning.

6.5.8 Training Instabilities

Large-scale training runs frequently encounter instabilities that can derail training entirely if not handled.

Loss Spikes

Sudden jumps in training loss, sometimes by several orders of magnitude, are a common occurrence in LLM training. They are typically caused by outlier batches with unusual gradient distributions, numerical overflow in attention computations, or learning rate being too high for the current loss landscape curvature. Loss spikes can often be recovered from (the loss returns to its pre-spike trajectory), but severe spikes may corrupt optimizer states and require rolling back to a checkpoint.

Gradient Clipping

The universal mitigation for gradient instability is gradient clipping: scaling the gradient vector so its global norm does not exceed a threshold (typically 1.0). This prevents any single batch from causing a catastrophically large parameter update.

z-Loss Regularization

PaLM introduced z-loss, an auxiliary loss term that penalizes large logits in the output layer: $L_{z} = \alpha \cdot log^{2}(Z)$, where $Z$ is the sum of exponentials in the Transformer architecture denominator. This prevents attention entropy collapse and reduces the frequency of loss spikes.

Exercises

import math
WARM, TOTAL, MIN = 2000, 100_000, 0.1
def lr_lambda(step):
  if step < WARM: return step / WARM
  p = (step - WARM) / (TOTAL - WARM)
  return MIN + (1 - MIN) * 0.5 * (1 + math.cos(math.pi * p))