Backpropagation by hand: MLPs, losses and layer normalization

Before this lesson: lesson 03 (Jacobians, Hessians, Taylor expansion), lesson 04 (chain rule, VJPs of $Wx$ and of elementwise maps, fan-out), lesson 05 (reverse mode). Outer products from linear-algebra lesson 01. The softmax Jacobian is derived in linear-algebra lesson 16. Time: about 90 minutes, plus the notebook. This is the heaviest lesson of the track; take a break after section 3.3.

---

Why this matters #

Every training step of every neural network runs a backward pass. Frameworks do it for you, but you need the formulas when you write a custom layer or a fused GPU kernel, when a loss turns into nan, when gradients vanish in a deep model, or when an interviewer asks you to "derive backprop for an MLP" or "the gradient of softmax cross-entropy" on a whiteboard. These are among the most common technical interview questions in ML.

The losses of this lesson are most of the loss code in practice: least squares in regression, logistic loss in binary classification, softmax cross-entropy in every classifier and every large language model (LLM). Layer normalization (LayerNorm) sits in every transformer block. In ML for science, weighted least squares is the data-misfit term of data assimilation and seismic inversion, and LayerNorm is inside transformer-based weather models.

1. Predict first #

2. The picture #

Backprop is reverse mode on a chain of layers. An MLP computes $x \to z^1 \to a^1 \to z^2 \to \cdots \to z^L \to \ell$. Lesson 05 showed that reverse mode runs the chain once forward (and stores the values), then once backward with the seed $\partial\ell/\partial\ell = 1$. On the way back we carry the error signal of each layer,

Each layer does two jobs on the backward sweep:

1. Pass the error down: $\delta^l$ is the VJP of the layer above. Multiply $\delta^{l+1}$ by $(W^{l+1})^\top$, then multiply entrywise by the slope $\varphi'(z^l)$ of the activation.
2. Pay the weights: the gradient of $W^l$ is the outer product $\delta^l (a^{l-1})^\top$ of the error signal and the input of the layer.

Interactive figure: loading…

Prediction minus target. Let $z$ be the score (or the vector of scores) that a model outputs, and $y$ the target. The three common losses have the same gradient form:

Loss	Formula in the score $z$	Prediction	$\partial\ell/\partial z$
squared error	$\tfrac12 (z - y)^2$	$z$	$z - y$
logistic, $y \in \{0, 1\}$	$\log(1 + e^z) - yz$	$\sigma(z) = 1/(1 + e^{-z})$	$\sigma(z) - y$
softmax cross-entropy, $C$ classes	$\log\sum_k e^{z_k} - z_y$	$p = \operatorname{softmax}(z)$	$p - e_y$

Here $e_y$ is the one-hot vector of the true class $y$: 1 in entry $y$, 0 elsewhere. If the prediction equals the target, the gradient is zero. For the logistic and softmax losses, a model that is sure and wrong gets gradient entries close to $\pm 1$, the largest possible size, so the next step corrects it strongly.

Why the same form? (generalized linear models)

Each loss has the form $\ell(z) = A(z) - y^\top z$ (plus a constant), with $A(z) = \tfrac12 z^2$, $A(z) = \log(1 + e^z)$ and $A(z) = \log\sum_k e^{z_k}$ (in the last case the target is the one-hot vector $e_y$). These functions $A$ are the log-partition functions of the Gaussian (unit variance), Bernoulli and categorical distributions. A standard fact about exponential families: $\nabla A(z)$ is the mean of the distribution (the prediction) and $\nabla^2 A(z)$ is its covariance: $1$, $\sigma(1 - \sigma)$ and $\operatorname{diag}(p) - pp^\top$. A covariance matrix is positive semidefinite (PSD), so every one of these losses is convex in the scores. These are the canonical-link generalized linear models of the ML fundamentals track (coming later).

Interactive figure: loading…

3. Definitions and main results #

3.1 The MLP and its notation #

We use column vectors for one sample (the same notation as lessons 03–05).

• Input $a^0 = x \in \R^{n_0}$.
• Hidden layers $l = 1, \dots, L-1$: the pre-activation $z^l = W^l a^{l-1} + b^l$ and the activation $a^l = \varphi(z^l)$, with $W^l \in \R^{n_l \times n_{l-1}}$, $b^l \in \R^{n_l}$ and an elementwise function $\varphi$ (ReLU, tanh, ...).
• Output layer: the logits $z^L = W^L a^{L-1} + b^L \in \R^C$, one score per class.
• Loss: softmax cross-entropy $\ell = -\log \operatorname{softmax}(z^L)_y$ for the true class $y \in \{1, \dots, C\}$.
• Error signal $\delta^l = \partial\ell/\partial z^l \in \R^{n_l}$.

$\odot$ is the entrywise product. Two facts from lesson 04: the VJP of $z = Wa$ is $u \mapsto W^\top u$, and the VJP of an elementwise map $a = \varphi(z)$ is $u \mapsto \varphi'(z) \odot u$.

3.2 Theorem 1 (backpropagation for an $L$-layer MLP) #

With the notation of 3.1:

• (a) $\delta^L = p - e_y$, where $p = \operatorname{softmax}(z^L)$;
• (b) $\delta^l = \varphi'(z^l) \odot (W^{l+1})^\top \delta^{l+1}$ for $l = L-1, \dots, 1$;
• (c) $\dfrac{\partial\ell}{\partial W^l} = \delta^l (a^{l-1})^\top$ (an outer product, $n_l \times n_{l-1}$) and $\dfrac{\partial\ell}{\partial b^l} = \delta^l$;
• (d) $\dfrac{\partial\ell}{\partial x} = (W^1)^\top \delta^1$.

Proof.

