3 — Attention, slowly

Chapter 2 ended with every token as a 768-number vector built from two lookups: what the token is and where it sits. Each vector was computed on its own. The vector for " it" doesn’t know what “it” refers to, and the vector for " bank" doesn’t know whether the sentence is about a river.

Attention is the layer that fixes this. For each position it builds a new vector as a weighted average of the vectors at that position and the positions before it, with the weights chosen by how relevant each one is. Relevant tokens contribute a lot; irrelevant ones contribute almost nothing. That one idea, plus a few pieces that make it trainable and stable, is the core of every GPT.

This chapter builds it on three words with 4-number vectors, small enough that every number can be checked by hand, and then grows it to GPT-2’s size. All runs were on an Apple M3 Pro with 18 GB of memory, with Python 3.12 and PyTorch 2.15.0 on the CPU. Seeds are given with each run.

Three words, by hand

The sentence is “the cat sat”. Instead of random vectors, these are picked by hand so the arithmetic stays in whole numbers:

import torch

x = torch.tensor([[1., 0., 0., 1.],   # the
                  [0., 1., 1., 0.],   # cat
                  [0., 1., 0., 1.]])  # sat

Each row is one token. The 1. with a dot makes the numbers floating point, which the later steps need.

Similarity by dot product

The first question is how relevant each token is to each other token. The simplest measure is the dot product: multiply two vectors element by element and add up the results.

It is large when the two vectors have large values in the same places, zero when they share nothing, and negative when they point in opposite directions. For “sat” and “the”:

and for “sat” with itself, 0 + 1 + 0 + 1 = 2. Doing this for every pair at once is a matrix multiplication of x with its own transpose (rows turned into columns):

scores = x @ x.T
print(scores)

tensor([[2., 0., 1.],
        [0., 2., 1.],
        [1., 1., 2.]])

Row i, column j is the dot product of token i with token j. Row 3 says that “sat” overlaps with “the” by 1, with “cat” by 1, and with itself by 2. These raw numbers are called attention scores.

Softmax turns scores into weights

Scores can be any size and any sign. A weighted average needs weights that are positive and add up to 1. Softmax does that conversion: raise e to the power of each score, then divide by the total.

The exponential makes everything positive and stretches the gaps, so a higher score gets disproportionately more weight. Dividing by the sum makes each row add up to 1. For the “sat” row, scores [1, 1, 2]:

PyTorch does every row at once:

weights = torch.softmax(scores, dim=-1)
print(weights)

tensor([[0.665, 0.090, 0.245],
        [0.090, 0.665, 0.245],
        [0.212, 0.212, 0.576]])

dim=-1 means “along the last dimension”, so each row is normalized separately. The last row matches the hand calculation. (Numbers in this chapter are printed to three decimals.)

A weighted sum of the vectors

The output for each token is the weighted average of all the token vectors, using that token’s row of weights. For “sat”:

Again, a matrix multiply does every row:

context = weights @ x
print(context)

tensor([[0.665, 0.335, 0.090, 0.910],
        [0.090, 0.910, 0.665, 0.335],
        [0.212, 0.788, 0.212, 0.788]])

These output vectors are often called context vectors. The new vector for “sat” is mostly “sat” but carries a fifth of “the” and a fifth of “cat”. Each output row now depends on the other tokens, which is exactly what chapter 2’s vectors lacked.

That is attention in its plainest form: scores from dot products, softmax, weighted sum. It has no parameters at all, which means it can’t learn anything. It also has two specific problems, and the rest of the chapter is about fixing them and a few others.

Queries, keys and values

The first problem is that one vector is doing three jobs. In x @ x.T the same vector for “sat” is used for what “sat” is looking for, what “sat” offers to tokens looking at it, and what “sat” hands over when it’s picked. Those don’t have to be the same. A verb looking for its subject is searching for something quite unlike what the verb itself would advertise. And the scores come out symmetric: “sat” scores “the” exactly as “the” scores “sat”, and every token scores highest on itself.

The fix gives each token three different vectors, each made by multiplying the token’s vector by its own learned matrix:

