Attention as Inference via Fenchel Duality

Attention has been widely adopted in many state-of-the-art deep learning models.1 While the significant performance improvements it brings have attracted great2 interest , attention is still poorly understood theoretically . This paper presents a new3 perspective to understand attention by showing that it can be seen as a solver of a4 family of estimation problems . In particular , we describe a convex optimization5 problem that arises in a family of estimation tasks commonly appearing in the de-6 sign of deep learning models . Rather than directly solving the convex optimization7 problem , we solve its Fenchel dual and derive a closed-form approximation of the8 optimal solution . Remarkably , the solution gives a generalized attention structure,9 and its special case is equivalent to the popular dot-product attention adopted in10 transformer networks . We show that T5 transformer has implicitly adopted the11 general form of the solution by demonstrating that this expression unifies the word12 mask and the positional encoding functions . Finally , we discuss how the proposed13 attention structures can be integrated in practical models.14 1 Introduction15 Attention-based deep neural networks are now integrated into cutting-edge language models that16 have revolutionized a broad range of tasks : machine translation [ 1 , 15 ] , sentiment classification [ 27 ] ,17 image captioning [ 29 ] and unsupervised representation learning [ 5 ] , etc . Especially , attention plays a18 pivotal role in the construction of the transformer architecture [ 25 ] , which has had a profound impact19 on the deep learning field.20 Despite great empirical success , the driving principles of attention are still poorly understood . This21 lack of understanding impedes practitioners from confidently and appropriately using attention layers22 and makes it challenging to develop new attention-based neural architectures.23 In this paper , we offer a new perspective for understanding attention by showing that it is in fact24 a solver for a certain type of optimization problem that corresponds to an inference task . We give25 several examples , all of which can be characterized as follows : given 1 ) an unreliable estimate of26 the mean of an unknown distribution p on Rd and 2 ) a preference distribution u on Rd encoding27 beliefs on p ’ s selection , the inference task is to get a better estimate of p ’ s mean given its unreliable28 estimate and u . We derive a convex optimization problem that is abstracted from the task and solve it29 by instead solving its Fenchel dual [ 22 , p.104 ] . Remarkably , the derived expression of the improved30 estimate of p gives a generalized attention structure whose special case is equivalent to the popular31 dot-product attention [ 15 ] that is also applied in the transformer network [ 25 ] . In addition , we show32 that our generalized attention expression has been implicitly adopted by T5 transformer [ 19 ] as the33 expression unifies the concept of word masks and its positional encoding functions . Extra examples34 are given to show how the generalized attention structures can be used in practice.35 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . 2 Related Works36 Since 2019 , several authors have investigated the properties and working mechanism of attention.37 This series of works mainly addresses whether the attention mechanism can serve as a proxy of38 saliency [ 9 , 18 , 23 , 24 , 26 , 28 ] . Most of these works obtain insights into the attention mechanism39 by performing empirical studies . The related methods include analyzing the behaviours of trained40 attention-based networks [ 4 ] , or pruning a few heads , or analyzing the effects of altering the attention41 weights [ 18 , 26 ] , or a mixture of these [ 9 , 24 ] .42 Apart from understanding attention empirically , some theoretical results presented by Brunner el43 al . [ 3 ] and Hahn [ 7 ] show that the self-attention layers are not identifiable . This implies there could44 exist multiple combinations of attention weights that can provide equally good final predictions.45 In particular , such non-uniqueness means that the use of attention may complicate interpretability.46 Another important approach to understand attention is to analyze its asymptotic behaviour when47 the number of heads and the network width approach infinity [ 8 , 30 ] . In this limiting case , the48 entire network can be seen as a Gaussian process [ 13 ] and its behaviours can be characterized by49 closed-form expressions that are not available in the finite case.50 Very recently ( since 2021 ) several theoretical works have appeared that study attention outside the51 asymptotic regime . Lu et al . [ 14 ] set up a simple attention-based classification model and derive a52 closed-form relationship between the word ’ s embedding norm and the product of its key and the53 query . They empirically show that such relationship also exists in a more complicated and practical54 configuration . Ramsauer et al . [ 20 ] construct an equivalence relationship between attention and55 a newly proposed Hopfield network with continuous states . In particular , they show that the new56 Hopfield network ’ s update rule is equivalent to the attention mechanism used in transformers [ 25 ] .57 3 A Motivating Example58 We first consider a seemingly unrelated example , to illustrate the key ingredients of this paper.59 Assume a probability distribution p on Rd has a spherical Gaussian prior u ∼ N ( µ , Id ) . Let hp60 denote the mean of the unknown p. Given an unreliable observation b of hp , what is the best guess61 of hp ? To solve this problem , we may formulate the following optimization problem62 p∗ = arg min p α 2 ∥∥∥∥b− ∫ ap ( a ) da∥∥∥∥2 +K ( p , u ) , ( 1 ) with α > 0 responsible for the relative strength of the two terms , where K ( p , u ) denotes the KL63 divergence between p and u . The basic idea behind ( 1 ) is that : although b is not reliable , it should64 not be too far from hp = ∫ ap ( a ) da . Also , as u encodes the preferred value of p , we add the65 KL divergence term to show preference for p that is close to u . As will be discussed later , such a66 formulation can be either obtained from the maximum likelihood principle or from the maximum67 entropy principle [ 10 , 11 ] . In particular , Rioux et al [ 21 ] develop ( 1 ) for image de-blurring by applying68 Maximum Entropy on the Mean ( MEM ) , an information-theoretic method due to Gamboa [ 6 ] but not69 yet widely known in machine learning.70 After obtaining the minimizer p∗ of ( 1 ) , its mean ∫ ap∗ ( a ) da gives our estimate of hp . Rioux et71 al . [ 21 ] prove , via Fenchel duality [ 22 , p.104 ] that the minimizer p∗ takes the form72 p∗ ( a ) = u ( a ) exp〈a , λ∗〉∫ u ( a′ ) exp〈a′ , λ∗〉 da′ , ( 2 ) where73 λ∗ = arg max λ∈Rd 〈b , λ〉 − 1 2α ‖λ‖2 − log ∫ Rd u ( a ) exp〈a , λ〉 da . ( 3 ) Note that ∫ Rd u ( a ) exp〈a , λ〉 da = exp ( 〈µ , λ〉 + 1 2 ‖λ‖ 2 ) as it is the moment generating function74 ( MGF ) of u ∼ N ( µ , Id ) . Substituting the expression into ( 3 ) followed by setting the derivative with75 respect to λ to zero yields λ∗ = αα+1 ( b − µ ) . By ( 2 ) , p ∗ ( a ) ∝ exp ( − 12 ‖a− µ‖ 2 + 〈a , λ∗〉 ) ∝76 exp ( − 12 ‖a− ( µ+ λ ∗ ) ‖2 ) . Substituting λ∗ = αα+1 ( b− µ ) into it implies that p ∗ follows a Gaussian77 distribution N ( 11+αµ+ α 1+αb , Id ) . Thus , our estimate of hp is 1 1+αµ+ α 1+αb.78 In this paper , we focus on a similar optimization problem that estimates hp assuming that u is instead79 a discrete distribution . We show that such optimization problems naturally and frequently arise in80 neural network designs . By solving the optimization problem , we derive a closed-form approximation81 for the estimate of hp , via Fenchel duality . The approximation then gives a generalized attention layer82 structure as shown in Fig 1 . A special case of it is equivalent to the familiar dot-product attention [ 15 ] 83 that is also adopted in transformers [ 25 ] . Moreover , we will show that T5 transformer [ 19 ] implicitly84 adopts our generalized attention expression.85 4 Setup of a Design Problem86 Throughout the rest of the paper , we87 consider a machine learning problem88 in which the objective is to predict an89 output quantity Y from a given input90 X . Additionally , Y may include K91 components , namely , be expressed as92 ( Y ( 1 ) , Y ( 2 ) , . . . , Y ( K ) ) . To be more93 concrete , we present a few example ma-94 chine learning problems and let them95 run through our development.96 Example : Translation Problem . In97 this problem , the input X is a sentence,98 or a sequence of words , in the source99 language , and output Y is the sequence100 of words in the target sentence , where101 Y ( k ) denotes the kth word.102 Example : Image Captioning . In this problem , the input X is a raw image and output Y is the103 sequence of words in the caption , where Y ( k ) denotes the kth word.104 Example : Filling in the blanks task . This task has been used to train the BERT model [ 5 ] . The105 input X is a sequence of words with certain percentage of words masked . The output Y are the106 predicted masked words , where Y ( k ) denotes the kth masked one.107 The objective of any of these problems and that we address in this paper is to learn a function F ,108 mapping from the space of X to the space of Y so that Y = F ( X ) . We will denote by F ( k ) the part109 of F responsible for predicting Y ( k ) ( Fig 1a ) , namely , Y ( k ) = F ( k ) ( X ) . Although we here express110 F as separate functions ( F ( 1 ) , F ( 2 ) , . . . , F ( K ) ) , we note that it is in fact possible that different F ( k ) ’ s111 share some component in common . We now focus on the design of F ( k ) .112 We restrict the architecture of F ( k ) to the form in Fig 1b with the main focus on the inference of h ( k ) .113 The extraction of feature h ( k ) is via two parallel modules f ( k ) evd and f ( k ) pref that directly operate on the114 input X followed by a function g ( k ) ( in Fig 1c ) , which we will design.115 The Design Problem We describe the problem of designing g as follows.116 Suppose that there is an unknown distribution p ( k ) on Rd whose mean vector is h ( k ) , namely,117 h ( k ) = ∫ Rd ap ( k ) ( a ) da . ( 4 ) Let u ( k ) be another distribution on Rd that is generated as the output of a network module f ( k ) pref . Here118 u ( k ) is referred to as the preference distribution , which serves as a prior guess of p ( k ) . Specifically119 u ( k ) puts non-zero probability masses on M “ template ” vectors t ( k ) 1 , t ( k ) 2 , . . . , t ( k ) M in R d , and their120 probabilities are respectively u ( k ) 1 , u ( k ) 2 , . . . , u ( k ) M ( which sum to 1 ) . Collectively , we will denote the121 set { t ( k ) 1 , t ( k ) 2 , . . . , t ( k ) M } of templates by T ( k ) .122 The preference distribution u ( k ) is considered as a good approximation of p ( k ) , in the sense that the123 support of p ( k ) is contained in the set T ( k ) of templates . Note that if Rd is the word embedding space124 for a large vocabulary , and if the size M of the template set T ( k ) is relative small , then restricting the125 support of p ( k ) to within T ( k ) imposes a strong constraint on p ( k ) .126 On the other hand , u ( k ) is not a sufficiently accurate approximation of p ( k ) , in the sense that u ( k ) may127 assign probabilities to T ( k ) somewhat differently . Such inaccuracy shifts the mean µ ( k ) of u ( k ) from128 the mean h ( k ) of p ( k ) . Suppose that there is another piece of information z ( k ) ∈ Rd that is generated129 by another network module f ( k ) evd and provides information regarding the mean shift . In particular , we130 assume that z ( k ) is a noisy version of the shift , more precisely,131 z ( k ) = h ( k ) − µ ( k ) + , ( 5 ) where ∼ N ( 0 , σ2I ) is the spherical Gaussian noise in Rd with covariance σ2I . We refer to z ( k ) as132 the evidence.133 Then the design problem is to construct a function , or a network block , g , which infers the un-134 known distribution p ( k ) and hence its mean h ( k ) based on the evidence z ( k ) and the preference135 distribution u ( k ) .136 This formulation of the design problem might seem peculiar at the first glance , but we will show137 via examples ( see Fig 2 ) that such a problem naturally arises in the construction of many machine138 learning models in practice.139 Example : Translation Problem . For the translation problem , consider the model implementation140 plotted in Fig 2a that is similar to the one proposed in [ 1 ] . We will focus on the part of the model141 responsible for inferring the kth word of the target sentence . In this model , h ( k ) corresponds to the142 constructed feature according to ( 4 ) that serves as an estimate of the context vector collecting the143 source sentence information . The estimated h ( k ) is then fed into a classifier f ( k ) out to predict the k th144 word . The preference distribution u ( k ) is generated by f ( k ) pref which takes the source sentence words as145 inputs . In particular , the support of u ( k ) consists of the source sentence word embeddings T ( called146 annotations in [ 1 ] ) which are pre-processed by two LSTM layers.1 The preference weight for each147 template depends on some positional encoding functions , which , in principle , should assign higher148 1In this model , given input X , all u ( k ) ’ s share the same support T. The superscripts of the templates are then omitted to show their independence from k. Similar comments apply to implementations of the other two running examples . weights to the templates appearing in the similar locations to the words we are inferring ( that is , h ( k ) 149 is assumed to rely on the templates near tk more heavily ) .150 Note that the inferred p ( k ) ’ s support must be a subset of u ( k ) ’ s as it is reasonable to assume that the151 target sentence words only depend on those appearing in the source sentence . Besides , although152 the preference weights specified by the positional encoding functions could provide some a priori153 information for the templates ’ weights in p ( k ) , they can not be accurate as their inferences do not154 consider the previously generated words Y ( i < t ) . This results in the mean µ ( k ) shifted from h ( k ) ,155 which is estimated by z ( k ) = f ( k ) evd that takes all the previously generated words Y ( i < t ) into account156 using another LSTM layer . Thus , h ( k ) and p ( k ) should not be far from z ( k ) + µ ( k ) and u ( k ) ,157 respectively.158 Example : Image Captioning . The caption generation model presented in Fig 2b has a similar159 architecture reported in [ 29 ] . This model shares the designs of f ( k ) evd and f ( k ) out with the translation160 model while f ( k ) pref instead extracts the templates from a raw image using a CNN network . In general,161 a word ’ s position in the caption is independent of the location of the object it describes in the image.162 Therefore , in this model , all templates extracted by the CNN share the same preference weight.163 As similar objects appear in an image would have similar features extracted by the CNN ( for example,164 a zebra and a horse ) , allowing similar templates not in T to participate in h ( k ) ’ s estimation would165 possibly mix in information not contained in the raw image and harm the word inference accuracy.166 Therefore , we could improve the estimate of h ( k ) by choosing p ( k ) similar to u ( k ) in the sense that167 p ( k ) ’ s support can not contain elements not in u ( k ) ’ s.168 Intuitively , as the generation process proceeds , the context h ( k ) should be updated to provide relevant169 information in the image to facilitate the next word inference . Such change is governed by the170 caption ’ s semantic evolution , which is captured by z ( k ) = f ( k ) evd that predicts the shift of the mean171 µ ( k ) from h ( k ) . For this reason , µ ( k ) + z ( k ) serves as an estimate of h ( k ) and should not be far away172 from it . Likewise , uk should be close to p ( k ) .173 Example : Filling in the blanks task . For the filling-in-the-blank tasks , let us consider a model174 architecture plotted in Fig 2c that is similar to the one used in BERT [ 5 ] . We focus on the inference of175 the kth masked word , which is assumed to be the jth word of the input sentence . In this model , f ( k ) pref176 and f ( k ) evd share the transformation layers ( TL ) that are commonly used in the NLP tasks to map one177 sequence of vector representations to another of the same length.2 Taking the output sequence , f ( k ) pref178 applies a linear map V to each of its elements to form T as the support of u ( k ) while the preference179 weights are specified by some positional encoding functions . At the same time , z ( k ) = f ( k ) evd estimates180 h ( k ) ’ s shift from the mean µ ( k ) due to the variation of the local information . For the same reasons181 discussed in the previous two examples , we need µ ( k ) + z ( k ) close to h ( k ) while p ( k ) is close to u ( k ) .182 Notably the formulation of the problem is based on the assumption that the network modules f ( k ) evd183 and f ( k ) pref are fixed and generate z ( k ) and u ( k ) satisfying the above assumed properties . In reality,184 f ( k ) evd and f ( k ) pref are in fact obtained via training . However , we argue that if g is made to satisfy our185 design objective , then we can at least interpret f ( k ) evd and f ( k ) pref obtained from training as serving to186 produce z ( k ) and u ( k ) with our desired properties.187 5 Formulation of an Optimization Problem188 The discussion made in the previous section implies that the key optimization problem we are about189 to focus on should ensure190 1. h ( k ) is not too far from µ ( k ) + z ( k ) , where h ( k ) is constructed by p ( k ) according to ( 4 ) and191 µ ( k ) is the mean of the preference distribution u ( k ) .192 2. p ( k ) is close to u ( k ) while p ( k ) ’ s support must be a subset of u ( k ) ’ s.193 2Typical implementation of such layers include convolution layers , recurrent layers and self-attention layers . These two desiderata prompt us to optimize:194 min p α 2 ∥∥∥∥ ( µ ( k ) + z ( k ) ) − ∫ Rd ap ( a ) da ∥∥∥∥2 +K ( p , u ( k ) ) ( 6 ) with α > 0 responsible for the relative strength of the two terms , where K ( p , u ( k ) ) denotes the KL195 divergence from p to u ( k ) . Remarkably , K ( p , u ( k ) ) has a finite value if and only if p ( k ) has non-zero196 values on the support of u ( k ) . Thus , both requirements in the second desideratum are satisfied by197 using the KL divergence as a measure for the closeness of p ( k ) and u ( k ) . Let p̃ ( k ) be the minimizer of198 ( 6 ) . The estimate of h ( k ) is199 ĥ ( k ) = ∫ Rd ap̃ ( k ) ( a ) da . ( 7 ) Naturally , this optimization problem can be derived from two different , though , related perspectives.3200 A maximum likelihood perspective . The optimization problem in ( 6 ) can be derived using the201 maximum log likelihood method by treating the KL-divergence term as a regularizer . According to202 ( 5 ) , the difference ( µ ( k ) +z ( k ) ) −h ( k ) follows a Gaussian distributionN ( 0 , σ2I ) . This implies the log203 likelihood function ` ( z ( k ) ) ∝ − 12σ2 ∥∥ ( µ ( k ) + z ( k ) ) − h ( k ) ∥∥2 . Maximizing it with the KL-divergence204 term as a regularizer is the same as minimizing205 1 2σ2 ∥∥∥ ( µ ( k ) + z ( k ) ) − h ( k ) ∥∥∥2 + ηK ( p , u ( k ) ) , ( 8 ) where η > 0 controls the strength of the regularization . Substituting ( 4 ) into ( 8 ) followed by206 rearrangement yields207 min p 1 2ησ2 ∥∥∥∥ ( µ ( k ) + z ( k ) ) − ∫ Rd ap ( a ) da ∥∥∥∥2 +K ( p , u ( k ) ) , ( 9 ) which is equivalent to ( 6 ) by setting α−1 = ησ2.208 A maximum entropy on the mean perspective Consider a problem that seeks a distribution p such209 that the expectation ∫ Rd ap ( a ) da is not far from µ ( k ) + z ( k ) . In particular , we require210 ∥∥∥∥ ( µ ( k ) + z ( k ) ) − ∫ Rd ap ( a ) da ∥∥∥∥2 ≤ 12α . ( 10 ) Note that , given z ( k ) , there are infinitely many p ’ s that satisfy the constraints , which makes it difficult211 to pick a “ best ” p for later use . A technique known in information theory as the maximum entropy on212 the mean ( MEM ) [ 6 , 21 ] solves this problem by picking the best guess of the ground truth p∗ that213 simultaneously satisfies ( 10 ) and minimizes the KL divergence to the preference distribution u ( k ) .214 That is,215 p̃ ( k ) = arg min p K ( p , u ( k ) ) subject to ∥∥∥∥ ( µ ( k ) + z ( k ) ) − ∫ Rd ap ( a ) da ∥∥∥∥2 ≤ 12α , ( 11 ) which is also the minimizer of ( 6 ) according to Equation ( 18 ) of [ 21 ] and Corollary 4.9 of [ 2 ] .216 6 Optimal Solution217 Rioux et al . proved that the optimization problem stated in ( 6 ) has the following Fenchel dual ( see218 Theorem 2 of [ 21 ] ) :219 Theorem 1 . The dual of ( 6 ) is given by220 max λ∈Rd { 〈 λ , µ ( k ) + z ( k ) 〉 − 1 2α ‖λ‖2 − log ∫ Rd u ( k ) ( a ) exp〈a , λ〉 da } . ( 12 ) Given a maximizer λ∗ of ( 12 ) , one can recover the minimizer p̃ ( k ) of ( 6 ) via221 p̃ ( k ) ( a ) = u ( a ) exp〈a , λ∗〉∫ Rd u ( a ′ ) exp〈a′ , λ∗〉 da′ . ( 13 ) 3In fact , there is also a Bayesian perspective to derive the problem , which will be discussed in Appendix A . By Theorem 1 , the estimated h ( k ) defined in ( 7 ) can be re-written as222 ĥ ( k ) = ∫ Rd a p̃ ( k ) ( a ) da = ∫ Rd a u ( k ) ( a ) exp〈a , λ∗〉∫ Rd u ( k ) ( a′ ) exp〈a′ , λ∗〉 da′ da , ( 14 ) where λ∗ is a maximizer of ( 12 ) .223 In general , λ∗ does not have a closed-form expression in terms of α , u ( k ) and z ( k ) , and a standard224 paradigm is to search for it using gradient ascent-based methods . In this paper , we will not search for225 λ∗ in this way ; instead , we will derive a closed-form expression to approximate it . Remarkably , this226 takes the form of the generalized attention presented in Fig 1.227 Note that the integration in ( 12 ) equals Eu ( k ) [ exp〈W , λ〉 ] , the expectation of the random variable228 exp〈W , λ〉 where W has the probability distribution u ( k ) . The expectation is just the moment229 generating function ( MGF ) M ( λ ) of W , and the value logM ( λ ) is called the cumulant of W [ 17,230 p.26 ] , which has an expansion [ 17 , ( 2.4 ) ] 231 logM ( λ ) = 〈 µ ( k ) , λ 〉 + 1 2 〈 λ , Σ ( k ) λ 〉 +O ( ‖λ‖3 ) , ( 15 ) where µ ( k ) = ∫ Rd a u ( k ) ( a ) da and Σ ( k ) = ∫ Rd ( a− µ ( k ) ) ( a− µ ( k ) ) T u ( k ) ( a ) da respectively232 denote the expectation and the variance-covariance matrix of W .233 Now we assume that α is small and we argue that this assumption is justified in practice . For instance,234 in the translation task , all of words in the dictionary can serve as candidate templates , which could235 be more than 10,000 , but u ( k ) reduces this size to the length of the source sentence ( usually less236 than tens of words ) . The inference of p ( k ) should strongly anchor around this prior information ; 237 consequently the information provided by z ( k ) should weigh less . On the other hand , z ( k ) can hardly238 provide an accurate estimate of the mean shift , since the generation of z ( k ) is often ignorant of the239 templates selected by u ( k ) ( for example , in the example translation and image captioning models ) or240 generated by a low-capacity module ( as in the example filling-in-the-blank model ) . For these reasons,241 one should de-emphasize the constraint imposed by z ( k ) and hence choose a small α.242 When α is picked to be small enough ( see ( 12 ) ) , the optimization of λ gets a large penalty on its L2243 norm and thus , ‖λ∗‖ is close to zero . Then , by ( 15 ) , we have244 log ∫ Rd u ( k ) ( a ) exp〈a , λ∗〉 da = logM ( λ∗ ) ≈ 〈µ ( k ) , λ∗〉+ 1 2 〈λ∗ , Σ ( k ) λ∗〉 . ( 16 ) Substituting ( 16 ) into ( 12 ) followed by setting the derivative with respect to λ to zero yields245 λ∗ = α ( Id + αΣ ( k ) ) −1z ( k ) , ( 17 ) where Id denotes the d× d identity matrix . As α is assumed close to zero , ( 17 ) is further reduced to246 λ∗ = αz ( k ) . ( 18 ) Plugging the expression into ( 14 ) gives the result stated as follows:247 Theorem 2 . For a small enough α > 0 , the estimated h ( k ) defined in ( 7 ) can be approximated by248 ĥ ( k ) = ∫ Rd a u ( k ) ( a ) exp ( α〈a , z ( k ) 〉 ) ∫ Rd u ( k ) ( a′ ) exp ( α〈a′ , z ( k ) 〉 ) da′ da . ( 19 ) For the case that u ( k ) is a discrete distribution with support { t ( k ) 1 , t ( k ) 2 , . . . , t ( k ) n } and the preference249 probability { u ( k ) 1 , u ( k ) 2 , . . . , u ( k ) n } , ( 19 ) becomes simply250 ĥ ( k ) = n∑ i=1 ti u ( k ) i exp ( α〈ti , z ( k ) 〉 ) ∑n j=1 u ( k ) j exp ( α〈tj , z ( k ) 〉 ) . ( 20 ) In Fig 3 , we set d = 2 and visualize the approximation of h ( k ) for different selections of α . We can251 observe that , as α decreases , ( 20 ) outputs a better approximation of ĥ ( k ) .252 Let α = τ− 1 2 , we rewrite Theorem 2 as follows for later reference.253 Corollary 1 . For a sufficiently large τ , the best guess of h ( k ) defined in ( 7 ) with α = τ− 12 equals254 ĥ ( k ) = n∑ i=1 t ( k ) i u ( k ) i exp ( 〈t ( k ) i , z ( k ) 〉/ √ τ ) ∑n j=1 u ( k ) j exp ( 〈t ( k ) j , z ( k ) 〉/ √ τ ) . ( 21 ) 7 Discussion255 In Section 6 , we derived an alternative expression of ĥ ( k ) defined in ( 7 ) by solving the Fenchel dual of256 the optimization problem stated in ( 6 ) . Although the expression is not in closed form , as we are only257 interested in the case when α is small , a closed-form approximation of ĥ ( k ) is derived in Theorem 2258 and reduced to the form stated in ( 21 ) when considering a discrete distribution u ( k ) .259 As we pointed out , the block g ( k ) in Fig 2a , Fig 2b and Fig 2c is expected to find the inferred p̃ ( k ) 260 minimizing ( 6 ) followed by plugging it into ( 7 ) to construct ĥ ( k ) . Therefore , one can complete261 the architecture designs of the three running examples by replacing g ( k ) with a network layer262 implementing ( 21 ) , namely , the structure in Figure 1 ( c ) .263 The relationship between the optimal solution and the attention models . Remarkably , the expres-264 sion stated in ( 21 ) gives a generalized attention block . By setting u ( k ) i = 1 n for all i , the expression265 is equivalent to the well known dot-product attention [ 15 ] , which is also applied in the transformer266 network [ 25 ] . The equivalence of the expression of ĥ ( k ) and the dot-product attention layer tells us:267 ( a ) by applying a dot-product attention layer in a model , we essentially ask the model to perform268 an optimization task defined in ( 6 ) and construct the output according to ( 7 ) . ( b ) the derivation of269 h ( k ) depends on two relatively independent pieces of information : a preference distribution given the270 global information and an estimate of the output ’ s deviation from the preference distribution ’ s mean271 according to some local information . This suggests that the design of attention-based model can be272 decomposed into two parts that respectively estimate these two values.273 The model consisting of a stack of attention layers . Although our discussion focuses on the case274 that contains a single attention layer , any attention layer L in an attention stack fits our frameworks275 ( see Fig 1 ) . In particular , all the attention layers closer to the input X than L can be grouped into the276 functions f ( k ) pref or f ( k ) evd . For those layers that take the current layer ’ s output as input , we can group277 them into f ( k ) out , where c ( k ) may contain the outputs of other attention layers working in parallel.278 T5 transformer implicitly adopts the generalized attention structure . We now show that T5279 transformer [ 19 ] can be seen as a realization of the generalized attention in ( 21 ) , where the preference280 weights u ( k ) unifies the concepts of word masks and T5 ’ s positional encoding functions . Consider281 the running example : filing in the blanks , with the preference distribution282 u ( k ) ( ti ) = { 0 if the ith word is masked exp ( bj−i ) /Z otherwise , ( 22 ) where Z is a normalizing constant and bj−i is a trainable scalar that only depends on the relative283 position of word i and word j ( which is the kth masked word that we are inferring ) . Substituting284 such u ( k ) into ( 21 ) with τ = d yields285 ĥ ( k ) = n∑ i=1 ti exp ( 〈ti , z ( k ) 〉√ d + bj−i + 1masked ( i ) ) ∑n l=1 exp ( 〈tl , z ( k ) 〉√ d + bj−l + 1masked ( l ) ) , ( 23 ) where 1masked ( i ) is an indicator function that equals −∞ if word i is masked and zero otherwise.286 The expression in ( 23 ) has the same structure as that adopted in T5 transformer , where the indicator287 function serves as the mask function to prevent the model from assigning weights to the masked288 words . In this way , the concepts of word masks and the positional encoding functions are unified289 by u ( k ) in ( 22 ) . Conversely , T5 transformer is a realization of the generalized attention with the290 preference weights u ( k ) specified in ( 22 ) .291 Generalized attention structures suggested by the optimal solution . While T5 transformer has292 implicitly adopted the generalized attention , ( 21 ) hints further generalizations could be made . For293 instance , in T5 transformer , the function outputting template ’ s preference weights only considers the294 word masks and the word ’ s relative positions . This function could be generalized to also consider the295 input sentence contexts , and the output weights encode the importance of each word before giving296 the local information stored in z ( k ) . The same idea could be applied to the image captioning example297 to replace the uniform preference weights . By adding a neural network taking the input image to298 generate non-uniform preference weights , we devise a mechanism to estimate the importance of each299 part of the image before the caption generation . In this way , the newly added network collects global300 information from the image to propose a preference distribution , which could be updated locally301 based on current generation stage encoded in z ( k ) .302 Moreover , although we mainly focus on the case when u ( k ) is discrete , we want to emphasize that the303 analysis performed in Section 6 also covers continuous u ( k ) . This hints that a continuous attention304 mechanism could also be implemented , which might prove to be useful in some applications.305 Limitations and other comments . The approximations performed in ( 15 ) and ( 18 ) have implicitly306 assumed that random variableW following distribution u ( k ) has bounded moments . For a discrete u ( k ) 307 with fixed support T = { t1 , t2 , . . . tn } , all the moments are bounded and we can always pick a small308 enough α ( or equivalently large enough scaling factor τ in Cor 1 ) to make the approximation meet309 our requirements . A concern may arise as the support T in our running examples are supplied by310 some neural layers , which could output templates of increasing norms as the training evolves . This311 problem could be alleviated by adding norm regularization or using normalized templates instead.312 8 Conclusion313 This paper presented a new perspective to understand the attention mechanism by showing that it can314 be treated as realizing a solver of a family of inference tasks . These tasks involve improving the noisy315 estimate of a distribution p ’ s mean by a preference distribution that encodes some beliefs of p ’ s value.316 We have used three running examples with the typical model architectures to show that such tasks317 naturally exist in neural network design . We then abstracted a convex optimization problem from318 these tasks and derived a closed-form approximation of the optimal solution by solving the problem ’ s319 Fenchel dual . We find that the closed-form approximation can be seen as a generalized attention320 layer and show that one of its special cases is equivalent to the dot-product attention adopted in321 transformers . We further performed an analysis on the general form and showed that T5 transformer322 implicitly adopts the generalized attention structure with attention weights unifying the concepts of323 the word masks and the positional encoding functions.324 This paper is the first work that presents a principled justification for the design of attention modules325 in neural networks . The generalized attention structure presented in this paper potentially opens a door326 to a wide design space . For example , the preference weights need not be derived from the positional327 encoding functions ; they could integrate a variety of information provided by other components of328 the network . Additionally , this research might have pointed to new ways to analyze the functioning329 of a neural network component , namely , via isolating the component from the complex network330 structure and asking : is there a “ local problem ” that is solved by the design of this component ? 331 Potential negative societal impacts . This paper presents a new perspective to understand attention332 and derived a generalized attention structure . Our work is foundational , which we believe does not333 have direct negative societal impacts . Due to the very wide range of applications of attention , such as334 self-driving [ 12 ] and healthcare [ 16 ] , our work may have unexpected negative impacts on these areas.335 References336 [ 1 ] Dzmitry Bahdanau , Kyunghyun Cho , and Yoshua Bengio . Neural Machine Translation by337 Jointly Learning to Align and Translate . In ICLR 2015 , pages 1–15 , 2014.338 [ 2 ] J. M. Borwein and A. S. Lewis . Partially finite convex programming , part i : Quasi relative339 interiors and duality theory . Mathematical Programming , 57 ( 1 ) :15–48 , 1992.340 [ 3 ] Gino Brunner , Yang Liu , Damián Pascual , Oliver Richter , Massimiliano Ciaramita , and Roger341 Wattenhofer . On Identifiability in Transformers . In International Conference on Learning342 Representations ( ICLR ) , 2019.343 [ 4 ] Kevin Clark , Urvashi Khandelwal , Omer Levy , and Christopher D. Manning . What Does BERT344 Look at ? An Analysis of BERT ’ s Attention . arXiv preprint arXiv:1906.04341 , 2019.345 [ 5 ] Jacob Devlin , Ming-Wei Chang , Kenton Lee , and Kristina Toutanova . BERT : Pre-training of346 deep bidirectional transformers for language understanding . In Proceedings of the 2019 Confer-347 ence of the North American Chapter of the Association for Computational Linguistics : Human348 Language Technologies , Volume 1 ( Long and Short Papers ) , pages 4171–4186 , Minneapolis,349 Minnesota , June 2019 . Association for Computational Linguistics.350 [ 6 ] Fabrice Gamboa . Methode du maximum d ’ entropie sur la moyenne et applications . Phd Thesis,351 1989.352 [ 7 ] Michael Hahn . Theoretical Limitations of Self-Attention in Neural Sequence Models . Transac-353 tions of the Association for Computational Linguistics , 8:156–171 , 2020.354 [ 8 ] Jiri Hron , Yasaman Bahri , and Jascha Sohl-dickstein Roman . Infinite attention : NNGP and355 NTK for deep attention networks . In International Conference on Machine Learning ( ICML ) ,356 2020.357 [ 9 ] Sarthak Jain and Byron C. Wallace . Attention is not Explanation . In Annual Conference of the358 North American Chapter of the Association for Computational Linguistics : Human Language359 Technologies ( NAACL-HLT ) , 2019.360 [ 10 ] E. T. Jaynes . Information theory and statistical mechanics . Phys . Rev. , 106:620–630 , May 1957.361 [ 11 ] E. T. Jaynes . Information theory and statistical mechanics . ii . Phys . Rev. , 108:171–190 , Oct362 1957.363 [ 12 ] Jinkyu Kim and John Canny . Interpretable learning for self-driving cars by visualizing causal364 attention . In 2017 IEEE International Conference on Computer Vision ( ICCV ) , pages 2961–365 2969 , 2017.366 [ 13 ] Jaehoon Lee , Jascha Sohl-dickstein , Jeffrey Pennington , Roman Novak , Sam Schoenholz , and367 Yasaman Bahri . Deep neural networks as gaussian processes . In International Conference on368 Learning Representations ( ICLR ) , 2018.369 [ 14 ] Haoye Lu , Yongyi Mao , and Amiya Nayak . On the dynamics of training attention models . In370 International Conference on Learning Representations , 2021.371 [ 15 ] Thang Luong , Hieu Pham , and Christopher D. Manning . Effective approaches to attention-based372 neural machine translation . In Proceedings of the 2015 Conference on Empirical Methods in Nat-373 ural Language Processing , pages 1412–1421 , Lisbon , Portugal , September 2015 . Association374 for Computational Linguistics.375 [ 16 ] Fenglong Ma , Radha Chitta , Jing Zhou , Quanzeng You , Tong Sun , and Jing Gao . Dipole:376 Diagnosis prediction in healthcare via attention-based bidirectional recurrent neural networks.377 In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery378 and Data Mining , KDD ’ 17 , page 1903–1911 , New York , NY , USA , 2017 . Association for379 Computing Machinery.380 [ 17 ] P. McCullagh . Tensor Methods in Statistics : Monographs on Statistics and Applied Probability.381 Chapman and Hall/CRC , Boca Raton , FL , first edition . edition , 1987.382 [ 18 ] Paul Michel , Omer Levy , and Graham Neubig . Are sixteen heads really better than one ? In383 Neural Information Processing Systems ( NIPS ) , volume 32 . Curran Associates , Inc. , 2019.384 [ 19 ] Colin Raffel , Noam Shazeer , Adam Roberts , Katherine Lee , Sharan Narang , Michael Matena,385 Yanqi Zhou , Wei Li , and Peter J. Liu . Exploring the limits of transfer learning with a unified386 text-to-text transformer . Journal of Machine Learning Research , 21 ( 140 ) :1–67 , 2020.387 [ 20 ] Hubert Ramsauer , Bernhard Schäfl , Johannes Lehner , Philipp Seidl , Michael Widrich , Lukas388 Gruber , Markus Holzleitner , Thomas Adler , David Kreil , Michael K Kopp , Günter Klambauer,389 Johannes Brandstetter , and Sepp Hochreiter . Hopfield networks is all you need . In International390 Conference on Learning Representations , 2021.391 [ 21 ] Gabriel Rioux , Rustum Choksi , Tim Hoheisel , Pierre Marechal , and Christopher Scarvelis . The392 maximum entropy on the mean method for image deblurring . Inverse Problems , oct 2020.393 [ 22 ] R. Tyrrell Rockafellar . Convex analysis . Princeton mathematical series ; 28 . Princeton University394 Press , Princeton , N.J , 1970.395 [ 23 ] Sofia Serrano and Noah A. Smith . Is attention interpretable ? In Annual Meeting of the396 Association for Computational Linguistics ( ACL ) , 2020.397 [ 24 ] Shikhar Vashishth , Shyam Upadhyay , Gaurav Singh Tomar , and Manaal Faruqui . Attention398 Interpretability Across NLP Tasks . In International Conference on Learning Representations399 ( ICLR ) , 2020.400 [ 25 ] Ashish Vaswani , Noam Shazeer , Niki Parmar , Jakob Uszkoreit , Llion Jones , Aidan N Gomez,401 Ł ukasz Kaiser , and Illia Polosukhin . Attention is all you need . In I. Guyon , U. V. Luxburg,402 S. Bengio , H. Wallach , R. Fergus , S. Vishwanathan , and R. Garnett , editors , Advances in Neural403 Information Processing Systems , volume 30 . Curran Associates , Inc. , 2017.404 [ 26 ] Elena Voita , David Talbot , Fedor Moiseev , Rico Sennrich , and Ivan Titov . Analyzing multi-head405 self-attention : Specialized heads do the heavy lifting , the rest can be pruned . In Proceedings of406 the 57th Annual Meeting of the Association for Computational Linguistics , July 2019.407 [ 27 ] Yequan Wang , Minlie Huang , Xiaoyan Zhu , and Li Zhao . Attention-based LSTM for aspect-408 level sentiment classification . In Empirical Methods in Natural Language Processing ( EMNLP ) ,409 pages 606–615 , Austin , Texas , November 2016 . Association for Computational Linguistics.410 [ 28 ] Sarah Wiegreffe and Yuval Pinter . Attention is not not explanation . In 2019 Conference on411 Empirical Methods in Natural Language Processing and 9th International Joint Conference on412 Natural Language Processing ( EMNLP-IJCNLP ) , 2019.413 [ 29 ] K. Xu , J. Ba , R. Kiros , K. Cho , A. Courville , R. Salakhudinov , R. Zemel , and Y. Bengio . Show,414 Attend and Tell : Neural Image Caption Generation with Visual Attention . In Proc . 32nd Int.415 Conf . Mach . Learn. , pages 257–261 , 2015.416 [ 30 ] Greg Yang . Tensor Programs I : Wide Feedforward or Recurrent Neural Networks of Any417 Architecture are Gaussian Processes . In Neural Information Processing Systems ( NIPS ) , 2019.418 Checklist419 1 . For all authors ... 420 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s421 contributions and scope ? [ Yes ] 422 ( b ) Did you describe the limitations of your work ? [ Yes ] We have created a sub-section,423 called Limitations and other comments , in Section 7 to discuss the potential limita-424 tions of our work.425 ( c ) Did you discuss any potential negative societal impacts of your work ? [ Yes ] We have426 discussed this at the end of the paper.427 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to428 them ? [ Yes ] 429 2 . If you are including theoretical results ... 430 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] 431 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] 432 3 . If you ran experiments ... 433 ( a ) Did you include the code , data , and instructions needed to reproduce the main exper-434 imental results ( either in the supplemental material or as a URL ) ? [ N/A ] Our work435 presents a theoretical framework and does not contain any experimental study.436 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they437 were chosen ) ? [ N/A ] 438 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-439 ments multiple times ) ? [ N/A ] 440 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type441 of GPUs , internal cluster , or cloud provider ) ? [ N/A ] 442 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 443 ( a ) If your work uses existing assets , did you cite the creators ? [ N/A ] 444 ( b ) Did you mention the license of the assets ? [ N/A ] 445 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ N/A ] 446 447 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re448 using/curating ? [ N/A ] 449 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable450 information or offensive content ? [ N/A ] 451 5 . If you used crowdsourcing or conducted research with human subjects ... 452 ( a ) Did you include the full text of instructions given to participants and screenshots , if453 applicable ? [ N/A ] 454 ( b ) Did you describe any potential participant risks , with links to Institutional Review455 Board ( IRB ) approvals , if applicable ? [ N/A ] 456 ( c ) Did you include the estimated hourly wage paid to participants and the total amount457 spent on participant compensation ? [ N/A ] 458

Summary: This work constructs an optimization problem whose approximate solution results in a general form of attention. The optimization objective is derived from three assumptions: 1. the prior of the attention distribution has a support of the memory bank; 2. the likelihood of the query vector in the attention mechanism is roughly conforming to an isotropic Gaussian centered at the difference between the mean of the attention distribution and the mean of a prior distribution; 3. the standard deviation of the isotropic Gaussian is large. Then the mean of the posterior of the attention distribution is a general form of attention, with the dot attention scores reweighted by the prior probabilities. Under this framework, when the prior is a uniform distribution, dot attention is recovered; when the prior is defined by positions and word masks, the attention used in T5 is recovered.