PaperlayerView exact revision on arXiv

Source-linked AI summary

The Lipschitz Constant of Self-Attention

Hyunjik Kim, George Papamakarios, Andriy Mnih

arXiv:2006.04710v2stat.MLcs.LG

TL;DR

The paper examines whether widely used self-attention satisfies Lipschitz constraints needed in several deep-learning applications. It proves that standard dot-product attention is not Lipschitz on unbounded inputs, develops an L2 alternative with explicit bounds, and evaluates invertible attention in character-level language modelling.

  • Problem

    Prior Lipschitz analyses focused on fully connected and convolutional networks, leaving the properties of widely used self-attention unclear despite applications requiring Lipschitz constraints.

  • Method

    The paper replaces the dot-product kernel with an L2 kernel, derives Lipschitz upper bounds, and uses them to formulate invertible self-attention in a Transformer.

  • Results

    Dot-product self-attention is not Lipschitz on unbounded inputs, whereas L2 self-attention is Lipschitz with bounds simplifying to O(log N) for p = ∞ and O(N log N) for p = 2.

  • Takeaways & Limitations

    L2-MHA loses little expressiveness, while normalising it for invertibility reduces expressiveness but stabilises optimisation and permits deeper architectures.

  • Takeaways & Limitations

    The theoretical upper bound may have a loose multiplicative constant, whose tightening remains an open question.

Abstract

from arXiv · show

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 for unbounded input domain, 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.

1. Introduction

The paper addresses the previously unanalysed Lipschitz properties of self-attention, motivated by the importance of Lipschitz constraints and self-attention's broad use. It proves that dot-product self-attention fails to be Lipschitz and introduces Lipschitz L2 self-attention with theoretical and practical evaluation.

  • Lipschitz constraints control output variation relative to input and support adversarial robustness, generalisation bounds, Wasserstein estimation, GAN stabilisation, and invertible networks.
  • Exact Lipschitz computation is hard even for fully-connected and convolutional networks, while self-attention had not previously been analysed.
  • Self-attention is a widely used Transformer component spanning language, vision, audio generation, and reinforcement learning.
  • Dot-product self-attention is not Lipschitz and is therefore unsuitable for applications requiring Lipschitz constraints.
  • L2 self-attention is proposed as a Lipschitz alternative, with an upper bound whose asymptotic tightness is assessed empirically.
  • The theory is applied to invertible self-attention in a Transformer architecture for character-level language modelling, comparing test log-likelihood and stability with dot-product attention.

2. Lipschitz Constant of Fully-Connected/Convolutional Layers

This section establishes the framework for Lipschitz analysis of neural networks by defining norm-dependent constants, relating them to Jacobians, and bounding compositions through products of component constants.

  • Lipschitz continuity requires a finite constant limiting output changes to K times input changes, with the smallest such K denoted Lip(f).
  • The paper primarily studies p = 2 and p = ∞ norms, so the Lipschitz constant is written Lip_p(f) to show norm dependence.
  • For differentiable functions, Lip_p(f) equals the supremum of the induced p-norm of the Jacobian over all inputs.
  • For linear maps represented by W, the Lipschitz constant is the induced operator norm ∥W∥_p; for p = 2, this is the largest singular value σmax(W).
  • Compositions of Lipschitz functions have constants bounded by the product of their individual constants, yielding products of weight norms for FCNs and CNNs with 1-Lipschitz nonlinearities.
  • The resulting FCN and CNN bound is not necessarily tight, and obtaining tighter bounds remains an active concern.

3. Lipschitz Constant of Self-Attention

The paper shows that standard dot-product self-attention is not Lipschitz on an unbounded domain, then replaces its kernel with an L2 formulation that is Lipschitz under tied query and key weights and admits explicit bounds.

  • 3.1. Dot-product self-attention is not Lipschitz: Dot-product multihead self-attention maps sequence inputs through softmax weights derived from pairwise dot products, followed by value projections and head combination.Each head uses learnable query, key, and value projections; the head outputs are concatenated and multiplied by an output matrix.
  • 3.1. Dot-product self-attention is not Lipschitz: For non-trivial parameters, DP-MHA is not Lipschitz for any vector p-norm with p ∈ [1, ∞].
  • 3.1. Dot-product self-attention is not Lipschitz: At xi = 0, the Jacobian can grow with the sample variance of the other inputs, which is arbitrarily large and makes the Jacobian unbounded.The proof uses the Jacobian criterion for Lipschitzness and identifies an unbounded diagonal Jacobian element.
  • 3.1. Dot-product self-attention is not Lipschitz: Bounding inputs can make dot-product self-attention Lipschitz on a compact domain, but the resulting Lipschitz constant depends on the input range.
  • 3.2. L2 self-attention: a Lipschitz formulation of self-attention: L2 self-attention replaces the dot-product kernel with an L2-distance-based, softmax-normalised kernel whose rows sum to 1.The formulation retains softmax normalisation to handle sequences of varying length and remains efficiently computable with matrix operations.
  • 3.2. L2 self-attention: a Lipschitz formulation of self-attention: Tying query and key weights is sufficient for L2-MHA to be Lipschitz, whereas arbitrary query and key weights are not sufficient.The formulation also includes an additional right multiplication by A_h for the theorem to hold.
  • 3.2. L2 self-attention: a Lipschitz formulation of self-attention: L2-MHA has bounds simplifying to O(log N) for p = ∞ and O(N log N) for p = 2, where N is sequence length.