• the query, what this token is looking for;
• the key, what this token can be found by;
• the value, what this token passes on when another token attends to it.

A library catalog is a fair analogy: the query is the search typed in, the keys are the index cards, and the values are the books handed over. Scores now compare one token’s query with every token’s key, and the weighted sum averages values instead of raw vectors.

By hand, with small matrices

In a real model these three matrices are learned. Here they are picked by hand so the numbers stay simple. Each maps 4 numbers down to 2:

W_q = torch.tensor([[1., 0.], [0., 1.], [0., 0.], [0., 0.]])   # keeps numbers 1 and 2
W_k = torch.tensor([[0., 0.], [0., 0.], [1., 0.], [0., 1.]])   # keeps numbers 3 and 4
W_v = torch.tensor([[1., 0.], [0., 0.], [0., 1.], [0., 0.]])   # keeps numbers 1 and 3

q, k, v = x @ W_q, x @ W_k, x @ W_v

q (queries)       k (keys)          v (values)
the  [1, 0]       the  [0, 1]       the  [1, 0]
cat  [0, 1]       cat  [1, 0]       cat  [0, 1]
sat  [0, 1]       sat  [0, 1]       sat  [0, 0]

The scores are now each query dotted with each key:

scores = q @ k.T

tensor([[0., 1., 0.],
        [1., 0., 1.],
        [1., 0., 1.]])

By hand, row 2 is the query of “cat”, [0, 1], against each key: with “the” [0, 1] it’s 1, with “cat” [1, 0] it’s 0, with “sat” [0, 1] it’s 1. The matrix is no longer symmetric (“sat” scores “the” at 1, but “the” scores “sat” at 0), and “cat” now scores itself lowest. What a token looks for and what it advertises have come apart.

With learned matrices this separation is what lets a head do something specific, like “a pronoun looks for the noun it refers to” or, as the end of this chapter shows in the real GPT-2, “look at the previous token”. The value matrix separately controls what gets copied once the head has decided where to look.

Linear layers

In PyTorch these matrices live in nn.Linear layers. nn.Linear(4, 2) holds a weight matrix and (optionally) a bias, and computes

Note the transpose: PyTorch stores the weight as (outputs × inputs), 2 × 4 here, the reverse of how W_q is written above. That detail comes back in chapter 6, because OpenAI’s GPT-2 files store their matrices the other way round. A linear layer created with bias=False has no b, which makes it exactly the plain matrix multiply used above.

Scaling by √d

Before the scores reach softmax, they are divided by the square root of the query and key length, written d (or d_k). The full formula, as written in “Attention Is All You Need”, is:

The reason shows up at real sizes. If the numbers in a query and a key are random with mean 0 and variance 1, their dot product is a sum of d such products, so its variance is d and its typical size (standard deviation) is √d. Measured over 100,000 random pairs (seed 123):

GPT-2’s heads use d = 64, so raw scores typically span ±8 or more. Softmax exaggerates gaps exponentially, and gaps that size make it nearly a hard pick. One query against six random keys, d = 64 (seed 1):

raw scores           [-9.799, -4.000, -0.607,  2.800, -3.593, -0.429]
softmax(raw)         [ 0.000,  0.001,  0.031,  0.930,  0.002,  0.037]
softmax(raw / 8)     [ 0.061,  0.126,  0.192,  0.294,  0.132,  0.196]

Unscaled, one key takes 93% of the weight and the rest get almost nothing, at the very start of training, before the model has learned anything about which tokens matter. Worse, a saturated softmax barely responds to its inputs, so training gets a weak signal for changing it. The gradient of one output weight with respect to the six scores (how much that weight would move if the scores were nudged) has a size of 0.041 for the raw scores and 0.173 for the scaled ones, about four times larger. The paper’s wording is that large dot products push softmax “into regions where it has extremely small gradients”, and dividing by √d_k counteracts this.

In the hand example, d = 2, so the scores are divided by √2 ≈ 1.414.

The causal mask

The second problem with the plain version: “the” gets 24.5% of its output from “sat”, a token that comes after it.

