DL S006 · Optimization Intuition + Numerical Stability

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 6 / 80 · Module M01 · ~2 hours

The story

We have (a) a way to compute how a loss depends on every weight (S004: gradients) and (b) a loss worth minimising (S005: cross-entropy). What remains is the actual algorithm for turning those gradients into better weights, step after step, for millions of steps, without the whole thing exploding, oscillating, or dying quietly in a corner.

That algorithm is called stochastic gradient descent, and its variants — momentum, RMSprop, Adam, AdamW — are the workhorses of every trained network on earth. The math is embarrassingly simple: "subtract the gradient times a small number." What's hard is the intuition for why one variant beats another on a given problem, and the numerical hygiene needed to make any of them stable in half precision on a modern GPU.

Here's the puzzle we'll walk through: SGD is provably convergent under mild conditions, but plain SGD trains modern LLMs terribly. Adam is not provably convergent (in fact its convergence proof was famously broken for years) and yet Adam trains almost every LLM you've ever heard of. Why? By the end of today you'll understand — and you'll have an intuition strong enough to pick your optimizer without reading a paper.

Prerequisites

• Session 004 — you need gradients to descend by.
• Session 005 — for the log-softmax and cross-entropy stability discussion.

1 · Gradient descent — the whole idea in one line

Given a loss L(θ) and its gradient g = ∇L(θ) at the current point, take a small step in the direction of steepest descent:

η is the learning rate — a small positive number, typically 1e-3 to 1e-1 depending on architecture. That's it. That's gradient descent.

What the learning rate physically means: if the loss surface near θ locally looks like a plane, one step reduces the loss by about η * ||g||². Bigger η → faster progress per step, but too big and you overshoot the minimum and diverge. Everyone's first hyperparameter search is learning rate. Everyone's.

1.1 A tiny 1D example

Minimise L(w) = (w - 3)². Gradient: 2 * (w - 3). Start at w = 0, η = 0.1.

Zooms toward 3 exponentially fast. Try η = 1.0:

Try η = 1.1:

These three regimes — converge / oscillate / diverge — happen at every scale, from toy 1D problems to 70B-parameter LLMs. The learning rate you'd pick to avoid divergence is bounded above by roughly 2/L_max, where L_max is the largest eigenvalue of the loss Hessian. You will never compute this. You will just try 1e-3 and see what happens.

2 · Stochastic vs batch vs mini-batch

Full-batch gradient descent computes g as the gradient over the ENTIRE dataset per step. That's the "true" gradient. It's slow and RAM-hungry.

Stochastic gradient descent computes g on ONE example per step. It's noisy but ~N times cheaper per step, and the noise actually helps escape saddle points and shallow minima.

Mini-batch SGD, what everyone actually uses, computes g over a small random batch (typically 32–512 examples). Gets you 95% of the noise benefit with dramatically better GPU utilization.

for epoch in range(n_epochs):
    for batch in dataloader:                    # mini-batch
        loss = model(batch).loss
        grads = compute_grads(loss)
        for p, g in zip(params, grads):
            p -= lr * g

That's a training loop. Sessions 011 and 017 will flesh it out; the essence is above.

2.1 Why noise is a feature, not a bug

Full-batch GD, on a non-convex loss, gets stuck in the nearest local minimum. Mini-batch SGD's noise sometimes pushes it up the hill and out of a bad basin, into a wider, flatter, better-generalising one. There's a decade of theory on "flat minima" and generalization; the TL;DR is: noise regularises. It's part of why 32-batch training often beats 4096-batch training on the same problem, even after adjusting the learning rate.

3 · Momentum — the ball that remembers

Vanilla SGD has a specific failure mode: on loss surfaces shaped like long, narrow valleys (which are basically every real DL loss surface), it oscillates across the walls of the valley rather than sliding down its length. Progress is slow.

Momentum adds a "velocity" state that accumulates past gradients:

Typical β = 0.9. The velocity v is an exponential moving average of gradients. In directions where the gradient consistently points one way (down the valley), velocity builds up and progress accelerates. In directions where the gradient oscillates (across the valley), successive gradients cancel and velocity stays small.

3.1 Code

class SGDMomentum:
    def __init__(self, params, lr=1e-2, beta=0.9):
        self.params = params
        self.lr = lr; self.beta = beta
        self.v = [np.zeros_like(p) for p in params]
 
    def step(self, grads):
        for i, (p, g) in enumerate(zip(self.params, grads)):
            self.v[i] = self.beta * self.v[i] + g
            p -= self.lr * self.v[i]

