Training Deep Networks

What separates a model that trains from one that doesn't — and how to make modern deep networks stable, fast, and generalizing well.

Table of contents

1. Challenges in training deep networks
2. Batch normalization
3. Layer normalization and alternatives
4. Dropout
5. Weight decay (L2 regularization in DL)
6. Data augmentation
7. Optimizers (DL perspective)
8. Learning rate schedules
9. Gradient clipping
10. Early stopping
11. Mixed precision training
12. Practical training checklist

1. Challenges in training deep networks

2. Batch normalization

2.1 Internal covariate shift

As weights change during training, the distribution of inputs to each layer shifts. Deeper layers face a non-stationary input distribution → must continuously re-adapt → slow training.

Ioffe & Szegedy (2015) proposed normalizing layer inputs to zero mean and unit variance within each mini-batch.

2.2 BatchNorm algorithm

For a mini-batch $\mathcal{B} = \{z_1, \ldots, z_B\}$ of pre-activations for one feature:

Step 1 — Normalize:

Step 2 — Scale and shift (learnable):

$\gamma$ (scale) and $\beta$ (shift) are learned parameters per feature. This allows the network to undo normalization if needed.

2.3 Backpropagation through BatchNorm

The gradient flows through the normalization:

Gradients w.r.t. parameters: $\partial\mathcal{L}/\partial\gamma = \sum_i \partial\mathcal{L}/\partial y_i \cdot \hat{z}_i$, $\partial\mathcal{L}/\partial\beta = \sum_i \partial\mathcal{L}/\partial y_i$.

2.4 Inference with BatchNorm

At test time, there is no mini-batch — use running statistics accumulated during training:

Running stats: $\mu_\text{running} \leftarrow (1-m)\mu_\text{running} + m\mu_\mathcal{B}$ (momentum $m \approx 0.1$).

2.5 Why BatchNorm works (multiple hypotheses)

1. Reduces internal covariate shift (original hypothesis, debated).
2. Smooths the loss landscape: reduces the Lipschitz constant of the loss and gradients, making training more stable.
3. Implicit regularization: mini-batch noise adds stochasticity similar to dropout.
4. Allows higher learning rates (less sensitivity to initialization).

2.6 BatchNorm placement

Standard: before the activation function (original paper), but after is also common and sometimes better in practice. Experiment based on task.

Option 1: Linear → BatchNorm → ReLU   (original)
Option 2: Linear → ReLU → BatchNorm   (used in PreResNets)

2.7 Limitations

• Small batch sizes: statistics are noisy → unstable. Batch size ≥ 16–32 typically needed.
• Variable-length sequences: batch statistics are complicated → prefer LayerNorm.
• Inconsistency between train and test if running stats are poorly calibrated.

3. Layer normalization and alternatives

3.1 Layer Normalization (Ba et al., 2016)

Normalize across features (not batch), so one sample at a time:

Key difference from BatchNorm: statistics are computed per sample over features. Works the same at train and test time. No batch dependency.

Used in: Transformers, RNNs, any setting where batch size is small or variable-length.

3.2 Instance Normalization

Normalize per sample and per channel. Removes style information → used in style transfer.

3.3 Group Normalization

Divide channels into groups; normalize within each group per sample. Works at batch size = 1. Used in detection/segmentation models.

3.4 RMS Norm (Root Mean Square Layer Norm)

No centering (no mean subtraction). Simpler, faster, used in LLaMA, T5.

4. Dropout

4.1 Algorithm (Srivastava et al., 2014)

During training, independently zero each activation with probability $p$ (drop rate):

The $1/(1-p)$ inverted dropout scaling ensures $\mathbb{E}[\tilde{a}_j] = a_j$ → at test time, use activations unchanged (no scaling needed).

4.2 Theoretical justifications

Ensemble interpretation: with $N$ neurons and drop rate $p$, there are $2^N$ possible sub-networks. Dropout trains all of them simultaneously with shared weights, and at test time approximates their geometric ensemble.

Noise injection: adds Gaussian-like multiplicative noise, preventing co-adaptation of neurons.

Weight uncertainty: interpretable as approximate Bayesian inference with a Bernoulli approximate posterior (Gal & Ghahramani, 2016).

4.3 Practical notes

• Where: applied to the outputs of hidden layers (before activation or after — convention varies). Rarely on first input or output layer.
• Rate: $p = 0.5$ for fully connected layers (standard). $p = 0.1$–$0.2$ for convolutional layers.
• During training: activate stochastic dropping. During evaluation/inference: model.eval() disables dropout.
• Larger dropout → more regularization → more training needed.

import torch.nn as nn

model = nn.Sequential(
    nn.Linear(512, 256),
    nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(256, 10),
)
model.train()   # enables dropout
model.eval()    # disables dropout (test time)

4.4 Spatial dropout (DropBlock)

For CNNs: standard dropout drops individual activations independently. Better: drop spatially contiguous blocks (DropBlock), forcing the network to use distributed representations.

5. Weight decay (L2 regularization in DL)

L2 penalty on weights:

Gradient update with weight decay:

The $(1-\eta\lambda)$ factor decays weights toward zero each step.

Important: AdamW vs Adam with L2. Standard Adam incorporates the L2 gradient into the adaptive moment estimates, which changes its effect. AdamW applies weight decay separately (directly to weights, not through gradient), correctly implementing L2 regularization with Adam.

