05 — Training & Optimization

Training a neural network means repeatedly adjusting its parameters until the loss stops falling. This chapter covers that loop and the techniques attached to it.

Those techniques — learning rate, batch size, momentum, Adam, warmup, schedules, weight decay, gradient clipping, initialization, dropout, mixed precision, parallelism — each address one fact about the shape of the loss surface being searched. That fact is Gradient descent and why the surface shape decides everything, and the rest of the chapter follows from it.

The aim is to look at a training curve, name the mechanism that produced its shape, derive the fix, and state for each technique when it helps, when it does nothing, and when it hurts.

This chapter assumes only that a neural network is a function with many adjustable parameters. Every derivative used is spelled out in words before it is used.

The words, defined once

These symbols recur throughout the chapter.

• A parameter (also called a weight), written w, is one adjustable number inside the model. A modest model has millions of them; the whole collection is what “training” changes.
• The loss, written L, is a single number measuring how wrong the model’s predictions were on some data. Lower is better; training is the attempt to make it small.
• The gradient, written g or dL/dw (read “dee L by dee w”), is the list of slopes of the loss with respect to every parameter. It answers “if I nudge this parameter up slightly, does the loss go up or down, and how fast?” It has exactly one entry per parameter, so the gradient is the same size as the model.
• The learning rate, written eta (the Greek letter eta, pronounced “ay-ta”), is how far you move along that slope on each adjustment. It is the single most consequential number in training.
• A step (or iteration) is one adjustment of every parameter.
• A batch (or mini-batch) is the group of training examples used to compute the gradient for one step. Its size is written B.
• An epoch is one full pass through the training dataset. With N examples and batch size B, an epoch is N/B steps.
• The loss surface (or loss landscape) is the imagined shape you get by plotting the loss as a height above every possible setting of the parameters. Training is walking downhill on it. It has as many horizontal directions as there are parameters, but every important idea in this chapter is visible in two.

What goes in, what comes out

One training step takes exactly two things as input: a batch of examples — the raw inputs x and their correct answers y — and the model’s current weights. It returns one thing: a new set of weights, very slightly better than the old ones. It also reports the loss for that batch, which is the number you plot.

Concretely, for an image classifier: in goes a batch of 32 images (x) with their 32 correct labels (y), plus the 25 million weights the model currently holds; out come 25 million slightly different weights, plus one number — the loss on those 32 images. Repeat a few hundred thousand times. Every algorithm in this chapter changes only the middle of that loop.

flowchart LR
    B["Training batch<br/>inputs x, labels y"] --> F["Forward pass<br/>compute the loss L"]
    F --> G["Backward pass<br/>g = dL/dw"]
    G --> T["Gradient transform<br/>momentum, rescaling, clipping"]
    T --> U["Weight update<br/>w := w - eta * g_hat"]
    U -.->|next step| B

    style T fill:#40916c,color:#fff
    style U fill:#2d6a4f,color:#fff

Walking the five boxes of that loop in order:

• Training batch. A group of examples is drawn from the dataset: the inputs x and the labels y.
• Forward pass. The model runs on x, produces predictions, and those predictions are compared against y to compute the loss L. “Forward” means data flowing from the input toward the loss.
• Backward pass. The chain rule from calculus is applied in reverse through the model, layer by layer, to get the gradient g = dL/dw. This reverse sweep is called backpropagation (backprop): it computes the slope of the loss with respect to every weight in one pass, which is why training a huge model is affordable at all.
• Gradient transform. The raw gradient is rarely used as-is. It is smoothed with momentum, rescaled per parameter, and clipped if it is too large. Momentum an exponential moving average of gradients, Rmsprop per parameter rescaling, Adam momentum rmsprop and the bias correction derived and Gradient clipping are all about this box; the transformed gradient is written g_hat (“g hat”).
• Weight update. The transformed gradient is subtracted from the weights, scaled by the learning rate: w := w - eta * g_hat. Then the dotted arrow marks the next step, and the loop restarts on a fresh batch.

The aim is not to memorize update rules but to look at a loss curve, name the geometric fact that produced it, and derive the fix. That is what Diagnosing a training run covers.

Notation. This chapter writes Greek letters as words — eta, beta, lambda, kappa, sigma, rho — while papers write them as symbols. They mean the same thing, and each is glossed at first use.

1. Gradient descent, and why the surface shape decides everything

How fast you can train is decided not by how steep the loss surface is, but by the ratio between its steepest and its shallowest direction.

The update is one line:

w <- w - eta * grad_w L

Read it as “the new weights are the old weights minus the learning rate times the gradient of the loss with respect to the weights”. Everything interesting lives in what happens when you apply that line to a surface whose curvature differs by direction.

Curvature means how quickly the slope itself changes as you move. On a road, the slope is how steep the hill is; the curvature is how fast that steepness changes. Mathematically it is the second derivative — the derivative of the derivative.

A two-parameter model you can do by hand

Real models have millions of parameters. To see the effect, use a model with exactly two, called x and y, and a loss that is a bowl — steep along x, shallow along y:

L(x, y) = 0.5 * (100 * x^2 + 1 * y^2)

The bottom of that bowl is at x = 0, y = 0, which is the answer training is trying to reach. Take the two derivatives:

dL/dx = 0.5 * 100 * 2x = 100x        curvature along x: d(100x)/dx = 100
dL/dy = 0.5 *   1 * 2y =   1y        curvature along y: d(1y)/dy   =   1

So the gradient is grad = (100x, y), and the curvature is 100 along x and 1 along y.

The ratio between the largest and smallest curvature is the condition number, written kappa (the Greek letter kappa, pronounced “cap-a”). It is a single number saying how lopsided the bowl is. kappa = 1 is a perfectly round bowl; here kappa = 100/1 = 100.

What one step does to each coordinate

Under plain gradient descent each coordinate updates independently, so you can analyze the two directions separately. Substitute the gradient into the update rule:

x <- x - eta * (100x)  =  x * (1 - 100*eta)
y <- y - eta * (  1y)  =  y * (1 -   1*eta)

Each line says the coordinate gets multiplied by a fixed factor on every step. Two cases:

• If that factor has magnitude below one, the coordinate shrinks toward zero on every step. That is what you want — it means you are converging.
• If it exceeds one in magnitude, the coordinate grows without limit. That is divergence.

The two conclusions

The learning rate is capped by the sharpest direction. For x to converge you need |1 - 100*eta| < 1. Work that out: -1 < 1 - 100*eta < 1 gives 0 < 100*eta < 2, so eta < 0.02. Above 0.02, x diverges — no matter how gentle the rest of the surface is. Notice where 0.02 came from: it is 2 / 100, that is, 2 divided by the sharpest curvature.

Progress is set by the flattest direction. Push eta all the way to the largest stable value, eta = 0.02. Then y gets multiplied by 1 - 1*0.02 = 0.98 on every step. To cut the y error by a factor of 1000 you need 0.98^n = 0.001, so:

n = ln(0.001) / ln(0.98) = (-6.9078) / (-0.020203) = 342 steps

That is 342 steps spent on a coordinate whose gradient was never the problem.

(ln is the natural logarithm, the inverse of e^x: ln(z) is the power you must raise e ≈ 2.718 to in order to get z. It appears here because “how many times must I multiply by 0.98 to reach 0.001” is exactly a logarithm.)

Put the two together. The learning rate is bounded above by the sharpest curvature. The step count is bounded below by the flattest. So the whole optimization problem is the ratio between them — which is kappa.

Every improvement in this chapter attacks kappa, by one of two routes: reshape the surface so it is less lopsided (RMSProp, normalization, good init), or use history to move faster along the flat directions (momentum).

Normalization is used throughout this chapter and has not been defined yet: it is a layer inserted into the model that rescales the numbers flowing through it to a fixed mean and spread before the next layer sees them, so no direction can drift to a wildly different scale from its neighbours. Normalization and why transformers use layer norm derives the variants.

What this analysis assumes

It treats the loss as a quadratic bowl with fixed curvature. A real loss surface is not a bowl.

It survives anyway, because near any point a smooth surface looks like a bowl, and the curvature ratio there still sets the local speed limit. That is why the conclusion holds for real networks even though the assumption is false globally.

Where the picture genuinely breaks down is when the curvature itself changes fast as you move. That is exactly the regime early training lives in, and exactly why warmup exists (Learning rate schedules).

And when the surface really is round (kappa = 1), none of this matters: a single well-sized step lands at the bottom, and every technique below buys you nothing.

What interviewers probe: “Why not just use a bigger learning rate?” The answer is not “it diverges” — it is “the maximum stable rate is 2/lambda_max, set by the sharpest direction, and past that the sharp coordinate oscillates outward regardless of what the rest of the surface looks like.” Here lambda_max (“lambda max”) is the largest curvature anywhere on the surface, the mathematical name for “the sharpest direction”. In the bowl above, lambda_max = 100 and 2/lambda_max = 0.02, which is the cap derived two paragraphs up.

2. Batch, mini-batch, stochastic — the variance/throughput trade

How many examples belong in one step? The answer is a compromise between the accuracy of the gradient and the speed of the hardware, not a preference.

The one asymmetry that decides it

You never compute the gradient on the whole dataset — that would be one step per epoch. You compute it on a random sample of B examples and treat the answer as if it were the real thing.

That estimate is unbiased, meaning that averaged over many random batches it lands on the true gradient, with no systematic lean in any direction. What it is not is exact: it jumps around from batch to batch.

How much it jumps is measured by variance, the average squared distance from the mean. For a batch of B examples the variance of the estimated gradient is sigma^2 / B, where sigma (the Greek letter sigma) is the standard deviation of a single example’s gradient. Taking the square root of both sides:

std(g_B) = sigma / sqrt(B)

Noise falls as 1/sqrt(B) while cost rises as B. Quadrupling the batch quadruples the compute per step and only halves the noise. Every decision in this section follows from that asymmetry.

The three regimes

The table below puts the three batch-size extremes side by side. The row to watch is “Hardware use” — it is the one that decides real runs, and it is the one the noise math does not tell you about.

Four of those entries need a sentence.

• Stochastic is the traditional name for a batch of one example. It is a historical term, not a description of anything special the algorithm does.
• “Kernel-launch bound” means the run spends its time on the fixed overhead of dispatching work to the accelerator — each dispatch is called a kernel launch — rather than on arithmetic. A batch of one wastes almost the whole device on this overhead.
• Activations are the intermediate values every layer produces on the forward pass, held in memory until the backward pass consumes them. There is one set per example in the batch, which is why memory grows as O(B) — read “order B”, meaning proportional to B.
• Escaping sharp minima refers to gradient noise acting as exploration. A narrow, steep-walled valley in the loss surface is easy to get shaken out of when each step is slightly wrong, and such valleys tend to generalize worse than wide flat ones. So a noisy gradient is not purely a cost.

The trade with real numbers

Take a dataset of 50,000 samples on a device whose throughput stops improving near a batch of 256. Read the “Noise” column against the “Epoch time” column: memory buys you noise reduction, and the exchange rate collapses partway down the table.

