Seeing attention.

Attention is the step inside a language model where each word decides which earlier words to read from, and how much. You can print those decisions and draw them as lines between words. The picture is one of the best windows into a model that exists, and one of the easiest to misread.

You’ve written attention from scratch, so you know the formula. The formula doesn’t tell you what the 144 separate attention units inside a trained GPT-2 actually ended up doing, or how to prove it. Finding one specific behaviour with a number, then switching it off to see what breaks, is the core move in research on what goes on inside trained models, and Anthropic’s research engineer postings ask for exactly this.

It builds on the GPT-2 you made in Train a small LM, end to end, and on The KV cache, by hand, which showed the same reading step from the serving side. Afterwards, FlashAttention and kernels explains why fast inference never builds the picture you’re about to draw.

Today

• Walk one word through attention by hand: score it against every earlier word, shrink the scores so they don’t blow up, block out the future, turn the scores into weights that add up to one, and blend. Say the shape of every array on the way, for each attention unit and each layer.
• Pull the real attention weights out of GPT-2 small, check them against your own from-scratch math to six decimal places, and ship them to the browser as JSON.
• Build an attention viewer in plain JavaScript and SVG: pick a layer and a unit, hover a word, see what it reads.
• Find, with a number rather than by squinting, the unit that always reads one word back and the unit that spots a repeat and predicts whatever came after it last time, using a string of random words said twice.
• Switch those units off and measure what breaks, and explain why a bright attention line is a clue, not a proof.

You should have built GPT-2 yourself (Train a small LM, end to end is the closest lesson here), and The KV cache, by hand helps. Budget one evening for the idea and the export, one for the viewer, and one for scoring units and switching them off. Everything runs on a laptop CPU.

The Problem

You built GPT-2. You wrote q @ k.transpose(-2, -1), you divided by something, you called masked_fill, the loss went down. Now someone asks the obvious question: when the model predicts the next token, what is it looking at?

You have the weights. You have the numbers. And yet a surprising number of small things go fuzzy the moment you try to answer. Which axis of the attention matrix is the query? Does row 7 sum to one, or column 7? Why does almost every head, in every layer, seem to stare at the first token? If head 5.5 puts 0.96 on some token, does that mean the model “used” that token? If you delete the head, does anything change?

So you dump output_attentions=True, get a tuple of twelve tensors of shape (1, 12, T, T), plot one as a heatmap, and see a pale triangle with a bright stripe down the left edge. That picture is real data. It is also, on its own, almost useless. This lesson is about turning it into something that answers questions: a viewer you can poke, scores that find specific behaviours across all 144 heads, and an experiment that tells you whether what you see actually matters.

The Idea

Attention is a lookup with soft edges

Think of a library where every book has a label on its spine and some content inside. You walk in with a question. You compare your question to every spine, and instead of taking one book, you take a little of each, more from the books whose labels match your question best.

That is one attention headattention headOne independent attention computation inside a layer. GPT-2 small has 12 per layer, 12 layers, so 144 heads, each with its own query, key and value weights.. Each token in the sequence plays all three roles at once:

• its query (q) is the question it brings,
• its key (k) is the label on its spine,
• its value (v) is the content it hands over if picked.

All three are just the token’s current vector multiplied by three learned matrices. “Current vector” matters: that vector is the residual streamresidual streamThe running vector for each token that every layer reads from and adds its output to. Layer 5 does not see raw words, it sees words plus everything layers 0-4 have written. at that layer, not the raw word. By layer 5, the vector for ” river” can already carry “I came right after ’ lamp’”, and its key can advertise that.

The shapes, once and slowly

Here is GPT-2 small, one layer, one sequence of T tokens. The model width is 768, there are 12 heads, so each head works in 768 / 12 = 64 dimensions, called d_headd_headThe size of each head’s query, key and value vectors. 64 in GPT-2 small. Heads are carved out of the model width, not added on top..

Two habits will save you hours. First, rows are queries, columns are keys. A[i, j] is “how much token i reads from token j”. Because of the mask, the matrix is lower-triangular: row 0 has one nonzero entry, row 10 has eleven. Second, rows sum to one, columns don’t. A row is a probability distribution over the past. A column is “how popular is this token as a source”, and it can sum to anything from 0 to T.

