L4 · 02 · Transformer Internals: Attention to a Tiny GPT

The transformer is the architecture behind every modern LLM. It's a stack of blocks, and each block does only two things: attention (tokens exchange information) and a feed-forward network (each token is processed alone). This page implements attention from scratch in plain Python, walks through a full block, and sketches a tiny GPT in PyTorch.

The whole architecture in one diagram

A decoder-only transformer (GPT-style): tokens in → vectors → N identical blocks → a probability over the vocabulary for the next token. Repeat to generate.

Attention: the one idea that matters

Each token asks: "which previous tokens are relevant to me?" It computes this with three learned projections per token:

• Query (Q) — "what am I looking for"
• Key (K) — "what do I contain" (an advertisement)
• Value (V) — "what do I actually pass along if you attend to me"

Score = Q · K for every pair → scale → softmax → weights → weighted sum of the values. That's it. Here it is in 15 lines:

import math


def dot(a, b):
    return sum(x * y for x, y in zip(a, b))


def softmax(xs):
    m = max(xs)
    e = [math.exp(x - m) for x in xs]
    s = sum(e)
    return [v / s for v in e]


def attention(Q, K, V, mask=None):
    """Q,K,V: lists of rows (seq_len x d). Returns seq_len x d."""
    d = len(Q[0])
    scores = [[dot(q, k) / math.sqrt(d) for k in K] for q in Q]
    if mask:                                    # causal: token i can't see j > i
        scores = [[s if not mask[i][j] else -1e9
                   for j, s in enumerate(row)] for i, row in enumerate(scores)]
    weights = [softmax(row) for row in scores]
    return [[sum(w * v[j] for w, v in zip(ws, V))
             for j in range(len(V[0]))] for ws in weights], weights


if __name__ == "__main__":
    # 3 tokens, 2-dim vectors (toy). Each token attends to all previous + itself.
    X = [[1, 0], [0, 1], [1, 1]]
    mask = [[j > i for j in range(3)] for i in range(3)]
    out, w = attention(X, X, X, mask)
    print("attention weights row 0:", [round(x, 2) for x in w[0]])   # [1.0, 0, 0]
    print("weights row 2:", [round(x, 2) for x in w[2]])             # splits across all
    # Token 0 can only see itself; token 2 sees all three.
    assert w[0][1] < 1e-6 and w[0][2] < 1e-6

Run it: token 0 puts 100% weight on itself (the mask hides the future); the last token spreads weight across all three. In a real model, Q, K, V aren't the input itself but learned linear projections of it.

Multi-head attention

One attention "head" learns one kind of relationship. Multiple heads run attention in parallel with separate projections, each on a smaller slice (d_model / n_heads), then concatenate. Head 1 might track subject–verb agreement, head 4 coreference, head 7 code indentation. Same cost, more kinds of pattern.

The rest of a block

A block, end to end

x = x + attention(norm(x))      # mix information across tokens
x = x + feedforward(norm(x))    # process each token independently
# after N blocks:
logits = unembed(norm(x))       # project to vocab size

That's a GPT block. Stack it 12–100+ times, train on next-token prediction, and you have a language model.

Training a tiny GPT (PyTorch sketch)

The canonical ~300-line implementation is nanoGPT. The core training loop is short:

import torch

model = TinyGPT(vocab_size, d_model=128, n_heads=4, n_layers=4)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)

for step in range(5000):
    xb, yb = get_batch("train")            # xb: tokens, yb: same shifted by 1
    logits = model(xb)                     # (batch, seq, vocab)
    loss = cross_entropy(logits.view(-1, V), yb.view(-1))
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 500 == 0: print(step, loss.item())

Everything on this page is in that loop: the model is the block stack, the loss is cross-entropy over every position, backward() is backprop, and the optimizer's step is gradient descent.

Why this matters even if you never train one

• Attention = context use. Long prompts cost more because attention is O(n²) over tokens — that's why context windows cost money.
• KV cache (next page): generation reuses the K/V of past tokens. Understanding it explains inference cost.
• Heads explain behaviour. "The model ignored the middle of the document" is an attention pattern, not magic.
• It's not a black box. Once you've written attention by hand, LLM behaviour stops being mysterious.