Quantifying Attention Flow in Transformers (2020)

Abstract: 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.

## introduction attention @xcite has become the key building block of neural sequence processing models, and visualizing attention weights is the easiest and most popular approach to interpret a model's decisions and to gain insights about its internals @xcite @xcite @xcite @xcite @xcite . although it is wrong to equate attention with explanation @xcite , it can offer plausible and meaningful interpretations @xcite @xcite . in this paper, we focus on problems arising when we move to the higher layers of a model, due to lack of token identifiability of the embeddings in higher layers @xcite . we propose two simple but effective methods to compute attention scores to input tokens (i.e., token attention) at each layer, by taking raw attentions (i.e., embedding attention) of that layer as well as those from the precedent layers. these methods are based on modelling the information flow in the network with a dag (directed acyclic graph), in which the nodes are input tokens and hidden embeddings, edges are the attentions from the nodes in each layer to those in the previous layer, and the weights of the edges are the attention weights. the first method, attention rollout, assumes that the identities of input tokens are linearly combined through the layers based on the attention weights. to adjust attention weights, it rolls out the weights to capture the propagation of information from input tokens to intermediate hidden embeddings. the second method, attention flow, considers the attention graph as a flow network. using a maximum flow algorithm, it computes maximum flow values, from hidden embeddings (sources) to input tokens (sinks). in both methods, we take the residual connection in the network into account to better model the connections between input tokens and hidden embedding. we show that compared to raw attention, the token attentions from attention rollout and attention flow have higher correlations with the importance scores obtained from input gradients as well as blank-out, an input ablation based attribution method. furthermore, we visualize the token attention weights and demonstrate that they are better approximations of how input tokens contribute to a predicted output, compared to raw attention. it is noteworthy that the techniques we propose in this paper, are not toward making hidden embeddings more identifiable, or providing better attention weights for better performance, but a new set of attention weights that take token identity problem into consideration and can serve as a better diagnostic tool for visualization and debugging. ## setups and problem statement in our analysis, we focus on the verb number prediction task, i.e., predicting singularity or plurality of a verb of a sentence, when the input is the sentence up to the verb position. we use the subjectverb agreement dataset @xcite . this task and dataset are convenient choices, as they offer a clear hypothesis about what part of the input is essential to get the right solution. for instance, given "the key to the cabinets" as the input, we know that attending to "key" helps the model predict singular as output while attending to "cabinets" (an agreement attractor, with the opposite number) is unhelpful. we train a transformer encoder, with gpt-2 transformer blocks as described in @xcite (without masking). the model has 6 layers, and 8 heads, with hidden/embedding size of 128. similar to bert (devlin et al., 2019) we add a cls token and use its embedding in the final layer as the input to the classifier. the accuracy of the model on the subject-verb agreement task is 0.96. to facilitate replication of our experiments we will make the implementations of the models we use and algorithms we introduce publicly available at https: //github.com/samiraabnar/attention_flow. we start by visualizing raw attention in figure 1a (like vig 2019). the example given here is correctly classified. crucially, only in the first couple of layers, there are some distinctions in the attention patterns for different positions, while in higher layers the attention weights are rather uniform. fig-ure 2 (left) gives raw attention scores of the cls token over input tokens (x-axis) at different layers (y-axis), which similarly lack an interpretable pattern.these observations reflect the fact that as we go deeper into the model, the embeddings are more contextualized and may all carry similar information. this underscores the need to track down attention weights all the way back to the input layer and is in line with findings of @xcite , who show that attention weights do not necessarily correspond to the relative importance of input tokens. to quantify the usefulness of raw attention weights, and the two alternatives that we consider in the next section, besides input gradients, we employ an input ablation method, blank-out, to estimate an importance score for each input token. blank-out replaces each token in the input, one by one, with unk and measures how much it affects the predicted probability of the correct class. we compute the spearman's rank correlation coefficient between the attention weights of the cls embedding in the final layer and the importance scores from blank-out. as shown in the first row of ## attention rollout and attention flow attention rollout and attention flow recursively compute the token attentions in each layer of a given model given the embedding attentions as input. they differ in the assumptions they make about how attention weights in lower layers affect the flow of information to the higher layers and whether to compute the token attentions relative to each other or independently. to compute how information propagates from the input layer to the embeddings in higher layers, it is crucial to take the residual connections in the model into account as well as the attention weights. in a transformer block, both self-attention and feed-forward networks are wrapped by residual connections, i.e., the input to these modules is added to their output. when we only use attention weights to approximate the flow of information in transformers, we ignore the residual connections. but these connections play a significant role in tying corresponding positions in different layers. hence, to compute attention rollout and attention flow, we augment the attention graph with extra weights to represent residual connections. given the attention module with residual connection, we compute values in layer l+1 as v l+1 = v l +w att v l , where w att is the attention matrix. thus, we have v l+1 = (w att + i)v l . so, to account for residual connections, we add an identity matrix to the attention matrix and re-normalize the weights. this results in @xmath0.5w att + 0.5i, where a is the raw attention updated by residual connections. furthermore, analyzing individual heads requires accounting for mixing of information between heads through a position-wise feed-forward network in transformer block. using attention rollout and attention flow, it is also possible to analyze each head separately. we explain in more details in appendix a.1. however, in our analysis in this paper, for simplicity, we average the attention at each layer over all heads. attention rollout attention rollout is an intuitive way of tracking down the information propagated from the input layer to the embeddings in the higher layers. given a transformer with l layers, we want to compute the attention from all positions in layer l i to all positions in layer l j , where @xmath1. in the attention graph, a path from node v at position k in l i , to node u at position m in l j , is a series of edges that connect these two nodes. if we look at the weight of each edge as the proportion of information transferred between two nodes, we can compute how much of the information at v is propagated to u through a particular path by multiplying the weights of all edges in that path. since there may be more than one path between two nodes in the attention graph, to compute the total amount of information propagated from v to u, we sum over all possible paths between these two nodes. at the implementation level, to compute the attentions from l i to l j , we recursively multiply the attention weights matrices in all the layers below. ## analysis and discussion now, we take a closer look at these three views of attention. figure 1 depicts raw attention, attention rollout and attention flow for a correctly classified example across different layers. it is noteworthy that the first layer of attention rollout and attention flow are the same, and their only difference with raw attention is the addition of residual connections. as we move to the higher layers, we see that the residual connections fade away. moreover, in contrast to raw attention, the patterns of attention rollout and attention flow become more distinctive in the higher layers. figures 2 and 3 show the weights from raw attention, attention rollout and attention flow for the cls embedding over input tokens (x-axis) in all 6 layers (y-axis) for three examples. the first example is the same as the one in figure 1 . the second example is "the article on nnp large @xmath2. the model correctly classifies this example and changing the subject of the missing verb from "article" to "articles" flips the decision of the model. the third example is "here the nns differ in that the @xmath3, which is a miss-classified example and again changing "nns" (plural noun) to "nnp" (singular proper noun) flips the decision of the model. for all cases, the raw attention weights are almost uniform above layer three (discussed before). in the case of the correctly classified example, we observe that both attention rollout and attention flow assign relatively high weights to both the subject of the verb, "article' and the attractor, "systems". for the miss-classified example, both attention rollout and attention flow assign relatively high scores to the "nns" token which is not the subject of the verb. this can explain the wrong prediction of the model. the main difference between attention rollout and attention flow is that attention flow weights are amortized among the set of most attended tokens, as expected. attention flow can indicate a set of input tokens that are important for the final decision. thus we do not get sharp distinctions among them. on the other hand, attention rollout weights are more focused compared to attention flow weights, which is sensible for the third example but not as much for the second one. furthermore, as shown in at last, to illustrate the application of attention flow and attention rollout on different tasks and different models, we examine them on two pretrained bert models. we use the models available at https://github.com/huggingface/ transformers . ## conclusion translating embedding attentions to token attentions can provide us with better explanations about models' internals. yet, we should be cautious about our interpretation of these weights, because, we are making many simplifying assumptions when we approximate information flow in a model with the attention weights. our ideas are simple and task/architecture agnostic. in this paper, we insisted on sticking with simple ideas that only require attention weights and can be easily employed in any task or architecture that uses self-attention. we should note that all our analysis in this paper is for a transformer encoder, with no casual masking. since in transformer decoder, future tokens are masked, naturally there is more attention toward initial tokens in the input sequence, and both attention rollout and attention flow will be biased toward these tokens. hence, to apply these methods on a transformer decoder, we should first normalize based on the receptive field of attention. following this work, we can build the attention graph with effective attention weights @xcite instead of raw attentions. furthermore, we can come up with a new method that adjusts the attention weights using gradient-based attribution methods @xcite .