Source-linked AI summary
Quantifying Attention Flow in Transformers
Samira Abnar, Willem Zuidema
TL;DR
Higher Transformer layers mix information across tokens, making raw attention weights unreliable as explanation probes. The paper introduces attention rollout and attention flow to approximate input-token attention, and both correlate better with blank-out importance and input gradients than raw attention, with attention flow more reliable than rollout. These methods offer complementary diagnostic views but rely on simplifying assumptions and require receptive-field normalization for decoder application.
Problem
Higher-layer contextualization creates a token-identifiability problem, making raw attention weights unreliable indicators of input-token relevance.
Method
The paper models cross-layer information propagation with an attention graph and proposes attention rollout and attention flow to compute attention to input tokens while incorporating residual connections.
Results
Attention rollout and attention flow correlate better with blank-out scores and input gradients than raw attention, while attention flow is more reliable than attention rollout.
Takeaways & Limitations
The proposed token-attention measures provide complementary diagnostic views for visualizing and debugging how input tokens contribute to model predictions.
Takeaways & Limitations
The methods approximate information flow using simplifying assumptions; applying them to decoders requires normalization for the masked receptive field.
Abstract
from arXiv · showhide
In the Transformer model, "self-attention" combines information from attended embeddings into the representation of the focal embedding in the next layer. Thus, across layers of the Transformer, information originating from different tokens gets increasingly mixed. This makes attention weights unreliable as explanations probes. In this paper, we consider the problem of quantifying this flow of information through self-attention. We propose two methods for approximating the attention to input tokens given attention weights, attention rollout and attention flow, as post hoc methods when we use attention weights as the relative relevance of the input tokens. We show that these methods give complementary views on the flow of information, and compared to raw attention, both yield higher correlations with importance scores of input tokens obtained using an ablation method and input gradients.
1 Introduction
The paper examines why raw attention weights become less informative for interpreting higher Transformer layers and proposes methods that track token-level information flow across layers.
- Attention weights are widely visualized to interpret neural sequence models, although attention should not automatically be equated with explanation.
- Higher-layer embeddings lose token identifiability as information becomes increasingly contextualized, complicating interpretations based on raw attention.
- The paper proposes attention rollout and attention flow to compute attention scores to input tokens using attention patterns from the current and preceding layers.
- Both methods model information flow as a directed acyclic graph whose nodes are input tokens and hidden embeddings, with attention weights as edge weights.
- The proposed token-attention measures are intended as diagnostic visualization and debugging tools rather than methods for improving hidden-state identifiability or model performance.
2 Setups and Problem Statement
The study evaluates raw attention and proposed alternatives on Transformer verb-number prediction, using ablation-based importance and gradients as comparison signals. Raw attention is often weakly aligned with token importance beyond early layers.
- The analysis uses verb-number prediction on a subject-verb agreement dataset, where subject nouns are relevant and attractor nouns can mislead the model.
- The trained six-layer, eight-head Transformer encoder achieves 0.96 accuracy on the subject-verb agreement task.
- Raw attention patterns distinguish positions mainly in the first couple of layers, becoming relatively uniform in higher layers for a correctly classified example.
- Blank-out estimates token importance by replacing each input token with UNK and measuring the change in predicted probability for the correct class.
- Raw final-layer CLS attention has rather low Spearman rank correlation with blank-out importance, except at the first layer, and shows the same issue against input gradients.
3 Attention Rollout and Attention Flow
Attention rollout and attention flow estimate token-level attention across Transformer layers by modeling information propagation through an attention graph. They incorporate residual connections and make different assumptions about how lower-layer attention contributes to higher-layer flow.
- Method comparison: The two methods differ in their assumptions about lower-layer effects and whether token attentions are computed relative to one another or independently.Both methods recursively compute token attentions for each layer from embedding attentions.
- Residual connections: Residual connections are represented by adding the identity matrix to each attention matrix and renormalizing, yielding A = 0.5Watt + 0.5I.This accounts for information carried between corresponding positions across layers.
- Attention rollout: Attention rollout recursively multiplies attention matrices across lower layers to estimate how information propagates from higher-layer embeddings to input tokens.It sums contributions across paths by multiplying the edge weights along each path.
- Attention flow: Attention flow models the attention graph as a flow network and uses maximum-flow values to approximate attention from layer nodes to input tokens.Its path weight is determined by the minimum edge weight rather than the product of edge weights.
- Complexity: Both methods are computable in polynomial time: O(d*n^2) for rollout and O(d^2*n^4) for flow, with d denoting depth and n the number of tokens.These complexity bounds characterize the computational cost of the two proposed procedures.
4 Analysis and Discussion
Across examples and models, attention rollout and attention flow produce more informative views of token importance than raw attention, while offering different levels of focus and averaging. Their patterns align better with model predictions and importance proxies, though attention flow is reported as more reliable overall.
- Attention patterns: Raw attention becomes almost uniform above layer three, whereas rollout and flow retain more distinctive higher-layer patterns.The comparison covers CLS attention across six layers and multiple examples.
- Attention patterns: In correctly classified agreement examples, rollout and flow highlight the subject and attractor; in a miss-classified example, both highlight the non-subject “NNS” token.The latter attribution is presented as explaining the wrong prediction.
- Method comparison: Attention flow spreads weights across the most attended tokens, while attention rollout produces more focused patterns.The paper characterizes flow weights as amortized and rollout weights as stricter and more concentrated.
- Method comparison: Both rollout and flow correlate better with blank-out scores and input gradients than raw attention, with attention flow more reliable than rollout.The paper also reports slightly better performance than raw attention for DistillBERT on SST-2, despite very low correlations with input gradients for all three methods.
- BERT pronoun resolution: In BERT pronoun-resolution examples, rollout and flow can agree with model predictions when final-layer raw attention does not, while in another example only flow is consistent.Figure 4 compares attention from the masked pronoun to candidate references and the predicted probabilities of “his” and “her”.
5 Conclusion
The proposed attention-rollout and attention-flow methods approximate information flow using attention weights, but rely on simplifying assumptions. Their decoder use is additionally constrained by causal masking, which biases both methods toward initial tokens unless receptive-field normalization is applied.
- Attention rollout and attention flow approximate token-level information flow from embedding attentions using simple, task- and architecture-agnostic methods.The methods require only attention weights and are intended for models using self-attention.
- These approximations should be interpreted cautiously because they make simplifying assumptions about information flow in the model.
- In Transformer decoders, causal masking biases rollout and flow toward initial input tokens, so receptive-field normalization is needed before applying them.
- The attention graph can also use effective attentions or gradient-adjusted weights instead of raw attentions.
A.1 Single Head Analysis
For multihead analysis, the method avoids assuming that heads never mix information by combining lower-layer head attentions into a single-head representation. Head-specific rollout is then computed using this aggregated lower-layer rollout.
- A.1 Single Head Analysis: Treating attention heads separately may assume that information does not mix between heads.
- A.1 Single Head Analysis: To avoid that assumption, attentions from all layers below the layer of interest are summed across heads and treated as single-head layers.
- A.1 Single Head Analysis: For head k at layer i, attention rollout is computed as ˜A(i, k) = A(i, k) ¯A(i), where ¯A(i) is the single-head rollout for layer i.