Attention Visualization › Attention Interpretation Caveats

Attention visualization is useful for forming hypotheses and building intuitions. It is also easy to misuse. Before acting on what you see in an attention heatmap, you should understand several well-documented limitations that have emerged from careful empirical work since 2019.

The most important caveat: high attention weight does not mean causal influence on the model's output. Attention weights tell you how information was routed, not why the model produced the output it did. A head might place high attention weight on a token for reasons entirely unrelated to the downstream prediction.

Researchers have demonstrated this directly. Jain and Wallace (2019) published "Attention is not Explanation," a study showing that in text classification tasks, alternative attention distributions that are very different from the learned distribution can produce the same predictions. If swapping the attention weights produces identical outputs, then the specific attention pattern cannot be the explanation for the output. Their experiments covered multiple NLP tasks and model types, and the finding was robust: attention weights and gradient-based feature importance measures frequently diverged.

Wiegreffe and Pinter (2019) responded with "Attention is not not Explanation," arguing that the question requires a more precise distinction. Under a different definition of "explanation" (one focused on faithfulness rather than counterfactual equivalence), attention weights can serve explanatory purposes in some settings. This debate has not been fully resolved, and it has productively pushed the field toward more careful definitions of what an "explanation" should provide.

The current consensus is pragmatic: treat attention as a hypothesis-generating tool, not a ground-truth explanation. The next chapter, Attention Analysis Limitations, covers this issue in depth with specific experiments.

Studies using gradient-based attribution methods frequently disagree with attention-based attributions. You can find sentences where attention strongly highlights certain tokens but gradients show those same tokens have minimal effect on the output probability. You can equally find tokens that gradients flag as highly influential but that receive low attention weights.

This divergence is informative rather than disqualifying. It suggests that attention and gradients are measuring different things. Attention measures information routing: which positions' representations were blended into other positions' representations. Gradients measure sensitivity: how much the output would change if the input representation at a given position were perturbed. Both views capture measurable properties of the computation, but neither alone is the complete story.

The observation that some heads look like they perform specific linguistic functions does not mean those heads are the only or even primary mechanism for those functions. When researchers ablate attention heads by setting their weights to uniform distributions, models often maintain performance far better than expected. The functions performed by the ablated heads appear to be distributed across other components or are simply not needed for the tasks tested.

This distributional redundancy means that even a correct identification of "head X tracks subject-verb relations" may not translate to "removing head X would impair subject-verb agreement." The model may just use other heads or MLP layers to compensate. Michel et al. (2019) showed that BERT models can be pruned to 50% of their heads with minimal performance loss on many tasks, and some heads could be pruned individually with zero impact. This resilience makes functional interpretation harder: if a function is not concentrated in one head, you cannot identify it by examining heads in isolation.

BertViz, the most widely used attention visualization library, displays multiple heads simultaneously by overlaying or averaging them. When you average attention across all heads, the result is often much more diffuse and less informative than any individual head. Conversely, cherry-picking the head that produces the most interpretable pattern may give a misleading impression of how the model works overall.

Averaging attention across layers is even more problematic because the semantics of attention weights change across layers. Layer 1 attention operates on raw token embeddings; layer 12 attention operates on highly transformed representations. The numerical value 0.3 at layer 1 and layer 12 cannot be compared directly. When aggregating across layers, the resulting "average attention" conflates computations that are fundamentally incomparable.

Some researchers proposed "attention rollout" (Abnar and Zuidema, 2020) as a more principled aggregation method. Attention rollout computes how information flows from the input through all layers by multiplying attention matrices across layers and accounting for residual connections. The result is an attention flow matrix that attempts to quantify how much each input token contributes to each output token's representation through the full stack of layers. Attention rollout is more informative than naive layer averaging, but it still does not account for the value projections, which shape what information is passed forward.

Attention patterns are sensitive to the specific input. The same head that tracks subject-verb relations on one sentence may show a completely different pattern on another. When you visualize attention on a single example, you are seeing one instantiation of a learned function, not a stable, general-purpose representation.

To characterize a head reliably, you need to examine it across hundreds or thousands of diverse examples and aggregate the patterns. Single-example visualization is useful for exploration but insufficient for claims about model behavior. A head that tracks coreference on three example sentences might track something entirely different on a fourth. The only way to know is to test systematically.