Original: T HE L IPSCHITZ C ONSTANT OF S ELF -A TTENTION
Revised: T HE L IPSCHITZ C ONSTANT OF S ELF -A TTENTION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Lipschitz constants of neural networks have been explored in various contexts in deep learning, such as provable adversarial robustness, estimating Wasserstein distance, stabilising training of GANs, and formulating invertible neural networks.
Revised: Lipschitz constants of neural networks have been explored in various contexts in deep learning, such as provable adversarial robustness, estimating Wasserstein distance, stabilising training of GANs, and formulating invertible neural networks.

Original: Such works have focused on bounding the Lipschitz constant of fully connected or convolutional networks, composed of linear maps and pointwise non-linearities.
Revised: Such works have focused on bounding the Lipschitz constant of fully connected or convolutional networks, composed of linear maps and pointwise non-linearities.

Original: In this paper, we investigate the Lipschitz constant of self-attention, a non-linear neural network module widely used in sequence modelling.
Revised: In this paper, we investigate the Lipschitz constant of self-attention, a non-linear neural network module widely used in sequence modelling.

Original: We prove that the standard dot-product self-attention is not Lipschitz, and propose an alternative L self-attention that is Lipschitz.
Revised: We prove that the standard dot-product self-attention is not Lipschitz, and propose an alternative L self-attention that is Lipschitz.

Original: We derive an upper bound on the Lipschitz constant of L2 self-attention and provide empirical evidence for its asymptotic tightness.
Revised: We derive an upper bound on the Lipschitz constant of L2 self-attention and provide empirical evidence for its asymptotic tightness.

Original: To demonstrate the practical relevance of our theoretical work, we formulate invertible self-attention and use it in a Transformer-based architecture for a character-level language modelling task.
Revised: To demonstrate the practical relevance of our theoretical work, we formulate invertible self-attention and use it in a Transformer-based architecture for a character-level language modelling task.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Lipschitz continuity is a strong form of continuity for functions.
Revised: Lipschitz continuity is a strong form of continuity for functions.

Original: Loosely speaking, a function is
Revised: Loosely speaking, a function is

Original: Lipschitz continuous if changing its input by a certain amount cannot change its output by more than K times that amount.
Revised: Lipschitz continuous if changing its input by a certain amount cannot change its output by more than K times that amount.

Original: 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 .
Revised: 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 .

Original: For example, f 1 ( x ) =
Revised: For example, f 1 ( x ) =

Original: (cid:112) | x | and f 2 ( x ) = exp( x ) for x ∈ R are not Lipschitz continuous, because their output can change arbitrarily fast as x approaches 0 and + ∞ respectively.
Revised: (cid:112) | x | and f 2 ( x ) = exp( x ) for x ∈ R are not Lipschitz continuous, because their output can change arbitrarily fast as x approaches 0 and + ∞ respectively.

Original: On the other hand, g 1 ( x ) = tanh( x ) and g 2 ( x ) = αx are Lipschitz continuous, because their rate of change (derivative) is bounded.
Revised: On the other hand, g 1 ( x ) = tanh( x ) and g 2 ( x ) = αx are Lipschitz continuous, because their rate of change (derivative) is bounded.

Original: 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.
Revised: 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.

Original: Such Lipschitz constraints are useful in several contexts.
Revised: Such Lipschitz constraints are useful in several contexts.

Original: For example, Lipschitz constraints can endow models with provable robustness against adversarial pertubations (Cisse et al., 2017; Tsuzuku et al., 2018;
Revised: For example, Lipschitz constraints can endow models with provable robustness against adversarial pertubations (Cisse et al., 2017; Tsuzuku et al., 2018;

Original: Anil et al., 2019), and guaranteed generalisation bounds (Sokoli ´ c et al., 2017).
Revised: Anil et al., 2019), and guaranteed generalisation bounds (Sokoli ´ c et al., 2017).

Original: Moreover, the dual form of the Wasserstein distance is deﬁned as a supremum over Lipschitz functions with a given Lipschitz constant, hence
Revised: Moreover, the dual form of the Wasserstein distance is deﬁned as a supremum over Lipschitz functions with a given Lipschitz constant, hence

Original: Lipschitz-constrained networks are used for estimating Wasserstein distances (Peyré & Cuturi, 2019).
Revised: Lipschitz-constrained networks are used for estimating Wasserstein distances (Peyré & Cuturi, 2019).

Original: Further, Lipschitz-constrained networks can stabilise training for GANs, an example being spectral normalisation (Miyato et al., 2018).
Revised: Further, Lipschitz-constrained networks can stabilise training for GANs, an example being spectral normalisation (Miyato et al., 2018).

Original: Finally, Lipschitz-constrained networks are also used to construct invertible models and normalising ﬂows.
Revised: Finally, Lipschitz-constrained networks are also used to construct invertible models and normalising ﬂows.

Original: For example, Lipschitz-constrained networks can be used as a building block for invertible residual networks and hence ﬂow-based generative models (Behrmann et al., 2019; Chen et al., 2019).
Revised: For example, Lipschitz-constrained networks can be used as a building block for invertible residual networks and hence ﬂow-based generative models (Behrmann et al., 2019; Chen et al., 2019).

Original: Additionally, Neural ODEs (Chen et al., 2018; Grathwohl et al., 2019) are typically deﬁned using vector ﬁelds parameterized via Lipschitz networks, so that the ﬂow generated by the vector ﬁeld is guaranteed to exist for all times.
Revised: Additionally, Neural ODEs (Chen et al., 2018; Grathwohl et al., 2019) are typically deﬁned using vector ﬁelds parameterized via Lipschitz networks, so that the ﬂow generated by the vector ﬁeld is guaranteed to exist for all times.

Original: Nonetheless, designing Lipschitz-continuous neural networks and computing (or even upperbounding) their Lipschitz constant is a hard problem.
Revised: Nonetheless, designing Lipschitz-continuous neural networks and computing (or even upperbounding) their Lipschitz constant is a hard problem.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Fully-connected and convolutional networks are not the only neural networks worthy of interest.
Revised: Fully-connected and convolutional networks are not the only neural networks worthy of interest.

Original: Recently, self-attention (Vaswani et al., 2017) has become a popular alternative to recurrent neural networks.
Revised: Recently, self-attention (Vaswani et al., 2017) has become a popular alternative to recurrent neural networks.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: In this work, we address this gap in the theory of self-attention by providing a thorough analysis of its Lipschitz properties.
Revised: In this work, we address this gap in the theory of self-attention by providing a thorough analysis of its Lipschitz properties.

Original: In particular, we make the following contributions:
Revised: In particular, we make the following contributions:

Original: • We prove that the widely used dot-product self-attention is not Lipschitz, and therefore not suitable to use in applications requiring Lipschitz constraints.
Revised: • We prove that the widely used dot-product self-attention is not Lipschitz, and therefore not suitable to use in applications requiring Lipschitz constraints.

Original: • We formulate L2 self-attention as an alternative, and show that it is Lipschitz.
Revised: • We formulate L2 self-attention as an alternative, and show that it is Lipschitz.

Original: • 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.
Revised: • 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.

Original: • 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.
Revised: • 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.

Original:  L IPSCHITZ C ONSTANT OF F ULLY -C ONNECTED /C ONVOLUTIONAL L AYERS
Revised:  L IPSCHITZ C ONSTANT OF F ULLY -C ONNECTED /C ONVOLUTIONAL L AYERS

Original: We ﬁrst deﬁne the notion of Lipschitz continuity, and proceed to deﬁne the Lipschitz constant.
Revised: We ﬁrst deﬁne the notion of Lipschitz continuity, and proceed to deﬁne the Lipschitz constant.

Original: Deﬁnition 2.1.
Revised: Deﬁnition 2.1.

Original: Given two metric spaces ( X , d X ) and ( Y , d Y ) , a function f : X → Y is called
Revised: Given two metric spaces ( X , d X ) and ( Y , d Y ) , a function f : X → Y is called

Original: Lipschitz continuous (or K - Lipschitz ) if there exists a constant K ≥ 0 such that
Revised: Lipschitz continuous (or K - Lipschitz ) if there exists a constant K ≥ 0 such that

Original: [Equation]
Revised: [Equation]

Original: The smallest such K is the Lipschitz constant of f , denoted Lip( f ) .
Revised: The smallest such K is the Lipschitz constant of f , denoted Lip( f ) .

Original: In this paper, we focus on the common case where X = R n , Y = R m , and d X , d Y are induced by a p -norm (cid:107) x (cid:107) p :=
Revised: In this paper, we focus on the common case where X = R n , Y = R m , and d X , d Y are induced by a p -norm (cid:107) x (cid:107) p :=

Original: ( (cid:80) i | x i | p ) 1 /p .
Revised: ( (cid:80) i | x i | p ) 1 /p .

Original: We will primarily consider the cases p = 2 and p = ∞ , where (cid:107) x
Revised: We will primarily consider the cases p = 2 and p = ∞ , where (cid:107) x

Original: (cid:107) ∞ :
Revised: (cid:107) ∞ :

Original: = max i | x i | .
Revised: = max i | x i | .

Original: To emphasise the dependence of the Lipschitz constant on the choice of p -norm, we will often denote it by Lip p ( f ) .
Revised: To emphasise the dependence of the Lipschitz constant on the choice of p -norm, we will often denote it by Lip p ( f ) .

Original: In this case, it follows directly from Deﬁnition 2.

Revised: In this case, it follows directly from Deﬁnition 2.

Original: that the
Revised: that the

Original: Lipschitz constant is given by
Revised: Lipschitz constant is given by

Original: [Equation]
Revised: [Equation]

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: To begin with, the following theorem suggests a way to bound Lip p ( f ) for a differentiable Lipschitz function f :
Revised: To begin with, the following theorem suggests a way to bound Lip p ( f ) for a differentiable Lipschitz function f :

Original: Theorem 2.1 (Federer, 1969) .
Revised: Theorem 2.1 (Federer, 1969) .

Original: Let f : R n → R m be differentiable and Lipschitz continuous under a choice of p -norm (cid:107) · (cid:107) p .
Revised: Let f : R n → R m be differentiable and Lipschitz continuous under a choice of p -norm (cid:107) · (cid:107) p .

Original: Let J f ( x ) denote its total derivative (Jacobian) at x .
Revised: Let J f ( x ) denote its total derivative (Jacobian) at x .

Original: Then
Revised: Then

Original: Lip p ( f ) =
Revised: Lip p ( f ) =

Original: sup x ∈ R n (cid:107)
Revised: sup x ∈ R n (cid:107)

Original: J f ( x ) (cid:107) p where (cid:107)
Revised: J f ( x ) (cid:107) p where (cid:107)

Original: J f ( x ) (cid:107) p is the induced operator norm on J f ( x ) .
Revised: J f ( x ) (cid:107) p is the induced operator norm on J f ( x ) .

Original: Hence if f is a linear map represented by a matrix W then
Revised: Hence if f is a linear map represented by a matrix W then

Original: [Equation]
Revised: [Equation]

Original: where (cid:107) W (cid:107) p is the operator norm on matrices induced by the vector p -norm, and σ max ( W ) is the largest singular value of W .
Revised: where (cid:107) W (cid:107) p is the operator norm on matrices induced by the vector p -norm, and σ max ( W ) is the largest singular value of W .

Original: Under this choice of norm, many common non-linearities (including relu , sigmoid , tanh , elu ) are 1 -Lipschitz.
Revised: Under this choice of norm, many common non-linearities (including relu , sigmoid , tanh , elu ) are 1 -Lipschitz.

Original: (cid:107) W (cid:107) 2 = σ max ( W ) is usually estimated via power iteration ; we provide details on how this is done in Appendix B.
Revised: (cid:107) W (cid:107) 2 = σ max ( W ) is usually estimated via power iteration ; we provide details on how this is done in Appendix B.

Original: 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:
Revised: 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:

Original: Lemma 2.1 (Federer, 1969) .
Revised: Lemma 2.1 (Federer, 1969) .

Original: Let g, h be two composable Lipschitz functions.
Revised: Let g, h be two composable Lipschitz functions.

Original: Then g ◦ h is also
Revised: Then g ◦ h is also

Original: Lipschitz with Lip( g ◦ h ) ≤ Lip( g ) Lip( h ) .
Revised: Lipschitz with Lip( g ◦ h ) ≤ Lip( g ) Lip( h ) .

Original: Corollary 2.1.
Revised: Corollary 2.1.

Original: For a fully-connected network ( FCN ) or a convolutional neural network ( CNN ) f = W K ◦ ρ K − 1 ◦ W K − 1 ◦ . .
Revised: For a fully-connected network ( FCN ) or a convolutional neural network ( CNN ) f = W K ◦ ρ K − 1 ◦ W K − 1 ◦ . .

Original: ◦ ρ 1 ◦ W 1 , we have Lip p ( f ) ≤ (cid:81)
Revised: ◦ ρ 1 ◦ W 1 , we have Lip p ( f ) ≤ (cid:81)

Original: k
Revised: k

Original: (cid:107) W k (cid:107) p under a choice of p -norm with 1 -Lipschitz non-linearities ρ k .
Revised: (cid:107) W k (cid:107) p under a choice of p -norm with 1 -Lipschitz non-linearities ρ k .

Original: The above bound is not necessarily tight; there are various works that compute tighter bounds for
Revised: The above bound is not necessarily tight; there are various works that compute tighter bounds for

Original: FCN and CNN (e.g. Virmaux & Scaman, 2018; Fazlyab et al., 2019; Latorre et al., 2020).
Revised: FCN and CNN (e.g. Virmaux & Scaman, 2018; Fazlyab et al., 2019; Latorre et al., 2020).

Original:  L IPSCHITZ C ONSTANT OF S ELF -A TTENTION 3. D OT - PRODUCT SELF - ATTENTION IS not L IPSCHITZ
Revised:  L IPSCHITZ C ONSTANT OF S ELF -A TTENTION 3. D OT - PRODUCT SELF - ATTENTION IS not L IPSCHITZ

Original: Moving on, we investigate whether self-attention is Lipschitz.
Revised: Moving on, we investigate whether self-attention is Lipschitz.

Original: We ﬁrst consider the widely used (scaled) dot-product multihead self-attention as formulated by Vaswani et al.
Revised: We ﬁrst consider the widely used (scaled) dot-product multihead self-attention as formulated by Vaswani et al.

Original: Let x 1 , . . .
Revised: Let x 1 , . . .

Original: , x N be a sequence of N elements, where x i ∈ R D for i = 1 , . . .
Revised: , x N be a sequence of N elements, where x i ∈ R D for i = 1 , . . .

Original: , N .
Revised: , N .

Original: We represent this sequence as a matrix X ∈ R N × D such that the i th row of X is the i th element of the sequence, i.e. X i : = x (cid:62) i .
Revised: We represent this sequence as a matrix X ∈ R N × D such that the i th row of X is the i th element of the sequence, i.e. X i : = x (cid:62) i .

Original: Dot-product multihead self-attention ( DP-MHA ) is a map from R N × D to R N × D consisting of H ‘heads’, where H is chosen to divide D .
Revised: Dot-product multihead self-attention ( DP-MHA ) is a map from R N × D to R N × D consisting of H ‘heads’, where H is chosen to divide D .

Original: Each head is a map from R N × D to R N × D/H deﬁned by
Revised: Each head is a map from R N × D to R N × D/H deﬁned by

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where W Q , W K , W V ∈ R D × D/H are learnable parameters speciﬁc to each head.
Revised: where W Q , W K , W V ∈ R D × D/H are learnable parameters speciﬁc to each head.

Original: 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.
Revised: 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.

Original: Finally, the outputs of all heads are concatenated into an N × D matrix and are right multiplied by W O ∈ R D × D , thus DP-MHA is deﬁned by
Revised: Finally, the outputs of all heads are concatenated into an N × D matrix and are right multiplied by W O ∈ R D × D , thus DP-MHA is deﬁned by

Original: [Equation]
Revised: [Equation]

Original: In what follows, we will prove that MHA as deﬁned above is not Lipschitz, assuming that the MHA map is non-trivial, i.e. W Q , W K , W V , W O (cid:54) = 0 .
Revised: In what follows, we will prove that MHA as deﬁned above is not Lipschitz, assuming that the MHA map is non-trivial, i.e. W Q , W K , W V , W O (cid:54) = 0 .

Original: It is sufﬁcient to show that a single head DP is not
Revised: It is sufﬁcient to show that a single head DP is not

Original: Lipschitz, since MHA is a linear combination of the outputs of each head.
Revised: Lipschitz, since MHA is a linear combination of the outputs of each head.

Original: Let us write Equation (4) as DP ( X )
Revised: Let us write Equation (4) as DP ( X )

Original: =
Revised: =

Original: PXW V , where P ∈ R N × N is the output of the softmax (we suppress the dependence of P on X to reduce clutter below).
Revised: PXW V , where P ∈ R N × N is the output of the softmax (we suppress the dependence of P on X to reduce clutter below).

