# Understanding the Regularity of Self-Attention
with Optimal Transport

Valérie Castin  DMA, École Normale Supérieure PSL  Pierre Ablin  Apple
Gabriel Peyré  DMA, École Normale Supérieure PSL  CNRS

###### Abstract

Transformers and their multi-head attention mechanism have completely changed
the machine learning landscape in just a few years, by outperforming state-of-
art models in a wide range of domains. Still, little is known about their
robustness from a theoretical perspective. We tackle this problem by studying
the local Lipschitz constant of self-attention, that provides an attack-
agnostic way of measuring the robustness of a neural network. We adopt a
measure-theoretic framework, by viewing inputs as probability measures
equipped with the Wasserstein distance. This allows us to generalize attention
to inputs of infinite length, and to derive an upper bound and a lower bound
on the Lipschitz constant of self-attention on compact sets. The lower bound
significantly improves prior results, and grows more than exponentially with
the radius of the compact set, which rules out the possibility of obtaining
robustness guarantees without any additional constraint on the input space.
Our results also point out that measures with a high local Lipschitz constant
are typically made of a few diracs, with a very unbalanced distribution of
mass. Finally, we analyze the stability of self-attention under perturbations
that change the number of tokens, which appears to be a natural question in
the measure-theoretic framework. In particular, we show that for some inputs,
attacks that duplicate tokens before perturbing them are more efficient than
attacks that simply move tokens. We call this phenomenon mass splitting.

###### Contents

  1. 1 Introduction
    1. 1.1 Contributions
    2. 1.2 Related Work
    3. 1.3 Notations
  2. 2 Generalizing Self-Attention to Probability Measures
  3. 3 Estimating the Lipschitz Constant of Self-Attention
    1. 3.1 Lipschitz Constant and Local Lipschitz Constant
      1. 3.1.1 Local Lipschitz Constant on Probability Measures versus on Matrices
    2. 3.2 Lower Bound on the Lipschitz Constant of Self-Attention
    3. 3.3 Tightness of the Lower Bound
      1. 3.3.1 Weighted Self-Attention
      2. 3.3.2 Experiments
    4. 3.4 The Lipschitz Constant of Masked Attention
  4. 4 Mass Splitting
    1. 4.1 Finite Mass Splitting: a Proxy for Mass Splitting
      1. 4.1.1 Duplication Is the Worst Case for Finite Mass Splitting
    2. 4.2 Numerical Evidence of the Mass Splitting Phenomenon
  5. A Optimal Transport Toolbox
    1. A.1 Pushforward, Wasserstein Distance
    2. A.2 Geodesics
  6. B Proofs of Section 3
    1. B.1 Proof of Lemma 3.1
    2. B.2 Proof of Lemma 3.3
    3. B.3 Proof of Proposition 3.7
    4. B.4 Weighted Self-Attention
    5. B.5 Proof of Proposition 3.5
  7. C Proofs of Section 4
    1. C.1 Effect of Duplication on the Jacobian of Self-Attention
    2. C.2 Concentrated Measures Are Not Splitted
    3. C.3 Proof of Theorem 4.2

##  1 Introduction

Introduced by [58], Transformers and their multi-head attention mechanism [2]
have completely changed the machine learning landscape in just a few years, by
outperforming state-of-art models on a wide variety of tasks, from natural
language processing [7, 47, 64] to computer vision [16, 67, 66, 30]. Despite
this great empirical success, however, little is known about the robustness of
Transformer architectures – and in particular of self-attention, their main
building block – from a theoretical perspective. We tackle this problem by
focusing on the Lipschitz properties of self-attention, especially on its
local Lipschitz constant, which controls how fast the output can change with
respect to the input in the neighborhood of each point of the domain.

Studying the (local) Lipschitz continuity of neural networks is of particular
interest for various questions [50]. Indeed, estimating the local Lipschitz
constant of a neural network provides guarantees of adversarial robustness, in
an attack-agnostic way [56, 13, 57, 1, 63]. Identifying inputs with a high
local Lipschitz constant and understanding which local perturbation triggers
the biggest change in the output also allows to robustify the network, for
example using adversarial training [20, 37, 40, 28]. The Lipschitz constant is
also involved in generalization bounds [54, 41, 3, 36]. From a different
perspective, Lipschitz constrained neural networks can be used to estimate
Wasserstein distances [45], stabilize training of GANs [38], and build
invertible neural networks [4, 11]. Finally, bounding the Lipschitz constant
of a neural network is an important step in the study of the associated neural
ODE [10], in particular of its well-posedness [35, 19].

Lipschitz continuity of feed-forward and convolutional neural networks has
been extensively studied and remains a hard problem, the main difficulty being
to estimate precisely the Lipschitz constant of a composition of several
linear maps and pointwise linearities [59, 17, 29]. Still, taken
independently, each linear map or activation function has a known Lipschitz
constant. This is however not the case for Transformers: the self-attention
map has an involved non-linear structure, which makes the estimation of its
local Lipschitz constant more challenging. [26] show that the self-attention
map is not globally Lipschitz continuous by proving a lower bound on its
Lipschitz constant on the closed ball $\,\overline{\\!{B}}(0,R)$, lower bound
that grows quadratically with $R$. On the other side, [19] and [60, 61] derive
an upper bound on the Lipschitz constant of self-attention on
$\,\overline{\\!{B}}(0,R)$, by viewing self-attention as a map acting on
probability measures. This approach, that builds on the fact that self-
attention, denoted $f$ in what follows, is permutation equivariant – i.e. does
not distinguish vectors according to their order in the input, which means
that for all permutations $\sigma$ of the set $\\{1,\dots,n\\}$ and inputs
$X\in\mathbb{R}^{n\times d}$, it holds

$f(X_{\circ\sigma})=f(X)_{\circ\sigma},$

with $Y_{\circ\sigma}:=(y_{\sigma(1)},\dots,y_{\sigma(n)})^{\top}$ for any
matrix $Y=(y_{1},\dots,y_{n})\in\mathbb{R}^{n\times d}$, allows them to use
tools from optimal transport to bound the Lipschitz constant of self-attention
in the sense of Wasserstein distance. Their upper bound grows more than
exponentially with $R$, so that the quadratic lower bound and the exponential
upper bound put together provide a very loose estimation of the Lipschitz
constant of self-attention on compact sets. It is therefore important to
understand which one of the two bounds can be improved, if not both.

###  1.1 Contributions

We make the following contributions.

  * $\bullet$

We bridge the gap between the matrix formalism and the measure-theoretic
formalism of self-attention by proving that they are equivalent when
restricting to empirical measures with a fixed number of samples (Lemma 3.3).
Such a connection between these two viewpoints on self-attention is absent in
[26], that use the matrix formalism, as well as in [60, 61] and [19], that use
measure-theoretic frameworks.

  * $\bullet$

We significantly improve the lower bound obtained by [26] on the Lipschitz
constant of self-attention on the closed ball $B_{R}$ (centered in 0 and of
radius $R$), that grows like $R^{2}$. More precisely, we find a family of two-
dirac measures $(\mu_{R})_{R>0}$ such that $\mu_{R}$ is supported in $B_{R}$
and the $W_{2}$ local Lipschitz constant of self-attention at $\mu_{R}$ is
equal to

$\frac{C}{2}R^{2}e^{CR^{2}}$

with $C$ depending on the parameters of self-attention, up to a constant
factor close to 1 (Theorem 3.4). This result shows that the upper bound
obtained by [19], which is of the form

$C^{\prime}R^{2}e^{C^{\prime}R^{2}}$

up to a constant factor close to 1, with $C^{\prime}\leq 16\,C$, cannot be
significantly improved, in the sense that it is necessarily exponential with
respect to $R^{2}$. Self-attention is therefore a very brittle function in the
Wasserstein sense, which is somewhat disappointing from the perspective of a
robustness analysis.

  * $\bullet$

We show numerically that our lower bound is tight up to a constant factor
close to 1, as well-chosen configurations with two unbalanced clusters – one
with a vanishing amount of mass, and the other concentrating the rest of the
mass, appear to maximize the local Lipschitz constant when fixing the support.
(Subsection 3.3)

  * $\bullet$

We generalize residual masked self-attention to probability measures by
introducing an innovative framework, where the order of points in input
measures is encoded in a supplementary coordinate, and find an upper bound of
the form

$1+CR^{2}e^{CR^{2}}$

on the Lipschitz constant of residual masked self-attention restricted to
measures supported in $[0,1]\times\,\overline{\\!{B}}(0,R)$. (Subsection 3.4)

  * $\bullet$

We evidentiate numerically, and theoretically in some special cases, a _mass
splitting_ phenomenon: in order to induce the greatest change in the output,
the best strategy to perturb an input with $n$ points is sometimes to
duplicate points and perturb the duplicated version of the output, which
amounts to a very small perturbation in the Wasserstein sense but can induce a
significant change in the output. (Section 4)

###  1.2 Related Work

##### Robustness and local Lipschitz constant estimation.

Neural networks are vulnerable to adversarial attacks [56], and most of the
methods proposed to measure and increase their robustness focus on specific
attacks [20, 43]. It turns out, however, that such methods can be defeated by
well-chosen unseen attacks [8]. Measures of robustness that are agnostic to
attack methods have therefore been proposed, often relying on the notion of
Lipschitz constant of networks [56, 31, 57]. As robustness lower bounds that
rely on the (global) Lipschitz constant tend to be too loose, tighter
constraints have been proposed involving the local Lipschitz constant [22,
63]. The problem of evaluating the local Lipschitz constant of deep networks
is now at the heart of several recent articles [57, 31], in particular for
Transformers [26, 60, 19]. From a more practical viewpoint, several Lipschitz-
constrained variants of the Transformer architecture have been proposed, to
increase robustness and reliability [25, 21, 65, 46].

##### Neural networks acting on measures.

[14] and [44] are the first to define neural networks acting on probability
measures, followed by several other articles [60, 68, 51, 19]. Modeling neural
networks as maps on probability measures has multiple applications, such as
studying Wasserstein regularity [60, 19], proving generalization bounds [68]
and doing a mean-field limit analysis of the dynamics of particles as they go
through the network [19]. The measure-theoretic approach is particularly
suited to the case of Encoder-only Transformers [15], as the self-attention
map is permutation equivariant, i.e., ignores the order of vectors in its
input. This property can be leveraged to model any infinitely deep Encoder as
a partial differential equation (PDE) on the space of measures, following the
principle of neural ODEs [10]. Analyzing this PDE then provides information
about the dynamics of tokens as they go through the Transformer, showing for
instance the emergence of clusters [19, 18]. In contrast, masked self-
attention, which is crucial in Decoder-only [33] and Encoder-Decoder [58]
architectures, is not permutation equivariant, so cannot be cast as naturally
into a measure-theoretic framework.

##### Regularity of self-attention and its variants.

[26] show that the self-attention map is not globally Lipschitz continuous by
proving a lower bound on its Lipschitz constant on the closed ball
$\,\overline{\\!{B}}(0,R)$, lower bound that grows quadratically with $R$. To
gain regularity, they define a new self-attention map, called L2 self-
attention, that is globally Lipschitz continuous on the set of inputs of
length $n$, for all $n\geq 1$. Their upper bound on the Lipschitz constant of
L2 self-attention however diverges with $n$. [19] and [60] prove a length-
independent upper bound on the Lipschitz constant of self-attention on
$\,\overline{\\!{B}}(0,R)$, by viewing self-attention as a map acting on
probability measures. Their upper bound grows more than exponentially with
$R$, so that the quadratic lower bound and the exponential upper bound put
together provide a very loose estimation of the Lipschitz constant of self-
attention on compact sets. Finally, [51] propose a modification of the
attention kernel that builds on the Sinkhorn-Knopp algorithm, and provide
empirical evidence of the better properties of this new choice of kernel with
respect to the classical one.

###  1.3 Notations

For any vector $x\in\mathbb{R}^{d}$, the Euclidean norm of $x$ is denoted
$\left\lvert x\right\rvert$. For $R>0$, the open ball centered in $x$ and of
radius $R$ is denoted $B(x,R)$, and we write $\,\overline{\\!{B}}(0,R)$ or
simply $B_{R}$ the corresponding closed ball. The simplex
$\\{a\in[0,1]^{n}\mid\sum_{i=1}^{n}a_{i}=1\\}$ is denoted $\Sigma_{n}$. The
set $\mathcal{P}(\mathbb{R}^{d})$ is the set of probability measures on
$\mathbb{R}^{d}$. The set of probability measures with compact support is
denoted $\mathcal{P}_{c}(\mathbb{R}^{d})$, and
$\mathcal{M}_{n}(\mathbb{R}^{d})$ is the set of empirical measures on
$\mathbb{R}^{d}$ with $n$ samples. $\delta_{x}$ denotes the Dirac measure at
point $x\in\mathbb{R}^{d}$. The vector $1_{n}\in\mathbb{R}^{n}$ is the vector
of dimension $n$ with all coordinates equal to 1. For any vector
$x\in\mathbb{R}^{d}$, we denote
$\operatorname{\mathrm{diag}}(x)\in\mathbb{R}^{d\times d}$ the associated
diagonal matrix. Coordinatewise multiplication for vectors and matrices is
denoted $\odot$, and $\otimes$ is the Kronecker product between matrices. For
a block matrix $M=(M_{ij})_{1\leq i,j\leq n}\in\mathbb{R}^{dn\times dn}$ with
blocks $M_{i,j}$ of shape $d\times d$, and a matrix $A\in\mathbb{R}^{d\times
d}$, we denote $M\otimes_{d}A$ the block matrix $(M_{ij}A)_{1\leq i,j\leq
n}\in\mathbb{R}^{dn\times dn}$. For a function
$g\colon\mathcal{E}\to\mathcal{F}$ and a subset
$\mathcal{X}\subset\mathcal{E}$, the restriction of $g$ to $\mathcal{X}$ is
denoted $g_{\lvert\mathcal{X}}$. The Jacobian of $g$ at point $x$ is denoted
$D_{x}g$.

In the whole paper, the Jacobian of self-attention $f\colon\mathbb{R}^{n\times
d}\to\mathbb{R}^{n\times d}$ is to be understood as the Jacobian of the
vectorized form of $f$, defined by
$\operatorname{\mathrm{vec}}(X)\in\mathbb{R}^{nd}\mapsto\operatorname{\mathrm{vec}}(f(X))\in\mathbb{R}^{nd}$
where vec denotes row-wise vectorization. For example,

$\begin{pmatrix}1&2&3\\\ 4&5&6\\\ 7&8&9\end{pmatrix}$

gives the order of coordinates in our choice of vectorization for a $3\times
3$ matrix. We will sometimes write $\operatorname{\mathrm{vec}}X$ as $\vec{X}$
to alleviate notations.

##  2 Generalizing Self-Attention to Probability Measures

