Optimization & Regularization

The model architecture decides what a network can represent. Optimization decides whether you can actually find good weights, and regularization decides whether those weights generalize past your training set. This page is the practical training toolbox — the knobs every deep-learning practitioner turns: the optimizer (SGD → momentum → Adam/AdamW), the learning-rate schedule (warmup + cosine decay), and the regularizers and stabilizers (dropout, batch/layer norm, weight init, gradient clipping).

What it is — intuition first

Imagine you're standing somewhere on a vast, foggy, mountainous landscape in the dark, and your only goal is to walk downhill to the lowest valley. You can't see the whole map. All you can feel is the slope under your feet right now. That slope is the gradient: the direction in which the ground rises fastest. To go down, you step the opposite way.

That is literally all gradient descent is:

new_weights = old_weights - learning_rate * gradient

• The gradient (∂Loss/∂weights) tells you which way is uphill (more loss).
• The learning rate is your step size — how big a stride you take.
• You subtract because you want to go down the loss.

Everything in this page is a refinement of that one line. Real landscapes are nasty: they have ravines (steep in one direction, flat in another), saddle points, plateaus, and noise (because you only feel the slope from a small random mini-batch of data, not the whole dataset). The "toolbox" exists because plain w -= lr * grad is too dumb to handle real terrain:

• Momentum stops you zig-zagging across narrow ravines.
• Adam gives each weight its own step size, so flat directions get big steps and steep directions get small ones.
• Learning-rate schedules make you take careful baby-steps at the start and the end, big strides in the middle.
• Regularization (dropout, weight decay, norms) stops you from memorizing the exact bumps of this particular foggy hill, so the route still works on a hill you've never seen.

The gradient itself comes from backpropagation (the chain rule, built from scratch in [[micrograd]]). This page assumes you have the gradient and asks: given the gradient, how do I update the weights well? If you're shaky on partial derivatives, the chain rule, or vectors/matrices, read [[math-for-ml-foundations]] first — it's the prerequisite. To watch all of this run end-to-end in a tiny real network, see [[neural-network-from-scratch]] and [[nn-zero-to-hero]].

Why it matters

A correct architecture with a bad optimization recipe trains slowly, unstably, or not at all. In practice the difference between "loss diverges to NaN at step 200" and "clean convergence to a good model" is almost never the architecture — it's the optimizer, the learning rate, the warmup, and the clipping. Three concrete reasons it matters:

1. It's where training actually breaks. Loss spikes, NaNs, dead neurons, and "my loss is stuck at exactly ln(vocab_size)" are all optimization/initialization problems. Knowing this toolbox is how you debug them.
2. It dominates the compute bill. A good schedule can hit the same loss in a fraction of the steps. Batch normalization's original claim was a 14× reduction in training steps to reach the same accuracy. At LLM scale, the optimizer recipe is part of the [[scaling-laws]] calculus.
3. It's reused everywhere. The exact same AdamW + warmup + cosine recipe trains a 2-layer MLP, [[nanogpt]], a 70B-parameter LLM, and the policy in [[grpo]]/[[rlhf-and-alignment]]. Learn it once; it pays off at every scale.

How it works — real mechanics

0. The base: SGD (stochastic gradient descent)

"Stochastic" just means we estimate the gradient from a random mini-batch, not the full dataset (which would be too expensive). This makes each step noisy — which is partly a curse (jitter) and partly a blessing (the noise helps escape sharp bad minima).

# one SGD step
g = grad(loss, w)          # gradient from current mini-batch
w = w - lr * g

Problems with raw SGD: it crawls through flat regions, overshoots steep ones, and zig-zags in "ravine" geometries (steep walls, gently-sloping floor — extremely common in deep nets).

1. Momentum — give the ball mass

Instead of stepping in the current gradient direction, keep a running velocity that accumulates past gradients. Like a heavy ball rolling downhill, it builds speed in consistent directions and damps the back-and-forth.

v = beta * v + g           # beta ~ 0.9 ; v = exponential moving avg of gradients
w = w - lr * v

