Deep Learning in Practice

Assumes you know from ML-101

This lesson builds directly on ML-101 Lesson 14 (Neural Networks & Backprop) and ML-101 Lesson 11 (Gradient Descent & Optimization). We assume you are already comfortable with the chain rule, the per-parameter update $\theta \leftarrow \theta - \eta \nabla_\theta \mathcal{L}$, the concept of a minibatch, and the high-level structure of a fully connected feed-forward network. We will not re-derive the chain rule or define what a tensor is.

The reader is also expected to be familiar with ML-101 Lesson 12 (Ensemble Learning): dropout, the standard motivation (an exponential ensemble of thinned subnetworks), and the empirical observation that it acts as regularisation. We will look under that observation and see that dropout in the linear regime is equivalent to $\ell_2$ regularisation at a specific rate that depends on the dropout probability — a result you cannot get without the matrix algebra developed here.

Learning Objectives

1. Derive the autograd graph for a two-layer MLP and trace one backwards pass, identifying exactly where each operator in the forward computation graph contributes a Jacobian row to the local gradient signal.
2. Explain why batch normalisation behaves differently at training and inference time, and derive the population-statistic update rule from the maximum-likelihood framing of the moving-average estimator.
3. Compare SGD with Polyak averaging, momentum, and Adam, and predict — from the loss surface curvature, not from folklore — which optimiser will reach a given loss tolerance in fewer parameter updates.
4. Diagnose a non-converging deep network by separating optimisation pathologies (vanishing/exploding gradients, dead ReLUs, learning-rate mismatch) from statistical pathologies (label noise, distribution shift, overfitting), and design a minimal experiment to distinguish them.
5. Justify, with the relevant inequality, why a residual connection $\mathbf{h} = f(\mathbf{x}) + \mathbf{x}$ is provably easier to optimise than the corresponding plain feed-forward block $\mathbf{h} = f(\mathbf{x})$ under the assumption of identity initialisation.

The training loop as a fixed-point iteration

In ML-101 Lesson 11 we wrote the parameter update as a single line: $\theta_{t+1} = \theta_t - \eta \nabla \mathcal{L}(\theta_t)$. In practice, that single line hides a loop with five distinct moving parts: a minibatch sampler, an autograd engine, an optimiser state object, a checkpoint policy, and an evaluation hook. Treating the loop as "just SGD" is the single most common source of silent bugs in deep-learning code. Let us write the loop honestly.

import torch
from torch.utils.data import DataLoader
from torch.optim import AdamW

def train(model, train_set, val_set, *, lr=1e-3, epochs=10, device="cuda"):
    model.to(device)
    loader = DataLoader(train_set, batch_size=64, shuffle=True, num_workers=4)
    opt = AdamW(model.parameters(), lr=lr, weight_decay=0.01)

    for epoch in range(epochs):
        model.train()                  # flip BN/dropout into train mode
        for x, y in loader:            # minibatch sampler
            x, y = x.to(device), y.to(device)
            opt.zero_grad(set_to_none=True)        # reset accumulated grads
            loss = loss_fn(model(x), y)            # forward: builds autograd graph
            loss.backward()                         # backward: populates .grad
            opt.step()                              # optimiser reads .grad, writes .theta
        evaluate(model, val_set, device=device)     # hook uses model.eval() mode

The loop is a fixed-point iteration over a distribution, not a function: each opt.step() updates $\theta$ using a gradient that is itself a random sample of the true population gradient $\nabla \mathcal{L}_{\text{pop}}(\theta)$. The convergence guarantee is therefore a stochastic one: under standard assumptions (Lipschitz-smooth loss, bounded gradient variance, bounded domain), Polyak–Ruppert averaging of the iterates converges at rate $\mathcal{O}(1/\sqrt{T})$ for the expected suboptimality, and $\mathcal{O}(1/T)$ for the averaged iterates' suboptimality. The constant in the rate depends on the curvature of the loss at the optimum, which is why the same learning rate that converges on MNIST diverges on ImageNet.

Two consequences matter in practice. First, the training loss is a biased estimator of the population loss, and the bias is approximately $\frac{1}{B}\sum_i \operatorname{tr}(H_i)$ with $B$ the batch size and $H_i$ the per-sample Hessian — small-batch noise is therefore not uniform: it is anisotropic in proportion to the local curvature. Second, the loop is mode-collapsing in expectation: the stationary distribution of SGD on a non-convex landscape concentrates mass near sharp minimisers of the training loss, which is the geometric reason why flat-minima-seeking methods like SAM help generalisation.

