Inside the model: interpretabilityAdvancedLesson 3116 min read

Attention maps and their limits

Attention weights show where a head looks. They do not show what it takes from there, or what the model does with it.

Lesson in motion

In 60 seconds

Attention maps and their limits

Attention weights show where a head looks. They do not show what it takes from there, or what the model does with it.

1/4
In simple words
In a classroom, a student may look at the clock many times. That does not mean the clock taught them anything. Looking is not the same as learning. Attention maps show where a model "looks", not what it learns.
Recall from the build-an-llm track: each attention head scores every earlier token with query · key, turns the scores into weights with softmax, and outputs a weighted sum of value vectors. An attention map is a heatmap of those weights. It is easy to draw and pleasant to look at, so it became one of the first tools people used to explain transformers.

What attention maps do show

  • Which earlier tokens a head reads from, for this input.
  • Clear patterns in some heads: always looking at the previous token, at the first token, or at a matching earlier word.
  • Hints for where to look next, for example a head that attends from a pronoun to a noun.

Why they are not full explanations

LimitWhy it matters
Weight is not contributionA token can get high attention but have a tiny value vector, so it adds almost nothing to the output.
Many heads, many layersOne map shows one head in one layer. The final output mixes all heads, all layers and the MLPs.
Tokens are already mixedAfter layer 1, the vector at position 5 contains information from earlier tokens too. "Attends to token 5" may really mean "reads something copied there".
No test of causeA map is an observation. It does not show what would happen if the head were removed.
Sinks and parking spotsIn many models, heads put a lot of weight on the first token or on punctuation, often when they have nothing useful to do.

Weight vs contribution, in numpy

High attention, small effectpython
import numpy as np

tokens = ["The", "farmer", "sold", "the", "mangoes"]
scores = np.array([3.0, 1.0, 0.5, 0.2, 1.5])    # query . key for each token
weights = np.exp(scores) / np.exp(scores).sum()  # softmax

# Value vectors: what each token actually passes on.
values = np.array([
    [0.05, 0.02, 0.01],   # "The": big weight, tiny value
    [1.2, -0.8, 0.9],
    [0.3, 0.1, -0.2],
    [0.04, 0.01, 0.02],
    [1.5, 1.1, -0.7],
])

contrib = weights[:, None] * values              # each token's share of the output
size = np.linalg.norm(contrib, axis=1)
for t, w, s in zip(tokens, weights, size):
    print(f"{t:8s} attention={w:.2f}  size of contribution={s:.2f}")
print("output vector:", np.round(contrib.sum(axis=0), 2))
Here "The" gets about two thirds of the attention, yet "mangoes" contributes far more to the output. An attention map alone would point you at the wrong token. Multiplying the weight by the size of the value vector is a simple improvement. Even that is only a step: it still ignores what later layers do with the result.
Misreading an attention map
  • "The model decided the answer because it looked at this word"
  • Averaging all heads into one map and trusting it
  • Showing a map as the explanation to users
Using an attention map well
  • "This head moves information from this word to that position"
  • Looking at single heads, and weighing by value size
  • Treating it as a clue, then testing with an intervention
Note
Attention maps are still useful. Some well-known mechanisms, such as induction heads (Unit 3), were first spotted because their attention patterns looked striking. The lesson is: use them to form a guess, then test the guess.
Try this. In the code above, change the value vector of "The" to be large. How does the output change? What does that tell you about reading attention weights alone?

Watch and read more

Lab

A short note, backed by numpy runs, showing when an attention map misleads.

~25 min

The problem

Use the code block from this lesson as your starter.

  1. Run it and note which token gets the most attention and which contributes most.
  2. Make a new version where the attention map does point to the most important token. What did you have to change?
  3. Now make a version where two tokens have equal attention but very different contributions.
  4. Reading task: find any figure of an attention heatmap (in a blog, textbook or paper on arXiv). Write three sentences: what the figure shows, what someone might wrongly conclude from it, and what test would check the claim.

You are done when

Hard questions

Try to answer before you reveal. If you can answer these, you understood the lesson.

Q1Even if you weigh attention by value size, why can you still be wrong about which input word mattered for the final answer?Reveal

Because the head's output is only one of many things added to the residual stream, and later layers can ignore it, cancel it or transform it. The MLP layers, which store much of the model's knowledge, are not shown at all. And the vector at each position is already a mixture of earlier tokens. Only a causal test, such as patching or ablating the head, shows whether it matters for the answer.

Q2Why might two different heads with very different attention maps have the same effect on the output?Reveal

Models often have redundancy: several heads can move the same information by different routes, for example one reads it directly from a word and another reads a copy stored at a nearby position. Their maps look different but their outputs to the residual stream can be similar. This is also why removing one head sometimes changes little: another head compensates.

Please sign in to continue.

Questions people ask

Are attention maps useless then?

No. They are cheap and often point at interesting heads. The problem is treating them as a complete explanation. Use them to generate guesses, and use interventions to check the guesses.

Why do heads attend so much to the first token?

Softmax weights must add up to 1, so a head has to put its attention somewhere even when nothing is relevant. Researchers have observed that many models use the first token, or other low-information tokens, as a kind of resting place. Weight there usually means "nothing to do", not "this is important".

What is attention rollout?

It is a method that combines attention maps across layers to estimate how much each input token flows to the output. It is better than looking at one layer, but it still ignores value sizes and MLP layers, so treat it as a rough picture.

Can I show attention maps to end users as explanations?

Be careful. Users will read them as "the reason for the answer", which they are not. If you show them, label them honestly as "where the model looked", not "why it decided".

Lesson test

5 questions. Get 3 right (60%) to pass and complete this lesson.

Sign in to take the test and save your progress