Transformers: The Architecture Behind Modern AI

Attention is all you need

In June 2017 eight researchers working at Google published "Attention Is All You Need". Their model, the transformer, threw recurrence away completely and built a translator from attention and simple feed-forward layers. It beat the best translation systems of the time while training in a fraction of the time, because every position in a sentence could be processed in parallel.

This lesson takes the architecture apart, piece by piece.

Giving attention a sense of order

Attention compares every token with every other token, but nothing in the calculation knows where a token is. Shuffle the words and every token gets exactly the same weights, so "dog bites man" and "man bites dog" would look identical.

The fix is to add a positional encoding to each token's embedding before the first layer. The original transformer used waves: position p gets a vector whose numbers are sines and cosines of p at different frequencies.

The first numbers oscillate quickly and the later ones slowly, like the second, minute and hour hands of a clock. Every position gets its own pattern, and nearby positions get similar patterns.

GPT-2 took a simpler route and learned a position vector for each of its 1,024 positions. Many recent models use rotary position embeddings (RoPE), which rotate each query and key by an angle that depends on its position, so the score between two tokens depends on how far apart they are.

Inside a transformer block

A transformer is a stack of identical blocks. Each block has two sublayers, and each sublayer is wrapped in a residual connection and a normalisation step:

1. Multi-head self-attention: the tokens exchange information.
2. Add and normalise: add the sublayer's input back onto its output (the residual connection from Level 8's ResNets), then apply layer normalisation, which rescales each token's vector to a mean of 0 and a standard deviation of 1.
3. Feed-forward network: the same small two-layer network, applied to every token separately. It widens each vector to 4 × d_model numbers, applies an activation such as ReLU or GELU, and narrows it back.
4. Add and normalise again.

The output has exactly the same shape as the input: one vector of d_model numbers for each token. That is what lets blocks stack. GPT-2 small stacks 12 of them, and GPT-3 stacks 96.

The residual stream

A helpful way to picture a transformer is as a residual stream. Each token has a lane, a vector of d_model numbers flowing up through the network, and every sublayer reads from the stream and adds its result back in.

• Attention is the only place where information moves between lanes, from one token to another.
• The feed-forward layers work on each lane on its own. Research suggests that much of a model's factual knowledge is stored in their weights.

By the top of the stack, the vector in the last token's lane holds everything needed to predict what comes next.

From vectors to the next word

After the last block, each token's vector is multiplied by an unembedding matrix with one column for every token in the vocabulary. That gives a score, called a logit, for every possible next token, and softmax turns the logits into probabilities, exactly as in Level 9. Many models reuse the embedding matrix here (weight tying), so the same numbers turn tokens into vectors at the bottom and vectors back into tokens at the top.

Here is where every one of GPT-2 small's parameters lives (d_model = 768, 12 blocks, 50,257 tokens, 1,024 positions):

About two thirds of the parameters are in the blocks, and two thirds of each block is its feed-forward network. Attention moves information around; the feed-forward layers do most of the storing.

Three families of transformer

The same block is used in three different arrangements:

Decoder-only models turned out to be the most general, because almost any task can be written as text in, text out: translation is "Translate into French: …", and classification is "Is this review positive or negative? …".

Why transformers won

• Parallel training: a whole sequence is processed in one pass of matrix multiplications, which GPUs are built for. An RNN has to go one step at a time.
• Short paths: any token reaches any other token in one attention step, so long-range connections are easy to learn.
• They scale: performance kept improving as models, data and computing power grew, as Level 11 explores.
• One architecture for everything: text, images, audio, protein sequences and robot actions can all be turned into sequences of tokens.

The Vision Transformer shows how far the idea stretches. It cuts an image into square patches and treats each patch as a token:

The weakness is cost. Every token attends to every token, so the attention matrix grows with the square of the number of tokens:

Hands-on: a transformer from scratch

Positional encodings

import math

