Backpropagation Through Transformers

Two passes, one graph: forward and backward

Every training step runs the same computation graph twice. The forward pass feeds inputs through the network layer by layer, producing intermediate activations and finally a scalar loss L that measures how wrong the prediction was. The backward pass traverses the identical graph in reverse, computing ∂L/∂θ for every parameter θ so an optimizer can nudge each one downhill.

The asymmetry to hold onto is this: the forward pass flows values left-to-right; the backward pass flows gradients right-to-left. Each node receives, from the layer above it, the gradient of the loss with respect to that node’s own output — call it the upstream gradient — and is responsible for two products: the gradient with respect to its parameters (used to update them) and the gradient with respect to its inputs (passed further down as the next layer’s upstream gradient). Because the backward pass reuses quantities computed on the way forward, the framework must remember the forward activations until they are consumed — a fact that becomes the memory story at the end of this article. Training is nothing more than this forward-then-backward loop, repeated millions of times.

The chain rule, from scalars to compositions

The whole edifice rests on one rule from calculus. If L = f(g(x)), then dL/dx = f’(g(x)) · g’(x) — the derivative of a composition is the product of the local derivatives along the chain. Stack more functions and you multiply more factors: for L = f₃(f₂(f₁(x))),

dL/dx = f₃’ · f₂’ · f₁’
        (each derivative evaluated at its own forward input)

A transformer is a deep composition — embedding, then a stack of attention and feed-forward blocks, then an output projection and a loss — so its gradient is a long product of per-layer factors. This is also the origin of vanishing and exploding gradients: a product of many factors each smaller than one decays toward zero, and a product of factors larger than one blows up. Much of transformer design — residual connections, normalization, careful initialization — exists to keep that long product well-behaved. Backpropagation is simply the chain rule applied systematically, storing shared sub-results so the same factor is never recomputed twice.

Jacobians: the derivative of a vector function

Transformer layers do not map scalars to scalars; they map vectors to vectors. The right generalization of the derivative is the Jacobian. If y = f(x) with x an n-vector and y an m-vector, the Jacobian J is the m × n matrix of partials J_ij = ∂y_i/∂x_j — row i tells you how output i responds to every input.

The multivariate chain rule is then just matrix multiplication of Jacobians. For y = f(x) followed by z = g(y), the Jacobian of the composition is J_g · J_f (shapes chain: [k×m][m×n] = [k×n]). If you tried to run backpropagation by literally forming and multiplying these Jacobians, you would drown: a layer mapping a d = 4096 vector to another 4096 vector has a 4096 × 4096 Jacobian — 16 million entries for one token at one layer. The genius of reverse-mode autodiff is that it never builds these matrices. To see why, we need the vector-Jacobian product.

The vector-Jacobian product (VJP) view

Here is the key move, and it is worth stating precisely because the two directions are easy to confuse. A Jacobian-vector product (JVP), J v, pushes an input perturbation forward through the Jacobian — that is forward-mode differentiation. Backpropagation is the transpose object: the vector-Jacobian product (VJP), ᵍ = ḡ J (equivalently Jᵀ ḡ), which pulls an output gradient backward through the Jacobian — that is reverse-mode.

Concretely, a layer receives the upstream row-gradient ḡ = ∂L/∂y (shape [1×m]) and must return ∂L/∂x = ḡ J (shape [1×n]). The point of the VJP view is that you never materialize J — you implement the action of ḡ J directly. For a matrix multiply the VJP is another matrix multiply; for softmax, as we will see, it collapses to an O(n) formula instead of an n×n matrix. Every layer in a deep-learning framework ships two functions: a forward y = f(x), and a backward that takes ḡ and returns the VJP. Chaining those backward functions from the loss down to the inputs is backpropagation.

Why reverse-mode is efficient for a scalar loss

Autodiff comes in two flavors, and which one wins depends on the shape of the problem. Forward-mode (JVPs) computes the derivative of all outputs with respect to one input per pass; reverse-mode (VJPs) computes the derivative of one output with respect to all inputs per pass. The cost of a full gradient scales with the number of passes you need.

