Universal Approximation of Mean-Field Models via Transformers

TL;DR

The paper addresses learning and simulating mean-field dynamics of permutation-equivariant particle systems using transformers. It introduces the concept of an expected transformer to lift finite-sequence models to the space of measures and proves universal approximation bounds for the mean-field vector field, linking finite-particle learning to infinite-dimensional dynamics. Empirical results on Cucker-Smale and fish milling data show transformers outperform baselines in learning the vector field and generalize to more particles, while theory guarantees convergence of the transformer-augmented dynamics to true mean-field evolution. The work advances data-driven modeling of collective behavior by providing rigorous links between finite transformer approximation and continuum mean-field dynamics, with implications for physics, biology, and engineering applications.

Abstract

This paper investigates the use of transformers to approximate the mean-field dynamics of interacting particle systems exhibiting collective behavior. Such systems are fundamental in modeling phenomena across physics, biology, and engineering, including opinion formation, biological networks, and swarm robotics. The key characteristic of these systems is that the particles are indistinguishable, leading to permutation-equivariant dynamics. First, we empirically demonstrate that transformers are well-suited for approximating a variety of mean field models, including the Cucker-Smale model for flocking and milling, and the mean-field system for training two-layer neural networks. We validate our numerical experiments via mathematical theory. Specifically, we prove that if a finite-dimensional transformer effectively approximates the finite-dimensional vector field governing the particle system, then the $L_2$ distance between the \textit{expected transformer} and the infinite-dimensional mean-field vector field can be uniformly bounded by a function of the number of particles observed during training. Leveraging this result, we establish theoretical bounds on the distance between the true mean-field dynamics and those obtained using the transformer.

Figures (4)

• Figure 1: Figure comparing training a two-layer neural network using gradient descent to update the weights and using a transformer to update the weights. The solid line is the median value over 100 trials, while the shaded region is the interquartile range (25th-75th percentile). Left: evolution of the training error during training. Center: evolution of the test error during training. Right: difference between the parameters learned by gradient descent and the transformer.
• Figure 2: Figure comparing the true dynamics of the Cucker-Smaler model versus those obtained from a transformer. The solid line is the median value over 100 trials, while the shaded region is the interquartile range (25th-75th percentile).
• Figure 3: Figure showing the trajectories of ten particles computed for the Cucker-Smale model using the true $\mathcal{F}$ versus the transformer in lieu of $\mathcal{F}$.
• Figure 4: Figure shows the error $\|\mathcal{T}_n - \mathcal{F}\|_*$ for the CS model. Here $(x,y)$ is held fixed while $(u,v)$ is varied in a $11\times 11$ grid.

Theorems & Definitions (25)

• Definition 4.1: Expected Transformer
• Definition 4.2: $1$-Wasserstein Distance
• Remark 4.4: Lipschitz Implies Linear Growth
• Remark 4.5: Example Models
• Definition 4.6
• Theorem 4.7: Universal Approximation
• Remark 4.8
• Theorem 4.9
• Corollary 4.10
• Definition 4.11