The Lipschitz Constant of Self-Attention

Hyunjik Kim, George Papamakarios, Andriy Mnih

Introduction

In deep learning, we often use Lipschitz continuity as a constraint for neural networks, to control how much a network’s output can change relative to its input. Such Lipschitz constraints are useful in several contexts. For example, Lipschitz constraints can endow models with provable robustness against adversarial pertubations (Cisse et al., 2017; Tsuzuku et al., 2018; Anil et al., 2019), and guaranteed generalisation bounds (Sokolić et al., 2017). Moreover, the dual form of the Wasserstein distance is defined as a supremum over Lipschitz functions with a given Lipschitz constant, hence Lipschitz-constrained networks are used for estimating Wasserstein distances (Peyré & Cuturi, 2019). Further, Lipschitz-constrained networks can stabilise training for GANs, an example being spectral normalisation (Miyato et al., 2018). Finally, Lipschitz-constrained networks are also used to construct invertible models and normalising flows. For example, Lipschitz-constrained networks can be used as a building block for invertible residual networks and hence flow-based generative models (Behrmann et al., 2019; Chen et al., 2019). Additionally, Neural ODEs (Chen et al., 2018; Grathwohl et al., 2019) are typically defined using vector fields parameterized via Lipschitz networks, so that the flow generated by the vector field is guaranteed to exist for all times.

Nonetheless, designing Lipschitz-continuous neural networks and computing (or even upper-bounding) their Lipschitz constant is a hard problem. Previous work mostly focused on fully-connected 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 non-linearities. Even in this case however, exact evaluation of the Lipschitz constant of fully-connected and convolutional networks is NP-hard (Virmaux & Scaman, 2018) and obtaining a tight upper bound remains a challenging task (Virmaux & Scaman, 2018; Fazlyab et al., 2019; Latorre et al., 2020).

Fully-connected and convolutional networks are not the only neural networks worthy of interest. Recently, self-attention (Vaswani et al., 2017) has become a popular alternative to recurrent neural networks. Self-attention is a key component of the Transformer (Vaswani et al., 2017), that has found success as a building block in models of various data modalities, starting with natural-language processing (Vaswani et al., 2017; Devlin et al., 2019; Brown et al., 2020) and extending to computer vision (Zhang et al., 2019; Parmar et al., 2019), audio generation (Huang et al., 2019), and reinforcement learning (Parisotto et al., 2020). However, so far no previous work has analysed the Lipschitz properties of self-attention, and thus it has been unclear whether self-attention is a viable option in applications that require Lipschitz constraints. In this work, we address this gap in the theory of self-attention by providing a thorough analysis of its Lipschitz properties. In particular, we make the following contributions:

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

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

We derive a theoretical upper bound on the Lipschitz constant of L2 self-attention, and provide empirical evidence of the asymptotic tightness of the bound.

As a practical demonstration of the theory, we use this bound to formulate invertible self-attention, and explore its use in a Transformer architecture for character-level language modelling. We compare its test log-likelihood and stability to dot-product self-attention.

Lipschitz Constant of Fully-Connected/Convolutional Layers

We first define the notion of Lipschitz continuity, and proceed to define the Lipschitz constant.

Given two metric spaces (X,dX)(\mathcal{X},d_{\mathcal{X}}) and (Y,dY)(\mathcal{Y},d_{\mathcal{Y}}), a function f:X→Yf:\mathcal{X}\rightarrow\mathcal{Y} is called Lipschitz continuous (or KK-Lipschitz) if there exists a constant K≥0K\geq 0 such that

The smallest such KK is the Lipschitz constant of ff, denoted Lip⁡(f)\operatorname{Lip}(f).

Next, we outline some basic results that are useful for estimating Lipschitz constants, also covered in related works (Virmaux & Scaman, 2018; Behrmann et al., 2019). We describe how these results are used to provide bounds on the Lipschitz constant of fully-connected networks (FCN) and convolutional neural networks (CNN), using the fact that both are compositions of linear maps and pointwise non-linearities. To begin with, the following theorem suggests a way to bound Lip⁡p(f)\operatorname{Lip}_{p}(f) for a differentiable Lipschitz function ff:

Hence if ff is a linear map represented by a matrix WW then

where ∥W∥p\|W\|_{p} is the operator norm on matrices induced by the vector pp-norm, and σmax⁡(W)\sigma_{\max}(W) is the largest singular value of WW. Under this choice of norm, many common non-linearities (including relu, sigmoid, tanh, elu) are 11-Lipschitz. ∥W∥2=σmax(W)\|W\|_{2}=\sigma_{\text{max}}(W) is usually estimated via power iteration; we provide details on how this is done in Appendix B.

Since we now know the Lipschitz constants of the components of both FCN and CNN, we can bound their Lipschitz constants by applying the following lemma:

Let g,hg,h be two composable Lipschitz functions. Then g∘hg\circ h is also Lipschitz with Lip⁡(g∘h)≤Lip⁡(g)Lip⁡(h)\operatorname{Lip}(g\circ h)\leq\operatorname{Lip}(g)\operatorname{Lip}(h).

For a fully-connected network (FCN) or a convolutional neural network (CNN) f=WK∘ρK−1∘WK−1∘…∘ρ1∘W1f=W_{K}\circ\rho_{K-1}\circ W_{K-1}\circ\ldots\circ\rho_{1}\circ W_{1}, we have Lip⁡p(f)≤∏k∥Wk∥p\operatorname{Lip}_{p}(f)\leq\prod_{k}\|W_{k}\|_{p} under a choice of pp-norm with 11-Lipschitz non-linearities ρk\rho_{k}.

The above bound is not necessarily tight; there are various works that compute tighter bounds for FCN and CNN (e.g. Virmaux & Scaman, 2018; Fazlyab et al., 2019; Latorre et al., 2020).

Lipschitz Constant of Self-Attention