Autograd: where the gradient actually comes from

PyTorch's autograd implements reverse-mode automatic differentiation over a dynamically built computation graph. The forward pass records every operator together with a reference to its inputs, and the backward pass walks the graph in reverse, applying the multivariate chain rule at every node. Let us make this concrete on a two-layer MLP with ReLU activations.

def forward_two_layer(x, W1, b1, W2, b2):
    h = x @ W1.T + b1          # (B, 200)
    h = torch.relu(h)           # non-linearity
    y_hat = h @ W2.T + b2       # (B, num_classes)
    return y_hat

Each operator in this forward pass contributes a Jacobian row to the local gradient signal during the backward pass. The chain rule that ML-101 introduced as $\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial y}\frac{\partial y}{\partial x}$ is, in the multivariate case, a chain of Jacobian-vector products. The gradient at any parameter $\theta_i$ is

where $J_k = \partial f_k / \partial f_{k-1}$ is the Jacobian of layer $k$ evaluated at the forward-pass activations, and the product is read right to left (the rightmost factor is closest to the loss). The forward pass stores the activations that each Jacobian is evaluated at; if you modify an activation in place between the forward and the backward, autograd will raise a RuntimeError because the value it needs has been overwritten. This is the deeper reason behind the "no in-place ops on ReLU outputs" rule.

A practical consequence is that gradient checking by finite differences must be done against the same forward pass that produced the cached values. The naïve checker

is correct in principle but unstable in floating point when $\epsilon$ is too small (catastrophic cancellation in the numerator) or too large (truncation of the Taylor remainder). The standard choice is $\epsilon \approx 10^{-3}\max(1, \lvert x \rvert)$ with double precision, and the check is considered passing when the relative error is below $10^{-5}$. Do not gradient-check against a network that uses BatchNorm in training mode — the per-batch running statistics make the forward pass non-deterministic between calls, and the finite-difference gradient will disagree with the autograd one by an amount unrelated to your code's correctness.

The most subtle autograd bug is a forgotten detach(). If a downstream loss includes a term that is supposed to be a target (a hard label, an index into a vocabulary) but is actually a tensor that carries a grad_fn, autograd will treat it as a parameter and backpropagate into it. This is usually caught by a RuntimeError: element 0 of tensors does not require grad, but in torch.no_grad() blocks it is silent and you lose gradient flow for the whole downstream computation.

Initialisation: why the variance matters

The default initialisation in PyTorch's nn.Linear is Kaiming uniform with $a = \sqrt{5}$, which sets the variance of the layer's output to the variance of its input under the assumption of a linear activation. The derivation is one you should be able to reproduce from scratch. Let $W \in \mathbb{R}^{n_\text{out} \times n_\text{in}}$ be a weight matrix and $\mathbf{x} \in \mathbb{R}^{n_\text{in}}$ an input vector. The pre-activation is $z_i = \sum_{j=1}^{n_\text{in}} W_{ij} x_j$. If $W_{ij}$ and $x_j$ are independent with $\mathbb{E}[W_{ij}] = 0$, $\operatorname{Var}(W_{ij}) = \sigma_w^2$, $\mathbb{E}[x_j] = 0$ and $\operatorname{Var}(x_j) = \sigma_x^2$, then

For the variance to be preserved across layers we require $\operatorname{Var}(z) = \operatorname{Var}(x)$, which forces $\sigma_w^2 = 2 / n_\text{in}$ for a ReLU activation (the factor of 2 compensates for the ReLU killing half the signal) and $\sigma_w^2 = 1 / n_\text{in}$ for a linear activation. This is the Kaiming–He initialisation for ReLU networks and the LeCun initialisation for purely linear networks.

The practical failure mode this initialisation prevents is signal collapse or explosion across depth. With $L$ layers and i.i.d. weights, the variance of the $L$-th layer's pre-activation is $\operatorname{Var}(z^{(L)}) = (\sigma_w^2 n_\text{in})^L \sigma_x^2$. If $\sigma_w^2 n_\text{in} \neq 1$ then either all activations saturate (in logistic / tanh networks) or blow up to infinity (in ReLU networks) within $\mathcal{O}(\log n)$ layers. You can observe this in real time by logging the per-layer activation mean and variance on a fixed batch during the first dozen steps; a healthy network keeps both roughly constant across layers.

