Visualizing Attention

Attention weights describe how one token mixes information from other tokens inside a particular head and layer. Tools such as BertViz turn those weights into lines, matrices, and heatmaps.

The visualization is useful for understanding architecture and information flow.

A normalized attention row

For query qiq_i and key kjk_j:

aij=softmaxj(qikjdk)a_{ij}=\operatorname{softmax}_j\left(\frac{q_i\cdot k_j}{\sqrt{d_k}}\right)

The dot product measures compatibility, scaling controls magnitude, and softmax creates nonnegative weights that sum to one.

Read One Row of an Attention Matrix

Choose a query token. Its row lights up across all key tokens, and the weighted-value summary updates below.

I
gave
Tater
a
toy
I
0.55
0.12
0.14
0.08
0.11
gave
0.16
0.22
0.31
0.12
0.19
Tater
0.06
0.18
0.34
0.08
0.34
a
0.12
0.21
0.18
0.31
0.18
toy
0.08
0.14
0.42
0.12
0.24

context(Tater) = Σ attention(0.06, 0.18, 0.34, 0.08, 0.34) × value vectors

An attention matrix is a collection of query-specific distributions. Reading one row is clearer than treating the whole heatmap as one explanation.
Three BertViz attention-head plots for the sentence 'The quick, brown fox jumps over the lazy dog', showing a diagonal identity-like pattern at layer 4, a broadcast pattern from a single token at layer 1, and attention converging onto the token 'lazy' at layer 1
Source: Vig (2019)
BertViz attention-head view for BERT on two related sentences, comparing dense cross-token attention at layer 0, more selective attention at layer 10, and attention filtered to only show Sentence A attending to Sentence B
Figure 2 from Vig (2019): attention-head view for BERT, for inputs "the cat sat on the mat" (Sentence A) and "the cat lay on the rug" (Sentence B). The left and center figures represent different layers/attention heads. The right figure depicts the same layer/head as the center figure, but with the Sentence A → Sentence B filter selected.
Grid of six BertViz plots showing GPT-2 attention from the pronouns 'She' and 'He' to gender-specific terms, names, and occupation words, with 'nurse' and 'doctor' receiving different attention weight depending on the pronoun
Figure 4 from Vig (2019): attention pattern in GPT-2 related to coreference resolution suggests the model may encode gender bias.

Always name the layer, head, query token, and aggregation method. Averaging across heads can erase specialized behavior, while selecting one striking head can exaggerate it.

Use attention visualization to inspect routing, then bring in causal or attribution evidence for prediction claims.

Checkpoint

Why is an attention heatmap incomplete without a layer and head label?