Transformers in Depth

You already know the story: in 2017 eight researchers threw out recurrence and let every token look at every other token at once (Attention Is All You Need). You also know the intuition: a large language model predicts the next token (What Is a Large Language Model?). This page closes the gap between those two pictures and working code. By the end you will have written a complete, if tiny, Transformer in NumPy, printed the shape of every array that passes through it, and counted the parameters of GPT-2 small to the last digit. Nothing here needs a GPU. It needs patience with matrix shapes, and the willingness to run every snippet yourself.

Queries, keys and values

Start with the problem attention solves. A sequence of tokens arrives as a matrix X with one row per token: T rows, each a vector of d_model numbers (the token’s embedding, plus something to say where it sits — more on that later). Each token needs to gather information from the others. In “Ali makan nasi lemak”, the vector for lemak is more useful if it has absorbed something from nasi; a pronoun needs to find its antecedent. A recurrent network (see RNNs, LSTMs and Sequence Models) passes information along one step at a time. Attention lets every token reach every other token directly, in one step, with learned weights deciding how much each one matters.

It does this by giving each token three different roles, each produced by multiplying the token’s vector by a learned matrix:

• Query q = x W^Q — what this token is looking for.
• Key k = x W^K — what this token offers to be found by.
• Value v = x W^V — what this token actually hands over if it is chosen.

The query–key–value vocabulary comes from databases and information retrieval: you send a query, it is compared with keys, and you receive the values of whichever keys match. The difference is that attention matches softly. Every key matches a little; the match scores are turned into weights that sum to one; and the result is a weighted average of all the values. Because everything is a smooth function of W^Q, W^K and W^V, backpropagation can adjust those matrices so that the right tokens end up paying attention to each other. Nobody writes the rules; the gradient finds them.

Why three separate matrices rather than comparing token vectors directly? Because “what I am looking for” and “what I am” are different questions. An adjective looking for its noun should match nouns, not other adjectives. Separate projections let the model learn asymmetric relationships: token A may attend strongly to B without B attending to A.

Scaled dot-product attention

Stack all the queries into a matrix Q (shape T × d_k), and likewise K and V. Then the whole operation is one line:

Read it right to left as a recipe: compare (Q Kᵀ), calm down (÷ √d_k), normalise (softmax, row by row), mix (× V). Row i of the result is a new vector for token i: a blend of every token’s value, weighted by how well its key matched token i’s query. Here it is in NumPy, on four tokens with random weights.

import numpy as np

np.random.seed(0)
tokens = ["Ali", "makan", "nasi", "lemak"]
d_model, d_k = 8, 4
X = np.random.randn(4, d_model)            # one row per token (embeddings)
W_q = np.random.randn(d_model, d_k) / np.sqrt(d_model)   # small, as at initialisation
W_k = np.random.randn(d_model, d_k) / np.sqrt(d_model)
W_v = np.random.randn(d_model, d_k) / np.sqrt(d_model)

def softmax(z):
    z = z - z.max(axis=-1, keepdims=True)  # for numerical stability
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V):
    scores = Q @ K.T / np.sqrt(K.shape[-1])
    weights = softmax(scores)
    return weights @ V, weights

Q, K, V = X @ W_q, X @ W_k, X @ W_v
out, w = attention(Q, K, V)
print("Q, K, V shapes:", Q.shape, K.shape, V.shape)
print("attention weights (rows = query token):")
for t, row in zip(tokens, w):
    print(f"  {t:<6}", np.round(row, 2), " sum =", round(row.sum(), 2))
print("output shape:", out.shape)
print("scores for lemak:", np.round(Q[3] @ K.T / np.sqrt(d_k), 2))

Q, K, V shapes: (4, 4) (4, 4) (4, 4)
attention weights (rows = query token):
  Ali    [0.89 0.02 0.08 0.01]  sum = 1.0
  makan  [0.51 0.15 0.21 0.13]  sum = 1.0
  nasi   [0.09 0.17 0.59 0.15]  sum = 1.0
  lemak  [0.96 0.03 0.01 0.  ]  sum = 1.0
output shape: (4, 4)
scores for lemak: [ 4.15  0.69 -0.52 -2.84]