For convolutional networks (S025–S030), SGD + momentum + a well-tuned schedule (S023) is still a very strong baseline. For transformers (S040+), Adam almost always wins.

4 · Adam — the default optimizer for everything

Adam combines two ideas: (1) momentum (like §3) and (2) per-parameter learning rate — each weight gets its own effective learning rate based on how large its gradients have been recently.

The full algorithm, in six lines, is:

Typical values: β₁ = 0.9, β₂ = 0.999, ε = 1e-8, η = 3e-4 (the "Karpathy constant" for transformers). t is the step number.

4.1 What each piece does

• m (1st moment) — momentum, same idea as §3.
• v (2nd moment) — an EMA of squared gradients. Big for weights that receive big gradients; small for weights that don't.
• √v̂ + ε in the denominator — divides the momentum-corrected step by the recent gradient magnitude. So a weight whose gradient has been 100× another weight's gets 100× smaller updates. Every weight moves at roughly the same effective scale.
• ε — prevents division by zero when v̂ is nearly zero for a fresh parameter.
• Bias correction (1 - βᵗ denominators) — early in training m and v are near zero because they started at zero; dividing by (1 - βᵗ) cancels the bias.

4.2 Code

class Adam:
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.999), eps=1e-8):
        self.params = params
        self.lr = lr; self.b1, self.b2 = betas; self.eps = eps
        self.m = [np.zeros_like(p) for p in params]
        self.v = [np.zeros_like(p) for p in params]
        self.t = 0
 
    def step(self, grads):
        self.t += 1
        for i, (p, g) in enumerate(zip(self.params, grads)):
            self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g
            self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * g * g
            m_hat = self.m[i] / (1 - self.b1 ** self.t)
            v_hat = self.v[i] / (1 - self.b2 ** self.t)
            p -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

Twenty lines. That's the entire optimizer that trained GPT-3, LLaMA, Mistral, and every transformer since 2019.

4.3 AdamW — the small tweak that matters for LLMs

Original Adam applies weight decay by adding λ*θ to the gradient before the moment updates. That interacts poorly with the per-parameter scaling — weights with tiny v get decayed less than they should. AdamW decouples: apply weight decay directly to θ outside the Adam step:

p -= self.lr * (m_hat / (np.sqrt(v_hat) + eps) + weight_decay * p)

Every modern LLM uses AdamW, not Adam. The difference matters at scale. Session 044 (training a small LLM end-to-end) will bake this in.

5 · SGD vs Adam — when to pick which?

The theoretical reason Adam works so well on transformers: attention weights and embedding weights have wildly different gradient magnitudes, and Adam's per-parameter scaling normalises them. On a plain CNN, all weights have similar gradient scales, so SGD is fine.

6 · The stability tricks — log-sum-exp, log-softmax, and friends

We hinted at this in S005. Now the full story.

6.1 The problem — softmax overflows

z = np.array([1000.0, 999.0, 998.0])
np.exp(z)         # [inf, inf, inf]
softmax(z)        # [nan, nan, nan]

e^1000 is way past float64's max. Softmax dies. Loss becomes NaN. Training over.

6.2 The shift trick

Softmax is invariant to adding a constant to all logits. So we subtract the max:

def softmax_stable(z):
    z = z - z.max(axis=-1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)

Now the largest entry of z - z.max() is 0; e^0 = 1. No overflow. Every framework does this. Every custom kernel MUST do this.

6.3 log-sum-exp — the same idea for log(sum(exp(x)))

Cross-entropy loss is -z_y + log Σ exp(z_j). That log Σ exp(z_j) term needs the same protection:

def logsumexp(z, axis=-1):
    m = z.max(axis=axis, keepdims=True)
    return m.squeeze(axis) + np.log(np.exp(z - m).sum(axis=axis))

Combining: log(softmax(z))_y = z_y - logsumexp(z). No exp-then-log dance. This is what F.log_softmax computes internally in PyTorch.

6.4 Never compute log(F.softmax(x)) — use F.log_softmax(x)

Same math, but the fused version is much more numerically stable. The naive version can produce log(0) = -inf when softmax rounds down to zero. F.log_softmax uses the shift trick throughout so it never does. If you catch a tutorial doing log(softmax(...)), mentally correct it.