Training has exactly one output that matters: the scalar loss L. And it has an enormous number of inputs: every weight, often billions of them. Reverse-mode delivers ∂L/∂θ for all parameters in a single backward pass, at a cost roughly two to three times the forward pass — independent of the parameter count. Forward-mode would need one pass per parameter, which is hopeless at billions. That is the whole reason deep learning runs on reverse-mode backpropagation: with a scalar objective and many parameters, seeding the backward pass with ∂L/∂L = 1 and propagating VJPs computes the entire gradient cheaply. The price you pay for this asymmetry is memory — reverse mode must keep the forward activations around to evaluate each local Jacobian on the way back — which is precisely the trade-off we return to at the end.

Gradients through a linear layer

The linear (dense) layer is the workhorse — attention’s Q/K/V/O projections and the feed-forward network are all linear layers — so its gradient is the one to know cold. Using the row-vector convention that matches how transformer activations are laid out ([batch, seq, d]), the forward pass is y = x W, with x: [N, d_in], W: [d_in, d_out], y: [N, d_out].

Given the upstream gradient ∂L/∂y (same shape as y), the two VJPs are clean and worth memorizing:

∂L/∂W = xᵀ (∂L/∂y)      shape [d_in, d_out]  (matches W)
∂L/∂x = (∂L/∂y) Wᵀ      shape [N, d_in]     (matches x)

Two sanity checks make these unforgettable. First, shapes must match: a parameter’s gradient always has the parameter’s shape, and an input’s gradient always has the input’s shape — there is only one way to arrange the matrix product that produces the right dimensions. Second, notice the transposes cross over: the forward multiplies by W, the input-gradient multiplies by Wᵀ. Because the loss is scalar and W is shared across all N rows (tokens), the xᵀ(∂L/∂y) product sums the per-token contributions — which is exactly why batching many tokens gives one averaged weight gradient.

Gradients through elementwise nonlinearities

Activation functions — ReLU, GELU, SiLU, tanh — act elementwise: y_i = σ(x_i), with each output depending only on the matching input. That independence makes the Jacobian diagonal, so its VJP is not a matrix multiply at all — it is a cheap elementwise (Hadamard) product:

∂L/∂x = (∂L/∂y) ⊙ σ’(x)
              (multiply the upstream gradient, position by position,
               by the local slope σ’ at the forward input x)

For ReLU(x) = max(0, x) the derivative is a gate: 1 where x > 0 and 0 where x < 0. So ReLU’s backward pass simply masks the upstream gradient — wherever a unit was clipped to zero on the way forward, its gradient is killed on the way back, and no learning signal flows through that dead path. Smooth activations like GELU pass a fractional slope everywhere, avoiding hard-zero gradients. The practical lesson: to backprop an elementwise nonlinearity you only need to have cached its input (or, for some functions, its output) from the forward pass — another activation the framework must hold in memory until the backward pass consumes it.

The softmax Jacobian, and why its VJP is cheap

Softmax turns a vector of scores into a probability distribution: s_i = exp(z_i) / Σ_j exp(z_j). Unlike an elementwise function, every output depends on every input (the normalizer couples them), so its Jacobian is full. The entries have a famously tidy form:

∂s_i/∂z_j = s_i (δ_ij − s_j)
   where δ_ij = 1 if i=j else 0

If you stopped here you would face an n × n Jacobian per attention row. But the VJP view rescues you. Given upstream g = ∂L/∂s, multiply it through the Jacobian and the structure collapses:

∂L/∂z_i = s_i ( g_i − Σ_j g_j s_j )
              = s_i ( g_i − (g · s) )

Read it off: compute the single scalar g · s (a weighted average of the upstream gradient under the softmax distribution), subtract it from each g_i, and scale by s_i. That is an O(n) operation with no matrix ever formed — the VJP view turned a quadratic object into a linear one. The intuition is that softmax’s gradient is centered: it only responds to how each score’s upstream signal differs from the distribution-weighted mean, which is exactly what keeps a probability vector on the simplex as it updates.

Layernorm: normalize forward, project backward

Layer normalization standardizes each token’s activation vector before scaling and shifting: with mean μ and standard deviation σ taken over the feature dimension, x̂ = (x − μ) / σ and y = γ ⊙ x̂ + β. Because μ and σ are computed from x, every output coordinate depends on every input coordinate — the backward pass is not elementwise.

