The Lipschitz Constant of Self-Attention

1 INTRODUCTION . Lipschitz continuity is a strong form of continuity for functions . Loosely speaking , a function is Lipschitz continuous if changing its input by a certain amount can not change its output by more than K times that amount . The constant K is a hard constraint on how rapidly the function ’ s output can vary , and the smallest such K is known as the function ’ s Lipschitz constant . For example , f1 ( x ) = √ |x| and f2 ( x ) = exp ( x ) for x ∈ R are not Lipschitz continuous , because their output can change arbitrarily fast as x approaches 0 and +∞ respectively . On the other hand , g1 ( x ) = tanh ( x ) and g2 ( x ) = αx are Lipschitz continuous , because their rate of change ( derivative ) is bounded . In deep learning , we often use Lipschitz continuity as a constraint for neural networks , to control how much a network ’ s output can change relative to its input . Such Lipschitz constraints are useful in several contexts . For example , Lipschitz constraints can endow models with provable robustness against adversarial pertubations ( Cisse et al. , 2017 ; Tsuzuku et al. , 2018 ; Anil et al. , 2019 ) , and guaranteed generalisation bounds ( Sokolić et al. , 2017 ) . Moreover , the dual form of the Wasserstein distance is defined as a supremum over Lipschitz functions with a given Lipschitz constant , hence Lipschitz-constrained networks are used for estimating Wasserstein distances ( Peyré & Cuturi , 2019 ) . Further , Lipschitz-constrained networks can stabilise training for GANs , an example being spectral normalisation ( Miyato et al. , 2018 ) . Finally , Lipschitz-constrained networks are also used to construct invertible models and normalising flows . For example , Lipschitz-constrained networks can be used as a building block for invertible residual networks and hence flow-based generative models ( Behrmann et al. , 2019 ; Chen et al. , 2019 ) . Additionally , Neural ODEs ( Chen et al. , 2018 ; Grathwohl et al. , 2019 ) are typically defined using vector fields parameterized via Lipschitz networks , so that the flow generated by the vector field is guaranteed to exist for all times . Nonetheless , designing Lipschitz-continuous neural networks and computing ( or even upperbounding ) their Lipschitz constant is a hard problem . Previous work mostly focused on fullyconnected and convolutional networks , not only because they are common in deep learning , but also because they are relatively simple to analyze , as compositions of linear maps and pointwise nonlinearities . Even in this case however , exact evaluation of the Lipschitz constant of fully-connected and convolutional networks is NP-hard ( Virmaux & Scaman , 2018 ) and obtaining a tight upper bound remains a challenging task ( Virmaux & Scaman , 2018 ; Fazlyab et al. , 2019 ; Latorre et al. , 2020 ) . Fully-connected and convolutional networks are not the only neural networks worthy of interest . Recently , self-attention ( Vaswani et al. , 2017 ) has become a popular alternative to recurrent neural networks . Self-attention is a key component of the Transformer ( Vaswani et al. , 2017 ) , that has found success as a building block in models of various data modalities , starting with natural-language processing ( Vaswani et al. , 2017 ; Devlin et al. , 2019 ; Brown et al. , 2020 ) and extending to computer vision ( Zhang et al. , 2019 ; Parmar et al. , 2019 ) , audio generation ( Huang et al. , 2019 ) , and reinforcement learning ( Parisotto et al. , 2020 ) . However , so far no previous work has analyzed the Lipschitz properties of self-attention , and thus it has been unclear whether self-attention is a viable option in applications that require Lipschitz constraints . In this work , we address this gap in the theory of self-attention by providing a thorough analysis of its Lipschitz properties . In particular , we make the following contributions : • We prove that the widely used dot-product self-attention is not Lipschitz , and therefore not suitable to use in applications requiring Lipschitz constraints . • We formulate L2 self-attention as an alternative , and show that it is Lipschitz . • We derive a theoretical upper bound on the Lipschitz constant of L2 self-attention , and provide empirical evidence of the asymptotic tightness of the bound . • As a practical demonstration of the theory , we use this bound to formulate invertible self-attention , and explore its use in a Transformer architecture for character-level language modelling . 2 LIPSCHITZ CONSTANT OF FULLY-CONNECTED/CONVOLUTIONAL LAYERS . We first define the notion of Lipschitz continuity , and proceed to define the Lipschitz constant . Definition 2.1 . Given two metric spaces ( X , dX ) and ( Y , dY ) , a function f : X → Y is called Lipschitz continuous ( or K-Lipschitz ) if there exists a constant K ≥ 0 such that dY ( f ( x ) , f ( x ′ ) ) ≤ KdX ( x , x′ ) for all x , x′ ∈ X . ( 1 ) The smallest such K is the Lipschitz constant of f , denoted Lip ( f ) . In this paper , we focus on the common case where X = Rn , Y = Rm , and dX , dY are induced by a p-norm ‖x‖p : = ( ∑ i |xi|p ) 1/p . We will primarily consider the cases p = 2 and p = ∞ , where ‖x‖∞ : = maxi |xi| . To emphasise the dependence of the Lipschitz constant on the choice of p-norm , we will often denote it by Lipp ( f ) . In this case , it follows directly from Definition 2.1 that the Lipschitz constant is given by Lipp ( f ) = sup x6=x′∈Rn ‖f ( x ) − f ( x′ ) ‖p ‖x− x′‖p . ( 2 ) Next , we outline some basic results that are useful for estimating Lipschitz constants , also covered in related works ( Virmaux & Scaman , 2018 ; Behrmann et al. , 2019 ) . We describe how these results are used to provide bounds on the Lipschitz constant of fully-connected networks ( FCN ) and convolutional neural networks ( CNN ) , using the fact that both are compositions of linear maps and pointwise non-linearities . To begin with , the following theorem suggests a way to bound Lipp ( f ) for a differentiable Lipschitz function f : Theorem 2.1 ( Federer , 1969 ) . Let f : Rn → Rm be differentiable and Lipschitz continuous under a choice of p-norm ‖ · ‖p . Let Jf ( x ) denote its total derivative ( Jacobian ) at x . Then Lipp ( f ) = supx∈Rn ‖Jf ( x ) ‖p where ‖Jf ( x ) ‖p is the induced operator norm on Jf ( x ) . Hence if f is a linear map represented by a matrix W then Lipp ( f ) = ‖W‖p : = sup ‖x‖p=1 ‖Wx‖p = { σmax ( W ) , if p = 2 maxi ∑ j |Wij | if p =∞ ( 3 ) where ‖W‖p is the operator norm on matrices induced by the vector p-norm , and σmax ( W ) is the largest singular value of W . Under this choice of norm , many common non-linearities ( including relu , sigmoid , tanh , elu ) are 1-Lipschitz . ‖W‖2 = σmax ( W ) is usually estimated via power iteration ; we provide details on how this is done in Appendix B . Since we now know the Lipschitz constants of the components of both FCN and CNN , we can bound their Lipschitz constants by applying the following lemma : Lemma 2.1 ( Federer , 1969 ) . Let g , h be two composable Lipschitz functions . Then g ◦ h is also Lipschitz with Lip ( g ◦ h ) ≤ Lip ( g ) Lip ( h ) . Corollary 2.1 . For a fully-connected network ( FCN ) or a convolutional neural network ( CNN ) f =WK ◦ ρK−1 ◦WK−1 ◦ . . . ◦ ρ1 ◦W1 , we have Lipp ( f ) ≤ ∏ k ‖Wk‖p under a choice of p-norm with 1-Lipschitz non-linearities ρk . The above bound is not necessarily tight ; there are various works that compute tighter bounds for FCN and CNN ( e.g . Virmaux & Scaman , 2018 ; Fazlyab et al. , 2019 ; Latorre et al. , 2020 ) . 3 LIPSCHITZ CONSTANT OF SELF-ATTENTION . 3.1 DOT-PRODUCT SELF-ATTENTION IS not LIPSCHITZ Moving on , we investigate whether self-attention is Lipschitz . We first consider the widely used ( scaled ) dot-product multihead self-attention as formulated by Vaswani et al . ( 2017 ) . Let x1 , . . . , xN be a sequence of N elements , where xi ∈ RD for i = 1 , . . . , N . We represent this sequence as a matrix X ∈ RN×D such that the ith row of X is the ith element of the sequence , i.e . Xi : = x > i . Dot-product multihead self-attention ( DP-MHA ) is a map from RN×D to RN×D consisting of H ‘ heads ’ , where H is chosen to divide D. Each head is a map from RN×D to RN×D/H defined by DP ( X ) : = softmax ( XWQ ( XWK ) > / √ D/H ) XWV , ( 4 ) where WQ , WK , WV ∈ RD×D/H are learnable parameters specific to each head . The input to the softmax is an N ×N matrix of dot products ( hence dot-product self-attention ) , and the softmax is applied to each row of this matrix . Finally , the outputs of all heads are concatenated into an N ×D matrix and are right multiplied by WO ∈ RD×D , thus DP-MHA is defined by MHA ( X ) : = [ DP1 ( X ) , . . . , DPH ( X ) ] WO . ( 5 ) In what follows , we will prove that MHA as defined above is not Lipschitz , assuming that the MHA map is non-trivial , i.e . WQ , WK , WV , WO 6= 0 . It is sufficient to show that a single head DP is not Lipschitz , since MHA is a linear combination of the outputs of each head . Let us write Equation ( 4 ) as DP ( X ) = PXWV , where P ∈ RN×N is the output of the softmax ( we suppress the dependence of P on X to reduce clutter below ) . P is a stochastic matrix , i.e . its entries are non-negative and its rows sum to 1 . Since the rows of X are the xi ’ s , a linear transformation of each xi by some matrix A is equivalent to right multiplication of X by A > . So right multiplication of X by WV is a linear map and thus Lipschitz . Therefore , we are interested in the mapping f ( X ) = PX ; this is not a linear mapping because P itself is a non-linear function of X . In fact , we show that f is not Lipschitz , thus proving the first main result of the paper : Theorem 3.1 . DP-MHA is not Lipschitz for any vector p-norm ‖ · ‖p with p ∈ [ 1 , ∞ ] . Summary of Proof . We use Theorem 2.1 , noting that if the supremum of the norm of the Jacobian is infinite , then the mapping is not Lipschitz . In particular , we show that when xi = 0 for some i , some elements of the Jacobian of f grow proportionally to the sample variance of x 6=i , which is unbounded . Proof . We show the proof for the case D = 1 ( i.e . X ∈ RN×1 , a column vector ) for readability . See Appendix C for the general case , which follows the same logic . The mapping f can be written as f ( X ) = PX = softmax ( aXX > ) X = f1 ( X ) > ... fN ( X ) > ∈ RN×1 , ( 6 ) where a = WKWQ ∈ R ( we assume a 6= 0 such that self-attention is non-trivial ) and fi ( X ) =∑N j=1 Pijxj with P > i : = softmax ( axiX ) . Hence f can be interpreted as a map of each xi to a point in the convex hull of x1 , ... , xN . Since f is a map from RN×1 to RN×1 , its Jacobian is Jf = J11 . . . J1N ... . . . ... JN1 . . . JNN ∈ RN×N , ( 7 ) where Jij = ∂fi ( X ) ∂xj ∈ R. By taking partial derivatives we can show that Jij = aX > P ( i ) [ ejiX + δijX ] + PijI where eij ∈ RN×N is a binary matrix with zeros everywhere except the ( i , j ) th entry , δij is the Kronecker delta , and P ( i ) : = diag ( Pi : ) − P > i : Pi : . So for i = j : Jii = aX > P ( i ) eiiX + aX > P ( i ) X + Pii ( 8 ) Let us investigate the scalarX > P ( i ) X . We observe that it is in fact a variance of a discrete distribution . Specifically : X > P ( i ) X = ∑ k Pikx 2 k − ( ∑ k Pikxk ) 2 = Var ( X ) , ( 9 ) where X is a discrete distribution with support at the inputs { x1 , . . . , xN } and probability mass function given by their softmax probabilities P ( X = xj ) = Pij . A consequence of this interpretation is that P ( i ) is positive semi-definite ( PSD ) since X > P ( i ) X = Var ( X ) ≥ 0 , with equality if and only if the xj are all equal . We use this observation to show that Jii is unbounded , and so ‖Jf‖p is unbounded , hence DP-MHA is not Lipschitz . Consider the case xi = 0 . Then P > i : = softmax ( XAxi ) = 1 N 1 , i.e . we have uniform attention regardless of x 6=i . The first term of Jii in Equation ( 8 ) disappears since eiiX = [ 0 , . . . , xi , . . . , 0 ] = 0 , and the last term becomes 1N I . Now consider the second term aX > P ( i ) X = aVar ( Xl ) . Note X is uniformly distributed , since P ( X = xj ) = Pij = 1/N . Hence the second term is equal to a times the sample variance of x1 , . . . , xN , which can be arbitrarily large . High-level intuition for proof . At xi = 0 , fi ( X ) = 1N ∑ k xk , the mean of the inputs . The rate of change of fi is governed by how fast the softmax saturates when xi is perturbed , which is determined by how spread out the x 6=i are . The more spread out they are ( the higher the sample variance ) , the greater the rate of saturation of the softmax , and the faster the rate of change of fi . Since the sample variance of x 6=i can be arbitrarily large , the rate of change of fi can also be arbitrarily large , i.e . the entries of the Jacobian ( and hence its p-norm ) can become arbitrarily large . In Appendix D , we show that adding bias terms to x > i W Q and x > j W K does not resolve the issue . The implications of this result are the following . ( 1 ) There can be undesirable behaviour ( e.g . training instabilities ) for the Transformer when some inputs are close to zero . ( 2 ) Dot-product self-attention ( and hence the standard Transformer ) is not a suitable choice when we require a Lipschitz neural network , such as for formulating invertible residual networks ( Behrmann et al. , 2019 ) . Therefore , to use self-attention and Transformers in such applications , a Lipschitz formulation of self-attention is required , together with an explicit ( ideally tight ) upper bound to its Lipschitz constant , to quantify how much the output can change with respect to changes in the input . One method to make dot-product self-attention Lipschitz is by ensuring its inputs are bounded . Indeed , if the input space is compact , e.g . [ 0 , 1 ] N×D , any continuously differentiable function is Lipschitz , including dot-product self-attention . However , as we further discuss in Section 6 , such an approach has its own challenges , since it makes the Lipschitz constant depend on the input range . Instead , in the next section we formulate a version of self-attention that is provably Lipschitz on all of RN×D , allowing us to derive an upper bound that holds for any subset of RN×D .

Summary: This paper studies the Lipschitz continuity properties of self-attention. It is proved that the widely-used dot-product self-attention is not Lipschitz continuous. A novel L2 self-attention is proposed and proven to be Lipschitz continuous. Experiments show that the upper bound of Lipschitz constant for L2 self-attention is asymptotically tight. Invertibility of MHA residual map is investigated to prove the Lipschitz continuity of L2 self-attention. Finally, experiments on Transformers with L2 self-attention are studied.