Why backprop, not forward-mode: the one fact that makes billion-parameter training possible

Everything about backpropagation covered so far in this series, the manual linear-layer and attention backward derivations, autograd’s DAG traversal, assumed reverse-mode differentiation as the obvious choice without asking why it’s the only real option. It isn’t obvious. It’s the direct, derivable consequence of one shape mismatch, and once you’ve seen the derivation, several things that look like separate topics turn out to be the same fact wearing different clothes.

The actual reason neural networks use backprop, not the elegant-sounding one#

Any automatic way of computing derivatives has to choose a direction to propagate derivative information, and there are exactly two choices. Forward-mode AD propagates alongside the computation: each intermediate value carries its derivative with respect to one specific input. Getting the full Jacobian of $f:\mathbb{R}^n\to\mathbb{R}^m$ this way needs one full forward pass per input, $n$ passes, each yielding one column of the Jacobian. Reverse-mode AD, backpropagation, propagates backward from the outputs instead, needing one full backward pass per output, $m$ passes, each yielding one row.

A neural network’s loss function is $f:\mathbb{R}^n\to\mathbb{R}^1$: $n$ is the parameter count, potentially hundreds of billions, and the output is a single scalar loss, $m=1$. Forward-mode would need $n$ passes, one per parameter, to get the full gradient: for a 100-billion-parameter model, 100 billion forward passes. Reverse-mode needs exactly $m=1$ pass, because there’s only one output to seed. This is not a minor efficiency difference. It’s the difference between computable and not computable in the history of the universe. Backpropagation isn’t used because it’s elegant. It’s used because it’s the only one of the two options with any chance of terminating for a function with billions of inputs and one output. The general rule, for completeness: forward-mode wins when $n \ll m$, few inputs, many outputs; reverse-mode wins when $m \ll n$. Neural network training is about as extreme an instance of the second case as exists anywhere in applied mathematics.

The exact 2x cost of a backward pass, not a rule of thumb#

For a linear layer with parameter count $N$, the forward pass $y=Wx$ is one matrix-vector product: $2N$ FLOPs, one multiply and one add per weight touched. The backward pass needs two matmul-shaped operations, not one: $\frac{\partial L}{\partial x}=W^T\frac{\partial L}{\partial y}$, another $2N$ FLOPs, and $\frac{\partial L}{\partial W}=\frac{\partial L}{\partial y}\otimes x$, an outer product, another $2N$ FLOPs. Total backward cost: $4N$, exactly double the forward pass’s $2N$. This is why an $N$-parameter model is commonly counted as roughly 6N matmul FLOPs per token, 2N forward plus 4N backward, and the reason isn’t a convention someone picked, it’s that backprop needs two separate gradients for every one forward computation: one with respect to the input, to keep propagating backward through earlier layers, and one with respect to the weight, to actually update it.

Attention’s version of this ratio is sharper than the naive guess. Its forward pass is 2 matmuls ($QK^T$, then times $V$). Its backward pass needs 5, not the naively-doubled 4, because computing $dQ$, $dK$, and $dV$ through the chained backward pass touches matmul-shaped operations five separate times rather than four. FlashAttention-2’s own paper states this directly: 5 matrix multiplies in the backward pass against 2 in the forward pass, a 2.5x ratio, not attention’s own version of the generic 2x.

LayerNorm’s backward pass, derived completely, not memorized#

This is worth deriving in full because most treatments just hand you the formula, and the derivation itself reveals something almost nothing states explicitly. For one feature vector $x\in\mathbb{R}^d$:

$\mu=\frac{1}{d}\sum_k x_k,\qquad v=\frac{1}{d}\sum_k(x_k-\mu)^2,\qquad s=\sqrt{v+\epsilon},\qquad \hat x_i=\frac{x_i-\mu}{s}$

The reason this is genuinely harder than an elementwise operation: $\hat x_i$ depends on $x_i$ directly, but also indirectly through $\mu$ and $v$, both of which depend on every $x_k$, not just $x_i$. That means $\frac{\partial \hat x_i}{\partial x_j}$ is nonzero even when $i\ne j$.

Differentiating $\hat x_i=(x_i-\mu)/s$ with respect to $x_j$ needs $\frac{\partial\mu}{\partial x_j}=\frac{1}{d}$ (immediate, since $\mu$ is an average) and $\frac{\partial s}{\partial x_j}$, which needs $\frac{\partial v}{\partial x_j}$. Working through $\frac{\partial v}{\partial x_j}=\frac{2}{d}\sum_k(x_k-\mu)(\delta_{kj}-\tfrac1d)$, half of this sum vanishes completely: $\sum_k(x_k-\mu)=0$ always, since deviations from a true mean sum to zero by definition. That leaves $\frac{\partial v}{\partial x_j}=\frac{2}{d}(x_j-\mu)=\frac{2s}{d}\hat x_j$. Propagating through $\frac{\partial s}{\partial x_j}=\frac{1}{2s}\frac{\partial v}{\partial x_j}=\frac{\hat x_j}{d}$, and substituting back, every instance of $(x_i-\mu)$ rewrites as $s\hat x_i$, and the algebra collapses to:

$\frac{\partial\hat x_i}{\partial x_j}=\frac{1}{s}\left(\delta_{ij}-\frac{1}{d}-\frac{\hat x_i\hat x_j}{d}\right)$

Implementing $\frac{\partial L}{\partial x_j}=\frac1s\left[g_j-\frac1d\sum_i g_i-\frac{\hat x_j}{d}\sum_i g_i\hat x_i\right]$ (where $g_i=\frac{\partial L}{\partial\hat x_i}$) and checking it against independent finite differences on a random 5-element vector gives a maximum discrepancy of $8.13\times10^{-7}$, the same order of numerical agreement as the attention backward pass checked earlier in this series.

The connection almost nothing makes explicit. Look at that boxed formula next to softmax’s Jacobian, already derived earlier in this series: $\frac{\partial\alpha_i}{\partial z_j}=\alpha_i(\delta_{ij}-\alpha_j)$, an identity term minus a rank-1 outer-product correction. LayerNorm’s Jacobian has precisely the same “diagonal minus outer product” shape, scaled differently. Softmax normalizes a probability distribution; LayerNorm normalizes a feature vector’s statistics. They look unrelated. They aren’t. Any operation with a normalization constraint that links every output to every other output produces a Jacobian of this exact shape, because the constraint itself, probabilities summing to 1, a fixed mean and variance, couples every pair of elements together in the derivative, not just each element to itself. Once you’ve seen one, you can predict the shape of the other before deriving it.

How round() ever gets a gradient at all#

This closes a real gap left open since quantization was first covered in this series. Quantization’s core operation is $q=\text{round}(x/s)$, and round() has derivative exactly zero almost everywhere, it’s a staircase function, undefined at the jump points. If autograd computed the true derivative of round(), every gradient through a quantized weight would be exactly zero, and quantization-aware training would be mathematically impossible.

This isn’t a theoretical concern. Running true gradient descent on $(\text{round}(x)-\text{target})^2$ for 20 steps leaves $x$ at exactly its starting value, $2.3$. The true gradient is zero at every single step, so there is no update, ever. Running the Straight-Through Estimator version instead, where the forward pass rounds normally but the backward pass is manually defined to pretend $\frac{d(\text{round}(x))}{dx}=1$, as if rounding were the identity function, moves $x$ to $6.5$ after the same 20 steps, which rounds to exactly the target.

The STE isn’t claiming round() is secretly differentiable. It’s a deliberate, documented substitution (Bengio et al., 2013): swap in a useful gradient, the identity, for the useless true one, zero, on the reasoning that round(x) and x stay close in value, so the identity’s derivative is a locally reasonable stand-in for how the rounded output responds to changes in the pre-rounded input. This is the exact mechanism underneath every quantization-aware training loop, and it’s rarely connected back to the “what is autograd actually computing” question this whole series has been built around. Most treatments simply state that QAT “uses the straight-through estimator” without ever showing what breaks without it.

Gradient checkpointing, and why FlashAttention is a specific instance of it#

A normal backward pass needs every intermediate activation from the forward pass, because computing $\frac{\partial L}{\partial W}$ at each layer needs that layer’s input, the outer product from the FLOPs derivation above requires $x$. Storing every activation in a deep network costs memory proportional to depth. Gradient checkpointing trades this away: store only a subset of activations, and when the backward pass needs one that wasn’t stored, recompute it by re-running the forward pass from the nearest checkpoint. The classical result here shows that checkpointing roughly every $\sqrt{n}$-th layer reduces memory from $O(n)$ to $O(\sqrt n)$, at the cost of only one additional forward pass’s worth of compute, a bounded overhead, not an unbounded one.

FlashAttention is this exact technique, applied within a single attention operation, confirmed by the paper’s own words, not inferred: storing only the output and the softmax normalization statistics lets the attention matrix be recomputed easily during the backward pass, which the authors themselves describe as “a form of selective gradient checkpointing.” FlashAttention never stores the full $N\times N$ attention matrix at all. It recomputes exactly the blocks it needs, on-chip, directly from $Q$, $K$, $V$ during the backward pass, achieving 10 to 20x memory savings and a 2 to 4x wall-clock speedup specifically because recomputing on-chip is faster than reading a huge matrix back from HBM. The general principle and the specific famous technique are the same idea, one is just a special case tuned to a specific memory-access pattern.

Being honest about what a secondhand account gets slightly wrong#

