09 · F · Reverse

Backpropagation and Jacobians

Reverse-mode AD, the p − y collapse, RMSNorm, attention, SwiGLU, and the skip whose Jacobian is I.

Chain rule and Jacobians

Training is gradient descent on $L (θ)$ . Backprop is reverse-mode automatic differentiation: one backward sweep computes $\nabla_{θ} L$ . If $u \in R^{m}$ , $v = g (u) \in R^{k}$ , and $L$ is scalar:

J_{g} = \frac{\partial v}{\partial u} \in R^{k \times m}, \frac{\partial L}{\partial u} = J_{g}^{⊤} \frac{\partial L}{\partial v}

(9.1)

$J_{g}$: Jacobian of the local map
$J_{g}^{⊤} v$: vector-Jacobian product (VJP) — what is actually computed

In practice nobody materializes giant Jacobians. Frameworks compute VJPs. For a matrix multiply $Y = X W$ , with incoming gradient $G = \partial L / \partial Y$ :

\frac{\partial L}{\partial X} = G W^{⊤}, \frac{\partial L}{\partial W} = X^{⊤} G

(9.2)

Why it mattersThose two identities run through almost the entire Transformer. Every linear map is this pair.

Softmax + cross-entropy collapses

This is the most important gradient in the whole system. Start from the definitions:

L_{t} = - k \sum (y_{t})_{k} lo g (p_{t})_{k}, p = softmax (z), J_{softmax} = diag (p) - p p^{⊤}

(9.3)

The upstream vector is $(\partial L / \partial p)_{k} = - y_{k} / p_{k}$ . One-hot at $c$ makes that $- e_{c} / p_{c}$ . The VJP through softmax is

(J^{⊤} v)_{j} = p_{j} v_{j} - p_{j} (p^{⊤} v)

(9.3b)

Here $p^{⊤} v = p_{c} \cdot (- 1/ p_{c}) = - 1$ and $p_{j} v_{j} = - 1 [j = c]$ , so

\frac{\partial L _{t}}{\partial z _{t}} = p_{t} - y_{t}

(9.4)

(\frac{\partial L _{t}}{\partial z _{t}})_{k} = (p_{t})_{k} - 1 [k = c], c = x_{t + 1}

(9.5)

Why it mattersYou do not backprop through softmax and log separately. The error signal is “probability assigned minus truth.” Overconfident wrong tokens get a large push; already-correct peaked predictions get a small one.

Into the residual stream and W_U

\frac{\partial L}{\partial H ~} = (\frac{\partial L}{\partial Z}) W_{U}^{⊤}, \frac{\partial L}{\partial W _{U}} = H ~^{⊤} (\frac{\partial L}{\partial Z})

(9.6)

Then the final Norm Jacobian maps $\partial L / \partial H ~$ back to $\partial L / \partial H^{(n)}$ .

RMSNorm Jacobian (one vector)

Let $r = RMS (u)$ , $u^= u / r$ , $v = γ ⊙ u^$ .

\frac{\partial r}{\partial u} = \frac{u}{d r}, \frac{\partial u ^}{\partial u} = \frac{1}{r} (I - \frac{u u ^{⊤}}{d r ^{2}})

(9.7)

Then $\partial L / \partial u$ follows by the product rule with $γ$ . LayerNorm is the same idea after projecting orthogonal to $1$ (mean removal).

Attention backward

Fix one head. Forward: $S = Q K^{⊤} / d_{h}$ , $A = softmax (S + M)$ , $O = A V$ . Given $G_{O} = \partial L / \partial O$ :

\frac{\partial L}{\partial A} = G_{O} V^{⊤}, \frac{\partial L}{\partial V} = A^{⊤} G_{O}

(9.8)

\frac{\partial L}{\partial S _{t, :}} = J_{softmax} (A_{t, :})^{⊤} (\frac{\partial L}{\partial A _{t, :}})^{⊤}

(9.9)

\frac{\partial L}{\partial Q} = \frac{1}{d _{h}} (\frac{\partial L}{\partial S}) K, \frac{\partial L}{\partial K} = \frac{1}{d _{h}} (\frac{\partial L}{\partial S})^{⊤} Q

(9.10)

RoPE’s backward pass is the inverse rotation. Then $Q = U W_{Q}$ etc. give

\frac{\partial L}{\partial W _{Q}} = U^{⊤} \frac{\partial L}{\partial Q}, \frac{\partial L}{\partial U} + = \frac{\partial L}{\partial Q} W_{Q}^{⊤}

(9.11)

The same for K, V. Multi-head: sum the $\partial L / \partial U$ contributions.

Residual Jacobian

Z = H + F (H) ⟹ \frac{\partial L}{\partial H} = \frac{\partial L}{\partial Z} + J_{F}^{⊤} \frac{\partial L}{\partial Z}

(9.12)

Why it mattersThe first term is the skip connection. That is why deep Transformers train: there is always a path whose Jacobian is I.

SwiGLU backward

Let $g = W_{G}^{⊤} u$ , $up = W_{up}^{⊤} u$ , $s = SiLU (g)$ , $h = s ⊙ up$ , $out = W_{D}^{⊤} h$ .

\frac{\partial L}{\partial h} \frac{\partial L}{\partial s} \frac{\partial L}{\partial g} = W_{D} \frac{\partial L}{\partial out} = \frac{\partial L}{\partial h} ⊙ up, \frac{\partial L}{\partial up} = \frac{\partial L}{\partial h} ⊙ s = \frac{\partial L}{\partial s} ⊙ SiLU^{'} (g)

(9.13)

Then ordinary matmul rules for $W_{G}, W_{up}, W_{D}, u$ .

Embedding backward

$\partial L / \partial H^{(0)}$ is scattered back into the rows of $W_{E}$ (and $W_{P}$ if used):

\frac{\partial L}{\partial ( W _{E} ) _{x_{t}, :}} + = \frac{\partial L}{\partial h _{t}^{(0)}}

(9.14)