Attention layers can be viewed as interacting particle systems, with tokens playing the role of particles. Passing to continuous depth and then to a mean-field description leads to evolution equations for probability distributions of token representations. In several cases, these equations formally resemble gradient flows on spaces of probability measures. Making this interpretation rigorous, however, requires identifying the correct geometry and its domain, especially for softmax attention and for distributions with unbounded support.

I will describe geometric frameworks for two basic attention mechanisms. For softmax attention, we introduce a local-mobility geometry that extends the Otto--Wasserstein construction and identify conditions under which the attention dynamics are genuine gradient flows. For linear attention, the associated Stein operator is highly degenerate; a quotient construction gives a well-defined geometry and differential calculus. These structures also allow us to introduce inertia and formulate Nesterov-type attention dynamics as damped Hamiltonian systems on spaces of probability measures.

For linear attention, the resulting infinite-dimensional dynamics reduce to finite-dimensional matrix equations. This yields conditions for global existence and finite-time growth of the second moment, and provides a quantitative connection between continuum blow-up and the discrete layer depth at which token norms become large. Finally, returning to finite-token models, the geometry suggests explicit corrections to both token and momentum updates. In decoder-only language-modeling experiments on TinyStories and OpenWebText, these geometric corrections improve validation negative log-likelihood relative to corresponding momentum-only controls at equal parameter counts and training-token budgets.

This is joint work with Levin Maier, Wuchen Li, and Gabriele Steidl.

Bio: Viktor Stein is a postdoctoral researcher in applied mathematics at the Technical University of Munich, where he works in Massimo Fornasier's Applied Numerical Analysis group, and a Junior Member of the Munich Center for Machine Learning. He completed his PhD in applied mathematics at TU Berlin in 2026 under the supervision of Gabriele Steidl. His research focuses on gradient flows on spaces of probability measures, optimal transport, particle methods, and geometry-inspired machine learning, with a particular interest in accelerated dynamics and mathematical models of neural-network architectures.