Two real incidents from earlier in this series are worth reconciling against their primary sources precisely, because the gap between a compressed summary and the actual paper is itself informative.

PaLM 540B’s roughly 20 loss spikes were resolved by rolling back to a checkpoint about 100 steps before the spike and skipping the specific data batch that triggered it, not primarily by “clipping gradients more aggressively.” The rollback and the spike count are solid. The fix that mattered was identifying and skipping the offending batch, not a global change to the clipping threshold.

OPT’s early-warning signal is centered on the dynamic loss scalar crashing toward zero correlated with the final-layer activation norm spiking, not primarily “gradient norm monitoring.” Gradient clipping does show up in the OPT paper as a stability change, lowering the threshold from 1.0 to 0.3, but it isn’t the early-warning signal itself, the loss scalar and activation norm are.

Neither correction changes the underlying lesson. Both show why reading the actual paper, rather than a compressed secondhand account, matters: a summary that’s “close in spirit” can still misattribute which specific signal or fix did the real work, and that’s exactly the detail that matters if you’re trying to reproduce the fix rather than just the story.

Common mistakes#

Assuming forward-mode and reverse-mode AD are interchangeable implementation choices: which one is remotely tractable depends entirely on the shape of $n$ versus $m$, and neural network training sits at the single most extreme point on that spectrum in favor of reverse-mode.

Guessing a backward pass costs a flat 2x the forward pass for every operation: it’s exactly 2x for a simple linear layer, but attention’s backward pass is 2.5x its forward pass, because it touches five matmul-shaped operations against the forward pass’s two, not the naively doubled four.

Treating LayerNorm and softmax as unrelated because they normalize different kinds of things: both produce a Jacobian with the identical diagonal-minus-outer-product shape, because both operations couple every output to every other output through a shared normalization constraint.

Stating that quantization-aware training “uses the straight-through estimator” without knowing what it’s actually standing in for: the true gradient of round() is zero almost everywhere, and without the STE’s identity substitution, gradient descent through a rounded value does not move it at all.

Repeating a compressed incident summary as if it matched the primary source exactly: the PaLM and OPT corrections above show the difference between “close in spirit” and “the actual mechanism,” and only checking the paper directly reveals which one you have.

What this takes to be frontier-job-ready#

The technical axis: given a custom architecture with a novel normalization or a non-differentiable operation, know to ask what the actual Jacobian is, the method used to derive LayerNorm’s above, rather than assuming autograd handles it correctly by default. A training loop making zero progress through one specific sub-module might be exactly the section-4 non-differentiability problem, not a learning-rate problem, and those two have completely different fixes.

The operational axis is the reconciliation work itself. Real incident reports rarely reduce to one clean root cause the way a compressed summary implies. Being able to read the actual paper and hold a looser, more accurate picture than a secondhand account is exactly the discipline that separates diagnosing a real incident from repeating what a slide deck claimed happened.

The autonomy axis: choosing whether a memory-bound architecture needs gradient checkpointing, a custom gradient override in the STE style, or neither, is a genuine judgment call with no fixed script. FlashAttention’s authors didn’t apply checkpointing because a rule told them to. They recognized attention’s specific memory-access pattern as a case where the general principle applied unusually well, which is the same pattern-recognition this series has pointed at as the actual differentiator, over and over.

Try it yourself#

Beginner. Using the FLOPs counting convention above, compute the exact forward and backward FLOPs for a linear layer with $N=10^9$ parameters, and confirm the backward pass is exactly double the forward pass’s cost.

Intermediate. Implement round() with a Straight-Through Estimator in a small autodiff experiment (or by hand, tracking the substitution manually), and reproduce the result above: true gradient descent leaves $x$ unchanged after 20 steps, while the STE version reaches the target.

Advanced. Derive the off-diagonal term of LayerNorm’s Jacobian independently, without looking at the boxed result above, starting only from the quotient rule and the fact that deviations from a true mean sum to zero. Then explain, in your own words, why this same “coupling through a shared constraint” argument predicts that any new normalization scheme you might invent will also produce a diagonal-minus-outer-product Jacobian, before you’ve derived it.

---

The one-sentence version: reverse-mode differentiation isn’t the elegant choice, it’s the only tractable one once a function has billions of inputs and one output, that same asymmetry is why a backward pass costs roughly double a forward pass’s FLOPs (and 2.5x for attention specifically), LayerNorm’s Jacobian has the identical diagonal-minus-outer-product shape as softmax’s because both operations couple every output through a shared normalization constraint, and the Straight-Through Estimator is the one honest cheat that lets a gradient flow through an operation, round(), whose true derivative is exactly zero almost everywhere. Every one of these pieces, LayerNorm’s Jacobian included, is a load-bearing part of the actual transformer architecture they assemble into.