Original: P is a stochastic matrix, i.e. its entries are non-negative and its rows sum to 1 .
Revised: P is a stochastic matrix, i.e. its entries are non-negative and its rows sum to 1 .

Original: Since the rows of X are the x i ’s, a linear transformation of each x
Revised: Since the rows of X are the x i ’s, a linear transformation of each x

Original: i by some matrix A is equivalent to right multiplication of X by A (cid:62) .
Revised: i by some matrix A is equivalent to right multiplication of X by A (cid:62) .

Original: So right multiplication of X by W V is a linear map and thus Lipschitz.
Revised: So right multiplication of X by W V is a linear map and thus Lipschitz.

Original: Therefore, we are interested in the mapping f ( X ) =
Revised: Therefore, we are interested in the mapping f ( X ) =

Original: PX ; this is not a linear mapping because P itself is a non-linear function of X .
Revised: PX ; this is not a linear mapping because P itself is a non-linear function of X .

Original: In fact, we show that f is not Lipschitz, thus proving the ﬁrst main result of the paper:
Revised: In fact, we show that f is not Lipschitz, thus proving the ﬁrst main result of the paper:

Original: Theorem 3.1.
Revised: Theorem 3.1.

Original: DP-MHA is not Lipschitz for any vector p -norm (cid:107) · (cid:107) p with p ∈ [1 , ∞ ] .
Revised: DP-MHA is not Lipschitz for any vector p -norm (cid:107) · (cid:107) p with p ∈ [1 , ∞ ] .

Original: Summary of Proof .
Revised: Summary of Proof .

Original: We use Theorem 2.1, noting that if the supremum of the norm of the Jacobian is inﬁnite, then the mapping is not Lipschitz.
Revised: We use Theorem 2.1, noting that if the supremum of the norm of the Jacobian is inﬁnite, then the mapping is not Lipschitz.

Original: In particular, we show that when x i = 0 for some i , some elements of the Jacobian of f grow proportionally to the sample variance of x (cid:54) =
Revised: In particular, we show that when x i = 0 for some i , some elements of the Jacobian of f grow proportionally to the sample variance of x (cid:54) =

Original: i , which is unbounded.
Revised: i , which is unbounded.

Original: Proof.
Revised: Proof.

Original: We show the proof for the case D = 1 (i.e. X ∈ R N × 1 , a column vector) for readability.
Revised: We show the proof for the case D = 1 (i.e. X ∈ R N × 1 , a column vector) for readability.

Original: See
Revised: See

Original: Appendix C for the general case, which follows the same logic.
Revised: Appendix C for the general case, which follows the same logic.

Original: The mapping f can be written as
Revised: The mapping f can be written as

Original: [Equation]
Revised: [Equation]

Original: where a = W K W Q ∈ R (we assume a (cid:54) = 

Revised: where a = W K W Q ∈ R (we assume a (cid:54) =

Original: such that self-attention is non-trivial) and f i ( X ) = (cid:80)
Revised: such that self-attention is non-trivial) and f i ( X ) = (cid:80)

Original: Nj =1 P ij x j with P (cid:62) i : = softmax ( ax i X ) .
Revised: Nj =1 P ij x j with P (cid:62) i : = softmax ( ax i X ) .

Original: Hence f can be interpreted as a map of each x
Revised: Hence f can be interpreted as a map of each x

Original: i to a point in the convex hull of x 1 , ..., x N .
Revised: i to a point in the convex hull of x 1 , ..., x N .

Original: Since f is a map from R N × 1 to R N × 1 , its Jacobian is
Revised: Since f is a map from R N × 1 to R N × 1 , its Jacobian is

Original: [Equation]
Revised: [Equation]

Original: where J ij =
Revised: where J ij =

Original: ∂f i ( X ) ∂x j ∈ R .
Revised: ∂f i ( X ) ∂x j ∈ R .

Original: By taking partial derivatives we can show that J ij =
Revised: By taking partial derivatives we can show that J ij =

Original: aX (cid:62) P ( i )
Revised: aX (cid:62) P ( i )

Original: [ e ji X + δ ij X ]
Revised: [ e ji X + δ ij X ]

Original: + P ij I where e ij ∈ R N × N is a binary matrix with zeros everywhere except the ( i, j ) th entry, δ ij is the Kronecker delta, and P ( i ) :
Revised: + P ij I where e ij ∈ R N × N is a binary matrix with zeros everywhere except the ( i, j ) th entry, δ ij is the Kronecker delta, and P ( i ) :

Original: =
Revised: =

Original: diag( P i : ) − P (cid:62) i : P i : .
Revised: diag( P i : ) − P (cid:62) i : P i : .

Original: So for i =
Revised: So for i =

Original: j :
Revised: j :

Original: [Equation]
Revised: [Equation]

Original: Let us investigate the scalar X (cid:62) P ( i ) X .
Revised: Let us investigate the scalar X (cid:62) P ( i ) X .

Original: We observe that it is in fact a variance of a discrete distribution.
Revised: We observe that it is in fact a variance of a discrete distribution.

Original: Speciﬁcally:
Revised: Speciﬁcally:

Original: [Equation]
Revised: [Equation]

Original: where X is a discrete distribution with support at the inputs { x 1 , .
Revised: where X is a discrete distribution with support at the inputs { x 1 , .

Original: , x N } and probability mass function given by their softmax probabilities P ( X = x j )
Revised: , x N } and probability mass function given by their softmax probabilities P ( X = x j )

Original: = P ij .
Revised: = P ij .

Original: A consequence of this interpretation is that P ( i ) is positive semi-deﬁnite (PSD) since X (cid:62) P ( i ) X =
Revised: A consequence of this interpretation is that P ( i ) is positive semi-deﬁnite (PSD) since X (cid:62) P ( i ) X =

Original: Var( X ) ≥ 0 , with equality if and only if the x j are all equal.
Revised: Var( X ) ≥ 0 , with equality if and only if the x j are all equal.

Original: We use this observation to show that J ii is unbounded, and so (cid:107) J f (cid:107) p is unbounded, hence DP-MHA is not Lipschitz.
Revised: We use this observation to show that J ii is unbounded, and so (cid:107) J f (cid:107) p is unbounded, hence DP-MHA is not Lipschitz.

Original: Consider the case x i = 0 .
Revised: Consider the case x i = 0 .

Original: Then P (cid:62) i : = softmax ( XAx i ) = 1 N 1 , i.e. we have uniform attention regardless of x (cid:54)
Revised: Then P (cid:62) i : = softmax ( XAx i ) = 1 N 1 , i.e. we have uniform attention regardless of x (cid:54)

Original: =
Revised: =

Original: i .
Revised: i .

