← Back to Table of Contents

Chapter 11 — Optimizers & Loss Functions for LLMs/VLMs

“The loss function defines what you want; the optimizer determines how you get there. Both choices profoundly shape model behaviour.”

Loss Functions

Cross-Entropy Loss — The Foundation

Every LLM is trained to minimise cross-entropy between the predicted token distribution and the true next token. This is also called negative log-likelihood (NLL):

where N is the number of labelled tokens (non-masked positions), and p_θ is the model’s predicted probability for the true token.

In code:

1
2
3
4
5
6
7
8
9
10
import torch.nn.functional as F

# logits: [B, T, V]  — model output (unnormalised)
# labels: [B, T]    — target token IDs (-100 = ignore)
loss = F.cross_entropy(
    logits.view(-1, vocab_size),   # [B*T, V]
    labels.view(-1),               # [B*T]
    ignore_index=-100,
    reduction="mean",
)

Tensor shapes:

Perplexity

Perplexity is the exponentiated cross-entropy loss — a human-interpretable measure of how “surprised” the model is:

Lower perplexity = better model. Typical values: GPT-2 ~30, Llama-3-8B ~7–9 on Wikitext-103.

Z-Loss (Auxiliary Stability Loss)

A regularisation term added during training to prevent the softmax logits from growing very large (logit explosion), which causes instability:

where z are the pre-softmax logits and α is a small coefficient (e.g., 1e-4). Used in PaLM, Gemini, and others.

Load-Balancing Loss (MoE)

For Mixture-of-Experts models, an auxiliary loss ensures experts are used roughly equally (preventing expert collapse):

where f_e is the fraction of tokens routed to expert e, and P_e is the average routing probability to e. This is added to the main cross-entropy loss. See Chapter 33 for full MoE coverage.

Reward Model Loss

Reward models (used in RLHF) are trained with a ranking loss — the chosen response should score higher than the rejected one:

where y_w is the winning (chosen) response, y_l is the losing (rejected) response, and r_θ is the scalar reward. This is the Bradley-Terry preference model.

DPO Loss

Direct Preference Optimisation (DPO) reformulates preference learning as a supervised loss:

No reward model is needed — the reference policy π_ref (frozen SFT model) implicitly regularises the policy. β controls how far the policy strays from the reference.

VLM-Specific Losses

Vision-Language Models combine image and text objectives:

Optimizers

SGD with Momentum

The classical optimizer — rarely used for LLM training in practice:

Memory: 1 momentum buffer per parameter → 1× params extra.

Adam — Adaptive Moment Estimation

Adam (Kingma & Ba, 2014) is the foundation of most LLM optimisers. It maintains running estimates of the first and second moments of gradients:

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

Default hyperparameters: β₁=0.9, β₂=0.95 (LLMs; note: 0.999 is the original paper default but 0.95 is standard for LLMs), ε=1e-8, η=1e-4 to 3e-4.

Memory cost: 2 optimizer states (m and v) per parameter, both in FP32:

This is why distributed training (FSDP/ZeRO) is essential — see Chapter 26.

AdamW — Adam with Weight Decay

The standard choice for LLM training. AdamW (Loshchilov & Hutter, 2017) decouples weight decay from the gradient update:

Adam (incorrect weight decay): θ_t = θ_{t-1} - η (m̂/√v̂ + λθ_{t-1})
AdamW (correct weight decay): θ_t = (1 - ηλ)θ_{t-1} - η · m̂/√v̂

The difference: in AdamW, weight decay is applied directly to weights, not scaled by the adaptive learning rate. This gives more consistent regularisation.

Typical settings for LLM training:

• lr = 1e-4 to 3e-4 (pre-training); lr = 2e-5 to 2e-4 (fine-tuning)
• weight_decay = 0.1
• β₁ = 0.9, β₂ = 0.95
• gradient_clip = 1.0

Adam-mini

A recent (2024) memory-efficient variant that reduces optimizer memory by using one learning rate per parameter group (e.g., per attention head) instead of per parameter:

• Memory: ~45% of Adam memory
• Performance: matches Adam/AdamW in practice
• Particularly useful for large models where optimizer states dominate memory

Adafactor

Designed for extreme memory efficiency — used to train early T5 models:

• Does not store full second moment vector; factorises it into rank-1 row/column factors
• Memory: ~O(√params) instead of O(params) for second moment
• Trade-off: can be less stable; often needs careful tuning
• Used in: T5, mT5, Switch Transformer

SOAP

SOAP (Shampoo with Adam in the Preconditioner, 2024) applies second-order information via matrix preconditioning:

• Precomputes Kronecker-factored curvature (similar to K-FAC/Shampoo)
• Memory overhead: similar to Adam
• Speed: can reach the same loss in ~40% fewer steps than AdamW
• Increasingly used for smaller training runs where step count is the bottleneck

Muon (Momentum + Orthogonalisation)

A newer optimiser (2024) that applies Nesterov momentum and then orthogonalises the gradient update:

• Applies to weight matrices only (not embeddings/biases)
• Reaches lower loss than AdamW at same compute in recent experiments
• Used in MicroGrad and some frontier training runs (reportedly used for DeepSeek V3’s internal layers)

Learning Rate Schedules

The learning rate schedule is as important as the optimiser choice for LLM training.

Cosine Decay with Warmup

The de facto standard for pre-training:

Typical warmup: 1000–2000 steps. Starting cold (no warmup) causes early instability.

WSD — Warmup, Stable, Decay

Used in MiniCPM and other recent models. Enables continual training by decoupling the decay phase:

1. Warmup: linear ramp from 0 to η_max
2. Stable: constant η_max for the bulk of training
3. Decay: cosine or linear decay to η_min

Key advantage: you can extend training (add more data) by extending the stable phase without restarting the decay. This enables the “annealing” trick — saving a checkpoint, then annealing on high-quality data.

Trapezoidal / Linear Decay

Used in models like OLMo 2. Simpler than cosine, easier to reason about:

• Warmup → flat → linear decay to 0

Constant LR (Fine-Tuning)

For SFT/LoRA fine-tuning over few epochs, a constant LR with short warmup often works well:

• lr = 2e-4 for LoRA, lr = 2e-5 for full fine-tune
• 5–10% of total steps for warmup

Gradient Clipping

A near-universal component of LLM training. Clips the global gradient norm to prevent instability from gradient spikes:

Standard value: max_norm = 1.0. Training loss spikes often coincide with large gradient norm events; clipping at 1.0 prevents most catastrophic updates.

1
2
# In PyTorch
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Optimizer Comparison Summary

Hyperparameter Sensitivity

What’s Next

With the training machinery understood — data, objectives, and optimisers — we’ll look at mid-training strategies that extend and refine pre-trained models.

← Previous: Chapter 10 — Pre-Training at Scale · Next: Chapter 12 — Mid-Training & Continued Pre-Training →

Last updated: April 2026