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 conceptually, I could read token embeddings and positional encoding 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; the dimensions are a toy so you can verify every line by hand. If you’ve built autograd from scratch, this is the derivation you want open while debugging — and when something looks wrong, numerical gradient checking is how I sanity-check each layer.
[!NOTE] Conceptual attention (queries, keys, values, no arithmetic) lives in how attention works. 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.
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<br/>(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 <bos> 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: [<bos>, i, am, <eos>].
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$:
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$):
Stack into the encoder input matrix $X_{\text{raw}} \in \mathbb{R}^{2 \times 4}$:
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:
For $d_{\text{model}} = 4$, the frequency divisors are:
Position $pos = 0$:
Position $pos = 1$:
Add to embeddings elementwise. For this toy I round the sum to keep later arithmetic tractable (same as the original derivation):
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:
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}$):
Matmul: $Q = X W_Q^{(1)}$
Row 0 of $X$: $[1,\, 0,\, 1,\, 0]$.
Row 1 of $X$: $[0,\, 1,\, 0,\, 1]$.
Same $W_K^{(1)}$, so $K^{(1)} = Q^{(1)}$.
Matmul: $V = X W_V^{(1)}$
Shapes: $Q, K, V \in \mathbb{R}^{2 \times 2}$.
How do I score, scale, softmax, and mix values?
Matmul: $Q K^{\top}$
Scale by $\sqrt{d_k}$
Softmax row 0
Softmax row 1
By symmetry:
Each row sums to $1$.
Matmul: $O^{(1)} = A^{(1)} V^{(1)}$
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$:
Shape: $(2,\, 4)$.
What is Add & Norm (layer norm)?
Residual connection plus layer norm:
With $\gamma = 1$, $\beta = 0$, $\epsilon = 0$.
Residual add after attention
Layer norm on row 0
Deviations from mean:
Normalized row 0:
Row 1 is antisymmetric:
How does the feed-forward network work?
Position-wise FFN:
$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$:
After ReLU (negative entries → 0):
After $W_2$ (sum pairs):
Second Add & Norm (same recipe) yields encoder output — the decoder’s memory:
Shape: $(2,\, 4)$.
How does the decoder run masked self-attention?
Shifted target <bos> i, embedded and positioned (same rounding trick):
Masked self-attention sets future positions to $-\infty$ before softmax:
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:
How does cross-attention connect encoder and decoder?
Cross-attention projections (chosen so every matmul is traceable):
Matmul: $Q = \text{DEC}_1 \, W_Q^{\text{cross}}$
Row 0 of $\text{DEC}_1$: $[1,\, 1,\, -1,\, -1]$.
Row 1: $[-1,\, -1,\, 1,\, 1]$.
Matmul: $K = \text{ENC} \, W_{K,V}^{\text{cross}}$
Row 0 of ENC: $[1,\, -1,\, 1,\, -1]$.
Row 1: $[-1,\, 1,\, -1,\, 1]$.
Matmul: $Q K^{\top}$
Scale and softmax
Row 0: $e^{1.414} \approx 4.10$, $e^{-1.414} \approx 0.243$, sum $\approx 4.34$.
Row 1 is symmetric:
Matmul: $O = A V$
FFN + second Add & Norm (same pattern as encoder). I score the final position — the one predicting am:
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$:
Softmax over vocab [<bos>, i, am, <eos>]:
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:
Label smoothing $\epsilon_{ls} = 0.1$ spreads mass onto every class:
Smoothed loss:
That is the scalar I backprop through.
How does backprop flow through the transformer?
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$:
Sum check (must be zero — softmax is shift-invariant):
Only am is negative → gradient pushes that logit up, pulls the rest down.
Output layer: $z = h W_{\text{out}}$
With $W_{\text{out}} = I$:
Layer norm backward
For $y_i = \gamma (x_i - \mu) / \sigma$ with $D$ dimensions:
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:
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}$.
Softmax Jacobian (row $i$):
Then:
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$:
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:
Learning rate schedule:
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$:
Second moment:
Bias correction:
Adam step direction:
Step 1 (warmup, real $d_{\text{model}} = 512$)
Weight update:
Barely moves — warmup is conservative on purpose.
Step 4000 (end of warmup)
Now you can see the weight shift. Run millions of steps on billions of tokens and am eventually beats <bos>.
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. Building blocks: token embeddings, positional encoding, autograd from scratch. Source: Attention Is All You Need (2017).
Support the writing
If this post helped, a coffee keeps the deep dives, paper picks, and news digests coming.
Buy Me a Coffee