Across a ravine the gradient flips sign every step, so those contributions cancel in v; along the valley floor the gradient is consistent, so they accumulate. Result: less zig-zag, faster progress. (Nesterov momentum is a slightly smarter variant that "looks ahead" before computing the gradient.)

2. Per-parameter adaptive rates — RMSProp → Adam

Different weights need different step sizes. A weight that consistently sees large gradients should take smaller steps (it's on a steep wall); one with tiny gradients should take bigger steps (it's on a plateau). RMSProp tracks a running average of squared gradients v and divides the step by √v, normalizing every parameter to roughly unit scale.

Adam (Kingma & Ba, 2014 — arXiv:1412.6980↗) is the workhorse: it combines momentum (1st moment m) with RMSProp's adaptive scaling (2nd moment v), plus a bias-correction trick. Here is the full update — this is the algorithm running inside almost every modern training loop:

# Adam — defaults: lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8
m = 0; v = 0; t = 0
for each step:
    t += 1
    g = grad(loss, w)
    m = beta1 * m + (1 - beta1) * g          # 1st moment: EMA of gradient  (momentum)
    v = beta2 * v + (1 - beta2) * (g * g)    # 2nd moment: EMA of squared gradient (scale)
    m_hat = m / (1 - beta1**t)               # bias correction (m,v start at 0 -> biased low early)
    v_hat = v / (1 - beta2**t)
    w = w - lr * m_hat / (sqrt(v_hat) + eps) # adaptive, momentum-smoothed step

Why the bias correction? m and v start at zero, so in the first few steps they're biased toward zero — dividing by (1 - beta^t) rescales them up to an unbiased estimate. (This early-step weirdness is exactly why warmup exists — see §4.)

Intuition for the final line: m_hat is the direction (momentum-smoothed), and dividing by √v_hat makes the magnitude roughly 1 per parameter regardless of gradient scale. Adam is "robust to the choice of learning rate" precisely because of this normalization — which is why lr=3e-4 ("the Karpathy constant") works as a default across wildly different models.

3. AdamW — fixing weight decay (the one you should actually use)

Weight decay is regularization: gently shrink every weight toward zero each step (w -= lr * wd * w), which prefers smaller, smoother weights that generalize better. For plain SGD, adding an L2 penalty ½·wd·‖w‖² to the loss is mathematically identical to weight decay.

Loshchilov & Hutter (2017 — [arXiv:1711.05101](https://arxiv.org/abs/1711.05101)) showed that for Adam this equivalence breaks. If you implement weight decay as an L2 term in the loss, that penalty's gradient gets divided by √v_hat along with everything else — so weights with large gradients get less decay, which is backwards. The fix is decoupled weight decay: apply the shrink directly to the weights, outside the adaptive machinery:

# AdamW — decoupled weight decay (wd applied to weights, NOT through the loss/grad)
m = beta1 * m + (1 - beta1) * g
v = beta2 * v + (1 - beta2) * (g * g)
m_hat = m / (1 - beta1**t)
v_hat = v / (1 - beta2**t)
w = w - lr * (m_hat / (sqrt(v_hat) + eps) + wd * w)   # <-- wd*w added OUTSIDE the sqrt(v) scaling

Practical consequences: (a) weight_decay becomes independent of the learning rate, so you can tune them separately; (b) Adam's generalization improves enough to match SGD+momentum on image tasks where it used to lose. AdamW is the default optimizer for essentially every modern LLM — it's what [[nanogpt]] and [[llm-from-scratch]] use. A common detail: apply weight decay to weight matrices but not to biases or LayerNorm gain/bias.

4. Learning-rate schedules: warmup + cosine decay

A constant learning rate is rarely optimal. The standard modern recipe is linear warmup → cosine decay:

Warmup ramps the LR from ~0 up to its peak over the first few hundred/thousand steps. Why? Early in training, Adam's v estimate is built from almost no data and is noisy/unreliable; large steps now can throw you into a bad region you never recover from. The NeurIPS 2024 analysis (Why Warmup the Learning Rate?↗) found the primary benefit is that warmup keeps training away from the divergence boundary, letting you ultimately use a higher peak LR than you otherwise could. Transformers in particular are notoriously warmup-sensitive — remove it and they often diverge.

Cosine annealing (Loshchilov & Hutter, SGDR — arXiv:1608.03983↗) then smoothly decays the LR from peak to ~0 following a half-cosine:

# after warmup_steps, decay from lr_max to lr_min over the run
import math
def lr_at(step, warmup, total, lr_max, lr_min=0.0):
    if step < warmup:
        return lr_max * step / warmup                      # linear warmup
    p = (step - warmup) / (total - warmup)                 # progress in [0,1]
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * p))  # cosine decay