Original: The ﬁrst term of J ii in Equation (8) disappears since e ii X = [0 , . . .
Revised: The ﬁrst term of J ii in Equation (8) disappears since e ii X = [0 , . . .

Original: , x i , . . .
Revised: , x i , . . .

Original: = 0 , and the last term becomes 1 N I .
Revised: = 0 , and the last term becomes 1 N I .

Original: Now consider the second term aX (cid:62) P ( i )
Revised: Now consider the second term aX (cid:62) P ( i )

Original: X
Revised: X

Original: = a Var( X l ) .
Revised: = a Var( X l ) .

Original: Note X is uniformly distributed, since P ( X = x j ) =
Revised: Note X is uniformly distributed, since P ( X = x j ) =

Original: P ij = 1 /N .
Revised: P ij = 1 /N .

Original: Hence the second term is equal to a times the sample variance of x 1 , . . .
Revised: Hence the second term is equal to a times the sample variance of x 1 , . . .

Original: , x N , which can be arbitrarily large.
Revised: , x N , which can be arbitrarily large.

Original: High-level intuition for proof.
Revised: High-level intuition for proof.

Original: At x i = 0 , f i ( X ) = 1 N (cid:80) k x k , the mean of the inputs.
Revised: At x i = 0 , f i ( X ) = 1 N (cid:80) k x k , the mean of the inputs.

Original: The rate of change of f i is governed by how fast the softmax saturates when x i is perturbed, which is determined by how spread out the x (cid:54)
Revised: The rate of change of f i is governed by how fast the softmax saturates when x i is perturbed, which is determined by how spread out the x (cid:54)

Original: =
Revised: =

Original: i are.
Revised: i are.

Original: 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 f i .
Revised: 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 f i .

Original: Since the sample variance of x (cid:54) = i can be arbitrarily large, the rate of change of f i can also be arbitrarily large, i.e. the entries of the Jacobian (and hence its p -norm) can become arbitrarily large.
Revised: Since the sample variance of x (cid:54) = i can be arbitrarily large, the rate of change of f i can also be arbitrarily large, i.e. the entries of the Jacobian (and hence its p -norm) can become arbitrarily large.

Original: In Appendix D, we show that adding bias terms to x (cid:62) i W Q and x (cid:62) j W K does not resolve the issue.
Revised: In Appendix D, we show that adding bias terms to x (cid:62) i W Q and x (cid:62) j W K does not resolve the issue.

Original: The implications of this result are the following.
Revised: The implications of this result are the following.

Original: (1) There can be undesirable behaviour (e.g. training instabilities) for the Transformer when some inputs are close to zero.
Revised: (1) There can be undesirable behaviour (e.g. training instabilities) for the Transformer when some inputs are close to zero.

Original: (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).
Revised: (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).

Original: 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.
Revised: 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.

Original: One method to make dot-product self-attention Lipschitz is by ensuring its inputs are bounded.
Revised: One method to make dot-product self-attention Lipschitz is by ensuring its inputs are bounded.

Original: Indeed, if the input space is compact, e.g. [0 , 1] N × D , any continuously differentiable function is
Revised: Indeed, if the input space is compact, e.g. [0 , 1] N × D , any continuously differentiable function is

Original: Lipschitz, including dot-product self-attention.
Revised: Lipschitz, including dot-product self-attention.

Original: 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.
Revised: 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.

Original: Instead, in the next section we formulate a version of self-attention that is provably Lipschitz on all of R N × D , allowing us to derive an upper bound that holds for any subset of R N × D .
Revised: Instead, in the next section we formulate a version of self-attention that is provably Lipschitz on all of R N × D , allowing us to derive an upper bound that holds for any subset of R N × D .

Original: 3. L2 SELF - ATTENTION : A L IPSCHITZ FORMULATION OF SELF - ATTENTION
Revised: 3. L2 SELF - ATTENTION : A L IPSCHITZ FORMULATION OF SELF - ATTENTION

Original: The pathology in dot-product self-attention arises because the softmax probabilities P i : are constant with respect to x (cid:54)
Revised: The pathology in dot-product self-attention arises because the softmax probabilities P i : are constant with respect to x (cid:54)

Original: = i
Revised: = i

Original: when x i = 0 .
Revised: when x i = 0 .

Original: This behaviour can be undesirable as we want P ij to vary according to x j , regardless of whether x i is zero or not.
Revised: This behaviour can be undesirable as we want P ij to vary according to x j , regardless of whether x i is zero or not.

Original: Hence we propose an alternative form of self-attention based on L2 distance:
Revised: Hence we propose an alternative form of self-attention based on L2 distance:

Original: [Equation]
Revised: [Equation]

Original: with the normalisation constant ensuring that (cid:80) j P ij = 1 .
Revised: with the normalisation constant ensuring that (cid:80) j P ij = 1 .

Original: We will refer to it as L2 self-attention .
Revised: We will refer to it as L2 self-attention .

Original: It is reminiscent of the standard squared-exponential kernel, but with softmax normalisation that ensures that each row of the kernel matrix sums to 1 .
Revised: It is reminiscent of the standard squared-exponential kernel, but with softmax normalisation that ensures that each row of the kernel matrix sums to 1 .

Original: Normalisation is usually necessary to deal with inputs of varying length N (Wang et al., 2018), hence we keep the softmax for L2 self-attention.
Revised: Normalisation is usually necessary to deal with inputs of varying length N (Wang et al., 2018), hence we keep the softmax for L2 self-attention.

Original: Similarly to dot-product self-attention, L2 self-attention can be computed efﬁciently with matrix operations; see Appendix E for details, with a comparison of wall-clock runtimes between different choices of attention.
Revised: Similarly to dot-product self-attention, L2 self-attention can be computed efﬁciently with matrix operations; see Appendix E for details, with a comparison of wall-clock runtimes between different choices of attention.

Original: We ﬁrst state the mathematical formulation of L2 multihead self-attention ( L2-MHA ) before proving the main result — the upper bound of its Lipschitz constant with respect to (cid:107) · (cid:107) p for p = 2 , ∞ .
Revised: We ﬁrst state the mathematical formulation of L2 multihead self-attention ( L2-MHA ) before proving the main result — the upper bound of its Lipschitz constant with respect to (cid:107) · (cid:107) p for p = 2 , ∞ .

Original: The full L2-MHA map F : R N × D
Revised: The full L2-MHA map F : R N × D

Original: → R N × D is deﬁned as
Revised: → R N × D is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: where f h ( X ) :
Revised: where f h ( X ) :

Original: = P h XA h .
Revised: = P h XA h .

Original: In the above, W V,h ∈ R D × D/H , W O ∈ R D × D , P h is deﬁned as in Equation (10) with W Q,h =
Revised: In the above, W V,h ∈ R D × D/H , W O ∈ R D × D , P h is deﬁned as in Equation (10) with W Q,h =

Original: W K,h ∈ R D × D/H , and A h := W Q,h W Q,h (cid:62) / (cid:112) D/H ∈ R D × D .
Revised: W K,h ∈ R D × D/H , and A h := W Q,h W Q,h (cid:62) / (cid:112) D/H ∈ R D × D .

Original: There are two changes from the usual form of multihead self-attention:
Revised: There are two changes from the usual form of multihead self-attention:

Original: (1) We require W Q,h = W K,h for each head
Revised: (1) We require W Q,h = W K,h for each head

Original: f h ( X ) to be Lipschitz.
Revised: f h ( X ) to be Lipschitz.

Original: In Lemma F.1 of Appendix F we show that L2-MHA is not Lipschitz for arbitrary W Q,h , W K,h , and that tying W Q,h = W K,h is sufﬁcient for L2-MHA to be Lipschitz, with intuition for why tying is sufﬁcient.
Revised: In Lemma F.1 of Appendix F we show that L2-MHA is not Lipschitz for arbitrary W Q,h , W K,h , and that tying W Q,h = W K,h is sufﬁcient for L2-MHA to be Lipschitz, with intuition for why tying is sufﬁcient.

Original: In each head of the self-attention f h ( X ) , right multiplication by A h has been included for the theorem below to hold (details are in the proof).
Revised: In each head of the self-attention f h ( X ) , right multiplication by A h has been included for the theorem below to hold (details are in the proof).

Original: In practice, there is little harm done by this extra linear transformation, since when the heads are combined together in F , each f h ( X ) is additionally transformed by W V,h , a free parameter.
Revised: In practice, there is little harm done by this extra linear transformation, since when the heads are combined together in F , each f h ( X ) is additionally transformed by W V,h , a free parameter.

Original: The second main result of the paper is the following:
Revised: The second main result of the paper is the following:

Original: Theorem 3.2.
Revised: Theorem 3.2.

Original: L2-MHA is Lipschitz, with the following bound on Lip ∞ ( F ) :
Revised: L2-MHA is Lipschitz, with the following bound on Lip ∞ ( F ) :

Original: [Equation]
Revised: [Equation]

Original: ∞ and the following bound on Lip 2 ( F ) :
Revised: ∞ and the following bound on Lip 2 ( F ) :

Original: [Equation]
Revised: [Equation]

Original: where φ ( x ) := x exp( x + 1) is an invertible univariate function on x > 0 , and N is the input sequence length.
Revised: where φ ( x ) := x exp( x + 1) is an invertible univariate function on x > 0 , and N is the input sequence length.

Original: Speciﬁcally, φ − 1 ( N − 1) =
Revised: Speciﬁcally, φ − 1 ( N − 1) =

Original: W 

Revised: W

Original: ( Ne ) where W 0 is the Lambert W -function, which grows sublogarithmically as O (log N − log log N ) (Corless et al., 1996).
Revised: ( Ne ) where W 0 is the Lambert W -function, which grows sublogarithmically as O (log N − log log N ) (Corless et al., 1996).

Original: Hence the above bounds can be simpliﬁed to O (log N ) for p = ∞ and O ( √ N log N ) for p = 2 .
Revised: Hence the above bounds can be simpliﬁed to O (log N ) for p = ∞ and O ( √ N log N ) for p = 2 .

Original: Proof.
Revised: Proof.

Original: See Appendix F, which uses the key observation that X (cid:62) P ( i )
Revised: See Appendix F, which uses the key observation that X (cid:62) P ( i )

Original: X is a covariance matrix (c.f.
Revised: X is a covariance matrix (c.f.

Original: Equation (9)) to bound (cid:107) J F (cid:107) p , the norm of the Jacobian of F .
Revised: Equation (9)) to bound (cid:107) J F (cid:107) p , the norm of the Jacobian of F .

Original: Appendix G shows how the argument can be modiﬁed to prove the analogous result for the case with masking in the self-attention.
Revised: Appendix G shows how the argument can be modiﬁed to prove the analogous result for the case with masking in the self-attention.

Original: These bounds are complemented by the concurrent work of Vuckovic et al.
Revised: These bounds are complemented by the concurrent work of Vuckovic et al.

Original: (2020), which provides a
Revised: (2020), which provides a

Original: O ( √ D log N ) bound on Lip 1 ( F ) using measure-theoretic tools.
Revised: O ( √ D log N ) bound on Lip 1 ( F ) using measure-theoretic tools.

Original:  A PPLICATION : I NVERTIBLE S ELF -A TTENTION 4. I NVERTIBLE RESIDUAL NETWORK
Revised:  A PPLICATION : I NVERTIBLE S ELF -A TTENTION 4. I NVERTIBLE RESIDUAL NETWORK

Original: Consider the residual function g ( x ) :
Revised: Consider the residual function g ( x ) :

Original: = x
Revised: = x

Original: + f ( x ) .
Revised: + f ( x ) .

Original: Behrmann et al.
Revised: Behrmann et al.

Original: (2019) give the following sufﬁcient condition for its invertibility: if f is a contraction with respect to some metric, i.e. if Lip( f ) < 1 , and the metric space on which f is deﬁned is complete, then g is invertible.
Revised: (2019) give the following sufﬁcient condition for its invertibility: if f is a contraction with respect to some metric, i.e. if Lip( f ) < 1 , and the metric space on which f is deﬁned is complete, then g is invertible.

Original: (A Euclidean space with a metric induced by a p -norm (cid:107) · (cid:107) p for p ∈ [1 , ∞ ] is always complete.)
Revised: (A Euclidean space with a metric induced by a p -norm (cid:107) · (cid:107) p for p ∈ [1 , ∞ ] is always complete.)

Original: Speciﬁcally, the inverse g − 1 ( y ) is the unique ﬁxed point of the recursion x i +1 :
Revised: Speciﬁcally, the inverse g − 1 ( y ) is the unique ﬁxed point of the recursion x i +1 :

Original: = y − f ( x i ) , since by the deﬁnition of the inverse we have y = g − 1 ( y )+ f ( g − 1 ( y )) .
Revised: = y − f ( x i ) , since by the deﬁnition of the inverse we have y = g − 1 ( y )+ f ( g − 1 ( y )) .

Original: Because f is a contraction, Banach’s Fixed Point Theorem guarantees that this ﬁxed point exists and is unique for all y , and that the recursion converges for all initial values x 0 (often set to y in practice) exponentially fast.
Revised: Because f is a contraction, Banach’s Fixed Point Theorem guarantees that this ﬁxed point exists and is unique for all y , and that the recursion converges for all initial values x 0 (often set to y in practice) exponentially fast.

Original: Hence the inverse can be computed to arbitrary accuracy (up to numerical precision in practice) by the above ﬁxed-point iteration.
Revised: Hence the inverse can be computed to arbitrary accuracy (up to numerical precision in practice) by the above ﬁxed-point iteration.

Original: Note that a composition of such invertible residual blocks is also invertible.
Revised: Note that a composition of such invertible residual blocks is also invertible.

Original: Behrmann et al.
Revised: Behrmann et al.

Original: (2019) use this observation to design invertible ResNets: they take f to be a CNN normalised by an upper bound on Lip( f ) given by Corollary 2.1, making the resulting function contractive .
Revised: (2019) use this observation to design invertible ResNets: they take f to be a CNN normalised by an upper bound on Lip( f ) given by Corollary 2.1, making the resulting function contractive .

Original: For the 2 -norm (cid:107) · (cid:107) 2 , a hyperparameter c < 1 is chosen and each linear map (convolution) W in the CNN is multiplied by c/ (cid:107) W (cid:107) 2 if c < (cid:107) W (cid:107) 2 where (cid:107) W (cid:107) 2 is estimated by power iteration (c.f. Appendix B).
Revised: For the 2 -norm (cid:107) · (cid:107) 2 , a hyperparameter c < 1 is chosen and each linear map (convolution) W in the CNN is multiplied by c/ (cid:107) W (cid:107) 2 if c < (cid:107) W (cid:107) 2 where (cid:107) W (cid:107) 2 is estimated by power iteration (c.f. Appendix B).

Original: This multiplicative factor determines the scale of the Lipschitz constant of the normalised function.
Revised: This multiplicative factor determines the scale of the Lipschitz constant of the normalised function.

Original: 4. I NVERTIBLE SELF - ATTENTION
Revised: 4. I NVERTIBLE SELF - ATTENTION

Original: [Figure]
Revised: [Figure]

Original: Output
Revised: Output

Original: The standard use case of self-attention is with a skip connection inside the
Revised: The standard use case of self-attention is with a skip connection inside the

Original: Transformer.
Revised: Transformer.

Original: A Transformer block is composed of residual blocks of multihead self-attention ( MHA ) and fully-connected ( FCN ) layers (Figure 1).
Revised: A Transformer block is composed of residual blocks of multihead self-attention ( MHA ) and fully-connected ( FCN ) layers (Figure 1).

Original: Hence similarly to invertible ResNets, we can normalise L2-MHA by the upper bounds given in Theorem 3.2 to obtain Contractive-L2-MHA f , with which we can obtain invertible self-attention g ( x ) = x
Revised: Hence similarly to invertible ResNets, we can normalise L2-MHA by the upper bounds given in Theorem 3.2 to obtain Contractive-L2-MHA f , with which we can obtain invertible self-attention g ( x ) = x

Original: + f ( x ) .
Revised: + f ( x ) .

Original: In the next section, we investigate the properties of invertible self-attention and how it compares with the standard dot-product self-attention; we replace
Revised: In the next section, we investigate the properties of invertible self-attention and how it compares with the standard dot-product self-attention; we replace

Original: DP-MHA in the Transformer with Contractive-L2-MHA , hence replacing the residual self-attention module with invertible self-attention.
Revised: DP-MHA in the Transformer with Contractive-L2-MHA , hence replacing the residual self-attention module with invertible self-attention.

Original: We are not interested in the modiﬁed Transformer per se, but rather in comparing the properties of invertible self-attention to standard self-attention — we only use the Transformer as a testbed for this purpose, since self-attention is commonly used in a Transformer.
Revised: We are not interested in the modiﬁed Transformer per se, but rather in comparing the properties of invertible self-attention to standard self-attention — we only use the Transformer as a testbed for this purpose, since self-attention is commonly used in a Transformer.

Original: Given the theoretical focus of the paper, we believe that a more challenging application of invertible self-attention, such as normalising ﬂow-based modelling, would be more suitable as a separate paper focused on that particular application.
Revised: Given the theoretical focus of the paper, we believe that a more challenging application of invertible self-attention, such as normalising ﬂow-based modelling, would be more suitable as a separate paper focused on that particular application.

Original: In Appendix H, we show that Dropout in the residual branch is also contractive.
Revised: In Appendix H, we show that Dropout in the residual branch is also contractive.

Original:  E XPERIMENTAL R ESULTS 5. A SYMPTOTIC TIGHTNESS OF THE UPPER BOUND ON Lip
Revised:  E XPERIMENTAL R ESULTS 5. A SYMPTOTIC TIGHTNESS OF THE UPPER BOUND ON Lip

Original: [Figure]
Revised: [Figure]

Original: ( F )
Revised: ( F )

Original: [Figure]
Revised: [Figure]

Original: A tight bound on the Lipschitz constant of selfattention is desirable for all listed applications in Section 1; it leads to tighter generalisation bounds, lighter constraints for provable robustness, and better expressiveness in residual ﬂow models.
Revised: A tight bound on the Lipschitz constant of selfattention is desirable for all listed applications in Section 1; it leads to tighter generalisation bounds, lighter constraints for provable robustness, and better expressiveness in residual ﬂow models.

Original: Hence we investigate the tightness of our bound on the Lipschitz constant of L2-MHA .
Revised: Hence we investigate the tightness of our bound on the Lipschitz constant of L2-MHA .

Original: The Lipschitz constant is a supremum over the space of inputs X ∈ R N × D (c.f. Equation (2)) and approximating it requires solving an intractable optimisation problem.
Revised: The Lipschitz constant is a supremum over the space of inputs X ∈ R N × D (c.f. Equation (2)) and approximating it requires solving an intractable optimisation problem.

Original: Hence it is infeasible to estimate accurately in general, especially when X is high-dimensional.
Revised: Hence it is infeasible to estimate accurately in general, especially when X is high-dimensional.

Original: However, we may compute a lower bound on the Lipschitz constant by maximising the norm of the Jacobian (cid:107) J f ( X ) (cid:107) with respect to X until convergence.
Revised: However, we may compute a lower bound on the Lipschitz constant by maximising the norm of the Jacobian (cid:107) J f ( X ) (cid:107) with respect to X until convergence.

Original: This local optimum will form a lower bound by Theorem 2.1, and we can expect this lower bound to be fairly tight for the low-dimensional case, provided the optimisation is thorough.
Revised: This local optimum will form a lower bound by Theorem 2.1, and we can expect this lower bound to be fairly tight for the low-dimensional case, provided the optimisation is thorough.

Original: We use this observation to provide empirical evidence for the asymptotic tightness of the upper bound on Lip ∞ ( f ) in Theorem 3.2.
Revised: We use this observation to provide empirical evidence for the asymptotic tightness of the upper bound on Lip ∞ ( f ) in Theorem 3.2.

Original: In Figure 2, we show the upper bound as well as the lower bound on Lip ∞ ( f ) obtained by optimising (cid:107) J f ( X ) (cid:107) ∞ with respect to X for L2-MHA f withdifferent random initialisations of X , with H = D = 1 and N varying between 100 and 1000 .
Revised: In Figure 2, we show the upper bound as well as the lower bound on Lip ∞ ( f ) obtained by optimising (cid:107) J f ( X ) (cid:107) ∞ with respect to X for L2-MHA f withdifferent random initialisations of X , with H = D = 1 and N varying between 100 and 1000 .

Original: See
Revised: See

Original: Appendix I for further details.
Revised: Appendix I for further details.

Original: Note that we use a log-scale for the x-axis, and recall that the upper bound is O (log N − log log N ) , dominated by the O (log N ) term for large N .
Revised: Note that we use a log-scale for the x-axis, and recall that the upper bound is O (log N − log log N ) , dominated by the O (log N ) term for large N .

Original: Hence the plot for the upper bound shows a linear trend.
Revised: Hence the plot for the upper bound shows a linear trend.

Original: We also observe that the slope of the lower bound is very similar, providing empirical evidence that the O (log N − log log N ) upper bound is asymptotically tight.
Revised: We also observe that the slope of the lower bound is very similar, providing empirical evidence that the O (log N − log log N ) upper bound is asymptotically tight.

Original: There are at least two possible explanations for the gap between the upper and lower bounds.
Revised: There are at least two possible explanations for the gap between the upper and lower bounds.

Original: (1) The lower bound is only a local optimum — the true Lipschitz constant is a global optimum across inputs, which can be difﬁcult to attain especially for high values of N .
Revised: (1) The lower bound is only a local optimum — the true Lipschitz constant is a global optimum across inputs, which can be difﬁcult to attain especially for high values of N .

Original: The multiplicative constant of the upper bound may be loose.
Revised: The multiplicative constant of the upper bound may be loose.

Original: Assuming asymptotic tightness, it remains an open question whether the multiplicative constant can be tightened.
Revised: Assuming asymptotic tightness, it remains an open question whether the multiplicative constant can be tightened.

Original: We show the analogous plot for Lip 2 ( F ) and discuss the results in Appendix K. Additionally in Appendix L, we show that optimising (cid:107) J f ( X ) (cid:107) ∞ w.r.t.
Revised: We show the analogous plot for Lip 2 ( F ) and discuss the results in Appendix K. Additionally in Appendix L, we show that optimising (cid:107) J f ( X ) (cid:107) ∞ w.r.t.

Original: X for DP-MHA f causes the norm to diverge, providing empirical veriﬁcation of Theorem 3.1, that
Revised: X for DP-MHA f causes the norm to diverge, providing empirical veriﬁcation of Theorem 3.1, that

Original: DP-MHA is indeed not Lipschitz.
Revised: DP-MHA is indeed not Lipschitz.

Original: 5. N UMERICAL INVERTIBILITY OF MHA RESIDUAL MAP
Revised: 5. N UMERICAL INVERTIBILITY OF MHA RESIDUAL MAP

Original: [Figure]
Revised: [Figure]

Original: Recall from Section 4.1 that g ( x ) = x
Revised: Recall from Section 4.1 that g ( x ) = x

Original: + f ( x ) is invertible if f is contractive.
Revised: + f ( x ) is invertible if f is contractive.

Original: Hence if f is
Revised: Hence if f is

Original: Contractive-L2-MHA , g is necessarily invertible.
Revised: Contractive-L2-MHA , g is necessarily invertible.

Original: However, technically we do not disprove the invertibility of DP-MHA , since the converse does not hold in general i.e. if f is
Revised: However, technically we do not disprove the invertibility of DP-MHA , since the converse does not hold in general i.e. if f is

Original: DP-MHA , which we have shown is not Lipschitz hence not contractive, it may still be the case that g is invertible.
Revised: DP-MHA , which we have shown is not Lipschitz hence not contractive, it may still be the case that g is invertible.

Original: To verify that DP-MHA (with the skip connection) is not invertible in practice, we compare the numerical invertibility of the residual map g ( x ) = x + cf ( x ) between the cases where f is L2-MHA and DP-MHA in Figure 3.
Revised: To verify that DP-MHA (with the skip connection) is not invertible in practice, we compare the numerical invertibility of the residual map g ( x ) = x + cf ( x ) between the cases where f is L2-MHA and DP-MHA in Figure 3.

Original: For each, we take MHA with 8 heads and randomly initialised weights, and quantify the maximum reconstruction error across a batch of 128 inputs whose outputs are inverted via the ﬁxed-point iteration described in Section 4.1.
Revised: For each, we take MHA with 8 heads and randomly initialised weights, and quantify the maximum reconstruction error across a batch of 128 inputs whose outputs are inverted via the ﬁxed-point iteration described in Section 4.1.

Original: We use N = 64 , D = 64 , and c ∈ { 0 . 5 , 0 .
Revised: We use N = 64 , D = 64 , and c ∈ { 0 . 5 , 0 .

Original: 9 } (see Appendix J for analogous results for a wider range of N and D and for DP-MHA with trained weights).
Revised: 9 } (see Appendix J for analogous results for a wider range of N and D and for DP-MHA with trained weights).

Original: To highlight the difference between the two types of self-attention, recall in the proof of
Revised: To highlight the difference between the two types of self-attention, recall in the proof of

Original: Theorem 3.1 (showing that DP-MHA is not Lipschitz) that when one of the inputs x i is 0 , some terms of the Jacobian grow with the sample variance of x (cid:54) =
Revised: Theorem 3.1 (showing that DP-MHA is not Lipschitz) that when one of the inputs x i is 0 , some terms of the Jacobian grow with the sample variance of x (cid:54) =

Original: i .
Revised: i .

Original: Hence we check numerical invertibility at a set of N inputs where x i = 0 and x (cid:54)
Revised: Hence we check numerical invertibility at a set of N inputs where x i = 0 and x (cid:54)

Original: =
Revised: =

Original: i are chosen uniformly at random.
Revised: i are chosen uniformly at random.

Original: In Figure 3, we see that DP-MHA is not invertible whereas L2-MHA is invertible for sufﬁciently small c .
Revised: In Figure 3, we see that DP-MHA is not invertible whereas L2-MHA is invertible for sufﬁciently small c .

Original: This shows how not having the theoretical guarantee of f being contractive can cost us invertibility in practice.
Revised: This shows how not having the theoretical guarantee of f being contractive can cost us invertibility in practice.

Original: We note that the ﬁgure shows local invertibility at the sampled inputs, as opposed to global invertibility across the whole input space, yet this clearly highlights the difference between the two choices of self-attention.
Revised: We note that the ﬁgure shows local invertibility at the sampled inputs, as opposed to global invertibility across the whole input space, yet this clearly highlights the difference between the two choices of self-attention.

Original: Experiments with the globally invertible self-attention obtained by normalising with the Lipschitz upper bound are provided in the next section.
Revised: Experiments with the globally invertible self-attention obtained by normalising with the Lipschitz upper bound are provided in the next section.

Original: 5. E XPRESSIVENESS OF L2-MHA AND INVERTIBLE SELF - ATTENTION
Revised: 5. E XPRESSIVENESS OF L2-MHA AND INVERTIBLE SELF - ATTENTION

Original: A natural question to ask is: how does the expressiveness of L2-MHA and Contractive-L2-MHA (that leads to invertible self-attention with the skip connection) compare with the original DP-MHA ?
Revised: A natural question to ask is: how does the expressiveness of L2-MHA and Contractive-L2-MHA (that leads to invertible self-attention with the skip connection) compare with the original DP-MHA ?

Original: We expect that the Lipschitz constraint will limit the expressiveness of the Transformer, and would like to ﬁnd out by how much.
Revised: We expect that the Lipschitz constraint will limit the expressiveness of the Transformer, and would like to ﬁnd out by how much.

Original: We investigate this by comparing the performance of the original
Revised: We investigate this by comparing the performance of the original

Original: Transformer and the Transformer with invertible self-attention (c.f. Figure 1) at character-level language modelling on the Penn Treebank dataset (Marcus et al., 1993).
Revised: Transformer and the Transformer with invertible self-attention (c.f. Figure 1) at character-level language modelling on the Penn Treebank dataset (Marcus et al., 1993).

Original: We compare the test negative log-likelihood (NLL) of a baseline LSTM, the original Transformer ( DP-MHA ), and a series of models between the original Transformer and the Transformer with invertible self-attention ( Contractive-L2-MHA ), making one change at a time and tuning the hyperparameters on a validation set.
Revised: We compare the test negative log-likelihood (NLL) of a baseline LSTM, the original Transformer ( DP-MHA ), and a series of models between the original Transformer and the Transformer with invertible self-attention ( Contractive-L2-MHA ), making one change at a time and tuning the hyperparameters on a validation set.

Original: For Contractive-L2-MHA , we normalise L2-MHA by the bound on Lip ∞ ( F ) as it is tighter than the bound on Lip 2 ( F ) .
Revised: For Contractive-L2-MHA , we normalise L2-MHA by the bound on Lip ∞ ( F ) as it is tighter than the bound on Lip 2 ( F ) .

Original: See Appendix I for experimental details
Revised: See Appendix I for experimental details

Original: [Figure]
Revised: [Figure]

Original: The results are shown in Figure 4.
Revised: The results are shown in Figure 4.

Original: The ﬁrst plot shows the best performing LSTM reaching a test
Revised: The ﬁrst plot shows the best performing LSTM reaching a test

Original: NLL of around 1 . 0 , and the second plot shows the best performing Transformer reaching a slightly improved performance for 3 – 5 layers of Transformer blocks.
Revised: NLL of around 1 . 0 , and the second plot shows the best performing Transformer reaching a slightly improved performance for 3 – 5 layers of Transformer blocks.

Original: We observe instabilities in training for a higher number of layers, requiring careful tuning of the learning rate schedule for stability at the cost of performance, a commonly observed phenomenon in the literature of deep Transformer architectures (Bapna et al., 2018; Parisotto et al., 2020).
Revised: We observe instabilities in training for a higher number of layers, requiring careful tuning of the learning rate schedule for stability at the cost of performance, a commonly observed phenomenon in the literature of deep Transformer architectures (Bapna et al., 2018; Parisotto et al., 2020).

Original: The third plot shows results for the Transformer with
Revised: The third plot shows results for the Transformer with

Original: DP-MHA replaced with L2-MHA but without tying W Q and W K , and we observe a very similar test performance.
Revised: DP-MHA replaced with L2-MHA but without tying W Q and W K , and we observe a very similar test performance.

Original: The fourth plot shows the change when we further tie the query and key weights (making W Q = W K ); we see that there is a small degradation in performance.
Revised: The fourth plot shows the change when we further tie the query and key weights (making W Q = W K ); we see that there is a small degradation in performance.

Original: Here the number of trainable parameters has been reduced, but in Appendix M we show that matching parameter count does not help performance, suggesting that the reduction in performance when tying queries and keys is not solely due to having fewer parameters.
Revised: Here the number of trainable parameters has been reduced, but in Appendix M we show that matching parameter count does not help performance, suggesting that the reduction in performance when tying queries and keys is not solely due to having fewer parameters.

Original: We note that performance saturates at around 5 layers for each Transformer model so far.
Revised: We note that performance saturates at around 5 layers for each Transformer model so far.

Original: On the rightmost plot we show results when further dividing self-attention in each block by the upper bound on Lip ∞ ( F ) , to obtain invertible self-attention.
Revised: On the rightmost plot we show results when further dividing self-attention in each block by the upper bound on Lip ∞ ( F ) , to obtain invertible self-attention.

Original: This does give reduced performance for the same number of layers, but we can attain similar performance with more layers, no longer saturating at 5 layers.
Revised: This does give reduced performance for the same number of layers, but we can attain similar performance with more layers, no longer saturating at 5 layers.

Original: Thus we conclude the following.
Revised: Thus we conclude the following.

Original: (1) Replacing the dot-product with the L2 distance incurs hardly any loss in expressiveness.
Revised: (1) Replacing the dot-product with the L2 distance incurs hardly any loss in expressiveness.

Original: (2) Tying the query and key weights to obtain Lipschitz self-attention incurs a small loss in expressiveness.
Revised: (2) Tying the query and key weights to obtain Lipschitz self-attention incurs a small loss in expressiveness.

Original: (3) Dividing by the upper bound on Lip ∞ ( F ) to obtain invertible self-attention incurs a noticeable loss in expressiveness, but also has a stabilising effect on the optimisation of the Transformer, thus allowing one to compensate for the apparent loss in expressiveness by increasing the number of layers.
Revised: (3) Dividing by the upper bound on Lip ∞ ( F ) to obtain invertible self-attention incurs a noticeable loss in expressiveness, but also has a stabilising effect on the optimisation of the Transformer, thus allowing one to compensate for the apparent loss in expressiveness by increasing the number of layers.

Original: We show further experimental results that compare the training stability of DP-MHA and (Contractive)-L2-MHA in Appendix N.
Revised: We show further experimental results that compare the training stability of DP-MHA and (Contractive)-L2-MHA in Appendix N.

Original:  C ONCLUSION AND D ISCUSSION
Revised:  C ONCLUSION AND D ISCUSSION

Original: We have shown that the widely used dot-product self-attention is not Lipschitz, and that the proposed
Revised: We have shown that the widely used dot-product self-attention is not Lipschitz, and that the proposed

Original: L2 self-attention is Lipschitz, by deriving an O (log N − log log N ) Lipschitz bound for p = ∞ and
Revised: L2 self-attention is Lipschitz, by deriving an O (log N − log log N ) Lipschitz bound for p = ∞ and

Original: an
Revised: an

Original: O ( √ N (log N − log log N )) bound for p = 2 , where N is the input sequence length.
Revised: O ( √ N (log N − log log N )) bound for p = 2 , where N is the input sequence length.

Original: We also provided empirical evidence of the asymptotic tightness of the bound for p = ∞ .
Revised: We also provided empirical evidence of the asymptotic tightness of the bound for p = ∞ .

Original: Finally we demonstrated that Lipschitz-constrained self-attention can be used to formulate invertible self-attention, which we experimentally evaluated on a character-level language modelling task.
Revised: Finally we demonstrated that Lipschitz-constrained self-attention can be used to formulate invertible self-attention, which we experimentally evaluated on a character-level language modelling task.

Original: Our approach to Lipschitz self-attention has been to replace the dot-product kernel with an L2 kernel.
Revised: Our approach to Lipschitz self-attention has been to replace the dot-product kernel with an L2 kernel.

Original: An alternative would be to constrain the inputs of self-attention to be bounded; if the input space is compact, e.g. [0 , 1] N × D , any continuously differentiable function is Lipschitz, including dot-product self-attention.
Revised: An alternative would be to constrain the inputs of self-attention to be bounded; if the input space is compact, e.g. [0 , 1] N × D , any continuously differentiable function is Lipschitz, including dot-product self-attention.

Original: However, while being simple to implement, this solution has its own difﬁculties.
Revised: However, while being simple to implement, this solution has its own difﬁculties.

Original: First, it makes the Lipschitz constant depend on the range of the input, and thus obtaining a tight bound would require non-trivial mathematical work.
Revised: First, it makes the Lipschitz constant depend on the range of the input, and thus obtaining a tight bound would require non-trivial mathematical work.

Original: We stress that a guarantee that the function is
Revised: We stress that a guarantee that the function is

Original: Lipschitz does not tell us anything about its Lipschitz constant; without a tight Lipschitz bound, the true Lipschitz constant can be very large, at which point it is unhelpful that the function is Lipschitz.
Revised: Lipschitz does not tell us anything about its Lipschitz constant; without a tight Lipschitz bound, the true Lipschitz constant can be very large, at which point it is unhelpful that the function is Lipschitz.

Original: Second, since self-attention is typically applied at multiple layers within a model (e.g. Transformer), the input to each self-attention will live in a different compact set that depends on the parameters of the previous layers, complicating the analysis for subsequent layers.
Revised: Second, since self-attention is typically applied at multiple layers within a model (e.g. Transformer), the input to each self-attention will live in a different compact set that depends on the parameters of the previous layers, complicating the analysis for subsequent layers.

Original: A solution is to constrain the inputs of each layer to be in the same compact set, e.g. by passing them through a sigmoid non-linearity.
Revised: A solution is to constrain the inputs of each layer to be in the same compact set, e.g. by passing them through a sigmoid non-linearity.

Original: This however can have undesirable side effects such as vanishing gradients when the sigmoids are saturated (Hochreiter, 1998).
Revised: This however can have undesirable side effects such as vanishing gradients when the sigmoids are saturated (Hochreiter, 1998).

Original: Despite these difﬁculties, this could be a worthwhile alternative route for obtaining Lipschitz self-attention to explore in the future.
Revised: Despite these difﬁculties, this could be a worthwhile alternative route for obtaining Lipschitz self-attention to explore in the future.

Original: Having a provably Lipschitz self-attention module at our disposal makes it possible to use Transformerbased architectures in applications requiring Lipschitz constraints, while enjoying theoretical guarantees.
Revised: Having a provably Lipschitz self-attention module at our disposal makes it possible to use Transformerbased architectures in applications requiring Lipschitz constraints, while enjoying theoretical guarantees.

Original: A natural application of Lipschitz self-attention is for residual ﬂows (Behrmann et al., 2019), and for parameterising Neural ODEs (Chen et al., 2018) where a Lipschitz vector ﬁeld guarantees the existence of a unique solution to the ODE for all times.
Revised: A natural application of Lipschitz self-attention is for residual ﬂows (Behrmann et al., 2019), and for parameterising Neural ODEs (Chen et al., 2018) where a Lipschitz vector ﬁeld guarantees the existence of a unique solution to the ODE for all times.

Original: These models can be used for density estimation and generative modelling of sets.
Revised: These models can be used for density estimation and generative modelling of sets.

Original: Another interesting direction for future work would be to analyse different variants of self-attention based on kernels other than dot-product and L2, as
Revised: Another interesting direction for future work would be to analyse different variants of self-attention based on kernels other than dot-product and L2, as

Original: Tsai et al.
Revised: Tsai et al.

Original: (2019) do from an experimental perspective, for which we believe the mathematical tools developed in this paper may aid the analysis.
Revised: (2019) do from an experimental perspective, for which we believe the mathematical tools developed in this paper may aid the analysis.

Original: A C HAIN R ULE FOR VECTOR VALUED F UNCTIONS
Revised: A C HAIN R ULE FOR VECTOR VALUED F UNCTIONS

Original: In this section, we list some useful identities for deriving the Jacobians of the expressions in the paper.
Revised: In this section, we list some useful identities for deriving the Jacobians of the expressions in the paper.

Original: Suppose λ is a scalar, u, v, x are column vectors, and f ( u ) is a vector valued function.
Revised: Suppose λ is a scalar, u, v, x are column vectors, and f ( u ) is a vector valued function.

Original: We use the standard convention that for a ∈ R m , b ∈ R n , we have ∂a∂b ∈ R m × n .
Revised: We use the standard convention that for a ∈ R m , b ∈ R n , we have ∂a∂b ∈ R m × n .

Original: Then we have the following chain rule identities:
Revised: Then we have the following chain rule identities:

Original: • ∂∂x
Revised: • ∂∂x

Original: [ λu ] = λ
Revised: [ λu ] = λ

Original: [Equation]
Revised: [Equation]

Original: +
Revised: +

Original: [Equation]
Revised: [Equation]

Original: • ∂f ( u ) ∂x =
Revised: • ∂f ( u ) ∂x =

Original: ∂f ( u ) ∂u ∂u∂x • ∂∂x
Revised: ∂f ( u ) ∂u ∂u∂x • ∂∂x

Original: [ u (cid:62) v ] =
Revised: [ u (cid:62) v ] =

Original: u (cid:62) ∂v∂x + v (cid:62) ∂u∂x
Revised: u (cid:62) ∂v∂x + v (cid:62) ∂u∂x

Original: Note ∂λ∂x is a row vector, so u ∂λ∂x is a matrix.
Revised: Note ∂λ∂x is a row vector, so u ∂λ∂x is a matrix.

Original: B P OWER I TERATION
Revised: B P OWER I TERATION

Original: Although (cid:107) W (cid:107) ∞ can be computed efﬁciently in O ( nm ) time for W ∈ R m × n , naïvely computing (cid:107) W (cid:107) 2 = σ max ( W ) :
Revised: Although (cid:107) W (cid:107) ∞ can be computed efﬁciently in O ( nm ) time for W ∈ R m × n , naïvely computing (cid:107) W (cid:107) 2 = σ max ( W ) :

Original: = (cid:112) λ max ( W (cid:62) W ) requires O ( n 3 ) operations.
Revised: = (cid:112) λ max ( W (cid:62) W ) requires O ( n 3 ) operations.

Original: (By λ max ( A ) we denote the greatest eigenvalue of a symmetric matrix A .)
Revised: (By λ max ( A ) we denote the greatest eigenvalue of a symmetric matrix A .)

Original: We can however obtain an underestimate ˜ σ ( W ) via power iteration :
Revised: We can however obtain an underestimate ˜ σ ( W ) via power iteration :

Original: [Equation]
Revised: [Equation]

Original: with each iteration taking O ( n 2 ) time.
Revised: with each iteration taking O ( n 2 ) time.

Original: Then using K (cid:28) n iterations gives us an underestimate ˜ σ K in
Revised: Then using K (cid:28) n iterations gives us an underestimate ˜ σ K in

Original: O ( Kn 2 ) time.
Revised: O ( Kn 2 ) time.

Original: Since this is an underestimate, the resulting approximation to the Lipschitz constant of the linear map will not be an upper bound.
Revised: Since this is an underestimate, the resulting approximation to the Lipschitz constant of the linear map will not be an upper bound.

Original: However the number of power iterations is usually chosen so that ˜ σ is accurate enough — K = 5 is shown to be sufﬁcient in the context of fully connected networks or convolutions considered by Behrmann et al.
Revised: However the number of power iterations is usually chosen so that ˜ σ is accurate enough — K = 5 is shown to be sufﬁcient in the context of fully connected networks or convolutions considered by Behrmann et al.

Original: The iteration will converge if W (cid:62) W has an eigenvalue that is strictly greater in magnitude than its other eigenvalues, and the starting vector b 0 has a nonzero component in the direction of an eigenvector associated with the dominant eigenvalue.
Revised: The iteration will converge if W (cid:62) W has an eigenvalue that is strictly greater in magnitude than its other eigenvalues, and the starting vector b 0 has a nonzero component in the direction of an eigenvector associated with the dominant eigenvalue.

Original: This happens with probability 1 if b 0 is chosen at random, and the convergence is geometric with ratio | λ 2 /λ max | where λ 2 is the eigenvalue with second largest magnitude (Mises & Pollaczek-Geiringer, 1929).
Revised: This happens with probability 1 if b 0 is chosen at random, and the convergence is geometric with ratio | λ 2 /λ max | where λ 2 is the eigenvalue with second largest magnitude (Mises & Pollaczek-Geiringer, 1929).

Original: C P ROOF OF T HEOREM 3.1 FOR G ENERAL D
Revised: C P ROOF OF T HEOREM 3.1 FOR G ENERAL D

Original: Theorem 3.1.
Revised: Theorem 3.1.

Original: DP-MHA is not Lipschitz for any vector p -norm (cid:107) · (cid:107) p with p ∈ [1 , ∞ ] .
Revised: DP-MHA is not Lipschitz for any vector p -norm (cid:107) · (cid:107) p with p ∈ [1 , ∞ ] .

Original: Proof.
Revised: Proof.

Original: The mapping f can be written as
Revised: The mapping f can be written as

Original: [Equation]
Revised: [Equation]

Original: (12) where A = W K W Q (cid:62) / (cid:112) D/H ∈ R D × D and f i ( X ) = (cid:80) Nj =
Revised: (12) where A = W K W Q (cid:62) / (cid:112) D/H ∈ R D × D and f i ( X ) = (cid:80) Nj =

Original: 1 P ij x j with P (cid:62) i : = softmax ( XAx i ) .
Revised: 1 P ij x j with P (cid:62) i : = softmax ( XAx i ) .

Original: Hence f can be interpreted as a map of each x
Revised: Hence f can be interpreted as a map of each x

Original: i to a point in the convex hull of x 1 , ..., x N .
Revised: i to a point in the convex hull of x 1 , ..., x N .

Original: Since f is a map from R N × D to R N × D , its Jacobian is
Revised: Since f is a map from R N × D to R N × D , its Jacobian is

Original: [Equation]
Revised: [Equation]

Original: (13) where J ij =
Revised: (13) where J ij =

Original: ∂f i ( X ) ∂x j ∈ R D × D .
Revised: ∂f i ( X ) ∂x j ∈ R D × D .

Original: By taking partial derivatives we can show that J ij =
Revised: By taking partial derivatives we can show that J ij =

Original: X (cid:62) P ( i ) (cid:2)
Revised: X (cid:62) P ( i ) (cid:2)

Original: e ji XA (cid:62) + XAδ ij (cid:3)
Revised: e ji XA (cid:62) + XAδ ij (cid:3)

Original: + P ij I where e ij ∈ R N × N is a binary matrix with zeros everywhere except the ( i, j ) th entry, δ ij is the Kronecker delta, and P ( i ) :
Revised: + P ij I where e ij ∈ R N × N is a binary matrix with zeros everywhere except the ( i, j ) th entry, δ ij is the Kronecker delta, and P ( i ) :

Original: =
Revised: =

Original: diag( P i : ) − P (cid:62) i : P i : .
Revised: diag( P i : ) − P (cid:62) i : P i : .

Original: So for i =
Revised: So for i =

Original: j :
Revised: j :

Original: [Equation]
Revised: [Equation]

Original: For the last equality, note e ii X has all rows equal to zero except for the i th row given by x (cid:62) i .
Revised: For the last equality, note e ii X has all rows equal to zero except for the i th row given by x (cid:62) i .

Original: We can then verify that X (cid:62) P ( i ) e ii X simpliﬁes to P ii ( x i − (cid:80)
Revised: We can then verify that X (cid:62) P ( i ) e ii X simpliﬁes to P ii ( x i − (cid:80)

Original: k P ik x k ) x
Revised: k P ik x k ) x

Original: (cid:62) i .
Revised: (cid:62) i .

Original: For vector p -norms, (cid:107) J f (cid:107) p is bounded if and only if its entries are bounded, by deﬁnition of the operator norm.
Revised: For vector p -norms, (cid:107) J f (cid:107) p is bounded if and only if its entries are bounded, by deﬁnition of the operator norm.

Original: The entries of X (cid:62) P ( i ) XA are bounded for arbitrary A only if the entries of X (cid:62) P ( i ) X are bounded.
Revised: The entries of X (cid:62) P ( i ) XA are bounded for arbitrary A only if the entries of X (cid:62) P ( i ) X are bounded.

Original: So let us investigate the entries of this D × D matrix.
Revised: So let us investigate the entries of this D × D matrix.

Original: Writing out each term of the matrix, we observe that it is in fact a covariance matrix of a discrete distribution.
Revised: Writing out each term of the matrix, we observe that it is in fact a covariance matrix of a discrete distribution.

Original: Speciﬁcally:
Revised: Speciﬁcally:

Original: [Equation]
Revised: [Equation]

Original: (15) where X is a discrete distribution with support at the inputs { x 1 , . .
Revised: (15) where X is a discrete distribution with support at the inputs { x 1 , . .

Original: , x N } and probability mass function given by their softmax probabilities P ( X = x j )
Revised: , x N } and probability mass function given by their softmax probabilities P ( X = x j )

Original: = P ij .
Revised: = P ij .

Original: A consequence of this interpretation is that P ( i ) is positive semi-deﬁnite (PSD) since for D = 1 , Equation (15) becomes X (cid:62) P ( i ) X =
Revised: A consequence of this interpretation is that P ( i ) is positive semi-deﬁnite (PSD) since for D = 1 , Equation (15) becomes X (cid:62) P ( i ) X =

Original: Var( X ) ≥ 0 , with equality if and only if the x j are all equal.
Revised: Var( X ) ≥ 0 , with equality if and only if the x j are all equal.

Original: We use this observation to show that the terms of J ii are unbounded, and so DP-MHA is not Lipschitz.
Revised: We use this observation to show that the terms of J ii are unbounded, and so DP-MHA is not Lipschitz.

Original: Consider the case x i = 0 .
Revised: Consider the case x i = 0 .

Original: Then P (cid:62) i : = softmax ( XAx i ) = 1 N 1 , i.e. we have uniform attention regardless of x (cid:54)
Revised: Then P (cid:62) i : = softmax ( XAx i ) = 1 N 1 , i.e. we have uniform attention regardless of x (cid:54)

Original: =
Revised: =

Original: i .
Revised: i .

Original: The ﬁrst term of J ii in Equation (14) disappears since x i = 0 , and the last term becomes 1 N I .
Revised: The ﬁrst term of J ii in Equation (14) disappears since x i = 0 , and the last term becomes 1 N I .

Original: For the second term, the entries [ X (cid:62) P ( i ) X ] ll =
Revised: For the second term, the entries [ X (cid:62) P ( i ) X ] ll =

Original: Var( X l ) are unbounded since the latter is equal to the sample variance of x 1 l , . .
Revised: Var( X l ) are unbounded since the latter is equal to the sample variance of x 1 l , . .

Original: , x Nl , which can be arbitrarily large.
Revised: , x Nl , which can be arbitrarily large.

Original: D B IAS TERM IN DP S ELF -A TTENTION
Revised: D B IAS TERM IN DP S ELF -A TTENTION

Original: A natural question to ask is whether we can add bias terms b Q to x (cid:62) i W Q and b K to x (cid:62) j W K to resolve the issue of attention weights
Revised: A natural question to ask is whether we can add bias terms b Q to x (cid:62) i W Q and b K to x (cid:62) j W K to resolve the issue of attention weights

Original: P i : becoming uniform when x i = 0 .
Revised: P i : becoming uniform when x i = 0 .

Original: The answer is no in general.
Revised: The answer is no in general.

Original: It can again be shown that J ii is unbounded when x i is chosen such that x (cid:62) i W Q + b Q
Revised: It can again be shown that J ii is unbounded when x i is chosen such that x (cid:62) i W Q + b Q

Original: = 0 (such a choice is possible assuming W Q is full rank, a dense set in R D × D/H ).
Revised: = 0 (such a choice is possible assuming W Q is full rank, a dense set in R D × D/H ).

Original: Then P (cid:62) i : = 1 N 1 again, and the diagonal entries of X (cid:62) P ( i ) X are unbounded.
Revised: Then P (cid:62) i : = 1 N 1 again, and the diagonal entries of X (cid:62) P ( i ) X are unbounded.

Original: E E FFICIENT C OMPUTATION OF L2 S ELF -A TTENTION
Revised: E E FFICIENT C OMPUTATION OF L2 S ELF -A TTENTION

Original: Dot-product self-attention only requires a few matrix multiplications to compute the logits (i.e. the inputs to the softmax) between all pairs of inputs, without having to loop over pairs, hence it can be computed efﬁciently.
Revised: Dot-product self-attention only requires a few matrix multiplications to compute the logits (i.e. the inputs to the softmax) between all pairs of inputs, without having to loop over pairs, hence it can be computed efﬁciently.

Original: Similarly, we can show that L2 self-attention can also be computed in an efﬁcient manner.
Revised: Similarly, we can show that L2 self-attention can also be computed in an efﬁcient manner.

Original: Using the identity (cid:107) a − b (cid:107)
Revised: Using the identity (cid:107) a − b (cid:107)

Original: [Equation]
Revised: [Equation]

Original: −
Revised: −

Original: [Equation]
Revised: [Equation]

Original: +
Revised: +

Original: [Equation]
Revised: [Equation]

Original: b
Revised: b

Original: [Equation]
Revised: [Equation]

Original: 22 we can compute the logits of
Revised: 22 we can compute the logits of

Original: L2 attention between all pairs via matrix multiplications and computation of row-wise L2 norms, with negligible overhead compared to dot-product self-attention.
Revised: L2 attention between all pairs via matrix multiplications and computation of row-wise L2 norms, with negligible overhead compared to dot-product self-attention.

Original: Speciﬁcally, for L2 self-attention we can show that
Revised: Speciﬁcally, for L2 self-attention we can show that

Original: [Equation]
Revised: [Equation]

Original: where (cid:107) A (cid:107) 2 row applies the squared L2 norm to each row of A , so if A ∈ R m × n then (cid:107) A (cid:107) 2 row ∈
Revised: where (cid:107) A (cid:107) 2 row applies the squared L2 norm to each row of A , so if A ∈ R m × n then (cid:107) A (cid:107) 2 row ∈

Original: R m .
Revised: R m .

Original: In Table 1 we show the wall-clock training times for the Transformer models with different attention functions and a varying number of layers.
Revised: In Table 1 we show the wall-clock training times for the Transformer models with different attention functions and a varying number of layers.

Original: It is evident that the differences between the models are rather small.
Revised: It is evident that the differences between the models are rather small.

Original: [Table]
Revised: [Table]

Original: F P ROOF OF T HEOREM 3. 
Revised: F P ROOF OF T HEOREM 3.

Original: Recall the formulation of L2-MHA :
Revised: Recall the formulation of L2-MHA :

Original: [Equation]
Revised: [Equation]

Original: j
Revised: j

Original: P hij
Revised: P hij

Original: [Equation]
Revised: [Equation]

Original: where we have that W Q,h , W K,h , W V,h ∈ R D × D/H , W O ∈ R D × D , P h ∈ R N × N and A h :=
Revised: where we have that W Q,h , W K,h , W V,h ∈ R D × D/H , W O ∈ R D × D , P h ∈ R N × N and A h :=

Original: W Q,h W Q,h (cid:62) /
Revised: W Q,h W Q,h (cid:62) /

Original: (cid:112) D/H ∈ R D × D , and the softmax is applied to each row of the input matrix.
Revised: (cid:112) D/H ∈ R D × D , and the softmax is applied to each row of the input matrix.

Original: Recall
Revised: Recall

Original: Equation (16):
Revised: Equation (16):

Original: [Equation]
Revised: [Equation]

Original: F.
Revised: F.

Original: L2 SELF - ATTENTION IS not L IPSCHITZ FOR
Revised: L2 SELF - ATTENTION IS not L IPSCHITZ FOR

Original: [Equation]
Revised: [Equation]

Original: W K
Revised: W K

Original: Let us ﬁrst look at the case of H = 1 and suppress the index h to reduce clutter.
Revised: Let us ﬁrst look at the case of H = 1 and suppress the index h to reduce clutter.

Original: Consider the map ˜ f
Revised: Consider the map ˜ f

Original: ( X ) :
Revised: ( X ) :

Original: = PX , so f ( X ) =
Revised: = PX , so f ( X ) =

Original: ˜ f
Revised: ˜ f

Original: ( X ) A .
Revised: ( X ) A .

Original: We need ˜ f to be Lipschitz for f and hence F to be Lipschitz.
Revised: We need ˜ f to be Lipschitz for f and hence F to be Lipschitz.

Original: Note that P is deﬁned as:
Revised: Note that P is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: and the normalisation constant satisﬁes (cid:80) j P ij = 1 , for P ∈ R N × N , X ∈ R N × D .
Revised: and the normalisation constant satisﬁes (cid:80) j P ij = 1 , for P ∈ R N × N , X ∈ R N × D .

Original: For L2 self-attention, we may take partial derivatives and use the chain rule to show that the Jacobian of ˜ f is:
Revised: For L2 self-attention, we may take partial derivatives and use the chain rule to show that the Jacobian of ˜ f is:

Original: [Equation]
Revised: [Equation]

Original: 17) with
Revised: 17) with