7 · fp16 and the exploding activations

Modern GPUs are 2–4× faster in fp16 (or bf16) than fp32. But fp16 has a much smaller dynamic range:

fp32:  ~1e-38  to  ~3e38     ~7 decimal digits precision
fp16:  ~6e-5   to  ~6e4       ~3 decimal digits precision
bf16:  same range as fp32,    ~2 decimal digits precision

fp16 gradient values below ~1e-8 round to zero. That's a lot of very small but real gradients silently vanishing.

7.1 Loss scaling — the fp16 workaround

Multiply the loss by a big constant S (typically 2^16) before .backward(). Gradients get scaled up by S, exit the "zeroed" regime, propagate correctly, then get UNscaled before the optimizer step:

scaled_loss = S * loss
scaled_loss.backward()
for p in params:
    p.grad /= S
optimizer.step()

Modern PyTorch's torch.cuda.amp.GradScaler does this for you, including dynamic scaling that ramps S up or down based on whether gradients overflowed. Session 018 covers AMP in production detail.

7.2 bf16 — the fp16 successor

bf16 has the same exponent range as fp32 but fewer mantissa bits (less precision). Practically: bf16 has no under/overflow problems, at the cost of some precision. LLMs train fine in bf16 without any loss scaling. Every modern LLM training run uses bf16 (on A100/H100 hardware). Session 018 goes into this.

8 · Diagnosing "loss went to NaN at step X"

This will happen to you. Here's the checklist, in the order you should try.

Gradient clipping (Session 024) is the seat-belt for all of these:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Rescales all gradients so that their global L2 norm is ≤ 1. Prevents any single big-gradient step from blowing weights to infinity. Every LLM training script has this.

8b · The 2024–2026 optimizer landscape — what's replacing Adam

Adam has been the default optimizer for transformers since 2017. But 2023–2026 is the first period since then where serious contenders are landing on real training runs. Here's the field.

8b.1 · AdamW — the actual default (Loshchilov & Hutter, 2019)

Before we talk about replacements, know that vanilla Adam is almost never used in production. AdamW ("Decoupled Weight Decay Regularization") applies weight decay as $\theta \leftarrow \theta - \eta \lambda \theta$ separately from the gradient-based Adam update, instead of folding decay into the gradient. On transformers this consistently gives 1–3 perplexity points better than Adam-with-L2. Every 2024–2025 LLM training run — Llama, Mistral, Qwen, DeepSeek — uses AdamW.

8b.2 · Lion (Chen et al., 2023) — sign momentum

Symbolic Discovery of Optimization Algorithms used program search to find a new optimizer, and Lion ("EvoLved Sign Momentum") popped out:

$m \leftarrow \beta_1 m + (1 - \beta_1) g, \quad \theta \leftarrow \theta - \eta \, \text{sign}(\beta_2 m + (1 - \beta_2) g) - \eta \lambda \theta$

Just the sign of a momentum-blended gradient. Uses ~50% less memory than Adam (only one moment buffer instead of two). Google reports 1.5–2× wall-clock speedups on ViT and language models. Adoption in 2024–2025 is real but not universal — Lion needs a smaller learning rate (~10× smaller than AdamW) and can be finicky.

8b.3 · Sophia (Liu et al., 2023) — diagonal Hessian

Sophia approximates the diagonal of the Hessian using Hutchinson's estimator and preconditions the gradient by it, then clips. Reports ~2× fewer training steps to reach GPT-2 quality. Some open-source repos have adopted it; the big labs are more cautious.

8b.4 · Shampoo and SOAP (2024)

Shampoo tracks full second-moment matrices per layer (Kronecker-factored), giving a real second-order-ish preconditioner. Its 2024 refinement SOAP (Vyas et al.) rotates the Adam step into Shampoo's eigenbasis and gets 30–50% wall-time improvements on GPT-2-scale runs. Anthropic's public-facing engineers have hinted at second-order optimizers in training pipelines; DeepMind has used Shampoo in production.

8b.5 · Muon (Jordan et al., 2024) — the hidden-layer surprise

Muon applies Newton-Schulz iteration to orthogonalize the momentum matrix before stepping, only for hidden 2D weight matrices (embedding/head still use AdamW). Set NanoGPT speed records in 2024 (nanogpt-speedrun leaderboard) and is reported to be part of Kimi K2 / Moonshot AI's training pipeline. As of 2025 this is the most-buzzed alternative to AdamW for LLM pretraining.