Big steps in the middle (fast exploration), tiny steps at the end (careful settling into a minimum). SGDR's original twist was warm restarts — periodically jumping the LR back up to escape local minima — but for LLM pretraining people mostly use the single warmup-then-cosine-to-zero shape. This exact schedule is in [[nanogpt]] and [[mingpt]].

5. Normalization: BatchNorm and LayerNorm

Normalization layers keep activations at a sane scale (mean ~0, variance ~1) as they flow through a deep net, which makes optimization dramatically easier.

Batch Normalization (Ioffe & Szegedy, 2015 — arXiv:1502.03167↗) normalizes each feature across the mini-batch, then re-scales with learnable gamma (gain) and beta (shift):

x_hat = (x - mean_batch) / sqrt(var_batch + eps)
y     = gamma * x_hat + beta              # learnable: lets the net undo normalization if useful

At inference there's no batch, so BN uses running (EMA) statistics accumulated during training — a frequent source of train/test bugs. BN enabled much higher learning rates and acted as a mild regularizer (sometimes replacing dropout). The original paper blamed "internal covariate shift," but Santurkar et al. (2018 — [arXiv:1805.11604](https://arxiv.org/abs/1805.11604)) showed that explanation is essentially wrong: BN barely affects covariate shift, and what it actually does is smooth the loss landscape (improves Lipschitzness of loss and gradients), making gradients more predictive. A clean example of "the technique works, the original story for why doesn't."

Layer Normalization (Ba, Kiros & Hinton, 2016 — arXiv:1607.06450↗) normalizes across the features of a single example instead of across the batch:

mean, var over the feature dimension of ONE token/example
y = gamma * (x - mean)/sqrt(var + eps) + beta

Because it never looks at other examples, LayerNorm is batch-size independent and behaves identically at train and test time (no running stats, no inference bug). That's exactly why transformers and LLMs use LayerNorm (or its sibling RMSNorm), not BatchNorm — see [[attention-and-transformers]]. In [[nanogpt]] every block is LayerNorm → attention/MLP → residual.

6. Weight initialization

If you start with weights that are too big, activations and gradients explode; too small, they vanish — both kill learning in deep nets before it starts. The fix is to scale the initial random weights by the layer's fan-in:

• Xavier/Glorot init: variance ≈ 1/fan_in (or 2/(fan_in+fan_out)) — for tanh/sigmoid.
• Kaiming/He init: variance ≈ 2/fan_in — for ReLU (the 2 compensates for ReLU zeroing half the activations).

In [[makemore]] Karpathy shows live why a too-confident init makes the first training steps waste themselves just fixing the initialization (the "hockey-stick" loss curve), and how proper init + a normalization layer removes it. GPT-style models additionally scale residual-projection weights by 1/√(2·n_layers) so the residual stream doesn't blow up with depth.

7. Gradient clipping

Sometimes a single bad batch produces a huge gradient that, multiplied by the LR, blasts the weights into garbage (instant NaN). Gradient clipping caps the gradient's magnitude before the update:

total_norm = sqrt(sum(g_i**2 for all params))   # global L2 norm of the gradient
if total_norm > clip:                            # clip ~ 1.0 is standard
    scale = clip / (total_norm + 1e-6)
    for g in grads: g *= scale                   # rescale ALL grads, preserving direction

It preserves the direction of the update but caps its length. This is near-mandatory for RNNs and transformers, where rare loss spikes would otherwise derail a long run. [[nanogpt]] clips the global norm to 1.0.

Key ideas & tradeoffs

