A linear layer backwards

$n$ rows of input (tokens), $i$ features in, $o$ out. The backward pass hands this layer $G$, the gradient of the loss with respect to its output. It must produce gradients for $W$ and for $X$.

The answer must be shaped like $W$, and it can only be built from $X$ and $G$. The only way to multiply them into $i \times o$ is $X^\top G$, summing over the $n$ rows: every token contributes.

Shaped like $X$, built from $G$ and $W$: $GW^\top$. It’s the vector–Jacobian product of the last deck, one row per token. Shape matching is the fastest way to get these right, and to catch a missing transpose.

Check one entry by hand: $W_{21}$ multiplies $X_{n2}$ into $Y_{n1}$ for each row $n$, so its gradient is $X_{12}G_{11} + X_{22}G_{21} = 2 \cdot 1 + 1 \cdot 0 = 2$. That’s entry $(2, 1)$ of $X^\top G$.

The bias was broadcast to every row (Part 4), so its gradient adds up the rows of $G$: here $[1, 1, 1]$. Broadcasting forward means summing backward.

Each matrix product in the forward pass becomes two of the same size in the backward pass. So training costs about three forward passes: $2N + 4N = 6N$ FLOPs per token (next deck but one).

Write the shapes of $\partial L/\partial W$ and $\partial L/\partial X$.

# y = x @ W.T with W stored (out, in); g = dL/dy
grad_W = g.T @ x          # (out, in), same as W
grad_x = g @ W            # (n, in)

The same rules, transposed to match how the weight is stored.

• $Y = XW$, $G = \partial L/\partial Y$.
• $\partial L/\partial W = X^\top G$; $\partial L/\partial X = GW^\top$: match the shapes.
• Bias gradient: sum $G$ over rows.
• Backward = two matmuls per forward matmul.