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.
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.
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
| Limit | Why it matters |
|---|---|
| Weight is not contribution | A token can get high attention but have a tiny value vector, so it adds almost nothing to the output. |
| Many heads, many layers | One map shows one head in one layer. The final output mixes all heads, all layers and the MLPs. |
| Tokens are already mixed | After 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 cause | A map is an observation. It does not show what would happen if the head were removed. |
| Sinks and parking spots | In 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
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))- "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
- "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
Watch and read more
Lab
A short note, backed by numpy runs, showing when an attention map misleads.
The problem
Use the code block from this lesson as your starter.
- Run it and note which token gets the most attention and which contributes most.
- Make a new version where the attention map does point to the most important token. What did you have to change?
- Now make a version where two tokens have equal attention but very different contributions.
- 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.
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.