Original: [Equation]
Revised: [Equation]

Original: 18) where ∂L i : ∂x
Revised: 18) where ∂L i : ∂x

Original: [Equation]
Revised: [Equation]

Original: 19) and
Revised: 19) and

Original: P ( i
Revised: P ( i

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: P
Revised: P

Original: [Equation]
Revised: [Equation]

Original: −
Revised: −

Original: [Equation]
Revised: [Equation]

Original: P iN .
Revised: P iN .

Original: − P iN P i− P iN P i. . .
Revised: − P iN P i− P iN P i. . .

Original: P iN (1 − P iN ) 
Revised: P iN (1 − P iN ) 

Original:  ,
Revised:  ,

Original: [Equation]
Revised: [Equation]

Original: Recall that e ji ∈ R N × N is a binary matrix with zeros everywhere except the ( j, i ) th entry.
Revised: Recall that e ji ∈ R N × N is a binary matrix with zeros everywhere except the ( j, i ) th entry.

Original: Hence e ji X has all rows equal to zero except for the j th row given by x (cid:62) i .
Revised: Hence e ji X has all rows equal to zero except for the j th row given by x (cid:62) i .

Original: We can then verify:
Revised: We can then verify:

Original: [Equation]
Revised: [Equation]

Original: Also note P ( i ) is symmetric, and each row/colum sums to 0 , i.e. P ( i ) 1 = 1 (cid:62) P ( i ) = 0 .
Revised: Also note P ( i ) is symmetric, and each row/colum sums to 0 , i.e. P ( i ) 1 = 1 (cid:62) P ( i ) = 0 .

Original: Hence we may simplify the Jacobian terms as follows:
Revised: Hence we may simplify the Jacobian terms as follows:

Original: [Equation]
Revised: [Equation]

Original: and for i (cid:54) =
Revised: and for i (cid:54) =

Original: j :
Revised: j :

Original: [Equation]
Revised: [Equation]

Original: We are now ready to show that ˜ f is not Lipschitz for general W Q , W K :
Revised: We are now ready to show that ˜ f is not Lipschitz for general W Q , W K :

Original: Lemma F.1.
Revised: Lemma F.1.

Original: If W K ∈ R D × D/H is full rank (i.e. full column rank), and W K (cid:54) =
Revised: If W K ∈ R D × D/H is full rank (i.e. full column rank), and W K (cid:54) =

Original: W Q , then J ij has terms that are unbounded for i (cid:54) =
Revised: W Q , then J ij has terms that are unbounded for i (cid:54) =

Original: j , hence ˜ f is not Lipschitz.
Revised: j , hence ˜ f is not Lipschitz.

Original: Proof.
Revised: Proof.

Original: Let us investigate the expression ˜ K ij :
Revised: Let us investigate the expression ˜ K ij :

Original: [Equation]
Revised: [Equation]

Original: P
Revised: P

Original: [Equation]
Revised: [Equation]

Original: ) ∈
Revised: ) ∈