For a model that predicts the next token, that is cheating. Chapter 2 set up training so that one pass over a window makes a prediction at every position, each checked against the following token. If position t could see position t + 1, it could just copy the answer. Training loss would drop fast and the model would learn nothing useful, because when it generates text for real, the future token doesn’t exist yet.

The rule is simple: position t may attend to positions 0 to t and nothing later. In the score matrix that means everything above the diagonal is forbidden. Setting those scores to −∞ before softmax enforces it, since e^−∞ = 0:

mask = torch.triu(torch.ones(3, 3, dtype=torch.bool), diagonal=1)
print(mask)
scores = (q @ k.T) / 2**0.5
scores = scores.masked_fill(mask, float("-inf"))
print(scores)

tensor([[False,  True,  True],
        [False, False,  True],
        [False, False, False]])
tensor([[0.000,   -inf,   -inf],
        [0.707,  0.000,   -inf],
        [0.707,  0.000,  0.707]])

torch.ones(3, 3, dtype=torch.bool) is a 3 × 3 grid of True. torch.triu(..., diagonal=1) (“upper triangle”) keeps only the entries strictly above the diagonal and sets the rest to False, so True marks the forbidden pairs. masked_fill replaces every score where the mask is True with −∞.

Softmax then gives zero weight to the future, and the remaining weights in each row still add up to 1:

weights = torch.softmax(scores, dim=-1)
print(weights)
print(weights @ v)

tensor([[1.000, 0.000, 0.000],
        [0.670, 0.330, 0.000],
        [0.401, 0.198, 0.401]])
tensor([[1.000, 0.000],
        [0.670, 0.330],
        [0.401, 0.198]])

Row 2 by hand: the allowed scores are 0.707 and 0, so the weights are

“The” can only see itself, so it gets weight 1 on itself. The outputs are weighted sums of the value rows: for “cat”, 0.670 × [1, 0] + 0.330 × [0, 1] = [0.670, 0.330].

Masking before softmax matters. Zeroing weights after softmax would leave rows that no longer add to 1, and the future tokens would still have affected the sum in the denominator.

One head, as a module

Those steps fit in a small PyTorch module. A module is a class that holds parameters (here three nn.Linear layers) and defines forward, the computation; calling the module calls forward.

import math
import torch.nn as nn

def causal_mask(T, device=None):
    """True above the diagonal: the positions a token is NOT allowed to see."""
    return torch.triu(torch.ones(T, T, dtype=torch.bool, device=device), diagonal=1)

class OneHeadAttention(nn.Module):
    def __init__(self, d_in, d_head, qkv_bias=False):
        super().__init__()
        self.W_q = nn.Linear(d_in, d_head, bias=qkv_bias)
        self.W_k = nn.Linear(d_in, d_head, bias=qkv_bias)
        self.W_v = nn.Linear(d_in, d_head, bias=qkv_bias)

    def forward(self, x):                       # x: (B, T, d_in)
        q, k, v = self.W_q(x), self.W_k(x), self.W_v(x)
        scores = q @ k.transpose(-2, -1) / math.sqrt(k.shape[-1])   # (B, T, T)
        T = x.shape[1]
        scores = scores.masked_fill(causal_mask(T, x.device), float("-inf"))
        weights = torch.softmax(scores, dim=-1)
        return weights @ v                      # (B, T, d_head)

The input now has a batch dimension, (B, T, d_in), as chapter 2’s batches do. k.transpose(-2, -1) swaps the last two dimensions, the batched version of .T, so q @ k.transpose(-2, -1) produces a (T, T) score table for every item in the batch. Copying the hand-picked matrices into the layers (transposed, because of the storage order) and running “the cat sat” gives the same output as above:

m = OneHeadAttention(4, 2)
with torch.no_grad():
    m.W_q.weight.copy_(W_q.T); m.W_k.weight.copy_(W_k.T); m.W_v.weight.copy_(W_v.T)
print(m(x.unsqueeze(0)))