6. Data augmentation

6.1 Principle

Augmentation generates additional training examples by applying label-preserving transformations, effectively increasing dataset size and reducing variance.

6.2 Image augmentation

6.3 MixUp derivation

Encourages linear behavior between training examples. Shown to improve calibration and reduce memorization.

6.4 Test-Time Augmentation (TTA)

Apply augmentations at test time; average predictions. Improves accuracy at inference cost.

7. Optimizers (DL perspective)

Refer to Note 04 for full derivations. Summary of what to use in DL:

7.1 SGD with momentum

Needs careful LR tuning but often achieves best final generalization with proper schedules (especially in vision tasks).

7.2 Adam / AdamW

Default for most DL tasks. See Note 04 for full equations.

AdamW (Loshchilov & Hutter, 2019):

The $-\eta\lambda\theta_t$ term is weight decay applied directly — decoupled from the gradient step.

7.3 Lion (Learning by Infinite-Norm)

Uses sign of update → memory efficient (1 state vector vs 2 for Adam). Competitive on transformers.

7.4 SGD vs Adam comparison

8. Learning rate schedules

8.1 Warmup

Ramp LR from near 0 to target over first $W$ steps:

Why: in early training, gradients are noisy and large. A large LR can cause early divergence. Warmup allows Adam's moment estimates to stabilize before taking large steps. Essential for transformers.

8.2 Cosine annealing

Smooth decay. Often paired with warmup: warmup then cosine.

8.3 Cosine with warm restarts (SGDR)

Restart cosine schedule periodically (with growing period):

Each restart can escape shallow local minima. Ensemble of models at restart points is effective.

8.4 One-cycle policy

Howard & Gugger (2018): Increase LR from base to max (1/4 of training), then decrease to near zero (3/4). Also vary momentum inversely with LR.

Often allows training in 1/10th the epochs vs standard training. torch.optim.lr_scheduler.OneCycleLR.

8.5 Warmup + linear/cosine decay (LLM standard)

Standard schedule for training BERT, GPT, T5 and similar.

9. Gradient clipping

Prevent exploding gradients by capping gradient norm:

Direction unchanged, magnitude capped.

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Typical: max_norm = 1.0 for transformers, max_norm = 5.0 for RNNs.

Value clipping (alternative): clip each gradient component to $[-c, c]$. Simpler but changes direction.

10. Early stopping

Monitor validation loss; stop when it stops improving for $p$ consecutive epochs (patience):

best_val_loss = float('inf')
patience_count = 0
for epoch in range(max_epochs):
    train_one_epoch()
    val_loss = evaluate()
    if val_loss < best_val_loss - delta:
        best_val_loss = val_loss
        save_checkpoint(model)
        patience_count = 0
    else:
        patience_count += 1
    if patience_count >= patience:
        load_checkpoint(model)   # restore best
        break

Equivalence to L2 regularization: for gradient flow in linear networks, early stopping is equivalent to Ridge with $\lambda \propto 1/T$ (number of training steps).

11. Mixed precision training

11.1 FP16 / BF16

Store weights and activations in 16-bit floats during forward/backward pass:

• FP16: 1 sign, 5 exponent, 10 mantissa bits. Range $\approx 6\times10^{-5}$ to $65504$.
• BF16: 1 sign, 8 exponent, 7 mantissa bits. Same range as FP32, less precision.

Benefits: 2× memory, 2–8× throughput on modern GPUs (NVIDIA Tensor Cores).

11.2 Loss scaling

FP16 underflows for small gradients. Multiply loss by a large scalar $S$ before backward pass, then divide gradients by $S$ before weight update:

Loss scaler adjusts $S$ dynamically (increase if no overflow, decrease on overflow).

11.3 Master weights

Keep FP32 master weights for the optimizer step. Cast to FP16 for forward/backward. Prevents precision loss in parameter updates (small updates to large weights).

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for X, y in dataloader:
    optimizer.zero_grad()
    with autocast():              # FP16/BF16 forward pass
        loss = model(X, y)
    scaler.scale(loss).backward() # scaled backward
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)        # FP32 update
    scaler.update()

12. Practical training checklist

Setup:
  [ ] Set all random seeds (torch, numpy, random, CUDA)
  [ ] Use GPU (model.to(device), data.to(device))
  [ ] Mixed precision with GradScaler

Sanity checks before full training:
  [ ] Overfit a small batch (loss should reach near 0)
  [ ] Loss at init: cross-entropy should = log(n_classes); MSE should = label variance
  [ ] Gradient check: verify no NaN/Inf in gradients after first step

Architecture:
  [ ] BatchNorm / LayerNorm after linear/conv layers
  [ ] Appropriate activation (ReLU/GELU for hidden, softmax/sigmoid for output)
  [ ] Correct initialization (He for ReLU, Xavier for tanh/sigmoid)

Training loop:
  [ ] model.train() for training, model.eval() for validation
  [ ] optimizer.zero_grad() before each backward
  [ ] Gradient clipping (max_norm=1.0 for transformers)
  [ ] LR scheduler step at right point (per step vs per epoch)

Monitoring:
  [ ] Log train/val loss per epoch
  [ ] Log learning rate
  [ ] Check for NaN (torch.isnan(loss).any())
  [ ] Save best checkpoint based on val metric
  [ ] Early stopping with patience

*File: notes/12_training_deep_networks.md — next: notes/13_cnns.md*