Every row of weights sums to 1, as softmax guarantees. The last line shows the raw scaled scores behind the lemak row: 4.15 for Ali is far above the rest, and softmax’s exponential turns that lead into a weight of 0.96. Notice what the numbers do not mean: these weights are random, so lemak attending to Ali is an accident of the seed, not linguistics. Training is what turns W^Q and W^K into matrices whose matches mean something. Notice also the shapes: the weights are a 4 × 4 table — one row per query, one column per key — which is why attention’s cost grows with the square of the sequence length.

The softmax subtracts each row’s maximum before exponentiating. That changes nothing mathematically (the constant cancels in the division) but stops np.exp overflowing on large scores. Every serious implementation does this. The weight matrices are divided by √d_model when created, roughly as real initialisation schemes do (see Training Well), so the scores start at a sensible size.

Why divide by √d_k? If the entries of q and k are independent with mean 0 and variance 1, their dot product is a sum of d_k such products, so its variance is d_k and its standard deviation √d_k. Wide heads produce large scores; large scores push softmax towards a one-hot spike; and a spiked softmax has almost zero gradient, so learning stalls. Vaswani and colleagues gave exactly this argument in 2017. You can watch it happen:

import numpy as np

rng = np.random.default_rng(42)

def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

for d in [4, 64, 512]:
    q = rng.standard_normal((1000, d))          # 1,000 random queries
    k = rng.standard_normal((1000, 16, d))      # 16 random keys for each
    raw = np.einsum("nd,nkd->nk", q, k)         # 1000 x 16 dot products
    top_raw = softmax(raw).max(axis=1).mean()
    top_scaled = softmax(raw / np.sqrt(d)).max(axis=1).mean()
    print(f"d={d:<4} std of q.k = {raw.std():5.1f}   "
          f"avg top weight: unscaled {top_raw:.2f}, scaled {top_scaled:.2f}")

d=4    std of q.k =   2.0   avg top weight: unscaled 0.43, scaled 0.24
d=64   std of q.k =   8.0   avg top weight: unscaled 0.84, scaled 0.24
d=512  std of q.k =  22.5   avg top weight: unscaled 0.95, scaled 0.25

Unscaled, the typical dot product grows from about 2 to about 22.5 as the dimension rises, and the biggest of 16 weights climbs from 0.43 to 0.95 — attention collapses onto a single key before training has even begun. Scaled, the top weight stays near 0.24 at every width. The scaling keeps the softmax in its responsive range regardless of head size.

Causal masks and multiple heads

A model trained to predict the next token must not see it. If token 3 could attend to token 4 during training, it would simply copy the answer and learn nothing useful. GPT-style models therefore use a causal mask: before the softmax, every score where the key comes after the query is set to −∞, so its weight becomes exactly zero.

import numpy as np

np.random.seed(0)
T, d_model, d_k = 4, 8, 4
X = np.random.randn(T, d_model)
W_q, W_k = (np.random.randn(d_model, d_k) / np.sqrt(d_model) for _ in range(2))

scores = (X @ W_q) @ (X @ W_k).T / np.sqrt(d_k)
mask = np.triu(np.ones((T, T), dtype=bool), k=1)   # True above the diagonal
scores = np.where(mask, -np.inf, scores)           # the future is forbidden

e = np.exp(scores - scores.max(axis=-1, keepdims=True))
weights = e / e.sum(axis=-1, keepdims=True)
print(mask.astype(int))
print(np.round(weights, 2))

[[0 1 1 1]
 [0 0 1 1]
 [0 0 0 1]
 [0 0 0 0]]
[[1.   0.   0.   0.  ]
 [0.78 0.22 0.   0.  ]
 [0.11 0.2  0.69 0.  ]
 [0.96 0.03 0.01 0.  ]]

The same seed gives the same embeddings and projections as before, so you can compare. The first token can see only itself (weight 1). The third token’s weights used to be 0.09, 0.17, 0.59, 0.15; with the future masked out, the remaining three are renormalised to 0.11, 0.20, 0.69. The last token sees everything, so its row is unchanged. The upper triangle being forbidden is what lets one forward pass over T tokens produce T training examples at once: position 1 predicts token 2, position 2 predicts token 3, and so on, each without cheating. (Encoder models such as BERT omit the mask and see both directions; that suits filling in blanks but not generation.)

One attention pattern per layer is limiting: a token may need to look at its verb, the previous word and the start of the sentence all at once. Multi-head attention runs several attentions side by side, each with its own smaller projections, and concatenates the results:

