Attention Visualization and Interpretation

Understanding where a Transformer model "looks" when it makes a decision is crucial for building trust, diagnosing failures, and advancing model design. By visualizing attention patterns, you move beyond the black box, gaining direct insight into the linguistic and structural relationships the model has learned. This guide will equip you with the core techniques and critical perspective needed to effectively interpret attention in models like BERT and GPT.

The Building Block: Visualizing Attention Heads

At its core, an attention mechanism computes a weighted sum of values, where the weights (attention scores) signify the importance of each input element relative to a given query. In a Transformer, this happens in parallel across multiple attention heads in each layer, with each head potentially learning a different type of relationship. The raw output of a single attention head is an attention matrix, where each row shows how much a specific token (the query) attends to all other tokens (the keys).

Visualizing this matrix for a given input sentence is the first step. For a sentence like "The cat sat on the mat," you might see that the head governing verb-direct object relationships causes "sat" to attend strongly to "cat" and "mat." Another head handling prepositions might show "on" attending strongly to "mat." These individual visualizations, often as heatmaps, reveal the specialized roles heads can develop. However, a model's prediction is the result of a complex cascade of these operations across all layers, requiring methods to synthesize this information.

Aggregating Information with Attention Rollout

To understand the flow of information from input to output, you need to see how attention propagates through the network's depth. Attention rollout is a simple yet powerful algorithm for aggregating multi-layer attention. It provides a consolidated view of which input tokens most influenced a given output token by mathematically combining attention weights across all layers.

The process is iterative. Start with an identity matrix, which represents each token attending to itself. Then, for each Transformer layer, you average the attention weights across all heads in that layer to get a single attention matrix $A_l$. To incorporate this layer's effect, you multiply the current aggregated attention matrix by $A_l$ (using matrix multiplication). This sequence of multiplications, from the first layer to the last, gradually mixes information. The final matrix shows the aggregated attention from the output layer back to the input. For an output token at position $i$, the row $R_i$ in the final matrix indicates how much of its representation is derived from each input token, providing a clean, global view of information pathways.

Interactive Exploration with BertViz

While static heatmaps are useful, interactive tools dramatically enhance exploration. BertViz is an open-source toolkit designed for this purpose. It allows you to load a model like BERT or GPT-2 and visualize attention in three primary views:

1. Head View: Shows the attention patterns for one or more specific heads in a selected layer, exactly like the heatmaps described earlier.
2. Model View: Provides a birds-eye perspective of attention across all heads and all layers simultaneously. You can quickly scan to see which layers and heads are most active for your input.
3. Neuron View: (For GPT-2 style models) Visualizes how attention weights are computed by decomposing them into contributions from individual query, key, and value vectors.

Using BertViz, you can hover over attention lines to see weight values, click to isolate specific heads, and dynamically change the input text. This interactivity is invaluable for forming hypotheses about head functionality—for example, by testing if a specific head consistently attends from pronouns to their antecedents across different sentences.

The Critical Limitation: Attention is Not Explanation

A crucial pitfall in interpretability is conflating attention weights with causal explanation. High attention from token A to token B does not prove that B was the reason for the model's prediction. It shows association, not causation.

Several issues underlie this limitation. First, attention weights are a poor measure of feature importance. The values are normalized (via softmax) across a sequence, so a "high" weight might simply be the least low among uninformative options. Second, there are many equivalent attention patterns. Different sets of attention weights can produce the same final aggregated vector after the value-weighted sum. The model is learning a representation, not necessarily providing human-readable justifications. Therefore, while attention is a fascinating signal about model behavior, it should be treated as one piece of evidence, not the final verdict. Relying on it alone for explanation can lead to incorrect conclusions about how the model works.

Probing Learned Representations

If attention shows where the model looks, we need complementary tools to understand what it has learned. This is where probing classifiers come in. A probe is a simple supervised model (like a linear classifier) trained to predict a specific linguistic property (e.g., part-of-speech tags, syntactic depth, semantic sentiment) from a model's internal representations (e.g., the output vectors of a specific layer).

The logic is diagnostic: if a simple classifier can easily learn to predict a property from a given representation, then that property is likely linearly encoded within it. For instance, you might find that the output of BERT's 8th layer is highly predictive of syntactic tree depth, while the final layer's [CLS] token representation is excellent for coreference resolution. Probes don't show causal use, but they powerfully characterize the information content available at different stages of the network, helping you map functions to specific layers and heads identified in your attention visualizations.

Applying Interpretability for Model Debugging

The ultimate practical value of these techniques lies in diagnosing and fixing model errors. Interpretability for model debugging transforms abstract failures into concrete engineering tasks.

Imagine your sentiment analysis model misclassifies "The movie was unexpectedly not boring" as negative. Attention visualization might reveal that the [CLS] token at the output attends overwhelmingly to "not boring," which is correct, but the final classification is wrong. A probing classifier on the [CLS] vector might show it has learned a useful representation of sentiment. This discrepancy directs you to investigate the final classification head itself—perhaps it is under-trained or has a simplicity bias. Alternatively, attention rollout might show that in earlier layers, information from "unexpectedly" is lost, preventing the model from modulating the negation. This insight could lead to architectural adjustments, such as adding skip connections, or data augmentation with more complex negation examples. By triangulating evidence from attention, probes, and error analysis, you can move from seeing that a model is wrong to understanding why it is wrong and how to fix it.

Common Pitfalls

1. Treating attention weight magnitude as feature importance: A weight of 0.6 in one context may be insignificant, while a weight of 0.3 in another may be critical. Always consider the distribution of weights in the softmax output and the content of the value vectors being summed.
2. Over-interpreting a single head or layer: Transformers are highly redundant. A function you attribute to one head (e.g., subject-verb agreement) may be redundantly encoded in several others. Use model-wide views (like BertViz's Model View) and ablation studies to confirm the importance of any single component.
3. Assuming human-aligned semantics for attention heads: While some heads learn linguistically intuitive roles, many do not. They are statistical optimizers, not linguists. Validate your hypotheses about a head's function by testing it across a diverse set of inputs, not just a few convincing examples.
4. Neglecting the value vectors: Attention has three components: Query, Key, and Value. The output is a weighted sum of the value vectors. Two identical attention patterns can produce drastically different outputs if their value vectors differ. Visualization shows the weights, but the actual information retrieved is in the values.

Summary

• Attention visualization reveals the relational focus of a Transformer model at the granular level of heads and layers, with tools like attention rollout providing aggregated, global views of information flow.
• Interactive tools like BertViz are essential for dynamically exploring the complex, multi-dimensional space of attention patterns.
• Attention is not a direct explanation; it indicates association, not causation, and should be interpreted alongside other evidence.
• Probing classifiers are a key complementary technique that diagnoses what linguistic or task-specific information is encoded within the model's learned representations.
• The combined use of visualization, probing, and error analysis forms a powerful methodology for model debugging, turning inscrutable failures into actionable engineering insights.