Influence Patterns for Explaining Information Flow in BERT

1 Introduction . Previous works show that transformer models such as BERT [ 7 ] encode various linguistic concepts [ 24 , 40 , 13 ] , some of which can be associated with internal components of each layer , such as internal embeddings or attention weights [ 24 , 40 , 15 , 33 ] . However , exactly how information flows through a transformer from input tokens to the output predictions remains an open question . Recent attempts to answer this question include using attention-based methods , where attention weights are used as indicators of flow of information [ 5 , 2 , 49 ] , or layer-based approaches [ 15 , 14 , 33 ] , which identify important network units in each layer . In this paper , we examine the information flow question through an alternative lens of gradient-based attribution methods . We introduce influence patterns , abstractions of sets of gradient-based paths through a transformer ’ s entire computational graph . We also introduce a greedy search procedure for efficiently and effectively finding patterns representative of concept-critical information flow . Figure 1 provides an example of an influence pattern in BERT . We conduct an extensive empirical study of influence patterns for several NLP tasks : Subject-Verb Agreement ( SVA ) , Reflexive Anaphora ( RA ) , and Sentiment Analysis ( SA ) . Our findings are summarized below . • A significant portion of information flows in BERT go through skip connections and not attention heads , indicating that attention weights [ 2 ] alone are not sufficient to characterize information flow . In our experiment , we show that on average , important information flow through skip connections 3 times more often than attentions . • By visualizing the extracted patterns , we show how information flow of words interact inside the model and BERT may use grammatically incorrect cues to make predictions . • The consistency of influence patterns across instances of a task reflects BERT ’ s performance for that task . ∗Correspondence to kaijil @ andrew.cmu.edu 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . • Through ablation experiments , we find that influence patterns account for information flows in BERT on average 74 % and 25 % more accurately than prior attention-based and layer-based explanation methods [ 2 , 23 , 9 ] , respectively . 2 Background . We begin this section by introducing notations and architecture of BERT in Sec . 2.1 . We then introduce distributional influence as an axiomatic method to explain the output behavior of any deep model in Sec . 2.2 , which serves as a building block for our method to follow in the next section . 2.1 BERT . Throughout the paper we use x to denote a vector and |S| to denote the cardinality of a set S or number of nodes in a graph S. We begin with the basics of the BERT architecture [ 7 , 44 ] ( presented in Fig . 1 ) . Let L be the number of layers in BERT , H the hidden dimension of embeddings at each layer , and A the number of attention heads . The list of input word embeddings is x def= [ x1 , x2 , ... , xN ] , xi ∈ Rd . We denote the output of the l-th layer as hl1 : N . First layer inputs are h 0 1 : N def = x1 : N .We use a l , i j to denote the j-th attention head from the i-th embedding at l-th layer and slj to denote the skip connection that is “ copied ” from the input embedding from the previous layer then combined with the attention output . The output logits are denoted by y. Computation Graph of BERT A deep network can be viewed as a computational graph G def= ( V , F , E ) , a set of nodes , activation functions , and edges , respectively . In this paper , we assume the graph is directed , acyclic , and does not contain more than one edge per adjacent pair of nodes . A path p in G is a sequence of graph-adjacent nodes [ p1 , p2 , · · · , pt ] ; pt is the output node . Thus , the Jacobian passing through a path p evaluated at input x is ∏−1 i=1 ∂pi ( x ) /∂pi−1 ( x ) as per chain rule . We further denote the Jacobian of the output of node ni w.r.t the output of connected ( not necessarily directly ) predecessor node nj evaluated at x as ∂nj ( x ) /∂ni ( x ) . For computational and interpretability reasons , an ideal graph would contain as few nodes and edges as possible while exposing the structures of interest . For BERT we propose two graphs : embedding-level graph Ge corresponding to the nodes and edges shown in Fig . 1 ( left ) to explain how the influence of input embeddings flow from one layer of internal representations to another and to the eventual prediction ; and attention-level graph Ga ⊃ Ge that additionally includes attention head nodes and skip connection nodes as in Fig . 1 ( right ) , a finer decomposition to demonstrate how influence from the input embedding flows through the attention block ( or skip connections ) within each layer . 2.2 Explaining Deep Neural Networks . Gradient-based explanations [ 39 , 9 , 37 ] are well-studied in explaining the behavior of a deep model by attributing the model ’ s output to each input feature as its feature importance score . While existing approaches often explain with the most important features for a single class i , distributional Influence [ 23 ] generalizes gradient-based approaches to answer a broader set of questions , e.g . why [ MASK ] should be IS instead of ARE ( Figure 1 ) , by introducing quantity of interest ( QoI ) . Suppose a general network f : Rd → Rn , a QoI is a differentiable function q ( f ( x ) ) that outputs a scalar to incorporate the subject of an explanation . For example , to answer the aforementioned question , we can define q ( f ( x ) ) = f ( x ) IS − f ( x ) ARE where f ( x ) ∗ is the logit output of class ∗ . Formally , we introduce Distributional Influence . Definition 1 ( Distributional Influence ) Given a model f : Rd → Rn , an input x , a user-defined distribution D ( x ) , and a user-defined QoI q , Distributional Influence g ( x ; q , D ) is defined as : g ( x ; q , D ) def= Ez∼D ( x ) ∂q ( f ( z ) ) ∂z Remark 1 Distributional Influence also leverages a user-defined distribution D ( x ) to capture the network ’ s behavior in a neighborhood of the input of interest x . By introducing D ( x ) , we can prove that several popular attribution methods are specific cases of the distributional influence ; e.g . when D ( x ) is a Gaussian Distribution , g ( x ; q , D ) reduces to Smooth Gradient [ 37 ] ; when D ( x ) is a uniform distributions over a path c = { x + α ( x− xb ) , α ∈ [ 0 , 1 ] } from a user-defined baseline input xb to x , D ( x ) reduces to Integrated Gradient [ 37 ] . We use D ( x ) as the uniform distribution over a linear path described in Remark 1 in the rest of the paper because it provides several provable properties [ 39 ] to ensure the faithfulness of our explanations . We approximate the expectation in Def . 1 by sampling discrete points in the uniform distribution . 3 Tracing Influence Flow with Patterns . To explain how different concepts in the input flow to final predictions in BERT , it is important to show how the information from each input word flows through each intermediate layer and finally reaches the output embedding of interest , e.g . [ MASK ] for pretraining or [ CLS ] for fine-tuned models . Prior approaches have use the attention weights to build directed graphs from one embedding to another [ 2 , 49 ] . These approaches use heuristics to treat high attention weights as indicators of important information flow between layers . However , as more work starts to highlight the axiomatic justifications of gradient-based methods over attentions weights as an explanation approach [ 42 ] , we therefore explore an orthogonal direction in applying distributional influence in BERT to trace the information flow . Since distributional influence only attributes the output behavior over the input features , in this section , we generalize it to find important internal components that faithfully account for the output behavior . Tracing Influence by Patterns By viewing BERT as a computation graph ( Ge or Ga in 2.1 ) , we restate the problem : given a source node s and a target node t , we seek a significant pattern of nodes from s to t that shows how the influence from s traverses from node to node and finally reaches t. An exhaust way to rank all paths by the amount of influence flowing from s to t is possible in smaller networks , as is done by Lu et al . [ 27 ] . However , the similar approach lacks scalability to large models like BERT . Therefore , we propose a way to greedily narrow down the searching space from all possible paths to specific patterns . That is , our approach is two-fold : 1 ) we employ abstractions of sets of paths as the localization and influence quantification instrument ; 2 ) we discover influential patterns with a greedy search procedure that refines abstract patterns into more concrete ones , keeping the influence high . We begin with the formal definition of a pattern . Definition 2 ( Pattern ) A pattern π is a sequence of nodes [ π1 , π2 , · · · , π−1 ] such that for any pair of nodes πi , πi+1 adjacent in the sequence ( not necessarily adjacent in the graph ) , there exists a path from πi and πi+1 . A pattern π abstracts a set of paths , written γ ( π ) that follow the given sequence of nodes but are free to traverse the graph between those nodes in any way . Interpreting paths and patterns as sets , we define γ ( π ) def = { p ⊆ P : π ⊆ p } where P is the set of all paths from π1 to π−1 . If every sequence-adjacent pair of nodes is directly connected then the pattern abstracts a single path . To quantify the amount of information that flows from the input node to the target node over a particular pattern , we propose Pattern Influence , motivated by distributional influence . Definition 3 ( Pattern influence ) Given a computation graph and a user-defined distribution D , the influence of an influence pattern π , written I ( x , π ) is the total influence of all the paths abstracted by the pattern : I ( x , π ) def= ∑ p∈γ ( π ) Ez∼D ( x ) ∏−1 i=1 ∂pi ( z ) ∂pi−1 ( z ) . Proposition 1 ( Chain Rule ) I ( x , π ) = Ez∼D ( x ) ∏−1 i=1 ∂πi ( z ) ∂πi−1 ( z ) for any distribution D ( x ) . Prop . 1 ( proof in Appendix A ) simplifies the evaluation of the pattern influence by specifying the exact set of internal nodes through which influence flows in a computational graph . We hereby introduce a greedy way of finding the most influential pattern in both Ge and Ga . Guided Pattern Refinement ( GPR ) Starting with source and target nodes s and t along with a initialized pattern π0 = { s , t } representing all paths between s and t , we construct π1 by adding one2 sequence-adjacent node e0 ( there is a direct path between s and e0 ) from a guiding set E0 that maximizes the influence of the resulting pattern such that : e0 = arg max e∈E0 { I ( x , [ s , e , t ] ) , π1 = [ s , e0 , t ] , ( 1 ) This procedure is iterated until we exhaust the last guiding set . We show an example of GPR in a toy graph in Fig . 2 . For an embedding-level graph Ge , each guiding set El includes all embeddings hl1 : N at layer l. For the attention-level graph Ga , we refine on the embedding-level pattern πe by only expanding πe from addition nodes in Ga compared with Ge : we perform GPR iterations between πel , π e l+1 ∈ πe with the guiding set Ela which includes A attention heads al and the skip node sl ( Fig . 1 Right ) , until we reach the same last node in πe . The returned attention-level pattern πa thus abstracts a single path from the source to the target in Ga. As the attention-level analysis refines the embedding-level analysis , the produced attention-level pattern πa abstracts a strict subset of the paths of the attention-level graph that the embedding-level pattern πe abstracts . That is , πe ⊂ πa while γ ( πa ) ⊂ γ ( πe ) . The detailed algorithm of GPR and analysis of its optimality can be found in Appendix B.1 and B.2 .

Summary: This submission proposes a novel graph-based method to visualize the information model in transformer-based models such as BERT. The proposed method utilizes graph search algorithm to produce patterns which are generalized paths in graphs. The advantage of these patterns is that they are sparse and have fewer nodes than a conventionally defined path in the computation graph and are therefore potentially more interpretable. The authors demonstrate that such visualization can lead to interesting findings on several natural language processing tasks. Through quantitative evaluations, the authors also show that the proposed method is better than previous methods in terms of ablated accuracy and sparsity.   That being said, my main concern is that I doubt such visualization can bring much *practical values* for model development and model understanding. I detail my concerns below.