In code, nobody loops over heads. You compute all queries, keys and values with one big matrix, then reshape so the head becomes a batch dimension:

import numpy as np

np.random.seed(1)
T, d_model, n_heads = 4, 8, 2
d_head = d_model // n_heads
X = np.random.randn(T, d_model)
W_qkv = np.random.randn(d_model, 3 * d_model) / np.sqrt(d_model)  # all heads at once
W_o = np.random.randn(d_model, d_model) / np.sqrt(d_model)

def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

Q, K, V = np.split(X @ W_qkv, 3, axis=-1)               # each (4, 8)
split = lambda M: M.reshape(T, n_heads, d_head).transpose(1, 0, 2)
Q, K, V = split(Q), split(K), split(V)                  # each (2, 4, 4)
weights = softmax(Q @ K.transpose(0, 2, 1) / np.sqrt(d_head))   # (2, 4, 4)
heads = weights @ V                                     # (2, 4, 4)
concat = heads.transpose(1, 0, 2).reshape(T, d_model)   # (4, 8)
out = concat @ W_o

print("per-head Q:", Q.shape, " weights:", weights.shape, " out:", out.shape)
print("head 0 weights for last token:", np.round(weights[0, -1], 2))
print("head 1 weights for last token:", np.round(weights[1, -1], 2))

per-head Q: (2, 4, 4)  weights: (2, 4, 4)  out: (4, 8)
head 0 weights for last token: [0.27 0.18 0.26 0.28]
head 1 weights for last token: [0.21 0.23 0.36 0.2 ]

The reshape-and-transpose turns (4 tokens, 8 numbers) into (2 heads, 4 tokens, 4 numbers), so the matrix products run for both heads at once. The two heads already disagree about where the last token should look, and after training they typically specialise further. The original Transformer used 8 heads of 64 dimensions in a 512-wide model; because each head is narrower, multi-head attention costs about the same as one full-width head. Interpretability research has found heads that track the previous token, copy repeated patterns and more, though many heads resist a tidy description.

Where am I? Positional information

Look again at the attention formula: nothing in it knows about order. Shuffle the rows of X and the outputs are shuffled the same way, but otherwise identical. “Ali makan nasi” and “nasi makan Ali” would look alike to an attention layer. The causal mask supplies some order implicitly, but models need position explicitly, which is the job of a positional encoding. Three approaches matter.

import numpy as np

def sinusoidal(n_pos, d):
    pos = np.arange(n_pos)[:, None]                 # (n_pos, 1)
    i = np.arange(0, d, 2)[None, :]                 # even dimensions
    angle = pos / (10000 ** (i / d))
    pe = np.zeros((n_pos, d))
    pe[:, 0::2], pe[:, 1::2] = np.sin(angle), np.cos(angle)
    return pe

pe = sinusoidal(50, 16)
print(np.round(pe[:3, :6], 2))
for gap in [0, 1, 2, 5, 20]:
    print(f"pos 10 vs pos {10 + gap:<3} dot = {pe[10] @ pe[10 + gap]:.2f}")
print(f"pos 30 vs pos 35  dot = {pe[30] @ pe[35]:.2f}   (same gap of 5)")

[[ 0.    1.    0.    1.    0.    1.  ]
 [ 0.84  0.54  0.31  0.95  0.1   1.  ]
 [ 0.91 -0.42  0.59  0.81  0.2   0.98]]
pos 10 vs pos 10  dot = 8.00
pos 10 vs pos 11  dot = 7.49
pos 10 vs pos 12  dot = 6.37
pos 10 vs pos 15  dot = 6.14
pos 10 vs pos 30  dot = 5.78
pos 30 vs pos 35  dot = 6.14   (same gap of 5)

The first row (position 0) alternates 0 and 1 because sin 0 = 0 and cos 0 = 1. The dot product between two position vectors depends only on the gap between them, because sin a sin b + cos a cos b = cos(a − b): positions 10 and 15 score 6.14, exactly like 30 and 35. That relative-position structure is what the model can exploit, and RoPE builds the same idea directly into the attention scores. The similarity does not fall smoothly with distance (6.14 at a gap of 5, 6.37 at a gap of 2), which is one reason these vectors are a hint for the model to use rather than a ruler.

The block and the residual stream

