Understanding the Regularity of Self-Attention with Optimal Transport
Valérie Castin Pierre Ablin Gabriel Peyré
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.