Transformers and self-attention. Introduced by [58], Transformers are
attention-based models that can process several types of data, from text [7,
47, 64] to images [16] and audio recordings [12]. When fed to a Transformer,
each data point is converted to a sequence of tokens (tokenization), then to a
sequence of vectors (input embedding). Then, positional encoding [58] adds
information about the order of vectors in the vectors themselves. The
resulting sequence of vectors $(x_{1},\dots,x_{n})\in(\mathbb{R}^{d})^{n}$,
that can be seen as a $n\times d$ matrix, is said to belong to the feature
space. In encoder-only models [15, 52, 34, 16],
$X\coloneqq(x_{1},\dots,x_{n})^{\top}$ then goes through a succession of
layers, each one made of two blocks. First, a residual multi-head self-
attention block of the form

$X\mapsto X+\sum_{h=1}^{H}f^{(h)}(X)W_{h}^{\top}$

with $W_{h}\in\mathbb{R}^{d\times k}$ and $f^{(h)}$ (single-head) self-
attention maps, where the self-attention map is defined as

$f\colon X\in\mathbb{R}^{n\times
d}\mapsto\operatorname{\mathrm{softmax}}\left(\frac{1}{\sqrt{k}}XQ^{\top}KX^{\top}\right)XV^{\top}\in\mathbb{R}^{n\times
d}$ (1)

where $k$ is a positive integer, $Q,K,V$ are $k\times d$ matrices that depend
on $h$ and on the layer, and

$\operatorname{\mathrm{softmax}}(w_{1},\dots,w_{n})\coloneqq\left(\frac{\exp(w_{i})}{\sum_{j=1}^{d}\exp(w_{j})}\right)_{1\leq
i\leq n}$

for $w_{1},\dots,w_{n}\in\mathbb{R}^{d}$. Then, a residual multilayer
perceptron applied vectorwise, i.e. of the form
$(x_{1},\dots,x_{n})\mapsto(x_{1},\dots,x_{n})+(g(x_{1}),\dots,g(x_{n}))$ with
$g$ a multilayer perceptron. Decoder-only models [33, 42, 9, 24] have the same
structure, but replace self-attention with masked self-attention, defined as

$f^{masked}(X)_{i}\coloneqq\operatorname{\mathrm{softmax}}\left(\frac{1}{\sqrt{k}}x_{i}Q^{\top}KX_{1:i}^{\top}\right)X_{1:i}V^{\top}$
(2)

for all $1\leq i\leq n$, where $X_{1:i}=(x_{1},\dots,x_{i})^{\top}$. Finally,
encoder-decoder models [58, 62, 48, 49, 32] couple self-attention, masked
self-attention and cross-attention, that is not considered in this work. It is
important to notice that the self-attention map introduced in Equation (1) is
permutation equivariant, which means that for all permutations $\sigma$ of the
set $\\{1,\dots,n\\}$, it holds

$f(X_{\circ\sigma})=f(X)_{\circ\sigma},$

with $Y_{\circ\sigma}:=(y_{\sigma(1)},\dots,y_{\sigma(n)})^{\top}$ for any
matrix $Y=(y_{1},\dots,y_{n})\in\mathbb{R}^{n\times d}$. Intuitively, a
permutation equivariant map does not use any information on the order of its
input vectors. On the contrary, the masked self-attention map has a sequential
structure: it is not permutation equivariant. We show in Subsection 3.4 that
it can still be cast into a measure-theoretic framework, by adding a
coordinate to the input space to encode the order of points.

##### Pushforward and Wasserstein distances.

Before introducing the measure-theoretic framework, we need a few notions from
optimal transport. For a probability measure $\mu$ on $\mathbb{R}^{d}$ and a
measurable map $\varphi\colon\mathbb{R}^{d}\to\mathbb{R}^{d}$, the pushforward
of $\mu$ by $\varphi$, denoted $\varphi_{\sharp}\mu$, is the probability
measure given by

$\left(\varphi_{\sharp}\mu\right)(B)\coloneqq\mu(\varphi^{-1}(B))$

for any Borel set $B\subset\mathbb{R}^{d}$, where
$\varphi^{-1}(B)\coloneqq\\{x\in\mathbb{R}^{d}:\varphi(x)\in B\\}$.
Intuitively, $\varphi_{\sharp}\mu$ is obtained by transporting each element of
mass $\mu(\mathrm{d}x)$ from $x$ to $\varphi(x)$. We also have the following
useful property of the pushforward. For all probability measures
$\mu\in\mathcal{P}(\mathbb{R}^{d})$ and maps
$\varphi\colon\mathbb{R}^{d}\to\mathbb{R}^{d}$ and
$g\colon\mathbb{R}^{d}\to\mathbb{R}$, it holds, when both members of the
equation are well-defined:

$\int g(x)\mathrm{d}\left(\varphi_{\sharp}\mu\right)(x)=\int
g(\varphi(x))\mathrm{d}\mu(x).$

Let $p\geq 1$. Denote

$\mathcal{P}_{p}(\mathbb{R}^{d})\coloneqq\\{\mu\in\mathcal{P}(\mathbb{R}^{d}):\int\left\lvert
x\right\rvert^{p}\mathrm{d}\mu(x)<+\infty\\}.$

For $\mu$ and $\nu$ in $\mathcal{P}_{p}(\mathbb{R}^{d})$, the $p$-Wasserstein
distance between $\mu$ and $\nu$ is defined as

$W_{p}(\mu,\nu)\coloneqq\left(\inf_{\pi\in\Pi(\mu,\nu)}\int\left\lvert
x-y\right\rvert^{p}\mathrm{d}\pi(x,y)\right)^{1/p},$

where $\Pi(\mu,\nu)$ is the set of couplings between $\mu$ and $\nu$, i.e. of
probability measures $\pi\in\mathcal{P}(\mathbb{R}^{d}\times\mathbb{R}^{d})$
such that $\int\pi(\cdot,y)\mathrm{d}y=\mu$ and
$\int\pi(x,\cdot)\mathrm{d}x=\nu$. More details on the aforementioned objects
can be found for example in [53].

##### Neural networks acting on measures.

The self-attention map defined in Equation (1) is permutation equivariant,
which means that permuting $x_{i}$ and $x_{j}$ in the input of $f$ amounts to
permuting $f(X)_{i}$ and $f(X)_{j}$ in the output. This remark allows to view
self-attention as acting on point clouds, or on empirical measures [14, 68,
51, 19] by mapping the matrix input $X=(x_{1},\dots,x_{n})^{\top}$ to the
associated empirical measure

$m(X)\coloneqq\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}}.$ (3)

Following [51], we can then extend self-attention to more general measures in
the following way. We can rewrite Equation (1) as

$f(X)=(\Gamma_{X}(x_{1}),\dots,\Gamma_{X}(x_{n})),$ (4)

with

$\Gamma_{X}\colon
x\mapsto\frac{\sum_{i=1}^{n}\exp\left(\frac{1}{\sqrt{k}}x^{\top}Q^{\top}Kx_{i}\right)Vx_{i}}{\sum_{i=1}^{n}\exp\left(\frac{1}{\sqrt{k}}x^{\top}Q^{\top}Kx_{i}\right)}.$

Seeing the ratio of sums as a ratio of integrals against the empirical measure
$m(X)$, and Equation (4) as a pushforward, we can define measure-theoretic
self-attention as

$F:\mu\in\mathcal{P}_{c}(\mathbb{R}^{d})\mapsto\left(\Gamma_{\mu}\right)_{\sharp}\mu$
(5)

where $\mathcal{P}_{c}(\mathbb{R}^{d})$ is the set of compactly supported
probability measures on $\mathbb{R}^{d}$, with

$\Gamma_{\mu}\colon
x\in\mathbb{R}^{d}\mapsto\frac{\int\exp\left(\frac{1}{\sqrt{k}}x^{\top}Q^{\top}Ky\right)Vy\,\mathrm{d}\mu(y)}{\int\exp\left(\frac{1}{\sqrt{k}}x^{\top}Q^{\top}Ky\right)\mathrm{d}\mu(y)}$

for all $\mu$ such that the integrals are well-defined, which is in particular
the case when $\mu$ is compactly supported. One sees that measure-theoretic
self-attention is a generalization of self-attention. As opposed to measure-
theoretic self-attention, we will call $f$ matrix self-attention in the
following. In order to study the Lipschitz constant of measure-theoretic self-
attention, we equip $\mathcal{P}_{c}(\mathbb{R}^{d})$ with the $p$-Wasserstein
distance $W_{p}$ for any $p\geq 1$. 111Note that this choice of distance is
more suited to the problem than the Hausdorff distance on point clouds: in
dimension 1, for example, the Hausdorff distance between the point clouds
$(1,1,0)$ and $(1,0)$ is zero, but self-attention does not map these point
clouds to the same point cloud.

##  3 Estimating the Lipschitz Constant of Self-Attention

Let us start with some preliminaries on (local) Lipschitz constants.

###  3.1 Lipschitz Constant and Local Lipschitz Constant

We first introduce the notions of Lipschitz constant and Lipschitz ratio, in a
general framework that is suited for measure-theoretic self-attention as well
as for matrix self-attention.

######  Definition 3.1 (Lipschitz constant, Lipschitz ratio).

Let
$\varphi\colon(\mathcal{E},d_{\mathcal{E}})\to(\mathcal{F},d_{\mathcal{F}})$
be a map between two metric spaces. The Lipschitz constant of $\varphi$ is
defined as

$\operatorname{\mathrm{Lip}}(\varphi)\coloneqq\sup_{\underset{x\neq
y}{x,y\in\mathcal{E}}}\frac{d_{\mathcal{F}}(\varphi(x),\varphi(y))}{d_{\mathcal{E}}(x,y)}.$

When $\mathcal{E}=\mathcal{F}$ and $d_{\mathcal{E}}=d_{\mathcal{F}}$, we may
denote
$\operatorname{\mathrm{Lip}}^{d_{\mathcal{E}}}(\varphi)\coloneqq\operatorname{\mathrm{Lip}}(\varphi)$
to stress on the choice of distance. The Lipschitz ratio associated to a pair
of distinct points $x,y\in\mathcal{E}$ is

$\frac{d_{\mathcal{F}}(\varphi(x),\varphi(y))}{d_{\mathcal{E}}(x,y)}.$

Bounding the Lipschitz constant of a function
$\varphi\colon(\mathcal{E},d_{\mathcal{E}})\to(\mathcal{F},d_{\mathcal{F}})$
gives control of how fast the output can change with a small perturbation of
the input. Sometimes, however, the function $\varphi$ is not globally
Lipschitz continuous: its Lipschitz constant is infinite. In that case, it is
more informative to look at the local Lipschitz constant, defined in the
following way.

######  Definition 3.2 (Local Lipschitz constant).

Let $\varphi\colon\mathcal{E}\to\mathcal{F}$ be a map between two metric
spaces. We define the local Lipschitz constant of $\varphi$ at point
$x\in\mathcal{E}$ as

$\operatorname{\mathrm{Lip}}_{x}(\varphi)\coloneqq\lim_{\varepsilon\to
0^{+}}\operatorname{\mathrm{Lip}}\varphi_{\lvert B(x,\varepsilon)}.$

The limit exists, as $\operatorname{\mathrm{Lip}}\varphi_{\lvert
B(x,\varepsilon)}$ decreases with $\varepsilon$.

Definition 3.2 is interesting, as it captures more information than the global
Lipschitz constant. More precisely, we have the following connection between
the two notions.

######  Lemma 3.1.

Let $\varphi\colon\mathcal{E}\to\mathcal{F}$ be a map between two metric
spaces. We have

$\operatorname{\mathrm{Lip}}(\varphi)\geq\sup_{x\in\mathcal{E}}\operatorname{\mathrm{Lip}}_{x}(\varphi).$

Assume moreover that the space $\mathcal{E}$ admits geodesics (see Appendix A
for the definition of a geodesic), which is the case for $\mathbb{R}^{d}$ and
$\mathcal{P}_{2}(\mathbb{R}^{d})$ equipped with $W_{2}$. Then, this inequality
becomes an equality.

See Appendix B.1 for the proof of Lemma 3.1. It is also important to notice
that bounding the local Lipschitz constant of $\varphi$ is equivalent to
bounding its Lipschitz constant on all possible compact sets
$K\subset\mathcal{E}$. We will therefore focus on proving Lipschitz bounds on
closed balls of the form $B_{R}$. Finally, let us point out that the local
Lipschitz constant has a more explicit formula in the particular case where
$\varphi$ is a differentiable function defined between normed vector spaces,
which is the case for matrix self-attention.

######  Lemma 3.2 ([55]).

Let $\varphi\colon\mathbb{R}^{d_{1}}\to\mathbb{R}^{d_{2}}$ be a differentiable
map. For all $x\in\mathbb{R}^{d_{1}}$, the local Lipschitz constant of
$\varphi$ at $x$ is equal to $\left\lVert D_{x}\varphi\right\rVert_{2}$, where
$D_{x}\varphi$ is the differential of $\varphi$ at $x$, and
$\left\lVert\cdot\right\rVert_{2}$ is the operator norm induced by the
Euclidean norm. Moreover, if $D_{x}\varphi$ is non zero and $u$ is its first
right singular vector, we have

$\frac{\left\lvert\varphi(x+\varepsilon
u)-\varphi(x)\right\rvert}{\varepsilon\left\lvert
u\right\rvert}\underset{\varepsilon\to 0^{+}}{\rightarrow}\left\lVert
D_{x}\varphi\right\rVert_{2}.$ (6)

######  Remark 3.3.

We are not aware of any result that would express the local Lipschitz constant
of measure-theoretic self-attention $F$ in terms of a notion of differential
of $F$, as Lemma 3.2 does for matrix self-attention. This seems an interesting
perspective for future work.

Let us now compare Lipschitz constants of matrix self-attention and measure-
theoretic self-attention.

####  3.1.1 Local Lipschitz Constant on Probability Measures versus on
Matrices

Let $p\geq 1$. To study the Lipschitz properties of self-attention, we use the
following norm:

$\left\lVert X\right\rVert_{F,p}\coloneqq\left(\sum_{i=1}^{n}\left\lvert
x_{i}\right\rvert^{p}\right)^{1/p}$

for $X=(x_{1},\dots,x_{n})^{\top}$, as [26] do, while we equip the set
$\mathcal{P}_{c}(\mathbb{R}^{d})$, on which $F$ operates, with the
$p$-Wasserstein distance $W_{p}$. We have the following link between the
Lipschitz constants of $f$ and $F$.

######  Lemma 3.3.

Let $p\geq 1$. For any matrix $X\in\mathbb{R}^{n\times d}$, we have

$\operatorname{\mathrm{Lip}}_{X}^{\left\lVert\cdot\right\rVert_{F,p}}(f)=\operatorname{\mathrm{Lip}}_{m(X)}^{W_{p}}(F_{\lvert\mathcal{M}_{n}(\mathbb{R}^{d})})\leq\operatorname{\mathrm{Lip}}_{m(X)}^{W_{p}}(F).$

The proof can be found in Appendix B.2. The main idea behind this statement is
that if two matrices $X$ and $Y$ are close enough (in the sense of the
$\left\lVert\cdot\right\rVert_{F,p}$ norm), then we have $\left\lVert
X-Y\right\rVert_{F,p}=W_{p}(m(X),m(Y))$. Note however that