4. Application: Invertible Self-Attention

The paper constructs invertible self-attention by placing a contractive L2-MHA module in a Transformer residual branch. This follows the general invertibility of residual functions whose added component is contractive.

  • A residual function g(x) = x + f(x) is invertible when f is contractive, with Lip(f) < 1, on a complete metric space.Its inverse can be obtained as the unique fixed point of the recursion x_i+1 := y − f(x_i).
  • Compositions of invertible residual blocks remain invertible, motivating normalized neural residual architectures.Invertible ResNets use linear-map normalization to make the residual function contractive.
  • Contractive-L2-MHA is obtained by normalizing L2-MHA using the upper bounds from Theorem 3.2, yielding invertible self-attention with a skip connection.Dropout remains contractive in both evaluation and training settings, while LayerNorm is outside the residual branch and can be replaced by an invertible normalization.
  • The modified Transformer replaces dot-product multihead self-attention with Contractive-L2-MHA to test invertible self-attention against the standard module.The Transformer serves as a testbed because self-attention is commonly used within Transformer blocks.
  • The application study does not target the modified Transformer itself; a more challenging normalizing-flow application is left for separate work.This scope boundary reflects the paper’s theoretical focus on comparing invertible and standard self-attention.

5. Experimental Results

Experiments examine the tightness of the L2-MHA Lipschitz bound, numerical invertibility of residual maps, and expressiveness and stability in character-level language modelling.

  • Asymptotic tightness: The lower bound’s slope closely matches the upper bound, supporting asymptotic tightness of the O(log N − log log N) bound.The experiment varies N from 100 to 1000 with H = D = 1 and optimises the Jacobian norm from 50 random initialisations.
  • Asymptotic tightness: The gap between bounds may reflect local optimisation or a loose multiplicative constant in the upper bound.The true Lipschitz constant is a global optimum, whereas the empirical lower bound is only a local optimum.
  • Numerical invertibility: For sampled inputs, DP-MHA is not invertible whereas L2-MHA is invertible for sufficiently small c.This comparison concerns local numerical invertibility of g(x) = x + cf(x), not global invertibility over the whole input space.
  • Expressiveness: Replacing dot-product attention with L2-MHA causes hardly any expressiveness loss, while tying query and key weights causes a small loss.The best models in each self-attention setting have similar generalisation performance, and matching parameter counts does not remove the tied-weight degradation.
  • Expressiveness: Normalising by the Lip∞ upper bound causes a noticeable expressiveness loss but stabilises optimisation, allowing compensation through deeper Transformers.Higher-depth Transformers can otherwise show training instabilities requiring careful learning-rate tuning.
  • Training stability: Trained L2-MHA outputs remain in a range similar to inputs, whereas trained DP-MHA produces outputs with a much wider range.The paper attributes the wider DP-MHA output range to its large Jacobian norm at the sampled inputs.

6. Conclusion and Discussion

The paper establishes Lipschitz bounds for L2 self-attention, demonstrates their practical use in invertible self-attention, and identifies limitations of bounded-input alternatives.

  • Practical implications: Lipschitz-constrained self-attention is used to formulate invertible self-attention and evaluated in a Transformer for character-level language modelling.The evaluation compares test log-likelihood and stability with dot-product self-attention.
  • Practical implications: L2-MHA is more stable during training, allowing fixed learning rates for stable training of deep architectures.This is reported as an experimental finding in the conclusion.
  • Discussion: Replacing the dot-product kernel with an L2 kernel avoids making the Lipschitz constant depend on a bounded input range.Constraining inputs to a compact set would require range-dependent bounds and complicate analysis across multiple self-attention layers.
  • Discussion: Provably Lipschitz self-attention enables Transformer architectures in applications requiring Lipschitz constraints, including residual flows and Neural ODE parameterisations.For Neural ODEs, a Lipschitz vector field guarantees a unique solution for all times.

C. Proof of Theorem 3.1 for General D

