Looking Inside a Trained Model
Visualising attention maps
An attention map shows exactly which earlier words a model looked at while reading each word, and real maps are often messier than the tidy diagrams used to teach the concept.
- 9 min read
- 3 reading levels
- Published
Read these first
On this page 5
One lesson, three depths. Pick the one that fits you today — you can switch any time.
Beginner — No maths. Plain English.
An attention map shows exactly which earlier words a model looked at, for every word it read.
Imagine filming someone's eyes as they read a sentence. Their gaze jumps around, resting longest on the words that matter most for what they are currently thinking about.
An attention map is that recording, for a model instead of a person. Earlier lessons described attention as a concept. This one actually looks at the numbers, for a real model, on real text.
Why it exists
Knowing attention exists is not the same as seeing it. A diagram in a lesson is drawn to teach a clean idea. A real model's attention pattern is measured, not drawn, and measured things are messier than diagrams.
Looking at real attention maps checks whether the tidy story actually holds up. Sometimes it does. Sometimes a real model's attention looks nothing like the clean example used to explain the concept.
That gap between the clean story and the messy reality is worth seeing directly, at least once.
How it works
Sentence: "The chai was hot so it burned my tongue"
For the word "it", check attention back to every earlier word:
The 0.000
chai 0.000
was 0.000
hot 0.000
so 1.000 <-- almost all attention lands here
it 0.000This is a real measurement from a real small model, not an invented example. The strongest pattern found was not "it" attending to "chai", the sentence's actual subject. It was a much simpler pattern: attending to the word right before it.
Where you have already seen it
- AI interpretability tools. Browser-based tools that let you type a sentence and see a live attention map are showing exactly this.
- Chatbot debugging. Engineers checking why a model misread a sentence sometimes inspect its attention pattern for a specific answer.
- Research papers. Attention-map diagrams are a common figure in papers explaining what a model is doing on a specific example.
Remember this
- An attention map shows exactly which earlier words a model attended to, as real numbers.
- Real attention patterns are often messier and simpler than clean teaching diagrams suggest.
- The most reliable, easy-to-find patterns are often simple ones, like attending to the previous word.
What to learn next
- Attention — the concept this lesson measures directly, in a real model.
- Induction heads — a more complex, two-step attention pattern built from simpler ones like this.
- Multi-head attention — why a model has 144 separate attention patterns to search through, not one.
Developer — Code and libraries.
This searches every layer and head of GPT-2 for the one that most consistently attends to the previous token, then displays that head's actual attention weights.
Setup
pip install transformers torchThe first run downloads gpt2, roughly 500 MB.
Finding and displaying a real attention pattern
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
tok = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2", attn_implementation="eager")
model.eval()
text = "The chai was hot so it burned my tongue"
inputs = tok(text, return_tensors="pt")
tokens = [tok.decode([t]) for t in inputs["input_ids"][0]]
with torch.no_grad():
out = model(**inputs, output_attentions=True)
n_layers, n_heads = len(out.attentions), out.attentions[0].shape[1]
print(f"{n_layers} layers x {n_heads} heads, sequence length {len(tokens)}\n")
# search for the head whose attention lands most on the previous token
best_score, best_layer, best_head = -1, None, None
for layer in range(n_layers):
for head in range(n_heads):
attn = out.attentions[layer][0, head]
prev_token_weight = torch.diagonal(attn, offset=-1).mean().item()
if prev_token_weight > best_score:
best_score, best_layer, best_head = prev_token_weight, layer, head
print(f"strongest previous-token head: layer {best_layer}, head {best_head} "
f"(avg weight on previous token = {best_score:.3f})\n")
attn = out.attentions[best_layer][0, best_head]
it_pos = tokens.index(" it")
print(f"attention FROM 'it' (position {it_pos}) TO every earlier word:")
for pos in range(it_pos + 1):
bar = "#" * int(attn[it_pos, pos].item() * 40)
print(f" {tokens[pos]!r:>8} {attn[it_pos, pos].item():.3f} {bar}")12 layers x 12 heads, sequence length 10
strongest previous-token head: layer 4, head 11 (avg weight on previous token = 1.000)
attention FROM 'it' (position 6) TO every earlier word:
'The' 0.000
' ch' 0.000
'ai' 0.000
' was' 0.000
' hot' 0.000
' so' 1.000 #######################################
' it' 0.000 Out of 144 attention heads, one attends to the previous token with weight 1.000, essentially every time. That is a real, clean, reliably-findable pattern, sometimes called a previous-token head.
Line by line
attn_implementation="eager" is required here. Some faster attention implementations do not expose their internal weights at all, so requesting output_attentions=True without this silently returns nothing useful.
torch.diagonal(attn, offset=-1) reads exactly the "attention to the previous position" entries out of the full attention matrix, one clean way to search for this specific, well-known pattern automatically.
Notice "chai" split into two tokens, ' ch' and 'ai'. This is the tokenizer at work, from earlier in the course. Position indexing has to account for this, or the word lookup silently breaks.
Common mistakes
Expecting every head to have a clean, nameable pattern. Most of the 144 heads in GPT-2 do not look like this one. Many attention patterns are diffuse, or only meaningful for specific kinds of input.
Assuming a strong attention weight means a strong causal effect. A head can attend heavily to a position without that position's information mattering much for the final output. Activation patching, from earlier in this section, is the way to check that directly.
Forgetting attn_implementation="eager". Without it, output_attentions=True can silently produce an empty result rather than a clear error, wasting time chasing the wrong bug.
Try it yourself
Search instead for the head with the strongest attention specifically from "it" to "ch" (the start of "chai"), across all layers and heads, the coreference pattern the concept lesson used as its teaching example.
The best match found this way scores only around 0.17, and even that head's strongest attention actually goes to "The", not "chai". The tidy textbook story and the real, searched-for result do not fully agree here, and that gap is worth sitting with rather than smoothing over.
What to learn next
- Induction heads — a two-step pattern built by composing simple heads like the one found here.
- Activation patching — testing whether a visible attention pattern actually matters for the output.
- The logit lens — the natural next step, reading the model's guess right after seeing where it looked.
Researcher — Mathematics and papers.
What an attention map actually is
For a single attention head at layer l, the map is the row-normalised score matrix:
A_l,h = softmax( (Q_l,h K_l,h^T) / sqrt(d_head) )A_l,hhas shapen x nfor sequence lengthn; rowisums to 1 and gives the distribution query positioniattends over.- A full model produces
L * Hsuch matrices,144for GPT-2 small (12layers times12heads), each a candidate for containing an interpretable pattern.
Automated search over this matrix, the kind used in the developer block, is standard practice for locating known head types. Previous-token heads use torch.diagonal(A, offset=-1). First-token or "attention sink" heads show column 0 dominating every row. Duplicate-token heads attend to earlier occurrences of the current token.
Why attention weights are not explanations
Jain & Wallace (2019), Attention is not Explanation, showed that for many models, one can construct an alternative attention distribution, over the same input, that yields a near-identical output prediction. If attention weights faithfully explained the model's reasoning, this should not generally be possible; the existence of alternative, equally-predictive attention patterns undermines a strong "attention equals explanation" reading.
Wiegreffe & Pinter (2019), Attention is not not Explanation, offer a partial rebuttal. Their tests show attention is not entirely arbitrary either, and can carry genuine explanatory signal, depending on the test and model. The honest current position sits between the two paper titles. Attention weights are informative but not sufficient on their own, and should be corroborated with a causal method like activation patching before being treated as an explanation.
Attention sinks
Xiao et al. (2023), Efficient Streaming Language Models with Attention Sinks, documented a widespread pattern. A large fraction of attention heads devote substantial weight to a sequence's very first token, regardless of its actual content. It functions as a kind of "no-op" the model can attend to when it has nothing more useful to attend to. This directly informed efficient long-context serving techniques, which deliberately preserve these sink tokens when evicting older context.
Complexity
Extracting attention maps costs no extra computation beyond a normal forward pass with output_attentions=True, since the weights are already computed internally, only not returned by default. Storage is the real cost at scale: O(L * H * n^2) floating point numbers for a full set of maps. That grows substantial at long context lengths, a direct consequence of the quadratic attention cost covered elsewhere in this course.
Key references
- Jain, S. & Wallace, B. (2019). Attention is not Explanation. arXiv:1902.10186
- Wiegreffe, S. & Pinter, Y. (2019). Attention is not not Explanation. arXiv:1908.04626
- Vig, J. (2019). A Multiscale Visualization of Attention in the Transformer Model. ACL Demo — the BertViz tool, a widely used interactive attention visualiser.
- Xiao, G. et al. (2023). Efficient Streaming Language Models with Attention Sinks. arXiv:2309.17453
Current state and open problems
Attention visualisation remains a useful first-pass tool, cheap, requires no extra training, and often surfaces genuine, well-documented patterns like previous-token heads, induction heads and attention sinks.
The open problem is the same one Jain & Wallace raised. An attention weight alone is observational, not causal. The field has increasingly moved toward pairing visualisation with activation patching or ablation, to confirm a visible pattern is actually load-bearing rather than stopping at the map itself. That is exactly the discipline this whole section has tried to model, lesson by lesson.
What to learn next
- Attention — the underlying mechanism this lesson visualises directly.
- Activation patching — confirming a visible pattern is actually causally load-bearing.
- Quadratic attention cost — why the storage cost noted above scales the way it does.