A common 101-level misconception is that initialisation is "a small thing that only matters for very deep networks". It is not. Even a three-layer MLP trained with $\sigma_w = 1$ (the default for torch.randn-initialised weights) on MNIST will plateau at a training loss an order of magnitude above the Kaiming-initialised baseline. Initialisation sets the operating point of the optimiser; the optimiser then walks downhill from there. A bad operating point is not something momentum or learning-rate schedules can fix.

Normalisation layers and the train/inference gap

Batch normalisation performs

where $\mu_{\mathcal{B}}$ and $\sigma_{\mathcal{B}}^2$ are the per-feature mean and variance over the current minibatch $\mathcal{B}$, and $\gamma, \beta$ are learned per-feature affine parameters. The gradient with respect to the inputs is well known; the subtle part is what happens at inference time.

At inference we cannot condition on a single minibatch — the deployed network receives one example at a time and we need a deterministic forward pass. BatchNorm therefore switches to using a running estimate of the population per-feature statistics,

where $\hat{\mu}$ and $\hat{\sigma}^2$ are updated during training via an exponential moving average:

This is a Robbins-Monro stochastic-approximation step for the population mean and second moment, with step-size $m$ (the momentum argument to nn.BatchNorm, confusingly — note this is not the optimiser momentum of Lesson 5). The step-size schedule must satisfy $\sum_t m_t = \infty$ and $\sum_t m_t^2 < \infty$ for convergence; the constant-$m$ schedule does not satisfy the second condition in theory but is the universal practical choice because $\hat{\mu}$ converges geometrically.

The train/inference gap — the difference between the statistics used in each mode — is the source of a famous production bug. If you call model.train() once and then model.eval() only at the very end of training, the running statistics are correctly accumulated. If you forget to flip back to model.eval() for evaluation, the running statistics are not updated but the per-batch statistics are used, so the evaluation loss matches a hypothetical network that sees one example at a time using batch-statistics — i.e. it is silently biased. The model.eval() / model.train() calls are not stylistic; they re-route the forward pass through different code paths inside BatchNorm, Dropout, and any custom layer that subclasses nn.Module and uses the self.training flag. A useful diagnostic is to record the running $\hat{\mu}$ values during training and verify they track the per-batch $\mu_{\mathcal{B}}$; if they diverge catastrophically, the running-average momentum is wrong (typically inverted: people pass momentum=0.1 when the API expects momentum=0.9).

Layer normalisation and group normalisation sidestep the train/inference gap entirely because they compute statistics over the feature axis rather than the batch axis. They are the right choice whenever the batch size is small (e.g. per-GPU micro-batches in data-parallel training) or when the deployment regime is genuinely online.

Optimiser trajectories in deep loss landscapes

The loss landscape of a deep network is non-convex in general, and the trajectory an optimiser takes through it is determined by two things: the local curvature (the Hessian, or its stochastic estimates) and the noise scale (the minibatch sampling variance). Let us compare the trajectories of SGD, momentum, and Adam on the same loss starting from the same initial point.

The table is not folklore; it is a consequence of how each update rule responds to a quadratic perturbation $\mathcal{L}(\theta + \delta) \approx \mathcal{L}(\theta) + g^\top \delta + \tfrac{1}{2}\delta^\top H \delta$. For SGD, the equilibrium noise amplitude along an eigenvector of $H$ with eigenvalue $\lambda$ is proportional to $\eta / \lambda$, so high-curvature directions ("sharp" directions) are visited with low amplitude — the trajectory prefers flat basins. For Adam, the per-coordinate rescaling divides out the Hessian's eigenvalue information, so the noise amplitude is approximately $\eta$ along every direction regardless of $\lambda$, and sharp minimisers are as reachable as flat ones. This is the geometric reason why Adam-trained networks often generalise worse than SGD-trained ones on the same compute budget, and it is the reason the original Adam paper recommends $\eta = 10^{-3}$ rather than the $10^{-1}$–$10^{-2}$ typical of SGD.

