Problems

STATEMENT

Run one forward pass and one backward pass through a small network by hand, producing every intermediate and every parameter gradient. Then check the result against central differences, which is the only way to know a hand-derived gradient is right rather than merely consistent with itself.

GIVEN

A network with two inputs, two sigmoid hidden units and one linear output. One example, so every activation is a row.

$\vec{x} = \begin{pmatrix} 1 & 2 \end{pmatrix}, \qquad y = 1$

The hidden nonlinearity is $\sigma(z) = 1/(1 + e^{-z})$, the head carries none, and the loss is $L = \tfrac{1}{2}(\hat{y} - y)^2$.

FIND

Every intermediate of the forward pass; the adjoints $\delta^{(2)}$ and $\delta^{(1)}$; the gradients with respect to all nine parameters; and the largest disagreement between those gradients and a central-difference estimate.

STRATEGY

Forward first, keeping everything — by Definition 3 the backward pass consumes the stored values, so discarding them means starting again. Then apply (I.6.2) from the sink: the loss adjoint is 1, and each node passes its own adjoint times its local Jacobian to its producers. For this architecture that is exactly the recurrence (I.6.3) and the outer product (I.6.4).

SOLUTION

Step 1 — forward. The pre-activations are $\vec{z}^{(1)} = \vec{x}\mat{W}^{(1)} + \vec{b}^{(1)}$:

$z^{(1)}_1 = 1(0.1) + 2(0.2) + 0 = 0.5, \qquad z^{(1)}_2 = 1(0.3) + 2(0.4) + 0.1 = 1.2$

$\vec{a}^{(1)} = \sigma(\vec{z}^{(1)}) = \begin{pmatrix} 0.62245933 & 0.76852478 \end{pmatrix}$

$\hat{y} = z^{(2)} = 0.62245933(0.5) + 0.76852478(-0.6) + 0.2 = 0.05011480$

$L = \tfrac{1}{2}(0.05011480 - 1)^2 = 0.45114095$

Step 2 — the adjoint at the sink. $\bar{L} = 1$ by definition, and the head is linear, so

$\delta^{(2)} = \frac{\partial L}{\partial z^{(2)}} = \hat{y} - y = -0.94988520$

The sign says the prediction is below the target, which it is.

Step 3 — the second layer’s parameters. By (I.6.4) the weight gradient is the outer product of what entered the layer with what came back:

The bias gradient is the adjoint, because a bias enters its node with coefficient one.

Step 4 — travel back through the layer. By (I.6.3), multiply by $(\mat{W}^{(2)})^{\mathsf{T}}$ and gate by the slope of the nonlinearity:

$\delta^{(2)}(\mat{W}^{(2)})^{\mathsf{T}} = \begin{pmatrix} -0.47494260 & 0.56993112 \end{pmatrix}$

For the logistic function $\sigma' = \sigma(1 - \sigma)$, so it is read off the stored activations rather than recomputed:

$\sigma'(\vec{z}^{(1)}) = \begin{pmatrix} 0.23500371 & 0.17789444 \end{pmatrix}$

Step 5 — the first layer’s parameters. The same outer product, one layer down:

The second row is exactly twice the first, because the second input is twice the first. A weight’s gradient carries the size of the thing that fed it.

Step 6 — the check. Central differences with $h = 10^{-6}$, evaluating the loss twice per parameter and taking $(L(\theta + h) - L(\theta - h))/2h$, agree with all nine analytic gradients. The largest disagreement is $5.79 \times 10^{-11}$.

That agreement is worth more than any of the individual numbers. The analytic gradient came from the chain rule; the numeric one came from evaluating the forward function eighteen times and never differentiating anything. Two unrelated computations landing on the same nine numbers is evidence the derivation is correct.

Check — sanity

The bias gradients equal the adjoints. $\partial L/\partial b^{(2)} = \delta^{(2)}$ and $\partial L/\partial \vec{b}^{(1)} = \delta^{(1)}$, to the last digit. A bias adds into its node with coefficient one, so its local Jacobian is the identity and (I.6.2) leaves the adjoint unchanged. Any discrepancy here is an arithmetic slip, not a modelling question.