tensor([[[1.000, 0.000],
         [0.670, 0.330],
         [0.401, 0.198]]], grad_fn=<UnsafeViewBackward0>)

torch.no_grad() tells PyTorch not to record these copies for training, copy_ overwrites a tensor in place, and unsqueeze(0) adds a batch dimension of size 1.

Seeing the leak

The mask’s effect can be measured directly. Take three random token vectors, change only the last one, and see which outputs move. This uses the full attention module built later in this chapter, with one head and 4-number vectors (seed 123), run once with its mask and once with the mask switched off:

attn = CausalSelfAttention(emb_dim=4, n_heads=1, context_length=8).eval()
x = torch.randn(1, 3, 4)
x2 = x.clone(); x2[0, 2] = torch.randn(4)     # change only the last token
print((attn(x) - attn(x2)).abs().sum(-1))     # how much each position's output moved

attn.mask.zero_()                             # mask all False: nothing hidden
print((attn(x) - attn(x2)).abs().sum(-1))

with the mask:     tensor([[0.000, 0.000, 0.299]])
without the mask:  tensor([[0.334, 0.171, 0.299]])

torch.randn draws numbers from a standard normal distribution, and clone makes an independent copy. With the mask, changing the third token moves only the third output; positions 1 and 2 are untouched, so they can’t have used it. Without the mask, every position changes, so every position was reading the future.

The same unmasked module also confirms the claim from chapter 2 about order. Feeding it the three rows in reverse order returns exactly the same three outputs in reverse order (torch.allclose reports True). Without position vectors and without a mask, attention can’t tell “dog bites man” from “man bites dog”. The mask does make position matter a little, since each token sees a different number of predecessors, but that is a weak signal. GPT-2 relies on its learned position vectors.

Dropout

Dropout is a guard against overfitting, used only during training. On each training step it sets a random fraction p of values to zero and scales the survivors up by 1 / (1 − p), so the expected total stays the same. Applied to attention weights, it means a head can’t come to depend on one particular connection always being there.

On a row of eight equal weights of 0.125, with p = 0.5 (seed 123):

drop = torch.nn.Dropout(0.5)
row = torch.full((8,), 0.125)
print(drop(row))
drop.eval()
print(drop(row))

tensor([0.250, 0.250, 0.000, 0.250, 0.250, 0.000, 0.000, 0.000])
tensor([0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125])

torch.full((8,), 0.125) makes eight copies of 0.125. In training mode, four weights were zeroed and the other four doubled to 0.25. .eval() switches the module to evaluation mode, where dropout does nothing; .train() switches it back. Every module has these two modes, and forgetting .eval() before generating text is a classic way to get noisy, inconsistent output.

Rates around 0.1 are typical (the Hugging Face configuration for GPT-2 lists 0.1 for attention dropout). OpenAI’s released GPT-2 code is for inference only and contains no dropout at all. When fine-tuning on small datasets in later chapters, dropout helps; when just running the model, it’s switched off.

Multiple heads

One head produces one set of weights per position: one opinion about where to look. Language needs several at once. To predict the token after “The Time Traveller smiled at”, it helps to know the previous word, the subject of the sentence, whether a quote is open, and more. One softmax row can’t put most of its weight in four places.

Multi-head attention runs several heads side by side, each with its own query, key and value matrices, and joins their outputs. To keep the cost the same, each head works in a smaller space. GPT-2 has 768-number vectors and 12 heads, so each head’s queries, keys and values are 768 / 12 = 64 numbers long. The original transformer did the same with 512 and 8 heads of 64.

The reshaping, on a tiny example

In practice the heads aren’t separate modules. One linear layer produces the queries for all heads at once, as one long vector per token, and the vector is then cut into equal slices, one per head. With 4-number vectors and 2 heads of 2 (seed 123):

q, all heads, shape (1, 3, 4):        q split into heads, shape (1, 2, 3, 2):
[[-0.129, -0.323, -0.782,  0.115],    head 0: [[-0.129, -0.323],    head 1: [[-0.782,  0.115],
 [ 0.545,  0.409,  0.659, -0.959],             [ 0.545,  0.409],             [ 0.659, -0.959],
 [ 0.156, -0.503, -0.200, -0.204]]             [ 0.156, -0.503]]             [-0.200, -0.204]]