def positional_encoding(pos, d_model):
    pe = []
    for i in range(d_model // 2):
        angle = pos / 10000 ** (2 * i / d_model)
        pe += [math.sin(angle), math.cos(angle)]
    return pe

print("pos  " + " ".join(f"{'d' + str(i):>5}" for i in range(8)))
for pos in range(6):
    print(f"{pos:3}  " + " ".join(f"{v:5.2f}" for v in positional_encoding(pos, 8)))

print()
here = positional_encoding(10, 16)
for gap in [0, 1, 2, 4, 8, 16]:
    there = positional_encoding(10 + gap, 16)
    sim = sum(a * b for a, b in zip(here, there))
    print(f"position 10 against position {10 + gap:2}: dot product {sim:5.2f}")

Column d0 changes fast and d6 barely moves, like clock hands. The dot product falls as positions get further apart, which is how attention can tell near from far.

A stack of transformer blocks

Each block below is the real recipe: layer norm, masked self-attention, add; layer norm, feed-forward, add. Each block has its own random weights.

import math, random

random.seed(7)
def rand_matrix(rows, cols, scale=0.5):
    return [[random.uniform(-scale, scale) for _ in range(cols)] for _ in range(rows)]
def matmul(A, B):
    return [[sum(a * b for a, b in zip(row, col)) for col in zip(*B)] for row in A]
def add(A, B):
    return [[a + b for a, b in zip(ra, rb)] for ra, rb in zip(A, B)]
def softmax(xs):
    m = max(xs)
    e = [math.exp(x - m) for x in xs]
    return [v / sum(e) for v in e]
def layer_norm(X):
    out = []
    for row in X:
        mean = sum(row) / len(row)
        sd = math.sqrt(sum((v - mean) ** 2 for v in row) / len(row) + 1e-5)
        out.append([(v - mean) / sd for v in row])
    return out
def attention(X, W_Q, W_K, W_V, W_O):
    Q, K, V = matmul(X, W_Q), matmul(X, W_K), matmul(X, W_V)
    d = len(Q[0])
    A = [softmax([sum(q * k for q, k in zip(Q[i], K[j])) / math.sqrt(d) if j <= i else -math.inf
                  for j in range(len(X))]) for i in range(len(X))]
    return matmul(matmul(A, V), W_O)
def feed_forward(X, W1, W2):
    hidden = [[max(0.0, v) for v in row] for row in matmul(X, W1)]   # ReLU
    return matmul(hidden, W2)

d_model, d_ff = 4, 16
def make_block():
    return {"attention": [rand_matrix(d_model, d_model) for _ in range(4)],
            "feed_forward": [rand_matrix(d_model, d_ff), rand_matrix(d_ff, d_model)]}

def block(X, p):
    X = add(X, attention(layer_norm(X), *p["attention"]))        # tokens share information
    X = add(X, feed_forward(layer_norm(X), *p["feed_forward"]))  # each token on its own
    return X

X = rand_matrix(3, d_model, 1.0)        # 3 tokens, 4 numbers each
for n, p in enumerate([make_block() for _ in range(3)], start=1):
    X = block(X, p)
    print(f"after block {n}: {len(X)} tokens x {len(X[0])} numbers, first token",
          [round(v, 2) for v in X[0]])

The shape never changes: 3 tokens of 4 numbers go in, and 3 tokens of 4 numbers come out, which is why the blocks can be stacked as deep as you like.

Counting GPT-2's parameters

def gpt2_parameters(n_layers, d, vocab=50257, context=1024):
    embeddings = vocab * d + context * d
    attention = 4 * d * d + 4 * d          # W_Q, W_K, W_V, W_O and their biases
    feed_forward = 8 * d * d + 5 * d       # d -> 4d -> d, with biases
    norms = 4 * d                          # two layer norms, each a scale and a shift
    return embeddings + n_layers * (attention + feed_forward + norms) + 2 * d

for name, layers, d in [("small", 12, 768), ("medium", 24, 1024),
                        ("large", 36, 1280), ("XL", 48, 1600)]:
    print(f"GPT-2 {name:6} {layers:2} blocks, d = {d:4}: {gpt2_parameters(layers, d):>13,} parameters")

These are exactly the sizes of the four GPT-2 models released in 2019.

Try this

1. Remove the two add(X, ...) calls in the block, keeping only the sublayer outputs. Stack 12 blocks instead of 3. What happens to the numbers?
2. GPT-3 has 96 blocks with d = 12,288 and the same 50,257-token vocabulary, with 2,048 positions. Use the function to estimate its size. How close is it to 175 billion?

Level check

• Explain why attention needs positional encoding, and describe two ways of adding it
• Draw and label a transformer block, including the residual connections
• Explain what layer normalisation and the feed-forward network do
• Describe the residual stream and which sublayer moves information between tokens
• Count the parameters of a transformer from its sizes
• Compare encoder-only, encoder-decoder and decoder-only transformers
• Explain why transformers replaced RNNs, and the cost of long contexts

Key terms you must be able to define

Transformer

A neural network built from stacked blocks of attention and feed-forward layers, introduced in 2017.

Positional encoding

A vector added to each token's embedding that tells the model where the token is.

Transformer block

Multi-head attention and a feed-forward network, each with a residual connection and layer normalisation.

Feed-forward network

Two linear layers with an activation between them, applied to each token separately.

Layer normalisation

Rescaling each token's vector to mean 0 and standard deviation 1, then applying a learned scale and shift.

Residual stream

The per-token vectors flowing up through the network, which every sublayer reads from and adds to.

Logit

The raw score for one vocabulary token, before softmax turns the scores into probabilities.

Weight tying

Reusing the embedding matrix as the unembedding matrix that produces the logits.

Encoder-only

A transformer in which every token sees every other, such as BERT, used for understanding tasks.

Decoder-only

A transformer with masked attention trained to predict the next token, such as GPT.

Vision Transformer (ViT)

A transformer that classifies images by treating patches of pixels as tokens.

Flashcards

Go further: notes for future AI engineers

Pre-norm and post-norm

The 2017 transformer normalised after each residual addition (post-norm). Deep post-norm stacks are hard to train without carefully warming up the learning rate, so most modern models normalise the input of each sublayer instead (pre-norm), as the code above does. That leaves the residual stream itself untouched: a clean path for gradients to travel down, exactly as in a ResNet. Many models also swap layer norm for the simpler RMSNorm, and ReLU for GELU or SwiGLU.

The KV cache

When generating, a decoder-only model would naively recompute the keys and values of every earlier token at every step. Because of the mask, those never change, so they are stored in a KV cache and only the newest token's query, key and value are computed. Each new token becomes cheap, but the cache grows with the context and is a large part of the memory needed to serve a model to many users.

Mixture of experts

Some large models replace each feed-forward network with many smaller expert networks and a router that sends each token to just one or two of them. The model can hold far more parameters than it uses on any single token, gaining knowledge without a matching rise in the computation per token.

Misconceptions to avoid