You write a ten-layer network. The forward pass is right, every gradient passes a finite-difference check, the optimizer is AdamW with a sensible learning rate. And the loss sits at 2.30 for an hour, not moving. Or it moves for a while and then turns into NaN. Nothing is wrong with your math. Something is wrong with the numbers flowing through the network.

Training deep networks well is a craft, but it is a craft with reasons. The signals going forward and the gradients going backward are both products of many layers, and products of many numbers either explode or vanish unless something keeps them near 1. Initialization, normalization, the choice of activation, and residual connections (§0.8) are all ways of keeping them near 1. This section derives the first three, builds an MLP and both kinds of normalization from scratch with their backward passes, covers the classic regularizers, explains the surprising shape of the test-error curve called double descent, and ends with a debugging recipe that will save you days.

An MLP from scratch

A multilayer perceptron (MLP) alternates linear maps and elementwise nonlinearities. For a batch $\mathbf{X} \in \mathbb{R}^{N \times d_{\text{in}}}$, a two-layer classifier with hidden width $H$ and $K$ classes computes

with $\mathbf{W}_1 \in \mathbb{R}^{d_{\text{in}} \times H}$, $\mathbf{W}_2 \in \mathbb{R}^{H \times K}$, an activation $\phi$ such as ReLU, and the mean softmax cross-entropy CE. In plain words, mix the inputs into $H$ new features, bend each feature with a nonlinearity, then mix those into $K$ class scores.

The backward pass is the chain of VJPs from §0.2, run from the loss back to the input.

where $\mathbf{P}$ is the softmax of $\mathbf{Z}_2$ and $\mathbf{Y}$ the one-hot targets. Every line is either the linear-layer rule or the elementwise rule. For ReLU, $\phi'(z)$ is 1 where $z > 0$ and 0 elsewhere, so the backward pass zeroes the gradient wherever the unit was off.

Why should such a simple model be able to fit anything? The universal approximation theorem (Cybenko, 1989, and Hornik, 1991) says that a single hidden layer with enough units can approximate any continuous function on a closed, bounded region as closely as you like. It is an existence result, though. It says nothing about how many units you need (possibly exponentially many), or whether gradient descent will find them. Depth buys enormous efficiency. Telgarsky (2016) constructed functions that a deep network represents with a modest number of units but that any shallow network needs exponentially many units to approximate. That is why we stack layers, and stacking layers is what creates the problems in the rest of this section.

import numpy as np

def mlp_loss_and_grads(X, y, params):
    raise NotImplementedError

Activations

The nonlinearity is what makes depth worth anything, since a stack of linear layers is just one linear layer. But its derivative multiplies every backward signal passing through it, so its shape matters a great deal for training.

Here $\Phi$ is the standard normal CDF, and GELU is often computed with the approximation $0.5z\big(1 + \tanh(\sqrt{2/\pi}(z + 0.044715z^3))\big)$. A few of these entries deserve a sentence each.

Sigmoid's derivative is at most 1/4. Backpropagating through ten sigmoid layers shrinks the gradient by at least $4^{10} \approx 10^6$, before even counting the weights. That alone made deep sigmoid networks nearly untrainable, and it is one reason ReLU took over around 2010 to 2012 (Glorot et al., 2011, and Krizhevsky et al., 2012).

ReLU units can die. If a unit's pre-activation is negative for every input (after a big update pushed its bias down, say), its gradient is zero everywhere and it can never recover. A large learning rate can kill a sizable fraction of a layer. Leaky ReLU and the smooth variants avoid this.

Smooth activations won in Transformers. GELU (Hendrycks & Gimpel, 2016) and SiLU are smooth near zero and have a small negative lobe. The gated feed-forward layers of modern LLMs, such as SwiGLU, multiply a SiLU branch by a linear branch, and they are covered in §5.3.

Initialization, derived

Before the first gradient step, the weights are random. How large should they be? This question decides whether a deep network can train at all, and the answer comes from the variance rules of §0.4.

First, a non-answer. Setting every weight to zero (or to the same constant) fails for a reason that has nothing to do with scale. Every hidden unit in a layer then computes the same function and receives the same gradient, so they stay identical forever, and the layer behaves like a single unit. Random initialization breaks the symmetry. The question is only how large the random values should be.

Keeping the forward signal steady