$\operatorname{\mathrm{Lip}}_{X}^{\left\lVert\cdot\right\rVert_{F,p}}(f)\neq\operatorname{\mathrm{Lip}}_{m(X)}^{W_{p}}(F)$

in general, as we will see in Section 4.

###  3.2 Lower Bound on the Lipschitz Constant of Self-Attention

We have the following lower bound on the Lipschitz constant of measure-
theoretic self-attention.

######  Theorem 3.4.

Let $p\geq 1$, and $R>0$. The self-attention map

$F\colon\mu\mapsto(\Gamma_{\mu})_{\sharp}\mu$

is $W_{p}$-Lipschitz continuous on the space $\mathcal{P}(B_{R})$ of
probability measures supported in $B_{R}$. Moreover, assume that:

  1. (i)

$k=d$ and $V=I_{d}$,

  2. (ii)

$A\coloneqq\frac{1}{\sqrt{d}}K^{\top}Q$ is symmetric,

and denote $\gamma_{1}\geq\dots\geq\gamma_{d}$ the ordered eigenvalues of $A$.
Then, the following claims hold true.

  1. 1.

If $\gamma_{1}\geq-8\gamma_{d}$, there exists a function
$\theta\colon[0,+\infty)\to[0,+\infty)$ such that
$\theta(R)\to_{R\to+\infty}1$ and:

$\operatorname{\mathrm{Lip}}^{W_{2}}(F_{\lvert\mathcal{P}(B_{R})})\geq\theta(R)\frac{C}{2}R^{2}e^{CR^{2}}$

with $C\coloneqq\frac{\gamma_{1}}{8}$.

  2. 2.

If $\gamma_{1}<-8\gamma_{d}$, there exists a function
$\theta\colon[0,+\infty)\to[0,+\infty)$ such that
$\theta(R)\to_{R\to+\infty}1$ and:

$\operatorname{\mathrm{Lip}}^{W_{2}}(F_{\lvert\mathcal{P}(B_{R})})\geq\theta(R)\frac{C^{\prime}}{2}R^{2}e^{C^{\prime}R^{2}}$

with $C^{\prime}\coloneqq\left\lvert\gamma_{d}\right\rvert$.

Moreover, in both cases, the right-hand side is equivalent when $R\to+\infty$
to the local Lipschitz constant of measure-theoretic self-attention at a well-
chosen pair of two-dirac probability measures, i.e. of the form
$p\delta_{x}+(1-p)\delta_{y}$.

The fact that $F$ is Lipschitz continuous on $\mathcal{P}(B_{R})$ comes from
the upper bound proven by [19], which tells us that

$\operatorname{\mathrm{Lip}}^{W_{2}}(F_{\lvert\mathcal{P}(B_{R})})\leq
2R^{2}\left\lVert V\right\rVert_{2}\left\lVert
A\right\rVert_{2}\left(1+e^{2R^{2}\left\lVert A\right\rVert_{2}}\right).$

Note that the lower bounds given in Theorem 3.4 are much closer to this upper
bound than to the lower bound obtained by [26], which grows like $R^{2}$ up to
a constant factor.

Assumption $(i)$ amounts to studying self-attention before multiplication by
$V$, as [26] do. In other words, the upper and lower bounds that we have on
the Lipschitz constant of the function
$\tilde{F}\colon\mu\mapsto(\tilde{\Gamma}_{\mu})_{\sharp}\mu$ with

$\tilde{\Gamma}_{\mu}(x)\coloneqq\frac{\int\exp\left(\frac{1}{\sqrt{k}}x^{\top}A^{\top}y\right)y\,\mathrm{d}\mu(y)}{\int\exp\left(\frac{1}{\sqrt{k}}x^{\top}A^{\top}y\right)\mathrm{d}\mu(y)}$

are nearly sharp. The situation is more complicated with a general matrix $V$,
as the matrices $V$ and $A$ may not interact properly: to take an extreme
example, probability measures leading to a high local Lipschitz constant for
$\tilde{F}$ may for instance have a support included in the kernel of $V$.
Moreover, our numerical experiments suggest that assumption $(ii)$ is too
strong: requiring the diagonalizability of $A$ instead seems to be enough for
the result to hold.

From the perspective of a robustness analysis, having a more than exponential
explosion of the Lipschitz constant with the radius is bad news. Indeed, to
ensure for example that the Lipschitz constant stays smaller than 1 on a ball
of radius 20 – which is a realistic value for the case of a trained Vision
Transformer [16] on CIFAR-10 [27] for example – with $V=I_{d}$, one should
take $A$ with an operator norm of the order of $10^{-2}$, which would shrink
all keys and queries to very small values, thus losing precision in the model.
In future work, it would therefore be interesting to find a relevant
constraint on the input space, that provides better guarantees on the
Lipschitzianity of self-attention.

Let us now precise which two-dirac measures have a local Lipschitz constant
equal to the lower bound. The following result implies Theorem 3.4.

######  Proposition 3.5.

Let $R>0$. Denote $\mathcal{P}^{2\mathrm{d}}(\mathbb{R}^{d})$ the set of two-
dirac probability measures. Assume that

  1. (i)

$k=d$ and $V=I_{d}$,

  2. (ii)

$A\coloneqq\frac{1}{\sqrt{d}}K^{\top}Q$ is symmetric.

Denote $\gamma_{1}\geq\dots\geq\gamma_{d}$ the eigenvalues of $A$, and
$u_{1},\dots,u_{d}$ corresponding unit eigenvectors.

  1. 1.

If $\gamma_{1}\geq-8\gamma_{d}$, define

$p_{R}\coloneqq e^{-2CR^{2}}\mbox{ and }\mu_{R}\coloneqq
p_{R}\delta_{Ru_{1}}+(1-p_{R})\delta_{\frac{R}{2}u_{1}}$

for $R>0$, with $C\coloneqq\frac{\gamma_{1}}{8}$. Then

$\operatorname{\mathrm{Lip}}_{\mu_{R}}^{W_{2}}(F)\geq\operatorname{\mathrm{Lip}}_{\mu_{R}}^{W_{2}}(F_{\lvert\mathcal{P}^{2\mathrm{d}}(\mathbb{R}^{d})})\sim_{R\to+\infty}\frac{C}{2}R^{2}e^{CR^{2}}.$

  2. 2.

If $\gamma_{1}<-8\gamma_{d}$, define

$p_{R}\coloneqq e^{-2C^{\prime}R^{2}}\mbox{ and }\mu_{R}\coloneqq
p_{R}\delta_{Ru_{d}}+(1-p_{R})\delta_{-Ru_{d}}$

for $R>0$, with $C^{\prime}\coloneqq\left\lvert\gamma_{d}\right\rvert$. Then

$\operatorname{\mathrm{Lip}}_{\mu_{R}}^{W_{2}}(F)\geq\operatorname{\mathrm{Lip}}_{\mu_{R}}^{W_{2}}(F_{\lvert\mathcal{P}^{2\mathrm{d}}(\mathbb{R}^{d})})\sim_{R\to+\infty}\frac{C^{\prime}}{2}R^{2}e^{C^{\prime}R^{2}}.$

The proof can be found in Appendix B.5. Note that Theorem 3.4 is an immediate
consequence of Proposition 3.5, using Lemma 3.1. The condition
$\gamma_{1}\geq-8\gamma_{d}$ can be interpreted as follows. Negative
eigenvalues induce a faster growth of the local Lipschitz constant than
positive eigenvalues: if $\gamma_{i}$ is a negative eigenvalue of $A$
associated to a unit eigenvector $u_{i}$, choosing

$p_{R}\coloneqq e^{-2\left\lvert\gamma_{i}\right\rvert R^{2}}\mbox{ and
}\mu_{R}\coloneqq p_{R}\delta_{Ru_{i}}+(1-p_{R})\delta_{-Ru_{i}}$

gives

$\operatorname{\mathrm{Lip}}_{\mu_{R}}^{W_{2}}(F_{\lvert\mathcal{P}^{2\mathrm{d}}(\mathbb{R}^{d})})\sim_{R\to+\infty}\frac{\left\lvert\gamma_{i}\right\rvert}{2}R^{2}e^{\left\lvert\gamma_{i}\right\rvert
R^{2}},$ (7)

whereas if $\gamma_{i}$ is a positive eigenvalue of $A$ associated to a unit
eigenvector $u_{i}$, choosing

$p_{R}\coloneqq e^{-\frac{\gamma_{i}}{8}R^{2}}\mbox{ and }\mu_{R}\coloneqq
p_{R}\delta_{Ru_{i}}+(1-p_{R})\delta_{\frac{R}{2}u_{i}}$

gives

$\operatorname{\mathrm{Lip}}_{\mu_{R}}^{W_{2}}(F_{\lvert\mathcal{P}^{2\mathrm{d}}(\mathbb{R}^{d})})\sim_{R\to+\infty}\frac{\gamma_{i}}{16}R^{2}e^{\frac{\gamma_{i}}{8}R^{2}}.$
(8)

We can see that when $\gamma_{1}>-8\gamma_{d}$, either $A$ has only positive
eigenvalues, or Equation (8) with $\gamma_{1}$ provides a faster growth of the
local Lipschitz constant than Equation (7) with $\gamma_{d}$. On the contrary,
when $\gamma_{1}<-8\gamma_{d}$, either $A$ has only negative eigenvalues, or
Equation (7) provides a better rate. Finally, in the edge case where
$\gamma_{1}=-8\gamma_{d}$, both equations grow at the same rate. From a high-
level perspective, the optimal strategy to make the local Lipschitz constant
explode is therefore to concentrate most of the mass on one point that goes to
infinity and is aligned with a well-chosen eigenvector of $A$, and to keep the
exponentially vanishing complement of mass on another point that also goes to
infinity, either twice faster in the same direction (if
$\gamma_{1}\geq-8\gamma_{d}$), or at the same speed in the opposite direction
(if $\gamma_{1}<-8\gamma_{d}$). This strategy differs from the one of [26],
who have one particle stuck at 0, and the others maximizing their empirical
variance, which leads to a local Lipschitz constant that is quadratic with the
radius (see Lemma C.4).

###  3.3 Tightness of the Lower Bound

To analyse the tightness of the lower bound given in Theorem 3.4, let us look
for the probability measure supported in $\mathcal{P}(B_{R})$ that induces the
highest local Lipschitz constant, and tackle this problem numerically. We
focus on probability measures with finite support:

$\mathcal{P}^{\mathrm{fin}}\coloneqq\\{\sum_{i=1}^{n}a_{i}\delta_{x_{i}}\mid
n\in\mathbb{N}^{*},\ x_{1},\dots,x_{n}\in\mathbb{R}^{d}\mbox{ and
}(a_{1},\dots,a_{n})\in\Sigma_{n}\\},$

with
$\Sigma_{n}:=\\{(a_{1},\dots,a_{n})\in\mathbb{R}_{+}^{n}:\sum_{i=1}^{n}a_{i}=1\\}$
the simplex of $\mathbb{R}^{n}$. This is still very general, as continuous
measures can be approximated arbitrarily well by elements of the set
$\mathcal{P}^{\mathrm{fin}}$, and Transformers can only handle a finite number
of points in practice. As we cannot compute numerically the local Lipschitz
constant of $F_{\lvert\mathcal{P}^{\mathrm{fin}}}$ at a measure
$\sum_{i=1}^{n}a_{i}\delta_{x_{i}}$, let us use a lower bound instead, by
introducing a generalization of matrix self-attention, called weighted self-
attention.

####  3.3.1 Weighted Self-Attention

Weighted self-attention implements a matrix representation of probability
measures of the form $\sum_{i=1}^{n}a_{i}\delta_{x_{i}}$, just as matrix self-
attention does for empirical measures. It is defined by the following formula.

######  Definition 3.4 (Weighted self-attention).

For any vector $a\in\Sigma_{n}$, denote

$\mathcal{P}_{a}(\mathbb{R}^{d})\coloneqq\\{\sum_{i=1}^{n}a_{i}\delta_{x_{i}}\mid
x_{1},\dots,x_{n}\in\mathbb{R}^{d}\\}.$

We define the matrix version of the restriction of self-attention to
$\mathcal{P}_{a}(\mathbb{R}^{d})$ in the following way:

$f_{a}\colon X\in\mathbb{R}^{n\times
d}\mapsto\operatorname{\mathrm{softmax}}_{a}(XQ^{\top}KX^{\top})XV^{\top},$

where

$\operatorname{\mathrm{softmax}}_{a}(w_{1},\dots,w_{n})\coloneqq\left(\frac{a_{i}e^{w_{i}}}{\sum_{j}a_{j}e^{w_{j}}}\right)_{1\leq
i\leq n}.$

The function $f_{a}$ is called weighted self-attention associated to the
coefficients $a$.

Definition 3.4 is designed so that for any matrix $X\in\mathbb{R}^{n\times
d}$, it holds

$F(m_{a}(X))=m_{a}(f_{a}(X)),$

with

$m_{a}(X)\coloneqq\sum_{i=1}^{n}a_{i}\delta_{X_{i}}.$

Weighted self-attention proves very useful for numerical experiments: it
provides a representation that is much more convenient than the traditional
matrix setting. For instance, measures of the form
$\mu_{R}=e^{-2R^{2}}\delta_{R}+(1-e^{-2R^{2}})\delta_{-R}$, which are of
interest in our work, require a very large number of samples $n$ to be well
approximated by an empirical measure, especially when $R$ is large – in fact,
we need $n$ to be of the order of $e^{2R^{2}}$ to put the right amount of mass
on the first dirac, whereas weighted self-attention allows us to represent
$\mu_{R}$ with just two vectors, coupled with a vector of coefficients
$a=(e^{-2R^{2}},1-e^{-2R^{2}})$. To study the Lipschitz continuity of $f_{a}$,
we equip the space $\mathbb{R}^{n\times d}$ with a new scalar product
$\left\langle\cdot,\cdot\right\rangle_{a}$, defined as follows:

$\left\langle X,Y\right\rangle_{a}\coloneqq\sum_{i=1}^{n}a_{i}\left\langle
x_{i},y_{i}\right\rangle,$

where $X=(x_{1},\dots,x_{n})^{\top}$ and $Y=(y_{1},\dots,y_{n})^{\top}$. This
allows us to connect the local Lipschitz constant of $f_{a}$ for the norm
$\left\lVert\cdot\right\rVert_{a}$ associated to the scalar product
$\left\langle\cdot,\cdot\right\rangle_{a}$ to the local Lipschitz constant of
measure-theoretic self-attention in the Wasserstein 2 sense.

######  Lemma 3.6.

Let $X\in\mathbb{R}^{n\times d}$ be an input matrix, and $a\in\Sigma_{n}$ a
vector of coefficients. Then, we have

$\operatorname{\mathrm{Lip}}_{m_{a}(X)}^{W_{2}}F_{\lvert\mathcal{P}_{a}(\mathbb{R}^{d})}=\left\lVert
D_{X}^{a}f_{a}\right\rVert_{2,a},$

