Training Deep Networks in Practice

Knowing backpropagation is not the same as being able to train a deep network. Stack twenty layers naively and the loss may sit flat, explode to NaN, or drop nicely and then the model fails on validation data. The techniques in this lesson, mostly developed between 2010 and 2017, are what turned "deep learning" from a fragile research craft into a reliable engineering discipline, and every one of them is still in daily use in training everything from small classifiers to large language models. The final section gives a training-loop template you can reuse.

Intuition: a signal passing through many hands

Think of a deep network as a long chain of people passing a message. If each person whispers a little more quietly than they heard, by the twentieth person the message is inaudible; if each one shouts a little louder, it becomes a deafening roar. The same happens to activations on the forward pass and to gradients on the backward pass, because each layer multiplies the signal by its weights and activation derivatives. Training deep networks is largely about keeping that signal at a sensible volume all the way through: good initialization sets the starting volume, normalization re-levels it at every layer, residual connections add a direct line that bypasses the chain, and gradient clipping caps the occasional shout.

Weight initialization

Why scale matters

Consider a linear layer $z = \sum_{i=1}^{n_{\text{in}}} w_i x_i$ with independent, zero-mean weights and inputs. Then

$\operatorname{Var}(z) = n_{\text{in}} \operatorname{Var}(w)\operatorname{Var}(x).$

If $n_{\text{in}}\operatorname{Var}(w) > 1$, variance grows layer after layer (exploding); if it is $< 1$, it shrinks (vanishing). Through $L$ layers the factor compounds as $(n_{\text{in}}\operatorname{Var}(w))^L$. With $n_{\text{in}} = 512$ and weights drawn with standard deviation 1, the factor per layer is 512, so activation variance passes FP16's maximum within two layers and FP32's within about fifteen.

Xavier (Glorot) and He initialization

• Xavier/Glorot (Glorot and Bengio, 2010) balances the forward condition $n_{\text{in}}\operatorname{Var}(w) = 1$ and the backward condition $n_{\text{out}}\operatorname{Var}(w) = 1$ by compromise: $\operatorname{Var}(w) = \frac{2}{n_{\text{in}} + n_{\text{out}}}$. As a uniform distribution, $w \sim U[-a, a]$ with $a = \sqrt{6/(n_{\text{in}} + n_{\text{out}})}$. Derived for symmetric, roughly linear-near-zero activations such as tanh.
• He/Kaiming (He et al., 2015) accounts for ReLU zeroing half its inputs, which halves the second moment of the activations: $\operatorname{Var}(w) = \frac{2}{n_{\text{in}}}$. Use it for ReLU-family networks.

Worked numbers. A ReLU layer with $n_{\text{in}} = 512$: He standard deviation $= \sqrt{2/512} = 0.0625$. A tanh layer mapping 512 to 256 units: Xavier uniform bound $a = \sqrt{6/768} \approx 0.0884$.

Biases are usually initialized to zero. PyTorch's nn.Linear uses a Kaiming-uniform-style default scaled for general use; torch.nn.init.kaiming_normal_ and xavier_uniform_ let you choose explicitly. Transformers commonly use small normal initializations (standard deviation around 0.02) with scaled-down initialization for residual output projections, alongside normalization layers.

Vanishing and exploding gradients

By the chain rule, the gradient at layer 1 of an $L$-layer network contains a product of $L$ Jacobians, each including the weight matrix and the activation derivative. If their typical magnitude is below 1, the product vanishes exponentially; above 1, it explodes. The sigmoid's derivative is at most 0.25, so ten sigmoid layers can scale gradients by as little as $0.25^{10} \approx 9.5 \times 10^{-7}$ even before weights are considered. Early layers then barely learn, which is why sigmoid hidden layers gave way to ReLU-family activations.

Symptoms. Vanishing: loss plateaus early, gradient norms of early layers are orders of magnitude smaller than late ones. Exploding: loss spikes or becomes NaN/inf, gradient norms jump. Recurrent networks, which apply the same weights at every time step, are especially prone (see the sequences lesson).

Remedies, in order of importance: ReLU/GELU activations, proper initialization, normalization layers, residual connections, gradient clipping, and a sensible learning rate.

Normalization layers

Batch normalization

Batch normalization (Ioffe and Szegedy, 2015) normalizes each feature using statistics of the current mini-batch $\mathcal{B} = \{x_1,\dots,x_m\}$:

$\mu_\mathcal{B} = \frac1m\sum_i x_i, \quad \sigma^2_\mathcal{B} = \frac1m\sum_i (x_i - \mu_\mathcal{B})^2, \quad \hat x_i = \frac{x_i - \mu_\mathcal{B}}{\sqrt{\sigma^2_\mathcal{B} + \epsilon}}, \quad y_i = \gamma \hat x_i + \beta,$

