Attention Maps: What They Show and What They Do Not

An attention map shows how much each position attended to each other position. It does not show why the model produced its output, and treating it as an explanation has been shown to be unsafe in a specific, reproducible way. This page is about that gap.

What an attention weight actually is

In one head, for one query position, the model computes a score against every key position, softmaxes those scores into weights that sum to one, and takes the weighted sum of the value vectors. The attention map is that weight matrix - per head, per layer, per example.

In a model with 32 layers and 32 heads you have 1,024 of these matrices for a single forward pass, which is the first thing people underestimate. The weight tells you the coefficient on a value vector. It does not tell you the magnitude of that value vector, what the head does with the result, whether the next 31 layers use it, or whether the same information also arrived through the residual stream by another route. Every one of those is required to get from “attended to” to “because of”.

If you have read how attention works, this page is the part that comes after.

A useful reframing: attention weights describe the connectivity of one layer’s information routing. They are a wiring diagram for a single step, and a wiring diagram tells you what could flow, not what did.

The “attention is not explanation” result

Jain and Wallace put this on the record at NAACL 2019 in a paper titled exactly that. The argument runs through two experiments, and both are simple enough to reproduce on any classifier with an attention layer.

Experiment one: correlation with other attribution methods

If attention weights explain the prediction, they should agree with other measures of which input tokens mattered - gradient magnitudes, or the effect of deleting a token. Across a range of text classification tasks the agreement was frequently poor. Two methods both claiming to identify the important tokens picked different tokens.

Experiment two: adversarial attention

This is the harder result. Hold the model fixed, take a single example, and search for an alternative attention distribution - one that puts its mass on entirely different tokens - such that the model’s output is essentially unchanged. On many examples such a distribution exists and is easy to find.

If two contradictory attention maps produce the same prediction, at most one of them is the explanation, and nothing about either map tells you which.

The mechanism behind this is not mysterious. The value vectors at different positions are often similar, especially after several layers of mixing, so the weighted sum is robust to redistributing weight among them. An explanation that survives changing what it points at is not doing explanatory work.

The reply, and what survived it

Wiegreffe and Pinter replied at EMNLP 2019 with “Attention is not not Explanation”, and the reply is as important as the original because it sharpens what the claim can be. Their objections, in substance:

• The adversarial distributions are not free. They were found by optimising per-example, outside the model. A distribution the model itself could arrive at, under its own training, is a stronger test - and when they trained models adversarially to use different attention, performance degraded.
• “Explanation” was underspecified. If you mean the unique account of the prediction, attention is not it. If you mean a plausible account consistent with the model’s behaviour, attention can be one, and the existence of another does not refute it.
• The right baseline is a uniform-attention model. Freeze attention to uniform, retrain, and see how much accuracy you lose. If you lose a lot, the learned attention was doing work, even if the specific map is not the unique explanation.

Serrano and Smith, also in 2019, asked a third version of the question - if you erase the most-attended representations, does the prediction change more than if you erase randomly chosen ones? Often yes, but far less reliably than the attention ranking would suggest.

The consensus that came out of this exchange is narrow and worth stating precisely: attention weights are evidence about information routing and are not, on their own, a causal account of an output. They can be part of one, alongside an intervention.

Why the weight is only half the story

A concrete refinement that came later, and that fixes some of the most misleading pictures: what a head writes to a position is the attention weight times the value vector, so a large weight on a near-zero value vector contributes nothing. Weighting each attention coefficient by the norm of the corresponding transformed value vector produces a substantially different, and better-behaved, picture. In particular it largely dissolves the “attention sink” artefact, where heads dump enormous weight on the first token or on punctuation: those positions receive weight, but what is written from them is small. If you are going to visualise attention at all, visualise the weighted version.

Combining layers: rollout and flow

A single layer’s map answers a question nobody asked. What people want is: how much did input token j influence the representation at position i after all the layers? Multiplying the per-layer attention matrices together is the obvious move and it is wrong, because it ignores the residual connection: information also reaches the next layer without passing through attention at all.

Abnar and Zuidema proposed the standard correction at ACL in 2020. Attention rollout mixes each layer’s head-averaged attention matrix with the identity before multiplying, so the residual path is represented:

A_hat[l] = 0.5 * A[l] + 0.5 * I  # A[l] averaged over heads, rows renormalised
rollout = A_hat[L - 1] @ ... @ A_hat[1] @ A_hat[0]
# row i of the rollout matrix is the claimed contribution of every input position
# to the representation at position i after L layers.

Their second method, attention flow, treats the layers as a graph and computes a maximum flow from input positions to the target, which is more expensive and gives different, often smoother, attributions. Both are improvements on naive multiplication and neither escapes the critique above. They still average over heads, which discards the fact that different heads do different jobs; they still weight by coefficients rather than by what was written; and they still describe routing rather than causation.

Rollout is a better picture of connectivity, not a repair of the explanation claim. Treat a rollout map the way you would treat a network diagram: useful for finding where to look, never a statement about why.

Getting the weights out

Two lines, and the shape is the thing to know. With HuggingFace transformers, pass output_attentions=True and you get a tuple with one tensor per layer, each of shape (batch, heads, query_positions, key_positions).

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "gpt2"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(
    name,
    attn_implementation="eager"  # required: fused kernels return no weights
)
model.eval()

text = "The doctor asked the nurse a question because she"
ids = tok(text, return_tensors="pt")

with torch.no_grad():
    out = model(**ids, output_attentions=True)

attn = out.attentions  # tuple, one per layer
print(len(attn), attn[0].shape)  # 12 torch.Size([1, 12, 10, 10])

# layer 5, head 3, what the final position attended to
row = attn[5][0, 3, -1]
for tokid, w in zip(ids["input_ids"][0], row):
    print(f"{tok.decode([tokid]):>12} {w.item():.3f}")

attn_implementation="eager" matters. Flash and SDPA attention kernels never materialise the full weight matrix, so output_attentions=True either returns None or is silently ignored depending on version. If your attention tuple is empty, this is why.

What attention maps are still good for

They remain genuinely useful, just not as explanations.

Finding candidate components

A head whose weights consistently point one token back, or to the previous occurrence of the current token, is a lead worth chasing with an intervention. This is how induction heads were first noticed - and note that they were only established as causal by ablation afterwards.

Debugging long-context behaviour

If a model is ignoring the middle of a long document, the attention pattern will show you it. That is a claim about routing, which is what attention legitimately reports, and it connects directly to the lost-in-the-middle effect.

Verifying an architectural change

Sliding-window or sparse attention should produce a visibly banded map. If it does not, the configuration is not doing what you think.

What they are not good for

Showing a user why a decision was made, justifying a decision to an auditor, or supporting a claim that the model “focused on” anything. For that you want an intervention - activation patching answers the causal question that attention maps only appear to.

Related

• Probing Classifiers: Finding What a Layer Knows
• Activation Patching and Causal Tracing
• Saliency Maps in Vision Models