One row, five steps

Take a single query: “sat” in <eot> The cat sat on the. Step through it below. The numbers are a toy (they aren’t from a real model), but the arithmetic is exactly what the model does.

Step 1: dot products. For every key j, compute q·k_j, a single number. Big positive means “my question matches your label”. There’s one per key, including keys in the future, because in training you compute the whole T×T grid in one matrix multiply.

Step 2: scale. Divide by √d_head, √64 = 8. Why? If the 64 entries of q and k are roughly independent with variance 1, their dot product is a sum of 64 terms and has variance 64, so a typical size of about 8. Dividing by 8 brings scores back to “about 1”, whatever d_head you chose.

Step 3: mask. A language model predicting token i+1 must not see token i+1 or anything after it. So set every score with j > i to −∞. This is the causal maskcausal maskA rule that each position may only attend to itself and earlier positions. Implemented by setting future scores to minus infinity (or a huge negative number) before softmax.. In code it’s scores.masked_fill(future, -inf). HF’s GPT-2 uses the most negative float instead of a literal −∞, which gives the same zeros after softmax.

Step 4: softmax. Exponentiate each score and divide by the sum. Every weight is positive, the row adds to 1, and e^−∞ is 0, so masked keys get exactly nothing. Note what softmax forces: the row must add to one even if nothing in the past is useful. Hold that thought, it explains the brightest stripe in every attention plot.

Step 5: mix. The head’s output for “sat” is the weighted sum of value vectors, Σ_j A[i,j] · v_j. If “cat” got 0.56, then 56% of the output is cat’s value vector. Then the 12 heads’ outputs are glued side by side (12 × 64 = 768), multiplied by W_O, and added back into the residual stream.

That’s everything a single head does. Now multiply: 12 heads per layer, 12 layers, and each one is a full T×T matrix. For a 15-token sentence that’s 144 little triangles. You need a way to look at them.

What real heads do

Below are four of GPT-2 small’s 144 heads, all looking at the same 15 tokens: the <|endoftext|> token, then seven arbitrary words, then the same seven words again. I ran GPT-2 small through Hugging Face transformers on this Mac, took the attention weights, and rounded to two decimals. Nothing else is processed.

Spend a few minutes here before reading on. Each head has a personality you can name.

L4 H11: the previous-token head. Every query sends almost all its weight one step back: “river” reads “lamp”, “seven” reads “river”. Averaged over a longer random sequence it puts 0.99 on the previous token. Its job is bookkeeping. It copies “the token before me was X” into each position’s residual stream, so later layers can ask about it.

L0 H1: the same-token head. In the first half each token reads itself. In the second half, the second “river” splits about 50/50 between itself and the first “river”. It’s matching on token identity: “where else have I appeared?”

L5 H7: the sink head. About 0.9 of every row lands on position 0, whatever the query is. This head has, on this input, nothing to do, and the softmax insists on spending the full budget of 1.0 somewhere.

L5 H5: the induction head. This is the one to stare at. In the first half every query parks on position 0, like the sink head. In the second half something new happens: the second “lamp” (position 8) reads “river” (position 2), with weight 0.96. Not “lamp”, the token after the earlier “lamp”. The second “river” reads “seven”. Every repeated token reads the token that followed its previous occurrence.

Induction heads: a two-head circuit

That two-step trick has a name. An induction headinduction headA head that, at a token A, attends to the token that came right after an earlier A, and copies it. Implements the rule “A B … A → predict B”. implements “A B … A → B”. It never works alone; it’s the second half of a circuitcircuitA small set of components (heads, MLP neurons) that together implement one identifiable behaviour, found by tracing how each one’s output feeds the next..

The previous-token head changes what B’s key looks like. That’s called K-compositionK-compositionWhen one head’s output, written into the residual stream, changes the keys a later head sees. The later head can then match on information the first head moved there., one of three ways heads compose in Elhage et al.’s “Mathematical Framework for Transformer Circuits” (the other two: changing a later head’s queries, Q-composition, or its values, V-composition). You can’t see composition in any single attention plot. You see two separate patterns, a diagonal stripe in layer 4 and an off-by-one stripe in layer 5, and you have to know to connect them.

