Training deep nets & optimizers

Module 3 gave every weight a gradient — the direction that most steeply raises the loss. The obvious move is to step the opposite way, scaled by the gradient's size. But that naive step has a fatal weakness: where the surface flattens, the gradient shrinks to almost nothing, and the step shrinks with it. Progress stalls exactly where you most need to keep moving. In the million-dimensional landscapes real networks live in, these flat saddle points vastly outnumber true minima — so the optimizer's real job is to keep moving through the flats. That is what separates a net that trains from one that sits stuck, and it is the whole subject of this chapter.

4.1 Where plain gradient descent gets stuck

Plain stochastic gradient descent (SGD) takes a step directly proportional to the negative gradient. That feels reasonable until the surface goes flat. Near a saddle point — a spot that curves down in one direction but up in another — the gradient is nearly zero, so SGD's step is nearly zero, and it can sit there for hundreds of iterations barely moving. For a long time people blamedlocal minima for training failures. The real culprit, Dauphin and colleagues showed, is these sprawling flat saddles: in high dimensions a random critical point is overwhelmingly likely to be a saddle (some directions up, some down) rather than a true bowl, because a minimum requires every direction to curve up at once, and that gets exponentially rare as dimensions grow.

There is a second reason we say stochastic: real training never sees the whole dataset at once. It draws a mini-batch — a handful of examples — and estimates the gradient from those. That estimate is noisy, which is a feature as much as a bug: the noise kicks the optimizer off razor-thin saddles and out of sharp, brittle minima, nudging it toward the flat, wide basins that generalize better. The batch size sets the trade: bigger batches give a cleaner gradient and better GPU utilization but less helpful noise; smaller batches are noisier but often generalize better.

4.2 Momentum and Adam: memory and adaptation

Momentum keeps a running velocity: each step nudges the velocity by the gradient and then moves along it. On a flat stretch the gradient is small but consistent in direction, so the velocity keeps building — a heavy ball that coasts across the saddle instead of stopping on it. The velocity is an exponentially weighted moving average of past gradients: recent steps count most, older ones fade by a factor β each step. A β of 0.9 means the ball effectively remembers about the last ten gradients — enough inertia to smooth out the jitter and roll straight through noise.