with learnable scale $\gamma$ and shift $\beta$ (so the network can undo the normalization if useful) and a small $\epsilon$ for numerical stability. During training it also keeps running averages of $\mu$ and $\sigma^2$, which are used at inference instead of batch statistics.

Worked example. One feature across a batch of four: $(1, 2, 3, 6)$. $\mu = 3$. $\sigma^2 = (4 + 1 + 0 + 9)/4 = 3.5$, $\sigma \approx 1.871$. Normalized (ignoring $\epsilon$): $(-1.069, -0.535, 0, 1.604)$. With $\gamma = 2$, $\beta = 0.5$: $(-1.638, -0.569, 0.5, 3.707)$.

Batch norm allows higher learning rates and made very deep CNNs practical. The original paper motivated it by reducing "internal covariate shift"; later work argues its main benefit is smoothing the optimization landscape, and the precise mechanism is still debated. Its weaknesses come from depending on the batch: it behaves poorly with small batches, differs between training and inference modes (forgetting model.eval() is a classic bug), and is awkward for variable-length sequences.

Layer normalization

Layer normalization (Ba, Kiros and Hinton, 2016) computes $\mu$ and $\sigma^2$ over the features of a single example instead of over the batch, then applies the same $\gamma,\beta$ affine step. It behaves identically in training and inference and is independent of batch size, which is why it is standard in Transformers and RNNs. RMSNorm, a simplification that skips mean-centering and divides by the root-mean-square, is common in recent large language models. Rule of thumb: batch norm for CNNs with reasonable batch sizes, layer norm (or RMSNorm) for Transformers and sequence models.

Residual connections

A residual block (He et al., 2015, ResNet) computes $y = x + F(x)$, where $F$ is a small stack of layers. The layers learn a correction to the identity rather than a whole new mapping. Two effects follow. First, if extra layers are not useful they can learn $F \approx 0$, so adding depth should not hurt. Second, the backward pass gets a direct path: $\frac{\partial y}{\partial x} = I + \frac{\partial F}{\partial x}$, so gradients can flow through the identity term even if $\partial F/\partial x$ is small. Residual connections enabled networks with over 100 layers and are in every Transformer block. Modern Transformers typically use pre-norm blocks, $y = x + F(\text{Norm}(x))$, which tend to train more stably than placing the normalization after the addition.

Regularization: dropout and weight decay

Dropout (Srivastava et al., 2014) randomly zeros each activation with probability $p$ during training. Implementations use inverted dropout, scaling surviving activations by $\frac{1}{1-p}$ so the expected value is unchanged and nothing needs rescaling at test time. With $p = 0.5$, kept activations are doubled. Dropout discourages co-adapted features and acts like an ensemble of thinned networks. It is switched off by model.eval(). Typical rates: 0.1 to 0.5 in MLPs; around 0.1 or none in large Transformers trained on huge datasets, where overfitting is less of a concern.

Weight decay shrinks weights toward zero each step, discouraging large weights, discussed with AdamW below. Other regularizers include data augmentation (often the most effective for vision and audio), label smoothing and early stopping.

Optimizers in practice: AdamW

Adam keeps exponential moving averages of the gradient $g_t$ and its square, with bias correction for their zero initialization:

$m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t, \qquad v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2,$

$\hat m_t = \frac{m_t}{1-\beta_1^t}, \qquad \hat v_t = \frac{v_t}{1-\beta_2^t}, \qquad \theta_t = \theta_{t-1} - \eta\left(\frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon} + \lambda\,\theta_{t-1}\right).$

Here $\eta$ is the learning rate, $\beta_1, \beta_2$ are decay rates (defaults 0.9 and 0.999), $\epsilon$ a small constant, and $\lambda$ the weight-decay coefficient. The $\lambda\theta$ term written separately like this is the decoupled weight decay of AdamW (Loshchilov and Hutter). In plain Adam, weight decay is implemented by adding $\lambda\theta$ to the gradient, which then gets divided by $\sqrt{\hat v}$, so parameters with large gradient history are regularized less than intended. AdamW applies decay directly to the weights, which generalizes better and makes $\lambda$ easier to tune; it is the default optimizer for Transformers.

Worked step. On step 1, $\hat m_1 = g_1$ and $\hat v_1 = g_1^2$, so the Adam part of the update is $\eta\, g_1/(|g_1| + \epsilon) \approx \eta\,\text{sign}(g_1)$. With $\eta = 10^{-3}$, a parameter whose gradient is 0.02 and one whose gradient is 50 both move by about $10^{-3}$. Adam's steps are roughly scale-free per parameter, which is why one learning rate works across layers with very different gradient magnitudes.

Practical defaults: AdamW with $\eta$ around $10^{-3}$ for small models and $10^{-4}$ to $3\times10^{-4}$ for Transformers, $\lambda$ around 0.01 to 0.1, and no weight decay on biases and normalization parameters. SGD with momentum (0.9) and a tuned schedule remains competitive for CNNs and sometimes generalizes slightly better.

