Training Deep Networks

Train a 3-layer MLP on MNIST with random init and SGD: 98% test accuracy in 30 seconds. Train a 30-layer MLP with the same setup: loss stuck at $\log 10 \approx 2.30$, accuracy at chance level. Same model class, same optimizer, same data, ten times more parameters. What changed?

The signal that has to propagate is the gradient at the loss with respect to the first layer's weights. By the chain rule, that gradient is the product of Jacobians from layer 30 back to layer 1 — each Jacobian is the matrix of partial derivatives of one layer's outputs with respect to its inputs, the matrix generalization of the scalar local derivatives you chained in another chapter. A Jacobian stretches some directions and shrinks others; its singular values are those per-direction stretch factors, and the set of them is its spectrum. If each Jacobian has a typical singular value $\lambda$, the gradient norm shrinks (or grows) by $\lambda^{30}$. For $\lambda = 0.9$, that's $\lambda^{30} \approx 0.04$. For $\lambda = 1.1$, that's $\lambda^{30} \approx 17$. Either way, the first-layer gradient is wrong by orders of magnitude relative to the last-layer gradient. The optimizer sees a gradient signal that is dominated by the layers nearest the output, and the early layers barely move.

This is the vanishing/exploding gradients problem (Bengio et al. 1994). It is not philosophical. It is the product of a bunch of singular values, and you can compute it. The whole edifice of modern deep learning — better activations, better init, normalization, residuals — is a stack of fixes for that one product.

There are five mechanical fixes, and each gets its own section below:

1. Activations that do not saturate (ReLU and descendants).
2. Initialization that keeps the per-layer Jacobian's expected singular values near 1 (Xavier, Kaiming, μP).
3. Normalization that re-centers the activation distribution at every layer (BatchNorm, LayerNorm, RMSNorm).
4. Residual connections that give the gradient a shortcut path independent of the layer's Jacobian (ResNet, transformer blocks).
5. Adaptive optimizers that rescale the update per-parameter so slow-moving directions still move (Adam, AdamW, Muon).

Stack all five and you can train hundreds of layers. Drop any one and you usually cannot.

A randomly-initialized linear layer maps input variance to output variance. If the weights are drawn from a normal with variance $\sigma^2$, then for input $x \in \mathbb{R}^{d_{in}}$ with components of unit variance, the output $y = Wx$ has components of variance $d_{in} \sigma^2$. To preserve variance across the layer, set $\sigma^2 = 1/d_{in}$.

That is the core of every init scheme. The variants differ in two ways: whether they account for fan-out and the non-linearity.

Xavier / Glorot initialization (Glorot and Bengio 2010) targets unit variance for both forward and backward passes. The compromise: $\sigma^2 = 2/(d_{in} + d_{out})$. It assumes the activation has unit derivative near zero, which is true for tanh but not for ReLU. Use it with tanh or sigmoid.

Kaiming / He initialization (He et al. 2015) targets unit variance only in the forward direction, and accounts for ReLU killing half the activations on average. The fix is to scale up by $\sqrt{2}$: $\sigma^2 = 2/d_{in}$. Use it with ReLU and its variants (GELU, SiLU, SwiGLU).

μP — maximal update parameterization (Yang and Hu 2021) is the modern frontier. The insight is that the right per-parameter learning-rate scaling depends on the layer's fan-in and fan-out, in a way that lets you tune hyperparameters at small width and transfer them to large width without retuning. For Llama-scale training this is the difference between "we ran 50 small ablations to find the right LR for the 70B model" and "we ran 5 small ablations and the LR transferred". μP is more involved to implement, and unless you are training at frontier scale, Kaiming is fine.

PyTorch's defaults: nn.Linear uses Kaiming uniform with a=sqrt(5), which is roughly Kaiming-for-LeakyReLU. It's not optimal but it does not actively break anything. Override it for serious work — nn.init.kaiming_normal_ against the explicit sqrt(2/fan_in) rule:

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
        if m.bias is not None:
            nn.init.zeros_(m.bias)