Read the columns. “Noise (relative)” is sqrt(32/B), taking B = 32 as the reference: at B = 512 that is sqrt(32/512) = 0.25, at B = 4096 it is sqrt(32/4096) = 0.088. “Epoch time” is 50,000 / throughput.

Now compare the last two rows. Going from 512 to 4096 you pay 8x the memory and get:

• 0.25 / 0.088 = 2.8x less noise
• 2.5 s down to 2.4 s = 4% less epoch time

The device was already saturated at 512. Past that you are buying only variance reduction, at the worst exchange rate on the table.

Scaling the learning rate with the batch

Growing the batch changes what learning rate you should use, because a bigger batch gives a more trustworthy gradient and you can afford a longer step. There are two rules. Which one applies depends on the optimizer, and that pairing is the part people get wrong.

Linear scaling (eta ~ B, read “eta proportional to B”) goes with SGD, plain stochastic gradient descent. The argument: k sequential steps on batch B move about as far as one step on batch kB taken at k times the rate — provided the gradient is roughly constant over those k steps. That proviso is why linear scaling always ships with warmup: the assumption is worst at the very start of training, when the gradient changes fastest.

Square-root scaling (eta ~ sqrt(B)) goes with Adam. Adam divides each step by sqrt(v), which already absorbs part of the noise reduction a bigger batch gives you. Scaling linearly on top of that double-counts it and diverges. (Adam and its v are derived in Adam momentum rmsprop and the bias correction derived; for now, treat it as an optimizer that already rescales each step by recent gradient size.)

When batch size helps, does nothing, and hurts

Growing B helps while the device is underused. You get more samples processed per second for free.

It does nothing once the device is saturated, which the throughput column shows arriving between 512 and 4096. Past that you are paying memory purely for a smaller gradient error you probably did not need.

It hurts in two ways, at both ends. Very large batches remove the noise that was doing useful exploration, so runs settle into sharper minima and generalize worse — unless you compensate with the scaling rules and warmup. Very small batches hurt too: the fixed per-step overhead dominates and the gradient is so noisy that progress is erratic.

The whole analysis also assumes examples are drawn independently from the same distribution. If your data loader hands out a sorted or grouped stream — all the cats, then all the dogs — batches are not representative and the “unbiased” claim fails at its root.

3. Momentum — an exponential moving average of gradients

The cheapest fix for a lopsided loss surface is to average recent gradients: consistent directions accumulate and oscillating ones cancel.

Momentum replaces the raw gradient with a running average of recent gradients:

v <- beta * v + g
w <- w - eta * v

Two new symbols. v is the velocity, a running total that carries over between steps — it starts at zero and is updated in place, exactly like a running sum in a loop. beta (the Greek letter beta, pronounced “bay-ta”, typically 0.9) is how much of the old velocity survives each step.

This kind of average — each new value mixed into a running total with a fixed weight, so old values fade out geometrically — is an exponential moving average, abbreviated EMA. Unrolling the recursion makes the fading explicit:

v_t = g_t + beta*g_(t-1) + beta^2*g_(t-2) + beta^3*g_(t-3) + ...

The gradient from ten steps ago still counts, at 0.9^10 = 0.35 of its original weight. The gradient from fifty steps ago counts at 0.9^50 = 0.005, which is effectively nothing.

The mechanism: two limiting cases

Momentum’s entire behavior falls out of what the EMA does to two kinds of gradient sequence.

Case 1 — a consistent direction accumulates. Suppose g_t = g, the same value every step. Then v is a geometric series:

v = g * (1 + beta + beta^2 + beta^3 + ...) = g / (1 - beta)
at beta = 0.9:   v = g / 0.1 = 10g

Case 2 — an alternating direction cancels. Suppose g_t flips sign every step: +g, -g, +g, -g, .... After a while v settles into a repeating two-step cycle, alternating between two values call them v_a and v_b. Write down what one step does to each:

v_a = beta*v_b + g        (a step where the gradient is +g)
v_b = beta*v_a - g        (a step where the gradient is -g)

Substitute the second line into the first and solve for v_a:

v_a = beta*(beta*v_a - g) + g
    = beta^2*v_a - beta*g + g
=>  v_a - beta^2*v_a = g - beta*g
=>  v_a * (1 - beta^2) = g * (1 - beta)

Now use 1 - beta^2 = (1 - beta)(1 + beta) and cancel the (1 - beta) on both sides:

v_a = g / (1 + beta) = g / 1.9 = 0.526g      at beta = 0.9

Put the two cases together. The same EMA amplifies consistent gradients by 1/(1-beta) = 10x and attenuates oscillating ones to 1/(1+beta) = 0.53x. The ratio between those two treatments is 10 / 0.526 = 19x, or in closed form (1+beta)/(1-beta) = 1.9/0.1 = 19.

That 19x swing, favoring the direction that keeps agreeing with itself, is what solves the ravine problem.

Why that solves the ravine

A ravine is the shape a lopsided loss surface makes: a long narrow valley, steep across and nearly flat along. It is the kappa = 100 bowl from Gradient descent and why the surface shape decides everything, seen from above.

In a ravine, the two directions behave in exactly the two ways just analyzed. The sharp direction’s gradient flips sign every step as you bounce across the valley floor — case 2. The flat direction’s gradient points the same way for hundreds of steps — case 1.

So momentum cancels precisely the component that was wasting your steps and accumulates precisely the one that was making progress. It did not need to know which direction was which; the EMA sorted them by their own behavior.

The speedup, quantified

For a quadratic bowl with condition number kappa, tuned to the best beta, the known results are:

plain GD  converges at rate  (kappa - 1)/(kappa + 1)             = 0.980 per step
momentum  converges at rate  (sqrt(kappa)-1)/(sqrt(kappa)+1)     = 0.818 per step

steps for 1000x error reduction:   GD 342,   momentum 35

The “rate” is the factor the remaining error is multiplied by each step, so smaller is faster. At kappa = 100: plain GD gives 99/101 = 0.980, and momentum gives (10-1)/(10+1) = 9/11 = 0.818. The step counts are ln(0.001)/ln(rate) in each case, the same calculation as in §1.

Read the two step counts as orders of magnitude, not as exact figures. 342 is ln(0.001)/ln(0.980) at the displayed three digits, but the underlying rate here is 99/101 = 0.98020, which gives 345. It is the same 342 as in §1, where the rate was exactly 0.98 because eta was pinned at 0.02; carried into this setting it is a rounding of a slightly different quantity. The comparison that matters — roughly 340 against roughly 35 — is unaffected.

Momentum turns the dependence on the condition number from kappa into sqrt(kappa). Here that is a 10x speedup, and 10 = sqrt(100) is not a coincidence — it is the whole result. That is why beta = 0.9 is a default rather than the first thing you tune.

The trap: momentum silently multiplies your learning rate

From case 1, the effective step in a consistent direction is eta/(1-beta), which at beta = 0.9 is 10 * eta. Turning momentum on multiplies your effective learning rate by ten.

Adding momentum to a tuned SGD run without lowering eta reliably diverges in the first epoch. If §1 told you eta = 0.02 is the stability cap, then with beta = 0.9 you are effectively running at 0.2.

When momentum helps, does nothing, and hurts

Momentum assumes the gradient direction is persistent — that the direction which was right ten steps ago is still roughly right now.

It helps on an ill-conditioned but smooth surface, which is the normal case. There momentum is close to free: one extra buffer the size of the model, and a sqrt(kappa) speedup.

It does nothing on a well-conditioned surface. At kappa = 1, sqrt(kappa) = kappa = 1 and the two convergence rates coincide, so there is nothing to win.

It hurts when the persistence assumption breaks: near a sharp turn in the surface the accumulated velocity carries you past the corner and up the far wall. And it hurts immediately and predictably if you enable it without dividing your learning rate by roughly 1/(1-beta), because you have silently multiplied your step size tenfold.

4. RMSProp — per-parameter rescaling

The second fix for a lopsided surface gives every parameter its own learning rate, derived from how large its own gradients have recently been. Momentum helps when directions disagree over time; RMSProp helps when they differ in scale.

The name expands to root mean square propagation. “Root mean square” describes what it tracks: square the gradient, average it, take the square root — which gives a measure of typical gradient magnitude that ignores sign.

s <- rho * s + (1 - rho) * g^2          # elementwise
w <- w - eta * g / (sqrt(s) + eps)

Three symbols to name:

• s is a per-parameter running average of squared gradients. It starts at zero, like v in momentum.
• rho (the Greek letter rho, pronounced “row”, usually 0.9 or 0.999) controls how far back that average reaches — the same role beta played in momentum.
• eps (epsilon, a tiny constant, typically 1e-8) keeps the division safe when s is near zero.

“Elementwise” means each parameter keeps its own s, independent of every other parameter. If the model has 25 million weights, s has 25 million entries.

Why dividing by sqrt(s) makes the step scale-free

Suppose one parameter’s gradient is roughly constant at some value c. Then the running average of g^2 converges to c^2, so s -> c^2. Substitute that into the update:

update = eta * c / sqrt(c^2) = eta * c / |c| = eta * sign(c)

The c cancels. The update magnitude becomes independent of the gradient magnitude — every parameter moves about eta per step regardless of its scale.

That is a per-coordinate rescaling of the surface, which is a direct attack on kappa: whatever the curvature ratio between two directions was, after this rescaling both directions move at the same speed.

The case it was built for: embeddings

An embedding is a lookup table that maps each discrete symbol to a row of numbers the model can compute with. The symbol is usually a token — a chunk of text a few characters long, the unit a language model reads and writes.

A row of that table is only updated on steps where its symbol actually appears in the batch. Common tokens appear constantly and accumulate large gradients; rare tokens appear almost never and receive tiny, infrequent ones. Here is what that does to the step size, with eta = 0.01:

The SGD column is just eta * grad: 0.01 * 10.0 = 0.1 and 0.01 * 0.01 = 0.0001. The RMSProp column is eta * sign(grad), so both rows get 0.01.

Read the bottom row. Under SGD, a 1000x gap in gradient magnitude is a 1000x gap in learning speed — the rare token’s row takes a thousand times longer to move anywhere. Under RMSProp the gap is gone.

That is why adaptive methods dominate on sparse features, embedding tables, and anything with heterogeneous gradient scales — which is to say, on transformers.

One note on eps: it is not a rounding detail. It bounds the step when s is near zero, and s starts at zero, so eps governs the entire first-step behavior of the algorithm.

When per-parameter rescaling helps, does nothing, and hurts

It assumes two things: that differences in gradient scale between parameters are a nuisance to be normalized away, and that each parameter’s typical scale is stable over the window the average covers.

It helps enormously where those hold — sparse features, embedding tables, and any architecture mixing layers with very different natural scales.

It does nothing when every parameter already sees gradients of similar magnitude, which is roughly what a well-normalized network with good initialization gives you. There, dividing by sqrt(s) is close to dividing everything by the same constant, which is just a change of learning rate.

It hurts in two situations. First, when the scale difference is signal rather than nuisance — when a parameter genuinely should move less because it matters less — rescaling erases that information. Second, consider a parameter whose true gradient is zero and whose observed gradients are pure noise: then sqrt(s) is the size of the noise, the division cancels it, and the parameter takes a full-sized step in a random direction every step, forever.