8b.6 · Learning-rate schedules — the other half of the story

Optimizer choice matters, but so does the learning-rate schedule. 2024–2025 defaults:

• Linear warmup + cosine decay (Llama, most academic) — 500–2000 warmup steps, then cosine from lr_max to lr_max * 0.1.
• WSD (Warmup-Stable-Decay) (Hu et al., 2024, MiniCPM) — warmup, then flat, then rapid decay only at the end. Lets you checkpoint mid-flat and "annealing-finetune" for downstream tasks. Now standard in Chinese-lab pretraining.
• Muennighoff "overtrained" scaling ("Scaling Data-Constrained Language Models", 2023) — push past Chinchilla-optimal token counts ("train longer, model gets better even past compute-optimal"). Standard practice in 2024 for inference-first models like Llama 3.1.

8b.7 · Gradient clipping and loss scaling (still critical in 2025)

Every production training run does:

• Global-norm gradient clipping: if ||g|| > c: g ← g * c / ||g||. Standard c = 1.0. Prevents rare large-gradient batches from destroying weeks of training.
• Loss scaling for FP16: multiply loss by ~2^15 before backward, divide gradients back after. FP16 mantissa has too few bits for small gradients; scaling shifts them into the representable range. In BF16 (used by nearly all 2024–2025 pretraining), this is unnecessary — BF16 has the exponent range of FP32.
• Skip-nan batch: if any gradient is NaN, skip the step entirely and log it. Essential for robust multi-week runs.

Further reading:

• Loshchilov & Hutter. "Decoupled Weight Decay Regularization (AdamW)" (2017/2019).
• Chen et al. "Symbolic Discovery of Optimization Algorithms (Lion)" (2023).
• Vyas et al. "SOAP: Improving and Stabilizing Shampoo" (2024).
• Jordan. "Muon: An Optimizer for Hidden Layers in Neural Networks" (2024).
• Hu et al. "MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies" (2024, WSD schedule).

9 · Diagram — the optimizer step

This loop is the same in every deep learning framework, in every language, forever. Everything else — data loaders, distributed training, mixed precision, activation checkpointing, KV caches — is engineering to make this loop run faster or on bigger models. The loop itself is what you just learned.

10 · War stories

11 · Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: Redo the SGD vs Adam mental simulation on a 2D quadratic tomorrow. Revisit §8b (the 2024–2026 optimizer landscape) on day 7.

Next session (S007): Module M01 is done. M02 opens with a single artificial neuron — forward pass, loss, gradient, one round of training — all in ~20 lines of NumPy. The math from M01 stops being abstract.

Sticky note: z - z.max() before every exp; global grad-norm clip at 1.0; BF16 if you have it, else FP16 + loss scaling.

Recall — from memory

Stretch — for one extra hour

Implement all three optimizers (SGD, SGD+momentum, Adam) as classes with the same interface, and train each on the same tiny MLP on a random 2D classification problem. Plot the loss curves for all three; verify Adam converges fastest and SGD+momentum converges deepest (given long enough). This is the classic "optimizer visualisation" plot; it will lock in your intuition.

In your own words

Explain in one sentence why Adam usually beats vanilla SGD on transformers:

Spaced-review pointer

• From S004 — gradients are the input to every optimizer here. Adam's g² in particular ties directly to gradient variance.
• From S005 — log-sum-exp is the same trick you first met when computing cross-entropy stably.

Next-session teaser

Module M01 is done. You have the mathematical foundation: broadcasting (S003), chain rule (S004), cross-entropy (S005), and gradient descent (S006). Session 007 opens Module M02 (Neural Nets from Scratch), where we finally connect these pieces to build a single artificial neuron — forward pass, loss, gradient, and one round of training on a 2D toy problem, all in ~20 lines of numpy. The math from M01 stops being abstract there.

What to bring back tomorrow — sticky note

• Diagram: the forward → grad → optimizer.step loop from §9.
• Equation: Adam's 6-line update; especially θ ← θ - η*m̂/(√v̂ + ε).
• Snippet: z - z.max() before every exp. It saves so many training runs.

Previous: ← DL S005 · Probability and Information · Next: DL S007 · A Single Artificial Neuron →