Transformers from First PrinciplesPost 5 of 18 · 18 of 18 written
  1. How to use this guide
  2. 01The machinery: how a neural network computes and learns
  3. 02Making deep networks trainable
  4. 03Turning text into numbers
  5. 04Attention from first principles
  6. 05The Transformer block
  7. 06Position: telling the model where each token is
  8. 07Architecture families and training objectives
  9. 08Pretraining at scale
  10. 09Inference and serving
  11. 10Efficiency variants: MoE, long context and beyond attention
  12. 11Post-training: from text predictor to assistant
  13. 12Adapting models: prompting, fine-tuning, LoRA and distillation
  14. 13Embeddings, retrieval and RAG
  15. 14Beyond text: vision, speech and multimodal models
  16. 15Inside the model: interpretability
  17. 16Evaluation, failure modes and security
  18. 17From model to system
  19. 18Fluency kit

5. The Transformer block

Transformers · 5 of 184 min readNov 2026

A transformer is a stack of identical blocks. Each block does two things in turn: multi-head attention moves information between positions, then an MLP processes each position on its own. Both read from a shared vector per token - the residual stream - and add their results back into it.

Multi-head attention

One attention operation gives each token one weighted average, so it can follow one kind of relationship at a time. Language needs many at once: the subject of this verb, the previous token, the matching bracket, the earlier mention of this name. Multi-head attention runs h attention operations in parallel, each in a subspace of size d_h = d / h, then combines them:

headi=Attention(XWQ(i), XWK(i), XWV(i))\text{head}_i = \text{Attention}\big(X W_Q^{(i)},\, X W_K^{(i)},\, X W_V^{(i)}\big)
MHA(X)=Concat(head1,…,headh) WO\text{MHA}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\, W_O

The total cost equals one big head of width d, because each head works in d/h dimensions. W_O (d × d) mixes the heads’ outputs back into the residual stream.

In code nobody loops over heads. One matrix multiply produces Q, K and V for all heads (d → 3d), and a reshape splits them:

qkv = x @ W_qkv                               # (B, T, 3d)
q, k, v = qkv.split(d, dim=-1)                 # each (B, T, d)
q = q.view(B, T, h, d_h).transpose(1, 2)       # (B, h, T, d_h); same for k, v
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)  # (B, h, T, d_h)
y = y.transpose(1, 2).reshape(B, T, d) @ W_o   # (B, T, d)

The MLP (feed-forward network)

After attention, each token passes independently through the same two-layer network:

MLP(x)=W2 σ(W1x)\text{MLP}(x) = W_2\, \sigma(W_1 x)

W₁ expands d to a hidden size (classically 4d), σ is the nonlinearity, and W₂ projects back to d. Most modern LLMs use a gated version, SwiGLU, with three matrices and a hidden size near (8/3)·d so the parameter count stays the same:

SwiGLU(x)=W2(SiLU(W1x)⊙W3x)\text{SwiGLU}(x) = W_2\left( \text{SiLU}(W_1 x) \odot W_3 x \right)

The division of labour: attention is communication between tokens; the MLP is computation within a token. MLPs hold about two thirds of the parameters and appear to store much of the model’s factual association (Part 15).

Residual stream and pre-norm

Each block applies:

x←x+Attention(Norm(x))x \leftarrow x + \text{Attention}(\text{Norm}(x))
x←x+MLP(Norm(x))x \leftarrow x + \text{MLP}(\text{Norm}(x))

This is the pre-norm arrangement: normalize each sublayer’s input but leave the residual path untouched. The 2017 original used post-norm, x ← Norm(x + F(x)), which puts a normalization on the main path; it trains less stably at depth and needs careful warmup. Pre-norm is now standard, with one final norm after the last block.

[embedded content: decoder-only forward pass · one block shown, repeated L times]

The lines on the left are the residual (skip) path: each sublayer’s output is added to an untouched copy of its input at the ⊕, so the stream flows straight from embedding to unembedding.

The full forward pass of a decoder-only model

  1. Token ids (B, T) from the tokenizer.
  2. Embedding lookup → (B, T, d). With learned absolute positions, a position vector is added here; with RoPE, position is applied inside attention instead.
  3. L blocks, each: norm → causal multi-head attention → add; norm → MLP → add. The shape stays (B, T, d) throughout.
  4. Final norm.
  5. Unembedding → logits (B, T, V).
  6. Training: cross-entropy between the logits at position t and the actual token at t + 1, averaged over all positions - T training examples from one pass. Inference: take the logits at the last position, sample a token, append it, repeat.

Counting parameters

Per block, attention has four d × d matrices (Q, K, V, O), which is 4d², and the MLP has two d × 4d matrices, which is 8d². So:

N≈12 L d2+VdN \approx 12\, L\, d^2 + V d

Check against GPT-2 small (L = 12, d = 768, V = 50,257): 12 × 12 × 768² ≈ 85M in the blocks, plus 38.6M token embeddings and 0.8M position embeddings, ≈ 124M - the published size. For GPT-3 (L = 96, d = 12,288): 12 × 96 × 12,288² ≈ 174B, matching the 175B headline. Doing this in your head is a fluency marker.

Design knobs you will see in model cards

Knob Typical choice What it trades
Depth L vs width d deep-and-narrow vs wide-and-shallow depth adds serial latency; width parallelizes better
Heads h and head size d_h d_h = 64 or 128 more heads = more relation types, each smaller
KV heads fewer than query heads (GQA) smaller KV cache at inference (Part 9)
MLP type and ratio SwiGLU, about (8/3)·d or larger quality vs parameters
Norm RMSNorm, pre-norm stability and speed
Position RoPE extrapolation and long context (Part 6)
Bias terms often removed slightly fewer parameters, more stable
Tied embeddings common in small models saves V × d parameters

Questions you should now be able to answer

  • Why does multi-head attention cost the same as single-head attention of width d?
  • What can attention do that the MLP cannot, and vice versa?
  • Why is pre-norm more stable than post-norm? Point to the residual path.
  • Estimate the parameter count for L = 32, d = 4096, V = 32,000, with a 4d MLP and no GQA.
  • Trace the tensor’s shape from token ids to logits.