Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

The Lipschitz Constant of Self-Attention

Hyunjik Kim, George Papamakarios, Andriy Mnih

TL;DR

This work analyzes the Lipschitz properties of self-attention, showing that standard dot-product multi-head attention is not Lipschitz on unbounded input domains. It then introduces L2 self-attention, derives explicit upper bounds on its Lipschitz constants in both $\ell_{\infty}$ and $\ell_2$ norms, and demonstrates asymptotic tightness experimentally. Building on these results, the authors formulate invertible self-attention by constraining the attention mechanism (Contractive-L2-MHA) and validate it within a Transformer-style architecture on a character-level language modelling task, noting improved training stability. The work highlights practical pathways for integrating Lipschitz constraints into Transformer models, enabling applications such as invertible residual networks, Neural ODEs, and density estimation, while outlining future directions for alternative kernels and bounded-input strategies.

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 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.

The Lipschitz Constant of Self-Attention

TL;DR

This work analyzes the Lipschitz properties of self-attention, showing that standard dot-product multi-head attention is not Lipschitz on unbounded input domains. It then introduces L2 self-attention, derives explicit upper bounds on its Lipschitz constants in both and norms, and demonstrates asymptotic tightness experimentally. Building on these results, the authors formulate invertible self-attention by constraining the attention mechanism (Contractive-L2-MHA) and validate it within a Transformer-style architecture on a character-level language modelling task, noting improved training stability. The work highlights practical pathways for integrating Lipschitz constraints into Transformer models, enabling applications such as invertible residual networks, Neural ODEs, and density estimation, while outlining future directions for alternative kernels and bounded-input strategies.

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 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.

Paper Structure

This paper contains 32 sections, 13 theorems, 79 equations, 11 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Lipschitz Constant of Fully-Connected/Convolutional Layers
  3. Lipschitz Constant of Self-Attention
  4. Dot-product self-attention is not Lipschitz
  5. L2 self-attention: a Lipschitz formulation of self-attention
  6. Application: Invertible Self-Attention
  7. Invertible residual network
  8. Invertible self-attention
  9. Experimental Results
  10. Asymptotic tightness of the upper bound on $\boldsymbol{\mathop{\mathrm{Lip}}\nolimits_{\infty}(F)}$
  11. Numerical invertibility of MHA residual map
  12. Expressiveness of L2-MHA and invertible self-attention
  13. Training Stability of DP-MHA vs L2-MHA
  14. Conclusion and Discussion
  15. Useful Identities for deriving Jacobian expressions
  16. ...and 17 more sections

Key Result

Theorem 2.1

Let $f: \mathbb{R}^n \rightarrow \mathbb{R}^m$ be differentiable and Lipschitz continuous under a choice of $p$-norm $\|\cdot\|_p$. Let $J_f(x)$ denote its total derivative (Jacobian) at $x$. Then $\mathop{\mathrm{Lip}}\nolimits_p(f) = \sup_{\mathbf{x}\in \mathbb{R}^n} \|J_f(\mathbf{x})\|_p$ where $

Figures (11)

  • Figure 1: Transformer block.
  • Figure 2: Lower and upper bound on $\mathop{\mathrm{Lip}}\nolimits_{\infty}(f)$ for L2-MHA $f$, with $H=D=1$ and varying $N$.
  • Figure 3: Invertibility of $g(\mathbf{x})= \mathbf{x} + c f(\mathbf{x})$ where $f$ is L2-MHA (left) and DP-MHA (right).
  • Figure 4: Test NLL curves during training for various LSTM/Transformer models on PTB character level language modelling.
  • Figure 5: Histogram showing distribution of inputs/outputs of trained L2-MHA and DP-MHA
  • ...and 6 more figures

Theorems & Definitions (25)

  • Definition 2.1
  • Theorem 2.1: federer1969geometric
  • Lemma 2.1: federer1969geometric
  • Corollary 2.1
  • Theorem 3.1
  • proof
  • Theorem 3.2
  • proof
  • proof
  • Lemma 6.1
  • ...and 15 more