In what follows, we will prove that MHA\mathit{MHA} as defined above is not Lipschitz, assuming that the MHA\mathit{MHA} map is non-trivial, i.e. WQ,WK,WV,WO≠0W^{Q},W^{K},W^{V},W^{O}\neq 0. It is sufficient to show that a single head DP\mathit{DP} is not Lipschitz, since MHA\mathit{MHA} is a linear combination of the outputs of each head. Also note that PP is a stochastic matrix, i.e. its entries are non-negative and its rows sum to 11. Since the rows of XX are the xi\mathbf{x}_{i}’s, a linear transformation of each xi\mathbf{x}_{i} by some matrix AA is equivalent to right multiplication of XX by A⊤A^{\top}. So right multiplication of XX by WVW^{V} is a linear map and thus Lipschitz. Therefore, we are interested in the mapping f(X)=PXf(X)=PX; this is not a linear mapping because PP itself is a non-linear function of XX. In fact, we show that ff is not Lipschitz, thus proving the first main result of the paper:

DP-MHA is not Lipschitz for any vector pp-norm ∥⋅∥p\|\cdot\|_{p} with p∈[1,∞]p\in[1,\infty].

Summary of Proof. We use Theorem 2.1, noting that if the supremum of the norm of the Jacobian is infinite, then the mapping is not Lipschitz. In particular, we show that when xi=0\mathbf{x}_{i}=\mathbf{0} for some ii, some elements of the Jacobian of ff grow proportionally to the sample variance of x≠i\mathbf{x}_{\neq i}, which is unbounded.

δij∈{0,1}\delta_{ij}\in\{0,1\} is the Kronecker delta

See Appendix A for useful identities in deriving the above Jacobian.

Let us investigate the scalar X⊤P(i)XX^{\top}P^{(i)}X. We observe that it is in fact a variance of a discrete distribution. Specifically:

We use this observation to show that JiiJ_{ii} is unbounded, and so ∥Jf∥p\|J_{f}\|_{p} is unbounded, hence DP-MHA is not Lipschitz. Consider the case xi=0x_{i}=0. Then

High-level intuition for proof. At xi=0x_{i}=0, fi(X)=1N∑kxkf_{i}(X)=\frac{1}{N}\sum_{k}x_{k}, the mean of the inputs. The rate of change of fif_{i} is governed by how fast the softmax saturates when xix_{i} is perturbed, which is determined by how spread out the x≠ix_{\neq i} are. The more spread out they are (the higher the sample variance), the greater the rate of saturation of the softmax, and the faster the rate of change of fif_{i}. Since the sample variance of x≠ix_{\neq i} can be arbitrarily large, the rate of change of fif_{i} can also be arbitrarily large, i.e. the entries of the Jacobian (and hence its pp-norm) can become arbitrarily large. In Appendix D, we show that adding bias terms to xi⊤WQ\mathbf{x}_{i}^{\top}W^{Q} and xj⊤WK\mathbf{x}_{j}^{\top}W^{K} does not resolve the issue.

The implications of this result are the following. (1) There can be undesirable behaviour (e.g. training instabilities) for the Transformer when some inputs are close to zero and others have large magnitude. (2) Dot-product self-attention (and hence the standard Transformer) is not a suitable choice when we require a Lipschitz neural network, such as for formulating invertible residual networks (Behrmann et al., 2019). Therefore, to use self-attention and Transformers in such applications, a Lipschitz formulation of self-attention is required, together with an explicit (ideally tight) upper bound to its Lipschitz constant, to quantify how much the output can change with respect to changes in the input.

2 L2 self-attention: a Lipschitz formulation of self-attention

The pathology in dot-product self-attention arises because the softmax probabilities Pi:P_{i:} are constant with respect to x≠i\mathbf{x}_{\neq i} when xi=0\mathbf{x}_{i}=0. This behaviour can be undesirable as we want PijP_{ij} to vary according to xj\mathbf{x}_{j}, regardless of whether xi\mathbf{x}_{i} is zero or not. Hence we propose an alternative form of self-attention based on L2 distance:

with the normalisation constant ensuring that ∑jPij=1\sum_{j}P_{ij}=1. We will refer to it as L2 self-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 11. Normalisation is usually necessary to deal with inputs of varying length NN (Wang et al., 2018), hence we keep the softmax for L2 self-attention. Similarly to dot-product self-attention, L2 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 require WQ,h=WK,hW^{Q,h}=W^{K,h} for each head fh(X)f^{h}(X) to be Lipschitz. In Lemma F.1 of Appendix F we show that L2-MHA is not Lipschitz for arbitrary WQ,hW^{Q,h}, WK,hW^{K,h}, and that tying WQ,h=WK,hW^{Q,h}=W^{K,h} is sufficient for L2-MHA to be Lipschitz, with intuition for why tying is sufficient.

In each head of the self-attention fh(X)f^{h}(X), right multiplication by AhA_{h} 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 FF, each fh(X)f^{h}(X) is additionally transformed by WV,hW^{V,h}, a free parameter.

The second main result of the paper is the following:

L2-MHA is Lipschitz, with the following bound on Lip⁡∞(F)\operatorname{Lip}_{\infty}(F):

and the following bound on Lip⁡2(F)\operatorname{Lip}_{2}(F):

where ϕ(x)≔xexp⁡(x+1)\phi(x)\coloneqq x\exp(x+1) is an invertible univariate function on x>0x>0, and NN is the input sequence length.

Specifically, ϕ−1(N−1)=W0(Ne)\phi^{-1}(N-1)=W_{0}(\frac{N}{e}) where W0W_{0} is the Lambert WW-function, which grows sub-logarithmically as O(log⁡N−log⁡log⁡N)O(\log N-\log\log N) (Corless et al., 1996). Hence the above bounds can be simplified to O(log⁡N)O(\log N) for p=∞p=\infty and O(Nlog⁡N)O(\sqrt{N}\log N) for p=2p=2.

See Appendix F, which uses the key observation that X⊤P(i)XX^{\top}P^{(i)}X is a covariance matrix (c.f. Equation (7)) to bound ∥JF∥p\|J_{F}\|_{p}, the norm of the Jacobian of FF. 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(Dlog⁡N)O(\sqrt{D\log N}) bound on Lip⁡1(F)\operatorname{Lip}_{1}(F) using measure-theoretic tools.

Application: Invertible Self-Attention

