Recipe

Build and train a small Transformer

Use when: you have Sections 1-5's pieces (embeddings, positional encoding, self-attention, the FFN, residual and norm, masked self-attention, cross-attention) and need to assemble them into one trainable encoder-decoder model

Goal: a complete, trained encoder-decoder Transformer that maps an input sequence to an output sequence

Taught in Training a small Transformer

1. 1

   Turn each input token into its learned embedding vector (Chapter 17). Then add its positional encoding (Section 1) elementwise. Do this once for the encoder's input sequence. Do it again for the decoder's target sequence, shifted right by one position — so position t's input is target token t-1.

   Example

   For the toy task above, input token "7" at position 1 gets embedding \mathrm{Embed}(7). Then PE(1) is added elementwise: x_1 = \mathrm{Embed}(7) + PE(1).

   Why: Attention has no built-in notion of order (Section 1). Position must already be part of the vector before the first block runs.

2. 2

   Wrap self-attention and the FFN in pre-norm RMSNorm and residual connections (Sections 3-4): x_1 = x + \mathrm{SelfAttn}(\mathrm{RMSNorm}(x)), then x_2 = x_1 + \mathrm{FFN}(\mathrm{RMSNorm}(x_1)).

   Why: This lets a token's representation gather context from the whole input (attention). Then it transforms that gathered information (the FFN). Neither sublayer's gradient vanishes at depth.

3. 3

   Repeat the previous step's block shape N times. Each block has its own independently-learned weights. Block k's output feeds directly into block k+1 as its input.

   • a small toy task (a few hundred training examples, short sequences) — N=2 or 3 is usually enough to converge.
   • a task needing longer-range structure or a larger vocabulary — increase N (or the model width d) before assuming more training data alone will fix underfitting.

   Why: depth lets the representation compose several rounds of "gather context, then transform it" rather than only one.

4. 4

   Add masked self-attention, cross-attention against the finished encoder output, and the FFN. Each is wrapped the same pre-norm way (Section 5): x_1 = x + \mathrm{MaskedSelfAttn}(\mathrm{RMSNorm}(x)), x_2 = x_1 + \mathrm{CrossAttn}(\mathrm{RMSNorm}(x_1), \text{encoder output}), x_3 = x_2 + \mathrm{FFN}(\mathrm{RMSNorm}(x_2)).

   Why: The decoder needs two things before it can transform information (the FFN): what it has already generated (masked self-attention), and the full input it's responding to (cross-attention).

5. 5

   Repeat the previous step's block shape N times, same as the encoder stack. The SAME fixed encoder output feeds cross-attention at every decoder block.

   Why: matching depth on both sides is a common default, though nothing requires the encoder and decoder to use the same N.

6. 6

   Apply one final linear layer. It maps the last decoder block's output at each position to a vector of vocabulary-size logits. Then apply softmax to get a probability distribution over the next token.

   Why: This is what turns a d-dimensional decoder vector into an actual predicted next token.

7. 7

   Feed the encoder the full input. Feed the decoder the shifted-right target (teacher forcing). Compute cross-entropy loss between each predicted position and the true next token. Run Chapter 11's training loop with a learning-rate warm-up-then-decay schedule.

   Why: A Transformer's early gradients come from randomly-initialized attention weights. They are large and poorly calibrated. Warm-up prevents an early, oversized update from derailing training, before attention has found a reasonable starting point.