Head 0 gets the first two numbers of each token’s query and head 1 the last two. The code for this split is

q.view(B, T, n_heads, head_dim).transpose(1, 2)

view reinterprets the same numbers with a new shape without copying them: (1, 3, 4) becomes (1, 3, 2, 2), each token’s 4 numbers now 2 heads × 2. transpose(1, 2) swaps the token and head dimensions to give (1, 2, 3, 2), so each head has its own (tokens × 2) matrix. Every operation from the single-head version then runs on all heads at once, because matrix multiplication in PyTorch treats leading dimensions as a batch. The scores become (B, H, T, T), one table per head:

head 0 weights:          head 1 weights:
[[1.000, 0.000, 0.000],  [[1.000, 0.000, 0.000],
 [0.263, 0.737, 0.000],   [0.153, 0.847, 0.000],
 [0.381, 0.297, 0.322]]   [0.344, 0.340, 0.316]]

The two heads weigh the same tokens differently. Afterward the head outputs are put back side by side (transpose back, then reshape to (B, T, C)) and passed through one more linear layer, the output projection, which lets information from different heads mix.

Three matrices fused into one

GPT-2 goes one step further and stores the query, key and value matrices as a single linear layer from 768 to 3 × 768 = 2,304 outputs, which OpenAI’s code calls c_attn. That is the same as three separate 768 → 768 layers placed side by side: the first 768 outputs are the queries, the next 768 the keys, the last 768 the values. One big matrix multiply is faster than three small ones, and matching GPT-2’s layout makes chapter 6 a straight copy. GPT-2 also gives these projections a bias, which the hand example left out.

The full module

Everything above, at any size:

class CausalSelfAttention(nn.Module):
    def __init__(self, emb_dim, n_heads, context_length, dropout=0.0, qkv_bias=True):
        super().__init__()
        assert emb_dim % n_heads == 0, "emb_dim must split evenly across heads"
        self.n_heads = n_heads
        self.head_dim = emb_dim // n_heads
        self.qkv = nn.Linear(emb_dim, 3 * emb_dim, bias=qkv_bias)   # GPT-2: c_attn
        self.proj = nn.Linear(emb_dim, emb_dim)                     # GPT-2: c_proj
        self.attn_drop = nn.Dropout(dropout)
        # Built once at full size and sliced per call. Not saved in state_dict.
        self.register_buffer("mask", causal_mask(context_length), persistent=False)

    def forward(self, x, return_weights=False):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=-1)                      # each (B, T, C)
        # (B, T, C) -> (B, T, H, D) -> (B, H, T, D): one T-by-D matrix per head
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)  # (B, H, T, T)
        scores = scores.masked_fill(self.mask[:T, :T], float("-inf"))
        weights = torch.softmax(scores, dim=-1)
        out = self.attn_drop(weights) @ v                            # (B, H, T, D)

        out = out.transpose(1, 2).reshape(B, T, C)                   # heads side by side
        out = self.proj(out)
        return (out, weights) if return_weights else out

A few new pieces:

• .split(C, dim=-1) cuts the 2,304-wide output into three 768-wide tensors: queries, keys, values.
• register_buffer attaches a tensor to the module that isn’t a parameter. The mask moves to the GPU along with the module, but training never changes it. persistent=False keeps it out of the saved weights, since it can always be rebuilt. It is built once at the full context length, and self.mask[:T, :T] takes the top-left corner for a shorter input.
• Dropout is applied to the weights after softmax, as in the previous section. Dropout on the layer’s output belongs to the transformer block in chapter 4.
• return_weights=True also returns the attention weights, which is how the maps at the end of this chapter were drawn.

As a check, PyTorch has its own fused implementation, torch.nn.functional.scaled_dot_product_attention, which does the scale, mask, softmax and weighted sum in one call when given is_causal=True. Feeding it the same queries, keys and values as the module and comparing outputs at GPT-2 size, the largest difference was 2.4 × 10⁻⁷, the size of ordinary 32-bit rounding error. The module computes the same thing.