Olsson et al. (2022) argue induction heads are a big part of how models do in-context learningin-context learningGetting better at predicting later tokens in a prompt by using earlier tokens in the same prompt, with no weight updates.. They form fairly abruptly early in training, in a “phase change” visible as a bump in the loss curve, and at the same moment the model gets much better at using long contexts. They also do more than literal copying: some induction-like heads match on meaning (“fuzzy” or translated copies), not just exact tokens.

Why random tokens

To find induction heads, you want a test where only induction can help. Natural text is a bad test: a head that copies “York” after “New” might be using induction, or might just know English. So: pick 50 random token ids, then repeat them. The first half is pure noise, unpredictable by any means. The second half is perfectly predictable, but only if the model looks back.

This gives you a score. If the repeat length is T, then in the second half, query position i holds the same token as position i − T, and the “token after the earlier copy” is at i − T + 1. So:

• previous-token score of a head = average of A[i, i-1] over positions,
• induction score of a head = average of A[i, i-T+1] over second-half positions.

A head that attended uniformly would score about 1/i, around 0.01-0.02 here. Real previous-token and induction heads score above 0.8. Compute these for all 144 heads and the right ones stand out like lit windows.

The sink: where attention goes when it has nowhere to go

Back to the stripe down the left edge. Over my 101-token random sequence, the average head put 0.42 of its attention on position 0. Some heads put over 0.9 there. Is position 0 that important?

Mostly not. A head that has nothing useful to read still has to spend a full 1.0. The cheapest way to do “nothing” is to put the weight on a token whose value vector is small or harmless, consistently. Position 0 is the one token every query can see, so models learn to use it as a parking spot. Xiao et al. (2023) named this an attention sinkattention sinkA token, usually the first, that receives a large share of attention from many heads regardless of content, acting as a place to dump attention the softmax forces each row to spend.. They found it the hard way: when you run a model with a sliding window and evict the first token’s keys and values from the KV cache, quality collapses, even though that token is meaningless. Keeping just the first four tokens’ KV restored it. That’s StreamingLLM.

The sink also teaches the first caveat of this lesson. Here’s what I measured on the 15-token sequence:

The value normvalue normThe length of a token’s value vector. A token can get high attention weight yet contribute little if its value vector is short. at the sink is about five times smaller than elsewhere. So 92% of the attention becomes a smaller share of what actually gets mixed in. (This “weight times norm” view is a rough proxy from Kobayashi et al., 2020; the true contribution also goes through W_O and depends on direction, not just length.) The heatmap lied by omission: it shows weights, and weights are only half of the product.

What a bright line does and does not tell you

Collecting the caveats, because the viewer you’re about to build will make all of them very easy to forget.

1. Weights are not contributions. Output = Σ weight × value. A huge weight on a tiny value adds little. Even the value isn’t the end: it’s multiplied by W_O, and then added to a residual stream that already holds a lot. A head’s whole output can be small next to the stream it joins.
2. Weights are not the cause. Jain & Wallace (2019) showed that for many text classifiers you could find very different attention distributions that gave the same prediction, and that attention often disagreed with gradient-based importance. Wiegreffe & Pinter (2019), “Attention is not not Explanation”, pushed back: swapping weights after training isn’t a fair test, and attention can be informative if you test it properly. Both agree on the practical point: a weight is a hypothesis to test, not a finding.
3. Tokens are not words at deeper layers. At layer 8, position 5’s vector carries information mixed in from positions 0-4. “Head 8.4 looks at ‘orange’” really means “looks at whatever is in position 5’s residual stream now”. Attention rollout (Abnar & Zuidema, 2020) tries to account for this by chaining matrices across layers; it’s a rough approximation.
4. One head is not one behaviour. A head can be a previous-token head on one input and a sink on another. Head 5.5 is an induction head when there is a repeat; otherwise it parks on the sink.
5. Heads back each other up. As you’re about to measure, delete one induction head and another covers for it.

The test that actually answers “does the model use this?” is the ablationablationDeleting or neutralising a component (here: a head’s output) during the forward pass and measuring how much the model’s behaviour gets worse.: switch the head off and see if the loss goes up. Zero-ablation sets the head’s output to zero. mean-ablationmean-ablationReplacing a head’s output with its average output over many inputs, instead of zero. Less likely to knock the model into a state it never sees in training. replaces it with its typical value. Neither is perfect, but either one is a causal test, and an attention plot is not.

