Applied Analysis Seminar Quantifying concentration phenomena of self-attention dynamics

In this talk, I will study the evolution of tokens in transformers at inference time which is described by a mean-field continuity equation. Leveraging ideas from the convergence analysis of interacting multi-particle systems, with particles corresponding to tokens, we prove that the token distribution rapidly concentrates onto the push-forward of the initial distribution under a projection map induced by the key, query, and value matrices, and remains metastable for moderate times. We quantify this behavior by bounding the Wasserstein distance of the token distribution and the metastable one in terms of an inverse temperature parameter β>0 and the inference time. For the proof, we use Lyapunov-techniques and stability estimates in Wasserstein space together with a quantitative Laplace principle. Our result implies that for time scales of order log β the token distribution concentrates at the identified limiting distribution. Numerical experiments confirm this and, beyond that, complement our theory by showing that for positive temperature and large time the dynamics enter a different terminal phase, dominated by the spectrum of the value matrix.

This is joint work with Albert Alcalde (Erlangen), Konstantin Riedl (Oxford), and Tim Roith (Munich).