Attention Visualization › Visualizing Attention Patterns

Raw numbers are difficult to interpret. Heatmap visualizations of the attention matrix make patterns visible at a glance, and the choice of visualization format significantly affects what patterns you notice and what you miss.

The most common visualization maps the attention matrix to a color grid where rows correspond to query tokens and columns to key tokens. Cell $(i, j)$ is colored proportionally to the attention weight that token $i$ places on token $j$. Reading across a row tells you the full attention distribution for a single query token: where does this token "look" when updating its representation? Reading down a column tells you which queries attend heavily to this key token: which tokens find this position useful?

Patterns often visible in these heatmaps include diagonal bands (each token attending mostly to itself), strong attention toward the [CLS] token or punctuation tokens, and occasional long-range connections between semantically related tokens. The fact that certain patterns recur across different sentences is evidence that a head has learned a general function rather than fitting a single example.

The color scale in an attention heatmap matters more than it might seem. By default, matplotlib normalizes the colormap to span the range of values present in the matrix. If one cell has an attention weight of 0.9 and all others are below 0.1, the single dominant cell will appear very dark and everything else will appear nearly identical shades of light blue. This can make diffuse patterns look more uniform than they are.

Two useful alternatives are: (1) fixing the color scale from 0 to 1 (vmin=0, vmax=1), which puts all heads on a common scale and makes sparse versus dense patterns visually distinct, and (2) using a logarithmic color scale, which amplifies small differences in low-weight cells and can reveal secondary patterns that would otherwise be washed out. Each choice tells a different story, so it is worth producing both when exploring a new head.

Because each head learns an independent set of query, key, and value projections, different heads in the same layer often develop distinct attention patterns. Visualizing multiple heads side by side reveals this specialization and demonstrates that the layer is computing multiple things simultaneously.

Looking across heads, you can see that attention is not monolithic. Some heads produce sharp, focused patterns while others spread attention diffusely. This diversity is intentional: the multi-head mechanism exists precisely to allow different heads to capture different types of relationships simultaneously. If all heads learned the same pattern, most of the model's parameter budget for attention would be wasted. The fact that heads specialize is evidence that the training objective is pushing each head toward a distinct functional niche, even though no explicit diversity loss is imposed during training.

A tempting operation is to average all attention matrices in a layer into a single summary heatmap. This produces a view of "what the layer attends to on average." However, averaged attention is almost always more diffuse than any individual head and often has lower information content than the most specialized heads. The averaging smears together patterns that operate at very different scales and for different purposes.

A better aggregation strategy for identifying important tokens is to take the maximum attention weight received by each key token across all heads. This shows you which tokens any head found relevant, rather than which tokens all heads found moderately relevant. For understanding model structure, examining individual heads and comparing them directly produces deeper insight than any aggregate.