Build It

You’ll make three files in one folder: an export script, a scoring and ablation script, and a small web page. Use Hugging Face’s GPT-2 so you can trust the weights. If your own from-scratch GPT-2 loads the same weights, Task 2 doubles as a test of your implementation.

mkdir -p ~/attn-viz && cd ~/attn-viz
python3 -m venv .venv && source .venv/bin/activate
pip install torch transformers

# check.py
import torch
from transformers import GPT2LMHeadModel, GPT2TokenizerFast

tok = GPT2TokenizerFast.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2", attn_implementation="eager").eval()

ids = [tok.bos_token_id] + tok(" The cat sat on the mat")["input_ids"]
with torch.no_grad():
    out = model(torch.tensor([ids]), output_attentions=True)

print(len(out.attentions), out.attentions[0].shape)   # 12, (1, 12, T, T)
A = out.attentions[5][0, 5]
print(A.sum(-1))           # each row sums to 1
print(A.triu(1).abs().max())  # nothing above the diagonal

out = model(x, output_attentions=True, output_hidden_states=True)
blk = model.transformer.h[5]
h = blk.ln_1(out.hidden_states[5])[0]          # (T, 768): layer 5's input, after its LayerNorm
qkv = ...                                       # (T, 2304) from blk.attn.c_attn
q, k, v = ...                                   # split, then (12, T, 64)
scores = ...                                    # (12, T, T)
# mask, softmax
print((A_mine - out.attentions[5][0]).abs().max())

{
  "tokens": ["<eot>", " lamp", " river", "..."],
  "n_layers": 12, "n_heads": 12,
  "attn": [[ [[1.0], [1.0, 0.0], [0.0, 1.0, 0.0]], "...12 heads" ], "...12 layers"]
}

const NS = "http://www.w3.org/2000/svg";
function el(tag, attrs) { /* createElementNS + setAttribute loop */ }

async function main() {
  const data = await (await fetch("attn.json")).json();
  let layer = 0, head = 0, focus = null;   // focus = {side: "q"|"k", i}
  function draw() { /* clear svg, draw tokens, draw edges filtered by focus */ }
  // wire selects and hover, then draw()
}
main();

torch.manual_seed(0)
T = 50
rand = torch.randint(1000, 10000, (4, T))   # 4 sequences; skip the rarest ids
ids = torch.cat([torch.full((4, 1), 50256), rand, rand], dim=1)   # (4, 101)

ablate = {}   # layer -> list of heads

def make_hook(layer):
    def pre_hook(module, args):
        # args[0]: (B, T, 768). Return a modified tuple, or None to leave it alone.
        ...
    return pre_hook

for l, blk in enumerate(model.transformer.h):
    blk.attn.c_proj.register_forward_pre_hook(make_hook(l))

Measure

Fill this in from Tasks 4 and 5. My numbers are from GPT-2 small on an Apple-silicon Mac CPU, torch.manual_seed(0), token ids drawn from 1000-9999, T = 50. Your exact numbers will differ a little with seed and batch size; the head identities and the shape of the results should not.

What to notice:

• Only a handful of heads per pattern score high. The scores are bimodal: most heads near zero, a few near one. That’s why scores beat eyeballing 144 plots.
• Induction heads sit in the middle layers (5-7), after the previous-token head (4). That order is forced by the circuit: the key-side information has to exist before something can match on it.
• Ablation is far from linear. One head: +0.05. Five heads: +4.3. The heads back each other up until too many are gone.
• The top previous-token head is not alone either. Killing 4.11 alone barely touches the loss and only dents 5.5’s induction score, from 0.94 to 0.82; it takes four previous-token heads to bring it to 0.63. Weaker previous-token heads in layers 2-6 carry the same information.
• Zero-ablation is a blunt instrument. With the top 5 heads zeroed, the first-half loss also rose, from 12.0 to 14.7, on tokens no induction head could help with. Zeroing pushed the residual stream somewhere it never goes in training and the model got worse at everything. Compare with your mean-ablation numbers; the gap between the two is a measure of how off-distribution zero-ablation is.

Break It