where $D^{a}$ is the Jacobian in the space $(\mathbb{R}^{n\times
d},\left\langle\cdot,\cdot\right\rangle_{a})$ and
$\left\lVert\cdot\right\rVert_{2,a}$ is the corresponding operator norm.

This is a nice property, as we can compute numerically the local Lipschitz
constant of $f_{a}$, just as for matrix self-attention. Lemma 3.6 can be
proven with the same steps as for Lemma 3.2. Now that we are equipped with
weighted self-attention, we can perform a gradient ascent on its local
Lipschitz constant, in order to estimate the Lipschitz constant of $F$
restricted to $\mathcal{P}(B_{R})$.

####  3.3.2 Experiments

To gain some intuition about which probability measures with finite support
have a high local Lipschitz constant, we perform a joint gradient ascent on
$X$ and $a$ to maximize the local Lipschitz constant $\left\lVert
D_{X}^{a}f_{a}\right\rVert_{2,a}$ of $f_{a}$. More precisely, we do a
projected gradient ascent on $X$, to ensure that all lines of $X$ stay in the
ball $B_{R}$, and a mirror ascent on $a$, because $a$ has to belong to the
simplex $\Sigma_{n}$ at each step. We also choose $R=1$ for the experiments.
So the iterations are

$\displaystyle X_{0}=\pi(Z)\mbox{ with }Z_{ij}\sim\mathcal{N}(0,1)\mbox{ for
all }i,j\ \ \mbox{ and }\ \ a_{0}=\frac{1}{n}1_{n},$ (9) $\displaystyle
X_{k+1}=\pi\left(X_{k}-\eta\nabla_{X_{k}}\left\lVert
D_{X_{k}}^{a_{k}}f_{a_{k}}\right\rVert_{2,{a_{k}}}\right),$ (10)
$\displaystyle\tilde{a}_{k+1}=a_{k}\odot\exp\left(\eta\nabla_{a_{k}}\left\lVert
D_{X_{k+1}}^{a_{k}}f_{a_{k}}\right\rVert_{2,{a_{k}}}\right),$ (11)
$\displaystyle a_{k+1}=\tilde{a}_{k+1}/\sum_{i=1}^{n}(\tilde{a}_{k+1})_{i},$
(12)

where $\pi$ is the projection on the closed unit ball
$\,\overline{\\!{B}}(0,1)$ and $\odot$ denotes coordinatewise multiplication.
We use autodifferentiation in JAX [6]. Numerics show that the optimization
problem is non-convex: the algorithm can get stuck in local minima.
Nevertheless, an interesting dynamics of clustering of the particles can be
observed along the gradient ascent, as shown in Figure 1. This suggests that
there are measures with a high local Lipschitz constant that are made of only
a few diracs – which confirms our theoretical results, where the highest
Lipschitz constant is obtained with two-dirac measures. With this in mind, one
can focus on two-dirac measures, and perform a grid search on the position of
diracs and their weights to find the two-dirac measure supported in $B_{R}$
with the highest local Lipschitz constant. Doing this, we recover a measure
that is very close to $\mu_{R}$, defined in Proposition 3.5.

Figure 1: Evolution of the input (left column) and output (right column)
measures along a gradient ascent on the local Lipschitz constant of weighted
self-attention (see Equation (9)), with $Q^{\top}K=-I_{2}$ and $V=I_{2}$.
Darker points concentrate more mass: we set a transparency parameter equal to
$\sqrt{a_{i}}$ for each point $x_{i}$. We can see two distinct clusters
emerge, diametrally opposed, with one of them concentrating most of the mass.
We also plot the first right singular vector of the Jacobian of self-
attention, that shows in which direction to move input points to obtain the
biggest change in the output.

Remarkably, the experiments are much more difficult to interpret if, instead
of doing a gradient ascent on the local Lipschitz constant of only one point
cloud, one performs a gradient ascent on the Lipschitz ratio of two point
clouds, by making both point clouds vary at each iteration while constraining
them to stay in the ball $\,\overline{\\!{B}}(0,1)$. Should the optimization
converge to the global maximum, this second method would give a good
approximation of the Lipschitz constant of self-attention restricted to
empirical measures supported in $\,\overline{\\!{B}}(0,1)$ and with a fixed
number of diracs. Unfortunately, we observe that iterations systematically get
trapped in local minima, and the particles do not cluster.

###  3.4 The Lipschitz Constant of Masked Attention

Let us conclude this section with an analysis of the smoothness of _masked_
self-attention. We focus on the residual form of masked self-attention,
defined for $X=(x_{1},\dots,x_{n})^{\top}\in\mathbb{R}^{n\times d}$ as

$f^{m}(X)_{i}\coloneqq
x_{i}+\operatorname{\mathrm{softmax}}\left(\frac{1}{\sqrt{k}}x_{i}Q^{\top}KX_{1:i}^{\top}\right)X_{1:i}V^{\top}$
(13)

for all $1\leq i\leq n$, where
$f^{m}(X)\eqqcolon(f^{m}(X)_{1},\dots,f^{m}(X)_{n})^{\top}$ and
$X_{1:i}\coloneqq(x_{1},\dots,x_{i})^{\top}$. Residual masked self-attention,
which is a central building block of the architecture of Decoders [58],
processes inputs in a sequential way, so it is not permutation equivariant. As
a consequence, the map $f^{m}$ does not directly induce a map on empirical
measures as for traditional self-attention. Still, we can introduce a
convenient measure-theoretic framework for studying masked self-attention, in
the following way. Instead of viewing an input
$(x_{1},\dots,x_{n})\in(\mathbb{R}^{d})^{n}$ as the empirical measure
$\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}}$, let us add a coordinate
$s_{i}\in[0,1]$ — the choice of the interval is arbitrary — to each input
vector, to encode the position of the $x_{i}$. We can then represent
$(x_{1},\dots,x_{n})$ as the empirical measure over
$[0,1]\times\mathbb{R}^{d}$

$\bar{\mu}=\frac{1}{n}\sum_{i=1}^{n}\delta_{(s_{i},x_{i})},$

with $0\leq s_{1}<\dots<s_{n}\leq 1$ any choice of increasing real numbers
between 0 and 1, for example $s_{i}=i/n$ for $i=1,\dots,n$. Masked self-
attention can be computed from the measure $\bar{\mu}$ only, where the
position of the vectors $x_{i}$ is recovered by looking at the order of the
first coordinates $s_{i}$. To extend this idea to more general probability
measures, we define the measure-theoretic generalization of residual masked
self-attention as taking as input probability measures on the product space
$[0,1]\times\mathbb{R}^{d}$.

######  Definition 3.5 (Measure-theoretic masked residual self-attention).

For any probability measure $\bar{\mu}\in{\cal
P}_{c}([0,1]\times\mathbb{R}^{d})$, denote $\mu$ the second marginal of
$\bar{\mu}$, i.e.

$\mu(A)\coloneqq\int_{0}^{1}\int_{A}\mathrm{d}\bar{\mu}(s,x)$

for all Borel sets $A\subset\mathbb{R}^{d}$. We define residual masked self-
attention on ${\cal P}_{c}([0,1]\times\mathbb{R}^{d})$ as

$F^{m}\colon\bar{\mu}\mapsto\left(\mathrm{Id}+\Gamma_{\bar{\mu}}\right)_{\sharp}\bar{\mu}$

with

$\Gamma_{\bar{\mu}}(s,x)\coloneqq\left(0,\frac{\int_{\mathbb{R}^{d}}G(x,y)Vy\mathbf{1}_{\tau\leq
s}\mathrm{d}\bar{\mu}(\tau,y)}{\int_{\mathbb{R}^{d}}G(x,y)\mathbf{1}_{\tau\leq
s}\mathrm{d}\bar{\mu}(\tau,y)}\right).$

This definition generalizes Equation (13) in the following sense: denoting
$\mathrm{ord}$ the transformation

$\mathrm{ord}\colon X=(x_{1},\dots,x_{n})^{\top}\in\mathbb{R}^{n\times
d}\mapsto\frac{1}{n}\sum_{i=1}^{n}\delta_{(i/n,x_{i})}\in{\cal
P}_{c}([0,1]\times\mathbb{R}^{d}),$

we have

$F^{m}(\mathrm{ord}(X))=\mathrm{ord}(f^{m}(X))$

for all $X\in\mathbb{R}^{n\times d}$.

######  Remark 3.6.

For empirical measures, $\mathrm{ord}$ gives a canonical way of representing
$\mu:=1/n\sum_{i=1}^{n}\delta_{x_{i}}$ by $\bar{\mu}\in{\cal
P}_{c}([0,1]\times\mathbb{R}^{d})$, via $X=(x_{1},\dots,x_{n})^{\top}$. In the
case where $\mu$ has infinitely many diracs, we simply assume that the
coupling $\bar{\mu}$ is already given as input.

Let us now bound from above the Lipschitz constant of measure-theoretic masked
self-attention. To this aim, we introduce the following distance on ${\cal
P}_{c}([0,1]\times\mathbb{R}^{d})$, for $p\geq 1$.

######  Definition 3.7.

Let $\bar{\mu}$ and $\bar{\nu}$ be two probability measures in ${\cal
P}_{c}([0,1]\times\mathbb{R}^{d})$, and $p\geq 1$. If $\bar{\mu}$ and
$\bar{\nu}$ have the same marginal with respect to $s$, i.e.

$\int_{s_{1}}^{s_{2}}\int_{\mathbb{R}^{d}}\mathrm{d}\bar{\mu}(s,x)=\int_{s_{1}}^{s_{2}}\int_{\mathbb{R}^{d}}\mathrm{d}\bar{\nu}(s,x)$

for all $0\leq s_{1}<s_{2}\leq 1$, denote $\theta$ this marginal distribution,
and write with the disintegration theorem [5]

$\mathrm{d}\bar{\mu}(\tau,x)\eqqcolon\mathrm{d}\theta(\tau)\mathrm{d}\mu^{\tau}(x)$

and

$\mathrm{d}\bar{\nu}(\tau,x)\eqqcolon\mathrm{d}\theta(\tau)\mathrm{d}\nu^{\tau}(x).$

The measures $\mu^{\tau}$ and $\nu^{\tau}$ can be seen intuitively as the
restriction of $\mu$ and $\nu$ to the mass located at position $\tau$,
rescaled to obtain probability measures. We then measure the distance between
$\bar{\mu}$ and $\bar{\nu}$ with

$d_{p}(\bar{\mu},\bar{\nu})\coloneqq\left(\int_{0}^{1}W_{p}(\mu^{\tau},\nu^{\tau})^{p}\mathrm{d}\theta(\tau)\right)^{1/p}.$

If $\bar{\mu}$ and $\bar{\nu}$ do not have the same first marginal, we set

$d_{p}(\bar{\mu},\bar{\nu})\coloneqq+\infty.$

The distance $d_{p}$ derives from the conditional optimal transport problem,
studied by [23]. We have the following bound on the $d_{p}$-Lipschitz constant
of masked self-attention.

######  Proposition 3.7.

Let $R>0$ and $p\geq 1$. The measure-theoretic residual masked self-attention
map $F^{m}$ is Lipschitz continuous on the space of measures supported in
$[0,1]\times B_{R}$, with a Lipschitz constant upper-bounded by

$1+\left\lVert V\right\rVert_{2}(1+3\left\lVert
A\right\rVert_{2}R^{2})e^{2\left\lVert A\right\rVert_{2}R^{2}}.$

Here, we consider that two measures with different first marginals induce an
infinite Lipschitz ratio.

The proof can be found in Appendix B.3. This result shows that similar upper
bounds hold for traditional self-attention and for masked self-attention – in
particular, the dependency in the radius is more than exponential.

We have seen that the local Lipschitz constant of self-attention grows
extremely fast when increasing the support, and that it seems to be the case
for masked self-attention as well. Another interesting perspective is to
understand which local perturbations induce the largest variations of the
output of $F$. In other words, in which direction should one perturb a
probability measure $\mu$ to induce a maximal variation of the output of $F$?
This question brings into play the notion of mass splitting, that is the
object of the following section.

##  4 Mass Splitting

Let $a\in\Sigma_{n}$ be a vector of coefficients, and consider a measure

$\mu=\sum_{i=1}^{n}a_{i}\delta_{x_{i}}\in\mathcal{P}_{a}(\mathbb{R}^{d}).$

Does it hold that

$\operatorname{\mathrm{Lip}}_{\mu}^{W_{2}}(F)=^{?}\operatorname{\mathrm{Lip}}_{\mu}^{W_{2}}\left(F_{\lvert\mathcal{P}_{a}(\mathbb{R}^{d})}\right)?$
(14)

If this equality holds true, it means intuitively222The following claim is not
rigorously true, but gives the right intuition. that among all probability
measures in $B^{W_{2}}(\mu,\varepsilon)$, the one (or one of those) that
induces the highest Lipschitz ratio with respect to $\mu$ is itself in
$\mathcal{P}_{a}(\mathbb{R}^{d})$, for $\varepsilon>0$ small enough. In other
words, by slightly moving the diracs of $\mu$ in the right direction, one can
perturb the output of $F$ in the worst way possible among all perturbations
that stay in the ball $B^{W_{2}}(\mu,\varepsilon)$. As we will see in what
follows, Equation (14) is not always true. We call this phenomenon _mass
splitting_.

######  Definition 4.1 (Mass splitting, magnitude).

Let $F\colon\mathcal{U}\to\mathcal{F}$, where $\mathcal{F}$ is a metric space,
and $\mathcal{U}$ is an open set of the space of probability measures
$\mathcal{P}(\mathbb{R}^{d})$. For a probability measure
$\mu=\sum_{i=1}^{n}a_{i}\delta_{x_{i}}\in\mathcal{U}$ with finite support, we
say that $F$ splits the mass at $\mu$, or equivalently that $\mu$ is splitted
by $F$, if

$\operatorname{\mathrm{Lip}}_{\mu}^{W_{2}}(F)>\operatorname{\mathrm{Lip}}_{\mu}^{W_{2}}\left(F_{\lvert\mathcal{P}_{a}(\mathbb{R}^{d})}\right),$

where
$\mathcal{P}_{a}(\mathbb{R}^{d})\coloneqq\\{\sum_{i=1}^{n}a_{i}\delta_{x_{i}}\mid
x_{1},\dots,x_{n}\in\mathbb{R}^{d}\\}$. The gap
$\operatorname{\mathrm{Lip}}_{\mu}^{W_{2}}(F)-\operatorname{\mathrm{Lip}}_{\mu}^{W_{2}}\left(F_{\lvert\mathcal{P}_{a}(\mathbb{R}^{d})}\right)>0$
is called magnitude of the splitting.

In what follows, we will take $F$ to be measure-theoretic self-attention, and
restrict our theoretical analysis to the case of empirical measures instead of
probability measures with finite support, for simplicity. Our aim is to
understand what kind of perturbations induce the highest change in the output
of $F$ at a splitted measure: do the diracs of $\mu$ split in a finite number
of points? Or on the contrary, is the worst perturbation continuous with
respect to Lebesgue?