Consider one layer $\mathbf{y} = \mathbf{W}\mathbf{x}$ for a single input vector $\mathbf{x} \in \mathbb{R}^{n_{\text{in}}}$ and $\mathbf{W} \in \mathbb{R}^{n_{\text{out}} \times n_{\text{in}}}$, with weights drawn independently with mean 0 and variance $\sigma_w^2$, independent of $\mathbf{x}$, and no bias (biases start at zero). Each output is $y_i = \sum_{j=1}^{n_{\text{in}}} W_{ij}x_j$, a sum of $n_{\text{in}}$ uncorrelated terms (the weights are independent with mean zero, which makes the terms uncorrelated even when the inputs are not). Each term has mean $\mathbb{E}[W_{ij}]\mathbb{E}[x_j] = 0$ and variance $\mathbb{E}[W_{ij}^2 x_j^2] = \sigma_w^2\,\mathbb{E}[x_j^2]$. Adding them up,

In plain words, each layer multiplies the signal's second moment by $n_{\text{in}}\sigma_w^2$, times whatever the activation does. Over $L$ layers that factor is raised to the $L$th power, so unless it is 1, the signal explodes or vanishes exponentially with depth.

For a linear or tanh network operating near zero, the input to the next layer has zero mean and $\mathbb{E}[x^2] = \operatorname{Var}(x)$, so the condition for preserving variance is

This is often called LeCun initialization (LeCun et al., 1998).

Keeping the backward signal steady

The backward pass sends $\mathbf{g}_x = \mathbf{W}^\top\mathbf{g}_y$, where each input gradient sums $n_{\text{out}}$ terms. The same calculation gives $\operatorname{Var}(g_x) = n_{\text{out}}\,\sigma_w^2\operatorname{Var}(g_y)$, so preserving gradient variance needs $\sigma_w^2 = 1/n_{\text{out}}$. Unless the layer is square, you cannot satisfy both. Glorot & Bengio (2010) proposed the compromise that averages the two fan counts,

If you draw from a uniform distribution $\mathcal{U}(-a, a)$ instead, its variance is $a^2/3$, so the same variance needs $a = \sqrt{6/(n_{\text{in}} + n_{\text{out}})}$.

The ReLU correction

ReLU breaks one assumption above. Its output is never negative, so it is not zero-mean, and it zeroes half its inputs. Let $y$ be a pre-activation, symmetric around zero (as it is when the weights are), and $x = \max(0, y)$ the next layer's input. Then $x^2$ equals $y^2$ for the positive half and 0 for the negative half, and by symmetry

Plug this into the forward rule. The pre-activation variance of the next layer is $n_{\text{in}}\sigma_w^2\cdot\tfrac12\operatorname{Var}(y)$, so preservation requires

In plain words, ReLU throws away half the signal's energy at every layer, so the weights must be $\sqrt{2}$ times larger to compensate. He et al. (2015) derived this and used it to train 30-layer plain networks that Xavier initialization could not. They also showed that preserving the backward variance instead, $\sigma_w^2 = 2/n_{\text{out}}$, works just as well, because the two choices differ only by a ratio of layer widths that does not compound with depth.

The general pattern is $\sigma_w = \text{gain}/\sqrt{n_{\text{in}}}$, with a gain of 1 for linear layers, $\sqrt2$ for ReLU, and $5/3$ for tanh in PyTorch's calculate_gain (a little above 1, to compensate for tanh squashing its larger inputs). Two practical notes. PyTorch's default nn.Linear initialization is uniform with bound $1/\sqrt{n_{\text{in}}}$, which gives variance $1/(3n_{\text{in}})$, smaller than He, so deep ReLU networks often benefit from explicit Kaiming init. And Transformers usually use a small fixed standard deviation such as 0.02, with the output projections of residual branches scaled down by $1/\sqrt{2L}$ for $L$ layers (GPT-2 did this), so that the residual stream's variance does not grow with depth. That residual argument is the subject of §5.7, and the question of how initialization should change with width leads to muP in §9.5.

Watch the derivation come true. The widget pushes a batch through a deep stack of layers and plots the standard deviation of the activations at each layer, and of the gradients on the way back.

With ReLU, Xavier's variance is half too small, and after 30 layers the signal has shrunk by a factor of about $2^{15}$, as the derivation predicts. He initialization keeps it near 1. With tanh, Xavier keeps the signal alive (it decays slowly, because tanh compresses large values), and a tiny constant scale kills it in a handful of layers.

