The transformer block

Level 2 assembles the block one piece at a time, starting with the picture of a team that meets, then works alone.

Everyday picture A team works in rounds. Each round starts with a meeting, where everyone listens to everyone else and takes notes on what's relevant to them (that's attention). Then comes desk work: each person goes back to their own desk and thinks through their notes alone, without talking to anyone (that's the feed-forward network). Nobody throws away their old notes; they only add to them (the residual connection). A large model runs dozens of these rounds: GPT-2 small has 12, big models have 80 or more.

Tiny worked example Take one token whose vector is x = (1, 2). The meeting produces a correction (0.1, −0.3); adding it gives (1.1, 1.7). Desk work then adds (−0.2, 0.4), giving (0.9, 2.1). The token's vector has been edited twice, never replaced.

flowchart TD
  IN[Input vectors<br/>one per token] --> N1[Layer norm]
  N1 --> AT[Multi-head attention<br/>the meeting: mix across tokens]
  AT --> R1((Add))
  IN --> R1
  R1 --> N2[Layer norm]
  N2 --> FF[Feed-forward network<br/>desk work: each token alone]
  FF --> R2((Add))
  R1 --> R2
  R2 --> OUT[To the next block]

The math and the code

x = [1.0, 2.0]
# what Attn(LN(x)) returned
attn = [0.1, -0.3]
# x ← x + Attn(LN(x))
x = [x_j + a_j for x_j, a_j in zip(x, attn)]
[round(x_j, 1) for x_j in x]  # → [1.1, 1.7]
# what FFN(LN(x)) returned
ffn = [-0.2, 0.4]
# x ← x + FFN(LN(x))
x = [x_j + f_j for x_j, f_j in zip(x, ffn)]
[round(x_j, 1) for x_j in x]  # → [0.9, 2.1]

In code: TransformerBlock holds one primer.ml.attention.MultiHeadAttention, one FeedForward and the two norms' learned gains and biases.

Why it matters The block's output has the same shape as its input, so blocks stack like Lego. And because the input is never overwritten, switching both sub-layers off gives back the input unchanged (a scenario in the spec), which is why very deep stacks train at all.

Everyday picture A teacher grading on a curve: subtract the class average from each score, then divide by how spread out the scores are. After that, "2 above average" means the same thing in every class. LayerNorm does this to the numbers inside one token's vector, so no token's numbers can drift huge or tiny as they pass through dozens of blocks.

Tiny worked example Normalize (1, 2, 3, 4). The mean (average) is 2.5. The variance (average squared distance from the mean) is (2.25 + 0.25 + 0.25 + 2.25) / 4 = 1.25, so the spread (standard deviation, its square root) is 1.118. Result: (−1.342, −0.447, 0.447, 1.342). Multiply the input by 100 and the result is identical.

flowchart LR
  X["one token: (1, 2, 3, 4)"] --> M["subtract the mean 2.5<br/>(−1.5, −0.5, 0.5, 1.5)"]
  M --> S["divide by the spread 1.118<br/>(−1.342, −0.447, 0.447, 1.342)"]
  S --> G["× learned gain, + learned bias<br/>(start as 1 and 0)"]

The math and the code

import math
x = [1, 2, 3, 4]
d = len(x)
# μ = (1/d) Σ x_j
mu = sum(x) / d
# σ² = (1/d) Σ (x_j − μ)²
sigma2 = sum((x_j - mu) ** 2 for x_j in x) / d
mu, sigma2  # → (2.5, 1.25)
# learned gain and bias start at 1 and 0
eps, gamma, beta = 1e-5, 1.0, 0.0
[round(gamma * (x_j - mu) / math.sqrt(sigma2 + eps) + beta, 3) for x_j in x]  # → [-1.342, -0.447, 0.447, 1.342]

Why it matters Without normalization, activations drift layer after layer until training blows up or stalls. It normalizes per token (not per batch, unlike the BatchNorm used in image networks), so it works for any batch size and sequence length.

Everyday picture Back at their desk, each person spreads their notes out over a desk four times wider than the notebook, underlines what matters, and writes a short summary back into the notebook. Nobody talks to anyone.

Tiny worked example With model width 8, the network widens each token to 32 numbers, applies GELU to each, and narrows back to 8. Parameters: 8·32 + 32 + 32·8 + 8 = 552. GELU on single numbers: GELU(1) = 0.841, GELU(10) = 10.0, GELU(−3) = −0.004, GELU(0) = 0.

flowchart LR
  X["token vector<br/>d numbers"] --> W1["× W1 + b1<br/>widen to 4·d"]
  W1 --> G["GELU on each number<br/>keep positives, squash negatives"]
  G --> W2["× W2 + b2<br/>back to d"]
  W2 --> Y["correction for this token"]

The math and the code A matrix multiply turns a list of numbers into a list of numbers: each output number is the dot product of with one column of (multiply matching numbers, add them up).

import math
def gelu(z):
    return 0.5 * z * (1 + math.tanh(math.sqrt(2 / math.pi) * (z + 0.044715 * z ** 3)))
[round(gelu(z), 3) for z in (1, 10, -3, 0)]  # → [0.841, 10.0, -0.004, 0.0]
d = 8
# W1 (8 × 32), b1, W2 (32 × 8), b2
d * 4 * d + 4 * d + 4 * d * d + d  # → 552

Why it matters Without a nonlinearity like GELU, the two matrix multiplies would collapse into one, and stacking layers would add nothing.

Everyday picture An assembly line. Token ids go in one end, get turned into vectors, pass through a row of identical workstations (the blocks), get a final polish (the last norm), and come out the far end as a score for every word in the dictionary.

Tiny worked example A model with a 50-token vocabulary, width 32, 2 blocks and 16 positions. The ids [3, 1, 4, 1, 5] become a 5 × 32 grid, pass through both blocks as 5 × 32, and come out as a 5 × 50 grid of logits (raw scores): row 5 scores every possible 6th token. It has 27,328 parameters.

flowchart LR
  IDS["ids [3,1,4,1,5]"] --> TE["token table lookup<br/>(50 × 32)"]
  POS["positions 0..4"] --> PE["position table lookup<br/>(16 × 32)"]
  TE --> ADD(("+"))
  PE --> ADD
  ADD --> B1["block 1"] --> B2["block 2"] --> LN["final layer norm"]
  LN --> OUT["× token tableᵀ<br/>logits (5 × 50)"]
  TE -. "same table (tied)" .- OUT

The math and the code

# LN(h) for the last position
ln_h = [1.0, -1.0]
# one row per vocabulary entry
E = [[1.0, 0.0], [0.0, 1.0], [-1.0, 1.0]]
# LN(h) Eᵀ: a dot product with each row
[sum(h_k * e_k for h_k, e_k in zip(ln_h, row)) for row in E]  # → [1.0, -1.0, -2.0]

In code: TinyGPT.hidden runs everything up to the final norm, and TinyGPT.n_params counts the 27,328 weights.

Why it matters This is the whole forward pass of GPT-2, Llama and Claude-style models in miniature. Real models differ in size, not in kind.

Everyday picture An editor reads the whole page before commenting on any sentence (encoder). A storyteller tells the story word by word and can't peek at words they haven't said yet (decoder). A translator reads the whole source sentence, then writes the translation word by word (encoder-decoder).

Tiny worked example With 3 tokens, an encoder's attention may use all 9 (query, key) pairs; a decoder only the 6 on or below the diagonal. In the spec, changing token 5 changes token 1's output in an encoder block, and leaves it untouched in a decoder block.

flowchart TB
  subgraph ENC["Encoder-only (BERT)"]
    e1[every token sees every token] --> e2[a vector per token<br/>classify, embed, tag]
  end
  subgraph DEC["Decoder-only (GPT, Claude, Llama)"]
    d1[each token sees only the past] --> d2[predict the next token<br/>generate, chat, agents]
  end
  subgraph ED["Encoder-decoder (T5, original transformer)"]
    x1[encoder reads the source<br/>bidirectionally] --> x2[decoder writes the output causally,<br/>attending to the encoder too]
  end

In code: TransformerBlock with its causal flag on is a decoder block and with it off an encoder block; mask_patterns runs one of each on the same input to draw the figure.

Why it matters Embedding models (primer.ml.embeddings) are usually encoders; chat models and agents are decoders.

Everyday picture Counting the bricks in a Lego model from its blueprint, without opening the box.

Tiny worked example For width d, one block holds: attention 4·d² (W_q, W_k, W_v, W_o), feed-forward 8·d² (two d × 4d matrices), plus small bias and norm terms. With d = 32: 12 · 1,024 = 12,288, plus 160 (biases) plus 128 (norms) = 12,576 per block. Two blocks, plus tables of 50 × 32 and 16 × 32 and a final norm of 64, gives 27,328. The same recipe on GPT-2 small (d = 768, 12 blocks, 50,257 tokens, 1,024 positions) gives 124,402,944. GPT-2 also puts a bias on its attention projections, which the tiny model leaves out: 4·d = 3,072 more per block, 36,864 in all, and that brings it to exactly its published 124,439,808.

flowchart LR
  B["one block"] --> A["attention<br/>4·d²"]
  B --> F["feed-forward<br/>8·d²"]
  B --> N["norms and biases<br/>~13·d (tiny)"]
  A & F --> T["≈ 12·d² per block"]
  T --> ALL["× L blocks + embeddings<br/>vocab·d + positions·d"]

The math

L, d, V = 12, 768, 50_257
# 12·L·d² and V·d
blocks, table = 12 * L * d ** 2, V * d
print(f"{blocks:,} + {table:,} = {blocks + table:,}")  # → 84,934,656 + 38,597,376 = 123,532,032
# the share it leaves out
round((124_439_808 - (blocks + table)) / 124_439_808, 3)  # → 0.007

In code: param_breakdown splits a GPT-2-shaped model's count into embeddings, attention, feed-forward and norms for the figure, and TransformerBlock.n_params counts one real block's weights.

Why it matters Parameters × bytes per parameter is the memory a model needs just to load (see primer.ml.inference for the arithmetic).

Everyday picture A hospital triage desk. Instead of one general practitioner seeing every patient, the desk sends each patient to the two most relevant specialists out of eight. The hospital employs eight doctors' worth of expertise, but each patient only takes up two doctors' time.

Tiny worked example A router scores one token against 4 experts: (2.0, 1.0, 0.5, −1.0). Keep the top 2 (experts 0 and 1) and softmax just those two: e² / (e² + e¹) = 7.39 / 10.11 = 0.731, and 0.269. The token's output is 0.731 × expert 0's output + 0.269 × expert 1's output; experts 2 and 3 never run for this token.

flowchart LR
  X[token vector] --> R["router<br/>one score per expert"]
  R --> K["keep top 2<br/>softmax over them"]
  K -->|0.731| E0[expert 0: an FFN]
  K -->|0.269| E1[expert 1: an FFN]
  K -.->|skipped| E2[expert 2]
  K -.->|skipped| E3[expert 3]
  E0 --> S(("weighted sum"))
  E1 --> S
  S --> Y[output]

The math and the code

import math
r, k = [2.0, 1.0, 0.5, -1.0], 2
# TopK(r)
top_k = sorted(range(len(r)), key=lambda i: r[i], reverse=True)[:k]
top_k  # → [0, 1]
exps = [math.exp(r[i]) for i in top_k]
# g: softmax over the kept scores only
[round(e / sum(exps), 3) for e in exps]  # → [0.731, 0.269]
d, n_experts = 16, 8
# one expert is one feed-forward network
expert = d * 4 * d + 4 * d + 4 * d * d + d
# W_r
router = d * n_experts
# one, all held, touched per token
expert, n_experts * expert + router, k * expert + router  # → (2128, 17152, 4384)

A router left alone tends to play favourites, overloading some experts while others starve. Training adds a small load-balancing loss (Switch Transformer):

n = 4
def balance(f, P):
    # n Σ f_i P_i
    return n * sum(f_i * P_i for f_i, P_i in zip(f, P))
# every expert gets a quarter
balance([0.25] * n, [0.25] * n)  # → 1.0
# everything goes to expert 0
balance([1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0])  # → 4.0

In code: MixtureOfExperts.n_params and MixtureOfExperts.active_params give the 17,152 and 4,384 above, and MixtureOfExperts.tokens_per_expert counts the bars in the figure.

Why it matters MoE lets a model hold far more knowledge (parameters) for the same compute per token, which is why many frontier models use it. The price is memory: every expert must be loaded even though few run per token.

Everyday picture Every parameter does one multiply and one add for every token that passes through it, like a toll booth that charges two coins per car.

Tiny worked example A 7B-parameter model generating one token: 2 × 7e9 = 1.4e10 operations (14 GFLOPs). Training a 70B model on a trillion tokens: 6 × 70e9 × 1e12 = 4.2e23 operations.

flowchart LR
  F["forward pass<br/>2 FLOPs per parameter per token"] --> L[loss]
  L --> B["backward pass<br/>≈ 4 FLOPs per parameter per token"]
  B --> T["training total ≈ 6·N per token"]
  F --> I["inference total ≈ 2·N per token"]

The math

def C_infer(N):
    # per generated token
    return 2 * N
def C_train(N, D):
    # over all D training tokens
    return 6 * N * D
print(f"{C_infer(7e9):.2g}  {C_train(70e9, 1e12):.2g}  {C_train(175e9, 300e9):.3g}")  # → 1.4e+10 4.2e+23 3.15e+23

Why it matters These two lines let you estimate GPU-hours, serving cost and training budgets on the back of an envelope.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.