###  4.1 Finite Mass Splitting: a Proxy for Mass Splitting

We lack tools to analyze mass splitting as it is defined. Let us therefore
focus on the following sufficient condition for mass splitting, which is
easier to handle and serves as a proxy for studying mass splitting.

######  Definition 4.2 (Finite mass splitting).

Let $F\colon\mathcal{U}\to\mathcal{F}$, where $\mathcal{F}$ is a metric space,
and $\mathcal{U}$ is an open set of the space of probability measures
$\mathcal{P}(\mathbb{R}^{d})$. For an empirical measure
$\mu=1/n\sum_{i=1}^{n}\delta_{x_{i}}\in\mathcal{U}$, we say that $F$ finitely
splits the mass at $\mu$, or equivalently that $\mu$ is finitely splitted by
$F$, if

$\sup_{N\geq
1}\operatorname{\mathrm{Lip}}^{W_{2}}\left(F_{\lvert\mathcal{M}_{nN}(\mathbb{R}^{d})}\right)>\operatorname{\mathrm{Lip}}^{W_{2}}\left(F_{\lvert\mathcal{M}_{n}(\mathbb{R}^{d})}\right).$

It is straightforward that if $\mu$ is finitely splitted, then $\mu$ is
splitted. Let us now focus on the specific case of self-attention. We have the
following equivalent formulation of finite mass splitting, which is a direct
consequence of Lemma 3.2.

######  Lemma 4.1.

Let $F$ be measure-theoretic self-attention. Let $X\in\mathbb{R}^{d}$. Then
$F$ finitely splits the empirical measure $m(X)$ if and only if we have

$\sup_{N\geq 1}\left\lVert D_{1_{N}\otimes X}f\right\rVert_{2}>\left\lVert
D_{X}f\right\rVert_{2}.$

####  4.1.1 Duplication Is the Worst Case for Finite Mass Splitting

The following result shows that to reach the magnitude of finite mass
splitting at an empirical measure $m(X)$ with $X\in\mathbb{R}^{n\times d}$, it
suffices to split only one well-chosen dirac of $\mu$ in just two equal parts.

######  Theorem 4.2.

Let $X=(x_{1},\dots,x_{n})^{\top}$ be a $n\times d$ matrix, and denote
$Y^{N}\coloneqq 1_{N}\otimes X$ the $N$-replication of $X$ for any positive
integer $N$. Assume that $V=I_{d}$. Then it holds

$\left\lVert D_{Y^{N}}f\right\rVert_{2}=\left\lVert
D_{Y^{2}}f\right\rVert_{2}$

for all $N\geq 2$. Furthermore, if the mass splits at $X$, then the maximal
magnitude of splitting at $X$ can be reached by splitting equally the mass of
the $i$-th particle of $X$ along the first singular vector of the matrix
$(\operatorname{\mathrm{Var}}^{(i)}Z)A$, with

$i\coloneqq\operatorname{\mathrm{argmax}}_{i^{\prime}\in\\{1,\dots,n\\}}\left\lVert(\operatorname{\mathrm{Var}}^{(i^{\prime})}Z)A\right\rVert_{2}^{2},$

where $\operatorname{\mathrm{Var}}^{(i)}Z$ is the variance of the random
variable $Z$ following the law $\mathbb{P}(Z=x_{j})\coloneqq P_{ij}$.

The proof of Theorem 4.2 can be found in Appendix C, and relies on a formula
for the Jacobian of $Y^{2}$ in terms of $D_{X}f$. Theorem 4.2 gives us a
simple way of checking numerically if a given empirical measure $m(X)$ is
finitely splitted.

###  4.2 Numerical Evidence of the Mass Splitting Phenomenon

Focusing on two-dirac measures, we compare

$\left\lVert D_{X}^{a}f\right\rVert_{2}\leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ \mbox{and}\leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ \left\lVert D_{1_{2}\otimes
X}^{a}f\right\rVert_{2}$

with $X=(x,y)$ for $(x,y)$ on a grid, and for various choices of coefficients
$a=(p,1-p)$. It turns out that splitted measures lie in well-defined areas, as
we can see on Figure 2 in dimension 1, and that there are infinitely many of
them – at least for suitable choices of parameters.

Figure 2: Visualisation of two-dirac measures splitted by self-attention. A
point $(x,y)$ belongs to the graph if the measure
$p\delta_{x}+(1-p)\delta_{y}$ is splitted by $f$, namely if $\left\lVert
D_{(x,y,x,y)^{\top}}^{a}f\right\rVert_{2}>\left\lVert
D_{(x,y)^{\top}}^{a}f\right\rVert_{2}$ with $a=(p,1-p)$. The colorbar
indicates the magnitude of the splitting. The first row corresponds to
$A=0.01$, the second one to $A=-0.01$, and $V=1$ in both cases. We make the
repartition of mass vary from 0.25 to 0.95, and observe that it can change the
magnitude of the splitting, and the shape of the zone containing the splitted
measures.

## Conclusion and Perspectives

We have shown that self-attention is a highly irregular function in the
Lipschitz sense, and more precisely that its Lipschitz constant on the space
of measures supported in $B_{R}$ grows like

$CR^{2}e^{CR^{2}}$

up to a factor close to $\frac{1}{2}$, where $C$ depends on the parameters of
self-attention. It is important to notice that the particular measures that
are responsible for the explosion of this local Lipschitz bound, that are two-
dirac measures with an extremely unbalanced display of mass, require a massive
number of samples to be well approximated by empirical measures. Therefore,
one can ask the question of deriving input length-dependent bounds –
hopefully, such bounds could grow significantly slower than our general bound.

Moreover, we have pointed out a mass splitting phenomenon for self-attention,
which opens promising perspectives for future work. Can we understand better
which type of functions are subject to mass splitting? Is finite mass
splitting equivalent to regular mass splitting? The link between mass
splitting and adversarial attacks seems also interesting to explore. The idea
of perturbing a duplication of the input brings into play a new family of
attacks, but these take place in the feature space, i.e. just before the
encoder, and not in the token space. Is it possible to invert the positional
encoding and the tokenization embedding to derive attacks in the token space?
From a different perspective, it could also be efficient to do virtual
adversarial training [39] to robustify the network against this kind of
attacks.

## References

  * [1] Cem Anil, James Lucas and Roger Grosse  “Sorting out Lipschitz function approximation”  In _International Conference on Machine Learning_ , 2019, pp. 291–301  PMLR
  * [2] Dzmitry Bahdanau, Kyunghyun Cho and Yoshua Bengio  “Neural machine translation by jointly learning to align and translate”  In _arXiv preprint arXiv:1409.0473_ , 2014
  * [3] Peter L Bartlett, Dylan J Foster and Matus J Telgarsky  “Spectrally-normalized margin bounds for neural networks”  In _Advances in neural information processing systems_ 30, 2017
  * [4] Jens Behrmann et al.  “Invertible residual networks”  In _International conference on machine learning_ , 2019, pp. 573–582  PMLR
  * [5] Vladimir Igorevich Bogachev and Maria Aparecida Soares Ruas  “Measure theory”  Springer, 2007
  * [6] James Bradbury et al.  “JAX: composable transformations of Python+NumPy programs”, 2018  URL: http://github.com/google/jax
  * [7] Tom Brown et al.  “Language models are few-shot learners”  In _Advances in neural information processing systems_ 33, 2020, pp. 1877–1901
  * [8] Nicholas Carlini and David Wagner  “Towards evaluating the robustness of neural networks”  In _2017 ieee symposium on security and privacy (sp)_ , 2017, pp. 39–57  Ieee
  * [9] Lili Chen et al.  “Decision transformer: Reinforcement learning via sequence modeling”  In _Advances in neural information processing systems_ 34, 2021, pp. 15084–15097
  * [10] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt and David K Duvenaud  “Neural ordinary differential equations”  In _Advances in neural information processing systems_ 31, 2018
  * [11] Ricky TQ Chen, Jens Behrmann, David K Duvenaud and Jörn-Henrik Jacobsen  “Residual flows for invertible generative modeling”  In _Advances in Neural Information Processing Systems_ 32, 2019
  * [12] Yu-An Chung et al.  “W2v-BERT: Combining Contrastive Learning and Masked Language Modeling for Self-Supervised Speech Pre-Training”, 2021  arXiv:2108.06209 [cs.LG]
  * [13] Moustapha Cisse et al.  “Parseval networks: Improving robustness to adversarial examples”  In _International conference on machine learning_ , 2017, pp. 854–863  PMLR
  * [14] Gwendoline De Bie, Gabriel Peyré and Marco Cuturi  “Stochastic deep networks”  In _International Conference on Machine Learning_ , 2019, pp. 1556–1565  PMLR
  * [15] Jacob Devlin, Ming-Wei Chang, Kenton Lee and Kristina Toutanova  “Bert: Pre-training of deep bidirectional transformers for language understanding”  In _arXiv:1810.04805_ , 2018
  * [16] Alexey Dosovitskiy et al.  “An image is worth 16x16 words: Transformers for image recognition at scale”  In _arXiv:2010.11929_ , 2020
  * [17] Mahyar Fazlyab et al.  “Efficient and accurate estimation of lipschitz constants for deep neural networks”  In _Advances in Neural Information Processing Systems_ 32, 2019
  * [18] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy and Philippe Rigollet  “A Mathematical Perspective on Transformers”
  * [19] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy and Philippe Rigollet  “The emergence of clusters in self-attention dynamics”  In _arXiv preprint arXiv:2305.05465_ , 2023
  * [20] Ian J Goodfellow, Jonathon Shlens and Christian Szegedy  “Explaining and harnessing adversarial examples”  In _arXiv preprint arXiv:1412.6572_ , 2014
  * [21] Kavya Gupta and Sagar Verma  “CertViT: Certified Robustness of Pre-Trained Vision Transformers”  In _arXiv preprint arXiv:2302.10287_ , 2023
  * [22] Matthias Hein and Maksym Andriushchenko  “Formal guarantees on the robustness of a classifier against adversarial manipulation”  In _Advances in neural information processing systems_ 30, 2017
  * [23] Bamdad Hosseini, Alexander W Hsu and Amirhossein Taghvaei  “Conditional Optimal Transport on Function Spaces”  In _arXiv preprint arXiv:2311.05672_ , 2023
  * [24] Michael Janner, Qiyang Li and Sergey Levine  “Offline reinforcement learning as one big sequence modeling problem”  In _Advances in neural information processing systems_ 34, 2021, pp. 1273–1286
  * [25] Xiaojun Jia et al.  “Revisiting and Exploring Efficient Fast Adversarial Training via LAW: Lipschitz Regularization and Auto Weight Averaging”  In _arXiv preprint arXiv:2308.11443_ , 2023
  * [26] Hyunjik Kim, George Papamakarios and Andriy Mnih  “The Lipschitz Constant of Self-Attention”, 2021  arXiv:2006.04710 [stat.ML]
  * [27] Alex Krizhevsky, Vinod Nair and Geoffrey Hinton  “The CIFAR-10 dataset”  In _online: http://www. cs. toronto. edu/kriz/cifar. html_ 55.5, 2014
  * [28] Alexey Kurakin, Ian Goodfellow and Samy Bengio  “Adversarial machine learning at scale”  In _arXiv preprint arXiv:1611.01236_ , 2016
  * [29] Fabian Latorre, Paul Rolland and Volkan Cevher  “Lipschitz constant estimation of neural networks via sparse polynomial optimization”  In _arXiv preprint arXiv:2004.08688_ , 2020
  * [30] Juho Lee et al.  “Set transformer: A framework for attention-based permutation-invariant neural networks”  In _International conference on machine learning_ , 2019, pp. 3744–3753  PMLR
  * [31] Klas Leino, Zifan Wang and Matt Fredrikson  “Globally-robust neural networks”  In _International Conference on Machine Learning_ , 2021, pp. 6212–6222  PMLR
  * [32] Mike Lewis et al.  “Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension”  In _arXiv:1910.13461_ , 2019
  * [33] Peter J Liu et al.  “Generating wikipedia by summarizing long sequences”  In _arXiv:1801.10198_ , 2018
  * [34] Yinhan Liu et al.  “Roberta: A robustly optimized bert pretraining approach”  In _arXiv:1907.11692_ , 2019
  * [35] Yiping Lu et al.  “Understanding and improving transformer from a multi-particle dynamic system point of view”  In _arXiv preprint arXiv:1906.02762_ , 2019
  * [36] Ulrike Luxburg and Olivier Bousquet  “Distance-Based Classification with Lipschitz Functions.”  In _J. Mach. Learn. Res._ 5.Jun, 2004, pp. 669–695
  * [37] Takeru Miyato et al.  “Distributional smoothing with virtual adversarial training”  In _arXiv preprint arXiv:1507.00677_ , 2015
  * [38] Takeru Miyato, Toshiki Kataoka, Masanori Koyama and Yuichi Yoshida  “Spectral normalization for generative adversarial networks”  In _arXiv preprint arXiv:1802.05957_ , 2018
  * [39] Takeru Miyato, Shin-ichi Maeda, Masanori Koyama and Shin Ishii  “Virtual adversarial training: a regularization method for supervised and semi-supervised learning”  In _IEEE transactions on pattern analysis and machine intelligence_ 41.8 IEEE, 2018, pp. 1979–1993
  * [40] Seyed-Mohsen Moosavi-Dezfooli, Alhussein Fawzi and Pascal Frossard  “Deepfool: a simple and accurate method to fool deep neural networks”  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 2574–2582
  * [41] Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester and Nati Srebro  “Exploring generalization in deep learning”  In _Advances in neural information processing systems_ 30, 2017
  * [42] OpenAI  “GPT-4 Technical Report”, 2023  arXiv:2303.08774 [cs.CL]
  * [43] Nicolas Papernot et al.  “Distillation as a defense to adversarial perturbations against deep neural networks”  In _2016 IEEE symposium on security and privacy (SP)_ , 2016, pp. 582–597  IEEE
  * [44] Tomas Pevny and Vojtech Kovarik  “Approximation capability of neural networks on spaces of probability measures and tree-structured domains”  In _arXiv preprint arXiv:1906.00764_ , 2019
  * [45] Gabriel Peyré and Marco Cuturi  “Computational optimal transport: With applications to data science”  In _Foundations and Trends® in Machine Learning_ 11.5-6 Now Publishers, Inc., 2019, pp. 355–607
  * [46] Xianbiao Qi et al.  “Lipsformer: Introducing lipschitz continuity to vision transformers”  In _arXiv preprint arXiv:2304.09856_ , 2023
  * [47] Alec Radford et al.  “Language models are unsupervised multitask learners”  In _OpenAI blog_ 1.8, 2019, pp. 9
  * [48] Alec Radford et al.  “Robust speech recognition via large-scale weak supervision”  In _International Conference on Machine Learning_ , 2023, pp. 28492–28518  PMLR
  * [49] Colin Raffel et al.  “Exploring the limits of transfer learning with a unified text-to-text transformer”  In _The Journal of Machine Learning Research_ 21.1 JMLRORG, 2020, pp. 5485–5551
  * [50] Mihaela Rosca, Theophane Weber, Arthur Gretton and Shakir Mohamed  “A case for new neural network smoothness constraints”  PMLR, 2020
  * [51] Michael E Sander, Pierre Ablin, Mathieu Blondel and Gabriel Peyré  “Sinkformers: Transformers with doubly stochastic attention”  In _International Conference on Artificial Intelligence and Statistics_ , 2022, pp. 3515–3530  PMLR
  * [52] Victor Sanh, Lysandre Debut, Julien Chaumond and Thomas Wolf  “DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter”  In _arXiv:1910.01108_ , 2019
  * [53] Filippo Santambrogio  “Optimal transport for applied mathematicians”  In _Birkäuser, NY_ 55.58-63 Springer, 2015, pp. 94
  * [54] Jure Sokolić, Raja Giryes, Guillermo Sapiro and Miguel RD Rodrigues  “Robust large margin deep neural networks”  In _IEEE Transactions on Signal Processing_ 65.16 IEEE, 2017, pp. 4265–4280
  * [55] Gilbert Strang  “Calculus”  SIAM, 1991
  * [56] Christian Szegedy et al.  “Intriguing properties of neural networks”  In _arXiv preprint arXiv:1312.6199_ , 2013
  * [57] Yusuke Tsuzuku, Issei Sato and Masashi Sugiyama  “Lipschitz-margin training: Scalable certification of perturbation invariance for deep neural networks”  In _Advances in neural information processing systems_ 31, 2018
  * [58] Ashish Vaswani et al.  “Attention is all you need”  In _Advances in neural information processing systems_ 30, 2017
  * [59] Aladin Virmaux and Kevin Scaman  “Lipschitz regularity of deep neural networks: analysis and efficient estimation”  In _Advances in Neural Information Processing Systems_ 31, 2018
  * [60] James Vuckovic, Aristide Baratin and Rémi Tachet Combes  “A Mathematical Theory of Attention”  In _ArXiv_ abs/2007.02876, 2020
  * [61] James Vuckovic, Aristide Baratin and Remi Tachet des Combes  “On the regularity of attention”  In _arXiv preprint arXiv:2102.05628_ , 2021
  * [62] Changhan Wang et al.  “Fairseq S2T: Fast speech-to-text modeling with fairseq”  In _arXiv preprint arXiv:2010.05171_ , 2020
  * [63] Tsui-Wei Weng et al.  “Evaluating the robustness of neural networks: An extreme value theory approach”  In _arXiv preprint arXiv:1801.10578_ , 2018
  * [64] Thomas Wolf et al.  “Huggingface’s transformers: State-of-the-art natural language processing”  In _arXiv preprint arXiv:1910.03771_ , 2019
  * [65] Wenqian Ye, Yunsheng Ma, Xu Cao and Kun Tang  “Mitigating Transformer Overconfidence via Lipschitz Regularization”  In _arXiv preprint arXiv:2306.06849_ , 2023
  * [66] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby and Lucas Beyer  “Scaling vision transformers”  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2022, pp. 12104–12113
  * [67] H Zhao et al.  “Point transformer. arXiv”  In _arXiv preprint arXiv:2012.09164_ , 2020
  * [68] Aaron Zweig and Joan Bruna  “A functional perspective on learning symmetric functions with neural networks”  In _International Conference on Machine Learning_ , 2021, pp. 13023–13032  PMLR