The proof shows that dot-product self-attention has an unbounded Jacobian, including in general dimension and under multihead composition, while L2 attention remains efficiently computable.

  • Theorem 3.1: DP-MHA is not Lipschitz for any vector p-norm with p ∈ [1, ∞].The proof uses an unbounded Jacobian as the criterion for non-Lipschitzness.
  • Proof: When xi = 0, attention becomes uniform regardless of the other inputs, while covariance terms equal sample variances that can grow arbitrarily large.These unbounded variance terms make entries of the Jacobian unbounded.
  • Proof: The single-head non-Lipschitz result extends to multihead self-attention because the multihead output is a linear combination of head outputs.Thus, adding heads does not remove the unbounded behavior.
  • Bias terms: Adding bias terms does not generally resolve the problem: choosing x_i so that x_i^T W^Q + b^Q = 0 again yields uniform attention and unbounded diagonal covariance terms.The construction assumes W^Q is full rank.
  • Efficient computation: L2 self-attention logits can be computed through matrix multiplications and row-wise L2 norms with negligible overhead compared with dot-product self-attention.The identity ||a − b||_2^2 = ||a||_2^2 − 2a^T b + ||b||_2^2 enables this computation.

F. Proof of Theorem 3.2

The proof bounds the Jacobian of L2 self-attention by exploiting covariance structure, permutation equivariance, and a restriction tying key and query projections.

  • Restriction: W^K = W^Q is imposed because otherwise L2 self-attention is not Lipschitz when W^K is full rank.If W^K ≠ W^Q, off-diagonal Jacobian terms can be unbounded.
  • p = ∞ bound: For p = ∞, permutation equivariance reduces the global Jacobian bound to the maximum norm of a single block row.All block rows have the same optimized maximum under permutation equivariance.
  • Core inequality: The key scalar bound is φ^-1(N − 1) = O(log N − log log N), where φ(c) = c exp(c + 1).This bound controls covariance-related terms appearing in the Jacobian.
  • p = 2 bound: For p = 2, block-row spectral norms are combined using block-matrix inequalities to obtain an O(N(log N − log log N)) bound.The section notes that this spectral-norm bound may be loose and that tightening it remains open.
  • Masking: Masking preserves the unmasked Lipschitz bound when each output attends to itself, because masked Jacobian blocks are zero.The self-attention restriction (i, i) ∈ M is required for the proof of the scalar bound.

H. Experimental Details

The experiments optimize inputs to test bound tightness and compare standard, Lipschitz, and invertible Transformer self-attention on character-level Penn Treebank language modelling.

  • Bound tightness: 50 random initialisations of X are used for each N to optimize the input-dependent Lipschitz objective until convergence with Adam.The free parameters are fixed to identity matrices, and initial entries are sampled uniformly within a randomly selected range.
  • Language modelling: Character-level language modelling is evaluated on the Penn Treebank dataset with original and Lipschitz/invertible Transformer self-attention.Examples are variable-length character sequences, batched by length, with maximum sequence length 288.

I. Numerical Invertibility of MHA Residual Map

Figures 6 and 7 examine numerical invertibility of residual maps using dot-product and L2 multihead self-attention across trained weights and varying N and D.

  • Figure 6 confirms that numerical invertibility does not hold for trained dot-product multihead self-attention.The trained weights come from a one-layer Transformer (DP) model.
  • The experiments extend the trained-weight comparison beyond the single setting shown for Figure 6.
  • Figure 7 compares numerical invertibility of g(x) = x + cf(x) for L2-MHA and DP-MHA across different N and D values.

J. Behaviour of Lower Bound on Lip2(F )

The empirical lower bound for Lip2(F) is difficult to optimize and appears to grow as O(log N), leaving the asymptotic tightness of the theoretical bound unresolved.

  • O(log N) is the apparent growth rate of the empirical lower bound on Lip2(F) for L2-MHA.The trend is less clear than in the corresponding analysis, and the top five optimized values show variance.
  • The optimization may produce only a loose lower bound because it is difficult even for small N.
  • If the lower bound is tight, the O(N log N) upper bound in Theorem 3.2 may be improvable to O(log N) or O(log N −log log N) for p = ∞.
  • The optimized Jacobian norm for DP-MHA keeps increasing with X, supporting the theoretical result that DP-MHA is not Lipschitz.Oscillations are attributed to the momentum term of the Adam optimizer.

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

The section tests whether tying keys and queries reduces performance because of fewer parameters by comparing a parameter-matched shared L2-MHA model with the unshared model.

  • Doubling the columns of W Q from D/H to 2D/H gives the shared model the same parameter count as the unshared model.This experiment tests whether the performance reduction from tying keys and queries is attributable to parameter count.
  • Figure 10 reports results for shared L2-MHA with the parameter count matched to the unshared model.
  • LayerNorm is Lipschitz, but the paper derives a loose upper bound with respect to the ∞-norm.
  • The LayerNorm analysis assumes D > 2 and inputs whose coordinates are not all equal, excluding the trivial constant-input or D = 1 cases.
  • The derivation bounds the maximum absolute row sum of the LayerNorm Jacobian to obtain an ∞-norm Lipschitz bound.
  • The LayerNorm optimization uses scale invariance to set max_d |z_d| = z_D = 1 and then characterizes maximizers with two coordinate values.
  • The analysis yields a crude lower bound based on minimizing a − b, while the conjectured global maximizer has one coordinate value and the remaining coordinates equal to b.
Loading 2006.04710v2…