Consider the residual function g(x)≔x+f(x)g(x)\coloneqq\mathbf{x}+f(\mathbf{x}). Behrmann et al. (2019) give the following sufficient condition for its invertibility: if ff is a contraction with respect to some metric, i.e. if Lip⁡(f)<1\operatorname{Lip}(f)<1, and the metric space on which ff is defined is complete, then gg is invertible. (A Euclidean space with a metric induced by a pp-norm ∥⋅∥p\|\cdot\|_{p} for p∈[1,∞]p\in[1,\infty] is always complete.) Specifically, the inverse g−1(y)g^{-1}(\mathbf{y}) is the unique fixed point of the recursion xi+1≔y−f(xi)\mathbf{x}^{i+1}\coloneqq\mathbf{y}-f(\mathbf{x}^{i}), since by the definition of the inverse we have y=g−1(y)+f(g−1(y))\mathbf{y}=g^{-1}(\mathbf{y})+f(g^{-1}(\mathbf{y})). Because ff is a contraction, Banach’s Fixed Point Theorem guarantees that this fixed point exists and is unique for all y\mathbf{y}, and that the recursion converges for all initial values x0\mathbf{x}^{0} (often set to y\mathbf{y} in practice) exponentially fast. Hence the inverse can be computed to arbitrary accuracy (up to numerical precision in practice) by the above fixed-point iteration.

Note that a composition of such invertible residual blocks is also invertible. Behrmann et al. (2019) use this observation to design invertible ResNets: they take ff to be a CNN normalised by an upper bound on Lip⁡(f)\operatorname{Lip}(f) given by Corollary 2.1, making the resulting function contractive. For the 22-norm ∥⋅∥2\|\cdot\|_{2}, a hyperparameter c<1c<1 is chosen and each linear map (convolution) WW in the CNN is multiplied by c/∥W∥2c/\|W\|_{2} if c<∥W∥2c<\|W\|_{2} where ∥W∥2\|W\|_{2} is estimated by power iteration (c.f. Appendix B). This multiplicative factor determines the scale of the Lipschitz constant of the normalised function.

2 Invertible self-attention

The standard use case of self-attention is with a skip connection inside the Transformer. A Transformer block is composed of residual blocks of multihead self-attention (MHA) and fully-connected (FCN) layers (Figure 1). Hence similarly to invertible ResNets, we can normalise L2-MHA by the upper bounds given in Theorem 3.2 to obtain Contractive-L2-MHA ff, with which we can obtain invertible self-attention g(x)=x+f(x)g(\mathbf{x})=\mathbf{x}+f(\mathbf{x}). Since Dropout is also part of the residual branch along with Contractive-L2-MHA, we should check that it is also contractive. At test time, Dropout multiplies inputs by the dropout keep probability p<1p<1, so it is a contraction with Lipschitz constant pp at 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 a diagonal binary matrix MM, and for such matrices we can verify ∥M∥p≔sup⁡∥x∥p=1∥Mx∥p≤1\|M\|_{p}\coloneqq\sup_{\|x\|_{p}=1}\|Mx\|_{p}\leq 1. Notice that LayerNorm is not part of the residual branch, hence its Lipschitz continuity is not relevant for invertibility; rather, we can replace it with an invertible normalisation such as ActNorm (Kingma & Dhariwal, 2018). However, architectures that place LayerNorm inside the residual branch (termed pre-LN as opposed to the traditional post-LN in Figure 1) have become more prevalent in the literature (Wang et al., 2019; Xiong et al., 2020), and in this case it makes sense to investigate its Lipschitz continuity. We show that LayerNorm is Lipschitz in Appendix N, with a bound on its Lipschitz constant.

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 Contractive-L2-MHA, 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 a testbed for this purpose, since self-attention is commonly used in a Transformer. Given the theoretical focus of the paper, we believe that a 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.

Experimental Results

We use this observation to provide empirical evidence for the asymptotic tightness of the upper bound on Lip⁡∞(f)\operatorname{Lip}_{\infty}(f) in Theorem 3.2. In Figure 2, we show the upper bound as well as the lower bound on Lip⁡∞(f)\operatorname{Lip}_{\infty}(f) obtained by optimising ∥Jf(X)∥∞\|J_{f}(X)\|_{\infty} with respect to XX for L2-MHA ff with 50 different random initialisations of XX, with H=D=1H=D=1 and NN varying between 100100 and 10001000. See Appendix H for further details. Note that we use a log-scale for the x-axis, and recall that the upper bound is O(log⁡N−log⁡log⁡N)O(\log N-\log\log N), dominated by the O(log⁡N)O(\log N) term for large NN. 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)O(\log N-\log\log N) upper bound is asymptotically tight.

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 a global optimum across inputs, which can be difficult to attain especially for high values of NN. (2) The multiplicative constant of the upper bound may be loose. Assuming asymptotic tightness, it remains an open question whether the multiplicative constant can be tightened. We show the analogous plot for Lip⁡2(F)\operatorname{Lip}_{2}(F) and discuss the results in Appendix J. Additionally in Appendix K, we show that optimising ∥Jf(X)∥∞\|J_{f}(X)\|_{\infty} w.r.t. XX for DP-MHA ff causes the norm to diverge, providing empirical verification of Theorem 3.1, that DP-MHA is indeed not Lipschitz.

2 Numerical invertibility of MHA residual map

Recall from Section 4.1 that g(x)=x+f(x)g(\mathbf{x})=\mathbf{x}+f(\mathbf{x}) is invertible if ff is contractive. Hence if ff is Contractive-L2-MHA, gg is necessarily invertible. However, technically we do not disprove the invertibility of DP-MHA, since the converse does not hold in general i.e. if ff is DP-MHA, which we have shown is not Lipschitz hence not contractive, it may still be the case that gg is invertible. 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)g(\mathbf{x})=\mathbf{x}+cf(\mathbf{x}) between the cases where ff is L2-MHA and DP-MHA in Figure 3. For each, we take MHA with 88 heads and randomly initialised weights, and quantify the maximum reconstruction error across a batch of 128128 inputs whose outputs are inverted via the fixed-point iteration described in Section 4.1. We use N=64N=64, D=64D=64, and c∈{0.5,0.7,0.9}c\in\{0.5,0.7,0.9\} (see Appendix I for analogous results for a wider range of NN and DD 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-MHA is not Lipschitz) that when one of the inputs xi\mathbf{x}_{i} is , some terms of the Jacobian grow with the sample variance of x≠i\mathbf{x}_{\neq i}. Hence we check numerical invertibility at a set of NN inputs where xi=0\mathbf{x}_{i}=0 and x≠i\mathbf{x}_{\neq i} are chosen uniformly at random.

