On the Dynamics of Training Attention Models

1 INTRODUCTION . Attention-based neural networks have been broadly adopted in many natural language models for machine translation ( Bahdanau et al. , 2014 ; Luong et al. , 2015 ) , sentiment classification ( Wang et al. , 2016 ) , image caption generation ( Xu et al. , 2015 ) , and the unsupervised representation learning ( Devlin et al. , 2019 ) , etc . Particularly in the powerful transformers ( Vaswani et al. , 2017 ) , attention is its key ingredient . Despite its great successes established empirically , the working mechanism of attention has not been well understood ( see Section 2 ) . This paper sets up a simple text classification task and considers a basic neural network model with the most straightforward attention mechanism . We study the model ’ s training trajectory to understand why attention can attend to the discriminative words ( referred to as the topic words ) . More specifically , in this task , each sentence is treated as a bag of words , and its class label , or topic , is indicated by a topic word . The model we consider involves a basic attention mechanism , which creates weighting factors to combine the word embedding vectors into a “ context vector ” ; the context vector is then passed to a classifier . In this setting , we prove a closed-form relationship between the topic word embedding norm and the inner product of its key and the query , referred to as the “ score ” , during gradient-descent training . It is particularly remarkable that this relationship holds irrespective of the classifier architecture or configuration . This relationship suggests the existence of a “ synergy ” in the amplification of the topic word score and its word embedding ; that is , the growths of the two quantities promote each other . This , in turn , allows the topic word embedding to stand out rapidly in the context vector during training . Moreover , when the model takes a fixed linear classifier , this relationship allows rigorous proofs of this “ mutual promotion ” phenomenon and the convergence of training to the topic words . Our theoretical results and their implications are corroborated by experiments performed on a synthetic dataset and real-world datasets . Additional insights are also obtained from these experiments . For example , low-capacity classifiers tend to give stronger training signals to the attention module . The “ mutual promotion ” effect implied by the discovered relationship can also exhibit itself as “ mutual suppression ” in the early training phase . Furthermore , in the real-world datasets , where perfect delimitation of topic and non-topic words does not exist , interesting training dynamics is observed . Due to length constraints , all proofs are presented in Appendix . 2 RELATED WORKS . Since 2019 , a series of works have been published to understand the working and behaviour of attention . One focus of these works pertains to understanding whether an attention mechanism can provide meaningful explanations ( Michel et al. , 2019 ; Voita et al. , 2019 ; Jain & Wallace , 2019 ; Wiegreffe & Pinter , 2019 ; Serrano & Smith , 2020 ; Vashishth et al. , 2020 ) . Most of these works are empirical in nature , for example , by analyzing the behaviours of a well-trained attention-based model ( Clark et al. , 2019 ) , or observing the impact of altering the output weights of the attention module or pruning a few heads ( Michel et al. , 2019 ; Voita et al. , 2019 ) , or a combination of them ( Jain & Wallace , 2019 ; Vashishth et al. , 2020 ) . Apart from acquiring insights from experiments , Brunner et al . ( 2019 ) and Hahn ( 2020 ) show theoretically that the self-attention blocks lacks identifiability , where multiple weight configurations may give equally good end predictions . The non-uniqueness of the attention weights therefore makes the architecture lack interpretability . As a fully connected neural network with infinite width can be seen as a Gaussian process ( Lee et al. , 2018 ) , a few works apply this perspective to understanding attention with infinite number of heads and infinite width of the network layers ( Yang , 2019 ; Hron et al. , 2020 ) . In this paper , we restrict our study to the more realist non-asymptotic regime . 3 PROBLEM SETUP . Learning Task To obtain insights into the training dynamics of attention models , we set up a simple topic classification task . Each input sentence contains m non-topic words and one topic word indicating its topic . Note that a topic may have multiple topic words , but a sentence is assumed to include only one of them . Assume that there are J topics that correspond to the mutually exclusive topic word sets T1 , T2 , · · · , TJ . Let T = ⋃J j=1 Tj be the set of all topic words . The non-topic words are drawn from a dictionary Θ , which are assumed not to contain any topic word . The training set Ψ consists of sentence-topic pairs , where each pair ( χ , y ) is generated by ( 1 ) randomly pick a topic y ∈ { 1 , 2 , · · · , J } ( 2 ) pick a topic word from set Ty and combine it with m words drawn uniformly at random from Θ to generate the sentence ( or the bag of words ) χ . In this task , one aims to develop a classifier from the training set that predicts the topic y for a random sentence χ generated in this way . We will consider the case that |Θ| > > |T| , which implies that a topic word appears much more frequently in the sentences than a non-topic word . Attention Model For this task , we consider a simple attention mechanism similar to the one proposed by Wang et al . ( 2016 ) . Each word w is associated with two parameters : an embedding νw ∈ Rd and a key κw ∈ Rd ′ . Based on a global query q ∈ Rd′ , the context vector of sentence χ is computed by ν̄ ( χ ) = ∑ w∈χ νw exp ( qTκw ) Z ( χ ) , where Z ( χ ) = ∑ w′∈χ exp ( q Tκw′ ) . Then ν̄ ( χ ) is fed into a classifier that predicts the sentence ’ s topic in terms of a distribution over all topics.1 Denote the loss function by l ( χ , y ) . Our upcoming analysis implies this attention model , although simple , may capture plenty of insight in understanding the training of more general attention models . Problem Statement Our objective is to investigate the training dynamics , under gradient descent , of this attention model . In particular , we wish to understand if there is an intrinsic mechanism that allows the attention model to discover the topic word and accelerates training . Moreover , we wish to investigate , beyond this setup , how the model is optimized when there is no clear delimitation between topic and non-topic words , as in real-world data . 1The condition that the attention layer directly attends to the word embeddings merely serves to simplify the analysis in Section 4 but this condition is not required for most results presented in Sections 4 and 5 . More discussions are given in Appendix A in this regard . 4 THEORETICAL ANALYSIS . It is common to fix some parameters when we train a model with limited resources . Also Lemma 1 . Assume q 6= 0 when initialized . Fixing it does not affect the attention block ’ s capacity . Thus , our upcoming discussion focuses on the case in which the query is fixed . Doing so also allows us to establish a closed-form expression connecting the word ’ s embedding and the inner product of its key and the query . In Appendix B , extra discussions and experimental results reveal that the trainability of the query does not affect the fundamental relationship we are about to present . For a topic word t , let Ψt denote the training samples involving it . Then , by gradient descent , ∆νt = τ |Ψ| ∑ ( χ , y ) ∈Ψt ∇ν̄ ( χ ) l ( χ , y ) exp ( qTκt ) Z ( χ ) ( 1 ) ∆κt = τ |Ψ| ∑ ( χ , y ) ∈Ψt q ( νt − ν̄ ( χ ) ) T ∇ν̄ ( χ ) l ( χ , y ) exp ( qTκt ) Z ( χ ) , ( 2 ) where τ denote the learning rate . As it will turn out , an important quantity in this setting is the inner product qT kw of query q and the key kw , which we denote by sw , and refer to it as the score of the word w. Denoting vw = ||q||2νw , η = τ ||q||2 , v̄ ( χ ) = ∑ w∈χ exp ( sw ) Z vw , and h ( v̄ ( χ ) ; y ) = ∇ν̄ ( χ ) l ( χ , y ) , for a topic word t , the dynamics simplifies to ∆vt = η |Ψ| ∑ ( χ , y ) ∈Ψt h ( v̄ ( χ ) ; y ) exp ( st ) Z ( χ ) ( 3 ) ∆st = η |Ψ| ∑ ( χ , y ) ∈Ψt ( vt − v̄ ( χ ) ) T h ( v̄ ( χ ) ; y ) exp ( st ) Z ( χ ) . ( 4 ) In the rest of the paper , whenever we refer to the embedding of word t , we actually mean vt not νt . Our analysis assumes the word embeddings are sampled i.i.d . from a distribution with mean zero and variance σ 2 d , where σ 2 is assumed close to zero . The word keys and the query are also sampled from zero mean distributions with a possibly different variance . We assume that this variance is so small that the initial word scores are approximately zero . This assumption of the initial configurations corresponds to the attention model starting as a word-averaging model , and allows us to investigate how the model deviates from this initial setting with training . We also assume the derivative h ( v̄ ( χ ) ; y ) of ` is Lipschitz continuous in v̄ ( χ ) throughout training . Further the assumption in Section 3 that the number of non-topic words |Θ| is much larger than the number of topic words |T| implies that with a sufficient number of training samples , the occurrence rate of a topic word is significantly higher than the non-topic ones . This then justifies the following assumption we will use throughout our analysis . Assumption 1 . The scores and the embeddings of the non-topic words are nearly unchanged compared to their counterparts for the topic words . Hence , our upcoming analysis will treat the scores and embeddings of the non-topic words as constants . Assumption 1 will be validated by experimental results presented in Section 5 . By selecting a sufficiently small η , we can take the gradient-descent updates in Eq ( 3 ) and Eq ( 4 ) to its continuous-time limit and get2 dvt dt = η |Ψ| ∑ ( χ , y ) ∈Ψt h ( v̄ ( χ ) ; y ) exp ( st ) Z ( χ ) ( 5 ) dst dt = η |Ψ| ∑ ( χ , y ) ∈Ψt ( vt − v̄ ( χ ) ) T h ( v̄ ( χ ) ; y ) exp ( st ) Z ( χ ) . ( 6 ) 2Reversely , Eq ( 3 ) is a discretized approximation of Eq ( 5 ) : vt ( t+1 ) −vt ( t ) = ∫ t+1 t dvt ( t ′ ) dt′ dt ′ ≈ 1· dvt ( t ) dt = ∆vt ( t ) . The approximation becomes accurate if vt ( t + 1 ) is close to vt ( t ) , which can be achieved by choosing a sufficiently small η . Likewise , Eq ( 4 ) is a discretized approximation of Eq ( 6 ) . We can then characterize the update of the score and the embedding of a topic word as a continuoustime dynamical system stated in Lemma 2 . The same technique has been used to analyze the training of neural networks in other contexts ( Saxe et al. , 2014 ; Greydanus et al. , 2019 ) . Lemma 2 . For sufficiently small η and σ2 , the score st and embedding vt of topic word t satisfy dvt dt = η|Ψt| |Ψ| 〈 h ( v̄ ( χ ) ; y ) exp ( st ) Z ( χ ) 〉 Ψt , ( 7 ) dst dt = [ ( vt − 〈v̄ ( χ \ t ) 〉Ψt ) T dvt dt ] 〈 exp ( st ) + Z ( χ \ t ) Z ( χ \ t ) 〉−1 Ψt , ( 8 ) where Z ( χ \ t ) = ∑w∈χ\ { t } exp ( sw ) , v̄ ( χ \ t ) = ∑w∈χ\ { t } vw exp ( sw ) Z ( χ\t ) , and 〈 · 〉Ψt denotes taking sample mean over the set Ψt . Eq ( 7 ) implies the speed of moving vt along the direction of 〈h ( v̄ ( χ ) ; y ) 〉Ψt is controlled by the attention weight exp ( st ) Z ( χ ) . Eq ( 8 ) shows that vt increases if and only if vt has a greater projection on 〈h ( v̄ ( χ ) ; y ) 〉Ψt than the weighted average of the non-topic word counterparts . Consider a simplified case where 〈h ( v̄ ( χ ) ; y ) 〉Ψt is fixed . Since the change of vt is much faster than the non-topic word counterparts , vt will have a larger projection on 〈h ( v̄ ( χ ) ; y ) 〉Ψt after a few epochs of training . Then st increases as well as its attention weight , which in turn speeds up the extension of the embedding vt . This observation reveals a mutual enhancement effect between the score increment and the embedding elongation . In fact such an effect exists in general , as stated in the theorem below , irrespective of whether 〈h ( v̄ ( χ ) ; y ) 〉Ψt is fixed . Theorem 1 . In the setting of Lemma 2 , from epoch t0 to t1 , the topic word score st and its embedding vt satisfy [ st ( t ) + exp ( st ( t ) ) 〈 1 Z ( χ \ t ) 〉 Ψt ] t1 t0 = [ 1 2 ||vt ( t ) − 〈v̄ ( χ \ t ) 〉Ψt || 2 2 ] t1 t0 . ( 9 ) Following from Lemma 2 , this theorem implies a positive relationship between the topic word score and the distance between vt and the non-topic word embedding average 〈v̄ ( χ \ t ) 〉Ψt . Remarkably this result makes no reference to 〈h ( v̄ ( χ ) ; y ) 〉Ψt , hence independent of it . This implies the identity in Eq ( 9 ) holds irrespective of the choice and setting of the classifier . Theorem 1 further implies a score and embedding norm ( “ SEN ” in short ) relationship for the topic words : Corollary 1 . In the context of Theorem 1 , by setting t0 = 0 and t1 = t , Eq ( 9 ) is reduced to ||vt ( t ) ||2 = √ 2 ( st ( t ) + exp st ( t ) m − 1 m ) , ( 10 ) The corollary indicates that ||vt ( t ) ||2 is monotonically increasing with st ( t ) . So , st increases if and only if the point vt departs from its initial location . That is , if the norm of the topic word embedding increases , it will be attended to . This result is independent of the configuration of all other network layers . Thus , if 〈h ( v̄ ( χ ) ; y ) 〉Ψt has a gradient field that pushes vt away from its original location , the topic word is expected to be attended to . This statement can be made precise , as in Theorem 2 , when the model uses a linear classifier . Theorem 2 . Assume the model has a fixed classifier in the form c ( v̄ ( χ ) ) = softmax ( UT v̄ ( χ ) ) , where the columns of U are linearly independent , and the model is trained using gradient descent with the cross-entropy loss . As training proceeds , the model will attend to the topic word in every input sentence and have its training loss approach zero . It is notable that the theorem holds broadly for any arbitrary fixed linear classifier ( subjective to the mild linear independence constraint of its parameter U ) . Additionally , we anticipate that this result holds for a much wider family of classifiers including trainable and even nonlinear ones . But rigorous proof appears difficult to obtain in such settings , and we will corroborate this claim in an experimental study in Section 5 . To sum up , in this section , we have shown two main results : ( a ) there is a closed-form positive relationship , the SEN relationship , between the topic word score and its embedding norm , which is independent of the configuration of the classifier . ( b ) the model , equipped with a fixed linear classifier stated in Theorem 2 , can be trained to have all topic words attended to .

Summary: This paper aims to prove and illustrate that attention components are defined during training by gradients that mutually amplify the embedding and score associated with crucial features. In particular, a word embedding with a high magnitude increases the gradient following the attention score for the same word, while a high attention score increases the gradient directed at the word's embedding. In addition to a proof that treats behavior during training as a dynamical system under a large suite of assumptions, they test the analytic predictions on a synthetic dataset following the same suite of assumptions. They then test on a natural language data set and discuss where it diverges from the analytic and synthetic findings,  concluding that the difference is a result of competition between different words associated with a label.