Learning-rate schedules

The learning rate is the single most important hyperparameter. Schedules change it during training:

• Warmup: increase linearly from near zero over the first few hundred to few thousand steps. Early on, Adam's second-moment estimates are noisy and the weights are random, so large steps can destabilize training; warmup is essential for Transformers.
• Cosine decay: after warmup, follow $\eta_t = \eta_{\min} + \tfrac12(\eta_{\max} - \eta_{\min})\left(1 + \cos\left(\pi\,\frac{t - T_w}{T - T_w}\right)\right),$ where $T_w$ is the number of warmup steps and $T$ the total.
• Step decay (divide by 10 at fixed epochs), linear decay, one-cycle, and reduce-on-plateau are common alternatives.

Worked example. $\eta_{\max} = 3\times10^{-4}$, $\eta_{\min} = 0$, $T_w = 1000$, $T = 10{,}000$. At step 500 (warmup): $\eta = 3\times10^{-4} \times 500/1000 = 1.5\times10^{-4}$. At step 3,250: progress $= 2250/9000 = 0.25$, so $\eta = 0.5 \times 3\times10^{-4} \times (1 + \cos(\pi/4)) \approx 0.5 \times 3\times10^{-4} \times 1.707 \approx 2.56\times10^{-4}$. At step 5,500: progress 0.5, $\eta = 1.5\times10^{-4}$. At step 10,000: $\eta = 0$.

To find a starting $\eta_{\max}$, run a short learning-rate range test: increase $\eta$ exponentially over a few hundred steps and pick a value somewhat below where the loss starts to rise sharply.

Early stopping

Track validation loss (or the metric you care about) every epoch, save a checkpoint whenever it improves, and stop when it has not improved for a set number of epochs (the patience, e.g. 5 to 10). Then restore the best checkpoint. Early stopping is a cheap, effective regularizer that limits how far the weights move from initialization. Use a validation set separate from the final test set.

Mixed precision: FP16 and BF16

Training in 16-bit floating point halves memory for activations and uses the much faster low-precision units on modern accelerators.

Mixed precision (Micikevicius et al., 2017) runs most matrix multiplications in 16-bit while keeping an FP32 master copy of the weights and doing sensitive operations (reductions, softmax, loss) in FP32. In PyTorch, torch.autocast chooses the precision per operation automatically.

With FP16, gradients smaller than about $6\times10^{-8}$ (the smallest subnormal) flush to zero, so you need loss scaling: multiply the loss by a large factor $S$ before backward(), so gradients are scaled into representable range, then divide by $S$ before the optimizer step. A dynamic scaler (GradScaler) raises $S$ while things are fine and lowers it, skipping the step, when it detects inf/NaN. BF16 has FP32's exponent range, so loss scaling is normally unnecessary; it is the default choice on hardware that supports it (NVIDIA GPUs from the Ampere generation onward, Google TPUs and others). Precision loss in BF16 (8 significant bits, about 2 to 3 significant decimal digits) is usually harmless for training because the master weights and accumulations stay in FP32.

Gradient clipping

Clip by global norm: if $\lVert g\rVert_2 > c$, rescale $g \leftarrow g \cdot \frac{c}{\lVert g\rVert_2}$. This preserves the gradient's direction and caps its size, preventing a single bad batch from destroying training. Example: $g = (3, 4, 12)$ has norm 13; with $c = 1$ it becomes $(0.231, 0.308, 0.923)$. A max norm of 1.0 is a common default for Transformers. With FP16 loss scaling, unscale gradients before clipping, or the threshold is applied to scaled values. Logging the pre-clip gradient norm is one of the best health signals you have: steady growth or frequent spikes indicate instability.

Debugging training

A systematic order of checks, following the spirit of Karpathy's widely cited "recipe":

1. Check the initial loss. With $K$ balanced classes and a properly initialized softmax classifier, the first loss should be about $\ln K$ (2.303 for 10 classes). Much larger means the output layer is badly scaled or labels are wrong.
2. Overfit a single batch. Take one batch of, say, 32 examples and train on it repeatedly. A correct model and pipeline drives loss to near zero within a few hundred steps. If it cannot, there is a bug (wrong labels, detached graph, missing zero_grad(), frozen parameters, a learning rate far off), and no amount of data will fix it.
3. Visualize inputs right before they enter the model, after augmentation and preprocessing. Many bugs are in the data.
4. Read the loss curves (training and validation on the same plot):
   • Training loss flat from the start: learning rate too low, vanishing gradients, or a bug.
   • Training loss explodes or goes NaN: learning rate too high, no clipping, FP16 overflow, or a log of zero.
   • Training loss falls, validation loss rises after a point: overfitting; add regularization, data or augmentation, or stop earlier.
   • Both high and close together: underfitting; increase capacity or train longer.
   • Noisy, spiky training loss: batch too small or learning rate too high.