model.apply(init_weights)

def init_kaiming_(model):
    for m in model.modules():
        if isinstance(m, nn.Linear):
            nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
            if m.bias is not None:
                nn.init.zeros_(m.bias)

# the variance the call computes, spelled out (draft.md §2):
#   if init == "kaiming":  sigma = np.sqrt(2.0 / d_in)
#   elif init == "xavier": sigma = np.sqrt(2.0 / (d_in + d_out))
#   W = rng.standard_normal((d_in, d_out)) * sigma; b = np.zeros(d_out)

A normalization layer takes a tensor, computes per-axis statistics, and rescales the tensor to have a stable distribution. The three flavors differ in which axis they normalize over.

BatchNorm (Ioffe and Szegedy 2015) normalizes across the batch axis. For a hidden activation $h \in \mathbb{R}^{n \times d}$ (batch of $n$ examples, dimension $d$), it computes $\mu, \sigma$ per dimension $d$ across the $n$ examples in the batch, then $\hat{h} = \gamma (h - \mu) / \sigma + \beta$ with learned $\gamma, \beta$. At inference, it uses running statistics (exponential moving average of $\mu, \sigma$ from training).

BatchNorm is great for CNNs on images: batches are large, batch statistics are stable, and the inductive bias matches. It is bad for transformers and RNNs — the sequence models, built from scratch in another chapter, that process a batch of token sequences where each token (each sequence position) carries its own feature vector: batch sizes can be small, sequence positions are variable, and you do not want positions to "see" each other through the normalization. It also has a training-vs-inference distribution mismatch (different statistics), and it breaks under distributed training without sync (each device uses its local statistics).

LayerNorm (Ba et al. 2016) normalizes across the feature axis, per example. Same formula, but $\mu, \sigma$ are computed per row of $h$ across the $d$ dimensions. No running stats; the same operation at train and inference time. Robust to batch size, robust to sequence length, plays well with distributed training. This is what every transformer uses.

RMSNorm (Zhang and Sennrich 2019) is LayerNorm without subtracting the mean: $\hat{h} = \gamma \cdot h / \sqrt{\text{mean}(h^2) + \epsilon}$. The math is that LayerNorm's mean-centering is mostly redundant with the learned $\beta$; dropping it costs nothing in accuracy and saves a kernel pass. Llama, Mistral, and most 2023+ LLMs use RMSNorm.

rmsn = nn.RMSNorm(normalized_shape=d, eps=1e-5)
y = rmsn(x)  # normalizes over the last dim, learnable gamma scale, no mean subtraction