##  Appendix A Optimal Transport Toolbox

This section gathers some useful definitions and lemmas from optimal
transport. In what follows, $\mathcal{X}$ is a Borel subset of
$\mathbb{R}^{d}$.

###  A.1 Pushforward, Wasserstein Distance

Let us start with the notion of pushforward.

######  Definition A.1 (Pushforward).

Set $\mu$ a probability measure on $\mathcal{X}$ and
$\varphi\colon\mathcal{X}\to\mathcal{X}$ a measurable function. The
pushforward of $\mu$ by $\varphi$, denoted $\varphi_{\sharp}\mu$, is the
probability measure given by

$\left(\varphi_{\sharp}\mu\right)(B)=\mu(\varphi^{-1}(B))$

for any Borel set $B\subset\mathcal{X}$, where
$\varphi^{-1}(B)\coloneqq\\{x\in\mathbb{R}^{d}:\varphi(x)\in B\\}$.

The pushforward measure $\varphi_{\sharp}\mu$ can be seen as the result of a
transportation of the mass of $\mu$ by $\varphi$. Intuitively,
$\varphi_{\sharp}\mu$ is obtained by transporting each element of mass
$\mu(\mathrm{d}x)$ from $x$ to $\varphi(x)$.

Another crucial tool is the notion of Wasserstein distance.

######  Definition A.2 (Wasserstein space, Wasserstein distance).

Let $p\geq 1$. Denote

$\mathcal{P}_{p}(\mathcal{X})\coloneqq\\{\mu\in\mathcal{P}(\mathcal{X}):\int_{\mathcal{X}}\left\lvert
x\right\rvert^{p}\mathrm{d}\mu(x)<\infty\\}$

the $p$-Wasserstein space. Then, the $p$-Wasserstein distance between two
probability measures $\mu,\nu\in\mathcal{P}_{p}(\mathcal{X})$ is defined as

$W_{p}(\mu,\nu)\coloneqq\left(\inf_{\pi\in\Pi(\mu,\nu)}\int\left\lvert
x-y\right\rvert^{p}\mathrm{d}\pi(x,y)\right)^{1/p}$

with $\Pi(\mu,\nu)$ the set of all couplings between $\mu$ and $\nu$, i.e. of
all probability measures $\pi\in\mathcal{P}(\mathcal{X}\times\mathcal{X})$
such that $\int\pi(\cdot,y)\mathrm{d}y=\mu$ and
$\int\pi(x,\cdot)\mathrm{d}x=\nu$.

Wasserstein distances have the following nice property, that is a direct
consequence of Jensen inequality.

######  Lemma A.1.

For every $p\geq 1$, it holds

$W_{1}\leq W_{p}.$

The distance $W_{1}$ has also a simple dual formulation.

######  Lemma A.2 ($W_{1}$ duality formula).

The distance $W_{1}$ can be rewritten with the so-called duality formula: for
all $\mu,\nu\in\mathcal{P}_{1}(\mathcal{X})$, it holds

$W_{1}(\mu,\nu)=\sup_{\operatorname{\mathrm{Lip}}(\varphi)\leq
1}\int\varphi\,\mathrm{d}(\mu-\nu),$ (15)

where the supremum is taken over all functions
$\varphi\colon\mathcal{X}\to\mathbb{R}$ with a Lipschitz constant bounded by
one.

The following result is useful to bound the Wasserstein distance between two
probability measures that are pushed forward by the same map.

######  Lemma A.3.

Let $p\geq 1$. Consider a measurable function
$\varphi\colon\mathcal{X}\to\mathcal{X}$, and probability measures
$\mu,\nu\in\mathcal{P}_{p}(\mathcal{X})$ such that
$\varphi_{\sharp}\mu\in\mathcal{P}_{p}(\mathcal{X})$ and
$\varphi_{\sharp}\nu\in\mathcal{P}_{p}(\mathcal{X})$. Then, it holds

$W_{p}(\varphi_{\sharp}\mu,\varphi_{\sharp}\nu)\leq\operatorname{\mathrm{Lip}}(\varphi)W_{p}(\mu,\nu).$

###### Proof.

We have

$\displaystyle W_{p}(\varphi_{\sharp}\mu,\varphi_{\sharp}\nu)^{p}$
$\displaystyle=\inf_{\pi^{\prime}\in\Pi(\varphi_{\sharp}\mu,\varphi_{\sharp}\nu)}\int\left\lVert
x-y\right\rVert^{p}\mathrm{d}\pi^{\prime}(x,y)$
$\displaystyle\leq\inf_{\pi\in\Pi(\mu,\nu)}\int\left\lVert\varphi(x)-\varphi(y)\right\rVert^{p}\mathrm{d}\pi(x,y)$
$\displaystyle\leq\operatorname{\mathrm{Lip}}(\varphi)^{p}\inf_{\pi\in\Pi(\mu,\nu)}\int\left\lVert
x-y\right\rVert^{p}\mathrm{d}\pi(x,y)$
$\displaystyle=\operatorname{\mathrm{Lip}}(\varphi)^{p}W_{p}(\mu,\nu)^{p},$

where the first inequality derives from the fact that every
$\pi\in\Pi(\mu,\nu)$ induces a coupling
$\pi^{\prime}\in\Pi(\varphi_{\sharp}\mu,\varphi_{\sharp}\nu)$ by setting

$\pi^{\prime}(B_{1}\times
B_{2})\coloneqq\pi(\varphi^{-1}(B_{1})\times\varphi^{-1}(B_{2}))$

for all Borel sets $B_{1},B_{2}\subset\mathcal{X}$, and that with this choice
of $\pi^{\prime}$ it holds

$\int\left\lVert
x-y\right\rVert^{p}\mathrm{d}\pi^{\prime}(x,y)=\int\left\lVert\varphi(x)-\varphi(y)\right\rVert^{p}\mathrm{d}\pi(x,y).$

∎

Let us now bound the Wasserstein distance between two different pushforwards
of the same probability measure.

######  Lemma A.4.

Let $p\geq 1$. Consider two measurable functions
$\varphi,\psi\colon\mathcal{X}\to\mathcal{X}$, and a probability measure
$\mu\in\mathcal{P}_{p}(\mathcal{X})$ such that
$\varphi_{\sharp}\mu\in\mathcal{P}_{p}(\mathcal{X})$ and
$\psi_{\sharp}\mu\in\mathcal{P}_{p}(\mathcal{X})$. Then, it holds

$W_{p}(\varphi_{\sharp}\mu,\psi_{\sharp}\mu)\leq\left\lVert\varphi-\psi\right\rVert_{L^{p}(\mu)}.$

###### Proof.

Recall that

$W_{p}(\varphi_{\sharp}\mu,\psi_{\sharp}\mu)^{p}=\inf_{\pi^{\prime}\in\Pi(\varphi_{\sharp}\mu,\psi_{\sharp}\mu)}\int\left\lVert
x-y\right\rVert^{p}\mathrm{d}\pi^{\prime}(x,y).$

Now consider the following coupling between $\varphi_{\sharp}\mu$ and
$\psi_{\sharp}\mu$, defined by the relation

$\pi^{\prime}(B\times C)\coloneqq\mu(\varphi^{-1}(B)\cap\psi^{-1}(C))$

for every Borel sets $B,C\subset\mathcal{X}$. In other words, we set
$\mathrm{d}\pi^{\prime}(y,z)\coloneqq\int_{\varphi^{-1}(y)\cap\psi^{-1}(z)}\mathrm{d}\mu$,
and $\mathrm{d}\pi^{\prime}(y,z)=0$ if
$\varphi^{-1}(y)\cap\psi^{-1}(z)=\emptyset$. With this definition of
$\pi^{\prime}$, we have

$\displaystyle W_{p}(\varphi_{\sharp}\mu,\psi_{\sharp}\mu)^{p}$
$\displaystyle\leq\int\left\lVert
x-y\right\rVert^{p}\mathrm{d}\pi^{\prime}(x,y)$
$\displaystyle=\int\left\lVert\varphi(x)-\psi(x)\right\rVert^{p}\mathrm{d}\mu(x).$

∎

###  A.2 Geodesics

The notion of geodesic is useful for the following section of the Appendix.

######  Definition A.3 (Geodesic).

Let $(\mathcal{E},d_{\mathcal{E}})$ be a metric space. A curve
$\gamma\colon[0,1]\to\mathcal{E}$ is called a geodesic if there exists a
constant $v\geq 0$ such that for all $t_{1},t_{2}\in[0,1]$ we have

$d_{\mathcal{E}}(\gamma(t_{1}),\gamma(t_{2}))=v\left\lvert
t_{2}-t_{1}\right\rvert.$

We say that the space $\mathcal{E}$ is a geodesic space if for any
$x,y\in\mathcal{E}$, there exists a geodesic between $x$ and $y$.

One important example of geodesic space is the 2-Wasserstein space
$\mathcal{P}_{2}(\mathbb{R}^{d})$.

######  Lemma A.5.

The space $\mathcal{P}_{2}(\mathbb{R}^{d})$ is a geodesic space.

For a proof of this result, see for instance [53].

##  Appendix B Proofs of Section 3

###  B.1 Proof of Lemma 3.1

###### Proof of Lemma 3.1.

The first inequality is straightforward. Let us assume that the space
$\mathcal{E}$ admits geodesics, and show that

$\operatorname{\mathrm{Lip}}(\varphi)=\sup_{x\in\mathcal{E}}\operatorname{\mathrm{Lip}}_{x}(\varphi).$

Recall that
$\operatorname{\mathrm{Lip}}_{x}(\varphi)\coloneqq\lim_{\varepsilon>0}\operatorname{\mathrm{Lip}}\left(\varphi_{\lvert
B(x,\varepsilon)}\right)$. Let us first show that for all $\varepsilon>0$, it
holds

$\operatorname{\mathrm{Lip}}(\varphi)=\sup_{x\in\mathcal{E}}\operatorname{\mathrm{Lip}}\left(\varphi_{\lvert
B(x,\varepsilon)}\right).$ (16)

Let $\varepsilon>0$ and $\eta>0$. Pick $x,y\in\mathcal{E}$ such that
$\frac{d_{\mathcal{E}}(\varphi(x),\varphi(y))}{d_{\mathcal{E}}(x,y)}\geq\operatorname{\mathrm{Lip}}(\varphi)-\eta$.
Let $\gamma\colon[0,1]\to\mathcal{E}$ a geodesic between $x$ and $y$, and
denote $z_{i}\coloneqq\gamma(i/n)$ for $0\leq i\leq n$, with $n$ chosen large
enough to have $d_{\mathcal{E}}(z_{i+1},z_{i})<\varepsilon$ for all $0\leq
i\leq n$. Then

$\displaystyle d_{\mathcal{E}}(\varphi(z_{n}),\varphi(z_{0}))$
$\displaystyle\leq\sum_{i=0}^{n-1}d_{\mathcal{E}}(\varphi(z_{i+1}),\varphi(z_{i}))$
$\displaystyle\leq\sup_{z\in\mathcal{E}}\operatorname{\mathrm{Lip}}\left(\varphi_{\lvert
B(z,\varepsilon)}\right)\sum_{i=0}^{n-1}d_{\mathcal{E}}(z_{i+1},z_{i})$
$\displaystyle=\sup_{z\in\mathcal{E}}\operatorname{\mathrm{Lip}}\left(\varphi_{\lvert
B(z,\varepsilon)}\right)d_{\mathcal{E}}(x,y)$

by definition of a geodesic. This implies that
$\sup_{z\in\mathcal{E}}\operatorname{\mathrm{Lip}}\left(\varphi_{\lvert
B(z,\varepsilon)}\right)\geq\operatorname{\mathrm{Lip}}(\varphi)-\eta$, and
therefore