The rows of $\partial L/\partial \mat{W}^{(1)}$ are in the ratio of the inputs. $-0.22322655 / -0.11161327 = 2$ exactly, matching $x_2/x_1 = 2$. The outer product in (I.6.4) says the gradient of a weight scales with the activation entering it, so this ratio is forced.

The sigmoid slopes are below $1/4$. $0.23500371$ and $0.17789444$, against the maximum $\sigma'(0) = 0.25$. Every layer of this kind multiplies the travelling adjoint by something no larger than a quarter, which is the observation I.6.B04 turns into a depth limit.

STATEMENT

Derive the two equations every implementation of backpropagation contains — the recurrence that carries an adjoint from one layer to the previous one, and the outer product that turns an adjoint into a weight gradient — from the graph rule (I.6.2) alone. Nothing else may be assumed.

GIVEN

One layer of an MLP, written as three nodes so that each has a Jacobian that can be stated:

with $\vec{a}^{(\ell-1)} \in \R^{1 \times d_{\ell-1}}$, $\mat{W}^{(\ell)} \in \R^{d_{\ell-1} \times d_\ell}$ and $\varphi$ applied elementwise. A scalar loss $L$ sits at the sink. Write $\delta^{(\ell)} = \partial L / \partial \vec{z}^{(\ell)}$.

FIND

Expressions for $\partial L/\partial \mat{W}^{(\ell)}$, $\partial L/\partial \vec{b}^{(\ell)}$ and $\delta^{(\ell-1)}$ in terms of $\delta^{(\ell)}$ and quantities the forward pass already stored.

STRATEGY

Work one node at a time and never write down a full Jacobian. Each step asks the same question: given the adjoint of this node’s output, what is the adjoint of its input? That is the vector–Jacobian product of Definition 7, and for these three operations it has a closed form in each case.

SOLUTION

Step 1 — through the nonlinearity. $\varphi$ acts elementwise, so $\partial a^{(\ell)}_j / \partial z^{(\ell)}_i$ is zero unless $i = j$. The Jacobian is diagonal, and multiplying by a diagonal matrix is multiplying elementwise:

This is the first place the forward pass is consumed: $\varphi'$ is evaluated at the stored $\vec{z}^{(\ell)}$, and without it there is nothing to multiply.

Step 2 — the bias. In $\vec{z} = \vec{a}\mat{W} + \vec{b}$ the bias enters with coefficient one, so $\partial z_j / \partial b_i = \delta_{ij}$ and the local Jacobian is the identity. By (I.6.2) the adjoint passes through unchanged:

Step 3 — the weights. Take one entry. Only the $j$-th output involves $W_{ij}$, and it involves it linearly with coefficient $a^{(\ell-1)}_i$:

So the sum over consumers in (I.6.2) collapses to a single term:

Every entry is a product of one input coordinate and one output adjoint, which is exactly an outer product:

The shapes confirm it: $(d_{\ell-1} \times 1)(1 \times d_\ell)$ gives $d_{\ell-1} \times d_\ell$, the shape of $\mat{W}^{(\ell)}$. A gradient always has the shape of the thing it differentiates.

Step 4 — back to the previous layer. Now differentiate with respect to the input rather than the parameters. Coordinate $a^{(\ell-1)}_i$ feeds every output $z^{(\ell)}_j$, so this time the sum does not collapse:

which is the $i$-th entry of $\delta^{(\ell)}(\mat{W}^{(\ell)})^{\mathsf{T}}$:

The transpose is not a convention. Forwards, $\mat{W}$ maps $d_{\ell-1}$ coordinates to $d_\ell$; backwards, the adjoint must travel from $d_\ell$ back to $d_{\ell-1}$, and the only way to do that with the same matrix is to use it the other way round.

Step 5 — combine. Applying Step 1 at layer $\ell - 1$ to the result of Step 4 gives the recurrence:

which is (I.6.3), and Steps 2 and 3 are (I.6.4). Both came from (I.6.2) and three local Jacobians, and no Jacobian was ever written out in full.