Original: R DH × DH for i (cid:54) =
Revised: R DH × DH for i (cid:54) =

Original: j , which is related to ˜ J ij as follows by Equation (22):
Revised: j , which is related to ˜ J ij as follows by Equation (22):

Original: [Equation]
Revised: [Equation]

Original: It sufﬁces to show that ˜ K ij is unbounded to show that ˜ J ij is unbounded, since W K is full rank and
Revised: It sufﬁces to show that ˜ K ij is unbounded to show that ˜ J ij is unbounded, since W K is full rank and

Original: P ij ∈ [0 , 1] .
Revised: P ij ∈ [0 , 1] .

Original: Let y (cid:62)
Revised: Let y (cid:62)

Original: j
Revised: j

Original: = x (cid:62)
Revised: = x (cid:62)

Original: i W Q − x (cid:62) j W K .
Revised: i W Q − x (cid:62) j W K .

Original: Then we have: y
Revised: Then we have: y

Original: [Equation]
Revised: [Equation]

Original: Hence ˜ K ij = − P ij ( y j − (cid:80) k
Revised: Hence ˜ K ij = − P ij ( y j − (cid:80) k

Original: P ik y k )
Revised: P ik y k )

Original: y (cid:62) j .
Revised: y (cid:62) j .

Original: Note y
Revised: Note y

Original: i can take an arbitrary value in R D/H , since
Revised: i can take an arbitrary value in R D/H , since

Original: W K (cid:54) =
Revised: W K (cid:54) =

Original: W Q and W K is full-rank.
Revised: W Q and W K is full-rank.

Original: For all j (cid:54) =
Revised: For all j (cid:54) =

Original: i , let us choose x j such that y j = − y i .
Revised: i , let us choose x j such that y j = − y i .

Original: This is possible for any value of y i since
Revised: This is possible for any value of y i since

Original: W K is full-rank.
Revised: W K is full-rank.

Original: Note y j = − y i and not
Revised: Note y j = − y i and not

Original: y i .
Revised: y i .

Original: We then have that (cid:107) y j (cid:107) 2 2 is equal for all j , hence
Revised: We then have that (cid:107) y j (cid:107) 2 2 is equal for all j , hence

Original: [Equation]
Revised: [Equation]

Original: 22 ) = 1 N for all j .
Revised: 22 ) = 1 N for all j .

Original: Then for i (cid:54) =
Revised: Then for i (cid:54) =

Original: j , ˜ K ij simpliﬁes to
Revised: j , ˜ K ij simpliﬁes to

Original: [Equation]
Revised: [Equation]