5. Adam = momentum + RMSProp, and the bias correction derived

Combining the two fixes gives Adam, the default optimizer for essentially all modern deep learning. This section also derives bias correction, the part most often stated without explanation.

Adam stands for adaptive moment estimation. A “moment” is a statistical average: the first moment is the mean of the gradient, the second moment is the mean of the squared gradient.

Adam keeps a running average of each. The first moment is exactly momentum from Momentum an exponential moving average of gradients. The second is exactly RMSProp from Rmsprop per parameter rescaling. Adam is those two ideas in the same update line, plus one correction.

The diagram below is the family tree of the optimizers in this chapter — read it as “what was added to what”.

flowchart TD
    SGD["Plain SGD<br/>w -= eta * g"] --> MOM["Momentum<br/>average of past gradients<br/>fixes ravines"]
    SGD --> RMS["RMSProp<br/>average of squared gradients<br/>fixes scale mismatch"]
    MOM --> ADAM["Adam<br/>both averages, plus<br/>bias correction"]
    RMS --> ADAM
    ADAM --> ADAMW["AdamW<br/>decay applied outside<br/>the adaptive rescaling"]

    style ADAM fill:#40916c,color:#fff
    style ADAMW fill:#2d6a4f,color:#fff

Walking the tree: Plain SGD takes the raw gradient. Momentum averages past gradients and fixes ravines. RMSProp averages squared gradients and fixes scale mismatch between parameters. Adam keeps both averages, plus a bias correction. AdamW changes only where weight decay is applied, moving it outside the adaptive rescaling — the subject of Adamw why decoupling weight decay matters.

Here is the whole algorithm, five lines. The first two are momentum and RMSProp verbatim; the middle two are the new part; the last is the update.

m <- beta_1 * m + (1 - beta_1) * g          # first moment  (momentum)
v <- beta_2 * v + (1 - beta_2) * g^2        # second moment (RMSProp)
m_hat = m / (1 - beta_1^t)                  # bias correction
v_hat = v / (1 - beta_2^t)
w <- w - eta * m_hat / (sqrt(v_hat) + eps)

Read the symbols aloud:

• m is the running average of the gradient.
• v is the running average of the squared gradient.
• beta_1 (“beta one”, default 0.9) is how much of the old m survives each step.
• beta_2 (“beta two”, default 0.999) is the same for v.
• t is the step number, counting from one.
• m_hat and v_hat (“m hat”, “v hat”) are the bias-corrected versions.

Those defaults have a consequence worth noticing before the derivation. m remembers roughly the last 1/(1-0.9) = 10 gradients; v remembers roughly the last 1/(1-0.999) = 1000. That is a hundredfold difference in memory between the two averages, and it is what makes the next two subsections interesting.

Where the bias comes from, and the correction derived

Both EMAs start at zero. On step 1, m is a weighted average in which the “history” contributing 90% of the weight does not exist and is silently treated as zero. That drags the estimate toward zero — a lie about the data.

The useful part is that the lie decays at a known rate, which is why it can be corrected exactly rather than tuned away.

Unroll the recursion for m_t. Assume the gradient is roughly stationary with mean g, meaning its statistics do not drift over the window being averaged:

m_t = (1 - beta_1) * SUM_{i=1..t} beta_1^(t-i) * g_i

That is the same unrolling as momentum in Momentum an exponential moving average of gradients, with the (1 - beta_1) factor out front. Take the expected value — E[...] means the average over many runs — and pull the constant g out of the sum:

E[m_t] = (1 - beta_1) * g * SUM_{i=1..t} beta_1^(t-i)

The sum is a finite geometric series. Substituting j = t - i, it is 1 + beta_1 + beta_1^2 + ... + beta_1^(t-1), whose closed form is (1 - beta_1^t)/(1 - beta_1). So:

E[m_t] = (1 - beta_1) * g * (1 - beta_1^t) / (1 - beta_1)
       = g * (1 - beta_1^t)

The (1 - beta_1) cancels, and what is left is the true mean g multiplied by (1 - beta_1^t).

The estimate is low by exactly the factor (1 - beta_1^t), so dividing by that factor makes the estimator unbiased. That division is m_hat = m / (1 - beta_1^t), which is the bias-correction line. The same argument on v gives E[v_t] = E[g^2] * (1 - beta_2^t).

Sanity-check it at t = 1 with beta_1 = 0.9: the raw update gives m_1 = 0.1 * g_1, which is a tenth of the gradient. The correction divides by 1 - 0.9^1 = 0.1, recovering g_1 exactly. Which is right — with one sample, the best estimate of the mean is that sample.

What skipping the correction costs

Because the Adam step depends on the ratio m / sqrt(v), the two biases partly cancel, and people conclude the correction is optional. Work out what is left.

Write the uncorrected step in terms of the corrected quantities. Since m = m_hat * (1 - beta_1^t) and v = v_hat * (1 - beta_2^t):

uncorrected step   eta * m / sqrt(v)         (1 - beta_1^t)
---------------- = --------------------- = -----------------
intended step      eta * m_hat/sqrt(v_hat)   sqrt(1 - beta_2^t)

So the leftover distortion is (1 - beta_1^t) / sqrt(1 - beta_2^t). The last column of the table below is that expression evaluated, and it reads “the step you actually take is this many times the step you intended”:

Check one row by hand, t = 1: the numerator is 1 - 0.9 = 0.1, the denominator is sqrt(1 - 0.999) = sqrt(0.001) = 0.0316, and 0.1 / 0.0316 = 3.16.

Two results matter.

The worst distortion is not step 1 — it peaks at step 12, at 6.57x. The reason is the mismatch in timescales: m de-biases over roughly 1/(1-beta_1) = 10 steps, while v de-biases over roughly 1/(1-beta_2) = 1000. The fast one recovers first, and while it does, the ratio is maximally wrong — a nearly-correct numerator divided by a still-badly-shrunk denominator.

The distortion is still 26% at step 1,000. Bias correction is not a first-few-steps patch. It matters for roughly 1/(1-beta_2) steps, which at the default beta_2 = 0.999 is a thousand.

The code below is the five update lines above, written out. The comment on the m_hat line is the derivation you just read.