Check — sanity

Every gradient has the shape of its parameter. $(d_{\ell-1} \times 1)(1 \times d_\ell) = d_{\ell-1} \times d_\ell$ for the weights, and $\delta^{(\ell)}$ is $1 \times d_\ell$ for the bias. This is the cheapest check there is and it catches a transposed derivation immediately.

The recurrence terminates. $\delta^{(L)}$ comes from the loss, and each application of the recurrence produces $\delta$ one layer earlier, so $L$ steps reach the input. Nothing in the derivation refers to a layer that has not been reached, which is what makes reverse topological order the right traversal.

Only stored quantities appear. $\vec{a}^{(\ell-1)}$, $\vec{z}^{(\ell-1)}$ and $\mat{W}^{(\ell)}$ — all present after the forward pass. No re-evaluation of $\varphi$ or of any matrix product is required, which is why the backward pass costs about what the forward pass cost rather than twice it.

STATEMENT

Tabulate the shape of every tensor in both directions through a batched MLP, and say where the batch dimension goes. One of the two gradients in (I.6.4) sums over the batch and the other does not; the shapes say which, before any calculus is done.

GIVEN

The network of I.5.B03 — widths $(784, 512, 256, 128, 10)$, four affine layers, ReLU after the first three, no activation on the head — now run on a batch of $B = 32$ examples, so the input is $\mat{A}^{(0)} \in \R^{32 \times 784}$. The loss is a scalar: a mean over the batch.

FIND

For every layer $\ell$: the shapes of $\mat{A}^{(\ell)}$, $\Delta^{(\ell)}$, $\partial L / \partial \mat{W}^{(\ell)}$ and $\partial L / \partial \vec{b}^{(\ell)}$. Then the total activation memory kept for the backward pass, in fp32.

STRATEGY

Two rules settle every row. A gradient has the shape of the thing it differentiates, so parameter gradients cannot carry a batch dimension — the parameters do not have one. An adjoint has the shape of the value it belongs to, so activation adjoints must carry it. Where those two rules disagree, a sum appears, and that is the only place a sum can appear.

SOLUTION

Step 1 — forward shapes. The batch dimension is untouched by every layer:

Step 2 — adjoints. $\Delta^{(\ell)} = \partial L / \partial \mat{Z}^{(\ell)}$ has the shape of $\mat{Z}^{(\ell)}$, so it carries the batch: $32 \times 512$, $32 \times 256$, $32 \times 128$, $32 \times 10$.

Step 3 — the weight gradient absorbs the batch. Per example the gradient is the outer product of (I.6.4). Over a batch it is the sum of those outer products, which is a single matrix product:

with shapes $(d_{\ell-1} \times 32)(32 \times d_\ell) \to d_{\ell-1} \times d_\ell$. The 32 is contracted away — it is the inner dimension of the product. The summation over the batch is not an extra step anybody writes; it is what the matrix product already does.

Step 4 — the bias gradient needs an explicit sum. Per example $\partial L / \partial \vec{b}^{(\ell)} = \delta^{(\ell)}$, of shape $1 \times d_\ell$. Over a batch there is no matrix product to hide the contraction in, so it is written out:

This is the asymmetry the Statement promised. Both gradients sum over the batch; the weight gradient’s sum is the inner dimension of a matrix product and the bias gradient’s is a reduction along an axis. Forgetting the second is a common implementation error and it is caught immediately by a shape check — $32 \times 512$ where $1 \times 512$ was required.

Step 5 — what must be kept. The backward pass consumes $\mat{A}^{(\ell-1)}$ for (I.6.4) and $\mat{Z}^{(\ell)}$ for $\varphi'$. For ReLU the mask $\mathbb{1}[\mat{Z} > 0]$ suffices, but counting stored floats:

At 4 bytes each that is $216\,320$ bytes, about $0.22$ MB.

Step 6 — compare with the parameters. I.5.B03 counted $567\,434$ parameters, $2.27$ MB in fp32. At $B = 32$ the activations are a tenth of the parameters. At $B = 1024$ they are $1\,730\,560$ floats, $6.92$ MB — three times the parameters. Activation memory is linear in the batch and parameter memory is constant in it, so which term dominates is a property of the batch size and not of the model.