$\sup_{z\in\mathcal{E}}\operatorname{\mathrm{Lip}}\left(\varphi_{\lvert
B(z,\varepsilon)}\right)\geq\operatorname{\mathrm{Lip}}(\varphi)$

by taking the limit $\eta\to 0^{+}$. The reverse inequality is
straightforward, which concludes the proof of Equation (16).

Now Equation (16) gives

$\operatorname{\mathrm{Lip}}(\varphi)=\sup_{x\in\mathcal{E}}\operatorname{\mathrm{Lip}}(\varphi_{\lvert
B(x,1)})\geq\sup_{x\in\mathcal{E}}\inf_{\varepsilon>0}\operatorname{\mathrm{Lip}}(\varphi_{\lvert
B(x,\varepsilon)})=\sup_{x\in\mathcal{E}}\operatorname{\mathrm{Lip}}_{x}(\varphi).$

It remains to prove the reverse inequality. Set $\eta>0$, and let $K_{\eta}$
be a compact set such that $\operatorname{\mathrm{Lip}}(\varphi_{\lvert
K_{\eta}})\geq\operatorname{\mathrm{Lip}}(\varphi)-\eta$. For all integers
$n\geq 1$, define $x_{n}\in K_{\eta}$ such that

$\operatorname{\mathrm{Lip}}(\varphi_{\lvert B(x_{n},1/n)})\geq\sup_{x\in
K_{\eta}}\operatorname{\mathrm{Lip}}(\varphi_{\lvert B(x_{n},1/n)\cap
K_{\eta}})-1/n.$

The sequence $(x_{n})$ lives in a compact set, so it admits a subsequential
limit: $x_{\varphi(n)}\to x^{*}$. Then, we have

$\sup_{x\in K_{\eta}}\lim_{\varepsilon\to
0+}\operatorname{\mathrm{Lip}}(\varphi_{\lvert
B(x,\varepsilon)})\geq\lim_{n\to+\infty}\operatorname{\mathrm{Lip}}(\varphi_{\lvert
B(x^{*},1/\varphi(n))})\geq\operatorname{\mathrm{Lip}}(\varphi_{\lvert
K_{\eta}}),$

which allows us to conlude by doing $\eta\to 0^{+}$. ∎

###  B.2 Proof of Lemma 3.3

###### Proof of Lemma 3.3.

Set $X\in\mathbb{R}^{n\times d}$. One can choose $\varepsilon_{1}>0$ small
enough (for example $\varepsilon_{1}<\min_{x_{i}\neq x_{j}}\left\lVert
x_{i}-x_{j}\right\rVert/2$) to have

$\left\lVert X-Y\right\rVert_{F,p}\leq\varepsilon_{1}\Rightarrow\left\lVert
X-Y\right\rVert_{F,p}=W_{p}(m(X),m(Y)).$

Indeed, if

$\left\lVert X-Y\right\rVert_{F,p}\leq\varepsilon_{1}<\min_{x_{i}\neq
x_{j}}\left\lVert x_{i}-x_{j}\right\rVert/2,$

then for all $1\leq i\leq n$, we have $\left\lVert
x_{i}-y_{i}\right\rVert\leq\varepsilon_{1}$, and thus $x_{i}$ is the nearest
neighbour (or one of the nearest neighbours) of $y_{i}$ among the $x_{j}$.
Similarly, one can choose $\varepsilon_{2}>0$ small enough to have

$\left\lVert
f(X)-f(Y)\right\rVert_{F,p}\leq\varepsilon_{2}\Rightarrow\left\lVert
f(X)-f(Y)\right\rVert_{F,p}=W_{p}(m(f(X)),m(f(Y))).$

Then, we can set $\varepsilon\leq\varepsilon_{1}$ small enough to have

$\left\lVert X-Y\right\rVert_{F,p}\leq\varepsilon\Rightarrow\left\lVert
f(X)-f(Y)\right\rVert_{F,p}\leq\varepsilon_{2},$

and it holds, for all $\eta\leq\varepsilon$ and all $Y$ such that $\left\lVert
X-Y\right\rVert_{F,p}\leq\eta$, that

$\left\lVert X-Y\right\rVert_{F,p}=W_{p}(m(X),m(Y))$

and

$\left\lVert f(X)-f(Y)\right\rVert_{F,p}=W_{p}(m(f(X)),m(f(Y))).$

Now for all $\eta\leq\varepsilon$, we have

$\displaystyle\operatorname{\mathrm{Lip}}^{\left\lVert\cdot\right\rVert_{F,p}}f_{\lvert
B_{\left\lVert\cdot\right\rVert_{F,p}}(X,\eta)}$ $\displaystyle=\sup_{Y\in
B_{\left\lVert\cdot\right\rVert_{F,p}}(X,\eta)}\frac{\left\lVert
f(X)-f(Y)\right\rVert_{F,p}}{\left\lVert X-Y\right\rVert_{F,p}}$
$\displaystyle=\sup_{Y\in
B_{\left\lVert\cdot\right\rVert_{F,p}}(X,\eta)}\frac{W_{p}(m(f(X)),m(f(Y)))}{W_{p}(m(X),m(Y))}$
$\displaystyle=\sup_{Y\in
B_{\left\lVert\cdot\right\rVert_{F,p}}(X,\eta)}\frac{W_{p}(F(m(X)),F(m(Y)))}{W_{p}(m(X),m(Y))},$

by definition of $\varepsilon$ and $F$. We conclude the proof by noticing
that:

  * $\bullet$

$Y\in B_{\left\lVert\cdot\right\rVert_{F,p}}(X,\eta)$ implies $m(Y)\in
B_{W_{p}}(m(X),\eta)$, which shows that
$\operatorname{\mathrm{Lip}}^{\left\lVert\cdot\right\rVert_{F,p}}f_{\lvert
B_{\left\lVert\cdot\right\rVert_{F,p}}(X,\eta)}\leq\operatorname{\mathrm{Lip}}^{W_{p}}F_{\lvert\mathcal{M}_{n}(\mathbb{R}^{d})}$,

  * $\bullet$

$\mu\in B_{W_{p}}(m(X),\eta)$ with $\mu\in\mathcal{M}_{n}(\mathbb{R}^{d})$
implies the existence of $Y\in\mathbb{R}^{n\times d}$ such that $\mu=m(Y)$ and
$\left\lVert X-Y\right\rVert_{F,p}=W_{p}(m(X),m(Y))$, so that $Y\in
B_{\left\lVert\cdot\right\rVert_{F,p}}(X,\eta)$. Indeed, take $Y$ such that
$\mu_{n}=m(Y)$ and then permute its coordinates so that $x_{i}$ becomes a
nearest neighbour of $y_{i}$. This shows the reverse inequality, and concludes
the proof.

∎

###  B.3 Proof of Proposition 3.7

###### Proof of Proposition 3.7.

Let $\bar{\mu}$ and $\bar{\nu}$ be two distinct measures in ${\cal
P}([0,1]\times B_{R})$. Assume that $\bar{\mu}$ and $\bar{\nu}$ have the same
first marginal, denoted $\theta$ (otherwise, we consider that they are
associated to an infinite Lipschitz ratio). We have

$\displaystyle d_{p}(F^{m}(\bar{\mu}),F^{m}(\bar{\nu}))$ $\displaystyle\leq
d_{p}((\mathrm{Id}+\Gamma_{\bar{\mu}})_{\sharp}\bar{\mu},(\mathrm{Id}+\Gamma_{\bar{\mu}})_{\sharp}\bar{\nu})+d_{p}((\mathrm{Id}+\Gamma_{\bar{\mu}})_{\sharp}\bar{\nu},(\mathrm{Id}+\Gamma_{\bar{\nu}})_{\sharp}\bar{\nu})$
$\displaystyle\leq\left(\int_{0}^{1}W_{p}\left((\mathrm{Id}+\Gamma_{\bar{\mu}}^{s})_{\sharp}\mu^{s},(\mathrm{Id}+\Gamma_{\bar{\mu}}^{s})_{\sharp}\nu^{s}\right)^{p}\right)^{1/p}$
$\displaystyle\phantom{space}+\left(\int_{0}^{1}W_{p}\left(\mathrm{Id}+\Gamma_{\bar{\mu}}^{s})_{\sharp}\nu^{s},(\mathrm{Id}+\Gamma_{\bar{\nu}}^{s})_{\sharp}\nu^{s}\right)\right)^{1/p},$

where we denote

$\Gamma_{\bar{\mu}}^{s}(x):=\frac{\int VyG(x,y)\mathbf{1}_{\tau\leq
s}\mathrm{d}\bar{\mu}(\tau,y)}{\int G(x,y)\mathbf{1}_{\tau\leq
s}\mathrm{d}\bar{\mu}(\tau,y)}.$

Using Lemma A.3, we get

$W_{p}\left((\mathrm{Id}+\Gamma_{\bar{\mu}}^{s})_{\sharp}\mu^{s},(\mathrm{Id}+\Gamma_{\bar{\mu}}^{s})_{\sharp}\nu^{s}\right)\leq\operatorname{\mathrm{Lip}}(\mathrm{Id}+\Gamma_{\bar{\mu}}^{s})W_{p}(\mu^{s},\nu^{s})$

where the Lipschitz constant is taken on $B_{R}$. A similar computation as for
traditional self-attention shows that for all $0\leq s\leq 1$ and $x\in B_{R}$
we have

$\left\lVert D_{x}\Gamma^{s}\right\rVert_{2}\leq\left\lVert
V\right\rVert_{2}\left\lVert A\right\rVert_{2}R^{2},$

so that

$\operatorname{\mathrm{Lip}}(\mathrm{Id}+\Gamma_{\bar{\mu}}^{s})\leq
1+\left\lVert V\right\rVert_{2}\left\lVert A\right\rVert_{2}R^{2}.$

To bound the second term in the previous inequality, we use Lemma A.4, to get

$W_{p}\left((\mathrm{Id}+\Gamma_{\bar{\mu}}^{s})_{\sharp}\nu^{s},(\mathrm{Id}+\Gamma_{\bar{\nu}}^{s})_{\sharp}\nu^{s}\right)\leq\left\lVert\Gamma_{\bar{\mu}}^{s}-\Gamma_{\bar{\nu}}^{s}\right\rVert_{L^{\infty}(B_{R})}.$

Again, a similar computation as for traditional self-attention shows that

$\left\lVert\Gamma_{\bar{\mu}}^{s}-\Gamma_{\bar{\nu}}^{s}\right\rVert_{L^{\infty}(B_{R})}\leq\left\lVert
V\right\rVert_{2}(1+2\left\lVert A\right\rVert_{2}R^{2})e^{2\left\lVert
A\right\rVert_{2}R^{2}}W_{p}(\mu^{s},\nu^{s}),$

which concludes the proof. ∎

###  B.4 Weighted Self-Attention

Introducing $f_{a}$ amounts to looking at $f$ on the space $\mathbb{R}^{nd}$
equipped with a new scalar product $\left\langle\cdot,\cdot\right\rangle_{a}$,
defined as follows.

$\left\langle X,Y\right\rangle_{a}\coloneqq\sum_{i=1}^{n}a_{i}\left\langle
x_{i},y_{i}\right\rangle,$

with $X=(x_{1},\dots,x_{n})^{\top}$ and $Y=(y_{1},\dots,y_{n})^{\top}$. This
change of scalar product also changes the definition of the norm, the operator
norm and the Jacobian, which are central in the analysis of the Lipschitz
properties of $f_{a}$. We have the following expressions for these quantities
in the case of weighted self-attention in its vectorized form.

######  Lemma B.1.

Let $a\in\Sigma_{n}$. Equip the vector space $\mathbb{R}^{nd}$ with the
following scalar product:

$\left\langle\vec{X},\vec{Y}\right\rangle_{a}\coloneqq\sum_{i=1}^{n}a_{i}\left\langle
x_{i},y_{i}\right\rangle$

where $\vec{X}$ and $\vec{Y}$ are the vectorized forms of the matrices
$X=(x_{1},\dots,x_{n})^{\top}$ and $Y=(y_{1},\dots,y_{n})^{\top}$, with row-
wise vectorization. We have the following properties.

  1. (i)

The norm associated with the scalar product
$\left\langle\cdot,\cdot\right\rangle_{a}$, denoted
$\left\lVert\cdot\right\rVert_{a}$, reads

$\left\lVert\vec{X}\right\rVert_{a}^{2}=\sum_{i=1}^{n}a_{i}\left\lVert
x_{i}\right\rVert^{2}=\left\lVert\operatorname{\mathrm{diag}}(\sqrt{a}\otimes
1_{d})\vec{X}\right\rVert^{2},$

where the square root is taken coordinatewise.

  2. (ii)

The scalar product $\left\langle\cdot,\cdot\right\rangle_{a}$ induces the
following action $\cdot_{a}$ of linear operators on vectors of
$\mathbb{R}^{nd}$:

$M\cdot_{a}\vec{X}=M\operatorname{\mathrm{diag}}(a\otimes 1_{d})\vec{X},$

for all $M\in\mathbb{R}^{nd\times nd}$.

  3. (iii)

The induced operator norm is

$\left\lVert
M\right\rVert_{2,a}=\sup_{\left\lVert\vec{X}\right\rVert_{a}=1}\left\lVert
M\cdot_{a}\vec{X}\right\rVert_{a}=\left\lVert\operatorname{\mathrm{diag}}(\sqrt{a}\otimes
1_{d})M\operatorname{\mathrm{diag}}(\sqrt{a}\otimes 1_{d})\right\rVert_{2}.$

  4. (iv)

If all coordinates of $a$ are non zero, the Jacobian of $f_{a}$ on the space
$(\mathbb{R}^{nd},\left\langle\cdot,\cdot\right\rangle_{a})$, denoted
$D^{a}f_{a}$, is equal to

$D_{\vec{X}}^{a}f_{a}=(D_{\vec{X}}f_{a})\operatorname{\mathrm{diag}}(a\otimes
1_{d})^{-1}.$

The proof consists in short computations we do not report here.

###  B.5 Proof of Proposition 3.5

We start with the following lemma, that derives from a simple computation.

######  Lemma B.2.

Let $\mu\coloneqq p\delta_{x}+(1-p)\delta_{y}$, with $p\in(0,1)$ and
$x,y\in\mathbb{R}^{d}$. Denote $a\coloneqq(p,1-p)$ and
$P=:\left(\begin{smallmatrix}\pi_{1}&1-\pi_{1}\\\
1-\pi_{2}&\pi_{2}\end{smallmatrix}\right)$, recalling that
$P\coloneqq\operatorname{\mathrm{softmax}}(XA^{\top}X^{\top})$ with
$X=\left(\begin{smallmatrix}x^{\top}\\\ y^{\top}\end{smallmatrix}\right)$.
Then

$D_{X}f_{a}=\begin{pmatrix}B_{1}&B_{2}\\\ B_{3}&B_{4}\end{pmatrix}$