1. Output. $\ell = -z^L_y + \log\sum_k e^{z^L_k}$. So $\partial\ell/\partial z^L_k = -[k = y] + e^{z^L_k}/\sum_j e^{z^L_j}$, which is $p_k - [k = y]$. Here $[k = y]$ is 1 if $k = y$ and 0 otherwise. This is (a). (Another loss changes only this step: for $\ell = \tfrac12\|z^L - t\|^2$ we get $\delta^L = z^L - t$.)
2. Layer Jacobian. $\ell$ depends on $z^l$ only through $z^{l+1} = W^{l+1}\varphi(z^l) + b^{l+1}$. By the chain rule, $\partial z^{l+1}/\partial z^l = W^{l+1}\operatorname{diag}(\varphi'(z^l))$.
3. Recursion. The gradient through a map is the VJP with its Jacobian. So $\delta^l = \big(W^{l+1}\operatorname{diag}(\varphi'(z^l))\big)^\top \delta^{l+1}$. A diagonal matrix times a vector is an entrywise product, so this is $\varphi'(z^l) \odot (W^{l+1})^\top \delta^{l+1}$: rule (b).
4. Weights. $\ell$ depends on $W^l$ only through $z^l$, and $z^l_k = \sum_j W^l_{kj} a^{l-1}_j + b^l_k$. So $\partial z^l_k/\partial W^l_{ij} = [k = i]\, a^{l-1}_j$. In the sum $\partial\ell/\partial W^l_{ij} = \sum_k \delta^l_k [k = i]\, a^{l-1}_j$ only $k = i$ survives, so $\partial\ell/\partial W^l_{ij} = \delta^l_i a^{l-1}_j$. In matrix form this is $\delta^l (a^{l-1})^\top$. In the same way, $\partial z^l_k/\partial b^l_i = [k = i]$ gives $\partial\ell/\partial b^l = \delta^l$. This is (c).
5. Input. $\ell$ depends on $x$ only through $z^1 = W^1 x + b^1$, so $\partial\ell/\partial x = (W^1)^\top \delta^1$. This is (d).
6. Order. Step 1 gives $\delta^L$. Step 3 gives each $\delta^l$ from $\delta^{l+1}$, for $l = L-1$ down to 1. Steps 4 and 5 then give all parameter and input gradients. $\blacksquare$

Read it. The backward pass is reverse accumulation on the chain $x \to z^1 \to \cdots \to z^L \to \ell$. Unrolled, it is a product of transposed layer Jacobians:

Step 1 differentiates the softmax and the loss together. So the $C \times C$ softmax Jacobian $\operatorname{diag}(p) - pp^\top$ is never formed; the backward pass costs $O(C)$ at the output.

What must be stored. Rule (b) needs $\varphi'(z^l)$ and rule (c) needs $a^{l-1}$. So the forward pass keeps every activation until the backward pass reaches its layer. For tanh, $\varphi'(z) = 1 - \tanh^2 z = 1 - a^2$, and for ReLU, $\mathbf 1[z > 0] = \mathbf 1[a > 0]$, so storing $a^l$ is enough.

Answer to question 1 of section 1. $\partial\ell/\partial W = \delta a^\top$ has the shape of $W$, $512 \times 256$. Its rank is at most 1, because it is an outer product: every column is a multiple of $\delta$. It is exactly 1 when $\delta \ne 0$ and $a \ne 0$.

3.3 Batch form #

Now take a batch of $B$ samples and the mean loss $\ell = \frac1B\sum_{i=1}^B \ell_i$. Put one sample per row: $A^0 = X \in \R^{B \times n_0}$ and

where $\mathbf 1 \in \R^B$ is the vector of ones. Row $i$ of $Z^l$ is $(z^l_i)^\top$, the pre-activation of sample $i$. This is the layout of PyTorch: nn.Linear stores $W^l$ with shape (out, in) and computes A @ W.T + b. Let $\Delta^l = \partial\ell/\partial Z^l \in \R^{B \times n_l}$. Its row $i$ is $(\delta^l_i)^\top / B$, because $\ell_i$ enters the mean with weight $1/B$.

1. $\Delta^L = (P - Y)/B$, with $P$ the row-wise softmax of $Z^L$ and $Y$ the one-hot labels as rows.
2. $\Delta^l = (\Delta^{l+1} W^{l+1}) \odot \varphi'(Z^l)$: rule (b), transposed and stacked as rows.
3. $\partial\ell/\partial W^l = (\Delta^l)^\top A^{l-1}$. By the outer-product view of the matrix product (linear-algebra lesson 01), this is $\sum_i (\text{row } i \text{ of } \Delta^l)^\top (\text{row } i \text{ of } A^{l-1})$: the sum of the per-sample outer products.
4. $\partial\ell/\partial b^l = (\Delta^l)^\top \mathbf 1$: the column sums over the batch.

The factor $1/B$ enters once, in $\Delta^L$. After that every step is a plain product or sum.

Shapes are your test. $\partial\ell/\partial W^l$ must have the shape of $W^l$ ($n_l \times n_{l-1}$). Some code stores the transposed matrix $W = (W^l)^\top$ with shape (in, out) and computes $Z = AW + \mathbf 1 b^\top$. Then step 3 reads $A^\top\Delta$ and step 2 multiplies by $W^\top$. Derive in one layout, then check the shapes in the other.

Cost and memory. Per layer, the forward pass does one matrix product. The backward pass does two of the same size: $\Delta^l W^l$ (the error for the layer below) and $(\Delta^l)^\top A^{l-1}$ (the weight gradient). So the backward pass costs about twice the forward pass, and a full training step about three forward passes. This is the "small constant" of reverse mode in lesson 05. The price is memory: all $A^{l-1}$ and $\varphi'(Z^l)$ stay stored until the backward pass uses them. Activation memory grows with depth and batch size; checkpointing (lesson 05) trades it for recomputation.

3.4 Theorem 2 (least squares) #

Now three losses on a linear model. Notation: $n$ samples with $d$ features; $X \in \R^{n \times d}$ has one sample $x_i^\top$ per row; $w \in \R^d$ are the weights; the scores are $z = Xw$; the loss is a mean over the samples. (Here the batch is the whole data set; for a mini-batch replace $n$ by $B$.) We use one fact from the optimization track: a twice continuously differentiable function on $\R^d$ is convex if and only if its Hessian is PSD at every point.

Theorem 2. $L(w) = \frac{1}{2n}\|Xw - y\|^2$ has

So $L$ is convex. It is strictly convex if and only if $X$ has full column rank.

Proof.

1. Let $r = Xw - y$ (the residual). For a step $h \in \R^d$, expand the square in $L(w + h) = \frac{1}{2n}\|r + Xh\|^2$. This gives $L(w + h) = L(w) + \tfrac1n r^\top X h + \tfrac{1}{2n} h^\top X^\top X h$.
2. $L$ is quadratic, so this Taylor expansion (lesson 03) is exact. The linear term is $\nabla L^\top h$, so $\nabla L = \frac1n X^\top r$. The quadratic term is $\frac12 h^\top \nabla^2 L\, h$, so $\nabla^2 L = \frac1n X^\top X$.
3. $h^\top X^\top X h = \|Xh\|^2 \ge 0$, so the Hessian is PSD and $L$ is convex.
4. $\|Xh\|^2 = 0$ exactly for $h$ in the null space of $X$. So the Hessian is positive definite if and only if that null space is $\{0\}$, that is, $X$ has full column rank. $\blacksquare$

$\nabla L = 0$ gives the normal equations $X^\top X w = X^\top y$. Solve them with QR or lstsq, not with $(X^\top X)^{-1}$ (linear-algebra lesson 12). A quadratic equals its own second-order Taylor model, so one Newton step $w - (\nabla^2 L)^{-1}\nabla L$ lands exactly on the minimizer.

3.5 Theorem 3 (logistic regression) #

Labels $y_i \in \{0, 1\}$ and scores $z_i = x_i^\top w$. The model says $P(y = 1) = \sigma(z)$. The negative log-likelihood of one sample is $\ell(z) = -y\log\sigma(z) - (1 - y)\log(1 - \sigma(z))$. It simplifies to $\log(1 + e^z) - yz$, because $\log\sigma(z) = z - \log(1 + e^z)$ and $\log(1 - \sigma(z)) = -\log(1 + e^z)$. So

with $S = \operatorname{diag}(s_i)$ and $s_i = \sigma(z_i)(1 - \sigma(z_i))$. $L$ is convex, and strictly convex if $X$ has full column rank.

Proof.

1. $\frac{d}{dz}\log(1 + e^z) = e^z/(1 + e^z) = \sigma(z)$. So $\frac{d\ell}{dz} = \sigma(z) - y$.
2. $\sigma'(z) = \sigma(z)(1 - \sigma(z))$ (lesson 01). So $\frac{d^2\ell}{dz^2} = s > 0$.
3. Chain rule through $z = Xw$ (Jacobian $X$): $\nabla L = \frac1n X^\top(\sigma(z) - y)$.
4. The gradient is $\frac1n \sum_i (\sigma(z_i) - y_i)\, x_i$. Its Jacobian in $w$ is $\frac1n \sum_i s_i\, x_i x_i^\top = \frac1n X^\top S X$.
5. For any $v$: $v^\top \nabla^2 L\, v = \frac1n \sum_i s_i (x_i^\top v)^2 \ge 0$, since every $s_i > 0$. If $X$ has full column rank and $v \ne 0$, some $x_i^\top v \ne 0$, so the sum is $> 0$. $\blacksquare$

Two remarks on the curvature. First, $s_i \le \tfrac14$, so $\nabla^2 L \preceq \frac{1}{4n} X^\top X \preceq \frac{\lambda_{\max}(X^\top X)}{4n} I$: the loss is smooth, and gradient descent with the fixed step $\eta = 4n/\lambda_{\max}(X^\top X)$ decreases it at every step (the optimization track proves this). Second, $s_i \to 0$ when $|z_i| \to \infty$: the curvature vanishes far from the origin. This matters on separable data.

Separable data: no minimizer. The data are linearly separable if some $w_{\text{sep}}$ has $x_i^\top w_{\text{sep}} > 0$ for every $y_i = 1$ and $x_i^\top w_{\text{sep}} < 0$ for every $y_i = 0$. Along the ray $t\,w_{\text{sep}}$, each sample loss is $\log(1 + e^{-t|x_i^\top w_{\text{sep}}|})$. It decreases to 0 as $t \to \infty$ but never reaches 0. So no finite minimizer exists, and gradient descent lets $\|w\|$ grow without bound. Strict convexity says only "at most one minimizer", not "a minimizer exists". An L2 penalty $\tfrac\lambda2\|w\|^2$ adds $\lambda I$ to the Hessian and gives a unique minimizer. (Newton's method on this loss is called IRLS, iteratively reweighted least squares.)

Interactive figure: loading…

Stable computation. $\log(1 + e^{800})$ overflows to inf in float64: $e^z$ overflows above $z \approx 709.8$ (in float32 above $z \approx 88.7$). Use the softplus form $\log(1 + e^z) = \max(z, 0) + \log(1 + e^{-|z|})$ (np.logaddexp(0, z)): the exponent is never positive. Compute $\sigma(z)$ as $1/(1 + e^{-z})$ for $z \ge 0$ and as $e^z/(1 + e^z)$ for $z < 0$.

3.6 Theorem 4 (softmax cross-entropy) #

One sample: logits $z \in \R^C$, true class $y$, $p_k = e^{z_k}/\sum_j e^{z_j}$. Write $\operatorname{lse}(z) = \log\sum_k e^{z_k}$ (the log-sum-exp function). Then

For a batch with logits $Z \in \R^{B \times C}$ (one row $z_b$ per sample), labels $y_b$ and the mean loss $L = \frac1B \sum_b \ell(z_b)$:

where $P$ is the row-wise softmax of $Z$ and $Y$ holds the one-hot labels as rows.

Proof.

1. $\partial \operatorname{lse}/\partial z_j = e^{z_j}/\sum_k e^{z_k} = p_j$: the gradient of log-sum-exp is the softmax.
2. $\partial(-z_y)/\partial z_j = -1$ if $j = y$ and 0 otherwise. So $\nabla_z \ell = p - e_y$.
3. Row $b$ of $Z$ enters only the term $\ell(z_b)$, with the factor $1/B$. So row $b$ of $\partial L/\partial Z$ is $(p_b - e_{y_b})^\top/B$. This is the matrix formula.
4. The Hessian is the Jacobian of $p$: $\partial p_i/\partial z_j = p_i([i = j] - p_j)$, that is, $\operatorname{diag}(p) - pp^\top$ (linear-algebra lesson 16).
5. For any $v$: $v^\top(\operatorname{diag}(p) - pp^\top)v$ equals $\sum_k p_k v_k^2 - \big(\sum_k p_k v_k\big)^2$. This is the variance of $v_k$ when the index $k$ is drawn with probabilities $p$. So it is $\ge 0$, and $\ell$ is convex in $z$. It is 0 for $v = \mathbf 1$: adding a constant to all logits changes nothing. $\blacksquare$

Properties you will use. The entries of $p - e_y$ sum to 0 (shift invariance). Each entry lies in $[-1, 1]$, however large the logits are. For softmax regression, $Z = XW$ with $W \in \R^{d \times C}$, the VJP of a linear map gives $\partial L/\partial W = X^\top (P - Y)/B$.

Stable computation. Let $m = \max_k z_k$. Since $\sum_k e^{z_k} = e^m \sum_k e^{z_k - m}$,

Every exponent is $\le 0$, so nothing overflows. The largest term is $e^0 = 1$, so the sum is $\ge 1$ and its log is never $\log 0$. Then compute $\ell = \operatorname{lse}(z) - z_y$ (this is $-$log_softmax). Do not compute $-\log p_y$ from the probabilities: even a stable softmax gives $p_y = 0$ when $z_y$ is more than about 745 below the largest logit ($e^{-746}$ rounds to 0 in float64), and then $-\log 0 = \infty$. This is why torch.nn.functional.cross_entropy takes logits, not probabilities. (See linear-algebra lesson 13 for overflow and underflow.)

Answer to question 2 of section 1. For $y = 1$: $e^{1000}$ overflows to inf, and inf/inf is nan, so the loss is nan. For $y = 2$: $p_2 = 1/\infty = 0$ and $-\log 0 = \infty$. The true losses are $\log(1 + e^{-1000}) \approx 0$ and $1000 + \log(1 + e^{-1000}) \approx 1000$. The log-sum-exp form gives both.

3.7 Theorem 5 (layer normalization backward) #

Layer normalization acts on one row $x \in \R^D$ (the $D$ features of one token). In this section $\sigma$ is a standard deviation, not the sigmoid:

$v$ is the biased variance (divide by $D$), as in PyTorch's layer_norm. $\varepsilon \ge 0$ is a small constant (PyTorch default $10^{-5}$); we need $v + \varepsilon > 0$. $\gamma, \beta \in \R^D$ are learned.

The picture: two invariances. Adding a constant $c$ to every entry of $x$ does not change $\hat x$. So a step along $\mathbf 1$ does not change the loss, and the gradient has no component along $\mathbf 1$: its entries sum to 0. With $\varepsilon = 0$, multiplying $x$ by a positive number does not change $\hat x$ either, so the gradient also has no component along $x - \mu\mathbf 1$ (that is, along $\hat x$). The backward pass removes these two components from the incoming gradient.

Theorem 5. Let $g = \partial L/\partial y$ be the upstream gradient and $\hat g = \gamma \odot g$. Then

where the last formula holds for each row, and mean averages over the $D$ features.

Proof of Theorem 5 (two simple maps and the chain rule)

1. Affine part. $y_j = \gamma_j \hat x_j + \beta_j$. So for one row $\partial L/\partial \beta_j = g_j$, $\partial L/\partial \gamma_j = g_j \hat x_j$ and $\partial L/\partial \hat x = \gamma \odot g = \hat g$. The same $\gamma, \beta$ act on every row, so their gradients are summed over the rows (fan-out, lesson 04).
2. Split the rest. $\hat x = N(\Pi x)$, with the centering matrix $\Pi = I - \tfrac1D \mathbf 1 \mathbf 1^\top$ (so $\Pi x = x - \mu \mathbf 1$) and $N(c) = c/\rho(c)$, $\rho(c) = \sqrt{\|c\|^2/D + \varepsilon}$. At $c = \Pi x$ we have $\|c\|^2/D = v$, so $\rho = \sigma$ and $c/\sigma = \hat x$.
3. Jacobian of $N$. $\partial \rho/\partial c_j = c_j/(D\rho)$. So by the quotient rule $\partial N_i/\partial c_j = [i = j]/\rho - c_i c_j/(D\rho^3)$. In matrix form, $J_N = \tfrac1\sigma\big(I - \tfrac1D \hat x \hat x^\top\big)$, a symmetric matrix.
4. VJP through $N$. $u = J_N^\top \hat g = \tfrac1\sigma\big(\hat g - \hat x\,(\hat x^\top \hat g)/D\big)$. Since $(\hat x^\top \hat g)/D = \operatorname{mean}(\hat g \odot \hat x)$, this is $u = \tfrac1\sigma\big(\hat g - \hat x\, \operatorname{mean}(\hat g \odot \hat x)\big)$.
5. VJP through $\Pi$. $\Pi$ is symmetric and $\Pi u = u - \operatorname{mean}(u)\mathbf 1$. Since $\operatorname{mean}(\hat x) = 0$, we get $\operatorname{mean}(u) = \operatorname{mean}(\hat g)/\sigma$. So $\partial L/\partial x = \Pi u$ is the formula. $\blacksquare$

Checks.

• Row sums: sum the formula over $j$. The first two terms give $\sum_j \hat g_j - D\operatorname{mean}(\hat g) = 0$. The third gives $\operatorname{mean}(\hat g \odot \hat x)\sum_j \hat x_j = 0$, because $\sum_j \hat x_j = 0$. So $\sum_j \partial L/\partial x_j = 0$, for every $\varepsilon$.
• Orthogonality to $\hat x$: using $\|\hat x\|^2 = Dv/(v + \varepsilon)$, one finds that $\hat x^\top \partial L/\partial x$ equals $\tfrac1\sigma (\hat x^\top \hat g)$ times $\varepsilon/(v + \varepsilon)$. This is exactly 0 for $\varepsilon = 0$ (scale invariance) and small when $v \gg \varepsilon$. When the row variance $v$ is close to $\varepsilon$, it is not small.

Cost. $O(D)$ per row. Cache $\hat x$ and $1/\sigma$ from the forward pass. Never build the $D \times D$ Jacobian.

Answer to question 3 of section 1. $\sum_j \partial L/\partial x_j = 0$. Layer norm of $x + c\mathbf 1$ equals layer norm of $x$, so $0 = \frac{d}{dc} L(x + c\mathbf 1)\big|_{c=0} = \mathbf 1^\top \nabla_x L$.

Interactive figure: loading…

3.8 Vanishing and exploding gradients #

The unrolled form of Theorem 1 shows that $\delta^1$ is $\delta^L$ multiplied by $L - 1$ factors $J_l^\top = \operatorname{diag}(\varphi'(z^l))\,(W^{l+1})^\top$. Norms of products are at most products of norms, so

If every factor has norm at most $q < 1$, then $\|\delta^1\| \le q^{L-1}\|\delta^L\|$: the gradient of the first layers shrinks exponentially with depth (vanishing gradients). The sigmoid has $\sigma' \le 1/4$, so a deep sigmoid network vanishes unless the weights are large. When the factors stretch vectors on average by more than 1, the norm typically grows exponentially (exploding gradients). Careful initialization (variance $1/n$ for tanh, $2/n$ for ReLU), residual connections ($y = x + F(x)$ adds an identity path, so its Jacobian is $I + DF$) and normalization layers keep this product near norm 1. The deep learning track treats these methods in detail.

Interactive figure: loading…

4. Worked example #

(a) One sample through a small MLP. Two inputs, three hidden ReLU units, one output, squared loss:

with $\ell = \tfrac12(z^2 - t)^2$ and $\varphi = \operatorname{ReLU}$, so $\varphi'(z) = \mathbf 1[z > 0]$.

Forward. Row by row, $z^1 = W^1 x$ is

Unit 2 is off. Then $z^2 = W^2 a^1 + b^2 = 1 - 0 + 2 + 0.5 = 3.5$ and $\ell = \tfrac12(3.5 - 1.5)^2 = 2$.

Backward.

1. Output error: $\delta^2 = z^2 - t = 2$.
2. Gradient of $W^2$: $\delta^2 (a^1)^\top = 2\,(1, 0, 1) = (2, 0, 2)$. Also $\partial\ell/\partial b^2 = 2$.
3. Pass the error down: $(W^2)^\top\delta^2 = (2, -2, 4)$. The mask $\varphi'(z^1) = (1, 0, 1)$ gives $\delta^1 = (2, 0, 4)$.
4. Gradient of $W^1$ (an outer product), and $\partial\ell/\partial b^1 = \delta^1 = (2, 0, 4)$:

5. Gradient of the input:

The row of $W^1$ for the dead unit gets a zero gradient. Each weight gradient has rank 1.

The same result as a Jacobian product. The row vector $\nabla_x\ell^\top$ is the product of the local Jacobians, from the loss back to the input:

(b) A batch of two. Add $x' = (0, 1)$ with target $t' = -0.5$.

Forward for sample 2. $z^{1\prime} = (1, 3, -1)$, $a^{1\prime} = (1, 3, 0)$, $z^{2\prime} = 1 - 3 + 0 + 0.5 = -1.5$ and $\ell' = \tfrac12(-1)^2 = 0.5$. The mean loss is $(2 + 0.5)/2 = 1.25$.

Backward for the batch (one row per sample):

and $\partial\ell/\partial b^2 = 1 - 0.5 = 0.5$. Then

and $\partial\ell/\partial b^1 = (0.5, 0.5, 2)$ (the column sums of $\Delta^1$).

Check. Sample 2 alone gives $\delta^{1\prime} = (-1, 1, 0)$. The mean of its outer product $\delta^{1\prime}x'^\top$ and the sample-1 gradient from (a) is the same matrix:

(c) Softmax cross-entropy for a batch, $B = 2$, $C = 3$.

• Row 1: $z_1 = (\ln 4, \ln 2, \ln 2)$ and $y_1 = 2$. Then $e^{z_1} = (4, 2, 2)$, $p_1 = (1/2, 1/4, 1/4)$ and $\ell_1 = -\log\tfrac14 = 1.3863$.
• Row 2: $z_2 = (0, 0, 0)$ and $y_2 = 3$. Then $p_2 = (1/3, 1/3, 1/3)$ and $\ell_2 = \log 3 = 1.0986$.
• The mean loss is $L = (1.3863 + 1.0986)/2 = 1.2425$, and

Each row sums to 0. The Hessian of $\ell_1$ is

with eigenvalues $0$, $0.25$ and $0.375$. The zero eigenvalue has the eigenvector $\mathbf 1$.

(d) Layer norm, $D = 4$, $\varepsilon = 0$ (to keep the numbers exact), $x = (1, 1, 5, 5)$, $\gamma = (1, 1, 2, 1)$, $\beta = 0$, upstream $g = (1, 2, 1, 0)$.

• Forward: $\mu = 3$, $v = (4 + 4 + 4 + 4)/4 = 4$, $\sigma = 2$, $\hat x = (-1, -1, 1, 1)$.
• $\partial L/\partial \beta = g = (1, 2, 1, 0)$.
• $\partial L/\partial \gamma = g \odot \hat x = (-1, -2, 1, 0)$.
• $\hat g = \gamma \odot g = (1, 2, 2, 0)$, so $\operatorname{mean}(\hat g) = 5/4$ and $\operatorname{mean}(\hat g \odot \hat x) = (-1 - 2 + 2 + 0)/4 = -1/4$.
• The bracket of Theorem 5 and the result:

Checks: the entries sum to 0, and $\hat x \cdot \partial L/\partial x = 0.25 - 0.25 + 0.5 - 0.5 = 0$.

5. In code #

NumPy checks examples (b), (c) and (d) with hand-written backward passes. The comments show the printed output.

import numpy as np
from scipy.special import logsumexp

# (a) + (b): the 2-3-1 ReLU network of section 4, batch form (one sample per row)
W1 = np.array([[2., 1.], [1., 3.], [0., -1.]]); b1 = np.zeros(3)
W2 = np.array([[1., -1., 2.]]);                 b2 = np.array([0.5])
X = np.array([[1., -1.], [0., 1.]]);  T = np.array([[1.5], [-0.5]])
B = len(X)

Z1 = X @ W1.T + b1;  A1 = np.maximum(Z1, 0)        # forward, store Z1 and A1
Z2 = A1 @ W2.T + b2; L = 0.5 * np.mean(np.sum((Z2 - T) ** 2, axis=1))
D2 = (Z2 - T) / B                                    # Delta^2: the only 1/B
dW2, db2 = D2.T @ A1, D2.sum(0)
D1 = (D2 @ W2) * (Z1 > 0)                            # Delta^1
dW1, db1 = D1.T @ X, D1.sum(0)
print(L, dW2, db2)        # 1.25 [[ 0.5 -1.5  1. ]] [0.5]
print(dW1.tolist())       # [[1.0, -1.5], [0.0, 0.5], [2.0, -2.0]]

# (c) softmax cross-entropy: (P - Y) / B, computed with log-sum-exp
Z = np.array([[np.log(4), np.log(2), np.log(2)], [0., 0., 0.]])
y = np.array([1, 2])                                  # classes 2 and 3 (code counts from 0)
lse = logsumexp(Z, axis=1)
loss = np.mean(lse - Z[np.arange(2), y])
P = np.exp(Z - lse[:, None]); Y = np.eye(3)[y]
print(round(loss, 4), ((P - Y) / 2).round(4).tolist())
# 1.2425 [[0.25, -0.375, 0.125], [0.1667, 0.1667, -0.3333]]

# naive vs stable at z = (1000, 0)
z = np.array([1000., 0.])
with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
    print(-np.log(np.exp(z) / np.exp(z).sum()))       # [nan inf]
print(logsumexp(z) - z)                               # [   0. 1000.]

# (d) layer norm backward (eps = 0), and its two checks
def layernorm_backward(x, gamma, g, eps=0.0):
    mu, var = x.mean(), x.var()                       # biased variance, as in PyTorch
    sigma = np.sqrt(var + eps); xh = (x - mu) / sigma
    gh = gamma * g
    return (gh - gh.mean() - xh * np.mean(gh * xh)) / sigma, xh

x = np.array([1., 1., 5., 5.])
gamma = np.array([1., 1., 2., 1.])
g = np.array([1., 2., 1., 0.])                        # upstream gradient dL/dy
dx, xh = layernorm_backward(x, gamma, g)
print(dx, dx.sum(), xh @ dx)                          # [-0.25  0.25  0.5  -0.5 ] 0.0 0.0

In PyTorch, autograd gives the same numbers. Note the weight layout of nn.Linear (out, in) and that cross_entropy takes logits:

import torch
import torch.nn.functional as F
torch.set_default_dtype(torch.float64)

# The same 2-3-1 network as nn.Linear layers: weight shape (out, in), Z = A W^T + b.
net = torch.nn.Sequential(torch.nn.Linear(2, 3), torch.nn.ReLU(), torch.nn.Linear(3, 1))
with torch.no_grad():
    net[0].weight.copy_(torch.tensor([[2., 1.], [1., 3.], [0., -1.]])); net[0].bias.zero_()
    net[2].weight.copy_(torch.tensor([[1., -1., 2.]]));               net[2].bias.fill_(0.5)
X = torch.tensor([[1., -1.], [0., 1.]]); T = torch.tensor([[1.5], [-0.5]])
loss = 0.5 * F.mse_loss(net(X), T)       # mse_loss is the mean of squares
loss.backward()
print(loss.item(), net[0].weight.grad.tolist())
# 1.25 [[1.0, -1.5], [0.0, 0.5], [2.0, -2.0]]

# Softmax cross-entropy takes LOGITS; its gradient is (P - Y) / B.
Z = torch.tensor([[1000., 0., 0.], [0., 0., 0.]], requires_grad=True)
ce = F.cross_entropy(Z, torch.tensor([1, 2]))       # finite: log-sum-exp inside
ce.backward()
print(round(ce.item(), 4), Z.grad.numpy().round(4).tolist())
# 500.5493 [[0.5, -0.5, 0.0], [0.1667, 0.1667, -0.3333]]

The loss of the first row is $\operatorname{lse}(1000, 0, 0) - 0 \approx 1000$, finite, and its gradient row is $(p - e_2)/2 \approx (1, -1, 0)/2$.

6. Where it shows up #

• Every autodiff framework. Each operation in PyTorch or JAX has a VJP rule like the ones in this lesson, and backward() runs them in reverse order. A custom torch.autograd.Function asks you to write the VJP by hand. Fused kernels (for example for LayerNorm or cross-entropy) use the closed forms of Theorems 4 and 5 to save memory and time.
• Classifiers and language models. The logit gradient of every classifier and LLM is $(P - Y)/B$. With a vocabulary of $10^5$ tokens, $P$ is the largest tensor of the step, so memory-efficient kernels compute $\operatorname{lse}$ and $p - e_y$ in chunks of the vocabulary.
• Soft targets. For any target distribution $q$ ($\sum_k q_k = 1$), the loss $-\sum_k q_k \log p_k = \operatorname{lse}(z) - q^\top z$ has the gradient $p - q$. This covers label smoothing and knowledge distillation. With a temperature $T$ (logits $z/T$) the gradient is $(\operatorname{softmax}(z/T) - q)/T$. When $q$ is the teacher's softmax at the same temperature, this scales like $1/T^2$ for large $T$, so distillation losses are often multiplied by $T^2$.
• Outer-product structure. A per-sample weight gradient $\delta a^\top$ has rank 1 and Frobenius norm $\|\delta\|\,\|a\|$. So per-sample gradient norms (needed for clipping in differentially private SGD) cost almost nothing. A batch gradient has rank at most $B$. K-FAC approximates the curvature of a layer with the Kronecker factors $\mathbb E[aa^\top]$ and $\mathbb E[\delta\delta^\top]$.
• Newton, Fisher and Gauss–Newton. The logistic Hessian $\frac1n X^\top S X$ does not depend on $y$, so it equals the Fisher information. For a nonlinear model with these losses, the Gauss–Newton matrix $J^\top \nabla^2_z \ell\, J$ uses exactly the score Hessians of this lesson: $I$, $S$ or $\operatorname{diag}(p) - pp^\top$ (lesson 08).
• ML for science. Variational data assimilation and seismic inversion minimize a weighted least squares loss $J(u) = \tfrac12 (Gu - d)^\top \Sigma_o^{-1}(Gu - d)$ $+\ \tfrac12 (u - u_b)^\top \Sigma_b^{-1}(u - u_b)$ (observation operator $G$, data $d$, background state $u_b$, covariances $\Sigma_o$ and $\Sigma_b$). The proof of Theorem 2 gives $\nabla J = G^\top \Sigma_o^{-1}(Gu - d) + \Sigma_b^{-1}(u - u_b)$. For a nonlinear $G$, the transposed Jacobian $DG(u)^\top$ is the adjoint model: one backward run gives the whole gradient. LayerNorm (or its variant RMSNorm) is in every block of transformer-based weather and Earth-system models.

7. Common mistakes #

• Missing transposes. The recursion uses $(W^{l+1})^\top\delta^{l+1}$, not $W^{l+1}\delta^{l+1}$, and $\partial\ell/\partial W^l = \delta^l(a^{l-1})^\top$ has the shape of $W^l$. Shapes catch both mistakes, unless the layer is square.
• The mask of the wrong layer. $\delta^l$ uses $\varphi'(z^l)$, the pre-activation of layer $l$ itself, not $\varphi'(z^{l+1})$.
• Mixing the two layouts. $z = Wa$ (one column per sample) and $Z = AW^\top$ or $Z = AW$ (one row per sample) give transposed formulas. Derive in one layout, then check shapes in the other.
• $1/B$ twice, or not at all. It enters once, in $\Delta^L$. Then $\partial\ell/\partial b$ is a sum over the batch, not another mean. A sum instead of a mean makes the gradient $B$ times larger.
• Naive logs and exps. $\log(1 + e^z)$ and $\log\sum_k e^{z_k}$ overflow when an exponent is above about 710 in float64 (about 89 in float32). Use logaddexp, softplus, log_sigmoid, logsumexp.
• $-\log$ of a probability. Even a stable softmax underflows for unlikely classes. Use $\operatorname{lse}(z) - z_y$; this is why cross_entropy takes logits.
• Sigmoid plus squared error for classification. $\tfrac12(\sigma(z) - y)^2$ has the gradient $(\sigma - y)\,\sigma(1 - \sigma)$, which vanishes when the model is confidently wrong. At $z = -5$ and $y = 1$ it is $-0.0066$; cross-entropy gives $\sigma - y = -0.9933$, about 150 times larger.
• Layer norm with $\mu$ and $\sigma$ treated as constants. That gives $dx = \hat g/\sigma$ and misses two terms. Test: the entries of $dx$ must sum to 0.
• Biased vs unbiased variance. Layer norm uses the biased variance. np.var is biased by default, but torch.var is unbiased by default (correction=1).
• Forgetting to sum $d\gamma$ and $d\beta$ over all rows (all tokens of all sequences in the batch).
• "Strictly convex, so a minimizer exists." Not on separable data: strict convexity gives at most one minimizer, not existence.

8. Check yourself #

1. $\ell = \tfrac12\|\sigma(Wx + b) - t\|^2$ with the elementwise sigmoid $\sigma$. Give $\delta = \partial\ell/\partial z$ for $z = Wx + b$, and $\partial\ell/\partial W$, $\partial\ell/\partial b$, $\partial\ell/\partial x$.
   Answer

   $\delta = (\sigma(z) - t)\odot\sigma(z)\odot(1 - \sigma(z))$, using $\sigma' = \sigma(1 - \sigma)$. Then $\partial\ell/\partial W = \delta x^\top$, $\partial\ell/\partial b = \delta$ and $\partial\ell/\partial x = W^\top\delta$.
2. A batch of $B = 8$ samples passes through a layer with $W^l \in \R^{512\times256}$. What is the largest possible rank of $\partial\ell/\partial W^l$?
   Answer

   8. $(\Delta^l)^\top A^{l-1}$ has inner dimension $B$: it is a sum of 8 outer products, so its rank is at most $\min(8, 512, 256) = 8$.
3. A residual block computes $y = x + F(x)$, and the loss depends on $x$ only through $y$. With $\bar y = \partial\ell/\partial y$, give $\partial\ell/\partial x$. Why does this help deep networks?
   Answer

   $\partial\ell/\partial x = \bar y + DF(x)^\top\bar y = (I + DF(x))^\top\bar y$: two paths, and at a fan-out the contributions add (lesson 04). The identity path passes $\bar y$ back unchanged, so the gradient does not have to go through a long product of layer Jacobians that can shrink to zero.
4. Label smoothing with $\alpha = 0.3$ uses the target $q = (1 - \alpha)e_y + \tfrac{\alpha}{C}\mathbf 1$. For $z = (\ln 4, \ln 2, \ln 2)$ and $y = 2$, what is $\nabla_z \ell$?
   Answer

   $q = (0.1, 0.8, 0.1)$ and $p = (0.5, 0.25, 0.25)$, so $\nabla_z \ell = p - q = (0.4, -0.55, 0.15)$. The entries still sum to 0.
5. Compute $\ell$ for $z = (1000, 1000, 998)$ and $y = 3$ without overflow.
   Answer

   $m = 1000$. Then $\operatorname{lse}(z) = 1000 + \log(e^0 + e^0 + e^{-2})$, which is $1000 + \log 2.1353 = 1000.7586$. So $\ell = 1000.7586 - 998 = 2.7586$.
6. In the logistic example of 3.5, change the labels to $y = (1, 1, 0)$. Is the Hessian at $w = (0, \ln 3)$ still positive definite? Does $L$ have a minimizer?
   Answer

   The Hessian $\tfrac1n X^\top S X$ does not depend on $y$, so it is the same positive definite matrix. But no minimizer exists: $w_{\text{sep}} = (0.5, 1)$ gives the scores $(0.5, 1.5, -0.5)$, positive exactly for the label-1 samples. So $L(t\,w_{\text{sep}}) \to 0$ as $t \to \infty$ without reaching 0 (for $t = 1, 10, 30$: $0.38$, $4.5\cdot10^{-3}$, $2.0\cdot10^{-7}$).

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. What is $\delta^L = \partial\ell/\partial z^L$ for softmax cross-entropy?

   Show answer

2. Write the backprop recursion for the error signal $\delta^l$ of a hidden layer.

   Show answer

3. What is $\partial\ell/\partial W^l$ for one sample, and what is its rank?

   Show answer

4. Batch form (one sample per row, mean loss): what are $\Delta^L$ and $\partial\ell/\partial W^l$?

   Show answer

5. Gradient and Hessian of $L(w) = \frac{1}{2n}\|Xw - y\|^2$?

   Show answer

6. Gradient and Hessian of the mean logistic loss with scores $z = Xw$?

   Show answer

7. Why has logistic regression no minimizer on linearly separable data?

   Show answer

8. Write the stable log-sum-exp formula and say why it does not overflow.

   Show answer

9. Why do the entries of the layer-norm input gradient sum to 0?

   Show answer

10. Why do gradients vanish in a deep sigmoid network?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Derive backprop for an MLP." — Define $\delta^l = \partial\ell/\partial z^l$. Then $\delta^L = p - e_y$, $\delta^l = \varphi'(z^l)\odot(W^{l+1})^\top\delta^{l+1}$, $\partial\ell/\partial W^l = \delta^l(a^{l-1})^\top$, $\partial\ell/\partial b^l = \delta^l$. Check every shape. Then give the batch form with $\Delta^L = (P - Y)/B$.
• "Derive the gradient of softmax cross-entropy." — Write $\ell = \operatorname{lse}(z) - z_y$. The gradient of $\operatorname{lse}$ is the softmax, so $\nabla\ell = p - e_y$; for a batch mean, $(P - Y)/B$.
• "Why subtract the max? Why does cross_entropy take logits?" — $e^z$ overflows, and $-\log p_y$ is $\infty$ when $p_y$ underflows. The shifted $\operatorname{lse}(z) - z_y$ avoids both; the shift is exact because the softmax is shift invariant.
• "Is logistic regression convex? Is a neural network?" — Logistic: yes, the Hessian $\frac1n X^\top S X$ is PSD. A network: no. Permuting the hidden units of a minimizer gives another minimizer with the same loss; for a convex loss the midpoint of the two would also be a minimizer, but the averaged weights are in general worse.
• "Why cross-entropy and not squared error for classification?" — It is the maximum-likelihood loss of the categorical model, it is convex in the logits, and its gradient $p - e_y$ stays large when the model is confidently wrong (squared error on a sigmoid output vanishes there).
• "What is the backward pass of layer norm?" — The formula of Theorem 5; it removes the mean and the $\hat x$ component of $\gamma\odot g$, then divides by $\sigma$.
• "Why do gradients vanish or explode?" — The gradient at the input side is a product of $L - 1$ layer Jacobians; initialization, residual connections and normalization keep it near norm 1.

10. Practice #

The notebook for this lesson runs in Colab. Each exercise has a check cell (it compares your gradient with a central-difference gradient) and a folded solution:

1. softmax_cross_entropy(Z, y): the stable mean loss and its gradient $(P - Y)/B$, finite at logits $\pm 1000$.
2. logistic_regression(X, y, w): loss, gradient and Hessian; then Newton's method on the breast cancer data set of scikit-learn.
3. layernorm_forward and layernorm_backward: check the gradient and its two invariances.
4. mlp_forward and mlp_backward for a 2-layer MLP on the scikit-learn digits; gradient check, then train a few epochs.
5. Measure the gradient norm per layer of a deep tanh MLP at two weight scales.

11. Watch and read #

Backpropagation calculus | Deep Learning Chapter 43Blue1Brown · Neural networks, chapter 4 · 10 min

The chain rule through one neuron per layer, then the indices for many neurons: the same recursion as Theorem 1.

Backpropagation, intuitively | Deep Learning Chapter 33Blue1Brown · Neural networks, chapter 3 · 13 min

The picture before the formulas: how each weight's gradient depends on the activation that feeds it.

Building makemore Part 4: Becoming a Backprop NinjaAndrej Karpathy · Neural Networks: Zero to Hero · 1 h 55 min

A full backward pass by hand, tensor by tensor, including cross-entropy and batch normalization; compare each step with sections 3.3 and 3.6.

Lecture 4 | Introduction to Neural NetworksStanford CS231n · Convolutional Neural Networks for Visual Recognition (2017) · 1 h 14 min

Computational graphs, local gradients and the vectorized backward pass of a linear layer.

• Kevin Murphy, Probabilistic Machine Learning: An Introduction, free PDF: §13.3 (backpropagation, VJPs of common layers), §13.4.2 (vanishing and exploding gradients), §10.2.3 (logistic regression gradient and Hessian), §11.2.2 (least squares).
• Michael Nielsen, Neural Networks and Deep Learning, chapter 2: the four equations of backpropagation with proofs, the same as Theorem 1.
• Goodfellow, Bengio and Courville, Deep Learning, chapter 6, §6.5: back-propagation and other differentiation algorithms.
• Stanford CS231n notes, Backpropagation, intuitions, and Justin Johnson's handout Derivatives, backpropagation, and vectorization (shapes of tensor derivatives).
• Andrej Karpathy, llm.c LayerNorm tutorial: forward and backward pass, what to cache. The lecture notebooks of Neural Networks: Zero to Hero.
• Nick Higham, What is the log-sum-exp function?: the max shift and its accuracy.
• Ba, Kiros and Hinton, Layer Normalization (2016), the original paper.
• Glorot and Bengio, Understanding the difficulty of training deep feedforward neural networks (AISTATS 2010): gradient norms per layer and the initialization that keeps them stable.
• Bright, Edelman and Johnson, Matrix Calculus (for Machine Learning and Beyond): the chain rule and forward vs reverse mode as bracketings of a Jacobian product.