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 studies the dynamics of attention in a task of simplified topic modeling, over the course of training for a specific model, where the context vector is the sum over words in a sentence of their embedding weighted by the exponential of the dot-product their key embedding with a global query vector, normalized. Due to the simplification of the topic modeling problem (two null-intersect sets of words: topic vs. non-topic), they consider the embeddings of the non-topic words to be fixed over the course of training for their theoretical analysis. The applicability of the theoretical result is close to zero, and a somewhat known property (e.g. in word2vec, Mikolov et al. 2013). The experimental results include two parts. One on a tiny synthetic dataset that matches the simplified topic modeling problem and serves as illustration. The other is on SST2 and SST5 (movie comments and ratings, sentiment analysis), where the results are poor (obviously, as the model is simple), e.g. yielding 79.59% on SST while the SOTA is 97.4, and BERT base is at 91.2. The analysis is interesting, but does not lead to new insights.