with

$\displaystyle B_{1}=\pi_{1}(1-\pi_{1})(x-y)(2x-y)^{\top}A+\pi_{1}I_{d},$
$\displaystyle B_{2}=-\pi_{1}(1-\pi_{1})(x-y)x^{\top}A+(1-\pi_{1})I_{d},$
$\displaystyle B_{3}=-\pi_{2}(1-\pi_{2})(y-x)y^{\top}A+(1-\pi_{2})I_{d},$
$\displaystyle B_{4}=\pi_{2}(1-\pi_{2})(y-x)(2y-x)^{\top}A+\pi_{2}I_{d}.$

Now, as we have assumed symmetry of the matrix $A$, we can write

$A=\sum\gamma_{i}u_{i}u_{i}^{\top},$

with $\gamma_{1}\geq\dots\geq\gamma_{d}$ the eigenvalues of $A$, and
$(u_{1},\dots,u_{d})$ an orthonormal basis of associated eigenvectors. Let us
prove Proposition 3.5 only in the first case, i.e. when
$p=p_{R}=e^{-2CR^{2}}$, and $x=Ru_{1}$ and $y=\frac{R}{2}u_{1}$. The other
case can be proven with the same method. Observing that $\pi_{1}\to 1$ and
$\pi_{2}\to 1/2$, we compute an equivalent of each block of $D_{X}f_{a}$ when
$R\to+\infty$.

$\displaystyle B_{1}$ $\displaystyle\to I_{d}$ $\displaystyle B_{2}$
$\displaystyle\sim-\frac{\gamma R^{2}}{2}e^{-\frac{\gamma
R^{2}}{4}}u_{1}u_{1}^{\top}\to 0_{d}$ $\displaystyle B_{3}$
$\displaystyle\sim\frac{\gamma R^{2}}{16}u_{1}u_{1}^{\top}$ $\displaystyle
B_{4}$ $\displaystyle\to\frac{1}{2}I_{d}.$

Equations (iii) and (iv) of Lemma B.1 allow us to conclude the proof, noticing
that

$\left\lVert\begin{pmatrix}\mathcal{O}(1)&{\cal O}(1)\\\ \frac{\gamma
R^{2}}{16}e^{\frac{\gamma R^{2}}{8}}u_{1}u_{1}^{\top}&{\cal
O}(1)\end{pmatrix}\right\rVert_{p}\sim\frac{\gamma R^{2}}{16}e^{\frac{\gamma
R^{2}}{8}}.$

##  Appendix C Proofs of Section 4

###  C.1 Effect of Duplication on the Jacobian of Self-Attention

Let us start by stating a formula for the Jacobian of $f$ at
$X\coloneqq(x_{1},\dots,x_{n})^{\top}$.

######  Lemma C.1 ([26]).

The Jacobian of (the vectorized form of) matrix self-attention $f$ is of the
following form:

$D_{X}f=(I_{n}\otimes V)\left(P\otimes
I_{d}+D\otimes_{d}A+R\otimes_{d}A^{\top}\right),$ (17)

with

  * $\bullet$

$P\coloneqq\operatorname{\mathrm{softmax}}(XA^{\top}X^{\top})$,

  * $\bullet$

$D$ a block-diagonal matrix, with $n$ diagonal blocks of shape $d\times d$
given by $\operatorname{\mathrm{Var}}^{(i)}Z$ for $1\leq i\leq n$, where
$\operatorname{\mathrm{Var}}^{(i)}Z$ is the variance of the random variable
$Z$ following the law $\mathbb{P}(Z=x_{j})=P_{ij}$,

  * $\bullet$

$R=(R_{ij})_{1\leq i,j\leq n}$ a block matrix with $R_{ij}\coloneqq
P_{ij}\left(x_{j}-\sum_{k=1}^{n}P_{ik}x_{k}\right)x_{i}^{\top}\in\mathbb{R}^{d\times
d}$.

We can then write the Jacobian of $f$ at
$Y\coloneqq\left(\begin{smallmatrix}1\\\ 1\end{smallmatrix}\right)\otimes X$
in function of $D_{X}f$.

######  Lemma C.2.

Let $X\in\mathbb{R}^{n\times d}$, and denote
$Y\coloneqq\left(\begin{smallmatrix}1\\\ 1\end{smallmatrix}\right)\otimes X$.
The point cloud $Y$ is simply the duplicated version of $X$, and is associated
to the same measure as $X$. Assume that $V=I_{d}$. Then, we have

$D_{Y}f_{2n}=\frac{1}{2}\begin{pmatrix}1&1\\\ 1&1\end{pmatrix}\otimes
D_{X}f_{n}+\frac{1}{2}\begin{pmatrix}1&-1\\\
-1&1\end{pmatrix}\otimes(D\otimes_{d}A),$ (18)

where the matrix $D$ is defined in Lemma C.1. Therefore, we have

$\left\lVert D_{Y}f_{2n}\right\rVert_{2}>\left\lVert
D_{X}f_{n}\right\rVert_{2}\leavevmode\nobreak\ \leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ \mbox{if and only
if}\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
\leavevmode\nobreak\ \left\lVert D\otimes_{d}A\right\rVert_{2}>\left\lVert
D_{X}f_{n}\right\rVert_{2}.$

###  C.2 Concentrated Measures Are Not Splitted

Lemma C.2 allows us to show that empirical measures that are concentrated
enough around their expectation cannot be splitted by matrix self-attention.

######  Proposition C.3.

Assume that $k=d$, and that $V=A=I_{d}$. Let a matrix $X\in\mathbb{R}^{n\times
d}$ whose rows belong to a ball of radius 1. Denote
$Y\coloneqq(\begin{smallmatrix}1\\\ 1\end{smallmatrix})\otimes X$. Then

$\operatorname{\mathrm{Lip}}_{X}f_{n}=\operatorname{\mathrm{Lip}}_{Y}f_{2n}.$

###### Proof.

We start with a useful lemma.

######  Lemma C.4.

Let $Z$ be a random variable supported in
$\,\overline{\\!{B}}(x_{0},r)\subset\mathbb{R}^{d}$, with any
$x_{0}\in\mathbb{R}^{d}$ and $r>0$. Then, denoting
$\operatorname{\mathrm{Var}}Z=\mathbb{E}[(Z-\mathbb{E}Z)(Z-\mathbb{E}Z)^{\top}]$
we have

$\left\lVert\operatorname{\mathrm{Var}}Z\right\rVert_{2}\leq r^{2},$

with equality when $\mu_{Z}=\frac{1}{2}(\delta_{x_{0}+x}+\delta_{x_{0}-x})$
with $x\in\mathbb{R}^{d}$ such that $\left\lVert x\right\rVert=r$.

###### Proof.

Let us assume without loss of generality that $x_{0}=0$. It is straightforward
to check that if $\mu_{Z}=\frac{1}{2}(\delta_{x}+\delta_{-x})$ then
$\left\lVert\operatorname{\mathrm{Var}}Z\right\rVert_{2}=r^{2}$. To show that
this is the maximal value the variance can take, we use the triangle
inequality:

$\displaystyle\left\lVert\mathbb{E}[(Z-\mathbb{E}Z)(Z-\mathbb{E}Z)^{\top}]\right\rVert_{2}$
$\displaystyle\leq\mathbb{E}\left\lVert(Z-\mathbb{E}Z)(Z-\mathbb{E}Z)^{\top}\right\rVert_{2}$
$\displaystyle=\mathbb{E}[\left\lVert Z-\mathbb{E}Z\right\rVert^{2}]$
$\displaystyle=\mathbb{E}[\left\lVert
Z\right\rVert^{2}]-\left\lVert\mathbb{E}Z\right\rVert^{2}.$

Now let us pick any $x\in B_{R}\setminus B(0,r)$. We have
$\mathbb{E}(Z-x)^{\top}(Z+x)\leq 0$, as the angle between the vectors $Z-x$
and $Z+x$ is at least $\pi/2$ for $Z$ in $B(0,r)$. By expanding this relation
we get

$E[\left\lVert Z\right\rVert^{2}]-\left\lVert x\right\rVert^{2}\leq 0,$

which yields the result.

∎

Let us now move on to the proof of the Proposition. We first show that for any
number of samples $n$ and any matrix $X\in\mathbb{R}^{n\times d}$, we have
$\left\lVert D_{X}f_{n}\right\rVert_{2}\geq 1$. Recall Lemma C.1:

$D_{X}f=P\otimes I_{d}+D+R.$

One can easily see that $R(1_{n}\otimes I_{d})=0$, where $1_{n}$ is the vector
of size $n$ with all coordinates equal to $1$, so that

$\displaystyle D_{X}f_{n}(1_{n}\otimes I_{d})$ $\displaystyle=P(1_{n}\otimes
I_{d})+D(1_{n}\otimes I_{d})$ $\displaystyle=1_{n}\otimes
I_{d}+\begin{pmatrix}(\operatorname{\mathrm{Var}}^{(1)}Z)\\\ \vdots\\\
(\operatorname{\mathrm{Var}}^{(n)}Z)\end{pmatrix}.$

Then for example $\left\lVert D_{X}f_{n}(1_{n}\otimes
e_{1})\right\rVert^{2}\geq\sum_{i=1}^{n}\left(1+D_{11}\right)^{2}\geq
n=\left\lVert 1_{n}\otimes e_{1}\right\rVert^{2}$, where $e_{1}$ is the first
vector of the canonical basis $(e_{1},\dots,e_{d})$ of $\mathbb{R}^{d}$,
because $(\operatorname{\mathrm{Var}}^{(i)}Z)_{11}\geq 0$ for any index
$i$.333This quantity is indeed the variance of a scalar random variable. This
proves that $\left\lVert D_{X}f_{n}\right\rVert_{2}\geq 1$.

Now thanks to Lemma C.2, to conclude that $X$ is not splitted by self-
attention, it is enough to show that for all $1\leq i\leq n$, we have

$\max_{1\leq i\leq
n}\left\lVert\operatorname{\mathrm{Var}}^{(i)}Z\right\rVert_{2}\leq\left\lVert
D_{X}f_{n}\right\rVert_{2},$ (19)

so that the maximal eigenvalue of $(D_{Y}f_{2n})^{\top}D_{Y}f_{2n}$ comes from
$(D_{X}f_{n})^{\top}D_{X}f_{n}$, and
$\operatorname{\mathrm{Lip}}_{Y}f_{2n}=\left\lVert
D_{Y}f_{2n}\right\rVert_{2}=\left\lVert
D_{X}f_{n}\right\rVert_{2}=\operatorname{\mathrm{Lip}}_{X}f_{n}$. And property
(19) follows from Lemma C.4, together with the fact that $\left\lVert
D_{X}f_{n}\right\rVert_{2}\geq 1$. ∎

###  C.3 Proof of Theorem 4.2

###### Proof of Theorem 4.2.

Using decomposition (17), we define the matrices $D^{N}$ and $P^{N}$ and
$R^{N}$ such that

$D_{Y^{N}}f=D^{N}\otimes_{d}A+P^{N}\otimes I_{d}+R^{N}\otimes_{d}A^{\top}$

for any positive integer $N$. One can check that

$\begin{cases}P^{N}&=\frac{1}{N}1_{N\times N}\otimes P^{1},\\\
D^{N}&=\operatorname{\mathrm{diag}}(1_{N})\otimes D^{1},\\\
R^{N}&=\frac{1}{N}1_{N\times N}\otimes R^{1},\end{cases}$

where $1_{N}$ and $1_{N\times N}$ are respectively is the vector of dimension
$N$ and the matrix of shape $N\times N$ with all coordinates equal to 1.
Hence,

$\displaystyle D_{Y^{N}}f$
$\displaystyle=\operatorname{\mathrm{diag}}(1_{N})\otimes
D^{1}\otimes_{d}A+\frac{1}{N}1_{N\times N}\otimes(P^{1}\otimes
I_{d}+R^{1}\otimes A^{\top})$ $\displaystyle=G_{N}\otimes
D^{1}\otimes_{d}A+\frac{1}{N}1_{N\times N}\otimes D_{X}f,$

with

$G_{N}=\frac{N-1}{N}\begin{pmatrix}1&-\frac{1}{N-1}&\cdots&-\frac{1}{N-1}\\\
-\frac{1}{N-1}&\ddots&\ddots&\vdots\\\ \vdots&\ddots&\ddots&-\frac{1}{N-1}\\\
-\frac{1}{N-1}&\cdots&-\frac{1}{N-1}&1\end{pmatrix}.$

Noticing that

$G_{N}1_{N\times N}=1_{N\times N}G_{N}=0,$

and that

$G_{N}^{2}=G_{N}\mbox{ and }\frac{1}{N^{2}}1_{N\times
N}^{2}=\frac{1}{N}1_{N\times N},$

we get

$D_{Y^{N}}f^{\top}D_{Y^{N}}f=G_{N}\otimes(D^{1}\otimes_{d}A)^{\top}(D^{1}\otimes_{d}A)+\frac{1}{N}1_{N\times
N}\otimes D_{X}f^{\top}D_{X}f.$

Computing the characteristic polynomial of $G_{N}$, one sees that this matrix
has two eigenvalues: $1$ and $\frac{N-1}{N}$. The eigenspace associated to $1$
has dimension $N-1$, and is equal to the set
$\\{x\in\mathbb{R}^{N}\mid\sum_{i=1}^{N}x_{i}=0\\}$, which is included in the
kernel of the matrix $1_{N\times N}$. This leads to the following
observations:

  * $\bullet$

The largest eigenvalue of the matrix $D_{Y^{N}}f^{\top}D_{Y^{N}}f$, which is
also equal to the squared operator norm $\left\lVert
D_{Y^{N}}\right\rVert_{2}^{2}$, is equal to

$\max\left(\left\lVert D_{X}f\right\rVert_{2}^{2},\max_{1\leq i\leq
n}\left\lVert(\operatorname{\mathrm{Var}}^{(i)}Z)A\right\rVert_{2}^{2}\right).$

  * $\bullet$

When there is mass splitting at $X$, namely when $\max_{1\leq i\leq
n}\left\lVert(\operatorname{\mathrm{Var}}^{(i)}Z)A\right\rVert_{2}^{2}>\left\lVert
D_{X}f\right\rVert_{2}^{2}$, we have $\left\lVert
D_{Y^{N}}f\right\rVert_{2}=\left\lVert D_{Y^{2}}f\right\rVert_{2}>\left\lVert
D_{X}f\right\rVert_{2}$ for all $N\geq 2$.

  * $\bullet$

If the mass splits at $X$, then the local Lipschitz constant at $X$ can be
reached by splitting equally the mass of the $i$-th particle of $X$ along the
first singular vector of the matrix $(\operatorname{\mathrm{Var}}^{(i)}Z)A$,
with

$i\coloneqq\operatorname{\mathrm{argmax}}_{i^{\prime}\in\\{1,\dots,n\\}}\left\lVert(\operatorname{\mathrm{Var}}^{(i^{\prime})}Z)A\right\rVert_{2}^{2}.$

∎