On Identifiability in Transformers

In this paper we delve deep in the Transformer architecture by investigating two of its core components : self-attention and contextual embeddings . In particular , we study the identifiability of attention weights and token embeddings , and the aggregation of context into hidden tokens . We show that , for sequences longer than the attention head dimension , attention weights are not identifiable . We propose effective attention as a complementary tool for improving explanatory interpretations based on attention . Furthermore , we show that input tokens retain to a large degree their identity across the model . We also find evidence suggesting that identity information is mainly encoded in the angle of the embeddings and gradually decreases with depth . Finally , we demonstrate strong mixing of input information in the generation of contextual embeddings by means of a novel quantification method based on gradient attribution . Overall , we show that self-attention distributions are not directly interpretable and present tools to better understand and further investigate Transformer models . 1 INTRODUCTION . In this paper we investigate neural models of language based on self-attention by concentrating on the concept of identifiability . Intuitively , identifiability refers to the ability of a model to learn stable representations . This is arguably a desirable property , as it affects the replicability and interpretability of the model ’ s predictions . Concretely , we focus on two aspects of identifiability . The first is related to structural identifiability ( Bellman & Åström , 1970 ) : the theoretical possibility ( a priori ) to learn a unique optimal parameterization of a statistical model . From this perspective , we analyze the identifiability of attention weights , what we call attention identifiability , in the self-attention components of transformers ( Vaswani et al. , 2017 ) , one of the most popular neural architectures for language encoding and decoding . We also investigate token identifiability , as the fine-grained , word-level mappings between input and output generated by the model . The role of attention as a means of recovering input-output mappings , and various types of explanatory insights , is currently the focus of much research and depends to a significant extent on both types of identifiability . We contribute the following findings to the ongoing work : With respect to attention indentifiability , in Section 3 , we show that – under mild conditions with respect to input sequence length and attention head dimension – the attention weights for a given input are not identifiable . This implies that there can be infinitely many different attention weights that yield the same output . This finding challenges the direct interpretability of attention distributions . As a supplement , we propose the concept of effective attention , a diagnostic tool that examines attention weights for model explanations by removing the weight components that do not influence the model ’ s predictions . ∗Equal contribution with authors in alphabetical order . Yang Liu initiated the transformer models study , perceived and performed the study of attention identifiability and effective attention , i.e. , Section 3 and Appendix A , and contributed to the token attribution discussions and calculations . With respect to token identifiability , in Section 4 , we devise an experimental setting where we probe the hypothesis that contextual word embeddings maintain their identity as they pass through successive layers of a transformer . This is an assumption made in much current research , which has not received a clear validation yet . Our findings give substance to this assumption , although it does not always hold in later layers . Furthermore , we show that the identity information is largely encoded in the angle of the embeddings and that it can be recovered by a nearest neighbour lookup after a learned linear mapping from hidden to input token space . In Section 5 we further investigate the contribution of all input tokens in the generation of the contextual embeddings in order to quantify the mixing of token and context information . We introduce Hidden Token Attribution , a quantification method based on gradient attribution . We find that selfattention strongly mixes context and token contributions . Token contribution decreases monotonically with depth , but the corresponding token typically remains the largest individual contributor . We also find that , despite visible effects of long term dependencies , the context aggregated into the hidden embeddings is mostly local . We notice how , remarkably , this must be an effect of learning . 2 BACKGROUND ON TRANSFORMERS . The Transformer ( Vaswani et al. , 2017 ) is currently the neural architecture of choice for natural language processing ( NLP ) . At its core it consists of several multi-head self-attention layers . In these layers , every token of the input sequence attends to all other tokens by projecting its embedding to a query , key and value vector . Formally , let Q ∈ Rds×dq be the query matrix , K ∈ Rds×dq the key matrix and V ∈ Rds×dv the value matrix , where ds is the sequence length and dq and dv the dimension of query and value vectors , respectively . The output of an attention head is given by : Attention ( Q , K , V ) = A · V with A = softmax ( QKT√ dq ) ( 1 ) The attention matrix A ∈ Rds×ds calculates , for each token in the sequence , how much the computation of the hidden embedding at this sequence position should be influenced by each of the other ( hidden ) embeddings . Self-attention is a non-local operator , which means that at any layer a token can attend to all other tokens regardless of the distance in the input . Self-attention thus produces socalled contextual word embeddings , as successive layers gradually aggregate contextual information into the embedding of the input word . We focus on a Transformer model called BERT ( Devlin et al. , 2019 ) , although our analysis can be easily extended to other models such as GPT , ( Radford et al. , 2018 ; 2019 ) RoBERTa ( Liu et al. , 2019 ) , XLNet ( Yang et al. , 2019b ) , or ALBERT ( Lan et al. , 2020 ) . BERT operates on input sequences of length ds . We denote input tokens in the sentence as xi , where i ∈ [ 1 , ... , ds ] . We use xi ∈ Rd with embedding dimension d to refer to the sum of the token- , segment- and position embeddings corresponding to the input word at position i . We denote the contextual embedding at position i and layer l as eli . Lastly , we refer to the inputs and embeddings of all sequence positions as matrices X and E , respectively , both in Rds×d . For all experiments we use the pre-trained uncased BERT-Base model as provided by Devlin et al . ( 2019 ) 1 . 3 ATTENTION IDENTIFIABILITY . We begin with the identifiability analysis of self-attention weights . Drawing an analogy with structural identifiability ( Bellman & Åström , 1970 ) , we state that the attention weights of an attention head for a given input are identifiable if they can be uniquely determined from the head ’ s output.2 We emphasize that attention weights are input dependent and not model parameters . However , their identifiability affects the interpretability of the output , i.e. , whether attention weights can provide the basis for explanatory insights on the model ’ s predictions ( cf . Jain & Wallace ( 2019 ) and Wiegreffe & Pinter ( 2019 ) ) . If attention is not identifiable , explanations based on attention may be unwarranted . 1https : //github.com/google-research/bert 2Cf . Appendix A.1 for more background on indentifiability . The output of a multi-head attention layer is the summation over each of the h single head outputs ( cf . Eq . 1 ) multiplied by the matrix H ∈ Rdv×d with reduced head dimension dv = d/h , Attention ( Q , K , V ) H = AEW VH = AT ( 2 ) where W V ∈ Rd×dv projects the embedding E into the value matrix V = EW V , and we define T = EW VH . Here , the layer and head indices are omitted for simplicity , since the proof below is valid for each individual head and layer in Transformer models . Intuitively , the head output is a linear combination of the T vectors using the attention as weighting coefficients . If the sequence length , i.e . the number of weighting coefficient , is larger than the rank of T , attention weights are not uniquely determined from the head output ; i.e. , they include free variables . In other words , some of the T rows are linear combinations of others . We now prove , by analyzing the null space dimension of T , that attention weights are not identifiable using the head or final model output . 3.1 UPPER BOUND FOR RANK ( T ) We first derive the upper bound of the rank of matrix T = EW VH . Note that rank ( ABC ) ≤ min ( rank ( A ) , rank ( B ) , rank ( C ) ) , therefore , rank ( T ) ≤ min ( rank ( E ) , rank ( W V ) , rank ( H ) ) ( 3 ) ≤ min ( ds , d , d , dv , dv , d ) = min ( ds , dv ) . The second step holds since rank ( E ) ≤ min ( ds , d ) , rank ( W V ) ≤ min ( d , dv ) and rank ( H ) ≤ min ( dv , d ) . 3.2 THE NULL SPACE OF T. The ( left ) null space LN ( T ) of T describes all vectors that are mapped to the zero vector by T : LN ( T ) = { x̃T ∈ R1×ds |x̃TT = 0 } ( 4 ) Its special property is that , for Ã = [ x̃1 , x̃2 , ... , x̃ds ] T where x̃Ti are vectors in this null space , ( A + Ã ) T = AT . ( 5 ) If the dimension of LN ( T ) is not zero , there exist infinitely many attention weights A + Ã yielding the exact same attention layer output and final model outputs . By applying the Rank-Nullity theorem , the dimension of the null space is : dim ( LN ( T ) ) = ds − rank ( T ) ≥ ds −min ( ds , dv ) = { ds − dv , if ds > dv 0 , otherwise ( 6 ) here we make use of the facts that dim ( LN ( T ) ) = dim ( N ( T T ) ) and rank ( T ) = rank ( T T ) where N ( T ) represents the null space of a matrix T . Equality holds if E , W V and H are of full rank and their matrix product does not bring further rank reductions . Hence , when the sequence length is larger than the attention head dimension ( ds > dv ) , selfattention is not unique . Furthermore , the null space dimension increases with the sequence length . In the next section we show that the attention weights are also non-identifiable ; i.e. , the non-trivial null space of T exists , even when the weights are constrained within the probability simplex .

Summary: This work studies the identifiability of attention, i.e., to what extent does the attention weight uniquely determine the output. The paper (1) formally establishes, using a rank argument, that at least for the cases where the attention head dimensions are smaller than the sources, infinite different attention weights can yield the same attention output. Based on such an observation, (2) a principled tool to extract the `effective` attention is introduced, with which several previous observations are challenged. This tool can certainly inspire future research. The paper further (3) explores how much token and context information is mixed at different layers of a transformer model, using both a carefully designed probing task and a gradient-based method.