class RMSNorm(nn.Module):
    def __init__(self, d: int, eps: float = 1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        rms = torch.sqrt((x * x).mean(dim=-1, keepdim=True) + self.eps)
        return self.gamma * x / rms

Where to put it. For a transformer block (Attn is the attention sub-layer, the part of the block you build in another chapter): pre-norm (x = x + Attn(LayerNorm(x))) is the modern default. Post-norm (x = LayerNorm(x + Attn(x))) was the original transformer; it needs careful warmup or it diverges. Pre-norm trains stably; post-norm gives slightly better final loss in some studies, but the engineering price is rarely worth it.

For a CNN: Conv → BN → ReLU is the standard pattern. Putting BN before the activation (rather than after) is empirically better and removes the need for bias in the conv layer.

A residual connection (also called a skip or shortcut connection — the "skip connections" from §1) wraps a sub-layer in an identity skip: $y = x + F(x)$ where $F$ is the sub-layer's computation. He et al. (2015, ResNet) introduced them to enable training of 100+ layer CNNs. The mechanical reason they work: the gradient of $y = x + F(x)$ with respect to $x$ is $1 + F'(x)$, so even if $F'(x) \to 0$, the gradient through the residual is $\geq 1$. The signal cannot vanish through a residual.

That is the training story, and it is sufficient for getting deep networks to converge. another chapter then leverages it as the foundation of mechanistic interpretability.

For a transformer block with pre-norm, the structure that becomes the centrepiece of another chapter:

class Block(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = FeedForward(d_model)
    def forward(self, x):
        x = x + self.attn(self.ln1(x))    # read, write
        x = x + self.ffn(self.ln2(x))     # read, write
        return x

The shape never changes. The same tensor flows through 12 (or 24, or 96) of these blocks and accumulates corrections at each one.

Three classic regularizers. Treating them as independent knobs is a mistake; they interact.

Dropout (Srivastava et al. 2014) randomly zeros a fraction $p$ of activations during training. Modern implementations (PyTorch's nn.Dropout, d2l) use inverted dropout: during training, each surviving activation $h$ is rescaled to $h/(1-p)$ so that the expected value is unchanged ($E[h'] = h$), and at inference dropout does nothing — no nodes are dropped and no rescaling is applied. Standard rates: $p=0.1$ for transformers, $p=0.2$ to $0.5$ for MLPs and convnets. The mental model: dropout trains an exponential ensemble of subnets, then averages them at inference.

self.dropout = nn.Dropout(p=0.1)
def forward(self, x):
    return self.dropout(self.fc(x))

Weight decay is L2 regularization on the weights, applied during the optimizer step. For SGD, the gradient of $\frac{\lambda}{2}\|W\|^2$ is $\lambda W$, so the update is $W \leftarrow W - \eta (g + \lambda W) = (1 - \eta \lambda) W - \eta g$. This shrinks weights toward zero at each step.

For Adam, the original implementation mixed weight decay into the gradient before the adaptive rescaling, which made it effectively weaker on parameters with large second moments. AdamW (Loshchilov and Hutter 2019) decouples weight decay from the gradient step. For LLM training, AdamW with weight_decay=0.1 is the canonical setting.

Crucial detail: weight decay is typically applied to weights but not biases or normalization parameters. ARENA's notes are explicit on this: "weight decay is often not applied to embeddings and layernorms in transformer models." The canonical pattern uses parameter groups:

def configure_optimizer(model, weight_decay: float = 0.1, lr: float = 3e-4):
    decay_params, nodecay_params = [], []
    for n, p in model.named_parameters():
        if p.dim() < 2 or "bias" in n or "norm" in n.lower() or "embed" in n.lower():
            nodecay_params.append(p)
        else:
            decay_params.append(p)
    return torch.optim.AdamW([
        {"params": decay_params, "weight_decay": weight_decay},
        {"params": nodecay_params, "weight_decay": 0.0},
    ], lr=lr, betas=(0.9, 0.95))

Label smoothing (Szegedy et al. 2016) replaces the one-hot target with a soft target: $0.9$ for the true class, $0.1 / (K-1)$ for the others. The effect is to discourage the model from producing overly-confident probabilities, which helps calibration and a bit of generalization.

loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)

How they interact:

• Dropout and weight decay both regularize; using them at full strength simultaneously typically over-regularizes. Standard combination: $p=0.1$ dropout + weight_decay=0.1 for transformers.
• Label smoothing + weight decay + dropout is the Imagenet competition stack. Each one buys 0.5-1% test accuracy on its own; combined, maybe 1.5-2% (sub-additive).
• For tiny models on small data, prefer weight decay over dropout (less compute overhead, more interpretable).
• For large models with abundant data, you can often drop all three and rely on data alone. Llama-3 uses no dropout.

The optimizer is the rule that turns gradients into parameter updates. The frontier has moved several times. The five that matter today:

SGD with momentum. The update is $v_t = \mu v_{t-1} + g_t$, $\theta_t = \theta_{t-1} - \eta v_t$. The momentum term $\mu \approx 0.9$ accumulates past gradients. Geometrically, it lets the optimizer build velocity through narrow valleys instead of zigzagging. Best for convnets when tuned well. Narrow LR sweet spot.

