# Multi-Headed Transformer Architectures as Time-dependent Wasserstein Gradient Flows

arXiv: https://arxiv.org/abs/2605.18870 · HTML: https://mlsanity.com/paper/2605.18870/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 63 references that are in the map.
- First posted on arXiv 15 May 2026.

## Builds on (most cited first, 30 of 37)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Robust Speech Recognition via Large-Scale Weak Supervision](https://mlsanity.com/paper/2212.04356/) (2022)
- [wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representations](https://mlsanity.com/paper/2006.11477/) (2020)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation](https://mlsanity.com/paper/1609.08144/) (2016)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [Fixup Initialization: Residual Learning Without Normalization](https://mlsanity.com/paper/1901.09321/) (2019)
- [Transformer Quality in Linear Time](https://mlsanity.com/paper/2202.10447/) (2022)
- [The Lipschitz Constant of Self-Attention](https://mlsanity.com/paper/2006.04710/) (2020)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [The emergence of clusters in self-attention dynamics](https://mlsanity.com/paper/2305.05465/) (2023)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [Why Transformers Need Adam: A Hessian Perspective](https://mlsanity.com/paper/2402.16788/) (2024)
- [Transformers for Modeling Physical Systems](https://mlsanity.com/paper/2010.03957/) (2020)
- [Theory, Analysis, and Best Practices for Sigmoid Self-Attention](https://mlsanity.com/paper/2409.04431/) (2024)
- [A Study on ReLU and Softmax in Transformer](https://mlsanity.com/paper/2302.06461/) (2023)
- [Stable ResNet](https://mlsanity.com/paper/2010.12859/) (2020)
- [Emergence of meta-stable clustering in mean-field transformer models](https://mlsanity.com/paper/2410.23228/) (2024)
- [Clustering in Causal Attention Masking](https://mlsanity.com/paper/2411.04990/) (2024)
- [A Unified Perspective on the Dynamics of Deep Transformers](https://mlsanity.com/paper/2501.18322/) (2025)
- [Robust Pruning at Initialization](https://mlsanity.com/paper/2002.08797/) (2020)
- [Measure-to-measure interpolation using Transformers](https://mlsanity.com/paper/2411.04551/) (2024)
- [Quantitative Clustering in Mean-Field Transformer Models](https://mlsanity.com/paper/2504.14697/) (2025)
- [Synchronization on circles and spheres with nonlinear interactions](https://mlsanity.com/paper/2405.18273/) (2024)

## Built on by (most cited first, 1 of 1)

- [On the Diverse Dynamical Behaviors Arising in Deep Linear Transformers](https://mlsanity.com/paper/2607.18584/) (2026)

## Abstract (arXiv)

In recent years, transformer architectures have revolutionized the field of language processing, opening the door to previously unforeseen possibilities. However, from a theoretical point of view, the mathematical models proposed in the literature often lack direct contact with the actual architectures and depend on strong simplifying assumptions. In this paper, we reduce this gap by modelling the data flow in multi-headed transformer architectures as time-dependent gradient flows for a suitable interaction energy capturing the design of the attention mechanism. The explicit dependence on time allows us to consider different weights for each head and for each layer, without imposing constraints on the initialization method. Moreover, we prove that, under a suitable integrability assumption on the evolution of the weights, each element of the ω-limit set of the gradient flows is a stationary point of the interaction energy at a limiting weight distribution. Finally, we analyse the stability of the gradient flows considering perturbations of both the initial data and the weights. Specifically, on the one hand, we study the robustness of the proposed models with respect to noisy inputs, establishing a continuous dependence of the gradient flows on the initial data and uniqueness of the flows. On the other hand, we prove the Gamma-convergence of the perturbed interaction energy to the unperturbed one, leading to the convergence of the corresponding gradient flows. We complement these theoretical results with numerical experiments that confirm the predicted energy-dissipation identity and clarify the asymptotic behavior of the dynamics in both the autonomous-like (Ornstein--Uhlenbeck) and the genuinely non-autonomous (oscillating-weights) regimes.