5. Change one thing at a time and keep a log of experiments (see the MLOps lessons).

A PyTorch training-loop template

import math, torch
from torch import nn

def make_optimizer(model, lr=3e-4, wd=0.05):
    decay, no_decay = [], []
    for name, p in model.named_parameters():
        if p.requires_grad:
            (no_decay if p.ndim < 2 else decay).append(p)   # no decay on biases/norms
    return torch.optim.AdamW([{"params": decay, "weight_decay": wd},
                              {"params": no_decay, "weight_decay": 0.0}], lr=lr)

def warmup_cosine(opt, warmup, total):
    def f(step):
        if step < warmup:
            return (step + 1) / warmup
        progress = (step - warmup) / max(1, total - warmup)
        return 0.5 * (1 + math.cos(math.pi * min(1.0, progress)))
    return torch.optim.lr_scheduler.LambdaLR(opt, f)

def train(model, train_loader, val_loader, loss_fn, epochs=20, patience=5,
          device="cuda", amp_dtype=torch.bfloat16, max_norm=1.0):
    model.to(device)
    opt = make_optimizer(model)
    sched = warmup_cosine(opt, warmup=500, total=epochs * len(train_loader))
    use_scaler = amp_dtype == torch.float16          # loss scaling only for FP16
    scaler = torch.amp.GradScaler(device, enabled=use_scaler)  # older: torch.cuda.amp
    best, bad_epochs = float("inf"), 0
    for epoch in range(epochs):
        model.train()
        for xb, yb in train_loader:
            xb, yb = xb.to(device), yb.to(device)
            with torch.autocast(device_type=device, dtype=amp_dtype):
                loss = loss_fn(model(xb), yb)
            opt.zero_grad(set_to_none=True)
            scaler.scale(loss).backward()
            scaler.unscale_(opt)                          # clip true gradients
            nn.utils.clip_grad_norm_(model.parameters(), max_norm)
            scaler.step(opt); scaler.update(); sched.step()
        model.eval(); total, n = 0.0, 0
        with torch.no_grad(), torch.autocast(device_type=device, dtype=amp_dtype):
            for xb, yb in val_loader:
                xb, yb = xb.to(device), yb.to(device)
                total += loss_fn(model(xb), yb).item() * len(xb); n += len(xb)
        val = total / n
        print(f"epoch {epoch} val_loss {val:.4f} lr {sched.get_last_lr()[0]:.2e}")
        if val < best:
            best, bad_epochs = val, 0
            torch.save(model.state_dict(), "best.pt")     # checkpoint the best model
        elif (bad_epochs := bad_epochs + 1) >= patience:
            break                                          # early stopping
    model.load_state_dict(torch.load("best.pt"))
    return model

What each piece does: parameter groups exclude biases and normalization weights from decay; LambdaLR multiplies the base learning rate by the warmup-cosine factor every step; autocast runs the forward pass in BF16 or FP16; the scaler is a pass-through when disabled (BF16) and handles loss scaling for FP16; gradients are unscaled before clipping; model.eval() and torch.no_grad() switch off dropout, use batch-norm running statistics and skip graph building during validation; the best checkpoint is saved and restored. On a CPU, set device="cpu" (CPU autocast supports BF16). The GradScaler constructor taking a device argument is the current API; older PyTorch versions use torch.cuda.amp.GradScaler().

Common pitfalls

• Forgetting model.train() / model.eval(), so dropout stays on at evaluation or batch norm uses batch statistics at inference.
• No warmup for Transformers, or a peak learning rate copied from a different model size or batch size.
• Using Adam with L2 regularization when you meant AdamW decoupled decay; and decaying biases and norm parameters.
• FP16 without loss scaling, or clipping before unscaling.
• Batch norm with tiny batches (for example 2 per GPU): statistics are too noisy; use group norm or layer norm, or synchronized batch norm.
• Judging by training loss alone; always plot validation loss on the same axes.
• Tuning many things at once without logging, making it impossible to know what helped.

In practice

• Default recipe for a new model: He or framework-default initialization, normalization layers, residual connections if deeper than a few layers, AdamW with warmup plus cosine decay, gradient clipping at 1.0, BF16 autocast where supported, early stopping on validation loss.
• Always run the overfit-one-batch test before a long run; it takes a minute and catches most bugs.
• Log training loss, validation loss, learning rate and gradient norm every few steps to an experiment tracker.
• When scaling batch size up, scale the learning rate up too (roughly linearly for SGD, more gently for Adam) and lengthen warmup, then re-verify.
• Fine-tuning pretrained models usually needs a learning rate 10 to 100 times smaller than training from scratch, short warmup, and few epochs.