Quantifying Context Mixing in Transformers (2023)

Abstract: Self-attention weights and their transformed variants have been the main source of information for analyzing token-to-token interactions in Transformer-based models. But despite their ease of interpretation, these weights are not faithful to the models’ decisions as they are only one part of an encoder, and other components in the encoder layer can have considerable impact on information mixing in the output representations. In this work, by expanding the scope of analysis to the whole encoder block, we propose Value Zeroing, a novel context mixing score customized for Transformers that provides us with a deeper understanding of how information is mixed at each encoder layer. We demonstrate the superiority of our context mixing score over other analysis methods through a series of complementary evaluations with different viewpoints based on linguistically informed rationales, probing, and faithfulness analysis.

## introduction transformers @xcite , with their impressive empirical success, have become a prime choice of architecture to learn contextualized representations across a wide range of modalities, such as language @xcite , vision @xcite , visionlanguage @xcite , and speech @xcite , mainly due to their ability to utilize pairwise interactions between input tokens at every timestep. to better understand the inner dynamics of transformers, we need to trace the information flow from the input embeddings up to the output representation (including quantifying the degree of context mixing, which we will define below). the attention weights from the multi-head attention mechanisms offer a straightforward starting point for understanding this flow, and these weights ('raw attention') have been used in many studies @xcite @xcite . however, the reliability and usefulness of raw attention weights has also been questioned @xcite . in particular, attention weights tend to concentrate on uninformative tokens in the context @xcite , and removing or altering them may lead to the same and sometimes even better model performance on downstream tasks @xcite @xcite . these findings suggest that transformers do not solely rely on self-attention, and other components in the encoder block play an essential role in information mixing. a number of methods have been proposed to compute some form of 'effective attention weights', with the goal of more faithfully tracing the relative contributions of different input tokens at various layers of the transformer (as we will discuss in section 2). these methods show an improvement over raw attention, but still ignore key components of the transformer encoder block. this is a particularly crucial shortcoming, given that most of the parameter budget in a transformer encoder is spent on position-wise feed-forward networks outside of the self-attention component, which can have a considerable impact on the degree of information mixing in the output representations. in this paper, we focus on context mixing: the property of transformers that in each node, at each layer, information from the context can be incorporated into the representation of the target token. we propose value zeroing, a novel approach to quantify the contribution each context token has in determining the final representation of a target token, at each layer of a transformer. value zeroing is based on the explaining-by-removing intuition @xcite shared by many posthoc interpretability methods, but it takes advantage of a specific feature of transformers: it zeroes only the value vector of a token t when computing its importance, but leaves the key and query vectors (and thus the pattern of information flow) intact. based on extensive experiments and three complementary approaches to evaluation, we demonstrate that importance scores we can obtain with value zeroing provide better interpretations than other analysis methods. firstly, we use a set of grammatical agreement tasks from the blimp corpus @xcite as a case study. transformer-based models do extremely well on the task of distinguishing grammatical from ungrammatical sentences, and the blimp corpus provides information on the cue words that determine the difference. we find that value zeroing, unlike earlier approaches, indeed reveals that transformers make use of relevant cues. secondly, we use information-theoretic probing @xcite as an independent approach to track information flow in transformer networks. the scores we obtain with value zeroing turn out to be highly correlated with layer-wise probing performance; that is, probing accuracy is higher in layers where relevant tokens are more effectively utilized by the model. thirdly, we assess the faithfulness of our method (jacovi and goldberg, 2020); compared to alternative analysis methods, we show that value zeroing is not only more plausible and human-interpretable, but also more faithful to models' decisions. ## related work while numerous studies have leveraged the weights assigned by the self-attention mechanism to gain intuition about the information mixing process in transformers @xcite @xcite @xcite , it is still a matter of debate whether attention weights are suitable for interpreting the model (see @xcite 's study for a full discussion). thus several post-processing interpretability techniques have been proposed to convert these weights into scores that provide a more detailed interpretation of the inner workings of transformers. we review the main approaches below. @xcite propose the attention-flow and attention-rollout methods to approximate information flow in transformers based on raw attention weights. the former treats raw attention weight matrices as a flow network and returns the maximum flow through each input token. the latter recursively multiplies the attention weight matrix at each layer by the preceding ones. there is, however, an unjustified assumption in the formulation of these methods that both multihead attention and residual connections contribute equally to the computation of the output. @xcite propose a method that incorporates the norm of the transformed value vectors and report a negative correlation between these norms and raw attention weights on frequent tokens, which partially explains the insufficiency of raw attention weights for context mixing estimation. @xcite extend this method to the whole self-attention block by incorporating residual connections (res) and layer normalization (ln) (two components with significant impact on both model performance and training convergence @xcite ), but demonstrate that res and ln components largely cancel out the mixing process. @xcite 's method, however, ignores the effect of the second sublayer in a transformer's encoder. @xcite and @xcite employ a gradient-based approach for analyzing the interaction of input representations, but the gradient measures the sensitivity between two vectors and ignores the impact of the input vector. in our experiments we show that despite their relative success in explaining model decisions, gradient-based approaches are not suitable for layer-wise analysis. more recently, the effectiveness of combining these approaches has also been investigated. @xcite propose a method that uses @xcite 's scores and incorporates the effects of the second layer normalization; they aggregate those scores using rollout @xcite to provide global token attributions. in the same vein, @xcite use rollout to aggregate a variant of @xcite 's scores: instead of relying on the euclidean norm of the transformed vector, they measure manhattan distance of each transformed vector to the context vector outputted from self-attention block. in both studies, however, the fact that these context vectors might undergo significant changes after passing through the second sublayer in the encoder layer is not taken into account. we will show (in section 7) that even at a global level, our scores provide better interpretation than prior methods. ## our proposed method to remedy for the limited scope of the existing methods, we introduce a new context mixing score that takes into account all components in a transformer encoder block. ## evaluation 1: cue alignment as cue words are the only indicators of the true labels in our dataset, we expect that when the model performs well, it overwhelmingly depends on these words to form the representation of a [mask] to-ken in a given context. to quantify the alignment between a context mixing score and the cue word, we first define a binary cue vector ξ according to the following condition: then we compare the cue vector and the prediction of a context mixing score s in two different ways: dot product. we quantify cue alignment as s • ξ, which measures the total score mass the model assigns to cue words to form the representation of the target token. ## evaluation 2: context mixing versus probing in this section, we investigate the relationship between cue word alignment and probing performance across layers. we hypothesize that if a layer aligns better with the cue word according to a reliable context mixing score, then the representation of the masked token on that layer can be used more effectively by a probing classifier to decode number agreement with the cue word. to verify our hypothesis, we obtain the representation of masked tokens in test examples across all layers. since all examples in our dataset share the same number agreement property, we asso1 2 3 4 5 6 7 8 9 10 11 ciate each masked representation with a singular or plural label. next, we perform an informationtheoretic probing analysis using minimum description length (mdl) to measures the degree to which representations encode number agreement. we chose mdl as our probe since it is theoretically justified and has been shown to provide more reliable results than conventional probes @xcite . to compute mdl, we employed the online coding of @xcite . since mdl can be affected by the number of data points (n), we measure compression as our evaluation metric which is defined as follows: figure 3 reports compression of probing classifiers based on representations obtained from both pre-trained and fine-tuned models across all layers. we also include the results for the embedding layer of the model (layer 0) which can serve as a non-contextualized baseline. we can see a jump in probing performance at layers 4 (1.52 → 1.72) and 9 (2.03 → 2.45) in the fine-tuned setup, the same layers for which we found a higher alignment with cue words in figures 1 and 2 . method ρ pt ρ ft rand -0.07 -0.02 attn 0.10 0.08 attn-norm 0.52 0.56 attn-norm + res -0.35 -0.24 attn-norm + res + ln 0.12 0.17 globenc¬ -0.01 -0.12 alti¬ -0.01 -0.12 gradxinput -0.96 -0.99 ig -0.86 -0.77 dl -0.97 -1.00 value zeroing 0.65 0.64 recall that based on figures 1 and 2 , according to the gradient-based methods the masked tokens pay more attention to the cue words only in earlier layers. however, we can see a highly negative correlation with probing results for these scores. due to the nature of the task the probing score goes up monotonically along the layers. at the same time, the gradient attribution score goes up monotonically as you get closer to the bottom embedding layers, suggesting that gradient-based methods are unreliable for layer-wise analysis and identifying important tokens in the context mixing process. ## evaluation 3: faithfulness analysis our experimental results in sections 5 and 6 show that the value zeroing score matches our prior linguistically-informed expectations. however, it is not always clear whether a plausible context mixing score that matches human expectations is also faithful to the model and reflects its decision making process @xcite @xcite . in this section we employ the notion of input ablation @xcite to evaluate the faithfulness of our context mixing score. the influence of a target token on a model's decision is often estimated as the drop in the model's predicted probability of the correct class after blanking out the target token from the input. a higher drop for an ablated token indicates that the token is more influential on the model's decision @xcite @xcite . we use this blank-out approach as a base for analyzing and comparing the faithfulness of context mixing scores. to estimate the blank-out scores in bert, we calculate the probability of its output y using a softmax function normalized over only the corresponding logit values of target t and foil words (cf. qualitative analysis. we also take a closer look at the aggregated scores for a qualitative comparison. in figure 4 , we illustrate different scores obtained from a fine-tuned bert model for a correctly classified example, where the model is asked to fill the masked token with one of the verbs were or is as target and foil classes, respectively. according to value zeroing scores, the model mainly relies on the main subject (pictures) as a cue word to form a contextualized representation of the [mask] token, while the word pictures is also important for the model's final decision based on overall, our faithfulness evaluation and qualitative analysis suggest that value zeroing can explain model decisions at a global level when it is aggregated across layers. the context mixing maps per layer are provided in appendix a.4, where some more meaningful patterns can be found in value zeroing scores (in both layer-wise and aggregated setups) in contrast to other context mixing scores. ## discussion although some desiderata such as plausibility @xcite and faithfulness (lakkaraju et al., 2019; jacovi and goldberg, 2020) are taken into account when developing explanation and analysis methods, evaluating them is still a challenge due to lack of a standard ground truth. evaluating context mixing scores, where token-totoken interactions in a context are also considered, is even more challenging. several studies have used gradient-based scores as an anchor of faithfulness, and measure how strongly context mixing scores correlate with them (jain and wallace, 2019; abnar and zuidema, 2020; modarressi et al., 2022). however, the reliability of gradient-based scores can be questioned, especially when different variations of them show considerable disagreement @xcite @xcite . thus, we suggest using controlled tasks for which we have strong prior expectations for evaluating these methods. in our study, we use a set of number agreement tasks to provide such priors, since the cue words are the only sources of information in the context for performing well in the task. another point worth discussing is the concern raised by @xcite that bert tends to preserve token representations rather than mixing them at each layer. we argue that their observation is due to the context-mixing ratio they defined by comparing the norm of residual effects against other token representations. in our view, this ratio is dominated by residuals and neglects the fact that a token representation carried by residual connections is indeed a contextualized representation outputted from previous layers. we keep the residuals intact within the encoder layer by zeroing only the value vectors and focusing on the context mixing performed by all tokens. ## conclusion in this paper, we propose value zeroing as a novel approach for quantifying the information mixing process in transformers to address the shortcomings of previous methods. we performed extensive complementary experiments and showed that our method outperforms others in three different evaluation setups. since our approach requires no supervision, it could be an interesting option for improving model efficiency by removing token representations across layers.