Original: whose entries are unbounded since y i can be any vector in R D/H (note we assume N ≥ 2 for self-attention to be well-deﬁned, hence 2 N − 2 (cid:54)
Revised: whose entries are unbounded since y i can be any vector in R D/H (note we assume N ≥ 2 for self-attention to be well-deﬁned, hence 2 N − 2 (cid:54)

Original: = 0 ).
Revised: = 0 ).

Original: The intuition for this result is as follows: a reason for DP-MHA not being Lipschitz is that for x i = 0 „
Revised: The intuition for this result is as follows: a reason for DP-MHA not being Lipschitz is that for x i = 0 „

Original: the attention weights P ij become uniform regardless of the values of x j for j (cid:54)
Revised: the attention weights P ij become uniform regardless of the values of x j for j (cid:54)

Original: =
Revised: =

Original: i .
Revised: i .

Original: A similar issue arises for L2-MHA with W Q (cid:54) =
Revised: A similar issue arises for L2-MHA with W Q (cid:54) =

Original: W K and full-rank W K , as shown above: given any x i , we can choose x j such that the P ij become uniform.
Revised: W K and full-rank W K , as shown above: given any x i , we can choose x j such that the P ij become uniform.

Original: F. L2 SELF - ATTENTION IS L IPSCHITZ FOR W Q = W K
Revised: F. L2 SELF - ATTENTION IS L IPSCHITZ FOR W Q = W K

Original: Hence we impose the restriction that W K = W Q .
Revised: Hence we impose the restriction that W K = W Q .

Original: With this assumption we have
Revised: With this assumption we have

Original: [Equation]
Revised: [Equation]

Original: (23) where A = W Q W Q (cid:62) / (cid:112) D/H ∈ R D × D and √ A is chosen such that A = √ A √ A (cid:62) , in particular √
Revised: (23) where A = W Q W Q (cid:62) / (cid:112) D/H ∈ R D × D and √ A is chosen such that A = √ A √ A (cid:62) , in particular √

Original: A := W Q / ( D/H ) 14 .
Revised: A := W Q / ( D/H ) 14 .

Original: The terms in the Jacobian of ˜ f simplify to:
Revised: The terms in the Jacobian of ˜ f simplify to:

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: for i (cid:54) =
Revised: for i (cid:54) =

Original: j
Revised: j

Original: [Equation]
Revised: [Equation]

Original: Let the Jacobian of f ( X ) be:
Revised: Let the Jacobian of f ( X ) be:

Original: J f =
Revised: J f =

Original: [Equation]
Revised: [Equation]

Original: Since f ( X ) =
Revised: Since f ( X ) =

Original: ˜ f
Revised: ˜ f

Original: ( X ) A , and by the chain rule ∂∂x j
Revised: ( X ) A , and by the chain rule ∂∂x j

Original: [ ˜ f i ( X ) A ]
Revised: [ ˜ f i ( X ) A ]

Original: = A (cid:62) ∂
Revised: = A (cid:62) ∂

Original: ˜ f i ( X ) ∂x j = A ∂ ˜ f
Revised: ˜ f i ( X ) ∂x j = A ∂ ˜ f

Original: i ( X ) ∂x j (by symmetry of A ), we have that J ij =
Revised: i ( X ) ∂x j (by symmetry of A ), we have that J ij =

Original: A ˜ J ij .
Revised: A ˜ J ij .

Original: Hence
Revised: Hence

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Noting Lip p ( f ) =
Revised: Noting Lip p ( f ) =

Original: sup X (cid:107) J f ( X ) (cid:107) p , we would like to upper bound (cid:107) J f (cid:107) p .
Revised: sup X (cid:107) J f ( X ) (cid:107) p , we would like to upper bound (cid:107) J f (cid:107) p .

Original: F.2. U PPER BOUND ON Lip ∞ ( F ) FOR L2-MHA
Revised: F.2. U PPER BOUND ON Lip ∞ ( F ) FOR L2-MHA

Original: Consider the choice p = ∞ , where (cid:107) J f (cid:107) ∞ is the maximum absolute row sum of J f .
Revised: Consider the choice p = ∞ , where (cid:107) J f (cid:107) ∞ is the maximum absolute row sum of J f .

Original: A key observation is that if we can bound the ∞ -norm of the Jacobian of f i , a single output of f (i.e. a single block row (cid:107) [ J i 1 , ..., J iN ] (cid:107) ∞ of J f ), then this is also a bound on (cid:107) J f (cid:107) ∞ due to permutation equivariance of self-attention; all block rows have the same maximal (cid:107) · (cid:107) ∞ when each is optimised over the input X .
Revised: A key observation is that if we can bound the ∞ -norm of the Jacobian of f i , a single output of f (i.e. a single block row (cid:107) [ J i 1 , ..., J iN ] (cid:107) ∞ of J f ), then this is also a bound on (cid:107) J f (cid:107) ∞ due to permutation equivariance of self-attention; all block rows have the same maximal (cid:107) · (cid:107) ∞ when each is optimised over the input X .

Original: Using this, we can prove that (cid:107) J f (cid:107) ∞ admits an upper bound that is O (log N − log log N ) .
Revised: Using this, we can prove that (cid:107) J f (cid:107) ∞ admits an upper bound that is O (log N − log log N ) .

Original: Below we state and prove lemmas that lead to the proof of this upper bound.
Revised: Below we state and prove lemmas that lead to the proof of this upper bound.

Original: First we analyse the term √ A (cid:62) X (cid:62) P ( i ) X √ A , that appears in the ﬁrst term of J ii .
Revised: First we analyse the term √ A (cid:62) X (cid:62) P ( i ) X √ A , that appears in the ﬁrst term of J ii .

Original: Note that for
Revised: Note that for

Original: Y :=
Revised: Y :=

Original: X √ A , so that the rows of Y are y (cid:62) i := x (cid:62) i √
Revised: X √ A , so that the rows of Y are y (cid:62) i := x (cid:62) i √

Original: A , we have
Revised: A , we have

Original: [Equation]
Revised: [Equation]

Original: (29) where P ( Y = y j ) =
Revised: (29) where P ( Y = y j ) =

Original: P ij = exp( −(cid:107)
Revised: P ij = exp( −(cid:107)

Original: y j − y i (cid:107) 22 ) / (cid:80) k exp( −(cid:107)
Revised: y j − y i (cid:107) 22 ) / (cid:80) k exp( −(cid:107)

Original: y k − y
Revised: y k − y

Original: i (cid:107) 22 ) .
Revised: i (cid:107) 22 ) .

Original: The last equality uses the observation in Equation (9).
Revised: The last equality uses the observation in Equation (9).

Original: The central inequality used throughout the proof of the main theorem is the following:
Revised: The central inequality used throughout the proof of the main theorem is the following:

Original: Lemma F.2.
Revised: Lemma F.2.

Original: Tr(Cov( Y ))
Revised: Tr(Cov( Y ))

Original: = (cid:80) j P ij (cid:107) y j − (cid:80) k
Revised: = (cid:80) j P ij (cid:107) y j − (cid:80) k

Original: P
Revised: P

Original: ik y k (cid:107) 22 ≤ (cid:80) j P ij (cid:107) y j − y
Revised: ik y k (cid:107) 22 ≤ (cid:80) j P ij (cid:107) y j − y

Original: i (cid:107) 22 ≤ φ − 1 ( N − 1) where φ ( c ) =
Revised: i (cid:107) 22 ≤ φ − 1 ( N − 1) where φ ( c ) =

Original: c exp( c + 1) is a one-dimensional invertible function on R ≥ 0 .
Revised: c exp( c + 1) is a one-dimensional invertible function on R ≥ 0 .

Original: Proof.
Revised: Proof.

Original: The ﬁrst equality holds since Tr(Cov( Y ))
Revised: The ﬁrst equality holds since Tr(Cov( Y ))

Original: = (cid:80) j Cov( Y ) jj
Revised: = (cid:80) j Cov( Y ) jj

Original: = (cid:80)
Revised: = (cid:80)

Original: j Var
Revised: j Var

Original: ( Y j ) = (cid:80)
Revised: ( Y j ) = (cid:80)

Original: j E
Revised: j E

Original: [( Y j −
Revised: [( Y j −

Original: E [ Y j ]) 2 ] .
Revised: E [ Y j ]) 2 ] .

Original: The next inequality holds since Var( Y j ) =
Revised: The next inequality holds since Var( Y j ) =

Original: Var( Y j ) =
Revised: Var( Y j ) =

Original: E [ Y 2 j ] − E [ Y j ] 2 ≤ E
Revised: E [ Y 2 j ] − E [ Y j ] 2 ≤ E

Original: [ Y 2 j ] where
Revised: [ Y 2 j ] where

Original: Y =
Revised: Y =

Original: Y − y i .
Revised: Y − y i .

Original: The ﬁnal inequality can be proved as follows.
Revised: The ﬁnal inequality can be proved as follows.

Original: We would like to bound
Revised: We would like to bound

Original: [Equation]
Revised: [Equation]

Original: where z j := (cid:107) y j − y i (cid:107) 2 (hence z i = 0 ).
Revised: where z j := (cid:107) y j − y i (cid:107) 2 (hence z i = 0 ).

Original: Deﬁne:
Revised: Deﬁne:

Original: [Equation]
Revised: [Equation]

Original: First note that as z j → ∞ , exp( − z 2 j ) → 0 exponentially fast, causing the product z 2 j exp( − z 2 j ) → 0 .
Revised: First note that as z j → ∞ , exp( − z 2 j ) → 0 exponentially fast, causing the product z 2 j exp( − z 2 j ) → 0 .

Original: Hence we expect the above quantity to be bounded and attain its maximum.
Revised: Hence we expect the above quantity to be bounded and attain its maximum.

Original: Let h ( z j ) :=
Revised: Let h ( z j ) :=

Original: exp( − z 2 j ) for notational conciseness, and note h ( z j ) > 0 .
Revised: exp( − z 2 j ) for notational conciseness, and note h ( z j ) > 0 .

Original: By taking partial derivatives with the chain rule, we have that for j (cid:54) =
Revised: By taking partial derivatives with the chain rule, we have that for j (cid:54) =

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: Hence the derivative is 0 if and only if z j = 0 or (1 − z 2 j ) (cid:80)
Revised: Hence the derivative is 0 if and only if z j = 0 or (1 − z 2 j ) (cid:80)

Original: k h
Revised: k h

Original: ( z k ) + (cid:80)
Revised: ( z k ) + (cid:80)

Original: k h ( z k ) z
Revised: k h ( z k ) z

Original: 2 k = 0 , the latter being equivalent to z 2 j = 1 + (cid:80)
Revised: 2 k = 0 , the latter being equivalent to z 2 j = 1 + (cid:80)

Original: k h ( z k ) z
Revised: k h ( z k ) z

Original: 2 k (cid:80) k h
Revised: 2 k (cid:80) k h

Original: ( z k )
Revised: ( z k )

Original: =
Revised: =

Original: 1 + g ( z ) .
Revised: 1 + g ( z ) .

Original: Hence at the maximum, the non-zero values among { z j } Nj =1 must be equal to one another.
Revised: Hence at the maximum, the non-zero values among { z j } Nj =1 must be equal to one another.

Original: It is clear now that the maximum value c is attained when z 2 j = 1 + c for j (cid:54)
Revised: It is clear now that the maximum value c is attained when z 2 j = 1 + c for j (cid:54)

Original: =
Revised: =

Original: i (and recall z i = 0 ).
Revised: i (and recall z i = 0 ).

Original: So h ( z j ) = exp( − 1 − c ) for j (cid:54)
Revised: So h ( z j ) = exp( − 1 − c ) for j (cid:54)

Original: =
Revised: =

Original: i .
Revised: i .

Original: Substituting this into g ( z ) , and rearranging, we obtain c exp( c + 1) =
Revised: Substituting this into g ( z ) , and rearranging, we obtain c exp( c + 1) =

Original: N − 1 .
Revised: N − 1 .

Original: Note φ ( x ) := x exp( x + 1) is increasing for x > 0 hence c
Revised: Note φ ( x ) := x exp( x + 1) is increasing for x > 0 hence c

Original: = φ − 1 ( N − 1) .
Revised: = φ − 1 ( N − 1) .

Original: Note φ (log N ) =
Revised: Note φ (log N ) =

Original: (log N ) exp(log N +1) ≥
Revised: (log N ) exp(log N +1) ≥

Original: N log N ≥ N − 1 for N ≥ 3 .
Revised: N log N ≥ N − 1 for N ≥ 3 .

Original: Since φ is increasing, we have φ − 1 ( N − 1) ≤ log( N ) for N ≥ 3 .
Revised: Since φ is increasing, we have φ − 1 ( N − 1) ≤ log( N ) for N ≥ 3 .

Original: In fact, it is known that φ − 1 ( N − 1) =
Revised: In fact, it is known that φ − 1 ( N − 1) =

Original: O (log N − log log N ) (Corless et al., 1996).
Revised: O (log N − log log N ) (Corless et al., 1996).

Original: Note the A term in f ( X ) =
Revised: Note the A term in f ( X ) =

Original: ˜ f
Revised: ˜ f

Original: ( X ) A allows us to use the above inequality, since Y (cid:62) P ( i ) Y
Revised: ( X ) A allows us to use the above inequality, since Y (cid:62) P ( i ) Y

Original: =
Revised: =

Original: Cov( Y ) now appears in the terms of J f :
Revised: Cov( Y ) now appears in the terms of J f :

Original: [Equation]
Revised: [Equation]

Original: Using the inequalities (cid:107) BC (cid:107) ≤ (cid:107) B (cid:107)(cid:107) C (cid:107) , (cid:107) B + C (cid:107) ≤ (cid:107) B (cid:107) + (cid:107) C (cid:107) and (cid:107)
Revised: Using the inequalities (cid:107) BC (cid:107) ≤ (cid:107) B (cid:107)(cid:107) C (cid:107) , (cid:107) B + C (cid:107) ≤ (cid:107) B (cid:107) + (cid:107) C (cid:107) and (cid:107)

Original: [ A 1 , . . .
Revised: [ A 1 , . . .

Original: , A N ] (cid:107) ≤ (cid:80) i (cid:107)
Revised: , A N ] (cid:107) ≤ (cid:80) i (cid:107)

Original: A i (cid:107) , we have:
Revised: A i (cid:107) , we have:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: , J iN ] (cid:107) ∞ ≤(cid:107)
Revised: , J iN ] (cid:107) ∞ ≤(cid:107)

Original: [Equation]
Revised: [Equation]

Original: ∞ +
Revised: ∞ +

Original: [Equation]
Revised: [Equation]

Original: j (cid:54) =
Revised: j (cid:54) =

Original: i (cid:107)
Revised: i (cid:107)

Original: [Equation]
Revised: [Equation]

Original: ∞ ≤
Revised: ∞ ≤

Original: [Equation]
Revised: [Equation]

Original: A
Revised: A

Original: [Equation]
Revised: [Equation]

Original: ∞ (cid:107)
Revised: ∞ (cid:107)

Original: [Equation]
Revised: [Equation]

Original: P ( i )
Revised: P ( i )

Original: [Equation]
Revised: [Equation]

Original: (cid:107) ∞ (cid:107)√ A
Revised: (cid:107) ∞ (cid:107)√ A

Original: [Equation]
Revised: [Equation]

Original: (cid:107) ∞ + P ii (cid:107)
Revised: (cid:107) ∞ + P ii (cid:107)

Original: A (cid:107) ∞ + 2 (cid:88)
Revised: A (cid:107) ∞ + 2 (cid:88)

Original: j (cid:54) =
Revised: j (cid:54) =

Original: i (cid:107)√
Revised: i (cid:107)√

Original: [Equation]
Revised: [Equation]

Original: (cid:107) ∞ (cid:107) P
Revised: (cid:107) ∞ (cid:107) P

Original: [Equation]
Revised: [Equation]