Real sizes

At GPT-2’s size (768 wide, 12 heads, context 1,024), on a batch shaped like chapter 2’s (seed 123):

attn = CausalSelfAttention(emb_dim=768, n_heads=12, context_length=1024, dropout=0.1).eval()
x = torch.randn(8, 256, 768)
out, w = attn(x, return_weights=True)
print(out.shape, w.shape)
print(sum(p.numel() for p in attn.parameters()))

torch.Size([8, 256, 768]) torch.Size([8, 12, 256, 256])
2362368

The output has the same shape as the input, (B, T, C), so attention layers can be stacked. The parameters break down as:

GPT-2 has 12 of these, one per block.

The weights tensor is the expensive part. It holds one number for every (query, key) pair in every head: 8 × 12 × 256 × 256 = 6.3 million numbers, 25 MB as 32-bit floats. That grows with the square of the context. At the full 1,024 tokens it would be 16 times larger, 403 MB for this batch and this one layer. On a freshly initialized module the weights are close to even: row 5 of head 0 puts between 0.15 and 0.21 on each of the six tokens it can see, nearly the 1/6 = 0.167 an exact average would give. The GPU track’s chapter on LLMs covers what that square does to memory traffic on real hardware, and how fused kernels such as FlashAttention avoid ever storing the full table.

What trained attention looks like

An untrained module spreads its weight almost evenly. A trained one doesn’t. To see the difference, GPT-2’s published weights (from the openai-community/gpt2 repository on Hugging Face) were loaded into this exact CausalSelfAttention module, one per layer. Reaching layers 4 and 5 means running the first layers in full, which needs the pieces chapter 4 builds; the maps below come from that run, with this chapter’s module computing the attention, and with return_weights=True returning the weights.

The input is “The Time Traveller smiled at us.” written twice, 16 tokens. In the labels, a dot marks a leading space.

The middle map is from head 11 of layer 4 (both counted from 0). Every row puts its entire weight, 1.00 to two decimal places, on the token just before it. It is a “previous token” head: after it, each position’s vector carries a copy of what came right before.

The right map, head 1 of layer 5, does something more interesting on the second copy of the sentence. At the second " Time" it puts 0.67 of its weight on " Trave", the token that followed " Time" the first time. At the second " Trave" it puts 0.76 on “ller”, and at the second “ller” it puts 0.72 on " smiled". The rule it follows is “find where this token appeared before, and look at what came next,” which is exactly what continuing a repeated pattern needs. Heads with this pattern are called induction heads in interpretability research, and the usual explanation of how they work is that they lean on a previous-token head in an earlier layer, like the one in the middle map, to know what came before each token. Most of the rest of this head’s weight goes to the very first token; each row has to add up to 1, so a head with nothing useful to find still has to put its weight somewhere.

The full model’s predictions line up with this. At the second " Time" its top prediction for the next token is " Trave", and at the second " Trave" it is “ller”. At the second “ller” it predicts " was" rather than " smiled", so the copying isn’t perfect after one repetition.

None of these behaviors were programmed. They are the same four matrices as the untrained map on the left, with different numbers in them, found by training on text.

Where this leaves things

Each position can now gather information from the positions before it, chosen by learned queries and keys, carried by learned values, across 12 heads that each look for something different, and never from the future. The output has the same (B, T, C) shape it came in with.

But attention on its own is a mixing step. Each output is a weighted average of value vectors, followed by one linear projection. It moves information between positions; it does very little with that information once it arrives. A model also needs a step that works on each position’s vector individually and can compute non-linear things from it. And a stack of twelve such layers needs help to train at all: the numbers flowing through it must stay in a sensible range, and the training signal has to reach the first layer without fading away.

Those are the rest of the transformer block: layer normalization, a small feed-forward network with the GELU activation, and residual connections around both. Chapter 4 builds them, stacks twelve blocks, adds the output head, and counts every parameter of a GPT-2-shaped model.