In Figure 3, we see that DP-MHA is not invertible whereas L2-MHA is invertible for sufficiently small cc. This shows how not having the theoretical guarantee of ff being contractive can cost us invertibility in practice. We note that the figure 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. Experiments with the globally invertible self-attention obtained by normalising with the Lipschitz upper bound are provided in the next section.

3 Expressiveness of L2-MHA and invertible self-attention

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? We expect that the Lipschitz constraint will limit the expressiveness of the Transformer, and would like to 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 1) at character-level language modelling on the Penn Treebank dataset (Marcus et al., 1993). 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. For Contractive-L2-MHA, we normalise F=F=L2-MHA by the bound on Lip⁡∞(F)\operatorname{Lip}_{\infty}(F) as it is tighter than the bound on Lip⁡2(F)\operatorname{Lip}_{2}(F). During training we backpropagate through these contractive blocks F/Lip⁡∞(F)F/\operatorname{Lip}_{\infty}(F) (including the denominator) to update the model parameters. We found that only backpropagating through the numerator (i.e. applying stop-gradient to denominator) gave slightly worse performance. See Appendix H for experimental details.

The results are shown in Figure 4. The first plot shows the best performing LSTM reaching a test NLL of around 1.01.0, and the second plot shows the best performing Transformer reaching a slightly improved performance for 33–55 layers of Transformer blocks. 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). The third plot shows results for the Transformer with DP-MHA replaced with L2-MHA but without tying WQW^{Q} and WKW^{K}, and we observe a very similar test performance. The fourth plot shows the change when we further tie the query and key weights (making WQ=WKW^{Q}=W^{K}); we see that there is a small degradation in performance. Here the number of trainable parameters has been reduced, but in Appendix L 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 55 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)\operatorname{Lip}_{\infty}(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 55 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)\operatorname{Lip}_{\infty}(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.

4 Training Stability of DP-MHA vs L2-MHA

In Figure 5, we compare the output variance of trained L2-MHA against trained DP-MHA, with weights from the one-layer Transformer (L2), WQ=WKW^{Q}=W^{K} model and (DP) model used for Figure 4 respectively. 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 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 at these inputs.

In practice, this leads to instabilities in 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 H. We investigate the behaviour of the different Transformer models on the above PTB task when using a fixed learning rate. We observe that DP-MHA fails to train at all beyond 10 layers, whereas both L2-MHA (WQ=WKW^{Q}=W^{K}) (i.e. Lipschitz L2-MHA but not contractive) and Contractive-L2-MHA shows stable training for up to 18 layers (see Appendix M for the training curves). This was the deepest model we could fit on a single GPU, and we expect to be able to train even deeper models with these two. In Table 1 we show the best Test NLL across training for each of the Transformer models. Note that for DP-MHA training becomes unstable beyond 10 layers, so we are only able to provide results up to 10 layers. The generalisation performance of the best model for each setting of self-attention is similar.

Conclusion and Discussion

We have shown that the widely used dot-product self-attention is not Lipschitz, and that the proposed L2 self-attention is Lipschitz, by deriving an O(log⁡N−log⁡log⁡N)O(\log N-\log\log N) Lipschitz bound for p=∞p=\infty and an O(N(log⁡N−log⁡log⁡N))O(\sqrt{N}(\log N-\log\log N)) bound for p=2p=2, where NN is the input sequence length. We also provided empirical evidence of the asymptotic tightness of the bound for p=∞p=\infty. 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. And finally, we also showed that L2-MHA is more stable during training, allowing the use of fixed learning rate for stable training of deep architectures.

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. N×D^{N\times D}, 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 a 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 a 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 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. 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. Despite these difficulties, this could be a worthwhile alternative route for obtaining Lipschitz self-attention to 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 guarantees. 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 a 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 on kernels other than 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.

Acknowledgements

We would like to thank Adam Kosiorek, Arnaud Doucet, Yee Whye Teh, Michalis Titsias, Emilien Dupont and Theophane Weber for helpful discussion and feedback.

References

Appendix A Useful Identities for deriving Jacobian expressions

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

∂∂x[λu]=λ∂u∂x+x∂λ∂x\frac{\partial}{\partial\mathbf{x}}[\lambda\mathbf{u}]=\lambda\frac{\partial\mathbf{u}}{\partial\mathbf{x}}+\mathbf{x}\frac{\partial\lambda}{\partial\mathbf{x}}

∂f(u)∂x=∂f(u)∂u∂u∂x\frac{\partial f(\mathbf{u})}{\partial\mathbf{x}}=\frac{\partial f(\mathbf{u})}{\partial\mathbf{u}}\frac{\partial\mathbf{u}}{\partial\mathbf{x}}

∂∂x[u⊤v]=u⊤∂v∂x+v⊤∂u∂x\frac{\partial}{\partial\mathbf{x}}[\mathbf{u}^{\top}\mathbf{v}]=\mathbf{u}^{\top}\frac{\partial\mathbf{v}}{\partial\mathbf{x}}+\mathbf{v}^{\top}\frac{\partial\mathbf{u}}{\partial\mathbf{x}}

Note ∂λ∂x\frac{\partial\lambda}{\partial\mathbf{x}} is a row vector, so u∂λ∂x\mathbf{u}\frac{\partial\lambda}{\partial\mathbf{x}} is a matrix.

Appendix B Power Iteration

The iteration will converge if W⊤WW^{\top}W has an eigenvalue that is strictly greater in magnitude than its other eigenvalues, and the starting vector b0b_{0} has a nonzero component in the direction of an eigenvector associated with the dominant eigenvalue. This happens with probability 11 if b0b_{0} is chosen at random, and the convergence is geometric with ratio ∣λ2/λmax⁡∣|\lambda_{2}/\lambda_{\max}| where λ2\lambda_{2} is the eigenvalue with second largest magnitude (Mises & Pollaczek-Geiringer, 1929).

Appendix C Proof of Theorem 3.1 for General D𝐷D

DP-MHA is not Lipschitz for any vector pp-norm ∥⋅∥p\|\cdot\|_{p} with p∈[1,∞]p\in[1,\infty].

For the last equality, note EiiXE_{ii}X has all rows equal to zero except for the iith row given by xi⊤\mathbf{x}_{i}^{\top}. We can then verify that X⊤P(i)EiiXX^{\top}P^{(i)}E_{ii}X simplifies to Pii(xi−∑kPikxk)xi⊤P_{ii}(\mathbf{x}_{i}-\sum_{k}P_{ik}\mathbf{x}_{k})\mathbf{x}_{i}^{\top}.

For vector pp-norms, ∥Jf∥p\|J_{f}\|_{p} is bounded if and only if its entries are bounded, by definition of the operator norm. The entries of X⊤P(i)XAX^{\top}P^{(i)}XA are bounded for arbitrary AA only if the entries of X⊤P(i)XX^{\top}P^{(i)}X are bounded. So let us investigate the entries of this D×DD\times D matrix. Writing out each term of the matrix, we observe that it is in fact a covariance matrix of a discrete distribution. Specifically:

Note that we have shown that single head dot-product self-atttention (H=1H=1) is not Lipschitz, but it is clear that this implies multihead self-attention DP-MHA is also not Lipschitz, since the output of multihead attention is a linear combination of the outputs of each head. ∎

Appendix D Bias term in DP Self-Attention

Appendix E Efficient Computation of L2 Self-Attention

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 efficiently. Similarly, we can show that L2 self-attention can also be computed in an efficient manner. Using the identity ∥a−b∥22=∥a∥22−2a⊤b+∥b∥22\|a-b\|_{2}^{2}=\|a\|_{2}^{2}-2a^{\top}b+\|b\|_{2}^{2} 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-attention we can show that

In Table 2 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.

Appendix F Proof of Theorem 3.2

Also note P(i)P^{(i)} is symmetric, and each row/colum sums to , i.e. P(i)\mathds1=\mathds1⊤P(i)=0P^{(i)}\mathds{1}=\mathds{1}^{\top}P^{(i)}=0. Hence we may simplify the Jacobian terms as follows:

Let yj⊤=xi⊤WQ−xj⊤WK\mathbf{y}_{j}^{\top}=\mathbf{x}_{i}^{\top}W^{Q}-\mathbf{x}_{j}^{\top}W^{K}. Then we have:

The intuition for this result is as follows: a reason for DP-MHA not being Lipschitz is that for xi=0\mathbf{x}_{i}=0,, the attention weights PijP_{ij} become uniform regardless of the values of xj\mathbf{x}_{j} for j≠ij\neq i. A similar issue arises for L2-MHA with WQ≠WKW^{Q}\neq W^{K} and full-rank WKW^{K}, as shown above: given any xi\mathbf{x}_{i}, we can choose xj\mathbf{x}_{j} such that the PijP_{ij} become uniform.

Hence we impose the restriction that WK=WQW^{K}=W^{Q}. With this assumption we have

Noting Lip⁡p(f)=sup⁡X∥Jf(X)∥p\operatorname{Lip}_{p}(f)=\sup_{X}\|J_{f}(X)\|_{p}, we would like to upper bound ∥Jf∥p\|J_{f}\|_{p}.

Consider the choice p=∞p=\infty, where ∥Jf∥∞\|J_{f}\|_{\infty} is the maximum absolute row sum of JfJ_{f}. A key observation is that if we can bound the ∞\infty-norm of the Jacobian of fif_{i}, a single output of ff, (i.e. a single block row ∥[Ji1,...,JiN]∥∞\|[J_{i1},...,J_{iN}]\|_{\infty} of JfJ_{f}) then this is also a bound on ∥Jf∥∞\|J_{f}\|_{\infty} due to permutation equivariance of self-attention; all block rows have the same maximal ∥⋅∥∞\|\cdot\|_{\infty} when each is optimised over the input XX. Using this, we can prove that ∥Jf∥∞\|J_{f}\|_{\infty} admits an upper bound that is O(log⁡N−log⁡log⁡N)O(\log N-\log\log N). Below we state and prove lemmas that lead to the proof of this upper bound.

First we analyse the term A⊤X⊤P(i)XA\sqrt{A}^{\top}X^{\top}P^{(i)}X\sqrt{A}, that appears in the first term of JiiJ_{ii}. Note that for Y≔XAY\coloneqq X\sqrt{A}, so that the rows of YY are yi⊤≔xi⊤A\mathbf{y}_{i}^{\top}\coloneqq\mathbf{x}_{i}^{\top}\sqrt{A}, we have

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

where zj≔∥yj−yi∥2z_{j}\coloneqq\|\mathbf{y}_{j}-\mathbf{y}_{i}\|_{2} (hence zi=0z_{i}=0). Define:

First note that as zj→∞z_{j}\rightarrow\infty, exp⁡(−zj2)→0\exp(-z_{j}^{2})\rightarrow 0 exponentially fast, causing the product zj2exp⁡(−zj2)→0z_{j}^{2}\exp(-z_{j}^{2})\rightarrow 0. Hence we expect the above quantity to be bounded and attain its maximum.

Let h(zj)≔exp⁡(−zj2)h(z_{j})\coloneqq\exp(-z_{j}^{2}) for notational conciseness, and note h(zj)>0h(z_{j})>0. By taking partial derivatives with the chain rule, we have that for j≠ij\neq i

Hence the derivative is if and only if zj=0z_{j}=0 or (1−zj2)∑kh(zk)+∑kh(zk)zk2=0(1-z_{j}^{2})\sum_{k}h(z_{k})+\sum_{k}h(z_{k})z_{k}^{2}=0, the latter being equivalent to zj2=1+∑kh(zk)zk2∑kh(zk)=1+g(z)z_{j}^{2}=1+\frac{\sum_{k}h(z_{k})z_{k}^{2}}{\sum_{k}h(z_{k})}=1+g(\mathbf{z}). Hence at the maximum, the non-zero values among {zj}j=1N\{z_{j}\}_{j=1}^{N} must be equal to one another. It is clear now that the maximum value cc is attained when zj2=1+cz_{j}^{2}=1+c for j≠ij\neq i (and recall zi=0z_{i}=0). So h(zj)=exp⁡(−1−c)h(z_{j})=\exp(-1-c) for j≠ij\neq i. Substituting this into g(z)g(z), and rearranging, we obtain cexp⁡(c+1)=N−1c\exp(c+1)=N-1. Note ϕ(x)≔xexp⁡(x+1)\phi(x)\coloneqq x\exp(x+1) is increasing for x>0x>0 hence c=ϕ−1(N−1)c=\phi^{-1}(N-1). ∎

Note ϕ(log⁡N)=(log⁡N)exp⁡(log⁡N+1)≥Nlog⁡N≥N−1\phi(\log N)=(\log N)\exp(\log N+1)\geq N\log N\geq N-1 for N≥3N\geq 3. Since ϕ\phi is increasing, we have ϕ−1(N−1)≤log⁡(N)\phi^{-1}(N-1)\leq\log(N) for N≥3N\geq 3. In fact, it is known that ϕ−1(N−1)=O(log⁡N−log⁡log⁡N)\phi^{-1}(N-1)=O(\log N-\log\log N) (Corless et al., 1996).

Using the inequalities ∥BC∥≤∥B∥∥C∥\|BC\|\leq\|B\|\|C\|, ∥B+C∥≤∥B∥+∥C∥\|B+C\|\leq\|B\|+\|C\| and ∥[A1,…,AN]∥≤∑i∥Ai∥\|[A_{1},\ldots,A_{N}]\|\leq\sum_{i}\|A_{i}\|, we have:

For the first equality, note that ∑jPij=1\sum_{j}P_{ij}=1. For the second equality, note that the summand for j=ij=i is because the term yi−yj=0\mathbf{y}_{i}-\mathbf{y}_{j}=\mathbf{0}. Each of the terms in the brackets are bounded by the following lemmas:

∥Y⊤P(i)Y∥∞≤ϕ−1(N−1)D/H\|Y^{\top}P^{(i)}Y\|_{\infty}\leq\phi^{-1}(N-1)\sqrt{D/H} (ϕ\phi defined as in Lemma F.2).

∑j∥Pij(yj−∑kPikyk)(yi−yj)⊤∥∞≤ϕ−1(N−1)D/H\sum_{j}\|P_{ij}(\mathbf{y}_{j}-\sum_{k}P_{ik}\mathbf{y}_{k})(\mathbf{y}_{i}-\mathbf{y}_{j})^{\top}\|_{\infty}\leq\phi^{-1}(N-1)\sqrt{D/H}.

Note ∥uv⊤∥∞=∥u∥∞∥v∥1\|\mathbf{u}\mathbf{v}^{\top}\|_{\infty}=\|\mathbf{u}\|_{\infty}\|\mathbf{v}\|_{1} for real vectors u,v\mathbf{u},\mathbf{v}. Hence

where aj=Pij∥yj−∑kPikyk∥∞a_{j}=\sqrt{P_{ij}}\|\mathbf{y}_{j}-\sum_{k}P_{ik}\mathbf{y}_{k}\|_{\infty}, bj=Pij∥yi−yj∥1b_{j}=\sqrt{P_{ij}}\|\mathbf{y}_{i}-\mathbf{y}_{j}\|_{1}.

Note aj≤cj≔Pij∥yj−∑kPikyk∥2a_{j}\leq c_{j}\coloneqq\sqrt{P_{ij}}\|\mathbf{y}_{j}-\sum_{k}P_{ik}\mathbf{y}_{k}\|_{2} since ∥u∥∞≤∥u∥2\|\mathbf{u}\|_{\infty}\leq\|\mathbf{u}\|_{2} for vector u\mathbf{u}. Hence ∥a∥2≤∥c∥2\|\mathbf{a}\|_{2}\leq\|\mathbf{c}\|_{2}.

Putting the above lemmas altogether, with the observation sup⁡X∥Jf(X)∥∞=sup⁡X∥[Ji1(X),…,JiN(X)]∥∞\sup_{X}\|J_{f}(X)\|_{\infty}=\sup_{X}\|[J_{i1}(X),\ldots,J_{iN}(X)]\|_{\infty} by permutation invariance of ∥Jf∥∞\|J_{f}\|_{\infty} (since ff is permutation equivariant and ∥⋅∥∞\|\cdot\|_{\infty} is the maximum absolute row sum), we have

where the last inequality holds for N≥3N\geq 3.

The full multihead attention map that combines the heads fh(X)f^{h}(X) is:

Note the Jacobian JgJ_{g} is a block matrix whose rows are JfhJ_{f^{h}}, hence ∥Jg∥∞=max⁡h∥Jfh∥∞\|J_{g}\|_{\infty}=\max_{h}\|J_{f^{h}}\|_{\infty}, and similarly ∥WV⊤∥∞=max⁡h∥WV,h⊤∥∞\|W^{V^{\top}}\|_{\infty}=\max_{h}\|W^{{V,h}^{\top}}\|_{\infty}. Hence we have

Combining this with Inequality (33), we have:

Let A be a block matrix with block rows A1,…,ANA_{1},\ldots,A_{N}. Then ∥A∥2≤∑i∥Ai∥22\|A\|_{2}\leq\sqrt{\sum_{i}\|A_{i}\|_{2}^{2}}, and equality holds if and only if the first right singular vectors of the AiA_{i} align.

Note that equality holds if and only if the first right singular vectors of the AiA_{i} align. ∎

Hence a bound on the spectral norm of each block row of JfJ_{f} can give us an O(N)O(\sqrt{N}) bound on ∥Jf∥2\|J_{f}\|_{2}, which may be loose, and it remains an open question as to whether this bound can be tightened.

To bound the ∥⋅∥2\|\cdot\|_{2} norm of each row of JfJ_{f}, we use the following lemmas:

∥Y⊤P(i)Y∥2≤ϕ−1(N−1)\|Y^{\top}P^{(i)}Y\|_{2}\leq\phi^{-1}(N-1)

∑j∥Pij(yj−∑kPikyk)(yi−yj)⊤∥2≤ϕ−1(N−1)\sum_{j}\|P_{ij}(\mathbf{y}_{j}-\sum_{k}P_{ik}\mathbf{y}_{k})(\mathbf{y}_{i}-\mathbf{y}_{j})^{\top}\|_{2}\leq\phi^{-1}(N-1)

Directly use Cauchy–Schwartz on cc and dd in the proof of Lemma F.4. ∎

Again using the inequalities ∥BC∥≤∥B∥∥C∥\|BC\|\leq\|B\|\|C\|, ∥B+C∥≤∥B∥+∥C∥\|B+C\|\leq\|B\|+\|C\| and ∥[A1,…,AN]∥≤∑i∥Ai∥\|[A_{1},\ldots,A_{N}]\|\leq\sum_{i}\|A_{i}\|, with the additional equality ∥B⊤∥2=∥B∥2\|B^{\top}\|_{2}=\|B\|_{2}, we have the bound:

To obtain the final result for the full multihead self-attention FF, we need a final lemma:

Let A be a block matrix with block columns A1,…,ANA_{1},\ldots,A_{N}. Then ∥A∥2≤∑i∥Ai∥22\|A\|_{2}\leq\sqrt{\sum_{i}\|A_{i}\|_{2}^{2}}.

where we are using the substitution xi=λiei\mathbf{x}_{i}=\lambda_{i}\mathbf{e}_{i}, and the last inequality holds by e.g. Cauchy–Schwartz inequality on [λ1,…,λN][\lambda_{1},\ldots,\lambda_{N}] and [∥A1∥2,…,∥AN∥2][\|A_{1}\|_{2},\ldots,\|A_{N}\|_{2}]. ∎

Since ∥fh(X)WV,h∥2≤∥Jfh∥2∥WV,h∥2\|f^{h}(X)W^{V,h}\|_{2}\leq\|J_{f^{h}}\|_{2}\|W^{V,h}\|_{2}, by Lemma F.8 we have that

Combining this with Inequality (34), we have:

Appendix 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 a set of mask indices M⊂{1,…,N}×{1,…,N}\mathcal{M}\subset\{1,\ldots,N\}\times\{1,\ldots,N\}, the logits (i.e. the inputs to the softmax) are set to −∞-\infty at the mask indices. That is,

Masking implies fi(X)f_{i}(X) is not a function of xj\mathbf{x}_{j} for (i,j)∈M(i,j)\in\mathcal{M}, hence Jij=0J_{ij}=0 for (i,j)∈M(i,j)\in\mathcal{M}. Thus fi(X)f_{i}(X) is equal to the iith output for self-attention with inputs restricted to {xj:(i,j)∉M}\{\mathbf{x}_{j}:(i,j)\notin\mathcal{M}\}, the unmasked inputs with respect to the iith output. Hence JijJ_{ij} will no longer contribute to the bound on ∥[Ji1,…,JiN]∥\|[J_{i1},\ldots,J_{iN}]\|, and hence the bound for the unmasked case will continue to hold as long as (i,i)∈M(i,i)\in\mathcal{M} i.e. xi\mathbf{x}_{i} attends to itself (this is necessary for the proof of Lemma F.2 to hold). The bound can in fact be tightened by replacing NN with ∣{xj:(i,j)∉M}∣|\{\mathbf{x}_{j}:(i,j)\notin\mathcal{M}\}|, the number of unmasked inputs with respect to the iith output.

Appendix H Experimental Details

For the experiment in Section 5.1, showing the asymptotic tightness of the upper bound on Lip⁡∞(F)\operatorname{Lip}_{\infty}(F) where FF is L2-MHA, we fix all free parameters of FF (namely WQ,WVW^{Q},W^{V}) to be the identity, and only optimise the input XX. We use 5050 random initialisations of XX for each NN, where Xij∼\pazocalU[−c,c]X_{ij}\sim\pazocal{U}[-c,c] for c∼\pazocalUc\sim\pazocal{U} (we observed that having cc itself be random improves optimisation). We display the top 55 results for each value of NN after optimising each random initialisation till convergence using Adam (Kingma & Ba, 2015) with a learning rate of 0.10.1.

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 Penn Treebank dataset (Marcus et al., 1993).We use the standard training-validation-test split, and the dataset can be found at e.g. https://github.com/harvardnlp/TextFlow/tree/master/data/ptb. 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 288288. 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 6464. The LSTM models have the dimensionality of the hidden state equal to the dimensionality DD of the cell state (the usual default implementation). The Transformer models are trained with a varying number of blocks (number of layers) with H=8H=8 heads and D=512D=512, tuning hyperparameters for dropout rate in {0,0.1,0.2}\{0,0.1,0.2\} and base learning rate γ∈{0.2,0.4,0.6,0.8,1.0,1.5,2.0}\gamma\in\{0.2,0.4,0.6,0.8,1.0,1.5,2.0\} with number of warmup iterations w∈{1000,2000,4000,8000}w\in\{1000,2000,4000,8000\} for the standard custom learning rate schedule in Vaswani et al. (2017):

where ϵt\epsilon_{t} is the learning rate at training iteration tt. Hence the learning rate linearly increases from to (Dw)−1/2(Dw)^{-1/2} over ww iterations, then decays proportionally to t−1/2t^{-1/2}. We use Glorot Uniform initialisation (Glorot & Bengio, 2010) for all weights (U[−1din+dout,1din+dout]U\left[-\sqrt{\frac{1}{d_{in}+d_{out}}},\sqrt{\frac{1}{d_{in}+d_{out}}}\right]), except for weights in L2-MHA that are initialised from U[−sD,sD]U\left[-\frac{s}{\sqrt{D}},\frac{s}{\sqrt{D}}\right], and ss is a hyperparameter. For D=512D=512, we used s=124s=\frac{1}{2^{4}}. All experiments were done in Tensorflow 1.14 (Abadi et al., 2016) on single Nvidia Tesla V100 GPUs.

Appendix I Numerical Invertibility of MHA Residual Map

Following Section 5.2, Figure 6 confirms 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. Figure 7 shows additional results for different values of NN and DD.

In Figure 8, we show the lower bound on Lip⁡2(F)\operatorname{Lip}_{2}(F) obtained by optimising ∥JF(X)∥2\|J_{F}(X)\|_{2} using the same optimisation procedure as for Figure 2 of Section 5.1. Here the optimisation is more difficult, evident in the variance of the top 55 values, and the trend is less clear, but it appears that Lip⁡2(f)\operatorname{Lip}_{2}(f) grows at a rate of O(log⁡N)O(\log N). The message is less clear here, and there are at least two possibilities:

The optimisation is difficult even for small values of NN, hence Figure 8 shows a loose lower bound.

If the lower bound is tight, this suggests that the O(Nlog⁡N)O(\sqrt{N}\log N) bound in Theorem 3.2 is not asymptotically tight, and could be improved to O(log⁡N)O(\log N) (or O(log⁡N−log⁡log⁡N)O(\log N-\log\log N) as for p=∞p=\infty).

Appendix K Optimising the norm of the Jacobian of DP-MHA

In Figure 9, we show how the norm of the Jacobian ∥Jf(X)∥∞\|J_{f}(X)\|_{\infty} for DP-MHA ff keeps increasing when being optimised with respect to XX. This is a useful sanity check validating our theoretical result of Theorem 3.1, that DP-MHA is not Lipshchitz. The oscillations are likely due to momentum term of Adam optimizer that was used to optimise the norm.

Appendix L Experiment tying keys and queries of L2-MHA but preserving parameter count

In Figure 4 of Section 5.3, we have shown that there is a 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 WQW^{Q} when the keys and queries are shared (i.e. from D/HD/H to 2D/H2D/H) so that the shared model has the same number of parameters as the unshared model. In Figure 10, 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. 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 insufficient explanation this phenomenon.

Appendix M Training curves for fixed learning rate DP-MHA vs L2-MHA

Appendix N The Lipschitz constant of LayerNorm

In this section, we show that LayerNorm is Lipschitz, with a loose bound on its Lipschitz constant w.r.t. to the ∞\infty-norm. LayerNorm is defined as follows:

In the trivial case where xdx_{d} are all equal or when D=1D=1, x=μ\mathbf{x}=\mu hence LN(x)=βLN(\mathbf{x})=\boldsymbol{\beta}, so its Lipschitz constant is 0. Thus let us assume D>2D>2 and not all xdx_{d} are equal.

First let us compute the derivative of μ\mu and σ2\sigma^{2} w.r.t xx:

Now the derivative of LN(x)d\text{LN}(\mathbf{x})_{d}, the ddth element of LN(x)\text{LN}(\mathbf{x}), w.r.t.x\mathbf{x} is

recalling that ∥⋅∥∞\left\|\cdot\right\|_{\infty} is the maximum absolute row sum.

Let zd≔xd−μz_{d}\coloneqq x_{d}-\mu. Hence ∑dzd=0\sum_{d}z_{d}=0, σ2=1D∑dzd2\sigma^{2}=\frac{1}{D}\sum_{d}z_{d}^{2} and

Noting that this expression is scale-invariant in z\mathbf{z}, we may assume WLOG max⁡d∣zd∣=zD=1\max_{d}|z_{d}|=z_{D}=1, since we are assuming not all xdx_{d} are equal and hence at least one zdz_{d} is non-zero.

Since all terms ∣zd∣≤1|z_{d}|\leq 1 are bounded, this continuous expression reaches a global maximum for some value of z\mathbf{z} with zD=1z_{D}=1.

It is easy to see that at the global maximum, zd≠0z_{d}\neq 0 ∀d\forall d: suppose this were to be true, WLOG z1=0z_{1}=0. Then let us see how the quantity (37) changes when z1=0z_{1}=0 is increased by 0<δ<10<\delta<1 and zD=1z_{D}=1 is decreased by δ\delta, keeping the sum constant. It is easy to see that the numerator ∑d∣zd∣\sum_{d}|z_{d}| stays constant, but the denominator ∑dzd2\sum_{d}z_{d}^{2} changes by 2δ2−2δ<02\delta^{2}-2\delta<0. Since for small δ\delta, the numerator of (37) stays constant but the denominator decreases, the quantity (37) increases, contradicting that the global max is obtained for z1=0z_{1}=0. Hence we may assume that zd≠0z_{d}\neq 0 ∀d\forall d.

Hence the quantity (37) (in particular, ∑d∣zd∣\sum_{d}{|z_{d}|}) is differentiable at the global maximum, at which the partial derivatives of the following Lagrangian are zero:

From now on let us write ∑\sum for ∑d<D\sum_{d<D} below to reduce clutter. Setting ∂L∂zk=0\frac{\partial\mathcal{L}}{\partial z_{k}}=0 and noting d∣zk∣dzk=sgn(zk)\frac{d|z_{k}|}{dz_{k}}=\text{sgn}(z_{k}), we obtain

Hence at the global maximum, zkz_{k} takes one of two values a>0a>0 and b<0b<0. Further we have that

If both aa and bb are among the zkz_{k}, we have that 1−λ(1+∑zd2)2a=−1−λ(1+∑zd2)2b\frac{1-\lambda(1+\sum z_{d}^{2})}{2a}=\frac{-1-\lambda(1+\sum z_{d}^{2})}{2b}. Solving for λ(1+∑zd2)\lambda(1+\sum z_{d}^{2}) and plugging it in back to Equation (38), we get:

Since a>0a>0, b<0b<0 and ∑zd=−1\sum z_{d}=-1, a−ba-b is minimised when only one of the zdz_{d} is aa and the rest are bb. Hence a crude lower bound on a−ba-b is 1D−2\frac{1}{D-2}, giving a bound:

However we conjecture that the true global maximum is attained when zd=−1D−1z_{d}=-\frac{1}{D-1} ∀d<D\forall d<D (i.e. all the zdz_{d} for d<Dd<D are equal to b<0b<0), for which it is easy to show that 1+∑d<D∣zd∣1+∑d<Dzd2=2(D−1)/D\frac{1+\sum_{d<D}|z_{d}|}{1+\sum_{d<D}z_{d}^{2}}=2(D-1)/D.