Original: ( y j −
Revised: ( y j −

Original: [Equation]
Revised: [Equation]

Original: For the ﬁrst equality, note that (cid:80) j P ij = 1 .
Revised: For the ﬁrst equality, note that (cid:80) j P ij = 1 .

Original: For the second equality, note that the summand for j = i is 0 because the term y i − y j = 0 .
Revised: For the second equality, note that the summand for j = i is 0 because the term y i − y j = 0 .

Original: Each of the terms in the brackets are bounded by the following lemmas:
Revised: Each of the terms in the brackets are bounded by the following lemmas:

Original: Lemma F.3.
Revised: Lemma F.3.

Original: (cid:107) Y
Revised: (cid:107) Y

Original: [Equation]
Revised: [Equation]

Original: P
Revised: P

Original: [Equation]
Revised: [Equation]

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: Y (cid:107) ∞ ≤ φ −
Revised: Y (cid:107) ∞ ≤ φ −

Original: [Equation]
Revised: [Equation]

Original: ( N − 1) (cid:112) D/H ( φ deﬁned as in Lemma F.2).
Revised: ( N − 1) (cid:112) D/H ( φ deﬁned as in Lemma F.2).

Original: Proof.
Revised: Proof.

Original: Recall that Y (cid:62) P ( i ) Y
Revised: Recall that Y (cid:62) P ( i ) Y

Original: =
Revised: =

Original: Cov( Y ) .
Revised: Cov( Y ) .

Original: Let σ ( Y m ) denote the standard deviation of Y m .
Revised: Let σ ( Y m ) denote the standard deviation of Y m .

Original: Then [Cov( Y )]
Revised: Then [Cov( Y )]

Original: [Equation]
Revised: [Equation]

Original: ≤ σ ( Y l ) σ ( Y m ) .
Revised: ≤ σ ( Y l ) σ ( Y m ) .

Original: Hence
Revised: Hence

Original: [Equation]
Revised: [Equation]

Original: since (cid:80) m σ ( Y m ) ≤ (cid:113) DH (cid:112)(cid:80) m σ 

Revised: since (cid:80) m σ ( Y m ) ≤ (cid:113) DH (cid:112)(cid:80) m σ

Original: ( Y m ) (by e.g. using the Cauchy–Schwartz inequality on [ σ ( Y 1 ) , . . .
Revised: ( Y m ) (by e.g. using the Cauchy–Schwartz inequality on [ σ ( Y 1 ) , . . .

Original: , σ ( Y D/H )] and 1 ) and max l σ ( Y l ) ≤ (cid:112)(cid:80) m σ 

Revised: , σ ( Y D/H )] and 1 ) and max l σ ( Y l ) ≤ (cid:112)(cid:80) m σ

Original: ( Y m ) , and the last inequality is from
Revised: ( Y m ) , and the last inequality is from

Original: Lemma F.2.
Revised: Lemma F.2.

Original: Lemma F.
Revised: Lemma F.

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Note (cid:107) ab (cid:62) (cid:107) ∞
Revised: Note (cid:107) ab (cid:62) (cid:107) ∞

Original: = (cid:107) a (cid:107) ∞ (cid:107) b (cid:107) 1 for real vectors a, b .
Revised: = (cid:107) a (cid:107) ∞ (cid:107) b (cid:107) 1 for real vectors a, b .

Original: Hence (cid:88)
Revised: Hence (cid:88)

Original: [Equation]
Revised: [Equation]

Original: where a j = (cid:112) P ij (cid:107) y j − (cid:80) k P ik y k (cid:107) ∞ , b j = (cid:112) P ij (cid:107)
Revised: where a j = (cid:112) P ij (cid:107) y j − (cid:80) k P ik y k (cid:107) ∞ , b j = (cid:112) P ij (cid:107)

Original: y i − y j (cid:107) 1 .
Revised: y i − y j (cid:107) 1 .

Original: Note a j ≤ c j := (cid:112) P ij (cid:107) y j − (cid:80)
Revised: Note a j ≤ c j := (cid:112) P ij (cid:107) y j − (cid:80)

Original: k P ik y k (cid:107) 2 since (cid:107) x (cid:107) ∞ ≤ (cid:107) x
Revised: k P ik y k (cid:107) 2 since (cid:107) x (cid:107) ∞ ≤ (cid:107) x

Original: (cid:107) 2 for vector x .
Revised: (cid:107) 2 for vector x .

Original: Hence (cid:107) a (cid:107) 2 ≤ (cid:107) c (cid:107) 2 .
Revised: Hence (cid:107) a (cid:107) 2 ≤ (cid:107) c (cid:107) 2 .

Original: Also b j ≤ (cid:113) DH d j := (cid:113) DH (cid:112) P ij (cid:107)
Revised: Also b j ≤ (cid:113) DH d j := (cid:113) DH (cid:112) P ij (cid:107)

Original: y i − y j (cid:107) 1 since (cid:107) x (cid:107) 1 ≤ (cid:113) DH (cid:107) x
Revised: y i − y j (cid:107) 1 since (cid:107) x (cid:107) 1 ≤ (cid:113) DH (cid:107) x

Original: (cid:107) 2 (e.g. by the Cauchy–Schwartz inequality on [ | x 1 | , . . .
Revised: (cid:107) 2 (e.g. by the Cauchy–Schwartz inequality on [ | x 1 | , . . .

Original: , | x D/H | ] and 1 ) for x ∈ R D/H .
Revised: , | x D/H | ] and 1 ) for x ∈ R D/H .

Original: Hence (cid:107) b (cid:107) 2 ≤ (cid:113) DH (cid:107) d (cid:107) 2 .
Revised: Hence (cid:107) b (cid:107) 2 ≤ (cid:113) DH (cid:107) d (cid:107) 2 .

Original: Note (cid:107) c
Revised: Note (cid:107) c

Original: (cid:107) 

Revised: (cid:107)

Original: = (cid:80) j P ij (cid:107) y j − (cid:80) k
Revised: = (cid:80) j P ij (cid:107) y j − (cid:80) k

Original: P
Revised: P

Original: ik y k (cid:107) 22 = Tr(Cov
Revised: ik y k (cid:107) 22 = Tr(Cov

Original: ( Y )) ≤ φ − 1 ( N − 1) from Lemma F.2, and (cid:107) d (cid:107) 

Revised: ( Y )) ≤ φ − 1 ( N − 1) from Lemma F.2, and (cid:107) d (cid:107)

Original: = (cid:80) j P ij (cid:107) y
Revised: = (cid:80) j P ij (cid:107) y

Original: i − y j (cid:107) 22 ≤ φ − 1 ( N − 1) also from Lemma F.2.
Revised: i − y j (cid:107) 22 ≤ φ − 1 ( N − 1) also from Lemma F.2.

Original: Hence (cid:107) a (cid:107) 2 (cid:107) b (cid:107) 2 ≤ (cid:113) DH (cid:107) c (cid:107) 2 (cid:107) d (cid:107) 2 ≤ (cid:113) DH
Revised: Hence (cid:107) a (cid:107) 2 (cid:107) b (cid:107) 2 ≤ (cid:113) DH (cid:107) c (cid:107) 2 (cid:107) d (cid:107) 2 ≤ (cid:113) DH

Original: φ − 1 ( N − 1) .
Revised: φ − 1 ( N − 1) .

Original: Putting the above lemmas altogether, with the observation sup X (cid:107) J f ( X ) (cid:107) ∞
Revised: Putting the above lemmas altogether, with the observation sup X (cid:107) J f ( X ) (cid:107) ∞

Original: =
Revised: =

Original: sup X (cid:107) [ J i 1 ( X ) , . .
Revised: sup X (cid:107) [ J i 1 ( X ) , . .

Original: , J iN ( X )] (cid:107) ∞ by permutation invariance of (cid:107) J f (cid:107) ∞ (since f is permutation equivariant and (cid:107) · (cid:107) ∞ is the maximum absolute row sum), we have
Revised: , J iN ( X )] (cid:107) ∞ by permutation invariance of (cid:107) J f (cid:107) ∞ (since f is permutation equivariant and (cid:107) · (cid:107) ∞ is the maximum absolute row sum), we have

Original: [Equation]
Revised: [Equation]

Original: (cid:62)
Revised: (cid:62)

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where the last inequality holds for N ≥ 3 .
Revised: where the last inequality holds for N ≥ 3 .

Original: The full multihead attention map that combines the heads f h ( X ) is:
Revised: The full multihead attention map that combines the heads f h ( X ) is:

Original: [Equation]
Revised: [Equation]

Original: where g : X (cid:55)→ [ f 1 ( X ) , . . .
Revised: where g : X (cid:55)→ [ f 1 ( X ) , . . .

Original: , f H ( X )] , W O ∈ R D × D
Revised: , f H ( X )] , W O ∈ R D × D

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: Note the Jacobian J g is a block matrix whose rows are J f h , hence (cid:107) J g (cid:107) ∞
Revised: Note the Jacobian J g is a block matrix whose rows are J f h , hence (cid:107) J g (cid:107) ∞

Original: =
Revised: =

Original: max h (cid:107) J f h (cid:107) ∞ , and similarly (cid:107) W V (cid:62) (cid:107) ∞ =
Revised: max h (cid:107) J f h (cid:107) ∞ , and similarly (cid:107) W V (cid:62) (cid:107) ∞ =

Original: max h (cid:107) W V,h (cid:62) (cid:107) ∞ .
Revised: max h (cid:107) W V,h (cid:62) (cid:107) ∞ .

Original: Hence we have
Revised: Hence we have

Original: [Equation]
Revised: [Equation]

Original: Combining this with Inequality (35), we have:
Revised: Combining this with Inequality (35), we have:

Original: Lip ∞
Revised: Lip ∞

Original: [Equation]
Revised: [Equation]

Original: F
Revised: F

Original: [Equation]
Revised: [Equation]

Original: ≤
Revised: ≤

Original: [Equation]
Revised: [Equation]

Original: F.2. U PPER BOUND ON Lip 2 ( F ) FOR L2-MHA
Revised: F.2. U PPER BOUND ON Lip 2 ( F ) FOR L2-MHA

Original: For p = 2 , we use the following lemma:
Revised: For p = 2 , we use the following lemma:

Original: Lemma F.5.
Revised: Lemma F.5.

Original: Let A be a block matrix with block rows A 1 , . . .
Revised: Let A be a block matrix with block rows A 1 , . . .

Original: , A N .
Revised: , A N .

Original: Then (cid:107) A (cid:107) 2 ≤
Revised: Then (cid:107) A (cid:107) 2 ≤

Original: (cid:112)(cid:80) i (cid:107)
Revised: (cid:112)(cid:80) i (cid:107)

Original: A i (cid:107) 22 , and equality holds if and only if the ﬁrst right singular vectors of the A i align.
Revised: A i (cid:107) 22 , and equality holds if and only if the ﬁrst right singular vectors of the A i align.

Original: Proof.
Revised: Proof.

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: Note that equality holds if and only if the ﬁrst right singular vectors of the A i align.
Revised: Note that equality holds if and only if the ﬁrst right singular vectors of the A i align.

Original: Hence a bound on the spectral norm of each block row of J f can give us an O ( √ N ) bound on (cid:107) J f (cid:107) 2 , which may be loose, and it remains an open question as to whether this bound can be tightened.
Revised: Hence a bound on the spectral norm of each block row of J f can give us an O ( √ N ) bound on (cid:107) J f (cid:107) 2 , which may be loose, and it remains an open question as to whether this bound can be tightened.

Original: To bound the (cid:107) · (cid:107) 2 norm of each row of J f , we use the following lemmas:
Revised: To bound the (cid:107) · (cid:107) 2 norm of each row of J f , we use the following lemmas:

Original: Lemma F.6.
Revised: Lemma F.6.

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: (cid:107)
Revised: (cid:107)

Original: [Equation]
Revised: [Equation]

Original: 2 = λ max (Cov( Y )) ≤ Tr(Cov( Y )) ≤ φ − 1 ( N − 1) , where the ﬁrst equality holds by symmetry of Cov( Y ) and the next holds by Cov( Y ) being positive semideﬁnite, so all its eigenvalues are non-negative, and hence the maximal eigenvalue is bounded by the sum of the eigenvalues, equal to its trace.
Revised: 2 = λ max (Cov( Y )) ≤ Tr(Cov( Y )) ≤ φ − 1 ( N − 1) , where the ﬁrst equality holds by symmetry of Cov( Y ) and the next holds by Cov( Y ) being positive semideﬁnite, so all its eigenvalues are non-negative, and hence the maximal eigenvalue is bounded by the sum of the eigenvalues, equal to its trace.

Original: The ﬁnal inequality is from Lemma F.2.
Revised: The ﬁnal inequality is from Lemma F.2.

Original: Lemma
Revised: Lemma

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Proof.
Revised: Proof.

Original: Directly use Cauchy–Schwartz on c and d in the proof of Lemma F.4.
Revised: Directly use Cauchy–Schwartz on c and d in the proof of Lemma F.4.

Original: Again using the inequalities (cid:107) BC (cid:107) ≤ (cid:107) B (cid:107)(cid:107) C (cid:107) , (cid:107) B + C (cid:107) ≤ (cid:107) B (cid:107) + (cid:107) C (cid:107) and (cid:107) [ A 1 , . .
Revised: Again using the inequalities (cid:107) BC (cid:107) ≤ (cid:107) B (cid:107)(cid:107) C (cid:107) , (cid:107) B + C (cid:107) ≤ (cid:107) B (cid:107) + (cid:107) C (cid:107) and (cid:107) [ A 1 , . .

Original: A N ] (cid:107) ≤ (cid:80) i (cid:107)
Revised: A N ] (cid:107) ≤ (cid:80) i (cid:107)

Original: A i (cid:107) , with the additional equality (cid:107) B (cid:62) (cid:107) 

Revised: A i (cid:107) , with the additional equality (cid:107) B (cid:62) (cid:107)

Original: = (cid:107) B (cid:107) 

Revised: = (cid:107) B (cid:107)

Original: , we have the bound:
Revised: , we have the bound:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: iN
Revised: iN

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Using Lemma F.5, we have that
Revised: Using Lemma F.5, we have that

Original: [Equation]
Revised: [Equation]

Original: To obtain the ﬁnal result for the full multihead self-attention F
Revised: To obtain the ﬁnal result for the full multihead self-attention F

Original: [Equation]
Revised: [Equation]

Original: we need a ﬁnal lemma:
Revised: we need a ﬁnal lemma:

Original: Lemma F.8.
Revised: Lemma F.8.

Original: Let A be a block matrix with block columns
Revised: Let A be a block matrix with block columns

Original: A 1 , . . . , A N .
Revised: A 1 , . . . , A N .

Original: Then (cid:107) A (cid:107) 2 ≤
Revised: Then (cid:107) A (cid:107) 2 ≤

Original: (cid:112)(cid:80) i (cid:107)
Revised: (cid:112)(cid:80) i (cid:107)

Original: A i (cid:107) 22 .
Revised: A i (cid:107) 22 .

Original: Proof.
Revised: Proof.

Original: [Equation]
Revised: [Equation]

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: where we are using the substitution x i
Revised: where we are using the substitution x i

Original: = λ i e i , and the last inequality holds by e.g. Cauchy–Schwartz inequality on [ λ 1 , . . .
Revised: = λ i e i , and the last inequality holds by e.g. Cauchy–Schwartz inequality on [ λ 1 , . . .

Original: , λ N ] and [ (cid:107) A 1 (cid:107) 2 , . . .
Revised: , λ N ] and [ (cid:107) A 1 (cid:107) 2 , . . .

Original: , (cid:107) A N (cid:107) 2 ] .
Revised: , (cid:107) A N (cid:107) 2 ] .

Original: Recall that
Revised: Recall that

Original: [Equation]
Revised: [Equation]

Original: Since (cid:107) f h ( X ) W V,h (cid:107) 2 ≤ (cid:107) J f h (cid:107) 2 (cid:107) W V,h (cid:107) 2 , by Lemma F.8 we have that
Revised: Since (cid:107) f h ( X ) W V,h (cid:107) 2 ≤ (cid:107) J f h (cid:107) 2 (cid:107) W V,h (cid:107) 2 , by Lemma F.8 we have that

Original: [Equation]
Revised: [Equation]

Original: and hence
Revised: and hence

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Combining this with Inequality (36), we have:
Revised: Combining this with Inequality (36), we have:

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: h
Revised: h

Original: [Table]
Revised: [Table]

Original: W Q,
Revised: W Q,

Original: [Table] [Table]
Revised: [Table] [Table]

Original: W
Revised: W

Original: [Table]
Revised: [Table]

Original: ,h
Revised: ,h

Original: [Table] [Table]
Revised: [Table] [Table]

Original: W O
Revised: W O

Original: [Table]
Revised: [Table]

Original: G T HE C ASE WITH M ASKING
Revised: G T HE C ASE WITH M ASKING

Original: Since self-attention is often used with masking , a natural question is how masking affects the derived bounds.
Revised: Since self-attention is often used with masking , a natural question is how masking affects the derived bounds.

Original: In self-attention (for any choice of attention function), masking is implemented as follows: given a set of mask indices M ⊂ { 1 , . . .
Revised: In self-attention (for any choice of attention function), masking is implemented as follows: given a set of mask indices M ⊂ { 1 , . . .

Original: , N } × { 1 , . . .
Revised: , N } × { 1 , . . .

Original: , N } , the logits (i.e. the inputs to the softmax) are set to −∞ at the mask indices.
Revised: , N } , the logits (i.e. the inputs to the softmax) are set to −∞ at the mask indices.

Original: That is,
Revised: That is,

Original: [Equation]
Revised: [Equation]

Original: where ˜ L ij is the original logit (e.g. for L2 self-attention, ˜ L ij = − ( x i − x j )
Revised: where ˜ L ij is the original logit (e.g. for L2 self-attention, ˜ L ij = − ( x i − x j )

Original: (cid:62) A ( x i − x j ) ).
Revised: (cid:62) A ( x i − x j ) ).

Original: Masking implies f i ( X ) is not a function of x j for ( i, j ) ∈ M ,
Revised: Masking implies f i ( X ) is not a function of x j for ( i, j ) ∈ M ,

Original: hence J ij = 0 for ( i, j ) ∈ M .
Revised: hence J ij = 0 for ( i, j ) ∈ M .

Original: Thus f i ( X ) is equal to the i th output for self-attention with inputs restricted to { x j : ( i, j ) / ∈ M} , the unmasked inputs with respect to the i th output.
Revised: Thus f i ( X ) is equal to the i th output for self-attention with inputs restricted to { x j : ( i, j ) / ∈ M} , the unmasked inputs with respect to the i th output.

Original: Hence J ij will no longer contribute to the bound on (cid:107)
Revised: Hence J ij will no longer contribute to the bound on (cid:107)

Original: [ J i 1 , . . .
Revised: [ J i 1 , . . .

Original: , J iN ] (cid:107) , and hence the bound for the unmasked case will continue to hold as long as ( i, i ) ∈
Revised: , J iN ] (cid:107) , and hence the bound for the unmasked case will continue to hold as long as ( i, i ) ∈

Original: M i.e. x i attends to itself (this is necessary for the proof of Lemma F.2 to hold).
Revised: M i.e. x i attends to itself (this is necessary for the proof of Lemma F.2 to hold).

Original: The bound can in fact be tightened by replacing N with |{ x j : ( i, j ) / ∈ M}| , the number of unmasked inputs with respect to the i th output.
Revised: The bound can in fact be tightened by replacing N with |{ x j : ( i, j ) / ∈ M}| , the number of unmasked inputs with respect to the i th output.

Original: H D ROPOUT IS C ONTRACTIVE
Revised: H D ROPOUT IS C ONTRACTIVE

Original: At test time, Dropout multiplies inputs by the dropout keep probability p < 1 , so it is a contraction with Lipschitz constant p at evaluation time.
Revised: At test time, Dropout multiplies inputs by the dropout keep probability p < 1 , so it is a contraction with Lipschitz constant p at evaluation time.

Original: At training time, Dropout amounts to setting some inputs to zero, while keeping other inputs constant.
Revised: At training time, Dropout amounts to setting some inputs to zero, while keeping other inputs constant.

Original: This can be expressed as right multiplication by a diagonal binary matrix M , and for such matrices we can verify (cid:107) M (cid:107) p :=
Revised: This can be expressed as right multiplication by a diagonal binary matrix M , and for such matrices we can verify (cid:107) M (cid:107) p :=

Original: sup (cid:107) x
Revised: sup (cid:107) x

Original: (cid:107) p =1 (cid:107)
Revised: (cid:107) p =1 (cid:107)

Original: Mx (cid:107) p ≤ 1 .
Revised: Mx (cid:107) p ≤ 1 .

Original: I E XPERIMENTAL D ETAILS
Revised: I E XPERIMENTAL D ETAILS

Original: For the experiment in Section 5.1, showing the asymptotic tightness of the upper bound on Lip ∞ ( F ) where F is L2-MHA , we ﬁx all free parameters of F (namely W Q , W V ) to be the identity, and only optimise the input X .
Revised: For the experiment in Section 5.1, showing the asymptotic tightness of the upper bound on Lip ∞ ( F ) where F is L2-MHA , we ﬁx all free parameters of F (namely W Q , W V ) to be the identity, and only optimise the input X .

Original: We use 50 random initialisations of X for each N , where X ij ∼ U [ − c, c ] for c ∼ U [0 , 10] (we observed that having c itself be random improves optimisation).
Revised: We use 50 random initialisations of X for each N , where X ij ∼ U [ − c, c ] for c ∼ U [0 , 10] (we observed that having c itself be random improves optimisation).

Original: We display the top 5 results for each value of N after optimising each random initialisation till convergence using Adam (Kingma & Ba, 2015) with a learning rate of 0 . 1 .
Revised: We display the top 5 results for each value of N after optimising each random initialisation till convergence using Adam (Kingma & Ba, 2015) with a learning rate of 0 . 1 .

Original: For the experiments in Section 5.3, we comparing the performance of the original Transformer and the Transformer with Lipschitz/invertible self-attention at character-level language modelling on the
Revised: For the experiments in Section 5.3, we comparing the performance of the original Transformer and the Transformer with Lipschitz/invertible self-attention at character-level language modelling on the

Original: Penn Treebank dataset (Marcus et al., 1993).
Revised: Penn Treebank dataset (Marcus et al., 1993).

Original: Each training example is a sentence represented as a variable-length sequence of characters, and examples are batched according to length such that padding is minimised, with the maximum sequence length set to 288 .
Revised: Each training example is a sentence represented as a variable-length sequence of characters, and examples are batched according to length such that padding is minimised, with the maximum sequence length set to 288 .

Original: All models are autoregressive, outputting the logits for the categorical likelihood predicting the next character, and are trained using maximum likelihood (cross-entropy loss) with a batch size of 64 .
Revised: All models are autoregressive, outputting the logits for the categorical likelihood predicting the next character, and are trained using maximum likelihood (cross-entropy loss) with a batch size of 64 .

Original: The LSTM models have the dimensionality of the hidden state equal to the dimensionality D of the cell state (the usual default implementation).
Revised: The LSTM models have the dimensionality of the hidden state equal to the dimensionality D of the cell state (the usual default implementation).

Original: The Transformer models are trained with a varying number of blocks (number of layers) with H = 8 heads and D = 512 , tuning hyperparameters for dropout rate in { 0 , 0 . 1 , 0 . 2 } and base learning rate γ ∈ { 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 , 1 . 0 , 1 . 5 , 2 . 0 } with number of warmup iterations w ∈ { 1000 , 2000 , 4000 , 8000 } for the standard custom learning rate schedule in Vaswani et al.
Revised: The Transformer models are trained with a varying number of blocks (number of layers) with H = 8 heads and D = 512 , tuning hyperparameters for dropout rate in { 0 , 0 . 1 , 0 . 2 } and base learning rate γ ∈ { 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 , 1 . 0 , 1 . 5 , 2 . 0 } with number of warmup iterations w ∈ { 1000 , 2000 , 4000 , 8000 } for the standard custom learning rate schedule in Vaswani et al.

Original: [Equation]
Revised: [Equation]

Original: where (cid:15) t is the learning rate at training iteration t .
Revised: where (cid:15) t is the learning rate at training iteration t .

Original: Hence the learning rate linearly increases from 0 to ( Dw ) − 1 / 2 over w iterations, then decays proportionally to t − 1 / 2 .
Revised: Hence the learning rate linearly increases from 0 to ( Dw ) − 1 / 2 over w iterations, then decays proportionally to t − 1 / 2 .

Original: We use Glorot Uniform initialisation (Glorot & Bengio, 2010) for all weights ( U (cid:104) − (cid:113) 1 d in + d out , (cid:113) 1 d in + d out (cid:105) ), except for weights in L2-MHA that are initialised from U (cid:104) − s √ D , s √ D (cid:105) , and s is a hyperparameter.
Revised: We use Glorot Uniform initialisation (Glorot & Bengio, 2010) for all weights ( U (cid:104) − (cid:113) 1 d in + d out , (cid:113) 1 d in + d out (cid:105) ), except for weights in L2-MHA that are initialised from U (cid:104) − s √ D , s √ D (cid:105) , and s is a hyperparameter.

Original: For D = 512 , we used s = 12 4 .
Revised: For D = 512 , we used s = 12 4 .

Original: All experiments were done in Tensorﬂow 1.14 (Abadi et al., 2016).
Revised: All experiments were done in Tensorﬂow 1.14 (Abadi et al., 2016).

Original: The code will be released upon de-anonymisation.
Revised: The code will be released upon de-anonymisation.

Original: In Table 2 we show the best Test NLL across training of Transformer models in Figure 4.
Revised: In Table 2 we show the best Test NLL across training of Transformer models in Figure 4.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: = W
Revised: = W

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: J N UMERICAL I NVERTIBILITY OF MHA R ESIDUAL M AP
Revised: J N UMERICAL I NVERTIBILITY OF MHA R ESIDUAL M AP

Original: Following Section 5.2, Figure 5 conﬁrms that numerical invertibility does not hold for trained weights for dot-product multihead self-attention (DP-MHA) (obtained from one-layer Transformer (DP) model used for Figure 4), similar to the randomly initialised weight case.
Revised: Following Section 5.2, Figure 5 conﬁrms that numerical invertibility does not hold for trained weights for dot-product multihead self-attention (DP-MHA) (obtained from one-layer Transformer (DP) model used for Figure 4), similar to the randomly initialised weight case.

Original: Figure 6 shows additional results for different values of N and D .
Revised: Figure 6 shows additional results for different values of N and D .

Original: [Figure]
Revised: [Figure]

Original: K B EHAVIOUR OF L OWER B OUND ON Lip 2 ( F )
Revised: K B EHAVIOUR OF L OWER B OUND ON Lip 2 ( F )

Original: N
Revised: N

Original: [Figure]
Revised: [Figure]

Original: In Figure 7, we show the lower bound on Lip 2 ( F ) obtained by optimising (cid:107) J F ( X ) (cid:107) 2 using the same optimisation procedure as for Figure 2 of Section 5.1.
Revised: In Figure 7, we show the lower bound on Lip 2 ( F ) obtained by optimising (cid:107) J F ( X ) (cid:107) 2 using the same optimisation procedure as for Figure 2 of Section 5.1.

Original: Here the optimisation is more difﬁcult, evident in the variance of the top 5 values, and the trend is less clear, but it appears that Lip 2 ( f ) grows at a rate of O (log N ) .
Revised: Here the optimisation is more difﬁcult, evident in the variance of the top 5 values, and the trend is less clear, but it appears that Lip 2 ( f ) grows at a rate of O (log N ) .

Original: The message is less clear here, and there are at least two possibilities:
Revised: The message is less clear here, and there are at least two possibilities:

Original: (1) The optimisation is difﬁcult even for small values of N , hence Figure
Revised: (1) The optimisation is difﬁcult even for small values of N , hence Figure

Original: shows a loose lower bound.
Revised: shows a loose lower bound.

Original: (2) If the lower bound is tight, this suggests that the O ( √ N log N ) bound in Theorem 3.2 is not asymptotically tight, and could be improved to O (log N ) (or O (log N − log log N ) as for p = ∞ ).
Revised: (2) If the lower bound is tight, this suggests that the O ( √ N log N ) bound in Theorem 3.2 is not asymptotically tight, and could be improved to O (log N ) (or O (log N − log log N ) as for p = ∞ ).

Original: L O PTIMISING THE NORM OF THE J ACOBIAN OF DP-MHA
Revised: L O PTIMISING THE NORM OF THE J ACOBIAN OF DP-MHA

Original: In Figure 8, we show how the norm of the Jacobian (cid:107) J f ( X ) (cid:107) ∞ for DP-MHA f keeps increasing when being optimised with respect to X .
Revised: In Figure 8, we show how the norm of the Jacobian (cid:107) J f ( X ) (cid:107) ∞ for DP-MHA f keeps increasing when being optimised with respect to X .

Original: This is a useful sanity check validating our theoretical result of Theorem 3.1, that DP-MHA is not Lipshchitz.
Revised: This is a useful sanity check validating our theoretical result of Theorem 3.1, that DP-MHA is not Lipshchitz.

Original: The oscillations are likely due to momentum term of Adam optimizer that was used to optimise the norm.
Revised: The oscillations are likely due to momentum term of Adam optimizer that was used to optimise the norm.

Original: [Figure]
Revised: [Figure]

Original: M E XPERIMENT TYING KEYS AND QUERIES OF L2-MHA BUT PRESERVING PARAMETER COUNT
Revised: M E XPERIMENT TYING KEYS AND QUERIES OF L2-MHA BUT PRESERVING PARAMETER COUNT

Original: In Figure 4 of Section 5.3, we have shown that there is a clear reduction in performance when tying the keys and queries.
Revised: In Figure 4 of Section 5.3, we have shown that there is a clear reduction in performance when tying the keys and queries.

Original: To test whether this can be attributed to the reduction in parameter count, we tried doubling the number of columns of W Q when the keys and queries are shared (i.e. from D/H to 2 D/H ) so that the shared model has the same number of parameters as the unshared model.
Revised: To test whether this can be attributed to the reduction in parameter count, we tried doubling the number of columns of W Q when the keys and queries are shared (i.e. from D/H to 2 D/H ) so that the shared model has the same number of parameters as the unshared model.

Original: In Figure 9, the third column shows results for shared L2-MHA , but with the same number of parameters as the unshared L2-MHA i.e. without tying the keys and queries.
Revised: In Figure 9, the third column shows results for shared L2-MHA , but with the same number of parameters as the unshared L2-MHA i.e. without tying the keys and queries.

Original: The performance is similar to the second column (tying with a reduced number of parameters), suggesting that there is an inherent limitation in expressiveness to tying the keys and queries, and that the reduction in number of parameters is an insufﬁcient explanation this phenomenon.
Revised: The performance is similar to the second column (tying with a reduced number of parameters), suggesting that there is an inherent limitation in expressiveness to tying the keys and queries, and that the reduction in number of parameters is an insufﬁcient explanation this phenomenon.

Original: [Figure]
Revised: [Figure]

Original: N S TABILITY E XPERIMENTS
Revised: N S TABILITY E XPERIMENTS

Original: In Figure 10 below, we compare the output variance of trained L2-MHA against trained DP-MHA , with weights from the one-layer Transformer (L2), W Q = W K model and (DP) model used for Figurerespectively.
Revised: In Figure 10 below, we compare the output variance of trained L2-MHA against trained DP-MHA , with weights from the one-layer Transformer (L2), W Q = W K model and (DP) model used for Figurerespectively.

Original: We take the same distribution of inputs as used for the numerical invertibility experiment in Section 5.2, and show the histogram of inputs and outputs after ﬂattening the input/output tensors.
Revised: We take the same distribution of inputs as used for the numerical invertibility experiment in Section 5.2, and show the histogram of inputs and outputs after ﬂattening the input/output tensors.

Original: We see that the range of outputs remains similar to the range of inputs for Lipschitz L2-MHA , whereas for DP-MHA the outputs have a much wider range, because the Jacobian norm is large for DP-MHA at these inputs.
Revised: We see that the range of outputs remains similar to the range of inputs for Lipschitz L2-MHA , whereas for DP-MHA the outputs have a much wider range, because the Jacobian norm is large for DP-MHA at these inputs.

Original: In practice, this leads to instabilities while training for DP-MHA , hence requiring careful tuning of the learning rate schedule for training deeper Transformer models: linear warmup and square root decay, as detailed in Appendix I.
Revised: In practice, this leads to instabilities while training for DP-MHA , hence requiring careful tuning of the learning rate schedule for training deeper Transformer models: linear warmup and square root decay, as detailed in Appendix I.

Original: In Figure 11, we show how training instabilities arise for
Revised: In Figure 11, we show how training instabilities arise for

Original: DP-MHA with deeper Transformer models if we use a ﬁxed learning rate .
Revised: DP-MHA with deeper Transformer models if we use a ﬁxed learning rate .

Original: We compare the training curves of DP-MHA , L2-MHA ( W Q = W K ) and Contractive-L2-MHA for a varying number of layers for the ﬁrst 20 epochs of training with a ﬁxed learning rate.
Revised: We compare the training curves of DP-MHA , L2-MHA ( W Q = W K ) and Contractive-L2-MHA for a varying number of layers for the ﬁrst 20 epochs of training with a ﬁxed learning rate.

Original: We see that DP-MHA fails to train beyond 10 layers, whereas both L2-MHA ( W Q = W K ) (i.e. Lipschitz L2-MHA but not contractive) and Contractive-L2-MHA shows stable training for up to 18 layers.
Revised: We see that DP-MHA fails to train beyond 10 layers, whereas both L2-MHA ( W Q = W K ) (i.e. Lipschitz L2-MHA but not contractive) and Contractive-L2-MHA shows stable training for up to 18 layers.

Original: This was the deepest model we could ﬁt on a single GPU, and we expect to be able to train even deeper models with L2-MHA ( W Q = W K ) and Contractive-L2-MHA .
Revised: This was the deepest model we could ﬁt on a single GPU, and we expect to be able to train even deeper models with L2-MHA ( W Q = W K ) and Contractive-L2-MHA .

Original: In Table 2 we show the best Test NLL across training of Transformer models with ﬁxed learning rate.
Revised: In Table 2 we show the best Test NLL across training of Transformer models with ﬁxed learning rate.

Original: Note that for DP-MHA training becomes unstable beyond 10 layers, so we are only able to provide results up to 10 layers.
Revised: Note that for DP-MHA training becomes unstable beyond 10 layers, so we are only able to provide results up to 10 layers.

Original: The generalisation performance of the best model for each setting is similar.
Revised: The generalisation performance of the best model for each setting is similar.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: Inputs
Revised: Inputs

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: (DP) 1.
Revised: (DP) 1.

Original: W Q = W K (L2) 1.
Revised: W Q = W K (L2) 1.

Original: (Contractive-L2) 1.
Revised: (Contractive-L2) 1.

Original: [Table]
Revised: [Table]

Revised: language modelling