Use the default attention implementation. Drop attn_implementation="eager" and ask for attentions. On my transformers version, out.attentions came back as an empty tuple with no error, and out.attentions[0] raised an IndexError three lines later. Fused kernels never build the matrix you want to plot, which is the whole reason they are fast. The lesson: the attention matrix is a debugging artifact, not something production inference ever holds in memory.

Forget the BOS token. Run the same 15 tokens without <|endoftext|> at position 0. Now “lamp” is position 0, and it becomes the sink: when I tried it, head 5.7 put 0.93 on “lamp”. Your sink heads now “attend to lamp”, and a naive reading says the model finds “lamp” fascinating. Same model, same words, a different story, purely from where the parking spot ended up.

Score induction on natural text. Use a repeated English paragraph instead of random tokens. The first-half loss drops to a normal English-text level (I got 2.9 nats on a 30-token sentence), because the model can predict English without looking back. Now a low second-half loss mixes two things, recall and plain language knowledge, and a head that attends to the “next” token might just be following grammar. The clean separation between “needs induction” and “doesn’t” is gone. Random tokens are the control that makes the measurement mean something.

Transpose the matrix in your viewer. Swap rows and columns when you draw. Now the previous-token head appears to look forward one step and the sink appears to be “position 0 reading everything”, which is impossible under the causal mask. If a plot shows a token attending to the future, your axes are flipped. The mask is a free correctness check: the upper triangle must be empty.

Explain It Back

Answer out loud, in plain words, as if in an interview. Then check.

Go Deeper

• Nelson Elhage, Neel Nanda, Catherine Olsson et al., “A Mathematical Framework for Transformer Circuits” (Anthropic, 2021, transformer-circuits.pub). Rewrites a transformer as heads reading from and writing to a residual stream, defines Q/K/V-composition, and derives induction heads from two-layer attention-only models. The single most useful thing to read after this lesson.
• Catherine Olsson, Nelson Elhage, Neel Nanda et al., “In-context Learning and Induction Heads” (Anthropic, 2022, transformer-circuits.pub). The evidence that induction heads drive in-context learning, the phase change in training, and the prefix-matching and copying scores you approximated here.
• Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, Mike Lewis, “Efficient Streaming Language Models with Attention Sinks” (2023, arXiv 2309.17453). Where attention sinks got their name, and the KV-cache trick that keeps the first four tokens.
• Sarthak Jain and Byron C. Wallace, “Attention is not Explanation” (NAACL 2019) and Sarah Wiegreffe and Yuval Pinter, “Attention is not not Explanation” (EMNLP 2019). Read them as a pair; the argument is the point.
• Kevin Clark, Urvashi Khandelwal, Omer Levy, Christopher D. Manning, “What Does BERT Look At? An Analysis of BERT’s Attention” (BlackboxNLP 2019). The classic catalogue of head types in a real model: positional heads, heads that dump attention on [SEP] and punctuation, heads that track syntax.
• Goro Kobayashi, Tatsuki Kuribayashi, Sho Yokoi, Kentaro Inui, “Attention is Not Only a Weight: Analyzing Transformers with Vector Norms” (EMNLP 2020). The weight-times-value-norm view used in the sink table, with the finding that high-weight special tokens often carry small norms.
• Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, Jacob Steinhardt, “Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small” (2022). A full circuit in the same model you used, found with ablation and path patching, including the backup heads that explain your redundancy result.
• TransformerLens (Neel Nanda and contributors, on GitHub) and BertViz (Jesse Vig). The standard library for hooking and ablating heads in GPT-2-style models, and the standard attention viewer. Use them to check your tools after you’ve built your own.

Connections

This builds directly on Train a small LM, end to end: the GPT-2 you built has every piece in Figure 001, and Task 1 is a test of whether your version matches. The KV cache, by hand is the same Q/K/V seen from serving: the keys and values you cache are exactly what each new query reads, and the attention sink is why sliding-window caches keep the first tokens.

FlashAttention and kernels explains why output_attentions needs the slow path: FlashAttention’s whole trick is never writing the T×T matrix you just plotted.

Next, the natural step is activation patching: instead of zeroing a head, copy its output from a run on a different input and see what the model now predicts. That turns “this head matters” into “this head carries this piece of information”.