The gradient with respect to the pre-norm input has three terms, and the shape is more instructive than the algebra. Writing ĝ = (∂L/∂y) ⊙ γ for the gradient just inside the normalization:

∂L/∂x = (1/σ) [ ĝ − mean(ĝ) − x̂ · mean(ĝ ⊙ x̂) ]

The two subtractions are the whole story: layernorm’s backward pass removes the mean of the incoming gradient and removes its projection onto the normalized direction x̂, then rescales by 1/σ. Intuitively, the forward pass made the output invariant to shifting and scaling the input, so the backward pass refuses to pass any gradient component that would merely shift or rescale — those directions cannot change the loss, so their gradient is subtracted away. Meanwhile γ and β get the simple gradients Σ (∂L/∂y) ⊙ x̂ and Σ (∂L/∂y). The 1/σ factor also explains why normalization stabilizes training: it keeps the backward signal at a controlled scale regardless of the input’s magnitude.

Residual connections: the gradient highway

If deep composition threatens gradients with that long, fragile product of factors, residual connections are the fix — and their effect is clearest in the backward pass. A residual block computes y = x + F(x), where F is the attention or feed-forward sublayer. Differentiate and the additive structure gives an additive gradient:

∂L/∂x = ∂L/∂y + (∂L/∂y) · F’(x)
           = (identity path) + (through-the-block path)

Look at the first term: the upstream gradient ∂L/∂y is passed straight down unchanged, added to whatever the sublayer contributes. Across a stack of L blocks, this identity path lets the loss gradient reach the earliest layer as (roughly) a sum of terms rather than a product of many Jacobians. A sum does not vanish just because it is deep — even if every F’(x) is tiny, the identity term carries a clean gradient all the way down. This is the ‘gradient highway’: residuals give the backward pass an unobstructed lane, which is exactly why transformers can be stacked dozens of layers deep and still train. Normalization placement (pre-norm vs post-norm) tunes how clean that lane stays, but the additive identity path is the load-bearing idea.

A worked example: backprop through a 2-layer stack

Abstractions settle once you push real numbers through them. Take a tiny two-layer network in the same row-vector convention as above. Forward: h = x W₁, a = ReLU(h), y = a W₂, and a squared-error loss L = ½(y − t)². Let

x  = [1, 2]            W₁ = [[ 1, 0],      W₂ = [[ 2],
                            [-1, 1]]             [-1]]   target t = 0

Forward. h = x W₁ = [1·1 + 2·(−1), 1·0 + 2·1] = [−1, 2]. ReLU clips the first unit: a = [0, 2]. Then y = a W₂ = 0·2 + 2·(−1) = −2, and L = ½(−2 − 0)² = 2.

Backward. Seed with ∂L/∂y = y − t = −2, then apply the linear-layer and ReLU VJPs in reverse:

∂L/∂y  = -2
∂L/∂W₂ = aᵀ (∂L/∂y) = [0, 2]ᵀ · (-2) = [[0], [-4]]
∂L/∂a  = (∂L/∂y) W₂ᵀ = -2 · [2, -1] = [-4, 2]
∂L/∂h  = (∂L/∂a) ⊙ ReLU’(h) = [-4, 2] ⊙ [0, 1] = [0, 2]
∂L/∂W₁ = xᵀ (∂L/∂h) = [1, 2]ᵀ · [0, 2] = [[0, 2], [0, 4]]
∂L/∂x  = (∂L/∂h) W₁ᵀ = [0, 2] · [[1, -1], [0, 1]] = [0, 2]

Everything you were told shows up concretely. The ReLU gate zeroed the first component of ∂L/∂h — the unit that was clipped forward gets no gradient backward, and correspondingly the first row of ∂L/∂W₁ is all zeros: that dead unit’s weights receive no update this step. Each parameter gradient matches its parameter’s shape, and the crossed transposes (W₂ᵀ, W₁ᵀ) route the signal back to the inputs exactly as the formulas promised.

Chaining it up: backprop through a whole block

