Transformer anatomy

Anatomy of a transformer

Select any part to see what it does, its tensor shapes, and how many parameters it holds.

Residual stream, 8 x 768

Block 1 of 12

11 more blocks with the same shape and their own weights

Multi-head self-attention

The only place where tokens exchange information. Each of 12 heads projects every token to a query, key, and value of size 64, scores each query against the keys of the tokens at or before it (the causal mask), and takes a softmax-weighted sum of their values. The heads' outputs are joined and projected back to 768 numbers by W_O, then added to the stream.

Shapes for Multi-head self-attention
W_Q	768 x 768
W_K	768 x 768
W_V	768 x 768
W_O	768 x 768
Scores, per head	8 x 8
Out	8 x 768

2,362,368 parameters per block, 28,348,416 across all 12, 23% of the model.

Where GPT-2 small's 124,439,808 parameters live

Token and position embeddings39.4M, 32%
Attention, all blocks28.3M, 23%
MLP, all blocks56.7M, 46%
Norms38.4K, under 0.1%

Parameters124M

Without embeddings85.1M

Forward FLOPs per token at T = 80.25 G

Challenges

0 of 3 complete

Pick a model and a sequence length where one head's T x T score matrix holds more numbers than all of a block's attention weights. (not yet)
Find a model whose MLPs hold more than three quarters of all its parameters. (not yet)
With attention selected, find a model whose key projection W_K is smaller than its query projection W_Q. (not yet)

Try this

Select every part of GPT-2 small and add up the parameters yourself. The total matches the released checkpoint exactly.
Raise the sequence length and watch which shapes change. Weights never do.
Compare where the parameters live in GPT-2 small and in Llama 3 70B.

Count GPT-2 small's parameters from its shapesPython with NumPy

import numpy as np

# GPT-2 small: vocabulary, context length, width, blocks.
V, ctx, d, L = 50257, 1024, 768, 12

# Every trained tensor in one block, listed by shape.
block = {
    "norms": [(d,), (d,), (d,), (d,)],         # two LayerNorms: scale and shift
    "attention": [(d, d)] * 4 + [(d,)] * 4,    # W_Q, W_K, W_V, W_O and biases
    "mlp": [(d, 4 * d), (4 * d,), (4 * d, d), (d,)],  # up, down, and biases
}

def count(shapes):
    return sum(int(np.prod(s)) for s in shapes)

per_block = sum(count(s) for s in block.values())
embeddings = count([(V, d), (ctx, d)])    # token and position tables
final_norm = count([(d,), (d,)])          # the unembedding reuses the token table

total = embeddings + L * per_block + final_norm
print(f"one block: {per_block:,}")
print(f"total: {total:,}")
print(f"MLP share of a block: {count(block['mlp']) / per_block:.0%}")
print(f"12 L d^2 shortcut: {12 * L * d * d:,}")

one block: 7,087,872
total: 124,439,808
MLP share of a block: 67%
12 L d^2 shortcut: 84,934,656

Parameter counts are computed from each model’s published configuration and checked against the released checkpoints in this site’s tests.