---
layout: post
title: "Transformer from Scratch: The Full Forward Pass, Backprop, and Weight Update Math"
date: 2026-06-05 13:00:00 +0530
tags: ["Transformers", "Backpropagation", "Attention", "Deep Learning", "Encoder-Decoder", "Layer Norm", "Adam", "From Scratch"]
category: "ai"
description: "One full transformer training step worked by hand — embeddings, positional encoding, attention, layer norm, the encoder-decoder, cross-entropy loss, backprop, and Adam."
keywords: "transformer from scratch, transformer math explained, backpropagation transformer, attention forward pass by hand, layer norm backprop, encoder decoder transformer, adam optimizer transformer, softmax cross entropy gradient, transformer weight update"
image: /assets/images/og-transformer-math.png
order: 11
---
Most transformer explainers stop at the forward pass and let the words "and then we train it" do all the work. I did that for a long time — I could sketch [how attention works](/blog/attention-in-transformers-explained/) conceptually, I could read [token embeddings](/blog/token-embeddings-explained/) and [positional encoding](/blog/positional-encoding-explained/) in isolation — but I could not trace one training step from input IDs to a weight that actually moves. This post fixes that.
Below is **one complete training step of the original 2017 encoder-decoder transformer**, solved like a single math problem. Every matrix multiply is written out. Every softmax row is summed. Every gradient lands on a number you can check with a calculator. The *formulas* are exactly those from [*Attention Is All You Need*](https://arxiv.org/abs/1706.03762); the *dimensions* are a toy so you can verify every line by hand. If you've built [autograd from scratch](/blog/build-autograd-from-scratch/), this is the derivation you want open while debugging — and when something looks wrong, [numerical gradient checking](/blog/numerical-gradient-checking-explained/) is how I sanity-check each layer.
> [!NOTE]
> Conceptual attention (queries, keys, values, no arithmetic) lives in [how attention works](/blog/attention-in-transformers-explained/). This post is the arithmetic.
---
## What machine am I solving?
The 2017 transformer is an encoder-decoder stack. The encoder reads the source sentence; the decoder reads a shifted target and attends to encoder memory. One layer on each side, one training step, every box below gets a number.
```mermaid
flowchart TB
subgraph ENC["Encoder ×N"]
EI["Input embed + PE"] --> ESA["Multi-head self-attention"]
ESA --> EAN1["Add & Norm"]
EAN1 --> EFF["Feed-forward"]
EFF --> EAN2["Add & Norm"]
end
subgraph DEC["Decoder ×N"]
DI["Output embed + PE"] --> DSA["Masked self-attention"]
DSA --> DAN1["Add & Norm"]
DAN1 --> DCA["Cross-attention
(Q=decoder, K,V=encoder)"]
DCA --> DAN2["Add & Norm"]
DAN2 --> DFF["Feed-forward"]
DFF --> DAN3["Add & Norm"]
end
EAN2 -->|"memory K,V"| DCA
DAN3 --> LIN["Linear → vocab"]
LIN --> SM["Softmax → probabilities"]
SM --> LOSS["Cross-entropy + label smoothing"]
```
---
## What toy setup am I using?
Task: translate **je suis** → **i am**. The decoder is fed the shifted target ` i` and must predict `i am`.
| Symbol | Toy | Real (base paper) |
|---|---:|---:|
| $d_{\text{model}}$ | 4 | 512 |
| heads $h$ | 2 | 8 |
| $d_k=d_v=d_{\text{model}}/h$ | 2 | 64 |
| $d_{ff}$ | 8 | 2048 |
| layers $N$ | 1 | 6 |
| vocab $V$ | 4 | 37,000 |
All formulas are exact at any scale; only the numbers shrink. Vocabulary order: `[, i, am, ]`.
---
## How do token embeddings enter the encoder?
Each token ID picks a row from the learned embedding matrix $E \in \mathbb{R}^{V \times d_{\text{model}}}$, then scales by $\sqrt{d_{\text{model}}}$.
For $d_{\text{model}} = 4$:
$$
\sqrt{d_{\text{model}}} = \sqrt{4} = 2
$$
Encoder tokens `je` (position 0) and `suis` (position 1). After lookup and scaling, I round to clean integers for hand arithmetic (shapes stay $2 \times 4$):
$$
\text{embed}(\text{je}) = 2 \cdot [0.5,\, 0,\, 0.5,\, 0] = [1,\, 0,\, 1,\, 0]
$$
$$
\text{embed}(\text{suis}) = 2 \cdot [0,\, 0.5,\, 0,\, 0.5] = [0,\, 1,\, 0,\, 1]
$$
Stack into the encoder input matrix $X_{\text{raw}} \in \mathbb{R}^{2 \times 4}$:
$$
X_{\text{raw}} =
\begin{bmatrix}
1 & 0 & 1 & 0 \\
0 & 1 & 0 & 1
\end{bmatrix}
$$
Row 0 is `je`. Row 1 is `suis`. Shape: $(2,\, 4)$.
---
## What is positional encoding and how do I add it?
Sinusoidal positional encoding from the paper:
$$
PE_{(pos,\, 2i)} = \sin\!\Big(\tfrac{pos}{10000^{2i/d_{\text{model}}}}\Big)
$$
$$
PE_{(pos,\, 2i+1)} = \cos\!\Big(\tfrac{pos}{10000^{2i/d_{\text{model}}}}\Big)
$$
For $d_{\text{model}} = 4$, the frequency divisors are:
$$
10000^{0} = 1
$$
$$
10000^{1/2} = 100
$$
Position $pos = 0$:
$$
PE_0 = [\sin 0,\, \cos 0,\, \sin 0,\, \cos 0] = [0,\, 1,\, 0,\, 1]
$$
Position $pos = 1$:
$$
PE_1 = [\sin 1,\, \cos 1,\, \sin 0.01,\, \cos 0.01] = [0.84,\, 0.54,\, 0.01,\, 1.00]
$$
Add to embeddings elementwise. For this toy I round the sum to keep later arithmetic tractable (same as the original derivation):
$$
X = X_{\text{raw}} + PE \approx
\begin{bmatrix}
1 & 0 & 1 & 0 \\
0 & 1 & 0 & 1
\end{bmatrix}
$$
Shape: $(2,\, 4)$. This $X$ is what enters encoder self-attention.
---
## How does encoder self-attention compute Q, K, and V?
Multi-head attention projects $X$ into queries, keys, and values per head:
$$
Q = X W_Q, \qquad K = X W_K, \qquad V = X W_V
$$
$$
\text{Attn}(Q,K,V) = \text{softmax}\!\Big(\tfrac{QK^{\top}}{\sqrt{d_k}}\Big) V
$$
Head 1 weights ($X$ is $4 \times 2$ per row-column convention: $X \in \mathbb{R}^{2 \times 4}$, $W_Q \in \mathbb{R}^{4 \times 2}$):
$$
W_Q^{(1)} = W_K^{(1)} =
\begin{bmatrix}
1 & 0 \\
0 & 1 \\
1 & 0 \\
0 & 1
\end{bmatrix}
$$
$$
W_V^{(1)} =
\begin{bmatrix}
1 & 0 \\
0 & 1 \\
0 & 0 \\
0 & 0
\end{bmatrix}
$$
### Matmul: $Q = X W_Q^{(1)}$
Row 0 of $X$: $[1,\, 0,\, 1,\, 0]$.
$$
Q[0,0] = 1 \cdot 1 + 0 \cdot 0 + 1 \cdot 1 + 0 \cdot 0 = 2
$$
$$
Q[0,1] = 1 \cdot 0 + 0 \cdot 1 + 1 \cdot 0 + 0 \cdot 1 = 0
$$
Row 1 of $X$: $[0,\, 1,\, 0,\, 1]$.
$$
Q[1,0] = 0 \cdot 1 + 1 \cdot 0 + 0 \cdot 1 + 1 \cdot 0 = 0
$$
$$
Q[1,1] = 0 \cdot 0 + 1 \cdot 1 + 0 \cdot 0 + 1 \cdot 1 = 2
$$
$$
Q^{(1)} =
\begin{bmatrix}
2 & 0 \\
0 & 2
\end{bmatrix}
$$
Same $W_K^{(1)}$, so $K^{(1)} = Q^{(1)}$.
### Matmul: $V = X W_V^{(1)}$
$$
V[0,0] = 1 \cdot 1 + 0 \cdot 0 + 1 \cdot 0 + 0 \cdot 0 = 1
$$
$$
V[0,1] = 1 \cdot 0 + 0 \cdot 1 + 1 \cdot 0 + 0 \cdot 0 = 0
$$
$$
V[1,0] = 0 \cdot 1 + 1 \cdot 0 + 0 \cdot 0 + 1 \cdot 0 = 0
$$
$$
V[1,1] = 0 \cdot 0 + 1 \cdot 1 + 0 \cdot 0 + 1 \cdot 0 = 1
$$
$$
V^{(1)} =
\begin{bmatrix}
1 & 0 \\
0 & 1
\end{bmatrix}
$$
Shapes: $Q, K, V \in \mathbb{R}^{2 \times 2}$.
---
## How do I score, scale, softmax, and mix values?
### Matmul: $Q K^{\top}$
$$
(QK^{\top})[0,0] = 2 \cdot 2 + 0 \cdot 0 = 4
$$
$$
(QK^{\top})[0,1] = 2 \cdot 0 + 0 \cdot 2 = 0
$$
$$
(QK^{\top})[1,0] = 0 \cdot 2 + 2 \cdot 0 = 0
$$
$$
(QK^{\top})[1,1] = 0 \cdot 0 + 2 \cdot 2 = 4
$$
$$
QK^{\top} =
\begin{bmatrix}
4 & 0 \\
0 & 4
\end{bmatrix}
$$
### Scale by $\sqrt{d_k}$
$$
\sqrt{d_k} = \sqrt{2} \approx 1.4142
$$
$$
S = \frac{QK^{\top}}{\sqrt{d_k}} =
\begin{bmatrix}
2.828 & 0 \\
0 & 2.828
\end{bmatrix}
$$
### Softmax row 0
$$
e^{2.828} \approx 16.92
$$
$$
e^{0} = 1
$$
$$
\text{sum}_0 = 16.92 + 1 = 17.92
$$
$$
A[0,0] = 16.92 / 17.92 \approx 0.944
$$
$$
A[0,1] = 1 / 17.92 \approx 0.056
$$
### Softmax row 1
By symmetry:
$$
A[1,0] = 0.056
$$
$$
A[1,1] = 0.944
$$
$$
A^{(1)} =
\begin{bmatrix}
0.944 & 0.056 \\
0.056 & 0.944
\end{bmatrix}
$$
Each row sums to $1$.
### Matmul: $O^{(1)} = A^{(1)} V^{(1)}$
$$
O^{(1)}[0,0] = 0.944 \cdot 1 + 0.056 \cdot 0 = 0.944
$$
$$
O^{(1)}[0,1] = 0.944 \cdot 0 + 0.056 \cdot 1 = 0.056
$$
$$
O^{(1)}[1,0] = 0.056 \cdot 1 + 0.944 \cdot 0 = 0.056
$$
$$
O^{(1)}[1,1] = 0.056 \cdot 0 + 0.944 \cdot 1 = 0.944
$$
$$
O^{(1)} =
\begin{bmatrix}
0.944 & 0.056 \\
0.056 & 0.944
\end{bmatrix}
$$
Head 2 uses its own $W_Q^{(2)}, W_K^{(2)}, W_V^{(2)}$ and lands identically here with my chosen weights. Concatenate heads → width $4$, mix with $W_O = I_4$:
$$
\text{MHA}(X) =
\begin{bmatrix}
0.944 & 0.056 & 0.944 & 0.056 \\
0.056 & 0.944 & 0.056 & 0.944
\end{bmatrix}
$$
Shape: $(2,\, 4)$.
---
## What is Add & Norm (layer norm)?
Residual connection plus layer norm:
$$
\text{LayerNorm}(x + \text{Sublayer}(x))
$$
$$
\mu = \frac{1}{d} \sum_{i=1}^{d} x_i
$$
$$
\sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i - \mu)^2
$$
$$
y_i = \gamma \cdot \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta
$$
With $\gamma = 1$, $\beta = 0$, $\epsilon = 0$.
### Residual add after attention
$$
X + \text{MHA}(X) =
\begin{bmatrix}
1.944 & 0.056 & 1.944 & 0.056 \\
0.056 & 1.944 & 0.056 & 1.944
\end{bmatrix}
$$
### Layer norm on row 0
$$
\mu_0 = (1.944 + 0.056 + 1.944 + 0.056) / 4 = 1.0
$$
Deviations from mean:
$$
[1.944 - 1,\, 0.056 - 1,\, 1.944 - 1,\, 0.056 - 1] = [0.944,\, -0.944,\, 0.944,\, -0.944]
$$
$$
\sigma_0^2 = (0.944^2 + 0.944^2 + 0.944^2 + 0.944^2) / 4 = 0.891
$$
$$
\sigma_0 = \sqrt{0.891} \approx 0.944
$$
Normalized row 0:
$$
\text{LN}_1[0,:] = [0.944/0.944,\, -0.944/0.944,\, 0.944/0.944,\, -0.944/0.944] = [1,\, -1,\, 1,\, -1]
$$
Row 1 is antisymmetric:
$$
\text{LN}_1[1,:] = [-1,\, 1,\, -1,\, 1]
$$
$$
\text{LN}_1 =
\begin{bmatrix}
1 & -1 & 1 & -1 \\
-1 & 1 & -1 & 1
\end{bmatrix}
$$
---
## How does the feed-forward network work?
Position-wise FFN:
$$
\text{FFN}(x) = \max(0,\, x W_1 + b_1)\, W_2 + b_2
$$
$W_1 = [I_4 \mid I_4]$ widens $4 \to 8$. $W_2 = [I_4;\, I_4]$ projects $8 \to 4$. Biases are zero.
### Row 0 through FFN
Input: $[1,\, -1,\, 1,\, -1]$.
After $W_1$:
$$
[1,\, -1,\, 1,\, -1,\, 1,\, -1,\, 1,\, -1]
$$
After ReLU (negative entries → 0):
$$
[1,\, 0,\, 1,\, 0,\, 1,\, 0,\, 1,\, 0]
$$
After $W_2$ (sum pairs):
$$
\text{FFN}[0,:] = [2,\, 0,\, 2,\, 0]
$$
Second Add & Norm (same recipe) yields encoder output — the decoder's **memory**:
$$
\text{ENC} =
\begin{bmatrix}
1 & -1 & 1 & -1 \\
-1 & 1 & -1 & 1
\end{bmatrix}
$$
Shape: $(2,\, 4)$.
---
## How does the decoder run masked self-attention?
Shifted target ` i`, embedded and positioned (same rounding trick):
$$
Y =
\begin{bmatrix}
1 & 1 & 0 & 0 \\
0 & 0 & 1 & 1
\end{bmatrix}
$$
Masked self-attention sets future positions to $-\infty$ before softmax:
```mermaid
flowchart LR
A["Raw scores"] --> B["future → −∞"]
B --> C["softmax → 0 there"]
C --> D["attend self + past only"]
```
Position 0 sees only itself → attention weights $[1,\, 0]$.
Position 1 sees positions 0 and 1 equally → weights $[0.5,\, 0.5]$.
After mixing, residual, and layer norm:
$$
\text{DEC}_1 =
\begin{bmatrix}
1 & 1 & -1 & -1 \\
-1 & -1 & 1 & 1
\end{bmatrix}
$$
---
## How does cross-attention connect encoder and decoder?
Cross-attention projections (chosen so every matmul is traceable):
$$
W_Q^{\text{cross}} =
\begin{bmatrix}
1 & 0 \\
0 & -1 \\
0 & 0 \\
0 & 0
\end{bmatrix}
$$
$$
W_{K,V}^{\text{cross}} =
\begin{bmatrix}
1 & 0 \\
0 & 1 \\
0 & 0 \\
0 & 0
\end{bmatrix}
$$
### Matmul: $Q = \text{DEC}_1 \, W_Q^{\text{cross}}$
Row 0 of $\text{DEC}_1$: $[1,\, 1,\, -1,\, -1]$.
$$
Q[0,0] = 1 \cdot 1 + 1 \cdot 0 + (-1) \cdot 0 + (-1) \cdot 0 = 1
$$
$$
Q[0,1] = 1 \cdot 0 + 1 \cdot (-1) + (-1) \cdot 0 + (-1) \cdot 0 = -1
$$
Row 1: $[-1,\, -1,\, 1,\, 1]$.
$$
Q[1,0] = (-1) \cdot 1 + (-1) \cdot 0 + 1 \cdot 0 + 1 \cdot 0 = -1
$$
$$
Q[1,1] = (-1) \cdot 0 + (-1) \cdot (-1) + 1 \cdot 0 + 1 \cdot 0 = 1
$$
$$
Q =
\begin{bmatrix}
1 & -1 \\
-1 & 1
\end{bmatrix}
$$
### Matmul: $K = \text{ENC} \, W_{K,V}^{\text{cross}}$
Row 0 of ENC: $[1,\, -1,\, 1,\, -1]$.
$$
K[0,0] = 1 \cdot 1 + (-1) \cdot 0 + 1 \cdot 0 + (-1) \cdot 0 = 1
$$
$$
K[0,1] = 1 \cdot 0 + (-1) \cdot 1 + 1 \cdot 0 + (-1) \cdot 0 = -1
$$
Row 1: $[-1,\, 1,\, -1,\, 1]$.
$$
K[1,0] = (-1) \cdot 1 + 1 \cdot 0 + (-1) \cdot 0 + 1 \cdot 0 = -1
$$
$$
K[1,1] = (-1) \cdot 0 + 1 \cdot 1 + (-1) \cdot 0 + 1 \cdot 0 = 1
$$
$$
K = V =
\begin{bmatrix}
1 & -1 \\
-1 & 1
\end{bmatrix}
$$
### Matmul: $Q K^{\top}$
$$
(QK^{\top})[0,0] = 1 \cdot 1 + (-1) \cdot (-1) = 2
$$
$$
(QK^{\top})[0,1] = 1 \cdot (-1) + (-1) \cdot 1 = -2
$$
$$
(QK^{\top})[1,0] = (-1) \cdot 1 + 1 \cdot (-1) = -2
$$
$$
(QK^{\top})[1,1] = (-1) \cdot (-1) + 1 \cdot 1 = 2
$$
$$
QK^{\top} =
\begin{bmatrix}
2 & -2 \\
-2 & 2
\end{bmatrix}
$$
### Scale and softmax
$$
S = QK^{\top} / \sqrt{2} =
\begin{bmatrix}
1.414 & -1.414 \\
-1.414 & 1.414
\end{bmatrix}
$$
Row 0: $e^{1.414} \approx 4.10$, $e^{-1.414} \approx 0.243$, sum $\approx 4.34$.
$$
A[0,0] \approx 4.10 / 4.34 \approx 0.944
$$
$$
A[0,1] \approx 0.243 / 4.34 \approx 0.056
$$
Row 1 is symmetric:
$$
A =
\begin{bmatrix}
0.944 & 0.056 \\
0.056 & 0.944
\end{bmatrix}
$$
### Matmul: $O = A V$
$$
O[0,0] = 0.944 \cdot 1 + 0.056 \cdot (-1) = 0.888
$$
$$
O[0,1] = 0.944 \cdot (-1) + 0.056 \cdot 1 = -0.888
$$
$$
O[1,0] = 0.056 \cdot 1 + 0.944 \cdot (-1) = -0.888
$$
$$
O[1,1] = 0.056 \cdot (-1) + 0.944 \cdot 1 = 0.888
$$
$$
O =
\begin{bmatrix}
0.888 & -0.888 \\
-0.888 & 0.888
\end{bmatrix}
$$
FFN + second Add & Norm (same pattern as encoder). I score the final position — the one predicting `am`:
$$
h = [1,\, 0,\, -1,\, 0]
$$
Shape: $(1,\, 4)$.
---
## How do I project to vocabulary and get probabilities?
Linear layer (weight-tied to embeddings in the full model). Here $W_{\text{out}} = I_4$, bias $b = 0$:
$$
z = h \, W_{\text{out}} + b
$$
$$
z[0] = 1 \cdot 1 + 0 \cdot 0 + (-1) \cdot 0 + 0 \cdot 0 = 1
$$
$$
z[1] = 1 \cdot 0 + 0 \cdot 1 + (-1) \cdot 0 + 0 \cdot 0 = 0
$$
$$
z[2] = 1 \cdot 0 + 0 \cdot 0 + (-1) \cdot 1 + 0 \cdot 0 = -1
$$
$$
z[3] = 1 \cdot 0 + 0 \cdot 0 + (-1) \cdot 0 + 0 \cdot 1 = 0
$$
$$
z = [1,\, 0,\, -1,\, 0]
$$
Softmax over vocab `[, i, am, ]`:
$$
e^{1} \approx 2.718
$$
$$
e^{0} = 1
$$
$$
e^{-1} \approx 0.368
$$
$$
\text{denom} = 2.718 + 1 + 0.368 + 1 = 5.086
$$
$$
p[0] = 2.718 / 5.086 \approx 0.535
$$
$$
p[1] = 1 / 5.086 \approx 0.197
$$
$$
p[2] = 0.368 / 5.086 \approx 0.072
$$
$$
p[3] = 1 / 5.086 \approx 0.197
$$
$$
p = [0.535,\, 0.197,\, 0.072,\, 0.197]
$$
The correct word `am` (index 2) gets only $0.072$. The model is confidently wrong — good setup for a nontrivial gradient.
---
## What is the cross-entropy loss with label smoothing?
Plain cross-entropy on `am`:
$$
\mathcal{L}_{\text{CE}} = -\log p_{\text{am}} = -\log 0.072 \approx 2.63
$$
Label smoothing $\epsilon_{ls} = 0.1$ spreads mass onto every class:
$$
q_k = (1 - \epsilon_{ls}) \cdot \mathbb{1}[k = \text{am}] + \epsilon_{ls} / V
$$
$$
\epsilon_{ls} / V = 0.1 / 4 = 0.025
$$
$$
q = [0.025,\, 0.025,\, 0.925,\, 0.025]
$$
Smoothed loss:
$$
\mathcal{L} = -\sum_{k} q_k \log p_k
$$
$$
\mathcal{L} = -\big(0.025 \cdot \log 0.535 + 0.025 \cdot \log 0.197 + 0.925 \cdot \log 0.072 + 0.025 \cdot \log 0.197\big)
$$
$$
\log 0.535 \approx -0.626
$$
$$
\log 0.197 \approx -1.627
$$
$$
\log 0.072 \approx -2.63
$$
$$
\mathcal{L} \approx -\big(0.025 \cdot (-0.626) + 0.025 \cdot (-1.627) + 0.925 \cdot (-2.63) + 0.025 \cdot (-1.627)\big)
$$
$$
\mathcal{L} \approx 2.53
$$
That is the scalar I backprop through.
---
## How does backprop flow through the transformer?
```mermaid
flowchart RL
Z["∂L/∂z = p − q"] --> WO["∂L/∂W_out"]
Z --> H["∂L/∂h"]
H --> LN["LayerNorm Jacobian"]
LN --> F["FFN (ReLU gate)"]
F --> AT["attention (softmax Jacobian)"]
AT --> EMB["∂L/∂embeddings"]
```
### Logits: softmax + cross-entropy collapse
The combined derivative is the elegant $p - q$:
$$
\frac{\partial \mathcal{L}}{\partial z} = p - q
$$
$$
\frac{\partial \mathcal{L}}{\partial z}[0] = 0.535 - 0.025 = 0.510
$$
$$
\frac{\partial \mathcal{L}}{\partial z}[1] = 0.197 - 0.025 = 0.172
$$
$$
\frac{\partial \mathcal{L}}{\partial z}[2] = 0.072 - 0.925 = -0.853
$$
$$
\frac{\partial \mathcal{L}}{\partial z}[3] = 0.197 - 0.025 = 0.172
$$
$$
\frac{\partial \mathcal{L}}{\partial z} = [0.510,\, 0.172,\, -0.853,\, 0.172]
$$
Sum check (must be zero — softmax is shift-invariant):
$$
0.510 + 0.172 + (-0.853) + 0.172 = 0.001 \approx 0
$$
Only `am` is negative → gradient pushes that logit up, pulls the rest down.
### Output layer: $z = h W_{\text{out}}$
$$
\frac{\partial \mathcal{L}}{\partial W_{\text{out}}} = h^{\top} \otimes \frac{\partial \mathcal{L}}{\partial z}
$$
$$
\frac{\partial \mathcal{L}}{\partial h} = \frac{\partial \mathcal{L}}{\partial z} \, W_{\text{out}}^{\top}
$$
With $W_{\text{out}} = I$:
$$
\frac{\partial \mathcal{L}}{\partial h} = [0.510,\, 0.172,\, -0.853,\, 0.172]
$$
### Layer norm backward
For $y_i = \gamma (x_i - \mu) / \sigma$ with $D$ dimensions:
$$
\frac{\partial \mathcal{L}}{\partial x_i} = \frac{\gamma}{\sigma} \Big[ \frac{\partial \mathcal{L}}{\partial y_i} - \frac{1}{D}\sum_j \frac{\partial \mathcal{L}}{\partial y_j} - \frac{x_i - \mu}{\sigma^2} \cdot \frac{1}{D}\sum_j \frac{\partial \mathcal{L}}{\partial y_j}(x_j - \mu) \Big]
$$
The two sum terms strip gradient along the mean and variance directions — layer norm is not "just divide by std."
### FFN backward
ReLU gates the gradient to positions where pre-activation was positive:
$$
\frac{\partial \mathcal{L}}{\partial (x W_1)} = \Big( \frac{\partial \mathcal{L}}{\partial \text{FFN}} \, W_2^{\top} \Big) \odot \mathbb{1}[x W_1 + b_1 > 0]
$$
Negative pre-activations get zero gradient. Dead ReLUs are real.
### Attention backward
Forward: $O = AV$, $A = \text{softmax}(S)$, $S = QK^{\top} / \sqrt{d_k}$.
$$
\frac{\partial \mathcal{L}}{\partial V} = A^{\top} \frac{\partial \mathcal{L}}{\partial O}
$$
$$
\frac{\partial \mathcal{L}}{\partial A} = \frac{\partial \mathcal{L}}{\partial O} \, V^{\top}
$$
Softmax Jacobian (row $i$):
$$
\frac{\partial \mathcal{L}}{\partial S_i} = A_i \odot \Big( \frac{\partial \mathcal{L}}{\partial A_i} - \sum_j \frac{\partial \mathcal{L}}{\partial A_{ij}} A_{ij} \Big)
$$
Then:
$$
\frac{\partial \mathcal{L}}{\partial Q} = \frac{1}{\sqrt{d_k}} \frac{\partial \mathcal{L}}{\partial S} \, K
$$
$$
\frac{\partial \mathcal{L}}{\partial K} = \frac{1}{\sqrt{d_k}} \frac{\partial \mathcal{L}}{\partial S}^{\top} \, Q
$$
$$
\frac{\partial \mathcal{L}}{\partial W_Q} = X^{\top} \frac{\partial \mathcal{L}}{\partial Q}
$$
Masked positions had zero forward attention weight → zero gradient. Causality is preserved in backward pass too. Gradients from all positions and heads **sum** into shared weight matrices.
### The one pattern underneath everything
If $Y = XW$ with $X : (N, d_{\text{in}})$, $W : (d_{\text{in}}, d_{\text{out}})$, $G = \partial L / \partial Y$:
$$
\frac{\partial L}{\partial W} = X^{\top} G
$$
$$
\frac{\partial L}{\partial X} = G \, W^{\top}
$$
Nearly every transformer gradient is this plus reshaping, layer-norm Jacobians, and the softmax Jacobian.
---
## How does Adam update a weight?
Adam hyperparameters from the paper:
$$
\beta_1 = 0.9
$$
$$
\beta_2 = 0.98
$$
$$
\epsilon = 10^{-9}
$$
Learning rate schedule:
$$
\text{lr}(t) = d_{\text{model}}^{-0.5} \cdot \min\!\big( t^{-0.5},\; t \cdot \text{warmup}^{-1.5} \big)
$$
$$
\text{warmup} = 4000
$$
Take one scalar: $W_{\text{out}}[0,0] = 1$, gradient $g = \partial \mathcal{L} / \partial W_{\text{out}}[0,0] = 0.510$.
First moment at step $t = 1$:
$$
m_1 = \beta_1 \cdot 0 + (1 - \beta_1) \cdot g = 0.9 \cdot 0 + 0.1 \cdot 0.510 = 0.0510
$$
Second moment:
$$
v_1 = \beta_2 \cdot 0 + (1 - \beta_2) \cdot g^2 = 0.98 \cdot 0 + 0.02 \cdot 0.510^2 = 0.00520
$$
Bias correction:
$$
\hat{m}_1 = m_1 / (1 - \beta_1^1) = 0.0510 / 0.1 = 0.510
$$
$$
\hat{v}_1 = v_1 / (1 - \beta_2^1) = 0.00520 / 0.02 = 0.260
$$
Adam step direction:
$$
\hat{m}_1 / (\sqrt{\hat{v}_1} + \epsilon) = 0.510 / (\sqrt{0.260} + 10^{-9}) = 0.510 / 0.510 \approx 1.0
$$
### Step 1 (warmup, real $d_{\text{model}} = 512$)
$$
\text{lr}(1) = 512^{-0.5} \cdot 4000^{-1.5} = 1.75 \times 10^{-7}
$$
Weight update:
$$
W_{\text{out}}[0,0] \leftarrow 1 - 1.75 \times 10^{-7} \cdot 1.0 = 0.99999983
$$
Barely moves — warmup is conservative on purpose.
### Step 4000 (end of warmup)
$$
\text{lr}(4000) = 512^{-0.5} \cdot 4000^{-0.5} \approx 0.000699
$$
$$
W_{\text{out}}[0,0] \leftarrow 1 - 0.000699 \cdot 1.0 = 0.999301
$$
Now you can *see* the weight shift. Run millions of steps on billions of tokens and `am` eventually beats ``.
---
## What invariants should I check while deriving?
I treat these as unit tests for the hand derivation:
- Write tensor shape beside every equation. Residual add requires identical shapes.
- Attention weights sum to $1$ across the key axis per query row.
- Masked positions get exactly zero probability forward and zero gradient backward.
- LayerNorm outputs have approximately zero mean and unit variance per token row.
- Softmax-cross-entropy gradients $\partial \mathcal{L} / \partial z$ sum to zero across the vocabulary.
Production code fuses kernels, uses mixed precision, checkpoints activations, and parallelizes — but none of that changes the calculus above. I use this derivation to debug toy implementations and to understand memory/compute flow before trusting a framework.
---
## Where does one training step leave me?
Forward: embeddings → positional encoding → encoder attention (every matmul) → layer norm → FFN → decoder masked attention → cross-attention → vocab projection → softmax.
Loss: $\mathcal{L} = 2.53$ with label smoothing.
Backward: $\partial \mathcal{L} / \partial z = p - q$ → chain rule through every weight.
Update: Adam ticks $W_{\text{out}}[0,0]$ from $1.000000$ toward $0.999301$ by step 4000.
Next batch runs the same machine on slightly better weights. The whole thing is a matmul, a softmax, a normalize, a ReLU, an outer product, and the chain rule — stacked $N$ deep on each side, iterated over trillions of tokens.
Conceptual version: [how attention works in transformers](/blog/attention-in-transformers-explained/). Building blocks: [token embeddings](/blog/token-embeddings-explained/), [positional encoding](/blog/positional-encoding-explained/), [autograd from scratch](/blog/build-autograd-from-scratch/). Source: [Attention Is All You Need (2017)](https://arxiv.org/abs/1706.03762).