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

2026-05-15

AI Summary

This work addresses the limitations of existing theoretical frameworks for multi-head Transformers, which rely on strong simplifying assumptions that fail to capture their true dynamics. For the first time, the data flow within the multi-head attention mechanism is modeled as a time-dependent Wasserstein gradient flow, incorporating a time-varying interaction energy that governs the unconstrained evolution of inter-head and inter-layer weights without requiring initialization constraints. Leveraging variational analysis, Γ-convergence, and non-autonomous dynamical systems theory, the study establishes continuous dependence of the gradient flow on initial conditions and perturbations, proves its convergence, and shows that its ω-limit set corresponds to critical points of the associated energy functional. Numerical experiments confirm the energy dissipation identity and asymptotic behavior under both autonomous and non-autonomous settings, systematically demonstrating the model’s stability against input noise and weight perturbations.

Abstract

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

Problem

• Transformer architectures
• Wasserstein gradient flows
• attention mechanism
• mathematical modeling
• time-dependent dynamics

Innovation

• Time-dependent Wasserstein gradient flows
• Multi-headed Transformers
• Attention mechanism
• Γ-convergence
• Energy-dissipation identity