A practical piece of advice that follows from the analysis: when transferring a network from Adam to SGD (a common fine-tuning trick to recover the flat-minima bias), drop the learning rate by an order of magnitude and expect the loss to spike for the first hundred steps before settling. Conversely, when transferring from SGD to Adam, the per-coordinate rescaling means the learning rate that worked for SGD is roughly the upper bound of what Adam can tolerate without diverging; start at $10^{-3}$ and tune upward.

Residual connections: a contraction argument

A residual block computes

where $f_i$ is some parametric function — typically two convolutional layers with a non-linearity in between. The plain version, $\mathbf{h}_{i+1} = f_i(\mathbf{h}_i; \theta_i)$, is the limit case where the residual function is constrained to include the identity. The original ResNet paper motivates this as "easier to learn the residual than the unreferenced mapping", but the cleaner argument is one of gradient magnitudes.

Consider the gradient of the loss with respect to an early-layer parameter $\theta_k$ under the residual form. By the chain rule,

Under the assumption of identity initialisation — i.e. $f_i(\mathbf{h}_i) = \epsilon \cdot \tilde{f}_i(\mathbf{h}_i)$ for small $\epsilon$ at initialisation — the Jacobians $J_{f_i}$ have small operator-norm contributions and the matrix $(I + J_{f_i})$ has eigenvalues bounded away from zero. The product over depth therefore has eigenvalues that are products of numbers bounded near one, not products of numbers that decay exponentially with depth. The plain network, by contrast, has $\prod_{i=k}^{L-1} J_{f_i}$, a product of matrices whose operator norms grow or shrink multiplicatively.

More precisely, under the assumption that $\lVert J_{f_i} \rVert_\text{op} \leq \alpha < 1$ for the plain network (a statement that the per-layer mapping is contractive), the gradient signal decays as $\alpha^{L-k}$ — a vanishing gradient. Under the residual assumption $\lVert J_{f_i} \rVert_\text{op} \leq \beta < 2$, the gradient signal decays at most as $\beta^{L-k}$, which can be made as slow as desired by initialising $f_i$ near zero. This is what "identity initialisation enables very deep networks" means precisely.

The 101-level misconception is that "skip connections help because they make the loss landscape smoother". That is empirically true but does not generalise: there are smooth landscapes that are still hard to optimise. The contraction argument above identifies the one specific assumption (that $f_i$ is initialised near identity and its Jacobian has bounded operator norm) whose violation breaks the optimisation benefit. In practice, if you zero-initialise the last layer of each residual block (so $f_i$ outputs zero at initialisation), the network starts as the identity function and the assumption is met by construction; the optimisation benefit then follows deterministically regardless of the loss landscape's shape.

import torch.nn as nn

class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.fc1 = nn.Linear(dim, dim)
        self.fc2 = nn.Linear(dim, dim)
        # Zero the LAST layer so f(h) = 0 at initialisation.
        # The block computes h + f(h) = h at step 0; optimisation
        # then has to "turn on" f slowly, never destroy h abruptly.
        nn.init.zeros_(self.fc2.weight)
        nn.init.zeros_(self.fc2.bias)

    def forward(self, h):
        return h + self.fc2(torch.relu(self.fc1(h)))

Key Takeaways

• The training loop is a stochastic fixed-point iteration; the bias of its minibatch gradient estimator is anisotropic and proportional to the local Hessian trace, which is why "large batch = small noise" is a misleading rule of thumb on ill-conditioned losses.
• Autograd implements reverse-mode differentiation over a dynamically constructed graph; the backward pass is a sequence of Jacobian-vector products, and gradient checking must respect the train-mode semantics of every layer it touches.
• Weight initialisation is not a small detail: it sets the operating point of the optimiser, and the Kaiming–He derivation is the precise condition that preserves activation variance across ReLU layers.
• BatchNorm has a train/inference gap because it uses per-batch statistics during training and running population statistics at inference; the model.eval() call is not stylistic, it re-routes the forward pass.
• Adam's per-coordinate rescaling removes the implicit sharpness bias of SGD, which is why SGD often generalises better on a fixed compute budget and why the AdamW decoupled-weight-decay form is preferred over Adam+L2.
• Residual connections work because they make the depth-wise Jacobian product a sum of bounded operators instead of a product of arbitrary ones, and the optimisation benefit follows from identity initialisation rather than from any property of the loss landscape.