2. Making deep networks trainable

Depth is powerful but fragile: a signal passing through 100 layers gets multiplied 100 times, so it tends to explode or vanish, and gradients do the same on the way back. Modern networks train because five fixes work together: careful initialization, residual connections, normalization, adaptive optimizers with tuned learning-rate schedules, and numerically safe precision.

Initialization: start where signals survive

If each layer multiplies the variance of its input by more than 1, activations blow up with depth; by less than 1, they shrink to nothing. Initialization draws random weights with a variance chosen so each layer roughly preserves variance.

Scheme	Weight variance	Designed for
Xavier / Glorot	2 / (fan_in + fan_out)	tanh, sigmoid
He / Kaiming	2 / fan_in	ReLU family
GPT-2 style	std ≈ 0.02; residual output projections scaled by 1/√(2L)	deep transformers

The 1/√(2L) trick exists because each of the L blocks adds two contributions (attention and MLP) into the residual stream. Scaling them keeps the stream’s total variance from growing with depth.

Residual connections: a highway for gradients

A residual block computes x_next = x + F(x) instead of x_next = F(x). The derivative now contains an identity term, ∂x_next/∂x = I + ∂F/∂x, so gradients flow back through the “+” untouched even when F’s derivatives are tiny. ResNet (2015) introduced this to train 100+ layer image networks; deep transformers don’t train without it.

A second view: each block makes a small edit to a running representation rather than replacing it. In transformers this running vector is called the residual stream, and it is the backbone of how we understand them (Part 15).

Normalization: keep each layer’s inputs in range

Method	Normalizes over	Notes
BatchNorm	each feature, across the batch	great for CNNs; awkward with variable-length sequences, small batches and one-token-at-a-time generation
LayerNorm	the features of each token separately	the original transformer default: subtract mean, divide by std, learned gain γ and bias β
RMSNorm	the features of each token, without mean subtraction	cheaper and works as well; used in Llama-family and most modern LLMs

Normalization makes each layer’s job independent of the scale of what came before. That permits larger learning rates and steadier training. Where you place it relative to the residual connection matters (pre-norm vs post-norm, Part 5).

Optimizers: from SGD to AdamW

Plain SGD uses one learning rate for every parameter. Two improvements:

• Momentum keeps a moving average of gradients, smoothing noise and speeding progress along consistent directions.
• Adaptive scaling (RMSProp, Adam) divides each parameter’s step by a running estimate of its gradient size, so rarely updated parameters, such as embeddings of rare tokens, still move.

Adam combines both:

m̂ and v̂ are bias-corrected, because the averages start at zero. The λθ term is decoupled weight decay, the “W” in AdamW: it shrinks weights directly rather than through the gradient, and it is the transformer standard. Common LLM settings are β₁ = 0.9, β₂ = 0.95 and weight decay 0.1 (approximate).

Systems consequence: Adam keeps two extra numbers per parameter, usually in FP32. For a 7B model that is 56 GB of optimizer state alone, which is why sharded optimizers (ZeRO, FSDP; Part 8) exist. Alternatives such as Adafactor (less memory), Lion and Muon exist; AdamW remains the reference point.

Learning-rate schedules

The learning rate is the single most important hyperparameter. Too high and training diverges; too low and compute is wasted.

• Warmup: ramp up from near zero over the first few hundred to few thousand steps. Early on, Adam’s variance estimates are unreliable and the network is far from a stable region, so large steps cause divergence.
• Decay: after warmup, lower the rate; cosine decay to about 10% of peak was the standard. Low rates at the end let the model settle, and the loss visibly drops.
• Warmup–stable–decay (WSD): hold the rate constant, then decay quickly at the end. Convenient because you can branch off and decay from any checkpoint.

Also standard: gradient clipping. If the global gradient norm exceeds a threshold (often 1.0), scale it down, so one bad batch can’t wreck training.

Regularization

Regularization fights overfitting: dropout randomly zeroes activations during training, and weight decay penalizes large weights. Large pretraining runs usually see each token about once, so overfitting is rare and dropout is often off. It matters more when fine-tuning on small datasets.

Precision: how many bits per number

Format	Bits	Range	Use
FP32	32	wide	master weights, optimizer states, sensitive ops
FP16	16	narrow; overflows easily	older mixed-precision training; needs loss scaling
BF16	16	same exponent range as FP32, less precision	default for training and inference today
FP8	8	very narrow	newest hardware, with careful scaling
INT8 / INT4	8 / 4	integers	quantized inference (Part 9)

Mixed precision runs the heavy matrix multiplies in BF16, halving memory and using fast tensor cores. A master copy of the weights stays in FP32 so tiny updates aren’t rounded away.

Training instabilities

Large runs suffer loss spikes: rare bad data, attention logits growing too large, numerical overflow. Common remedies are a lower learning rate, QK-norm (normalizing queries and keys before their dot product), z-loss (penalizing large output logits), skipping bad batches and restarting from a checkpoint. Knowing these names lets you read training reports.

Questions you should now be able to answer

• Why does a residual connection prevent vanishing gradients? Write the derivative.
• Why is LayerNorm preferred over BatchNorm for autoregressive text generation?
• How much memory does AdamW’s state take for a 13B-parameter model in FP32?
• Why is warmup needed specifically with Adam-style optimizers?
• Why is BF16 safer than FP16 for training?