The Lipschitz Constant of Self-Attention
Authors: Hyunjik Kim
George Papamakarios
Andriy Mnih
Venue:  (2021)
Keywords: Lipschitz constant
self-attention
theory

Abstract: 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. Such works have focused on bounding the Lipschitz constant of fully connected or convolutional networks, composed of linear maps and pointwise non-linearities. In this paper, we investigate the Lipschitz constant of self-attention, a non-linear neural network module widely used in sequence modelling. We prove that the standard dot-product self-attention is *not* Lipschitz, and propose an alternative L2 self-attention that *is* Lipschitz. We derive an upper bound on the Lipschitz constant of L2 self-attention and provide empirical evidence for its asymptotic tightness. 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.

Under review as a conference paper at ICLR 2021
THE LIPSCHITZ CONSTANT OF SELF-ATTENTION
Anonymous authors
Paper under double-blind review
ABSTRACT
Lipschitz constants of neural networks have been explored in various contexts in deep learning, such provable adversarial robustness, estimating Wasserstein distance, stabilising training of GANs, and formulating invertible neural networks. Such works have focused on bounding the Lipschitz constant of fully connected convolutional networks, composed of linear maps and pointwise non-linearities. In this paper, we investigate the Lipschitz constant of self-attention, a non-linear neural network module widely used in sequence modelling. We prove that the standard dot- dot-product self-attention is not Lipschitz, and propose an alternative L2 self-attention that is Lipschitz. We derive an upper bound on the Lipschitz constant of 1.2 self-attention and provide empirical evidence for its asymptotic tightness. To demonstrate the practical relevance of our theoretical work, we formulate invertible self-attention and use it in Transformer-bass architecture for a character-level language modelling task.
INTRODUCTION
L Lipschitz continuity is a strong form of continuity for functions. Loosely speaking, a function is Lipschitz continuous if changing its input by certain amount cannot change its output by more than K times that amount. The constant K is a hard constraint on how rapidly the function's output can vary, and the smallest such K is known as the function's Lipschitz constant. For example, (x) = / [a| and fz(z) = exp(x) for x E R are not Lipschitz continuous, because their output can change arbitrarily fast as approaches and 4:0 respectively. On the other hand, 91 (æ) = tanh(z) and g2 (2)
are Lipschitz continuous, because their rate of change (derivative) is bounded.
In deep learning, we often use Lipschitz continuity as a constraint for neural networks, to control how much network's 's output can change relative to its input. Such Lipschitz constraints are useful in several contexts. For example, Lipschitz constraints can endow models with provable robustness against adversarial pertubations (Cisse et al., 2017 Tsuzuku et al. 2018 Anil et al. 2019), and guaranteed generalisation bounds (Sokolić et al. (017). Moreover, the dual form of the Wasserstein distance is defined as supremum over Lipschitz functions with given Lipschitz constant, hence Lipschitz-constraineed networks are used for estimating Wasserstein distances (Peyré & Cuturi] 2019). Further, Lipschitz-conssrriined networks can stabilise training for GANs, an example being spectral normalisation [Miyato et al, (018). Finally, invertible models and normalising flows. For example, Lipschitz-constrrined networks can be used as a building block for invertible residual networks and hence flow-based generative models (Behrmann al. 2019: Chen et al. 2019). Additionally, Neural ODEs (Chen et al., 2018 Grathwohl al. 2019) are typically defined using vector fields parameterized via Lipschitz networks, so that the flow generated by the vector field is guaranteed to exist for all times.
networks are also used to construct
Nonetheless, designing Lipschitz-contimmuuu neural networks and computing (or even upper- bounding) their Lipschitz constant is a hard problem. Previous work mostly focused on fully- connected and convolutional networks, not only because they common in deep learning, but also because they are relatively simple to analyze, as compositions of linear maps and pointwise non- linearities. Even in this case however, exact evaluation of the Lipschitz constant of fully-connected and convolutional networks is NP-hard (Virmaux & Scamann 2018) and obtaining a tight upper bound remains challenging task (Virmaux & Scaman 2018 Fazlyab et al. 2019 Latorre al. 2020).
Fully colle-connected and convolutional networks are not the only neural networks worthy of interest. Recently, self-attention /Vaswani al. 2017) has become popular alternative recurrent neural

---

Under review as a conference paper at ICLR 2021
networks. Self-attention is key component of the Transformer (Vaswani al. 2017). that has found success as a building block in models of various data modalities, starting with natural -laguugg processing (Vaswani et al. 2017 Devlin et al. 2019; Brown et al. 2020) and extending computer vision (Zhang et al. 2019; Parmar et al. 2019), audio generation (Huang et al., 2019), and reinforce- ment learning (Parisotto al. 2020p. However, far previous work has analyzed the Lipschitz properties of self-attention, and thus it has been unclear whether self-attention is viable option in applications that require Lipschitz constraints. In this work, we address this gap in the theory of self-attention by providing a thorough analysis of its Lipschitz properties. In particular, we make the following contributions:
We prove that the widely used dot dot-product self-attention is not Lipschitz, and therefore not suitable to use in applications requiring Lipschitz constraints.
We formulate L2 self- self-attention as an alternative, and show that it Lipschitz.
We derive theoretical upper bound on the Lipschitz constant of L2 self-attention, and provide empirical evidence of the asymptotic tightness of the bound.
As a practical demonstration of the theory, we use this bound to formulate invertible self- attentonnn and explore its use in Transformer architecture for character-level language modelling.
2 LIPSCHITZ CONSTANT OF FULLY- FULLI-CONNECTTD//COVVOUUTOIOO
LAYERS
We first define the notion of Lipschitz continuity, and proceed to define the Lipschitz constant.
Definition 2.1. Given two metric spaces (X, dx) and ()) dy). a function f X + y is called Lipschitz continuous (or K K-Lipschirz) if there exists constant K 2 such that
(x(f f(x)) S (2) for all x x' E X.
(1)
The smallest such K is the Lipschitz constant of denoted Lip(S).
In this paper. we focus on the common case where X' = R". y' R", and dx, dy are induced by a p-norm ||>||, (Σ: //pp) We will primarily consider the cases p and p = oc, where || = max, | To emphasise the dependence of the Lipschitz constant on the choice of p-norm, we will often denote it by (f). In this case, follows directly from Definition 2.1 that the Lipschitz constant is given by
lf(x)) - f(x')ll
(f) = sup
(2)
xy || x'
Next, we outline some basic results that are useful for estimating Lipschitz constants, also covered in related works (Virmaux & Scaman 2018: [Bhrrman al. 2019). We describe how these results are used to provide bounds on the Lipschitz constant of fully cully-connected networks (FCN) and convolutional neural networks (CNN), using the fact that both are compositions of linear maps and pointwise non-linearee non-lincarities. To begin with, the following theorem suggests a way to bound Lipp(f) for a differentiable Lipschitz function f:
Theorem 2.1 (Federer 1969). Let 1. + =mm be differentiable and Lipschitz continuous under a choice of p-norm Ipp- Let Jy(2) denote its total derivative (Jacobian) at x. Then (fp,,()) supReR= (z)]pp where [|J the induced operator norm Jy(x).
Hence if f is linear map represented by a matrix W then
mmax (W), max; E, [asl if 00
if
(f) = ||W||p = sup
|
(3)
where ||W||p is the operator norm on matrices induced by the vector p-norm, and max (m) is the largest singular value of W. Under this choice of norm, many common non-linearities (including relu, si simmiid, tanh, elu) are 1-Lipschitz. || Fmax (W) is usually estimated via power
iteration; we provide details on how this is done in Appendix B
Since we now know the Lipsehitz constants of the components of both FCN and CNN, we can bound their Lipschitz constants by applying the following lemma:
Lemma 2.1 (Federer 1969). Let g. h be two composable Lipschitz with Lip(g h) Lip(g) Lip(h).
Then g o h is also
2

---

Under review as a conference paper at ICLR 2021
Corollary 2.1. For fully-connected network (PCN) or convolutional neural network (CNN) ..... P1 o W1. we have Lip, 1 & under a choice of p-norm f= f= PK-1 o WK-1 with 1-Lipschitz non-linearities Pk-
The above bound is not necessarily tight; there are various works that compute tighter bounds for FCN and CNN (e.g. Virmaux & Scaman 2018 Fazlyab et al. 2019 Latorre et al. 2020
LIPSCHITZ CONSTANT OF SELF-ATTENTION
3.1 DOT-PRODUCT SELF-ATTENTION IS not LIPSCHITZ
Moving on, we investigate whether self-attention is Lipschitz. We first consider the widely used (scaled) dot-product multihead self-attention as formulated by Vaswani et al. (201)). Let æ1. Z1...., N be sequence of N clements, where x; RD for i ,,...,. N. We represent this sequence matrix E N such that the ith row of X is the ith element of the sequence, i.e. X₁ = x! Dot-product multihead self-attention (DP- -HH)) is map from RNxD to N consisting of H 'heads', where H chosen divide D. Each head is map from RN R. NDD// defined by
DP(X) = softmax (XWQ(XWX)T/VDD)
XW
(4)
where W'O, W.K W R DxDHH are learnable parameters specific to each head. The input the softmax is an N x N matrix of dot products (hence dot-product self-attention), and the softmax is applied to each row of this matrix. Finally, the outputs of all heads are concatenated into an N D matrix and are right multiplied by Do RDD, thus DP -HA is defined by
MHA(X)= = [DP' (x), (x)] w
(5)
In what follows, we will prove that MHA as defined above map is non-trivial, i.e. Wº, W W W j O. It is sufficient to show that a single head DP is not
Lipschitz, assuming that the MHA
Lipschitz, since MHA is linear combination of the outputs of each head. Let us write Equation {4) as (X) = PXWV, where P E R N X is the output of the softmax (we suppress the dependence of P on X to reduce clutter below). P is stochastic matrix, i.e. its entries are non-negative and its rows sum to 1. Since the rows of X are the x 's, a linear transformation of each v; by some matrix A equivalent to right multiplication of X by AT. So right multiplication of X by W is linear map are interested in the mapping
and thus Lipschitz. Therefore, mapping because itself is a non-linear function of X. In fact, we show that f is not Lipschitz, thus proving the first main result of the paper:
f(X) = PX; this is not linear
Theorem 3.1. MP-MAA is not Lipschitz for any vector p-norm || Ilp with E [o].
Summary of Proof. We use Theorem 2.1 noting that if the supremum of the norm of the Jacobian is infinite, then the mapping is not Lipschitz. In particular, we show that when æ; 0 for some i, some elements of the Jacobian of f grow proportionally to the sample variance of Eyzi which is unbounded.
Proof. We show the proof for the case = 1 (i.c. X e Nx a column vector) for readability, See Appendix Coor the general case, which follows the same logic.
The mapping f can be written as
{Xx)
f(X) = PX softmax (axx))xx x=
E N
(6)
(X)
where a WKW@ E R (we assume a + 0 such that self-attention is non-trivial) and fi(X) = Σn Pijar with P.T = softmax (ar,X). Hence f can be interpreted as a map of each x; to point in the convex hull of .1, .N. Since is map from N to Nx1 its Jacobian is
Jn
JIN
J, =
E NRNN
JN1
JNN
By taking partial derivatives we show that Jij axT T p(0) lejiX X + 81 PoTT where eij E NNN is binary matrix with zeros everywhere
where Jíj 81((X)
Boy E R.
3

---

Under review as a conference paper at ICLR 2021
except the (i. entry, 82z is the Kronecker delta, and p() == diag(P.) - P P... So for =
Ju axx pl + ax plo X P.
Let us investigate the scalar X p()) X. We observe that is in fact variance of a discrete distribution. Specifically:
X x x Pa - (Ex
Puæxa) Var(x);
where X is a discrete distribution with support at the inputs {z1..... ,N } and probability mass function given by their softmax probabilities P(X- =2) = Pıj. A consequence of this interpretation is that p(f) is positive semi-definite (PSD) since X p()) X = Var(X)= 2 0, with equality if and only the zy are all equal.
We use this observation to show that Jus is unbounded, and so ||ا l|p is unbounded, hence -HAMA is not Lipschitz. Consider the case I2 = 0. Then P softmax (X AA;; = N II, i.e. we have uniform attention regardless of digs. The first term of Jss in Equation (8) disappears since eg = [O..... -i,,.... 0 = O, and the last term becomes N I. Now consider the second term ax p()) =
(Vrr(X)). Note X is uniformly distributed, since P(X = xz) Pij 1/N. Hence the second term is equal to a times the sample variance of T1, IN, which can be arbitrarily large.
High-level intuition for proof. At x; = 0. f(() = N Σx Ty. the mean of the inputs. The rate of change of f, is governed by how fast the softmax saturates when 2; is perturbed, which is determined by how spread out the Vigl are. The more spread out they are (the higher the sample variance), the greater the rate of saturation of the softmax, and the faster the rate of change of f.. Since the sample variance of Tzz can be arbitrarily large, the rate of change of ft can also be arbitrarily large, i.e. the entries of the Jacobian (and hence its p-norm) can become arbitrarily large. In Appendix D.
we show
that adding bias terms to z; Wº and 2; W does not resolve the issue.
The implications of this result are the following. (1) There can be undesirable behaviour (e.g. training instabilities) for the Transformer when some inputs are close to zero. (2) Dot- not-product self-attention (and hence the standard Transformer) is not a suitable choice when we require Lipschitz neural network, such as for formulating invertible residual networks (Behrmann et al. 2019). Therefore, to use self-attention and Transformers in such applications, a Lipschitz formulation of self-attention required, together with explicit (idcally tight) upper bound to its Lipschitz constant, quantify how much the output can change with respect to changes in the input.
One method to make dot-product self-attention Lipschitz is by ensuring its inputs are bounded. Indeed, if the input space is compact, e.g. [0, NxD, any continuously differentiable function is Lipschitz, including dot-product self-attention. However, approach has its own challenges, since it makes the Lipschitz constant depend on the input range. Instead, in the next section we formulate version of self-attention that is provably Lipschitz on all of R N DD allowing us to derive an upper bound that holds for any subset of XD
further discuss in Section such an
3.2 L2 SELF-ATTENTION A LIPSCHITZ FORMULATION OF SELF-ATTENTION
The pathology in dot-product self-attention arises because the softmax probabilities P,, are constant with respect to xyy; when 2; = 0. This behaviour can be undesirable as we want P.j to vary according to xy. regardless of whether x; zero or not. Hence we propose an alternative form of self-attention based on L2 distance:
Pij x exp(L.g) = exp (- |o_ W xrJ TTT / D/HH.
(10)
with the normalisation constant ensuring that Σ, P.j = 1. We will refer to it as L2 self eelf-attention 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. 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. Similarly to dot-product self-attention, 1.2 self-attention can be computed efficiently with matrix operations; see Appendix E for details, with a comparison of wall-clock runtimes between different choices of attention.
We first 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 || |lpp for p = 2, co. The
4

---

Under review as a conference paper at ICLR 2021
full L2-MHA map F: RRXD RND is defined as -
...... (").,""(X)WW... W where P(X)=P"XAA X
In the above, WVA.B DDDD/HH W RDDDD ph defined as in Equation (TO) with WQ..
W.K.h E DDDD/H and A = wQ.. /VD/H E DxD, There are two changes from the usual form of multihead self-attention:
(1) We require M = WKA for each head (X) to be Lipsehitz. In Lemma F.1 of Appendix show that L2-MHA is not Lipschitz for arbitrary W WK. and that tying W = W is sufficient for L2 -HA to be Lipschitz, with intuition for why tying is sufficient.
(2) In each head of the self-attention fdh (X), right multiplication by An has been included for the theorem below to hold (details are in the proof). In practice, there is little harm done by this extra linear transformation, since when the heads are combined together in F, each f (X) is additionally transformed by WVS.,, free parameter.
The second main result of the paper is the following: Theorem 3.2. L2-MHA Lipschitz, with the following bound on Lip (F):
Lip (F)
max A AN 1)+ DD/H
|| 11 max WW
2lc |
1mx
b
and the the following bound on (Fp;
VN
(N (VEn II |Wv.i |
(F) S VD/H
O, and N is the input where ((x) = cxx((x + 1) is invertible univariate function sequence length.
Specifically, LN 1) Wo N) where Wo is the Lambert W- W-function, which grows sub- logarithmically as Ollog N - log log N) (Corlles el al. 1996). Hence the above bounds can be simplified to Ollog N) for p and O(vN Noog for
Proof. See Appendix F which uses the key observation that X p(() X is a covariance matrix (c.f. Equation 9 bound Jr pp: the norm of the Jacobian of F. Appendix G shows how the argument can be modified to prove the analogous result for the case with masking in the self-attention.
These bounds are complemented by the concurrent work of Vuckovic et al. (2020), which provides a O(VDlog N) bound Lip, (F) using measure-theoretic tools.
4 APPLICATION: INVERTIBLE SELF-ATTENTION
4.1 INVERTIBLE RESIDUAL NETWORK
Consider the residual function g(x) x f(x). Behrmann al. (2019) give the following sufficient condition for its invertibility: if f is a contraction with respect to some metric, i.c. if Lip(f) < 1, and the metric space on which f defined is complete, then g invertible. (A Euclidean space with [1, oo] is always complete.) Specifically, the inverse
a metric induced by p-norm (v) is the unique fixed point of the recursion x: = y f(z"), since by the definition of the inverse we have y = g (g) + f(g (y)). Because is a contraction, Banach's 's Fixed Point Theorem guarantees that this fixed point exists and is unique for all y. and that the recursion converges for all initial values z0 (often set to practice) exponentially fast. Hence the inverse can be computed to arbitrary accuracy (up to numerical precision in practice) by the above fixed-poin iteration.
||p for
Note that a composition of such invertible residual blocks is also invertible. [ehhrmann et al. (2019) use this observation design invertible ResNets: they take f to be CNN normalised by upper bound on Lip(f) given by Corollary 2.1 making the resulting function contractive. For the 2-norm < is chosen and each linear map (convolution) W in the CNN is multiplied by c/|lW]|p 2|2 if c < |||]|2 where [ is estimated by power iteration (c.f. Appendix BB. This multiplicative factor determines the scale of the Lipschitz constant of the normalised function.
|| ll2. hyperparameter
5

---

Under review as a conference paper at ICLR 2021
4.2 INVERTIBLE SELF-ATTENTION
The standard use case of self-attention is with a skip connection inside the Transformer. A Transformer block composed of residual blocks of multihead self-attention (MHA) and fully-connected (FCN) layers (Figure D. Hence similarly invertible ResNets, we can normalise L2-MHA by the upper bounds given in Theorem 3.2 to obtain Cont. vveavvieieiaaiiieeeeeeeeeeeeeeee -MHA f. with which we can obtain invertible self-attention g(x) x f(x)).
Output
LayerNorm
Dropout
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 DP-MHA in the Transformer with Cont r -2 hence replacing the residual self-attention module with invertible self-attention. We are not interested in the modified Transformer per se, but rather in comparing the properties of invertible self-attention to standard self-attention - we only use the Transformer as testhed for this purpose, since self-attention is commonly used in a Transformer. Given the theoretical focus of the paper, we believe that more challenging application of invertible self-attention, such as normalising flow-based modelling, would be more suitable as a separate paper focused on that particular application. In Appendix residual branch is also contractive.
FCN
LayerNorm
Dropout
MHA
H we show that Dropout the
Input
EXPERIMENTAL RESULTS
Figure 1: A Trans-
former block.
5.1 ASYMPTOTIC TIGHTNESS OF THE UPPER BOUND ON Lip (F)
A tight bound on the Lipschitz constant of self-
attention is desirable for all listed applications
20
in Section T it leads to tighter generalisation
18
bounds, lighter constraints for provable robust-
16
ness, and better expressiveness in residual flow
14
models. Hence we investigate the tightness of
12
our bound the Lipschitz constant of L2 -HHA.
10
LB: top out of 50 random init LB: top out of 50 random init
The Lipschitz constant is supremum over the space of inputs X E N (c.f. Equation (2) and approximating it requires solving an intractable optimisation problem. Hence it is infeasible to estimate accurately in general, es- pecially when X is high-dimensional. However, we may compute lower bound on the Lipschitz constant by maximising the norm of the Jacobian (X) with respect to X until convergence. This local optimum will form 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.
100
200
500
1000
N
Figure 2: Lower and upper bound on Lip (f) for L2-MHA f, with H D and varying N.
We use this observation provide empirical evidence for the asymptotic tightness of the upper bound on Lip (f) in Theorem 3.2. In Figure _ we show the upper bound as well as the lower bound on Lip (f) obtained by optimising [|J [ different random initialisations of X, with H D = and N varying between 100 and 1000. See Appendix [T for further details. Note that we use a log-scale for the x-axis, and recall that the upper bound is Olog N - log log N), dominated by the Ollog N) term for large N. Hence the plot for the upper bound shows a linear trend. 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.
with respect to X for L2 2H- with 50
There are at least two possible explanations for the gap between the upper and lower bounds. (1) The lower bound is only a local optimum - the true Lipschitz constant is global optimum across inputs, which can be difficult to attain especially for high values of N. (2) The multiplicative constant of the upper bound may be loose. Assuming asymptotic tightness, it remains open question whether the multiplicative constant can be tightened. We show the analogous plot for Lipz(F) and discuss the results in Appendix K Additionally in Appendix L we show that optimising [X)(X)l w.r.t. X for DP-MHA f causes the norm to diverge, providing empirical verification of Theorem 3.1 that DP-MHA is indeed not Lipschitz.

---

Under review as a conference paper at ICLR 2021
5.2 NUMERICAL INVERTIBILITY OF MHA RESIDUAL MAP
Recall from Section 4.1 that g(z) x + f(a) is invertible if f is contractive. Hence if f is Contractive-12--11
Inverting L2-MHA
Inverting DP-MHA Residual Map
10'
Residual Map
g is necessarily in-
10a
vertible. However, technically we do not dis- prove the invertibilitty of DP-MHA, since the converse does not hold in general i.e. if f is DP-MHA, which we have shown is not Lipschitz hence not contractive, it may still be the case that g is invertible. To verify that DP -MH (with the skip connection) is not invertible in practice, we compare the numerical invertibility of the residual map g(z) I + (x) between the cases where f is L2 MHA and DP -MH in Fig- ure For each, we take MHA with heads and randomly initialised weights, and quantify the maximum reconstruction error across batch of 128 inputs whose outputs are inverted via the fixed-point iteration described in Section 4. (see Appendix | for analogous results for a wider range of N and 1) and for DP-MHA with trained weights). To highlight the difference between the two types of self-attention, recall in the proof of Theorem 3.1 (showing that DP- -MHMAA is not Lipschitz) that when one of the inputs z; is O, some terms of the Jacobian grow with the sample variance of Tigs. Hence a set of N inputs where x; = 0 and Tzz are chosen uniformly at random. In Figure 3 we see that DP-MHA is not invertible whereas L2-MHA is invertible for sufficiently small This shows how not having the theoretical guarantee of f being contractive can cost us invertibility in practice. We note that the figure shows local invertibility the sampled inputs, opposed global invertibility across the whole input space, yet this clearly highlights the difference between the two choices of self-attention. Experiments with the globally invertible self-attention obtained by normalising with the Lipschitz upper bound are provided in the next section.
10
10
10
10
100 10
10aa
0.90
0.70
max 100
10
C= 0.50
5 fixed paint iterations
10
15
20
25
0
5
10
15 20 25
fixed paint iterations
Figure 3: Invertibility of g(x) = + cxx) where f is L2-MHA (left) and DP-MHA (right).
We use N = 64, D = 64, and E {0.5, 0.7, 0.9)
check numerical invertibility at
5.3 EXPRESSIVENESS OF 1.2MHHA AND INVERTIBLE SELF-ATTENTION
A natural question to ask is: how does the expressiveness of L.2 MHA and Contracto incetivee-ie ve (that leads to invertible self-attention with the skip connection) compare with the original DP-MHA? We expect that the Lipschitz constraint will limit the expressiveness of the Transformer, and would like find out by how much. We investigate this by comparing the performance of the original Transformer and the Transformer with invertible self-attention (c.f. Figure IT] at character-level language modelling on the Penn Treebank dataset (Mareus et al., [993).. We compare the test negative log-likelihood (NLL) of baseline LSTM, the original Transformer (DP-MHA), and series of models between the original Transformer and the Transformer with invertible self-attention (Contra ract i 12-122MMHR) making one change at time and tuning the hyperparameters on validation set. For Cont raactiieeeee.ee L2 MHA, we normalise L2- MHA by the bound on Lip (P) as it is tighter than the bound on (F). See Appendix I for experimental details.
LSTM
Original Transformer (DP)
Transformer (L2)
Transformer (12), W
Transfommrr (Contractive-12,, p=>
1.8
Num Layers
Num Layers
Mum Layers
Num Layers
64
128
1.6
256
512
Text
1.2
1.0
20
40
60 opoch
80
100
20: 40 60 opoch
80
100
20
40 60 apoch
80
100
20
40 opach
60
80
100
30 100 epach
150
200
Figure 4: Test NLL curves during training for various LSTM/Transfomerr models. models.
7

---

Under review as a conference paper at ICLR 2021
The results are shown in Figure The first plot shows the best performing LSTM reaching test 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. We observe instabilities in training for 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). The third plot shows results for the Transformer with DP 8P-MAA replaced with L2 MHA but without tying W and W and we observe a very similar test performance. The fourth plot shows the change when we further tie the query and key weights (making W = W*)): ); we see that there is a small degradation in performance. 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. We note that performance saturates at around 5 layers for each Transformer model so far. 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. 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.
Thus we conclude the following, (1) Replacing the dot-product with the L2 distance incurs hardly any loss in expressiveness.
(2) Tying the query and key weights to obtain Lipschitz self-attention incurs a small loss in expressiveness. (3) Dividing by the upper bound on Lip (F) to obtain invertible self-attention incurs 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. We show further experimental results that compare the training stability of DP-MHA and (Cont ractivee -areetivvee----eeee L2-MHH in Appendix N
6 CONCLUSION AND DISCUSSION
We have shown that the widely used dot dot-product self-attention is not Lipschitz, and that the proposed L2 self-attention is Lipschitz, by deriving an O (log N 1 log log .\ Lipschitz bound for p = oc and an O(vN (log N log log N)) bound for 2, where N is the input sequence length. We also provided empirical evidence of the asymptotic tightness of the bound for p co. Finally we demonstrated that
constrained self-attention can be used to formulate invertible attennttenion, which we experimentally evaluated on a character-level language modelling task.
Our approach to Lipschitz self-attention has been to replace the dot-product kernel with an L2 kernel. An alternative would be to constrain the inputs of self-attention to be bounded; if the input space is compact, e.g. [0. N any continuously differentiable function is Lipschitz, including dot-product self-attention. However, while being simple to implement, this solution has its own difficulties. First, it makes the Lipschitz constant depend on the range of the input, and thus obtaining tight bound would require non-trivial mathematical work. We stress that a guarantee that the function is Lipschitz does not tell us anything about its Lipschitz constant; without tight Lipschitz bound, the true Lipschitz constant can be very large, at which point it is unhelpful that the function is Lipschitz. Second, since self-attention is typically applied at multiple layers within 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. 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. This however can have undesirable side effects such as vanishing gradients when the sigmoids are saturated (Hochreiter 1998). Despite these difficulties, this could be a worthwhile alternative route for obtaining Lipschitz self-attention explore in the future.
Having a provably Lipschitz self-attention module at our disposal makes it possible to use Transformer- based architectures in applications requiring Lipschitz constraints, while enjoying theoretical guaran- tees. A natural application of Lipschitz self-attention is for residual flows (Behrmann et al. 2019). and for parameterising Neural ODEs (Chen et al.. 2018) where a Lipschitz vector field guarantees the existence of unique solution to the ODE for all times. These models can be used for density estimation and generative modelling of sets. Another interesting direction for future work would be to analyse different variants of self-attention based kernels other than dot- dot-product and L2, as Tsai et al. (2019) do from an experimental perspective, for which we believe the mathematical tools developed in this paper may aid the analysis.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Abadi, M., Agarwal, A., Barham, P., Brevdo, E., Chen, Z., Citro, C., Corrado, G. S., Davis, A.,
Dean, J., Devin, M., al. Tensorflow: Large-scale machine learning on heterogeneous distributed systems. arXiv preprint arXiv:1603.044. 2016.
Anil, C., Lucas, J., and Grosse, R. Sorting out Lipschitz function approximation. In International Conference on Machine Learning, pp. 291-301, 2019.
Bapna, A., Chen, M. X., Firat, O., Cao, Y., and Wu, Y. Training deeper neural machine translation models with transparent attention. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 3028-32 2008-,33, 2018.
Behrmann, J., Grathwohl, W., Chen, R. T. Q.. Duvenaud, D., and Jacobsen, J.-H. Invertible residual networks. In Proceedings of the 36th International Conference on Machine Learning, 2019.
Brown, T. B., Mann. B., Ryder. N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A..
Shyam, P. Sastry, G., Askell, A., al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. Neural ordinary differential equations. In Advances in Neural Information Processing Systems. pp. 6571-6583, 2018.
Chen, R. Q., Behrmann, J., Duvenaud, D., and Jacobsen, J.-H. Residual flows for invertible generative modeling. In Advances in Neural Information Processing Systems, 2019.
Cisse, M., Bojanowski, P., Grave, E., Dauphin, Y., and Usunier, N. Parseval networks: Improving robustness to adversarial examples. Proceedings of the 34th International Conference on Machine Learning, pp. 854 863, 2017.
Corless, R. M., Gonnet, G. H., Hare, D. E., Jeffrey, D. J., and Knuth, D. E. On the Lambert W function. Advances Computational mathematics, (119)3329399)9 1996.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics, pp. 4171-41866 2019.
Fazlyab, M., Robey, A., Hassani, H., Morari, M., and Pappas, G. Efficient and accurate estimation of Lipschitz constants for deep neural networks. In Advances in Neural Information Processing Systems, pp. 11423-11434, 2019.
Federer, H. Geometric Measure Theory. Classics in Mathematics. Springer Berlin Heidelberg, 1969. ISBN 9783642620102.
Glorot, X. and Bengio, Y. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256 2010.
Grathwohl, W., Chen, R. T. Q. Betterncourt, J., Sutskever, I., and Duvenaud, D. FFJORD: Free-form continuous dynamics for scalable reversible generative models. In International Conference on Learning Representations, 2019.
Hochreiter, S. The vanishing gradient problem during learning recurrent neural nets and problem solutions. International Journal of Uncertainty, Fuzziness and Knowledge- Based Systems, 6(02): 107-116, 1998.
Huang, C.-Z. A., Vaswani, A., Uszkoreit, J., Simon, I., Hawthorne, C., Shazeer, N., Dai, A. M., Hoffman, M. D., Dinculescu, M., and Eck, D. Music Transformer. In International Conference on Learning Representations, 2019.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, 2015.
Latorre, F., Rolland, P., and Cevher, V. Lipschitz constant estimation of neural networks via sparse polynomial optimization. In International Conference Learning Representations, 2020.

---

Under review as a conference paper at ICLR 2021
Marcus, M. P., Marcinkiewicz, M. A., and Santorini, B. Building a large annotated corpus of English: The Penn Treebank. Computational Linguistics, 19(2):113 1993.
Mises, R. and Pollaczek-Geirrnger,, H. Praktische verfahren der gleichungsauflösung, ZAMM-Joura of Applied Mathematics and Mechanics/ZZitscrriff für Angewandte Mathematik und Mechanik, (2):152-164, 1929.
Miyato, T., Kataoka, T., Koyama, M., and Yoshida, Y. Spectral normalization for Generative Adversarial Networks. In International Conference on Learning Representations, 2018.
Parisotto, E., Song, H. F., Rac, W., Pascanu, R., Gulcehre, C., Jayakumar, S. M., Jaderberg, M., Kaufman, R. L.., Clark, A., Noury, S., Botvinick, M. M., Heess, N., and Hadsell, R. Stabilizing Transformers for reinforcement learning. In International Conference on Machine Learning, 2020.
Parmar, N., Ramachandran, P.. Vaswani, A., Bello, I., Levskaya, A., and Shlens. J. Stand-alone self-attention in vision models. In Advances in Neural Information Processing Systems, pp. 68-80, 2019.
Peyré, G. and Cuturi, M. Computational optimal transport. Foundations and Trends in Machine Learning, 1(5-6))355-607, 2019.
Sokolić, J., Giryes, R., Sapiro, G., and Rodrigues, M. R. Robust large margin deep neural networks. IEEE Transactions Signal Processing, 65(16)44265-44 4280, 2017.
Tsai, Y.-H. H., Bai, S., Yamada, M., Morency, L.-P., and Salakhutdinov, Transformer dissection: An unified understanding for Transformer's attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, pp. 4344 4353-333,3300 2019.
Tsuzuku, Y., Sato, I., and Sugiyama, M. Lipschitz-margin training: Scalable certification of perturba- tion invariance for deep neural networks. Advances in Neural Information Processing Systems, 2018.
pp.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,, Gomez, A. N., Kaiser, E.,, and Polosukhin, I. Attention is all you need. In Advances in Neural Information Processing Systems, pp. 5998-6008, 2017.
Virmaux, A. and Scaman, K. Lipschitz regularity of deep neural networks: analysis and efficient estimation. In Advances in Neural Information Processing Systems, pp. 3835-3844, 2018.
Vuckovic, J.. Baratin, A., and Tachet des Combes, R. A mathematical theory of attention. arXiv preprint drXiv:2007.02876, 2020.
Wang, X., Girshick, R., Gupta, A., and He, K. Non-local neural networks. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 7794-7803, 2018.
Zhang, H., Goodfellow, I., Metaxas, D.,, and Odena, A. Self-attention Generative Adversarial Networks. In Proceedings of the 36th International Conference on Machine Learning, pp. 7354 7363, 2019.
10

---

Under review as a conference paper at ICLR 2021
A CHAIN RULE FOR VECTOR VALUED FUNCTIONS
In this section, we list some useful identities for deriving the Jacobians of the expressions in the paper.
Suppose À is scalar, u, v, are column vectors, and [u) is vector valued function. We use the standard convention that for a E Ih b E TR" we have de E [gmxn, Then we have the following chain rule identities:
86
Ba
a
[uu]
Ba
u
ax
af(w)
(f(u) 0uc
8iz
Un
8m
Tv]
du dz
Bm
On
Note 0A is vector, so Om
0%
is matrix.
B POWER ITERATION
Although || can be computed efficiently in O(um) time for W E Rm naïvely computing o mmx (W) = Vmmax (W' W) requires O(n²) operations. (By (A) we denote the greatest eigenvalue of symmetric matrix A.) We can however obtain an underestimate (W) via power iteration:
||W||2
W WTwbe
b W
Wbb
bk+1
(w)
(11)
I Wha
by. by
with each iteration taking O(n²) time. Then using K & iterations gives an underestimate ěж in O(Kn²) time. Since this is an underestimate, the resulting approximation the Lipschitz constant of the linear map will not be an upper bound. However the number of power iterations is usually chosen so that ō is accurate enough - K = 5 is shown to be sufficient in the context of fully connected networks or convolutions considered by Behrmann et al. (2019).
The iteration will converge W has an eigenvalue that is strictly greater magnitude than its other eigenvalues, and the starting vector by has nonzero component in the direction of an cigenvector associated with the dominant cigenvalue. This happens with probability if by is chosen at random, and the convergence is geometric with ratio / where X2 is the eigenvalue with second largest magnitude (Mises & Pollaczek-Geiringer 1929).
C PROOF OF THEOREM 3.1 FOR GENERAL D
Theorem 3.1. DP-MEA Lipschitz for any vector p-norm || Il with p E o0).
Proof. The mapping f can be written as
f(x) PX- softmax x TX(
NRD
(12)
(x)
where wxwe /DD/H DXD and f((X) = En Paya, with P.T softmax (XAr;)
Hence ff can be interpreted as map of each x; point in the convex hull of Tı, Ny. Since is map from ND to NDD, its Jacobian is
Jn
JIN
NDNND
(13)
JNI
JNN
where Jay 8f.(X)
RRDx By taking partial derivatives show that Jij
dz;
p()() [exXA A Ad P.ş where eij NxN is binary matrix with zeros everywhere
11

---

Under review as a conference paper at ICLR 2021
except the (i entry, 8iz is the Kronecker delta, and p() diag(P.)) P P... So for =
Ju X xTo x p XA+ PrI
= Pa (z; En PR ai A x pl) XA+
(14)
For the last equality, note ex has all rows equal to zero except for the ith row given by x... We can then verify that X e simplifies (a Σ. P..
For vector p-norms, 1 is bounded if and only if its entries are bounded, by definition of the operator norm. The entries of X p(t) X A are bounded for arbitrary A only if the entries of X p()) X are bounded. us investigate the entries of this D D matrix. Writing out each term of the matrix, observe that it is in fact covariance matrix of discrete distribution. Specifically:
[X p(o Σx PATTITKM - <E PAETRI) (Ex PskTtk)) Cov X.),
(15)
where X. is discrete distribution with support at the inputs { f...... #N) and probability mass function given by their softmax probabilities P xz) = Pij. A consequence of this interpretation is that p() is positive semi defi-ief (PSD) since for D = 1, Equation (15) becomes X p() X
Var(X (x) 2 with equality and only if the 2y are all equal.
We use this observation to show that the terms of Já are unbounded, and so DP-MHA is not Lipschitz. Consider the case x; O. Then P softmax (X Axx))= = N Il, i.e. we have uniform attention regardless of Zyei. The first term of J in Equation [4] disappears since
O, and the last term becomes $ I. For the second term, the entries [X p() Var(X;) are unbounded since the NNI, which can be arbitrarily large.
latter is equal to the sample variance
D BIAS TERM IN DP AELE-ANTENTION
natural question ask whether
add bias terms 8 to x. W and 1 s. W resolve
the issue of attention weights P:: becoming uniform when d; = 0. The answer is no in general. again be shown that Ju is unbounded when 2; is chosen such that x; W b2 = 0 (such
choice is possible assuming WQ is full rank, a dense set in le D/HHH, Then P.T = & again, and
the diagonal entries of X p(t) X unbounded.
E EFFICIENT COMPUTATION OF L2 AELE-ATTENTION
Dot-product self-attention only requires few matrix multiplications to compute the logits (i.e. the inputs the softmax) between all pairs of inputs, without having to loop over pairs, hence it can be computed efficiently. Similarly, we can show that L2 self-attention can also be computed in an efficient manner. Using the identity ||a b113 = |all 2a + ||b|[ we can compute the logits of L2 attention between all pairs via matrix multiplications and computation of row-wise L2 norms, with negligible overhead compared to dot- product self-attention. Specifically, for L2 self self-attention can show that
|X
2XW2 (XWK)T + + 1 row
2T
P = softmax
(16)
VD/H
where || applies the squared L. norm to each row of A, so mm then || E m
In Table we show the wall-clock training times for the Transformer models with different attention functions and a varying number of layers. It is evident that the differences between the models are rather small.
12

---

Under review as a conference paper at ICLR 2021
Layer
Layers
Layers
Layers
Layers
Transformer (DP)
37
56
77
92
110
Transformer (I.2)
35
99
115
Transformer, Wº W (L2)
39
58
79
91
108
Transformer, (Contractive-L22)
37
60
81
102
127
Table I: Wall clock training times for one epoch of training (seconds)
F
PROOF OF THEOREM 3.2
Recall the formulation of L2-MHA:
F. NXD - NXD
F(X) [X(X) (" wo
(x(X) ph XAA
|| Nr.TwR.- - 23 WKK.K.
phj exp(Lin) = exp
Σ =1
VD/H
where have that WQth, W.k.h, W RDDD/H Wo RDDD ph NxN and AA W... WD/H DxD, and the softmax is applied each row of the input matrix. Recall Equation
16
XW 1 2X (XWKM)T + I|XW K.h
ph softmax
VD/H
F.1 L2 SELF-ATTENTION
LIPSCHITZ FOR GENERAL WΩ, W
Let us first look at the case of И = 1 and suppress the index h to reduce clutter. Consider the map f(X) = PX. so f(X) AA We need f to be Lipschitz for f and hence F to be Lipschitz. Note that defined as:
(_1ee, | W x WKWK
Pe x exp(T.i)) exp
VD/H
and the normalisation constant satisfies E, Piy P P NXN X
For 1.2 self-attention, we may take partial derivatives and use the chain rule to show that the Jacobian
of f is:
JIN
J :
E NDXND
(17)
JN1
JNN
with
Je x p(1)
The dx,
+ P E DXD
(18)
where
ÖL:
[XW* 1x, w.) W dij + -
- -
wk
(19)
dæ,
VD/H
and
P.a Pa) -PoPa
-PaPP Piz(1 P22)
-PiPPN PPePPN
p(() diag(P.)
P. P:
PaNPa
-PiNP
Pa (1 P.N)
13

---

Under review as a conference paper at ICLR 2021
exp (-1=;; W x; (xlkk)
P.
Σn exp (-1|z; W - 2g w
Recall that eji E RNNN is a binary matrix with zeros everywhere except the i)th entry. Hence e has all rows equal to zero except for the jth row given by x." We can then verify:
x = Piy(as - Σ Pa.x)aa.-
(20)
Also note p(1) is symmetrie, and each row/colum sums 0, i.e. p 1 T(0 0. Hence may simplify the Jacobian terms as follows:
Jas
2
x xXWW (XW* 1x 9 + +x x plo p WK Pal
VD/H 2
XTpl - Ix. + Pila,, Σ
P "
+ P
VD/H
h
2
x xppxXxWK'' + Pu((z Σ Pu (WQ WKIwE
+ P.I
VD/H
(21)
and for i j::
2
Jij =
X -cxxWW--x P.jI
VD/H
VD/H Puflar E
(T w = =[ P.yI.
(22)
We are now ready to show that f is not Lipschitz for general W w Lemma F.1. If E Dx D/H is full rank (i.e. full column rank), and W # W then Jig has terms that are unbounded for + j. hence f not Lipschitz.
Proof. Let us investigate the expression Kij = P. (x,- E. (x.
Te W WK) e
T R ft for + j. which is related J.y as follows by Equation (2).
WK J.p =
(VD.T
2
Kij + Popl )))) wX.
suffices to show that K, is unbounded show that J.y is unbounded, since WK is full rank and P. E [0. Let vJ_ xx WQQ xxJ W.k Then we have:
U1 Σ
Pa W x4 WWK xj SS-EPa(W Is I,- I- IN)
k
= -WQ x, WW x, -WWQ EPaw xx
--WK (xz Σ PaIx).
Hence Kí P33 Σ. P.. Note ys can take an arbitrary value RD/H since WK * WQ and W K is full-rank.
For all j + let choose z; such that yj -
Ve. This is possible for any value of yi since
W is full-rank. Note y - and not ye. We then have that || [1 is equal for all j, hence
14

---

Under review as a conference paper at ICLR 2021
Poj == Ex exp( like (2
N for all j. Then for i j. Kig simplifies
Riy (- 91 1 (N 2)(- (-vi)²
2N - 2
N2 Moll
whose entries are unbounded since yi can be any vector in RD/HH (note we assume N 2 for self-attention to be well-defined, hence 2N- - 2 + 0).
The intuition for this result is as follows: a reason for DP -MA not being Lipschitz is that for z; = O,, the attention weights Pt become uniform regardless of the values of 2, for # i. A similar issue arises for L2 MHA with WQ * W and full-rank W shown above: given any Ti we can choose x such that the P.j become uniform.
F.2 L2 SELE-ATTENTION LIPSCHITZ FOR W = W
Hence we impose the restriction that W = Wº. With this assumption we have
Py x exp (-1(x;--x;) -
VAIl2
(23)
where = wQwQ" VD/H DDD and VA is chosen such that A VAVA in particular VA = WQ/(D/H) H|) i. The terms the Jacobian of f simplify to:
Jn = 2x p()) XA + Pat (note p(0) = 0),
(24)
Jıy (z) E Pa (2,))
TA+ P. for fori
(25)
Let the Jacobian of S(X) be:
J11
JIN
J, =
NDNND
(26)
JN1
JNN
Since
= -(x)=((x)) and by the chain rule
8 Jz, [x((A = AT ATO T(x(x)
0fx(X) Jz, dzz (by symmetry
of A), we have that Jij = AJiz. Hence
Jn = 2AX p() XA + PAA (note p(0 = 0),
(27)
Jú 2P, A(2, Σ (a. - T,,T TA PijA for i jj.
(28)
Noting (f) = supx (X)]pp we would like upper bound ||
F.2.1 UPPER BOUND ON Lip (F) FOR L2-MHA
Consider the choice p- oc, where ||ا the maximum absolute row sum of Jy. key observation is that if we can bound the co-norm of the Jacobian of fi, single output of f (i.e.
single block row [[[ .... JIN]ll oc of Jy), then this is also bound on ||J due to permutation equivariance of self-attention; all block rows have the same maximal || Il>o when each is optimised over the input X. Using this, we can prove that ||S admits an upper bound that is Ollog N log log N). Below state and prove lemmas that lead the proof of this upper bound.
First we analyse the term VA T p()) X A, that appears in the first term of Jg. Note that for = VA, so that the rows of Y are == VA, we have
A X p()) x V A = p(oy = Cov(Y)
(29)
where P(Y y,) observation in Equation
=, exp( Hppp--ly-
|ll
911l Σn exp(- vill3). The last equality uses the
The central inequality used throughout the proof of the main theorem is the following:
15

---

Under review as a conference paper at ICLR 2021
Lemma E.2. Tr(Cov(Y)) = P,yll,,, - Ex
S & voll2 - x- 1) where
((c) cexp(e + 1) is one-dimensional invertible function on Rzo-
Proof. The first equality holds since Tr(Cov(Y)) = E, Cov Cor(Y)ss = E, Var(Y;) = Σ, B[(Y, -
E|y I. The next inequality holds since Var(W,) = Var(Y,) = B,W,, E[ S E( where
Y=Y Y y/- The final inequality can be proved as follows.
We would like to bound
E Pollus vill Ej 183 34 12 exp( 11y - 3x113)
Σ, 23 2xp(-23) exp(
(30)
Ex exp( (111 - yi vl)
where 2j 11ys villa (hence 24 = 0). Define:
E, 2 exp( -23) Ejti 33 exp(exppp -23)
s(=)
(31)
Σx eexp(=2]
1+ Enzs exp( =
First note that as 2j + oc, exp( -2 + O exponentially fast, causing the product 2) exp(-22) - 0. Hence we expect the above quantity to be bounded and attain its maximum.
Let h(2z) = exp(-23) for notational conciseness, and note h(zz) >
O. By taking partial derivatives
with the chain rule, we have that for j +
0g(2)
2y,h(z,) (Ex (2x)))
(1- 29) Σ h((x) + Σ h(zx)
(32)
0z,
Hence the derivative is 0 if and only if 2j = or (1 - 23) E. h(z6) + En h(x)==² = 0, the latter + 9(2). Hence at the maximum, the non-zero values
being equivalent to 2/ = + Σ, M(
Σ. h(xx)
among {2 must be equal to one another. It is clear now that the maximum value c is attained
when 2) = 1 + for j * i (and recall = 0). So h(2,) = exp(- -1 c) for j a i. Substituting this into g(2). ) and rearranging, we obtain cexp(e + 1) = N - 1. Note o(x) = x exp(2 + 1) is
increasing for x > 0 hence c = &= '(N - 1).
Note ologg N) (log N) exo N +1) 2 N log 2 N N 2 3. Since d is increasing, we have o -1 ((N - 1)= (N) for N 2 3. In fact, is known that क (N - 1) = Olog N w llgg log N) (Corless et al. 1996).
Note the A term in f(X)= = f(X) A allows us to use the above inequality, since YT p(iy = CVv( now appears the terms of Jy:
J/ 2/ p0) YIVA" +
(33)
Jaj = (3)
Σ (u - vz) VA + P.jA fj.
(34)
Using the inequalities |CC|| < ||I| |++C|| $ || ||||| and [[,,,..., AN]l S E 11,11 we have:
[|......, JeN]le 1xx
Jallo + Σ |
J+1
<2 VAll P(O Ilol 1100 + Alll
+2 2
Σ
| PiplUs Σ
Pawn)(v uz) Iooel 1mx
+ A
PigllAlloo
=2|1 22l/v/|ll>>111 VIVA 1000 (IX PYtp +
-5,) Σ |P Σ Pa - y)) ll0e) +A
st
| 6z I 1100
(IYTp69yllo
|w 1bx
Σ
|P (33 - Σ PuRyn)(u - 115) 11x0
+
VD/H
VD/H
3
16

---

Under review as a conference paper at ICLR 2021
For the first equality, note that E, Psj = 1. For the second equality, note that the summand for j i
is O because the term y/: Yj 0. Each of the terms in the brackets are bounded by the following
lemmas:
Lemma E.3. | p 10 S 'N D/H (o defined in Lemma F.2
Proof. Recall that y Cov(Y). Let o (mm) denote the standard deviation of Y.. Then [Cov( I {W(W))((Wm)) Hence
[Coo(Y))lo max Σ [[Cov(Y)]an maxx(Mi Σ (8mm)
D
2 m) / PH Tr(Cov(Y)) H
H
D (N 1),
$ V H o
since Σm (Ym) < VAVEE (by e.g. using the Cauchy-Schwartz inequality on (ooY)), ((()))....,.(....... and 1) and max (Vi) S v (Ymm, and the last inequality is from
Lemma
Lemma F.4. Σ, (i) Ex Pu - 1;) (N- -
Proof. Note llab 1000 ||aa| lb6l for real vectors a, Hence
Σ P1 - Σ Pa - Tloo = Σ Pilllys Σ
Pasalloollu - 2ylla
j
s ||l|2|16|l2.
where aj Pijllys En Puvylloor PPyl Us yo 111
Note a, S oj = / /Psjl [ly} En Punsalle ll2 since ||~||>> S ||>|| for vector æ. Hence ||a||2 S ||a||2-
Also bj S / d2 = v R Pujllus vilh ll since |||||| v / |0ll2 (e.g. by the Cauchy- -echw-SS
inequality on [|-]|]--- |:1.....,////// b/D/ and for RD/H Hence ||b||2 v 1ld22
Note ||c|| Ej Pujllus- - -E Pw = Tr(Cov(Y)) (Y))
x (N 1) from Lemma R2 and ||d|2
Σ, Pby |1 y, 12 S $-'(N - 1) also from Lemma R2 Hence ||||2||bll S v 2 ||d|a S
V 0 (N 1).
Putting the above lemmas altogether, with the observation supx (X) supx (X)......... (X)]]] by permutation invariance of || (since f is permuta- tion equivariant and || oloo is the maximum absolute row sum), we have
lww@wQ 1
| < | 1 -N - 1)
VD/H
46 4N- - -1)+ VD/H < 122 1ox
(35)
| bx0 I
Iloc
4log
VD/H
where the last inequality holds for N 2 3.
The full multihead attention map that combines the heads fM(X) is:
F + [X)(X) ..... S " (X)WV.HI w =gXX)WW wo
17

---

Under review as a conference paper at ICLR 2021
where g x (X),,
(x))), wo E ER and
WV1
W
E ER DHID
WV
Note the Jacobian Jg is a block matrix whose rows are yn, hence || max g1 lloo. and similarly ||ll 11 = max |WM. Iloc. Hence we have
Lip (F) S max I max [[ 1|x0 || 11x0
Combining this with Inequality (35. we have:
Li (F) < 46 AN
|o
11 max |WV.
|owl 1lo0
10x.
max
VD/H)
b
F.2.2 UPPER BOUND ON Lip, (F) FOR L22MHA
For = 2, we use the following lemma:
Lemma F.S. Let A be block matrix with block rows A1,..., AN. Then || S vEi and equality holds and only if the first right singular vectors of the A align.
Proof.
A1
A
||-11/2
:
sup ||
sup
Σ 11
sup Σ
|| = Σ 114:112
||=|||=
|=|
AN
AN
Note that equality holds if and only if the first right singular vectors of the A align.
Hence a bound on the spectral norm of each block row of Jj can give us an ((NN) bound on || which may be loose, and it remains an open question to whether this bound can be tightened.
To bound the || ||2 norm of each row of Jy. we use the following lemmas:
Lemma F6. || y I I 1)
S (N Proof. ||0 = (CCo(YY)l = Amax (Cvv(Y) < < first equality holds by symmetry Cov(Y) and the next holds by Cov(Y) being positive semi- definite, so all its cigenvalues are non-negative, and hence the maximal eigenvalue is bounded by the sum of the eigenvalues, equal its trace. The final inequality from Lemma F2
where the
Lemma F.7. E, P.,(, Σx Pu 9;) o NNN(NN-1
Cauchy-Schwartz Proof. Directly Cauchy- Schwartz and d in the proof of Lemma E.4
Again using the inequalities ||BC||| <
|B || CII. | + ClI <
|B
+ |C|| and |[[ Ax 5
E, ||ll with the additional equality B 1l
Bl12, we have the bound:
IIla,....,
JeN]l2
|| ||2 |
112
Myyppoyy: + Σ
[Pr Σ
Purn) - 11) + |wewee l
VD/H
VD/H
< 440 (N - 1) 1 lwwwe
DD/H
VD/H
46 ' - 1) +1)
DD/H
18

---

Under review as a conference paper at ICLR 2021
Using Lemma F5
we have that
V VNI
|J S
44 11 -N1)+1 VD/H
(36)
V
(4 + 1).
VD/H
To obtain the final result for the full multihead self-attention F, we need a final lemma:
Lemma E8. Let A be block matrix with block columns Ai.... AN. Then ||A|| S V. |
Proof.
2
Ø1
||Al|2 = [||A1..... ANlll.
sup
[A
Anl
sup
|| E A,ail2
Σ, 1|=:113=1
2
$
sup
|A llelll=1, x==1
sup
> A1|lA
sup
Σ A111A,112
Σ. 1|x.113=1
Σ, x==1 ;
$ |.
where using the substitution x, dee, and the last inequality holds by e.g. Cauchy-Schward inequality on [Ai...., Ay] and [||A1]12..... || [12].
Recall that
F: + ((())))))
(X)WV. wo
Since (X)
I 1| 122 12/1 by Lemma F.8 we have that
|[F " (X)
E 11 13
and hence
Lip S
Σ 10 1Bl ||
(37)
Combining this with Inequality 36
we have:
Lip
VN
(40- -NN-1)+1
1 (VEnIIWeM13
|w ||
VD/H
G THE CASE WITH MASKING
Since self-attention is often used with masking, a natural question is how masking affects the derived bounds. In self-attention (for any choice of attention function), masking is implemented as follows: given set of mask indices M C (1..... .) x (1,..., N.....N.. the logits (i.e. the inputs to the softmax) are set to -o the mask indices. That is,
if (f.j) & M
Lij
-o0 if (i,j) M
where T.g the original logit (e.g. for 1.2 self-attention, Tag -x; (x;-x;;) TT A(xr - xz)))
Masking implies fa(X) is not function of æj for (i,j) M,, hence J.j 0 for (i. E Thus f((X) is equal to the ith output for self-attention with inputs restricted to {z, (I.j) $ M). the unmasked inputs with respect to the ith output. Hence Jij will no longer contribute to the bound on [[[a....,, Jenllll and hence the bound for the unmasked case will continue to hold as long as (i, E i.e. i.e. 2; attends to itself (this is necessary for the proof of Lemma F.2 to hold). The bound can in fact be tightened by replacing N with Hay : (1, & M)]. the number of unmasked inputs with respect to the ith output.
19

---

Under review as a conference paper at ICLR 2021
H DROPOUT IS CONTRACTIVE
At test time, Dropout multiplies inputs by the dropout keep probability < 1, so it is a contraction with Lipschitz constant p evaluation time. At training time, Dropout amounts to setting some inputs to zero, while keeping other inputs constant. This can be expressed as right multiplication by diagonal binary matrix M, and for such matrices we can verify || == sup m1,==1 || S 1.
I EXPERIMENTAL DETAILS
For the experiment in Section 5.1 showing the asymptotic tightness of the upper bound on Lip (F) where F is L2-MHA, fix all free parameters of F (namely W WV) to be the identity, and only optimise the input X. We use 50 random initialisations of X for each N, where Xij ~ ' (-c, c) for ~ 20, 10] (we observed that having c itself be random improves optimisation). We display the top 5 results for each value of N after optimising each random initialisation till convergence using Adam
(Kingma & Ba 2015) with learning rate of 0.1.
For the experiments in Section 5.3 we comparing the performance of the original Transformer and the Transformer with Lipschitz/nveertbble self-attention at character-level language modelling on the Penn Treebank dataset (Marcus et al. 1993 Each training example is a sentence represented as 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. 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 batch size of 64. The LSTM models have the dimensionality of the hidden state equal to the dimensionality D of the cell state (the usual default implementation). The Transformer models are trained with varying number of blocks (number of layers) with H = heads and D = 512, tuning hyperparameters for dropout rate in (0, 0.1 0... and base learning rate & E (0.2, 0.4, 0.6, 0.8, 1.0, 1.5, 2.0} with number of warmup iterations w E (1000, 2000, 4000, 8000} for the standard custom learning rate schedule Vaswani al. 2017)
7 min(t (tin(t-Hu 1/2 tw 33/2),
VD
where is the learning rate at training iteration t. Hence the learning rate linearly increases from O to (Dw)-1/2 over w iterations, then decays proportionally to t 111 We use Glorot Uniform initialisation (Glorot & Bengio, 2010) for all weights (U
1
din douc ) except for
de Adout
[- vD and s is a hyperparameter. For = 512, weights in 1.2-MHA that are initialised from U
we used Κτ. All experiments were done in Tensorflow 1.14 (Abadi et al. 2016). The code will
be released upon de-anonymisation.
In Table 2
we show the best Test NLL across training of Transformer models in Figure
Number of Layers
2
3
4
6
8
10
12
14
(DP)
1.164
0.989
0.954
0.943
0.940
(L.2
1.134
0.987
0.958
0.942
0.946
W W (L2)
1.237
1.055
1.001
0.994
1.023
(Contractive-L.2
1.558
1.453
1.343
1.191
1 1.205
1.119
1.086
1.0495
1.015
1.005
Table 2: Test NLL for Transformer models on PTB character level language modelling
J NUMERICAL INVERTIBILITY OF MHA RESIDUAL MAP
Following Section 5.2 Figure 5 confirms that numerical invertibility does not hold for trained weights
for dot-product multihead self-attention (DH-MHA) (obtained from one-layer Transformer (DP)
We use the standard training--aliddation-tsssss split, and the dataset can he found at e.g. htttps://giiii com/ hom/ha /ttp////eeeeeee//e/oooroerrerrrr
20

---

Under review as a conference paper at ICLR 2021
model used for Figure 4. similar to the randomly initialised weight case. Figure shows additional results for different values of N and
Inverting trained DP-MHA
Residual Map (D=512,H==)
max reconstruction error
107
10a
10a
0.90
10-1
0.70
0.50
10-3
10 15 20
25
fixed point iterations
Figure 5: Invertibility of g(x) = r + (z) for trained DP-MHA f.
Inverting L2-MHA Residual Map
Inverting DP-MHA Residual Map D-64
max reconstructionn error
10
D-64
D-128
D-128
10a
10a
10
10
N-64
10
10:3
10-1
10.
10.6
10
max reconstruction erro
103
10a
10
10
128
10
10
10
0.90
10
10
c 0.70
10
c 0.50
10"
25 fixed point iterations
10
25 fixed point iterations
10 15 fixed point iterations
10 25 fixed point iterations
Figure 6: Numerical invertibility of g(z) = x + cf(x) where f is L2-MHA(left) or DP-MHA (right), for different values of N and D.
21

---

Under review as a conference paper at ICLR 2021
K
BEHAVIOUR OF LOwER BOUND ON Lipz
8
7
6
5
4
c*logN
3
top 1 out of 50 random init
top 5 out of 50 random init
2
10
100
300
N
Figure 7: Lower bound on (F)) where is L2-MHA, with D 1 and varying N, obtained by optimising (X)]l with respect to X, with 50 random initialisations of X for each N.
In Figure show the lower bound Lip, (F) obtained by optimising (X) ll2 using the same optimisation procedure as for Figure of Section 5.1 Here the optimisation more difficult, evident in the variance of the top values, and the trend less clear, but it appears that Tipp grows at rate of (log N). The message is less clear here, and there are at least two possibilities:
(1) The optimisation is difficult even for small values of N, hence Figure 7 shows loose lower bound.
(2) If the lower bound is tight, this suggests that the O(vN asymptotically tight, and could be improved to O(log N) (or Ollog N log log N) as for
NlVNlogg log N) bound in Theorem 3.2 is not
p oo).
L OPTIMISING THE NORM OF THE JACOBIAN DP-MHA
In Figure we show how the norm of the Jacobian [X)] for DP MHA keeps increasing when being optimised with respect to X. This is useful sanity check validating theoretical result of Theorem 3.1 that DP-MHA is nor Lipshchitz. The oscillations are likely due to momentum term of Adam optimizer that was used to optimise the norm.
Optimise ||ЛX]|| wrt for
DP-MHA (trained, D=512)
25000
w
20000
LU-XOII
15000
10000
50 100 150 200 250 300 optim iterations
Figure 8: Optimise
(X) w.r.t. X for trained MH-MHA f.
22

---

Under review as a conference paper at ICLR 2021
M EXPERIMENT TYING KEYS AND QUERIES OF L2-MHA BUT PRESERVING PARAMETER COUNT
In Figure 4 of Section 5.3 we have shown that there clear reduction in performance when tying the keys and queries. To test whether this can be attributed to the reduction in parameter count, we tried doubling the number of columns of W when the keys and queries are shared (i.e. from D H to 2D H1) so that the shared model has the same number of parameters as the unshared model. In Figure | the third column shows results for shared L2 L2-MHA, but with the same number parameters as the unshared L2-MHA i.e. without tying the keys and queries. The performance is similar to the second column (tying with a reduced number of parameters), suggesting that there is an inherent limitation in expressiveness tying the keys and queries, and that the reduction in number of parameters is an insufficient explanation this phenomenon.
Transformer (L2), W W*, same parameter count
Transformer (L2)
Transformer (L2), W = WK
Num Layers
Num Layers
Num Layers
1.8
1
1
2
2
1.6
3
3
Valid NNL
5
5
1.4
1.2
1.0
20
40
60
80
100
20
40
60
80
100
20
40
60
80
100
epoch
epoch
epoch
Figure 9: Experiment tying keys/queries but preserving parameter count.
N STABILITY EXPERIMENTS
In Figure below, we compare the output variance of trained L2-MHA against trained DP-MHA, with weights from the one-layer Transformer (L2). Wº W model and (DP) model used for Figure respectively. We take the same distribution of inputs as used for the numerical invertibility experiment Section 5.2] and show the histogram of inputs and outputs after flattening the input/output tensors. 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 these inputs.
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 [] In Figure [T] we show how training instabilities arise for DP-MHA with deeper Transformer models if we use fixed learning rate. We compare the training curves of DP-MHA, L2 MHA (WQ = WK) and Cont ract ive-tiieeeee -iree----eee -HH for a varying number of layers for the first 20 epochs of training with fixed learning rate. We see that DP-MHA fails to train beyond 10 layers, whereas both L2 -HA (W? W (i.c. Lipschitz L2-MHA but not contractive) and Contract. 122 shows stable training for up to 18 layers. This was the deepest model we could fit on a single GPU, and we expect to be able to train even deeper models with L2 MHA (WQQ = k and Contract i ive iie- L2-MHA MHA. In Table we show the best Test NLL across training of Transformer models with fixed learning rate. Note that for OP-MHA training becomes unstable beyond 10 layers, so we are only able to provide results up to layers. The generalisation performance of the best model for each setting is similar.
23

---

Under review as a conference paper at ICLR 2021
Inputs
Outputs: L2-MHA (trained)
Outputs: DP-MHA (trained)
200
Figure 10: Histogram showing distribution of inputs/outputs of L2-MHA and DP-MHA
Transformer (DP)
Transformer (L2, W*= Wo)
Transformer (Contractive-L2)
Num Layers
2
4
6
Train NLL
8
12
16
20
10
15
20
epoch
epoch
epoch
Figure 11: Train NLL for Transformer (DP), Transformer (L.2) and Transformer (Contractive-I.2)
Number of Layers
2
4
10
12
14
16
18
(DP)
1.061
1.032
1.021
1.017
1.025
W = W (L2)
1.168
1.040
1.023
1.024
1.019
1.008
1.018
1.027
1.034
(Contractive-L2)
1.246
1.135
1.103
1.079
1.072
1.060
1.039
1.029
1.031
Table Test NLL for Transformer models trained with fixed learning rate on PTB character level language modelling
24


Reviews:
[{"review_id": "qbbHAO79m2", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies the Lipschitz continuity properties of self-attention. It is proved that the widely-used dot-product self-attention is not Lipschitz continuous. A novel L2 self-attention is proposed and proven to be Lipschitz continuous. Experiments show that the upper bound of Lipschitz constant for L2 self-attention is asymptotically tight. Invertibility of MHA residual map is investigated to prove the Lipschitz continuity of L2 self-attention. Finally, experiments on Transformers with L2 self-attention are studied.\n\nThe problem studied in this paper is a novel problem which is unexplored and interesting. The paper is equipped with solid mathematical analysis and experimental results. It is well written and easy to follow.\n\nMy main concern about this paper is --- why do we need Lipschitz continuity in self-attention. In the introduction, the authors mention some contexts where Lipschitz continuity properties of neural networks are needed. But for self-attention, what are the benefits of Lipschitz continuity, or are there any potential scenarios where non Lipschitz continuous self-attention models do not work but only Lipschitz continuous self-attention models do work. In the experiments of Transformer, the L2 self-attention model behaves just comparable (or sometimes worse) than DP self-attention, which is not a strong evidence of the necessity of Lipschitz continuity of self-attention models.\n\nMinor questions:\n1. The non Lipschitz continuity of DP attention models seem to come from the \"quadratic\" terms of X. L2 self-attention looks like Gaussian kernel in its form. I wonder is there any general assumptions (like using bounded kernel?) under which self-attention models are Lipschitz. This will probably generalize the mathematical content of this paper.\n\n2. In Figure 4, could you provide the exact values of the valid NLL for each panel. It is difficult to distinguish the performance, especially the performance between panel 2 and panel 3, from the figures only.", "extra_scores": {}, "extra_text": {}}, {"review_id": "fgYXwuF9KEp", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The presented manuscript studies the Lipschitz constants of self-attention networks. It proofs that the widely used dot-product self-attention is not globally Lipschitz. The authors propose a so called L2 self-attention and derive its Lipschitz constant. An application in a transformer-based architecture shows only a minor performance reduction as price for the Lipschitz property.\n\nThe literature is properly reviewed and the contributions are clearly stated. The authors explain all concepts required to understand the theoretical derivation. However, I suggest that the difference between local and global Lipschitz continuity are introduced in detail as this becomes particularly important. \nFurthermore, the proof of Theorem 3.1. was difficult to follow for me and I think there should be a simpler way to formulate it. A few suggestions: i)  first give a high level overview of the steps of the proof ii) use lemmas to show intermediate results iii) proof it for N=1, D=1 first, which will simplify the notation\nOne point, which is shortly discussed by the authors, is that the dot-product self-attention is local Lipschitz. For most applications, the inputs are naturally bounded e.g. images. i.e. also the DP-MHA is Lipschitz.  So it unclear for me whether the L2-MHA will be need in these cases and maybe an experimental comparison would be convincing here.\nThe experimental section evaluates the proposed approach from three quite distinct directions. While 5.1. relies quite heavily on empirical evidence to support a theoretical claim, the example 5.2. and 5.3. get a clear message across.\nThe supplementary material is excessive, but was not further considered for this review.\n\nIn general, I think the theoretical analysis of deep learning methods is important to advance the field and I am therefore in favor of this paper. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "E8QK4yERZ2R", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper shows that the widely used dot-product self-attention is not Lipschitz with respect to any p-norm with $p \\ge 1$, and formulates L2 self-attention, which is Lipschitz.\n\nThe theorems and the proofs in the paper make sense to me, and the authors use experiments to show the tightness of Lipschitz constant derived from their theorem, the invertability using L2 self-attention, and the expressiveness of L2 self-attention on a NLP task compared with the transformer using the dot-product self-attention.\n\nI think this paper is interesting, given that there is almost no theoretical paper (compared with fully connected NN and convolutional neural network) studying the attention mechanism of neural network.\n\nHowever, my main concern is the “significance” of Lipschitz property. From a theoretical side, although there are many theoretical guarantees using the Lipschitz constant of the neural network or the general classification function, it may not lead to the best bounds. Actually, even if the original dot-product self-attention is not Lipschitz, I guess there is also some probability that most of the time (or along the trajectory) the self-attention is Lipschitz. This similar phenomenon happens in the Neural Tangent Kernel regime: although wide neural networks with ReLU have local minima and is not smooth everywhere, the function is easy to optimize along the trajectory.\n\nFrom the experimental/application side, I am not so familiar with the invertible NN literature, and I am not so convinced that invertible NN can directly lead to the applications mentioned in the paper. I think the most straight-forward way (for me and general audience) to show the advantage L2 self-attention is to find some application, which may be easy and the experiments can be very preliminary, and show that L2 self-attention works better on these tasks in the preliminary experiments.\n\nIn the light of my previous concern, I would like to ask a few questions:\n(1)\tDo you find any cases through experiments that the dot-product self-attention does not perform well but L2 self-attention can perform better? E.g. with the L2 self-attention mechanism, the training of GAN is more stable than that using the original self-attention, or the adversarial training using L2 self-attention performs better than the original self-attention?\n\n(2)\tIs it very natural to extend the invertible NN to the tasks mentioned in the introduction, i.e. flow based generative models?\n\n(3)\tAre there any references claiming some lower bounds of the performance in terms of the Lipschitz constant? If yes, I suggest to add those references into a related work section, or the introduction, to make the contribution of this paper clear.\n\n(4)\tIn the experiment, you mention that training the original transformer needs lots of tuning and the performance is not stable, while training transformer using L2 self-attention is much more stable. Could you please show more results on this phenomenon, because I think this phenomenon is interesting and can help to clarify question (1) I ask? E.g. the results of original transformer using different parameters with different initializations and the results of L2 transformer by choosing different parameters?\n\n(5)\tA suggestion: I think the proof of Theorem 3.1 can be put in the appendix and the conclusion can be more concise, and the main content can contain more experiment results and related works.\n\n\nI will adjust my score after some of my questions are answered.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "CrJvQnKP9l", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \nThis paper theoretically analyzes the Lipschitz constant of the self-attention module. In particular, the authors prove that the vanilla dot-product self-attention is not Lipschitz and propose a Lipschitz L2 self-attention whose Lipschitz constant is upper-bounded. The theoretical results and the asymptotic tightness of the derived bound are empirically verified by analytical experiments.\n\nReason for score: \nI vote for an accept. I really like the theoretical analysis for the Lipschitz continuity of self-attention. And the proposed Lipschitz L2 self-attention has several potential practical applications. My major concern is whether the advantage of the Lipschitz L2 self-attention could outweigh the disadvantages in practice (see cons below). Overall, I think the theoretical contributions of this paper are sufficient for an accept.\n\nPros: \n+ Studying the Lipschitz constant of self-attention is very interesting and motivated, since Lipschitz continuity is a desirable property for neural networks but most previous works only restrict to simple fully connected or convolution networks.\n+ The proposed L2 self-attention is a simple plug-in variant of the standard dot-product self-attention and could potentially be applicable to several Lipschitz-required situations, e.g., the analysis of robustness and the design of invertible models.\n+ Overall, the paper is clearly written and the derivations that I have checked are correct.\n\nCons: \n- From the experiments, the expressivity of the proposed L2 self-attention is a big concern. It is not clear that the advantages of the Lipschitz L2 self-attention would outweigh its disadvantages. Therefore, it would be great to do more extensive experiments to indeed verify the superiority of Lipschitz continuity, e.g., the improved robustness and the training stability or the applicability to invertible models.\n- In the experiments regarding the asymptotic tightness of the upper bound, only the setting where $H=D=1$ is shown. It would be more convincing to include experiments for larger dimensions as well.\n\nAdditional Comments/Questions: \n1. Another potential application of the Lipschitz/Contractive L2 self-attention is DEQ[1] to ensure the uniqueness and the stability of fixed points. It would be great to include some discussions (if it is correct).\n2. In Lemma E.1, is the equation between $\\tilde{J}_{ij}$ and $\\tilde{K}_{ij}$ correct? \n\n[1] https://arxiv.org/abs/1909.01377\n\n##########Post-Rebuttal Feedback########\nI appreciate the author's thorough response and I think the additional experiments on robustness/stability make the paper stronger, so I decide to raise my score to 7. For the future version of the paper, it would be great to see more comprehensive experiment results that show the improved robustness/stability in the main text.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper shows that L2 self-attention is Lipschitz and presents a new method for computing the Lipschitz constant. All reviewers are positive about the technical part of the paper. However, the major concern comes from the significance of the computed Lipschitz constant. The paper only presents some numerical results using simple toy examples, which is insufficient to justify the  importance of the proposed method. The paper would be a much stronger paper if better numerical results could be presented.
Decision: Reject