Dropout in one line: during training, randomly zero each unit with probability p (typically 0.5 for hidden layers, 0.1–0.2 for inputs). This forces redundant, robust features (no single neuron can be relied on) and approximates averaging an exponential ensemble of thinned networks (Srivastava et al., JMLR 2014↗). At test time you use the full network; "inverted dropout" scales the kept activations by 1/(1-p) during training so no test-time rescaling is needed.

The big mental model: optimizers (SGD/Adam) decide the step; schedules (warmup/cosine) decide its size over time; normalization (BN/LN) and init keep the signal well-scaled so the steps are meaningful; regularizers (dropout/weight decay) keep the solution general; clipping is the seatbelt. You almost always use several at once.

Honest caveats

• There is no universally best optimizer. Adam/AdamW dominates LLMs, but well-tuned SGD+momentum still wins some vision benchmarks on final accuracy. "Use AdamW" is a strong default, not a theorem.
• Stated reasons are often wrong even when the method works. BatchNorm's "internal covariate shift" story was largely debunked (Santurkar 2018↗); the real mechanism is landscape smoothing. Treat mechanistic explanations in original papers with healthy skepticism — the empirics are what shipped.
• Hyperparameters interact. LR, batch size, warmup length, weight decay, and clip threshold are not independent. Doubling the batch size usually wants a larger LR; changing the schedule changes the best weight decay. There is no free lunch; expect to sweep.
• Adam's memory cost is real at scale. Storing m and v doubles optimizer state vs. the model itself — a major factor in large-model training budgets and a driver of memory-saving variants (8-bit Adam, Adafactor, and newer optimizers like Lion/Muon). This connects to [[quantization]] of optimizer states.
• Regularization needs shift with scale. Dropout was essential in the small-data 2014 era; modern LLMs trained on trillions of tokens often use little or no dropout because the sheer data volume regularizes, and rely mainly on (small) weight decay. Don't cargo-cult p=0.5 into a foundation model.
• These papers are old; the field moved. Adam (2014) and BN (2015) are foundational but a decade old. Optimizer research is active (Lion, Sophia, Muon, schedule-free optimizers) and warmup is still being re-explained in 2024. Foundations are stable; "state of the art" is not.

How it connects to OpenAlice + the Academy ladder

This page is a mid-ladder foundation in the Atlas Academy: it's the bridge between "I can compute a gradient" and "I can train a real language model."

The ramp (build order):

1. [[math-for-ml-foundations]] — derivatives, the chain rule, vectors/matrices. Prerequisite.
2. [[micrograd]] — backprop from scratch: this is where the gradient comes from.
3. [[neural-network-from-scratch]] / [[nn-zero-to-hero]] — wire a real net and watch SGD/Adam run.
4. Optimization & regularization (this page) — the toolbox that makes step 3 actually converge.
5. [[makemore]] — see init, BatchNorm/LayerNorm, and LR tuning fix a real model live.
6. [[nanogpt]] / [[mingpt]] / [[llm-from-scratch]] — the exact recipe here (AdamW + warmup/cosine + grad-clip + LayerNorm + weight-init scaling) training a GPT. If you read nanoGPT's training loop, every line maps back to a section above.
7. [[attention-and-transformers]] — why LayerNorm (not BatchNorm) and why warmup is non-negotiable for transformers.
8. [[scaling-laws]] — how the optimizer recipe folds into the compute/data/parameters tradeoff at large scale.

Why OpenAlice cares. Alice is built on the conviction (NAO directive 2026-06-09) that autonomous coding is a base AGI capability, grown on the existing substrate and measured on real benchmarks. Whenever the lab fine-tunes or post-trains a model — LoRA/PEFT adapters ([[lora-and-peft]]), RLHF/alignment ([[rlhf-and-alignment]]), or RL post-training like [[grpo]] and [[agentic-rl-long-horizon-coding]] — this toolbox is the inner loop: AdamW with decoupled weight decay, a warmup→cosine LR schedule, gradient clipping to survive the high-variance gradients of RL, and the right normalization. The optimizer doesn't change when you move from pretraining to RLHF; only the loss does. Understanding optimization & regularization is therefore a hard prerequisite for every training-side contribution in the OpenAlice ecosystem — and a recurring source of "why won't this converge" bugs that this page is meant to help you debug.