Attention Visualization › Attention Visualization Tools

Several libraries have been developed to make attention visualization accessible without writing custom matplotlib code. Each makes different design choices that affect what aspects of attention are easy to see.

BertViz is the most widely used library for transformer attention visualization. It provides two main views that serve different analytical purposes.

The head view shows attention as arcs connecting query tokens to key tokens. The thickness of each arc represents the attention weight. You can show a specific token to see all attention flowing to and from it, or view all tokens simultaneously. The arc representation makes it easy to follow individual attention paths but can become cluttered when many tokens have non-trivial weights.

The model view provides a compact overview of all attention patterns across all layers, allowing you to compare layers at a glance and spot which layers contain the most interesting patterns. Each layer's attention pattern is shown as a small thumbnail that you can click to expand. This overview is particularly useful for identifying which layers are worth examining closely.

BertViz outputs interactive HTML, which works in Jupyter but not in static documents. For paper figures or non-interactive reports, matplotlib heatmaps like those shown earlier are more suitable.

One limitation of BertViz's head view is that the arc representation is optimized for short sequences. With sequences of 30 or more tokens, the arcs overlap heavily and the visualization becomes difficult to read. For long sequences, a matrix heatmap scales much better visually.

The transformers-interpret library provides integrated gradients alongside attention, allowing you to compare attention-based and gradient-based attributions for the same prediction. This comparison is valuable precisely because of the caveats discussed above: seeing where attention and gradients agree (and disagree) helps calibrate how much trust to place in attention alone.

When the two methods agree on a set of important tokens, that agreement is informative: both the routing mechanism (attention) and the sensitivity measure (gradients) point to the same tokens as relevant. When they disagree, the disagreement is a signal to investigate further. It may indicate that attention is routing information through an indirect path, that a token's influence operates through its effect on other tokens rather than directly, or that the gradient signal is dominated by a different computation than what attention shows.

Beyond BertViz and transformers-interpret, a growing set of tools implements more principled attention-based interpretability methods.

Attention rollout (Abnar and Zuidema, 2020) propagates attention weights backward through all layers, accounting for residual connections. Instead of looking at a single layer's attention, you get an estimate of the full information path from each input token to each output token. The formula multiplies the augmented attention matrices (attention matrix plus identity, to represent the residual connection) layer by layer:

$$\tilde{A}^{(l)} = 0.5 \cdot A^{(l)} + 0.5 \cdot I$$ $$R = \tilde{A}^{(1)} \cdot \tilde{A}^{(2)} \cdots \tilde{A}^{(L)}$$

where:

• $A^{(l)}$: the attention matrix at layer $l$, averaged across heads
• $I$: the identity matrix, representing the residual path that skips the attention computation
• $R$: the final rollout matrix, whose entry $(i, j)$ estimates how much input token $j$ influences output position $i$ through all layers
• The $0.5$ weight splits information equally between the attended path and the residual path

Attention rollout is still an approximation (it ignores value projections and nonlinearities), but it is more principled than naive layer averaging and has been shown to correlate better with gradient-based saliency maps on several benchmarks.

For research purposes, direct extraction using the Hugging Face output_attentions=True flag, as shown in this chapter, gives you maximum flexibility. Library wrappers like BertViz are useful for interactive exploration but can obscure the underlying data structure. When you need to aggregate patterns across a dataset, run quantitative analyses, or feed attention data into downstream analyses, working directly with the raw tensors is preferable.

The raw tensor approach also makes it easy to integrate attention analysis with other interpretability methods. You can compute attention entropy for the same sentences where you also compute gradient norms, compare the two, and feed both into a downstream analysis without being constrained by what a library exposes through its API.