Check — sanity

No parameter gradient carries a batch dimension. It cannot: the parameter does not have one, and a gradient has the shape of what it differentiates. Any implementation returning a $32 \times 784 \times 512$ weight gradient has computed per-example gradients and not yet reduced them — sometimes deliberately, which is what per-sample gradient methods want, but never by accident.

Every adjoint does carry it. Different examples have different activations and therefore different adjoints, and they stay separate until a parameter is reached. The batch dimension survives exactly as far as the parameters and no further.

The two contractions are the same contraction. Both gradients sum $B$ terms; only the notation differs. Writing the bias as a row of ones times $\Delta^{(\ell)}$ makes them identical in form, which is a good way to remember that the missing sum is missing rather than unnecessary.

STATEMENT

Backpropagation is exact. Show that exactness is not sufficiency: exhibit a network on which the gradient reaching the first layer is correct to the last bit and numerically zero, and establish the depth at which that happens. Then show that the obvious remedy — larger weights — makes it worse.

GIVEN

A chain of scalar sigmoid units, the simplest object with the property under study:

for $\ell = 1 \dots L$, with one shared weight $w$ and no biases. The adjoint leaving the loss is normalised to $1$.

FIND

A bound on $|\partial L/\partial a^{(0)}|$ in terms of $L$ alone; the measured ratio at $L = 10$ and $L = 40$ for $w = 1$; the smallest $L$ at which the ratio falls below fp32’s smallest normal number; and the effect of raising $w$.

STRATEGY

Apply the recurrence (I.6.3) down the chain and collect factors. Each layer contributes two: the weight, and the slope of the nonlinearity at the stored pre-activation. The second factor has a ceiling that does not depend on anything trainable, and a product of numbers below a ceiling is bounded by the ceiling raised to the depth.

SOLUTION

Step 1 — the bound. For the logistic function $\sigma' = \sigma(1-\sigma)$, maximised where $\sigma = 1/2$:

The run confirms $\sigma'(0) = 0.250000$. Now walk the chain back. Each layer multiplies the travelling adjoint by $w$ and by $\sigma'(z^{(\ell)})$, so

With $w = 1$ this is $4^{-L}$, and nothing about the data, the loss or the training procedure can raise it. It is a property of the architecture.

Step 2 — the measurement. Propagating the adjoint explicitly:

The measured values sit below the bound at every depth, as they must, and the gap widens because the units are not all sitting at their maximum slope.

Step 3 — the depth where it stops being a small number. fp32’s smallest normal value is $1.17549435 \times 10^{-38}$. The first depth at which the ratio falls below it is

At that depth the gradient of the first layer is not small. It is unrepresentable: it enters the subnormal range, where fp32 has almost no precision left, and one more layer flushes it to exactly zero. The parameter receives an update of zero and stays at its initialisation for the whole of training, while the last layers train perfectly well. The loss falls, then stalls, and nothing about the curve says which layers stopped.

Step 4 — why more weight makes it worse. The bound $(|w|/4)^L$ suggests $|w| > 4$ should grow the adjoint rather than shrink it. Measured:

Doubling $w$ to 2 helps a little. Raising it to 4 makes the gradient at depth 40 eighteen orders of magnitude smaller than at $w = 1$. The bound is not wrong; it is not tight, and the reason is that the two factors are not independent. A large $w$ drives $z$ away from zero, and a sigmoid away from zero is saturated, so $\sigma'(z)$ collapses. The weight enters the product linearly and takes the slope down exponentially with it.

This is what makes the failure a counterexample rather than a difficulty. There is no setting of the free parameter that rescues it: small $w$ loses to the factor $w$, large $w$ loses to saturation, and the best case is a shallow optimum near $w = 2$ that still vanishes, only more slowly.

Check — sanity

Every measured ratio is below its bound. $4^{-L}$ was derived from $\sigma' \le 1/4$ with no reference to the values, so a measurement above it would mean the propagation is wrong rather than that the bound is loose.