import numpy as np

def init_weights(fan_in, fan_out, scheme, dist, rng):
    raise NotImplementedError

Vanishing and exploding gradients

The backward pass through $L$ layers multiplies the gradient by $L$ Jacobians, $\nabla_{\mathbf{x}_0}\mathcal{L} = \mathbf{J}_1^\top\mathbf{J}_2^\top\cdots\mathbf{J}_L^\top\nabla_{\mathbf{x}_L}\mathcal{L}$. Each Jacobian of a dense layer with an activation is $\operatorname{diag}(\phi'(\mathbf{z}))\mathbf{W}$. If these matrices typically shrink vectors by a factor $r < 1$ (how much a matrix can stretch or shrink a vector is measured by its singular values, the subject of §0.10), the gradient at the first layer is about $r^L$ times the one at the last, and the early layers stop learning. That is the vanishing gradient problem. If they typically stretch by $r > 1$, gradients grow like $r^L$, steps become huge, and the loss diverges. That is the exploding gradient problem.

Initialization makes $r \approx 1$ at the start of training. It cannot guarantee it stays there, since the weights and activation statistics change as training proceeds. That is why deep networks also rely on normalization layers and residual connections, which keep the signal well scaled throughout training, not just at step zero. Recurrent networks are the extreme case, since they multiply by the same matrix at every time step, and §0.8 studies them.

Batch normalization

Batch normalization (Ioffe & Szegedy, 2015) takes a direct approach. If you want each layer's inputs to have zero mean and unit variance, normalize them to have exactly that, using statistics of the current minibatch. For a batch of activations $\mathbf{X} \in \mathbb{R}^{N \times D}$ ($N$ examples, $D$ features), BatchNorm normalizes each feature (column) over the batch,

where $\boldsymbol{\gamma}, \boldsymbol{\beta} \in \mathbb{R}^D$ are learned scale and shift parameters and $\epsilon$ (such as $10^{-5}$) prevents division by zero. In plain words, standardize each feature across the batch, then let the network rescale and shift it however it likes. The learned $\gamma$ and $\beta$ mean normalization never limits what the layer can represent. Setting $\gamma_j = \sqrt{\sigma_j^2 + \epsilon}$ and $\beta_j = \mu_j$ would undo it.

Train mode and eval mode

At test time you may see a single example, which has no batch statistics. So during training BatchNorm keeps exponential moving averages of $\mu$ and $\sigma^2$, and at evaluation time it normalizes with those fixed values instead. In PyTorch the update is $\text{running} \leftarrow (1 - m)\,\text{running} + m\cdot\text{batch}$ with momentum $m = 0.1$, and the running variance uses the unbiased batch variance (dividing by $N - 1$, as in §0.4), while the normalization itself uses the biased one. This train-eval split is the reason model.eval() exists, and forgetting it is one of the most common PyTorch bugs (§0.9).

The backward pass

This is a classic interview derivation, and it rewards care, because every output in a column depends on every input in that column through $\mu$ and $\sigma$. Work with one feature column and drop the index $j$. Write $s = \sqrt{\sigma^2 + \epsilon}$, and let $\mathbf{g} = \partial\mathcal{L}/\partial\mathbf{Y}$ be the upstream gradient for that column. The parameter gradients are immediate, $\partial\mathcal{L}/\partial\beta = \sum_n g_n$ and $\partial\mathcal{L}/\partial\gamma = \sum_n g_n\hat x_n$. For the input, the result is

In plain words, take the upstream gradient, remove its mean, remove its component along the normalized inputs, and scale by $\gamma/s$. Those two removals are the gradient's way of respecting normalization's invariances. Shifting all inputs by a constant, or scaling them, does not change $\hat{\mathbf{x}}$, so the gradient has no component in those directions.

Why BatchNorm helps, and its problems

Networks with BatchNorm train much faster, tolerate learning rates several times larger, and are far less sensitive to initialization. Why is debated. The original paper attributed it to reducing "internal covariate shift", the changing distribution of each layer's inputs during training. Santurkar et al. (2018) showed experimentally that deliberately injecting noise after BatchNorm, to increase that shift again, barely hurts, and argued that the real benefit is a smoother loss landscape with better-behaved gradients. Practitioners agree that it works more than they agree on why.

It also has real drawbacks. Its output for one example depends on the other examples in the batch, which couples examples, breaks down for small batches (the statistics get noisy), and complicates distributed training, where each device sees a different sub-batch. The train-eval discrepancy causes bugs. And for sequences of varying length, "the batch statistics of position 17" mixes padding with real tokens. These problems are why sequence models moved to the next layer.

import numpy as np

def batchnorm_forward(X, gamma, beta, eps=1e-5):
    raise NotImplementedError

def batchnorm_backward(dY, cache):
    raise NotImplementedError

import numpy as np

class BatchNorm1d:
    def __init__(self, D, momentum=0.1, eps=1e-5):
        raise NotImplementedError

    def __call__(self, X):
        raise NotImplementedError

Layer normalization

Layer normalization (Ba, Kiros & Hinton, 2016) normalizes each example over its own features instead of each feature over the batch. For one token's vector $\mathbf{x} \in \mathbb{R}^D$,

In plain words, standardize the features of each example on their own, then rescale per feature. Because nothing depends on other examples, LayerNorm behaves identically in training and evaluation, works with a batch of one, and handles variable-length sequences without caring about padding. It is the normalization of the Transformer, applied to every token's vector independently.

Its backward pass is the BatchNorm formula with the roles of the axes swapped. For one row, let $s = \sqrt{\sigma^2 + \epsilon}$ and $\hat x_j = (x_j - \mu)/s$ use that row's own mean and variance. The averages then run over the $D$ features of the row instead of the $N$ rows of a column, and since $\gamma$ is per feature it sits inside the sums,

A cheaper variant, RMSNorm (Zhang & Sennrich, 2019), skips the mean subtraction and divides by the root mean square, $\sqrt{\frac1D\sum_j x_j^2 + \epsilon}$. Most recent LLMs use it. Where to put the normalization (before or after each sublayer) matters enormously for training stability, which §5.2 explains.

import numpy as np

def layernorm_forward(X, gamma, beta, eps=1e-5):
    raise NotImplementedError

def layernorm_backward(dY, cache):
    raise NotImplementedError

Dropout

Dropout (Srivastava et al., 2014) randomly zeroes each activation with probability $p$ during training. A network can then never rely on any single unit, since it might vanish. This discourages fragile co-adaptations between units and acts like training a huge ensemble of thinned networks that share weights.

The detail that matters is the scaling. In inverted dropout, which every framework uses, the surviving activations are divided by $1 - p$ during training,

In plain words, keep each unit with probability $1 - p$ and scale up the survivors so the expected activation is unchanged. At evaluation time dropout does nothing at all, and because expectations already match, no rescaling is needed. Dropout was essential for the large fully connected layers of 2012-era networks. In LLM pretraining it is commonly set to zero, since a model that sees each token roughly once has little opportunity to overfit, while fine-tuning on small data sets often turns it back on.

import numpy as np

def dropout(x, p, training, rng):
    raise NotImplementedError

Weight decay

Weight decay adds $\tfrac\lambda2\lVert\theta\rVert^2$ to the loss, where $\theta$ is all the weights and $\lambda$ the decay strength. §0.4 showed this is MAP estimation under a Gaussian prior, and with SGD at learning rate $\eta$ it shrinks each weight by a factor $(1 - \eta\lambda)$ per step. With Adam, use the decoupled AdamW form from §0.6. A typical value for Transformers is $\lambda = 0.1$, applied to weight matrices but not to biases or normalization gains.

Overfitting, bias, variance, and double descent

A model overfits when it fits its training data much better than new data, learning the noise along with the signal. The classic way to reason about this is the bias-variance decomposition. Suppose $y = f(\mathbf{x}) + \varepsilon$ with noise of mean 0 and variance $\sigma^2$, and a learning procedure produces $\hat f_{\mathcal{D}}$ from a random training set $\mathcal{D}$. At a fixed test point, write $\bar f = \mathbb{E}_{\mathcal{D}}[\hat f_{\mathcal{D}}]$ for the average prediction. Then

In plain words, test error is irreducible noise, plus how wrong the model is on average (bias), plus how much it changes with the particular training set (variance). The classical picture says simple models have high bias and low variance, flexible ones the reverse, and the best model sits at a sweet spot in between, giving a U-shaped test-error curve. Regularization (weight decay, dropout, early stopping, data augmentation) trades a little bias for less variance.

Modern deep learning broke that picture. Networks with far more parameters than training examples fit their training data perfectly, yet generalize well. Belkin et al. (2019) showed the curve has a second part. Test error rises toward a peak at the interpolation threshold, where the model has just enough capacity to fit the training set exactly, and then falls again as capacity keeps growing. This is double descent. Nakkiran et al. (2019) found it in deep networks as a function of width, training time, and data set size.

The cell reproduces it with the simplest possible model, linear regression on random ReLU features, fit by the minimum-norm least-squares solution (which is what gradient descent from zero converges to). With 50 training points, watch what happens near 50 features.

Below 50 features, the classical U appears. At 50, the model must thread every noisy training point with a single exact solution, which needs huge coefficients, and test error spikes. Beyond 50 there are infinitely many exact fits, and the minimum-norm one gets smoother and better as features are added, until it beats every model on the classical side. The deep-learning lesson is that "more parameters than data" is not by itself a reason to expect overfitting. The implicit regularization of the training procedure matters as much as the parameter count.

A recipe for debugging a training run

When a network will not train, the bug is almost never where you first look. Karpathy's "A Recipe for Training Neural Networks" (linked below) captures the habits that find it fast. Here is a condensed version, in the order to apply it.

1. Check the loss at initialization. With $K$ balanced classes and a well-initialized output layer, the logits are near zero, the softmax is near uniform, and the loss should be $\ln K$ (2.303 for 10 classes, and 10.82 for a language model with GPT-2's 50,257-token vocabulary, as in §0.5). A much larger initial loss means the outputs are overconfident from the start, usually because the last layer's weights are too large.
2. Overfit a single small batch. Turn off regularization and train on one batch of, say, 32 examples. The loss should go to nearly zero within a few hundred steps. If it does not, there is a bug in the model, the loss, or the data pipeline, and no hyperparameter will fix it.
3. Verify the data. Look at actual inputs and labels right before they enter the model. Check shapes, value ranges, and that labels line up with inputs after shuffling and augmentation.
4. Start simple and add one thing at a time. A small model with a known-good setup first. Then add complexity, checking that each change helps.
5. Watch the right signals. Training and validation loss, the gradient norm, the update-to-weight ratio (around $10^{-3}$ per step is a common rule of thumb), and activation statistics per layer. Saturated or dead units show up there long before the loss tells you.
6. Tune the learning rate first. It is the most important hyperparameter. Sweep it on a log scale.
7. Regularize only once you overfit. Get the training loss down, then trade some of it for validation performance. And before you believe a small validation improvement, check that it is larger than the run-to-run noise across seeds (§0.11 shows how).

With a large output-layer scale the initial loss is far above $\ln 10$, a sign of overconfident random predictions that waste the first part of training undoing them. With a small scale it starts at $\ln 10$ as it should. Either way the network memorizes a batch of random labels, which confirms that the forward and backward passes are wired correctly. If this check failed, nothing else would matter.

In PyTorch

Read more

• Glorot & Bengio (2010). Understanding the difficulty of training deep feedforward neural networks and He, Zhang, Ren & Sun (2015), Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification. The variance analyses behind Xavier and He initialization (the ReLU correction is in He et al., Section 2.2).
• Telgarsky (2016). Benefits of depth in neural networks. A clean proof that some functions need exponentially more units when the network is shallow.
• Ioffe & Szegedy (2015). Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. The original method, including inference with population statistics.
• Ba, Kiros & Hinton (2016). Layer Normalization. Per-example normalization, designed for recurrent networks.
• Santurkar, Tsipras, Ilyas & Madry (2018). How Does Batch Normalization Help Optimization? Evidence against the covariate-shift story, and the smoothness explanation.
• Srivastava et al. (2014). Dropout: A Simple Way to Prevent Neural Networks from Overfitting. The method and its ensemble interpretation.
• Belkin, Hsu, Ma & Mandal (2019). Reconciling modern machine learning practice and the bias-variance trade-off and Nakkiran et al. (2019), Deep Double Descent: Where Bigger Models and More Data Hurt. The double-descent curve, in simple models and in deep networks.
• Andrej Karpathy (2019). A Recipe for Training Neural Networks. The debugging habits this section condenses.

Watch next

▶

▶

▶