Adam. Per-parameter adaptive rescaling using first and second gradient moments with bias correction. Each parameter effectively gets its own learning rate. Defaults: $\beta_1=0.9, \beta_2=0.999, \epsilon=10^{-8}, \eta=3 \times 10^{-4}$. The full update equations and derivation are in Ch 4 §12 (which owns optimizer math); the entry here covers only what changes when you scale Adam up to deep nets.

AdamW. Adam with decoupled weight decay. Original Adam folds weight-decay into the gradient before the adaptive rescaling, which means the effective regularization depends on the per-parameter variance estimate — that is wrong, decay should be uniform. AdamW applies weight decay as a separate $\theta \leftarrow (1 - \eta \lambda) \theta$ step outside the gradient. For any model with weight decay, AdamW is the default. Karpathy: "use AdamW. Always."

Lion (Chen et al. 2023). Uses only the sign of the momentum: $\theta_t = \theta_{t-1} - \eta \cdot \text{sign}(\beta_1 v_{t-1} + (1-\beta_1) g_t)$. The sign operation makes Lion memory-efficient (only one moment to track, not two) and somewhat faster. It works well on vision and small transformers. For large LLMs the results are mixed. Use it if you are memory-constrained.

Muon (Jordan et al. 2024). The current frontier for matrix-valued parameters (the 2D weights of linear layers). It orthogonalizes the momentum matrix via a Newton-Schulz iteration (a cheap iterative routine that pushes the matrix's singular values toward 1) before stepping. The intuition is that gradient-descent updates often align along a few dominant directions; orthogonalizing them spreads the update across all directions, accelerating training. Muon is used for the matrix params; biases, norms, and 1D params still use AdamW. State-of-the-art for transformer training in 2024-25.

# AdamW with parameter groups (the production default)
optimizer = torch.optim.AdamW([
    {"params": decay_params, "weight_decay": 0.1},
    {"params": nodecay_params, "weight_decay": 0.0},
], lr=3e-4, betas=(0.9, 0.95))

# Lion (third-party, https://github.com/lucidrains/lion-pytorch)
from lion_pytorch import Lion
optimizer = Lion(model.parameters(), lr=1e-4, weight_decay=0.1)

# Muon (third-party, https://github.com/KellerJordan/Muon)
from muon import Muon
optimizer_muon = Muon(matrix_params, lr=0.02, momentum=0.95)
optimizer_aw = torch.optim.AdamW(scalar_params, lr=3e-4)
# Use both: step both each iteration.

Rule of thumb: if you do not know what to use, use AdamW. If you have measured AdamW and want to try something new, try Muon for matrix params. Lion is a niche choice.

The learning rate $\eta$ matters more than any other hyperparameter. A schedule changes $\eta$ over training, exploiting that you want different behavior at different points.

Linear warmup. For the first $T_w$ steps, ramp $\eta$ from 0 to the target value. Without warmup, the very first gradient step is too aggressive: the model's outputs are near-random, the loss is high, the gradient is huge, and a single big step can destabilize the parameters before adaptive optimizers have time to build their second moments. Standard $T_w$ is 1-10% of total training steps.

Cosine decay. After warmup, smoothly decay $\eta$ to a minimum (usually $0.1 \times \eta_{max}$) following a half-cosine: $\eta(t) = \eta_{min} + 0.5 (\eta_{max} - \eta_{min})(1 + \cos(\pi t / T))$. The geometric reason: late in training, you want small steps to find the flat minimum's interior; early in training, you want large steps to make progress. Cosine is smooth and works almost universally. Default for LLM training.

1cycle policy (Smith 2017). Warmup to $\eta_{max}$, then anneal below the starting LR. The trick is that momentum is also scheduled (high momentum at low LR, low momentum at high LR). Often gives better results than warmup+cosine for fixed-budget training. Default for fastai-style training.

Warm restarts (Loshchilov and Hutter 2017). Periodic resets back to $\eta_{max}$. The intuition is that each restart escapes a local minimum; the final epoch settles into the best one. Useful for very long training runs.

The PyTorch API:

import torch.optim.lr_scheduler as sched

opt = torch.optim.AdamW(model.parameters(), lr=3e-4)

# Warmup + cosine (the most common modern combo, via SequentialLR)
warmup = sched.LinearLR(opt, start_factor=0.01, total_iters=1000)
cosine = sched.CosineAnnealingLR(opt, T_max=99_000, eta_min=3e-5)
scheduler = sched.SequentialLR(opt, schedulers=[warmup, cosine], milestones=[1000])

# Or all-in-one OneCycleLR
scheduler = sched.OneCycleLR(opt, max_lr=3e-4, total_steps=100_000,
                              pct_start=0.1, anneal_strategy="cos")

# In your training loop, step the scheduler EVERY iteration:
for batch in train_dl:
    opt.zero_grad(); loss = compute_loss(batch); loss.backward(); opt.step()
    scheduler.step()

If a gradient norm spikes (a bad batch, a numerical issue, a saturation), the resulting update can move the parameters far enough to break training. Gradient clipping bounds the L2 norm of the gradient before the update.

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

Standard max_norm: 1.0 for transformers. The clipping is rare in normal training (most gradients are much smaller than 1.0); the value is in catching the outlier batches that would otherwise destroy training.

There are two related techniques. Gradient norm penalty adds $\lambda \|\nabla L\|^2$ to the loss; useful for adversarial training, less so for standard training. Per-parameter clipping clips each parameter's gradient independently, which is faster but less principled.

Most production models start from a pretrained checkpoint — weights someone else already trained on a large dataset, saved and loaded with the same machinery you used for your own checkpoints in another chapter. By convention the backbone is the feature-extracting bulk of that network and the head is the small final layer that maps features to task outputs. Fine-tuning is the process of adapting that checkpoint to a new task.

The two patterns:

Linear probe. Freeze all pretrained weights, train only a new classifier head. Cheap (no backward pass through the backbone), works well when the pretrained features are already good for the target task. Standard baseline.

for p in model.backbone.parameters():
    p.requires_grad = False
model.head = nn.Linear(model.backbone.out_features, n_target_classes)
opt = torch.optim.AdamW(model.head.parameters(), lr=1e-3)

Full fine-tune. Unfreeze everything, train with a small learning rate. Most expressive, best results when you have enough data to avoid catastrophic forgetting.

opt = torch.optim.AdamW(model.parameters(), lr=1e-5, weight_decay=0.01)

Discriminative learning rates (Howard and Ruder 2018). Use different learning rates for different layers. Lower rates for the early layers (general features that you do not want to overwrite) and higher rates for the later layers and head (task-specific features).

opt = torch.optim.AdamW([
    {"params": model.backbone.early_layers.parameters(), "lr": 1e-6},
    {"params": model.backbone.late_layers.parameters(),  "lr": 1e-5},
    {"params": model.head.parameters(),                   "lr": 1e-3},
])

LoRA (low-rank adaptation, Hu et al. 2021) is the modern frontier for fine-tuning large models. Instead of updating the full weight matrix $W$, you add a low-rank correction: $W' = W + AB$ where $A \in \mathbb{R}^{d \times r}$, $B \in \mathbb{R}^{r \times d}$, $r \ll d$. Only $A$ and $B$ are trained; $W$ is frozen. With $r=8$, a 7B-parameter model has only ~20M trainable parameters. Train on a single GPU; merge at inference. Covered in detail in another chapter (efficient inference).

Catastrophic forgetting. When you fine-tune, you can erase the pretrained capabilities. Symptoms: pre-fine-tune the model on a benchmark, the benchmark scores collapse post-fine-tuning. Mitigations: use a low LR, train for few epochs, mix in some of the original pretraining data, or use LoRA (which leaves the pretrained weights unchanged).

You will tune at least: learning rate, weight decay, batch size, and either dropout or some architectural choice. Manual grid search wastes compute. Optuna automates it.

The two pieces: an objective function that takes a Trial and returns the validation metric, and a study that runs trials.

import optuna

def objective(trial: optuna.Trial) -> float:
    # Search space
    lr = trial.suggest_float("lr", 1e-5, 1e-2, log=True)
    weight_decay = trial.suggest_float("weight_decay", 1e-4, 1e-1, log=True)
    hidden_dim = trial.suggest_categorical("hidden_dim", [128, 256, 512])
    dropout = trial.suggest_float("dropout", 0.0, 0.5)

    # Train
    model = MLP(d_in=784, d_hidden=hidden_dim, d_out=10, dropout=dropout)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    val_acc = train_model(model, opt, n_epochs=10)

    # Report intermediate values for pruning
    trial.report(val_acc, step=10)
    if trial.should_prune():
        raise optuna.TrialPruned()
    return val_acc

study = optuna.create_study(direction="maximize",
                             sampler=optuna.samplers.TPESampler(),
                             pruner=optuna.pruners.MedianPruner())
study.optimize(objective, n_trials=100)
print(study.best_params)
print(study.best_value)

Three knobs that matter:

Sampler. TPESampler (Tree-structured Parzen Estimator) for most problems. GridSampler for small categorical search spaces where you actually want exhaustive coverage. CmaEsSampler for high-dimensional continuous spaces.

Pruner. Stops unpromising trials early. MedianPruner is the default and works well. HyperbandPruner is more aggressive and saves compute on long-training setups.

Search space size. Use log=True for learning rates and weight decays (the right scale is multiplicative). Use categorical for architectural choices (number of layers, activation type). Use int for things like hidden dimension (with step=64 or 128 if you want round numbers).

A training run can fail in many ways. The first 80% of debugging follows Karpathy's recipe with high reliability. Here is the checklist, in order:

Loss does not decrease at all (flat at $\log K$).

1. Verify loss at init: it should be $\log K$ for $K$-class classification. If it is something else, your output layer is mis-initialized or you have a bug in the loss function.
2. Print one batch right before it enters the model. Are the inputs and labels what you expect? Normalization correct? Labels not flipped?
3. Overfit one batch of 4 examples. If you cannot drive loss to zero on 4 examples with a model that has enough capacity, you have a bug. Find it before scaling up.
4. Check that optimizer.zero_grad is being called.
5. Check that requires_grad=True is set on parameters (it is by default for nn.Parameter, but if you froze something during transfer learning, did you unfreeze the right thing?).

Loss is decreasing but very slowly.

1. Is the learning rate too small? Try a 10× LR.
2. Are you using optim.SGD without momentum? Adam usually trains faster early.
3. Is your data normalized? Inputs with very different scales train slowly.
4. Are gradients vanishing at early layers? Log per-layer gradient norms.

Loss spikes to NaN.

1. Gradient explosion. Add clip_grad_norm_ with max_norm=1.0.
2. Mixed-precision overflow. Check GradScaler is enabled. Or switch to bfloat16.
3. Numerical instability in the loss. log(0) or 1/0. Add eps to denominators.
4. Bad data: an outlier example with extreme values. Print the batch right before the NaN.

Loss decreases but val accuracy does not improve.

1. You are overfitting. Increase regularization (dropout, weight decay, or more data).
2. Val/train data distribution mismatch. Plot both distributions, look for differences.
3. Bug in eval code. Run train and val on the same data — they should give the same loss.

Loss diverges late in training.

1. LR schedule has not annealed. Lower the final LR.
2. The optimizer's second moment estimate has accumulated bad statistics. Reset the optimizer.
3. Catastrophic forgetting during fine-tuning. Lower the LR or stop earlier.

This checklist is mostly a transcription of Karpathy's recipe (24-founder-blogs/karpathy-recipe), with additions from the cs336 lecture 2 troubleshooting section and the ARENA optimization notebook. Internalize it before you have to debug under time pressure.