Depth 1 nearly attains the bound. $0.235004$ against $0.25$: the first unit sees $z = 0.5$, close enough to zero that its slope is close to maximal. The later units see larger $z$, saturate more, and the product falls away from the bound — which is why the gap at depth 40 is two orders of magnitude.

The arithmetic of the exponent checks out. $4^{-59} = 8.6 \times 10^{-36}$, and the measured $5.9 \times 10^{-39}$ is about 1500 times smaller — consistent with a mean slope near $0.23$ rather than $0.25$, since $(0.23/0.25)^{59} \approx 0.0075$.

STATEMENT

Account for both resources a backward pass consumes. Show that the arithmetic is a small multiple of the forward pass and that the memory is not, find the batch size at which stored activations overtake parameters, and quantify what checkpointing buys.

GIVEN

A stack of $L = 48$ square affine layers of width $d = 1024$, in fp32. Two working regimes: a batch of 32 rows, and a batch of 32 sequences of length 512, which presents $32 \times 512 = 16\,384$ rows to every layer.

FIND

Forward and backward FLOPs and their ratio; activation memory against parameter memory in both regimes; the batch at which the two are equal; and the memory and arithmetic under $\sqrt{L}$ checkpointing.

STRATEGY

Count matrix products, since everything else is lower order. Forward, each layer does one. Backward, by (I.6.3) and (I.6.4), each layer does two — one to send the adjoint onward and one to form the weight gradient. For memory, count what the backward pass will read and cannot recompute for free.

SOLUTION

Step 1 — arithmetic. A product of a $B \times d$ matrix with a $d \times d$ matrix costs $2Bd^2$ floating-point operations, counting a multiply and an add for each of the $Bd^2$ pairs. Forward:

Backward does two such products per layer, so $2F$, and the step costs $3F$:

The ratio is 3 in both regimes and would be 3 for any $B$, $d$ and $L$: it is a property of the graph, not of the size. This is the cheap gradient principle (I.6.T1) at work, and the constant here — three — sits comfortably inside the $c \le 4$ the theorem allows.

Step 2 — memory. The parameters are $Ld^2 = 50\,331\,648$ values, $0.201$ GB. The activations kept for the backward pass are $LBd$ by (I.6.7):

The same model, the same weights, the same arithmetic per row — and the memory picture inverts completely. At batch 32 the parameters dominate by thirty times; at batch 16 384 the activations dominate by sixteen.

Step 3 — the crossover. Setting $LBd = Ld^2$ gives

The layer count cancels, and so does everything else. Activations overtake parameters when the batch reaches the width, which for this model is 1024 rows — a number a single sequence of moderate length already exceeds. That is why the second regime is the normal one and the first is the exception.

Step 4 — checkpointing. Keep every $\sqrt{L}$-th activation and recompute the rest during the backward pass. With $L = 48$, $\sqrt{48} \approx 7$: hold 7 segment boundaries, plus the 6 tensors of whichever segment is currently being recomputed, so 13 tensors instead of 48.

The price is recomputing each segment’s forward pass once, which adds one forward pass to a step that cost three:

A third more arithmetic for nearly four times less memory. On a device where memory is the binding constraint and arithmetic is not — which describes most training — that is not a close decision.

Check — sanity

The FLOP ratio is exactly 3 in both regimes. It has to be — it is a count of matrix products per layer, one forward and two back, and no dimension enters it. A measured ratio far from 3 on real hardware is a statement about memory bandwidth and kernel efficiency, not about the arithmetic.

The crossover has no free parameters. $LBd = Ld^2$ reduces to $B = d$ with $L$ cancelling, so the answer cannot depend on depth. A deeper model has more activations and more parameters, in the same proportion.

Checkpointed memory is near $2\sqrt{L}$. $2\sqrt{48} \approx 13.9$ against the 13 tensors counted, and $48/13 = 3.7$ matches the measured reduction. The $\sqrt{L}$ of (I.6.7) is where boundary storage and segment storage balance; away from it, one term or the other grows.