A real transformer block is just more of the same pieces wired in the order the forward pass visits them, run backward in reverse order. Take the feed-forward sublayer with a residual: y = x + W₂(GELU(x W₁)). The backward pass peels it off outermost-first — residual add, then the second linear, then GELU, then the first linear — each step a VJP we have already derived:

g_y                      # upstream gradient at the block output
g_res = g_y              # residual: identity path copies it straight through
g_h2  = g_y             # ...and also feeds the sublayer branch
g_a   = g_h2 W₂ᵀ       # back through the 2nd linear
g_pre = g_a ⊙ GELU’(·)  # back through the nonlinearity (elementwise)
g_x   = g_res + g_pre W₁ᵀ  # 1st linear, then ADD the residual path

Two habits from this generalize to the entire network. First, gradients accumulate at fan-out points: because x was used twice (once by the residual, once by the sublayer), its gradient is the sum of the gradients along both paths — a direct consequence of the multivariable chain rule. Second, the parameter gradients (∂L/∂W₁, ∂L/∂W₂) fall out as by-products using the cached forward activations. Stack this block, add attention (whose softmax VJP we derived), and repeat for every layer: that is the entire backward pass of a transformer, no new mathematics required — only the chain rule, applied node by node.

The catch: activations must be kept in memory

Reverse-mode’s efficiency has a price, and it is memory. Look back at every VJP: the linear layer’s ∂L/∂W = xᵀ(∂L/∂y) needs the forward input x; the nonlinearity’s VJP needs x (or the output) to evaluate its slope; softmax needs its output s; layernorm needs x̂ and σ. The backward pass evaluates each local Jacobian at its forward input, so the framework must store the forward activations from the moment they are produced until the backward pass consumes them.

That bill scales alarmingly. Activation memory is roughly O(batch × sequence_length × d_model × layers), and at training time it typically dwarfs the memory used by the parameters and their gradients — long sequences and deep stacks make it the binding constraint on how large a batch or context you can fit. This is the direct consequence of the forward-then-backward structure: you cannot free an activation until its gradient has flowed back through it. Which raises the obvious question — must we really keep all of them?

The bridge to gradient checkpointing

No — and the escape hatch trades compute for memory. Gradient (activation) checkpointing keeps only a sparse set of activations from the forward pass — say, the input to each transformer block — and discards the intermediate activations inside each block. When the backward pass reaches a block, it re-runs that block’s forward pass from the saved checkpoint to regenerate the intermediates it needs, uses them for the VJPs, and frees them again.

The arithmetic is favorable. Storing checkpoints only at block boundaries cuts activation memory from O(layers) toward O(√layers) with a common strategy, at the cost of roughly one extra forward pass — so total training compute rises by about 33% (one recompute on top of the usual forward-plus-backward) in exchange for a large memory saving. That trade is often exactly what lets a long-context or large model fit on the hardware at all. The conceptual point to carry away: checkpointing is only possible because backpropagation is deterministic and local — every activation can be recomputed from an earlier one by replaying the forward function. The chain rule tells you what to compute; checkpointing is a memory-management decision about when to compute it. Once you see the backward pass as VJPs consuming cached forward values, both the memory problem and its cleanest solution follow directly.

Common pitfalls when reasoning about transformer gradients

A handful of mistakes recur once you start tracing gradients by hand. Forgetting to sum at shared parameters. A weight matrix is applied to every token in the batch and sequence; its gradient is the sum over all those uses, which is why xᵀ(∂L/∂y) contracts the N dimension. Treat it as per-token and your shapes (and your updates) will be wrong.

Confusing forward and reverse mode. Backprop is the VJP ḡJ, not the JVP Jv; the transpose is not cosmetic. Assuming softmax and layernorm are elementwise. Both couple all coordinates, so their backward passes have the Σ-shaped correction terms — the −(g·s) in softmax, the mean-subtractions in layernorm — that a naive elementwise guess would miss. Ignoring the activation-memory cost until an out-of-memory error forces the issue; the gradient formulas tell you in advance which forward values must be kept. And overwriting activations in place: a forward op that mutates a tensor the backward pass still needs will silently corrupt the gradient. Every one of these is a direct reading of the chain rule — keep the forward-value / upstream-gradient / local-Jacobian trio in mind and the errors become hard to make.