Adam attacks the same problem differently, and adapts per weight. It tracks two moving averages: the mean gradient (m, a momentum term) and the meansquared gradient (v, the gradient's typical size). Then it divides the step by that size. On a flat saddle the gradient is tiny — but so is its recorded size, and the two nearly cancel, so Adam still takes a decisive, near-full-length step. A weight whose gradient has been a steady roar gets a smaller effective step; one whose gradient has been a whisper gets a larger one. Every coordinate moves at its own right pace. (RMSProp is Adam without the momentum term — just the divide-by-size half.)

m mean gradient — the momentum termv mean gradient size — the adaptive scaleη the base learning rate

Below, all three start on the same flat shoulder of a saddle: two basins sit left and right (the dark wells), with a flat ridge between them where the gradient nearly vanishes. Press race and watch their personalities. SGD drops to the ridge then crawls, almost frozen on the flat. Momentum builds speed and coasts off into a basin. Adam steps decisively off the saddle from the very start. Drag the learning rate, and drag momentum to zero to watch the momentum runner degrade back into a stalled SGD.

The story in your hands: SGD is the tortoise that falls asleep on the flat — its loss barely drops for a long stretch because the gradient there is a whisper.Momentum and Adam both escape and dive into a basin, which is exactly why nearly every large model today is trained with one of them. Now pullmomentum to zero: the momentum runner loses its velocity and stalls right alongside SGD — proof that it was the accumulated speed, not the gradient, carrying it across. Adam keeps escaping regardless, because its trick is normalization, not memory of speed.

4.3 The learning rate — the one knob that matters most

If you could tune only one hyperparameter, tune the learning rate. Too small and training crawls — you may need ten times the compute to reach the same loss. Too large and the steps overshoot the valley and bounce up the far wall, oscillating or diverging outright. The sweet spot is a narrow band, and it shifts as training proceeds: early on you want big strides to cover ground; late on you want tiny ones to settle precisely into the minimum without rattling around it.

That is what a schedule does — it changes the learning rate over time. Step decaycuts it by a factor every so many epochs. Cosine decay glides it smoothly from high to near zero.Warmup ramps it up from tiny for the first few hundred steps — essential for Transformers, whose freshly initialized attention would otherwise blow up under a full-size step on the very first batch. Below, the surface is an ill-conditioned ravine: steep across, gentle along. Watch a single ball descend under each schedule. Push η high and the ball zig-zags violently across the ravine; too high and it flies out. A schedule lets you start bold and end gentle — the best of both.

The ravine is the whole reason schedules exist. A constant rate large enough to make progress along the gentle axis is too large for the steep one, so the ball oscillates; a rate safe for the steep axis crawls along the gentle one. Decay resolves the tension in time instead of space: big early steps race down the gentle valley, then shrinking steps stop the cross-ravine bouncing so it can settle. The little inset tracesη over the run so you can see the shape of each schedule.

4.4 Initialization — keeping the signal alive across depth

Before the first step, every weight gets a random starting value — and the scale of that randomness decides whether a deep net can train at all. A signal passes through dozens of layers; at each one it is multiplied by a weight matrix. If those weights are a touch too large, the signal's size grows a little each layer and, compounded over depth, explodes to enormous numbers. A touch too small and itvanishes toward zero. Either way the gradients on the far side are garbage, and learning never gets going. This is the exponential tyranny of depth: a per-layer factor of 1.2 or 0.8 looks harmless until you raise it to the fortieth power.

The fix, from Glorot & Bengio (Xavier) and He and colleagues (He init), is to scale the random weights so the signal's variance is preserved layer to layer: draw them with a standard deviation of roughly √(1/n) for tanh, or √(2/n) for ReLU (which zeros half its inputs, so it needs the extra factor of two to compensate), where n is the number of inputs to the layer. Below, a28-layer stack passes a unit signal through, and each bar is the typical activation size at that depth. Drag the weight scale and flip the activation. Watch the bars avalanche to nothing or blow off the top — then land on the matched init and see them hold flat all the way down.

Two lessons hide in that picture. First, the right scale depends on the activation:tanh stays healthy near √(1/n), while ReLUneeds √(2/n) because it throws away its negative half. Second, and bigger: initialization is not a formality. It is the difference between a hundred-layer network that trains and one that is dead on arrival. The same reasoning is why residual connections(module 9) matter — they give the signal a clean highway that skips the multiply entirely.

4.5 Normalization & saturation — staying in the live zone

Good initialization sets the signal right at step zero, but as weights change during training the distribution of each layer's inputs drifts — a problem Ioffe & Szegedy calledinternal covariate shift. When the drift pushes a neuron's input far from zero, a saturating activation like sigmoid or tanh flattens out: its slope goes to zero, and a zero slope means a zero gradient — the neuron stops learning. It is saturated, stuck out on the tail where nothing moves. ReLU has its own version: any neuron whose input is negative outputs zero with zero slope, and if it gets stuck there it is a dead ReLU.

Normalization is the cure. Batch norm (and its cousin layer norm, which the Transformer uses) re-centers and re-scales each layer's inputs back to roughly zero mean and unit variance before the activation sees them — dragging every neuron back into the steep, high-gradient middle of its curve. Below, the curve is your chosen activation and the dashed line is its gradient. The shaded bell is the distribution of pre-activations flowing in. Drag drift to push that distribution out toward the flat tail and watch the saturated fraction climb — then hitnormalize and see the whole population snap back to the live zone where gradients are big.

Normalization does more than rescue saturated neurons — it smooths the loss surface, which lets the optimizer take bigger, safer steps and makes training far less sensitive to the learning rate and the initialization you started from. That is why batch norm was one of the single biggest jumps in trainability the field has seen, and why essentially every modern architecture normalizes somewhere. It is also, as we'll keep seeing, more arithmetic — one more pass over every activation, every step.

4.6 Regularization — training that generalizes

A network with millions of parameters can simply memorize the training set — nailing every training example while failing on anything new. Regularization is the family of tricks that hold it back to the patterns that generalize. Weight decay (L2) adds a small penalty on large weights, nudging every weight gently toward zero each step so the net prefers simple, smooth functions.Dropout randomly switches off a fraction of neurons on each forward pass, forcing the network to spread its knowledge redundantly instead of leaning on any single unit — a kind of ensemble baked into one model. Early stopping just halts training when validation loss stops improving, before memorization sets in. Together with a good optimizer, sane init, and normalization, these are the standard kit that makes a deep net both trainable and trustworthy on data it has never seen.