Attention moves information between tokens. The other half of each layer, a small two-layer neural network applied to each token separately (the MLP or feed-forward layer; see Neural Networks from Scratch), processes information within each token. Two more ingredients make deep stacks of these trainable. A residual connection adds a sub-layer’s output back to its input instead of replacing it (an idea from ResNets in vision; see Convolutional Networks). Layer normalisation rescales each token’s vector to mean 0 and variance 1 before a sub-layer reads it:

GPT-2 and most models since put the norm at the input of each sub-layer (“pre-norm”), giving the block two lines:

Follow one block, step by step:

Because each block only adds to x, you can view the whole network differently. Elhage and colleagues at Anthropic, in “A Mathematical Framework for Transformer Circuits” (2021), called x the residual stream: a shared channel that runs from the embedding to the output, which every attention layer and MLP reads from (through a layer norm and a projection) and writes to (by addition). Unrolled, the final vector is the embedding plus the sum of every sub-layer’s contribution. This guide’s framing Think of the stream less as a pipeline and more as a shared whiteboard: early layers write rough notes (“this token is a noun”, “the previous token was nasi”), later layers read those notes and add their own, and the unembedding reads the final board. This view explains why residual connections make deep networks trainable — the gradient has a direct additive path back to every layer — and it is the starting point of much interpretability research, which is covered in the Expert track.

Counting parameters

Knowing where the parameters live tells you where the memory, compute and capacity go. For a GPT-2-style model the count follows directly from the shapes: a token table (vocabulary × d), a position table (context × d), and per block the attention projections (4d² plus biases), the MLP (two matrices of d × 4d, so 8d² plus biases) and two layer norms. GPT-2 reuses the token table as its output layer (“weight tying”), so the unembedding costs nothing extra.

def gpt_params(vocab, n_ctx, d, n_layers, d_ff=None):
    d_ff = d_ff or 4 * d
    embed = vocab * d + n_ctx * d                 # token + learned position tables
    attn = d * 3 * d + 3 * d + d * d + d          # W_qkv, biases, W_o, bias
    mlp = d * d_ff + d_ff + d_ff * d + d          # two linear layers with biases
    norms = 2 * 2 * d                             # two layer norms (scale + shift)
    block = attn + mlp + norms
    total = embed + n_layers * block + 2 * d      # + final layer norm
    return embed, block, total                    # output layer shares the token table

for name, cfg in {"toy": (65, 64, 64, 2), "GPT-2 small": (50257, 1024, 768, 12)}.items():
    embed, block, total = gpt_params(*cfg)
    print(f"{name:<12} embeddings {embed:>11,}  per block {block:>10,}  total {total:>12,}")

toy          embeddings       8,256  per block     49,984  total      108,352
GPT-2 small  embeddings  39,383,808  per block  7,087,872  total  124,439,808

124,439,808 is exactly the figure usually quoted for the smallest GPT-2, the one Karpathy’s nanoGPT reproduces. A handy rule falls out: each block has about 12d² parameters (4d² attention + 8d² MLP); for d = 768 that is 7,077,888, within 0.2% of the exact 7,087,872. Notice too that almost a third of GPT-2 small is its embedding tables, because the vocabulary of 50,257 tokens is large relative to the model. In big models the blocks dominate and the embeddings become a rounding error — one reason vocabulary size is a design choice that depends on scale (see Tokenisation and Embeddings).

Parameters are not the whole memory story. During generation, models cache the keys and values of past tokens (the “KV cache”) so they need not be recomputed, and that cache grows with sequence length, layers and width. It is why long contexts are costly to serve, even though the parameter count does not change.

A full forward pass

Time to put it together: a one-block, two-head, character-level Transformer, with random weights, run on the five characters “nasi ” to predict the sixth. Every line maps to a piece above. Read it slowly with the printed shapes beside it.

import numpy as np

rng = np.random.default_rng(42)
text = "nasi lemak"
vocab = sorted(set(text))                       # 9 characters incl. space
stoi = {c: i for i, c in enumerate(vocab)}
V, T, d, H, d_ff = len(vocab), 5, 16, 2, 64
P = lambda *shape: rng.standard_normal(shape) * 0.1   # small random weights

tok_emb, pos_emb = P(V, d), P(T, d)
W_qkv, W_o, W_1, W_2 = P(d, 3 * d), P(d, d), P(d, d_ff), P(d_ff, d)