def adam_step(params, grads, state, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
    """One Adam update. state holds m, v, t per parameter."""
    state["t"] += 1
    t = state["t"]
    for name, g in grads.items():
        m = state["m"][name] = b1 * state["m"][name] + (1 - b1) * g
        v = state["v"][name] = b2 * state["v"][name] + (1 - b2) * g * g
        m_hat = m / (1 - b1 ** t)          # unbiased: E[m_t] = g * (1 - b1^t)
        v_hat = v / (1 - b2 ** t)
        params[name] -= lr * m_hat / (v_hat ** 0.5 + eps)
    return params

When Adam helps, does nothing, and hurts

Adam inherits both sets of assumptions: that gradient direction is persistent (from momentum), and that per-parameter gradient scale is a nuisance to normalize away (from RMSProp). Bias correction adds one more — that the gradient is stationary over the averaging window.

It helps on exactly the workloads where those hold and the surface is badly conditioned: transformers, embeddings, anything sparse. Its biggest practical win is removing the need to hand-tune a learning rate per layer.

It does nothing much on small, well-conditioned, densely-supervised problems. There, tuned SGD with momentum often matches or beats it — and it halves the optimizer state, because Adam stores both m and v for every parameter where momentum stores only m: 2N numbers against 1N.

That memory difference is concrete. On the 1B-parameter model in Scale’s table, at 4 bytes per fp32 number:

Adam:              m + v  = 2 * 1B * 4 bytes  =  8 GB optimizer state
SGD + momentum:    v only = 1 * 1B * 4 bytes  =  4 GB optimizer state

with weights (4 GB) and gradients (4 GB) counted too:
Adam 16 GB   against   SGD+momentum 12 GB

Half off the state, a quarter off the total.

Adam hurts in two well-known ways. It converges to slightly worse-generalizing solutions than SGD on some vision benchmarks. And its v makes it fragile to gradient spikes in a way plain SGD is not — the subject of Gradient clipping.

Bias correction itself has no downside. Skipping it is what hurts.

What interviewers probe: “What does bias correction actually do?” Weak: “it fixes the early steps.” Strong: “the EMAs start at zero, so E[m_t] = g(1 - beta_1^t); dividing by that factor makes them unbiased. The two corrections do not cancel — the leftover error peaks around step 10 at 6.5x the intended rate and is still 26% off at step 1,000, because v de-biases 100x slower than m.”

6. AdamW — why decoupling weight decay matters

The standard way of penalizing large weights breaks when it is combined with Adam. The fix is the W in AdamW.

Regularization means anything you add to training that is not about fitting the training data better, but about making the model generalize better to data it has not seen. Usually it works by discouraging the model from using more of its capacity than the data justifies.

The oldest form pushes weights toward zero. It comes in two variants that look identical and are not.

Two operations that are the same under SGD and different under Adam

L2 is named for the L2 or Euclidean norm, the square root of the sum of squared weights. L2 regularization adds a penalty term to the loss, so it enters training through the gradient:

g_total = g + lambda * w

lambda (the Greek letter lambda, pronounced “lam-da”) is the strength dial. Because the penalty rides in on the gradient, it flows into m and v like any other gradient contribution, and gets divided by sqrt(v_hat) along with everything else.

Weight decay skips the loss entirely. It shrinks the weights directly, outside the adaptive machinery:

w <- w - eta * m_hat/(sqrt(v_hat)+eps) - eta*lambda*w
                                         ^^^^^^^^^^^^
                                         added after, not divided by sqrt(v_hat)

Under plain SGD these are the same operation, because SGD does not divide by anything — subtracting eta*lambda*w and adding lambda*w to the gradient produce identical updates. Under Adam they are not the same, because Adam divides the gradient path by sqrt(v_hat) and the decay path is not on the gradient path.

The W in AdamW stands for exactly this decoupled weight decay.

How much the difference is worth

Take eta = 1e-3 and lambda = 0.01, and two parameters both currently at w = 1.0, differing only in their gradient history. sqrt(v_hat) is Adam’s estimate of typical gradient magnitude for each — high for a parameter receiving big gradients, low for one receiving small ones.

Compare the last two columns. AdamW gives both parameters the same decay, 1e-3 * 0.01 * 1.0 = 1.0e-5. L2-in-Adam gives parameter B a hundred times more, purely because B’s gradients happen to be smaller.

With L2 inside Adam, the regularization a parameter receives is inversely proportional to its gradient magnitude — so the parameters that are actively learning, and therefore most able to overfit, are the ones regularized least. That is precisely backwards.

Two follow-on consequences

First, lambda and eta become entangled under L2-in-Adam. The L2 term passes through sqrt(v_hat), which depends on gradient scale, which depends on the learning-rate schedule. Retune eta and your effective regularization moves with it, for no reason you asked for. AdamW’s decay is just eta * lambda * w, so it scales predictably and a lambda you tuned on one run transfers to the next.

Second, never decay biases or LayerNorm gains. A bias is the single additive constant each unit adds to its input. A LayerNorm gain is the single multiplier a normalization layer applies per feature after standardizing it (normalization layers are covered in Normalization and why transformers use layer norm).

Both have one degree of freedom per feature, so they add essentially no overfitting capacity — there is nothing there to regularize. And shrinking a norm gain directly attenuates the signal passing through that layer, which is damage rather than regularization.

The function below is how that exclusion is implemented in practice: it splits the model’s parameters into a decay group and a no-decay group, and hands both to the optimizer.

def param_groups(model, weight_decay=0.1):
    decay, no_decay = [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if p.ndim < 2 or name.endswith(".bias"):   # biases and norm gains
            no_decay.append(p)
        else:
            decay.append(p)
    return [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]

The test is p.ndim < 2: anything with fewer than two dimensions is a bias or a norm gain, since those are one-dimensional vectors while weight matrices are two-dimensional or more. That one check catches nearly everything you want to exclude; the name.endswith(".bias") clause catches the rest.

When decoupled decay helps, does nothing, and hurts

It assumes the model has more capacity than the data constrains, so shrinking the unconstrained directions costs nothing and buys generalization. (Regularization derives why “unconstrained directions” is the right way to think about what L2 removes.)

It helps whenever you use an adaptive optimizer and want regularization at all — which is why AdamW is the default for training transformers.

It does nothing under plain SGD, where L2 and weight decay differ only by a constant factor of eta and either choice is fine.

It hurts in three cases: when the model is underfitting, where any decay makes a too-weak model weaker; when lambda is large enough that the decay term dominates the data gradient, at which point weights collapse toward zero and the loss plateaus; and when it is applied to biases and normalization gains, where it attenuates signal for no regularizing benefit.

7. Learning-rate schedules

The learning rate should not stay fixed over a run, because the beginning and the end of training are genuinely different from the middle — and each difference has a specific mechanism behind it.

A schedule is a rule that sets the learning rate as a function of the step number, so eta becomes eta(t) rather than a constant.

Two shapes matter. Warmup raises eta from near zero to its full value over the first few thousand steps. Decay lowers it toward the end of the run. Most real schedules are one of each, glued together.

The tree below is the decision procedure for picking both. Follow it top to bottom; each leaf is a concrete choice.

flowchart TD
    A{"Adaptive optimizer<br/>Adam family?"} -->|yes| W["Always use warmup<br/>500-4000 steps"]
    A -->|no| W2["Warmup optional<br/>needed only at large batch"]
    W --> D{"Is the total step<br/>budget known?"}
    W2 --> D
    D -->|yes| C["Cosine decay to<br/>about 0.1 * eta_max"]
    D -->|no| P["Constant rate, then<br/>step down on plateau"]
    C --> CL["Clip the global<br/>gradient norm"]
    P --> CL

    style W fill:#bc6c25,color:#fff
    style CL fill:#2d6a4f,color:#fff

Read it top-down: two questions, then one box that applies either way.

Question 1: are you using an adaptive optimizer? An adaptive optimizer is one that rescales per parameter — anything in the Adam family. If yes, always use warmup, typically over 500 to 4,000 steps. With plain SGD, warmup is optional and needed only at large batch sizes.

Question 2: is the total step budget known in advance? If it is, cosine decay to about 0.1 * eta_max is the default. If it is not, hold a constant rate and step down on plateau — drop the rate when the validation loss stops improving.

Validation data is a slice of labelled examples held out of training and never used to update a weight. The loss measured on it estimates how the model will do on data it has not seen. The training loss cannot do that, because the model was fitted to it.

Both paths end in the same box. Clipping composes with every schedule, so whatever you chose above, clip the global gradient norm before every step (Gradient clipping).

Warmup, derived — and why transformers specifically need it

“The model is fragile early” is not a mechanism. There are two real ones.

1. The second-moment estimate is computed from almost no data.

At step t, Adam’s v has averaged roughly min(t, 1/(1-beta_2)) squared gradients. At t = 5 that is five samples. Five samples is not enough to estimate a variance, so 1/sqrt(v_hat) — the per-parameter step-size multiplier — jumps around wildly.

Say that plainly: the adaptive part of Adam, whose entire job is setting a sane per-parameter step size, is least reliable exactly when the weights are most random and the surface is least forgiving.

Warmup fixes it by keeping eta small over the window where v is noisy. The timescales line up: a 1,000-4,000 step warmup is the same order as 1/(1-beta_2) = 1000, which is how long v takes to fill up.

2. Transformers amplify a bad early step through the residual stream.

A transformer is the architecture behind modern language models. Two of its features matter here. Attention lets every position in the sequence look up information from every other position. The residual stream is a running sum that each layer adds its output into, rather than replacing. (Derivations: Attention derived as content based lookup and Residual connections the gradient highway.)

Attention computes a softmax over Q · K^T. Unpacking that: softmax is the function that turns a list of arbitrary scores into positive weights that sum to one; Q and K are the “query” and “key” matrices, produced by weight matrices W_Q and W_K, and their product Q · K^T scores how much each position should attend to each other position.

Now the failure. One oversized update to W_Q or W_K blows up those scores, which saturates the softmax into near-one-hot attention — nearly all the weight landing on a single position. A saturated softmax has near-zero gradient, because moving the scores a little no longer changes the output. So the layer stops learning, and there is no gradient left to undo the damage.

Architecture makes this better or worse. Post-LN, the original design that applies layer normalization after each block, makes it worse: the size of the update at initialization scales with depth. That is exactly why pre-LN, which normalizes before each block, became the default.

Here is the failure happening, on a 12-layer pre-LN transformer at eta = 3e-4 with no warmup. Watch two columns at once: train_loss going down and then back up, and grad_norm exploding and then dying.

step    train_loss    grad_norm
   0      10.9863        4.12      <- at the uniform floor, where any fresh init starts
  10       9.9034       71.60      <- v still tiny; steps ~6x intended (§5)
  20       8.4416      302.85
  40      11.2038      918.44      <- attention softmaxes saturated
  80      10.9701        0.07      <- gradient dead
 200      10.9700        0.06

grad_norm is the length of the whole gradient vector, a single number summarizing how large the update wants to be. It goes from 4 to 918 and then to 0.06 — that collapse is the softmax saturating.

The annotation at step 10 points back to the uncorrected-step table in Adam momentum rmsprop and the bias correction derived, where the distortion peaks at 6.57x around step 12. This run is being hit by that factor while its weights are still random.

Reading the trace against the uniform floor

The number to anchor on is ln(58000) = 10.9682.

A vocabulary is the set of distinct tokens the model can output; this one has 58,000. If the model predicts all 58,000 with equal probability — that is, if it knows nothing — the cross-entropy loss is ln(58000) = 10.9682. That is the uniform floor.

Step 0 sits at 10.9863, just above the floor, and it has to. With random weights the correct token’s score averages zero, so the expected loss is E[logsumexp(z)], and Jensen’s inequality gives E[logsumexp(z)] >= ln(V), with equality only when every score is identical.

You can check that by simulation. Draw 58,000 random scores at various spreads and average the resulting loss:

logit std        0        0.02      0.1       0.5       1.0       2.0
mean init loss   10.9682  10.9684   10.9732   11.0933   11.4682   12.9675

Every value is above 10.9682, and rising with the spread. (The closed form is ln(V) + s^2/2 for logit standard deviation s, which those numbers match to four digits.) A trace that opened below 10.9682 would be reporting something impossible, and that is a useful thing to notice in someone else’s logs.

By step 20 the run is at 8.44, genuinely below the floor. That is allowed: the floor is a statement about random weights, not a bound on the loss. Twenty steps were enough to learn the token frequencies, which beats uniform.

Then it is thrown back up to 10.97 and pinned there.

The model did not diverge to nan — it climbed back to the floor, collapsed to predicting uniform, and then had no gradient left to escape. (nan is “not a number”, the floating-point value produced by operations like 0/0 or inf - inf; once it appears it contaminates every subsequent computation.) A nan is at least loud. This is silent: the loss is a perfectly plausible-looking number that never moves again.

The identical run with a 2,000-step linear warmup reaches 6.1 by step 2,000 and 3.4 by step 20,000.

The four schedules

Each row below is one schedule shape. The “Formula” column gives eta as a function of the step t; the symbols are defined right after the table.

The symbols in those formulas: T_w is the warmup length in steps, T the total run length, eta_max and eta_min the peak and floor rates, and “anneal” means to lower gradually.

Check the cosine formula at the two ends. At t = 0, cos(0) = 1, so eta = eta_min + 0.5*(eta_max - eta_min)*2 = eta_max. At t = T, cos(pi) = -1, so the bracket is zero and eta = eta_min. It sweeps smoothly between them, spending most of the budget nearer the top.

Cosine’s real argument is what happens at the end. The final eta sets the width of the basin the optimizer can still escape from — a small step cannot climb out of a valley wider than itself. So the last few thousand steps become pure refinement inside one basin, which is what you want at the end of a run.

Annealing to exactly 0 is slightly worse than annealing to about 0.1 * eta_max, because a fully frozen rate stops averaging out gradient noise. You want the steps small, not zero.

When schedules help, do nothing, and hurt

Warmup assumes the early surface is both badly estimated and unusually dangerous. It helps for every adaptive optimizer, every large-batch run, and every transformer. It does nothing for a small-batch SGD run on a well-conditioned problem with good initialization, where the first steps are no more dangerous than the thousandth. It hurts only by wasting budget, and only when made far longer than 1/(1-beta_2).

Cosine decay assumes the total budget is known before the run starts. That assumption is what breaks in two directions: stop the run early and the rate never got low, so the final weights are mid-flight; extend a finished run and the rate is already at its floor, so the extension does almost nothing.

Plateau-based decay assumes a validation signal clean enough that a plateau is real rather than noise. On a noisy metric it fires early and freezes the run prematurely.

8. Gradient clipping

A single bad batch can destroy a run. Gradient clipping is the guard, and the less obvious question is why you still need it when the optimizer is Adam.

Gradient clipping caps how large an update is allowed to be.

||g||_2 (read “the L2 norm of g”) is the length of the gradient vector: square every entry, add them all up, take the square root. With 25 million parameters it is one number summarizing all 25 million.

If that length exceeds a threshold c, the whole vector is scaled down to length c:

if ||g||_2 > c:  g <- g * c / ||g||_2

Check that the scaling does what it claims. If ||g||_2 = 847 and c = 1.0, the vector is multiplied by 1/847, so its new length is 847 * (1/847) = 1. Every entry shrinks by the same factor, so the direction is untouched.

Clip by global norm, not per-tensor and not by value

Clip across all parameters at once. (A tensor is one rectangular array of numbers with a shape: a single layer’s weight matrix is a tensor, and so is a batch of images. “Per-tensor” clipping would mean rescaling each layer’s slice of the gradient on its own.)

Global-norm clipping is the only variant that preserves the gradient direction. Per-tensor clipping rescales different layers by different factors, and per-value clipping truncates individual entries — both change where the update points, which means you are no longer descending the gradient at all, just moving somewhere vaguely downhill.

What clipping prevents

Below is a real failure. Read grad_norm first: it is stable near 2.4, spikes to 847 on one step, and everything after that is the aftermath.

step   loss    grad_norm
1240   1.8802     2.31
1241   1.8790     2.44
1242   1.8814   847.02      <- one malformed sample in the batch
1243  11.4471     6.88      <- weights displaced outside the trained region
1244  11.4402     0.02
1900  11.4400     0.02      <- never recovers

Note that the loss at 1242 is still fine — 1.8814. The oversized gradient was computed on that step but applied at the end of it, so the damage shows up at 1243, where the loss jumps from 1.88 to 11.45 and stays there for the next 650 steps.

One sample, in one step, killed the run. With clip_grad_norm_(params, 1.0) the step at 1242 is scaled by 1/847 and the run does not notice.

Exploding and vanishing gradients

Clipping is also the standard defence against an exploding gradient: a gradient whose magnitude grows multiplicatively as it is propagated backward through many layers, so that by the time it reaches the early layers it is enormous. Multiply by 1.5 per layer across 50 layers and you have multiplied by 1.5^50 ≈ 6e8.

Its mirror image is the vanishing gradient: the same multiplicative process running the other way. Each layer shrinks the signal, so the early layers receive essentially nothing and never learn.

Clipping addresses only the exploding half — it can shrink a gradient that is too big, but it cannot invent one that is too small. The vanishing half is an initialization and architecture problem, treated in Initialization.

Why you still clip when using Adam

Adam already divides each coordinate by sqrt(v), so clipping looks redundant. Half of that intuition is right.

The half that is right: the immediate step is bounded. Take the isolated-spike case, where gradients have been near zero and one huge gradient g arrives. After that step, m ≈ (1-beta_1)*g and v ≈ (1-beta_2)*g^2. Substitute into the update:

step = eta * (1-beta_1)*g / sqrt((1-beta_2)*g^2)
     = eta * 0.1*g / (0.0316*g)          <- the g cancels
     = eta * 0.1/0.0316
     = 3.16 * eta

Whatever g was — 100, or a million — the step is 3.16 * eta. Adam does absorb the blow at the moment it lands.

The half that is wrong: the spike poisons v. The damage is not in the step you take on the spike step; it is in every step after it. Deriving how much is worth doing, because the obvious guess is wrong by three orders of magnitude.

How much a spike actually inflates v

The update is v <- beta_2*v + (1-beta_2)*g^2. Any single step enters v with weight 1 - beta_2 = 0.001. So a spike cannot dominate v no matter how large it is — it would need g^2 > 1000 merely to double it.

Start from equilibrium v = 1 and hit it with a 100x spike (g = 100, so g^2 = 10,000):

v        =  0.999*(1) + 0.001*(10000)  =  0.999 + 10  =  11.0
sqrt(v)  =  3.32

v went up 11x, not 10,000x. And because the step divides by sqrt(v), the effect is that every subsequent step is 3.32x too small, not too large.

How long the freeze lasts

The excess in v now decays geometrically at beta_2 per step. Starting from v = 11, the excess is 10, so v_k = 1 + 10*(0.999^k) after k further steps. The step penalty is sqrt(v_k):

step      0      220     1,203    3,861
v      11.00     9.02     4.00     1.21
step     3.3x    3.0x     2.0x     1.1x  too small

The freeze is deep but short: 3x for a couple of hundred steps, and essentially gone by four thousand.

Be careful about how you describe the middle of that decay, because two readings of “halved” disagree:

• The penalty goes 3.3x -> 2.0x between step 0 and step 1,203. That is a 40% drop, not a halving.
• The excess over 1x — the part that is actually a penalty, since 1x is normal — goes 2.32 -> 1.16 by about step 1,000. That is a halving.

Track the excess, because 1x is where the run is trying to get back to.

Either way, the practical picture is the same: a silent, localized stall. The loss curve flattens for part of an epoch and then resumes, with nothing in the logs to explain it. That is the real argument for clipping under Adam, and it is a far better interview answer than “it prevents exploding gradients.”

Order of operations

Unscale (under mixed precision) -> clip -> step.

Clipping scaled gradients clips against the wrong threshold: if the loss scaler multiplied everything by 1024, then a gradient of true norm 0.5 arrives as 512 and gets clipped at 1.0, destroying a perfectly healthy step. Mixed precision and the scaling it applies are covered in Scale.

When clipping helps, does nothing, and hurts

It assumes the occasional enormous gradient is an outlier — a corrupt sample, a numerical accident — rather than genuine signal about the loss surface.

It helps whenever that holds, which is nearly always in language and speech workloads, where a single malformed example can produce an arbitrarily large loss.

It does nothing when the threshold sits comfortably above the typical gradient norm. That is the intended state: a well-set clip fires on a fraction of a percent of steps.

It hurts when the threshold sits below the typical norm. Then every step gets rescaled, and you have quietly replaced your learning rate with one that shrinks exactly when the gradient grows. Training slows down precisely when the surface is steep, which is backwards. The symptom is a run that trains but never quite converges.

9. Initialization

Clipping guards a run once it is moving; whether it can move at all is decided before the first step.

Initialization is the choice of what values the weights hold before the first step. It is not a detail: for a deep network, most choices produce a model that cannot learn, and the reason is arithmetic rather than luck.

Why zero fails: symmetry

Take two hidden units in the same layer with identical incoming weights. Follow the consequences one link at a time:

h_1 = f(w_1 · x),  h_2 = f(w_2 · x),  with w_1 = w_2
=>  h_1 = h_2  for all x   =>   dL/dw_1 = dL/dw_2   =>   w_1 = w_2 forever

Here h_1 and h_2 are the outputs of the two units, f is the activation function, and w_1 · x is the dot product of that unit’s weights with its input.

Spoken out: identical weights means identical outputs on every input; identical outputs means backpropagation hands both units identical gradients; identical gradients means identical updates; so they stay identical forever. Nothing in training can ever tell them apart, because nothing in the network ever treats them differently.

A width-512 layer initialized to a constant has exactly one distinct unit; the other 511 are copies. With all weights at zero it is worse — forward activations are zero, so dL/dW for every layer below is zero and nothing learns at all. Biases may be zero; weights may not. Random initialization exists to break this symmetry, and that is its first job.

Why naive random fails: variance compounds with depth

Breaking symmetry is not enough, because the size of the random numbers matters exponentially.

One unit computes y = SUM_i w_i * x_i, a sum of n_in products. n_in is the fan-in, the number of inputs feeding one unit. If the w and x are independent and zero-mean, the variances of the terms add:

Var(y) = n_in * Var(w) * Var(x)

Define the per-layer gain k = n_in * Var(w). Then Var(y) = k * Var(x): each layer multiplies the variance of what flows through it by k. Stack L layers and the variance is multiplied by k^L.

Anything but k = 1 is exponential in depth. That is the whole argument, and the table shows what the exponent does. Take n_in = 256 and vary only the standard deviation you initialized with:

Follow the first row: Var(w) = 0.10^2 = 0.01, so k = 256 * 0.01 = 2.56, and 2.56^10 = 1.2e4 while 2.56^50 = 2.6e20. Third row: Var(w) = 1/256, so k = 256 * (1/256) = 1 exactly, and 1^L = 1 at any depth. Only that row survives.

These are the exploding and vanishing regimes named in Gradient clipping, arriving before the first gradient is even computed.

The middle row is the more dangerous one, because it does not crash. It produces a run whose loss sits flat at chance while every activation and gradient in the network is 1e-16 — there is no error message, just a model that never learns.

Xavier and He, derived from the same requirement

Both standard formulas fall out of one demand: keep the per-layer gain at exactly one, so signal neither grows nor shrinks with depth.

Forward pass. Set k = n_in * Var(w) = 1, which gives Var(w) = 1/n_in.

Backward pass. Run the same argument in reverse through the chain rule, and gradient variance is preserved when Var(w) = 1/n_out, where n_out (the fan-out) is the number of units the layer feeds.

Those two demands conflict unless the layer is square. So Glorot takes the compromise Var(w) = 2/(n_in + n_out), which is the harmonic-mean-flavoured average of the two: if n_in = n_out = n it reduces to 2/(2n) = 1/n, exactly matching both.

Glorot initialization is also called Xavier initialization, after its author’s first name — the two names are the same thing, and interviewers use both. The derivation assumes the activation is roughly linear near zero. That is true for tanh, weakly true for sigmoid, and false for ReLU.

ReLU (rectified linear unit) is max(0, z): it passes positive values through unchanged and replaces negatives with zero.

Because it zeroes half its inputs, for symmetric zero-mean z you get E[relu(z)^2] = 0.5 * E[z^2] — half the terms contribute nothing to the sum of squares. So the forward gain picks up an extra factor of 1/2. Compensate for it and you get He initialization:

Var(y) = 0.5 * n_in * Var(w) * Var(x)  =  1     =>     Var(w) = 2/n_in
std(w) = sqrt(2 / n_in)

He init is Xavier times sqrt(2), and the 2 is literally the reciprocal of the fraction of inputs ReLU keeps alive (ReLU keeps 1/2, so the correction is 2).

Get it wrong in the other direction and the cost is exponential. Use Xavier with ReLU in a 50-layer network: the per-layer gain is 0.5 instead of 1, so activations shrink by 0.5^50 = 9e-16. That is the vanishing regime, arriving purely from a missing sqrt(2).

The table below extends the same argument to the activations you will actually meet. The “Gain” column is the multiplier applied to the baseline sqrt(1/n_in).

Reading the rows:

• tanh and sigmoid are the classic S-shaped activation functions, squashing their input into (-1, 1) and (0, 1) respectively. Both are roughly linear near zero, which is what Xavier’s derivation assumes.
• Leaky ReLU passes negatives through multiplied by a small slope a instead of zeroing them. Its correction is sqrt(2/(1+a^2)) — at a = 0 that is sqrt(2), plain He, and as a -> 1 it goes to 1, plain Xavier, because the function has become linear.
• SELU (scaled exponential linear unit) is designed so that activations converge to zero mean and unit variance on their own, which is what “self-normalizing” means. It therefore wants the plain 1/n_in LeCun variance with no extra gain.

The activation functions themselves are derived in Activations and the dying relu mechanism.

The three functions below are the two standard formulas plus one case the standard formulas do not cover. Read the third one’s docstring carefully — it is the residual case, unpacked immediately after.

import math
from numpy.random import randn, uniform

def he_normal(n_in, n_out):
    return randn(n_in, n_out) * math.sqrt(2.0 / n_in)

def xavier_uniform(n_in, n_out):
    limit = math.sqrt(6.0 / (n_in + n_out))     # Var of U(-a,a) is a^2/3
    return uniform(-limit, limit, size=(n_in, n_out))

def residual_branch_scale(n_in, n_out, num_layers):
    """Each block ADDS to the stream, so variance grows linearly with depth
    even with perfect per-layer init. The 2 is the two residual adds per
    transformer block; dividing by sqrt(2L) makes the total variance O(1)
    in depth instead of O(L) -- bounded, NOT equal to one."""
    return he_normal(n_in, n_out) / math.sqrt(2 * num_layers)

The residual case, which per-layer init does not cover

Everything above analyzed one layer feeding the next. A residual architecture does something the analysis cannot see: each block adds its output into a running sum rather than replacing it, and a transformer block adds twice, once per sublayer (The transformer block).

Independent variances add. So a stream entering at variance 1 and receiving 2L unit-variance additions leaves L blocks at:

Var = 1 + 2L        at L = 32:   1 + 64 = 65

That grows linearly in depth. Perfect per-layer initialization does not prevent it, because the growth is in the additions, not the layers.

The fix is to divide each branch’s output by sqrt(2L), which divides each contribution’s variance by 2L. Then:

Var = 1 + 2L * (1/(2L)) = 1 + 1 = 2       independent of L

Say what that buys precisely, because the usual phrasing overclaims it. The scaling makes the growth O(1) in depth instead of O(L). It does not restore unit variance.

Simulated at L = 12, 32, 80: a two-branch-per-block stream ends at variance 2.00 every time. With one branch per block it ends at 1.50. Neither number is 1, and — this is the point — neither moves with depth. Depth-independence is the property you wanted; a final norm before the output layer handles the remaining constant factor.

When initialization schemes help, do nothing, and hurt

They assume three things: that weights are independent and zero-mean, that the activation behaves as the derivation says, and — critically — that nothing downstream re-normalizes the signal.

They help enormously in deep plain networks. There they are the difference between training and not training at all.

They do far less in a modern network with normalization layers after every block. A normalization layer rescales activations to unit variance regardless of what arrived, absorbing an initialization error that would otherwise compound. That is why architectures with normalization are famously forgiving about initialization scale.

They hurt in two ways. When the gain assumption is wrong for the activation actually used — Xavier with ReLU in a deep stack, losing a factor of 0.5 per layer, is the canonical example. And in residual architectures, if applied per-layer without the 1/sqrt(2L) branch scaling, because the per-layer analysis does not see the additive stream.

10. Regularization

There are five places in a training step where you can constrain a model, and two of the standard constraints turn out to be the same technique.

Regularization is any deliberate constraint that trades a little training accuracy for better performance on data the model has never seen.

Overfitting is the failure it targets: a model fitting the noise and idiosyncrasies of the training set rather than the pattern. Its signature is training loss that keeps falling while performance on held-out data gets worse.

The diagram below lists the five places in a training step where you can attach a constraint.

flowchart LR
    D["Data<br/>augmentation"]
    M["Model<br/>dropout and<br/>normalization layers"]
    L["Loss<br/>L1 · L2 · label smoothing"]
    O["Optimizer<br/>decoupled weight decay"]
    S["Schedule<br/>early stopping"]

The five boxes are deliberately unconnected: they are five independent places you can attach a constraint, laid out in the order a training step meets them. They are not a pipeline — nothing flows from one to the next, and you can use any subset.

The five are not interchangeable — each constrains a different thing:

• Data. Augmentation manufactures new training examples.
• Model. Dropout and normalization layers restrict what the architecture can represent or rely on.
• Loss. L1, L2 and label smoothing add penalties to the objective.
• Optimizer. Decoupled weight decay shrinks weights outside the loss (Adamw why decoupling weight decay matters).
• Schedule. Early stopping limits how far the fit can go.

L2 does not shrink everything equally

Take one weight, a quadratic data loss around its unpenalized optimum, plus an L2 penalty:

argmin  0.5*h*(w - w*)^2 + 0.5*lambda*w^2

Three symbols. argmin reads “the value of w that minimizes what follows” — as opposed to min, which would be the minimum value itself. w* (“w star”) is the value the weight would take with no penalty at all. h is the curvature of the loss along that direction, which is a measure of how strongly the data constrains that weight: high h means moving the weight costs a lot of loss, low h means the data barely cares.

To minimize, set the derivative to zero and solve:

d/dw [0.5*h*(w - w*)^2 + 0.5*lambda*w^2]  =  h*(w - w*) + lambda*w  =  0
=>  h*w - h*w* + lambda*w = 0
=>  w*(h + lambda) = h*w*
=>  w = w* * h / (h + lambda)

So the penalized weight is the unpenalized one multiplied by h/(h + lambda). That fraction is what does the work.

Evaluate it at lambda = 0.01 across three curvatures:

h = 1.0    ->  1.0   / 1.01   = 0.990   of w* retained
h = 0.01   ->  0.01  / 0.02   = 0.500   of w* retained
h = 0.001  ->  0.001 / 0.011  = 0.091   of w* retained

A strongly-constrained direction (h = 1.0) keeps 99% of its value. A barely-constrained one (h = 0.001) loses 91% of it.

L2 does not make weights small — it deletes the directions the data did not constrain and leaves the constrained ones alone. That is why it improves generalization rather than just shrinking the model.

L1 is sparse because its gradient does not vanish

L1 is the sum of absolute weight values rather than squared values. The difference in behavior is entirely about what happens near zero, and it is visible in one pair of derivatives:

d/dw of lambda*|w|     = lambda * sign(w)      constant magnitude
d/dw of 0.5*lambda*w^2 = lambda * w            shrinks to 0 as w -> 0

Read the right-hand column. L2’s pull is lambda * w, which weakens exactly as the weight approaches zero — so L2 asymptotes toward zero and never arrives. L1’s pull is lambda regardless of how small w is, so it keeps pushing all the way down and the weight lands on exactly zero.

Once there, it stays. |w| has a corner at zero with no single slope, so the derivative is replaced by a subgradient: any value between -lambda and +lambda is a valid slope there. That whole range acts as a dead zone, holding the weight at zero against any data gradient smaller than lambda.

The closed form is a soft threshold:

w = sign(w*) * max(0, |w*| - lambda/h)

Everything with |w*| < lambda/h becomes exactly 0; everything larger is shifted toward zero by lambda/h and keeps its sign.

Sparsity is a property of the gradient’s behavior at the origin, not of the penalty’s shape. (“Sparsity” means many weights being exactly zero, which is what makes L1 usable as a feature-selection tool — a weight of exactly zero means the feature is not used at all.)

Dropout as an implicit ensemble — and the scaling everyone gets wrong

Dropout randomly switches off units during training. Each unit is kept with probability p at train time.

That makes it an implicit ensemble. A layer of n units has 2^n possible on/off patterns, so it defines 2^n subnetworks that all share one set of weights, and each training step trains one randomly sampled member.

Watch the convention on p, because this chapter and the last one differ. Here p is the keep probability, so the correction factor is 1/p. Dropout writes the same technique with p as the drop probability and the factor 1/(1-p). Both conventions are standard and both appear in real codebases and papers; the framework APIs mostly use drop (PyTorch’s nn.Dropout(p) drops with probability p). Nothing below changes if you swap them — substitute p -> 1-p throughout — but a factor read off one page and applied on the other is inverted. When you see a bare p, find out which one it is before you use it.

At test time you want the ensemble average: the average prediction over all 2^n subnetworks. Running 2^n forward passes is not an option, so the cheap approximation is one full-network pass with the expected input to each unit matched.

Work out what that requires. A downstream unit sees SUM_i w_i * x_i * mask_i, where mask_i is 1 if unit i survived and 0 if it was dropped, so E[mask_i] = p:

train (mask sampled):   E[preactivation] = p * SUM w_i x_i
test  (no mask):            preactivation =     SUM w_i x_i     <- 1/p too large

A preactivation is the weighted sum arriving at a unit before its activation function is applied.

At p = 0.5 every test-time preactivation is 2x what the network was trained to handle. Through tanh that means saturation; through a softmax it means wildly overconfident logits (the raw pre-softmax scores).

There are two ways to fix the factor of p, and only one of them is used:

• Classic dropout multiplies the weights by p at test time.
• Inverted dropout divides by p at train time, so test needs no adjustment at all.

Every framework does the inverted version, which is why model.eval() — the call that switches a model from training to inference behavior — only has to disable the mask and nothing else.

The implementation is one line of arithmetic; the / p_keep is the inversion.

import torch

def dropout(x, p_keep, training):
    if not training or p_keep >= 1.0:
        return x
    mask = (torch.rand_like(x) < p_keep).float()
    return x * mask / p_keep      # inverted: E[out] == x, so eval() is a no-op

The bug that ships: forgetting model.eval() before validation. Dropout stays on, validation loss is noisier and worse, and the run reads as if it is not overfitting when it is. Its mirror image — forgetting model.train() afterwards — silently removes all regularization for the rest of training.

Early stopping IS weight decay

Early stopping means halting training when the validation loss stops improving, rather than running to convergence. It looks like a scheduling decision. It is arithmetically the same thing as L2.

Here is why. Reuse the two-parameter analysis from Gradient descent and why the surface shape decides everything: along a direction with curvature h, gradient descent multiplies the remaining error by (1 - eta*h) on every step. Starting from w = 0, the distance covered toward w* after t steps is:

w_t = w* * (1 - (1 - eta*h)^t)  ~  w* * (1 - exp(-eta*h*t))

The second form uses (1 - x)^t ≈ exp(-x*t) for small x, which holds here because eta*h is small.

Now compare that with the L2 solution derived above, w* * h/(h+lambda). Both expressions do the same thing: they are near 1 for large h and near 0 for small h. Both say “retain the high-curvature directions, suppress the low-curvature ones.”

Match the point where each transitions from retaining to suppressing, and you get:

lambda_effective ~ 1/(eta * t)

The longer you train, the weaker the implied regularization. Evaluate it at eta = 1e-3, tracking a well-constrained direction (h = 1) and a barely-constrained one (h = 0.001):

Check one cell: at t = 2,000 and h = 1, 1 - exp(-1e-3 * 1 * 2000) = 1 - exp(-2) = 0.865, which is the 86%. At the same t with h = 0.001, 1 - exp(-0.002) = 0.002, which is the 0.2%.

Read the two right-hand columns down. The h = 1 direction is essentially fully learned by 20,000 steps and stays there. The h = 0.001 direction is still only 18% learned at 200,000 steps. Training longer is how the weakly-constrained directions eventually get learned — which is exactly what L2 was there to prevent.

Stopping 10x earlier is approximately 10x more weight decay. That is why an early-stopped run and a heavily-decayed run often land on nearly the same validation number, and why tuning both hard at once wastes budget: you are turning one knob twice.

The rest, briefly

The four techniques above each earned a derivation. These four do not, but each is in wide use and each has one failure mode worth carrying, so the table gives the mechanism and the trap in a line apiece; the paragraphs after it unpack the vocabulary and the one derivation that does not fit in a cell.

Four terms in that table need a gloss.

A manifold is the low-dimensional surface that real data actually occupies inside its much larger raw space. A 224×224 RGB image is 150,528 numbers, but photographs of cats form a vanishingly thin sheet inside that 150,528-dimensional space, because almost every possible setting of those numbers is static. “Enlarges the dataset along exactly that manifold” means an augmentation produces new points that are still on the sheet — a rotated cat photo is a photo you could have taken — rather than points floating off it, which would teach the model nothing about photographs.

An invariance is a change to the input that must not change the label. Rotating a cat photo leaves it a cat photo. Augmentation works by generating exactly those changes, and it fails when the change is not actually label-preserving: a flipped digit 2 is not a 2. (OCR, optical character recognition, reads text from images; cropping away part of a character changes what it says.)

BatchNorm (batch normalization) standardizes each feature using the mean and variance of the current batch, which couples every example’s output to whichever other examples happened to share its batch.

LayerNorm (layer normalization) standardizes each example across its own features instead, so there is no coupling at all.

Label smoothing, derived

Label smoothing gets a derivation, because the derivation is what explains the failure it fixes.

Cross-entropy is the standard classification loss: the negative log of the probability the model assigned to the correct class. Predict the right answer with probability 1 and the loss is -ln(1) = 0; predict it with probability 0.5 and the loss is -ln(0.5) = 0.69.

With hard targets — exactly 1 for the true class and 0 for the rest — that loss can only be driven to zero by pushing the correct probability to exactly 1, which requires an infinite logit gap. There is no finite optimum. So the model inflates the gap forever and becomes arbitrarily overconfident.

Label smoothing replaces those hard targets with softened ones, spreading a small mass eps (epsilon) evenly over all K classes. At eps = 0.1 over K = 1000 classes:

target p_correct = 1 - 0.1 + 0.1/1000 = 0.9 + 0.0001 = 0.9001
target p_other   =         0.1/1000   =                1e-4
optimal logit gap = ln(0.9001 / 1e-4) = ln(9001) = 9.10

The true class keeps 1 - eps of the mass plus its own 1/K share of the redistributed eps. Every other class gets eps/K. The optimal logit gap is the log of the ratio between them, because that is what a softmax needs to produce those two probabilities — and it is now a finite number.

Smoothing replaces “maximize the margin forever” with “reach a margin of 9.1 and stop.”

That is also why it improves calibration — the property that a prediction stated at 70% confidence is right about 70% of the time. A model with no finite optimum ends up stating 99.99% on everything.

Why dropout and BatchNorm interact badly

The mechanism is a variance shift, and it is derivable in four lines.

Inverted dropout with keep probability p outputs x/p with probability p, and 0 otherwise. For zero-mean x:

E[out]    = p * (E[x]/p) = 0                  mean preserved
E[out^2]  = p * E[x^2]/p^2 = E[x^2]/p
Var_train = Var(x) / p                        variance INFLATED by 1/p

Note what happened: the 1/p correction fixed the mean exactly, which is what it was designed for. It did not fix the variance — squaring turns 1/p into 1/p^2, and only one factor of p comes back from the probability of surviving.

At inference the mask is off, so Var_test = Var(x). So dropout’s output has one variance at train time and a different one at test time.

That is harmless on its own. It stops being harmless when a BatchNorm layer downstream memorized the train-time value. BatchNorm’s running_var is the variance it accumulated over training batches and uses at inference, since at inference there may be no batch to compute statistics from. It therefore holds Var(x)/p, and at test it divides by that stale number:

output std at test = sqrt(Var(x)) / sqrt(Var(x)/p) = sqrt(p) = 0.707   at p = 0.5

Every BN layer downstream of a dropout layer systematically shrinks its output by sqrt(p) at inference — a bias that exists at test time only, and therefore never appears in your training loss.

It compounds. Stack ten such blocks and each layer is fed inputs already off-distribution relative to the statistics its own running_var encodes.

The symptom is unmistakable once you know it: training loss healthy, validation fine while train()-mode statistics are used, and test accuracy collapsing the moment you switch to eval().

Three fixes, in order of preference:

1. Use LayerNorm. It computes per-sample statistics at inference, so there is nothing stale to shift. This is a real reason transformers dodge the problem entirely.
2. Put dropout only after the last BN, typically in the classifier head, so no BN layer is downstream of it.
3. Do not use both. BN already regularizes through batch-statistic noise, so the dropout may be buying you nothing anyway.

When regularization helps, does nothing, and hurts. Every technique here assumes the model has more capacity than the data constrains — that there is something to give up. That assumption is what decides everything. Regularization helps when the training loss is far below the validation loss, which is the signature of exactly that surplus. It does nothing measurable when the model is already at capacity for the data, and it hurts when the model is underfitting, because you are removing capacity from a model that did not have enough; this is why adding dropout to a model whose training loss is still high makes things strictly worse. Beyond that shared assumption each technique has its own: augmentation assumes the transformation preserves the label, dropout assumes redundant co-adapted features and enough width to spare, label smoothing assumes you care about calibrated probabilities rather than well-separated embeddings, and BatchNorm’s regularizing noise assumes a batch large enough that its statistics are meaningful.

11. Diagnosing a training run

Diagnosis is a procedure, not an art: first a single test that separates bugs from everything else, then a decision tree mapping the shape of a loss curve onto the mechanism that produced it.

Step zero: overfit a single batch

Before anything else, take 8 samples, turn off shuffling, augmentation, dropout, and weight decay, and train until the loss reaches approximately zero.

Turning those four off is not incidental. Shuffling and augmentation change what the 8 samples are; dropout and weight decay are handicaps that fight memorization. On this test you want the model to memorize, so remove everything designed to stop it.

The function below is the gate. It returns (True, step) if the loss got under tol, (False, steps) if it never did.

def overfit_one_batch(model, batch, opt, criterion, steps=500, tol=1e-3):
    """Sanity gate. Must reach ~0 loss or there is a wiring bug.

    `criterion` is the loss function, e.g. torch.nn.CrossEntropyLoss().
    Pass it in rather than reaching for a global: the wrong loss is one of
    the bugs this gate exists to catch, so it has to be visible at the
    call site.
    """
    model.train()
    x, y = batch
    for step in range(steps):
        opt.zero_grad(set_to_none=True)
        loss = criterion(model(x), y)
        loss.backward()
        opt.step()
        if step % 50 == 0:
            print(f"{step:4d}  {loss.item():.4f}")
        if loss.item() < tol:
            return True, step
    return False, steps

Two shapes, side by side — a healthy run memorizing 8 samples over 10 classes, and a broken one that never moves:

healthy (8 samples, 10 classes)        broken
   0   2.3026  <- ln(10), chance          0   2.3026
  50   1.9034                            50   2.3025
 100   0.8412                           100   2.3026
 200   0.1107                           400   2.3025  <- cannot memorize 8 rows
 400   0.0009  <- memorized

Both columns start at the same place. ln(10) = 2.3026 is the loss of a model that assigns equal probability to all 10 classes — the same uniform floor as in Learning rate schedules, just with V = 10 instead of 58,000. That is what “chance” means here.

The difference is what happens next. The healthy run leaves 2.3026 immediately and is at 0.0009 by step 400. The broken run is still at 2.3025 at step 400: it has not moved at all, which means it is still emitting a uniform guess after 400 attempts to memorize 8 rows.

This isolates the bug because a model that cannot memorize 8 examples has no capacity-or-data explanations left. Every usual suspect is excluded by construction:

• Generalization is off the table — there is no held-out set.
• Dataset size is off the table — 8 rows is the dataset.
• Regularization is off the table — you turned it off.
• The data pipeline is reduced to one tensor you can print.

What remains is a short list, and it is worth memorizing:

• labels misaligned with inputs;
• the loss applied to the wrong axis, or a softmax applied before CrossEntropyLoss, which applies its own softmax internally and therefore double-applies it;
• requires_grad=False on parameters you meant to train;
• zero_grad() missing or misplaced, so gradients from old batches accumulate into new ones;
• an optimizer built over an empty or stale parameter list;
• eta = 0.

Running this first turns a two-day “why won’t it learn” into a ten-minute check, and it is the answer interviewers most want to hear to “your model isn’t training — what do you do?” A fuller debugging procedure, covering data and metric bugs as well, is in 09 — Model Debugging Playbook.

The decision tree

Once the single-batch gate passes, the shape of the loss curve is the evidence, and there are only a handful of shapes. The tree below maps each one onto the mechanism that produces it, so a run can be diagnosed from its plot before anything is instrumented; the table after it gives the first fix for each leaf.

flowchart TD
    S0{"Can it overfit<br/>8 samples to ~0 loss?"} -->|no| BUG["Wiring bug:<br/>labels, loss axis, frozen params,<br/>missing zero_grad, lr = 0"]
    S0 -->|yes| S1{"Train loss shape?"}
    S1 -->|"flat at ln(C)"| LR0["LR too low, dead ReLUs,<br/>or vanishing init"]
    S1 -->|"loss rises then nan"| LR1["LR too high, or<br/>unclipped gradient spike"]
    S1 -->|"spiky"| SP["Batch outliers, or<br/>batch too small"]
    S1 -->|"falling"| S2{"Val loss?"}
    S2 -->|"falling with it"| OK["Healthy — train longer<br/>or grow the model"]
    S2 -->|"rising"| OF["Overfitting"]
    S2 -->|"below train loss"| LK["Almost always dropout or<br/>BatchNorm accounting;<br/>only then leakage"]
    S2 -->|"flat from step 0"| BR["Val set broken,<br/>trivial, or mislabeled"]

    style BUG fill:#9d0208,color:#fff
    style OK fill:#2d6a4f,color:#fff
    style LK fill:#1d3557,color:#fff
    style OF fill:#bc6c25,color:#fff

Colour is a verdict, not decoration:

• Red is a bug in your code.
• Green is a healthy run.
• Orange is a real modelling problem in an otherwise correct run.
• Blue is a measurement artefact — the model may be fine and the number is simply not comparable.

Leakage and overfitting are deliberately not the same colour, because they are not the same kind of finding: overfitting is something the run did, leakage is something the split did.

Read the tree top-down. First question: can the model overfit 8 samples to near-zero loss? If it cannot, stop and hunt a wiring bug — labels, loss axis, frozen params, missing zero_grad, lr = 0.

If it can, look at the shape of the training loss. Three of the four shapes end the diagnosis immediately: flat at ln(C), rising then nan, or spiky. Only “falling” is worth looking at validation for.

For a falling training loss, the validation curve gives four branches:

• falling with it — healthy; train longer or grow the model;
• rising — overfitting;
• below the train loss — almost always dropout or BatchNorm accounting, and only then leakage;
• flat from step 0 — the val set is broken, trivial, or mislabeled.

Four terms in the tree: ln(C) is the natural logarithm of the number of classes C, the chance-level loss. LR is the learning rate. A dead ReLU is a unit whose input is negative for every example in the dataset, so it outputs zero always and its gradient is exactly zero forever. Vanishing init is the middle row of the initialization table in Initialization, where activations shrink toward zero with depth.

The table below expands every leaf of that tree, plus a few shapes the tree does not have room for. The “Mechanism” column is the part to be able to say out loud; the “First fix” column is what you do next.

Four terms in that table need defining. An LR range test is a short run in which the learning rate is swept across many orders of magnitude while watching where the loss first starts falling and where it blows up. Leakage is when information about the answer reaches the model through a feature it will not have at prediction time, or when the same rows appear in both the training and validation splits, making the validation number meaningless. A checkpoint is a saved copy of the weights at some step, so you can go back to it. GELU (Gaussian error linear unit) is a smooth alternative to ReLU that does not have a hard zero region, so it does not produce dead units.

Reading the overfitting trace

The trace below is the most common shape you will see. Watch the two number columns move apart:

epoch   train   val      note
  1     1.8241  1.7904   healthy
  3     0.9377  1.0210   val flattening
  6     0.4112  1.1802   val rising: overfitting starts here
 10     0.0914  1.6331   memorization

At epoch 1 the two curves are together. By epoch 3 train has dropped to 0.94 while val has only reached 1.02. By epoch 6 train is still falling and val has turned upward — the model is now learning things that are true of the training set and false of everything else. By epoch 10 train is at 0.09 and val is worse than it was at epoch 1.

The right checkpoint is epoch 3-4, not epoch 10, and not epoch 6 either: you want the epoch with the best val number, and val bottomed out around 1.02.

Two quantities in that trace are commonly confused, and it matters:

• The gap between the curves is the overfitting. Regularization narrows the gap.
• The level of the val curve is the performance. Capacity and data lower the level.

Confusing the two is why people add dropout to a model that is underfitting. Dropout narrows a gap; it cannot lower a level.

What this procedure assumes

The single-batch test assumes the model has enough capacity to memorize 8 examples, which every realistic network does. It also assumes your loss can reach zero. That second one is false for a task with genuinely ambiguous labels, or for a loss with an irreducible floor, and in those cases the test tells you less.

The decision tree assumes your validation split is trustworthy. If it is not, the entire right-hand side of the tree reports on the split rather than on the model — which is exactly why “flat from step 0” and “near zero from epoch 1” are branches rather than footnotes.

12. Scale

Eventually a run outgrows one device. Three techniques buy the headroom — lower-precision arithmetic, gradient accumulation, and parallelism — and each buys something more specific than “more memory”.

Mixed precision

Numbers in a computer are stored in a fixed number of bits split between two parts. The exponent sets the range of magnitudes representable. The mantissa sets the precision within that range. Spend bits on one and you have fewer for the other.

fp32 is the 32-bit standard format. fp16 and bf16 are 16-bit formats that halve memory and run much faster on modern accelerators. Mixed precision means doing the arithmetic in 16 bits while keeping a 32-bit master copy of the weights.

Why fp16 needs loss scaling

fp16 spends 5 bits on its exponent. That gives it a smallest normal value of 6.10e-5 and a largest of 65504.

Gradients in a trained network routinely sit near 1e-7. That is below fp16’s smallest normal value, so they flush toward zero and the parameter stops learning.

Loss scaling fixes it. Multiply the loss by a constant S before the backward pass; by the chain rule every gradient comes out multiplied by S too:

raw gradient          1.0e-7    ->  underflows fp16, becomes 0
scaled by S = 1024    1.02e-4   ->  representable
unscale before step   1.0e-7    ->  correct value, applied in fp32

1.0e-7 * 1024 = 1.024e-4, which is above 6.10e-5 and therefore survives. Dividing by S afterwards, in fp32, recovers the true value exactly.

Dynamic scaling automates the choice of S: raise it while steps succeed, halve it on any inf/nan, and skip the step that produced the overflow.

bf16 avoids the whole problem. It has fp32’s 8-bit exponent, so it needs no loss scaling at all. The price is mantissa bits — 7 against fp16’s 10. For training, range matters more than precision, which is the trade that made bf16 win.

What mixed precision does and does not save

Mixed precision does not shrink your optimizer state. The table below is a 1B-parameter model; the row to read is Total state, which is identical in both columns.

Follow the arithmetic. Mixed precision saves 2 GB on gradients — and then spends it again, because it must keep both a bf16 working copy of the weights (2 GB) and the fp32 master (4 GB), where fp32 needed only one 4 GB copy. Adam’s m and v stay in fp32 in both columns. Net change: zero.

The real wins are the last two rows: activations halve, and matrix multiplications run 2-8x faster. (“Matmul” is matrix multiplication, the dominant operation in training; tensor cores are the dedicated hardware units that execute it in reduced precision.)

Say that in an interview, because the common claim — “mixed precision halves memory” — is wrong about the half that usually dominates a large-model budget.

When mixed precision helps, does nothing, and hurts. It assumes the arithmetic tolerates reduced mantissa precision, which for the matrix multiplications that dominate training it does. It helps on hardware with reduced-precision units and on models whose memory is dominated by activations. It does nothing for optimizer-state memory, as the table shows, so it will not rescue a run whose problem is that m and v do not fit. It hurts with fp16 and no loss scaling, where small gradients silently become zero and the model quietly stops learning in its lower layers, and it hurts in any operation genuinely needing precision — large reductions, some loss functions — which is why those are kept in fp32 even under mixed precision.

Gradient accumulation

When the batch you want does not fit in memory, split it into micro-batches, run several of them, and add their gradients up before taking one step. Memory only ever holds one micro-batch of activations, but the gradient you finally apply is the one a large batch would have produced.

The batch you end up with is the product of three numbers:

effective_batch = micro_batch * accum_steps * num_devices
                = 8          * 16          * 4          = 512

In the loop below, ACCUM is accum_steps. Note where opt.step() sits — inside the if, so it fires once every 16 iterations, not every iteration.

opt.zero_grad(set_to_none=True)
for i, micro in enumerate(loader):
    loss = criterion(model(micro.x), micro.y)
    (loss / ACCUM).backward()          # <- the division is the whole point
    if (i + 1) % ACCUM == 0:
        unscale_(opt)                  # mixed precision: unscale BEFORE clipping
        clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        opt.zero_grad(set_to_none=True)

Why the division matters: .backward() adds into the existing gradients rather than replacing them, so after 16 micro-batches the buffer holds the sum of 16 gradients. You wanted the mean. Dividing each loss by ACCUM before its backward pass makes the sum come out as the mean.

Without / ACCUM the accumulated gradient is ACCUM times too large, which is exactly a silent 16x learning rate. The trace is unmistakable — a loss that dips, turns, and goes to nan within 40 steps:

step   loss
   0   2.3141
  10   2.2907
  20   4.7188
  30  11.8402
  40      nan

Two more gotchas.

BatchNorm does not accumulate. Its statistics are computed per micro-batch, so 16 micro-batches of 8 give you BN statistics from a batch of 8, not 512. Accumulation genuinely cannot emulate a large batch for BN. LayerNorm and GroupNorm are unaffected — GroupNorm normalizes over groups of features within a single example, so like LayerNorm it never looks across the batch.

The reduction has to match. If your loss sums over the batch rather than averaging it, / ACCUM is the wrong normalizer and you need / total_samples instead.

When accumulation helps, does nothing, and hurts

It assumes gradients from separate micro-batches can simply be added, which is true for every per-example loss and false for anything computed across the batch.

It helps whenever the batch size you want exceeds the memory you have. The only cost is wall-clock time.

It does nothing for speed. You run the same total arithmetic, just with fewer optimizer steps.

It hurts when the model contains BatchNorm, where the emulation is simply wrong, and when the loss reduction is a sum rather than a mean, where the normalizer is wrong and the effective learning rate moves with the micro-batch size.

Parallelism

Precision and accumulation buy headroom on one device. When one device is not enough, the run has to be split across several, and there are exactly three things you can split — the batch, a single weight matrix, or the stack of layers. Which one you split is decided by which dimension overflowed, not by preference, and the tree below is that decision. Throughout this subsection D is the number of devices taking part; N keeps its meaning from the vocabulary above — the number of examples in the dataset.

flowchart TD
    Q{"Does one replica fit<br/>in device memory?"} -->|yes| DDP["Data parallel<br/>replicate weights, shard the batch,<br/>all-reduce gradients once per step"]
    Q -->|no| Q2{"Which dimension<br/>does not fit?"}
    Q2 -->|"optimizer state"| ZERO["Shard states across ranks<br/>ZeRO / FSDP"]
    Q2 -->|"a single layer"| TP["Tensor parallel<br/>shard the matmul,<br/>communicate every layer"]
    Q2 -->|"layer count"| PP["Pipeline parallel<br/>shard by depth,<br/>pay the bubble"]

    style DDP fill:#2d6a4f,color:#fff
    style TP fill:#bc6c25,color:#fff
    style PP fill:#bc6c25,color:#fff

To shard is to split one logical object across several devices, each holding a piece. That word appears in three of the four leaves.

The first question is whether one replica — one complete copy of the model — fits in device memory. If it does, use data parallelism and stop. Everything below this point is a response to “it does not”.

If it does not fit, ask which dimension overflowed:

• Optimizer state. Shard the optimizer states across ranks (a rank is one participating device) using ZeRO (Zero Redundancy Optimizer) or FSDP (Fully Sharded Data Parallel). These keep only a slice of m, v and the master weights on each device and gather what they need on demand.
• A single layer. Use tensor parallelism — split the matrix multiply itself.
• The layer count. Use pipeline parallelism — give each device a consecutive run of layers.

The table compares the three across the dimensions that decide which one you can afford. The row that most often makes the decision is “Needs”.

Four terms in that table:

• DDP is distributed data parallel: every device holds a full copy of the model, each processes a different slice of the batch, and their gradients are averaged.
• An all-reduce is the collective operation that does that averaging — every device contributes a value, and every device ends up holding the total.
• NVLink is the high-bandwidth interconnect between GPUs (graphics processing units, the accelerator cards that run essentially all deep-learning training) inside a single machine. It is much faster than the network between machines, which is why “intra-node” and “inter-node” are different worlds.
• The bubble is idle time in a pipeline. While the first stage works on the first micro-batch, every later stage has nothing to do yet, and the same happens in reverse when the pipeline drains.

Two numbers worth having ready

DDP communication. A ring all-reduce moves 2*(D-1)/D times the gradient size per device, where D is the device count. Take 1B parameters in bf16, so 2 GB of gradients, across D = 8 GPUs:

2 * (7/8) * 2 GB = 3.5 GB per device per step
3.5 GB / 200 GB/s = 17 ms

Seventeen milliseconds is a large fraction of a step. That is why real implementations use gradient bucketing: start the all-reduce on each group of gradients as soon as it is ready, overlapping the communication with the rest of the backward pass rather than waiting for the whole thing to finish.

Pipeline bubble. With P stages and m micro-batches, the idle fraction is (P-1)/(m+P-1). At P = 8:

m = 8    ->  7/15  = 46.7% idle
m = 32   ->  7/39  = 17.9% idle
m = 128  ->  7/135 =  5.2% idle

Pipeline parallelism is only efficient when micro-batches greatly outnumber stages. That is why it is the last dimension you reach for, and why real systems combine all three (data x tensor x pipeline) rather than choosing one.

When each form of parallelism helps, does nothing, and hurts

Data parallelism assumes the model fits on one device and that communication can be overlapped with computation. It helps up to the point where the batch per device gets too small to saturate it. It does nothing for memory — every device still holds a full copy. It hurts when the interconnect is slow enough that the all-reduce cannot hide behind the backward pass.

Tensor parallelism assumes very fast links between the participating devices, because it communicates on every layer. Inside one machine over NVLink it helps. Across a slow network it hurts badly enough to be unusable.

Pipeline parallelism assumes many micro-batches. With m comparable to P the bubble dominates and you have bought almost nothing.

Cheat sheet

One row per failure you can see in a log or a plot. The “Mechanism” column is the sentence to say in an interview; the “Fix” column is what you actually change.

Next: 06 — Metrics — how to know whether any of this actually worked.