def layer_norm(x, eps=1e-5):
    mu, var = x.mean(-1, keepdims=True), x.var(-1, keepdims=True)
    return (x - mu) / np.sqrt(var + eps)        # scale/shift omitted (1 and 0)

def softmax(z):
    e = np.exp(z - z.max(-1, keepdims=True))
    return e / e.sum(-1, keepdims=True)

ids = np.array([stoi[c] for c in "nasi "])      # 5 input tokens
x = tok_emb[ids] + pos_emb[:T];                 print("embed     ", x.shape)
h = layer_norm(x) @ W_qkv;                      print("qkv       ", h.shape)
q, k, v = (m.reshape(T, H, d // H).transpose(1, 0, 2) for m in np.split(h, 3, -1))
s = q @ k.transpose(0, 2, 1) / np.sqrt(d // H); print("scores    ", s.shape)
s[:, np.triu(np.ones((T, T), bool), 1)] = -np.inf
a = (softmax(s) @ v).transpose(1, 0, 2).reshape(T, d)
x = x + a @ W_o;                                print("after attn", x.shape)
x = x + np.maximum(0, layer_norm(x) @ W_1) @ W_2; print("after mlp ", x.shape)
logits = layer_norm(x) @ tok_emb.T;             print("logits    ", logits.shape)
p = softmax(logits[-1])
print("next-char probs after 'nasi ':", {c: round(float(pi), 3) for c, pi in zip(vocab, p)})

embed      (5, 16)
qkv        (5, 48)
scores     (2, 5, 5)
after attn (5, 16)
after mlp  (5, 16)
logits     (5, 9)
next-char probs after 'nasi ': {' ': 0.211, 'a': 0.115, 'e': 0.112, 'i': 0.123, 'k': 0.062, 'l': 0.151, 'm': 0.076, 'n': 0.09, 's': 0.06}

What each step does:

• Embed: tok_emb[ids] picks one row per character (an embedding lookup is just row indexing), and the learned position rows are added. Shape (5, 16): five tokens, sixteen numbers each.
• QKV: one matrix multiply produces queries, keys and values for both heads: 3 × 16 = 48 columns. The reshape splits them into two heads of 8.
• Scores: (2, 5, 5) — for each head, a 5 × 5 table of query–key matches, upper triangle masked to −∞.
• After attention: heads are concatenated, projected by W_o and added to x. Shape unchanged.
• After MLP: widen to 64 with ReLU (GPT-2 uses the smoother GELU), project back to 16, add. Shape unchanged — the stream always keeps its width.
• Logits: a final layer norm, then multiplication by the transposed token table (weight tying again) gives a score for each of the 9 characters at each of the 5 positions. We only need the last row to predict what follows “nasi ”.

With random weights the prediction is close to uniform (1 ÷ 9 ≈ 0.111), and the small preference for a space is noise. That is the honest starting point of every language model. Pretraining is the process that turns these weights into ones where, after “nasi ”, the letter l becomes very likely. A real GPT differs from this toy in scale (dozens of blocks, thousands of dimensions), in batching (an extra leading dimension for many sequences at once), in learned layer-norm parameters and biases, and in engineering (fused kernels, mixed precision), but not in kind. Karpathy’s nanoGPT implements the full thing in about 300 lines of PyTorch for the model, and is the natural next read.

Key terms on this page

Sources & further reading

• Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser & Illia Polosukhin, “Attention Is All You Need”, NeurIPS (2017).
• Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei & Ilya Sutskever, “Language Models are Unsupervised Multitask Learners” (GPT-2, 2019) — pre-norm placement and the 124M configuration.
• Nelson Elhage, Neel Nanda, Catherine Olsson et al., “A Mathematical Framework for Transformer Circuits”, Transformer Circuits Thread (2021) — the residual stream view.
• Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen & Yunfeng Liu, “RoFormer: Enhanced Transformer with Rotary Position Embedding” (2021).
• Jimmy Lei Ba, Jamie Ryan Kiros & Geoffrey E. Hinton, “Layer Normalization” (2016); Kaiming He, Xiangyu Zhang, Shaoqing Ren & Jian Sun, “Deep Residual Learning for Image Recognition” (2015, CVPR 2016).
• Andrej Karpathy, nanoGPT (code, 2022 onwards) and the video “Let’s build GPT: from scratch, in code, spelled out” (2023).
• Jay Alammar, “The Illustrated Transformer” (2018) — a visual companion.