Universal Approximation of Mean-Field Models via Transformers
Authors: Shiba Biswal
Karthik Elamvazhuthi
Rishi Sonthalia
Venue: ICML 2025 poster (2025)
Keywords: Mean Field Equations
Transformers
Universal Approximation
Collective Behavior
Primary area: deep_learning
TL;DR: We empirically and theoretically show that transformers can approximate infinite dimensional mean field equations

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_\infty$ distance between the \textit{expected transformer} and the infinite-dimensional mean-field vector field can be 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.

Universal Approximation of Mean-Field Models via Transformers
Shiba Biswal Karthik Elamvazhuthi
1 Rishi Sonthalia"
Abstract
which are independent of the particles' identities and are permutation-equuvaiiine Examples include biological enti- ties (Lopez al., 2012), robots (Elamvazhuthi & Berman, 2019), traffic flow (Piccoli et al., 2009; Siri al., 2021), and parameters in two-layer neural networks (Mei et al., 2019). A common approach to simplifying the analysis of such systems is to consider the continuum limit as the number of particles n + oo, resulting in mean-field models rooted in statistical physics. Instead of specifying the dynamics of each agent, the particles are modeled using probability measures. This paper learns the mean-field dynamics of particles via particle trajectories using transformers.
This paper investigates the use of transformers to approximate the mean-field dynamics of interact- ing particle systems exhibiting collective behavior. Such systems are fundamental in modeling phe- nomena across physics, biology, and engineering. including opinion formation, biological networks, and swarm robotics. The key characteristic of these systems is that the particles are indistin- guishable, leading to namics. 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 ex- periments via mathematical theory. Specifically, transformer
dy-
Let Ω C Rd. Consider vector field F $ x P(Ω2) - Rd. The following equation describes the general mean-field behavior of interacting particles evolving on S2
de dt =(2,1))) =(0) 20 140
we prove that if effectively approximates the finite-dimensional vector field governing the particle system, then the L2 distance between the expected transformer and the dnfiite-dimensionall mean-field vector field can be uniformly bounded by 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 trans- former.
where S denotes the state of particle, μ E P(2) is the distribution of the particles time and 140 is the initial distribution. The above equation can be approximated by finite particle-level model. Let z (=1,-.., =n) E Ω represent the state of particles, and consider the following,
& = F 1/"),
where the empirical distribution of the n-particle system given by 1n == _ Σm 62, E DD?).. Several works learn the mean-field dynamics by learning the particle-level dynamics. Pham & Warin (2023) constructed neural net- work via binning as approximations of the vector field and proved approximation capabilities when restricted dy- namics on measures that admit densities. The works of Feng et al. (2024); Lu et al. (2019); Miller et al. (2023) present a kernel- method for identifying the dynam- ics of interacting particle systems. Whereas, Messenger & Bortz (2022) presents a form of the SINDy algorithm for identifying mean- mean-field dynamics of interacting particle systems. Kratsios al. (2022) constructed a probabilis- tic transformer that can satisfy additional constraints via maps on measures. Additionally, Furuya et al. (2025) use continuous version of transformers and attention ( (Vuckovic et al., 2020; Geshkovski et al., 2023) and provide universal approximation of measure theoretic maps. Similarly (Adu Gharesifard, 2024) have proven an approximation result for
Introduction
The identification of dynamical system models for physical processes fundamental application of machine learning (ML). Of particular interest are systems of particles exhibit- ing collective Schaviors-such ion dynamics, and consensus. These systems involve large number of particles or agents that follow identical dynamics,
swarming, flocking, opin-
Eqqul contribution 'Theoretical Division, Los Alamos National Laboratory, Los Alamos, NM, USA "Boston Col- lege, Boston, MA, USA. Correspondence to: Shiba Biswal <sbiswal gov>, Rishi Sonthalia
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Universal Approximation of Mcan-Field Models via Transformers
measure theoeretic maps arising from solutions of continuity equations on the sphere using a different architecture.
noted by wen y VX...XX The support of a measure
times
E P(RR), denoted by supp(v), is the smallest closed set S C Rdd such that j(R4 \ 5) = 0. Given mea- surable map X : R + Rd and measure μ E P(R)), the pushforward measure X# E P(RR)) is defined by X **)(AA) == 1 (X- for every Borel measurable set A CR..
This paper explores the use of transformers to approximate the dynamical systems that govern the collective behavior of interacting particles with permutation-equivariaa dynamics. We compare the transformer against models from Pham & Warin (2023) and Messenger & Bortz (2022) as well as two additional permutation equivariant baselines models on synthetic and real data from the Cucker-Smal Smale model for swarming and milling. We prove theoretical guarantees for the transformer by lifting it from gequence-to-sqquene map to a map on measures, by taking the expectation of transformer with respect to product measure. We refer to this as the expected transformer.
For a vector y RR, the i-th component is denoted by y/- The p-norm, p E [1, 0o), and co-nomm of y are, respectively, IulI Ei= Iui |yll max; lusl- Boldface letters, such as z, denote elements in Rd representing collection of n vectors in Rd. The j,j-th element is denoted by z. The corresponding p-norm (p E [1, oc)) and co-norm of z are, respectively, ||z|| = Ei= Ej=1 |z:y || = maxi.j [ayl
a
The main contributions are follows:
1. We empirically show that transformers approximate the vector field better compared to other network architec- tures (see Table 1).
A function f R + Rd is permutation equiv- ariant if for any permutation E Su, where S, is the symmetric group on n elements, and for any x (II,...,,Tn) E Rdxn, we have ((Bfo(())....TTTt))
2. We define a continuum version of the transformer as an expectation of finte-dimenniioal transformers (see (9)).
3. We establish approximation rate bounds of measure- valued maps by this expected transformer (see Theo- rem 4.7).
where fi(x) denotes the i-th component of the output. We denote by C* (R<) the space of k-times continuously differentiable functions Rd. The space of continuous functions with compact support is de- noted by (Rd)..
(f.c)(().......f
4. We mathematically show that the solution to the mean- field model can be approximated by approximating the vector field by the expected transformer (see Theo- rem 4. 14 and Figure 2).
2 Problem Formulation
The rest of the paper is organized as follows. Section 1.1 defines the notation used in the paper. Section 2 defines the problem, Section presents our numerical experiments, and Section 4 presents our theoretical results.
We briefly introduced the problem Section 1, we restate the equations again for clarity. Let Q c Rd. Consider a vector field F Ω x P(2)) + Rd The general mean-field behavior of interacting particles evolving on Q is given by
1.1 Notation
2(0) = 20 ~ 140- dz = F2, dt
(1)
Let Rd denote the d-dimensional Euclidean space, and let I., denote the set of positive integers. The diameter of subset A IR is defined as diam(A) := supillz yll v,yE E A} and let (2) be the closed ball of radius > 0 centered at 2 E Rd4
where E Ω denotes the state of each particle, /a E P(2) denotes the distribution of the particles, and 2g is the initial state that is distributed according (4. the initial distribu- tion. The inter-particle interactions are modeled through ju; specifically, the dynamics of each particle are influenced by the distribution of all other particles.
We denote by P(Rd)) the set of all Borel probability mea- sures on Idd Similarly, for a subset C Rd. P(2) de- notes the set of Borel probability measures on Q. The subset of probability measures with finite p-th moments, Mp(x)= := (fau dx(x)) 1/p is denoted by Pp(R)) {NE E P(RR)): My(() < oo). Let (o( be the set of probability measures with compact support. The set of empirical measures formed by finite sums of n Dirac- delta measures is denoted by D (RR) == (v E P(R'): :
Corresponding to (1), the continuity equation describes the evolution of the distribution Jr:
Op 2t + V (F=,(()))))) = 0, p(0) = 140-
(2)
For a finite final time > 0, ,, [0,7] + P())) denotes
the solution of the continuity equation (2) over the time interval [0, r].
v = + Dm 8z,, &4 E Rd). Similarly, the set (2) denotes the set of empirical measures on Q c Rd. For E P(RR)), the n-fold product measure on (Rdy" is de-
In this paper, we propose to use transformers to approximate the map (1) and the system (2). However, transformers
2

---

Universal Approximation of Mean-Field Models via Transformers
are defined on sequences of vectors in IR,, whereas the map F is defined on 52 x PS)), an infinite-dimensional space. Therefore, consider finite-dimensional approximation of (1) via a particle-level system. Specifically, consider a n-particle system where the state of each particle i E {I.... n is given by 2; E Q. We assume that each 2; is independently sampled from the distribution μ E P((2). Let z (21, (21...., 2n) E Ω denote the collection of particle states. The empirical distribution of the n-particle system is then given by
Next, consider the n-dimensional (finite) approximation of the model above.
dx;
vi, Osisn
dt
du: dt
R
(6)
0(| - x2,1( - v.) Σ ==1
To generate the data, we compute trajectories for 500 ran- dom initial conditions (z)(0), (O))), chosen uniformly at random from f = [0, 1] [0. 1]. For each initial condition, we solve the differential equation (6) for n 20 agents over time horizon [0. 100) using SciPy's solve ve i ivp. Hence, for each initial condition, we obtain position x;(() and ve- locity Us (t), for each time step t.
1
1/
82, E (S2).
(3)
The particle-level dynamics on Hd according to the map defined in (I), for single particle i, can be written
Fish Milling Model The second example that we consider is the fish milling model (D' DDOOrogga al., 2006) which is version of the Cucker-Smale model. In contrast the previous example, this example uses an actual experimental dataset (Katz et al., 2021). The dataset includes the position x E RR2, velocity e E R and fish identifier for n = 300 fish over 5000 time steps. While the experiment tracks 300 fish, real-world data is incomplete at certain time points. Specifically, the positions of all fish are not available at all times, and the fish move in milling pattern, some fish disappear and reappear. When a fish reappears, is assigned new index, and its previous identity is lost. On average, only about 220 fish are observed at any given time. As a result, we have 5000 data points, but the number of particles varies at each time step.
= F(z,, ",,, (0) Po
(4)
Note that the collection of random variables (2,) is permu- tation equivariant because the joint distribution of (z;) is invariant under any permutation of the indices.
3 Numerical Simulations
In this section, we present experiments demonstrating that transformers can learn mean-field dynamics and generalize to cases involving more particles than those seen during training. We also compare the transformer with Graph Neu- ral Networks (GNNs) and prior models from Pham & Warin (2023); Messenger & Bortz (2022).
Data Next, we will describe how we assemble the data from the Cucker-Smale and fish milling model for the ML architectures. We note that F maps [x o|" + [ 8) T
3.1 Learning the Vector Field
Our first goal focuses on learning the vector field F. To- wards this, in this experiment, we use two datasets: first, synthetic dataset generated from the Cucker-Smale model (Cucker & Smale, 2007), and second, real data of fish milling (Katz et al., 2021).
Let =t) [x((t) r())²² and ε(t) = [=(t) i())]
From the simulations described above, we obtain the col-
lection z(1) (z1 (4),... .. Zn (I)),, where each =+(1)
[1;(1) =,(t))² To construct ²(t), we use the centered- difference formula:
Cucker-Smale Model The first example that we consider is the well-studied 2d x n-dimensional Cucker-Smale (CS) equation that models consensus of n-agent system (Cucker & Smale, 2007). For d = 2. he vector field FF R x P(RR') > I is given by
2,) 21(t + At) 2 At)
241
where At is the time step. We assume that at each time step t, we may know =tt) for only a subset of the particles:as is the case for the fish milling model. Moreover, we do not know the functional form of F(z, Let I(t) == (i : (t), (t) are known]. and we redefine z(t) to mean collection of particle states for which we have data, and similarly for i(1),
F(I,v, p)) [ Sgs 20[ -il---llluwwduuuou w)du(u.e)
(5)
H
((r)= (s² + 1230
Where E 2² and v e R denote the position and velocity of each agent, respectively. Here, á a neg-nigae function, is the interaction potential that determines the inter-agent interaction, and H, s, b are parameters (set to 1 here).
z (t) (t(()) : I(t)} and k(t) = (2)(t) i I(t)].
The data so obtained are split into an 80-20-20 split for training, validation, and testing.
3

---

Universal Approximation of Mean-Field Models via Transformers
Model
Cucker-Smale
Fish Milling
Transformer
(1.9 0.3) 10
=:2=0.00 x 10
Cylindrical NN
(1.1 10-3
( x 10
GraphCony m
(2.0 10"4
(1.2 x 10
GraphConv 20
(2. (4.1 x 10-3
>1
TransformerConv m 3
10
(6.5 = 1.7) x 10
TransformerConv m 20
(3.3 x 10-6
(1.0 + 0.3) x 10-1
FNN
(6.7
10-6
N/A
Kernel
(3.0 0.1) 10-3
N/A
Table 1. Table showing the mean-squared error in approximating the map Fn. (14) for the different data and models.
0.54
0.3
5GD
SGD
0.22
Tansformer
0.50
erl"
Locs
0.10
Locs
0.08
0.28
18:00
0.99
9:29
0.4
0.32
0.00
:00
200
302
400
500
0.30
200
300
402
500
150
200
305
600
500
iteration
iteration
Iteration
Figure Figure comparing training a two-layer neural network using gradient descent update the weights using 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 training error during training. Center: evolution of the test error during training. Right: difference hetween the parameters learned by gradient descent and the transformer.
Next, we will describe the ML architectures that we use to approximate F. In partieular, the transformer will be compared against four additional permutation equivariant baselines.
(t) is constant for all t (i.e., for the CS data), evaluate fully connected feedforward net and kernel regression with sine, cosine, and polynomial basis. For both of these methods, we concatenate the vectors in z(t) and zt() to create the training data.
Transformer In Appendix D, we define a transformer network, T: Raad + gaxk (k not necessarily equal to d) At each time step t. we provide the transformer with input 2() and have it predict (tt). The model is trained using mean-squared error loss.
Training Details and Hyperparameter Search For each model and dataset, we conduct hyperparameter search to optimize depth, width, and learning rate. We train the mod- els using mini-batch Adam and cosine annealing learning rate schedule. For kemnel regression, we explore different numbers of basis functions for each type of basis function. See Appendix E for more training details.'
Graph Neural Network (GNN) Given set of particles (t), let Gt(t) and (t) represent the graphs based on the three nearest neighbors and the twenty nearest neighbors, respectively. We use two graph neural networks (GNNs) that, at each time step 4, take (G. ((), z())) as input, where is either 3 or 20, and aim to predict the node labels (t). This forms graph regression task where the target size depends on the number of nodes in the graph.
account for the randomness introduced by initialization and training, we conduct five trials for each hyperparameter setting. The best hyperparameter configuration is selected based on performance on the validation data.
Results Table presents the mean-squared error for the different methods, across the two data sets. We observe that the transformer has the lowest mean-squared error among all of the models. Therefore, we conclude (empirically) that transformers can effectively learn the vector field.
We employ two common GNN architectures: the Graph Convolutional Network (GraphConv) (Morris et al., 2019) and TransformerConv (Shi al., 2021). Note that, similar GAT (Veličković al., 2018), TransformerConv applies attention only local neighborhoods.
'odee can be found at: rsonthal/ MaatthallMMMa--Peee -eean--nnneddd Field- Field-Trans
Other Baselines The final three baselines we consider are: cylindlrical nets from Pham & Warin (2023). Additionally, if

---

Universal Approximation of Mean-Field Models via Transformers
3.2 Simulating Mean-Field Dynamies
34-3
Our second goal is to approximate the solution to the con- tinuity equation (2). We use two synthetic datasets: the Cucker-Smale dataset from Section 3.1, and the training dynamics for a two-layer neural network (Mei et al., 2019).
0.0
05
1.0
1.3
2.0
2.
60
0.5
1.0 2.0
25
50
Cucker-Smale After obtaining a transformer that ap- proximates the vector field in (5) (as described Sec- tion 3.1), we solve the differential equation (6) twice, using SciPy's solve_ivp fuucction for new set of initial con- dition. First, with the true vector field F, and second, using the transformer 7 in lieu of F. We note that solving (2) for Dirac valued initial distributions, (2,(0), is equivalent to solving (6) for initial condition (0).
Figure 2. Figure comparing the true dynamics of the Cucker- Smaler model versus those obtained from transformer. The solid line is the median value over 100 trials, while the shaded region is the interquartile range (25th-75th percentile).
Training 2-Layer Neural Network Consider a two-layer network f(x) = E'= (ro(""))) Let 0g (a,, w.) be the parameters (the states of this model) and @ (OI...., 0n). In this model, we treat each 0; as parti- cle, with its distribution evolving according to the following continuity equation, as derived by Mei et al. (2019),
The Passriis
Figure 3. Figure showing the trajectories of ten particles computed the Bucker-Smale model using the true JF versus trans- former in lieu of F.
j V NNVVV..VVVVO.. (O,,).
Here, E(t depends the learning rate schedule and V (1) == EEyy [y [x]
while the error in the velocity seems to plateau and even decrease slightly. Additionally, the initial interquartile range is small, but it grows over time.
+ / FE..y (x "))(("" (oo(TTW)]]] w)] du(áá, w).
Here, the vector field we wish to approximate is F(0, (1) Vo p)). To generate the data, let X be a d times 1000 dimensional matrix with i.i.d. Gaussian entries. Let ʃ be a fixed two-layer teacher network with sigmoid activation (o) and let Y f(X). We then initialize 50 two-layer net- works with hidden nodes, and train them for 500 epochs using the data X, Y. We use the MSE loss and gradient descent to train the network. This gives us 2500 training data points ( )). We set the hidden layer with n = 100 and use an input dimension of d 10.
Finally, we note that we trained the model with 20 particles but simulated the model with 30 particles. Hence we see that the model can generalize to more particles than that seen during training. This implies that the transformer is able learn the interaction kernel between the particles.
Next we simulate training a two-layer neural network us- ing Transformer. Figure 1 illustrates the training and test loss as we train the network. The blue line represents the model trained using GD, while the orange line corresponds to the model trained with the transformer. Note that the transformer model does not compute any gradients. We trained models with hundred different random initial- izations. The solid lines indicate the median and the shaded regions represent the interquartile range. The right-most plot in Figure shows the Frobenius norm of the difference between the parameters learned using GD and the trans- former. The figure demonstrates that the maximum norm of the difference is at most 5 10 10-5, even after 500 iterations. Therefore, observe that the transformer has learned to approximate the dynamics of GD for two-layer network.
Once the transformer trained, we test follows: We reinitialize two-layer network, providing its weights input to the transformer. The transformer's outputs are then treated as gradients, which we use to perform gradient descent. Finally, we compare the results with the dynamics obtained from performing true gradient descent.
Results Figure shows the evolution of 30 particles using the Cucker-Smale equations (6) along with the transformer approximation. Notably, the dynamics obtained using the transformer tracks the true dynamics near exactly. Figure 2 plots the 2-norm between the positions coordinates x, y and velocities u, v. The figure indicates that the error is generally quite small (< 10 4), although it increases over time. This increase appears to be linear for the position coordinates,
4 Theoretical Results
Our experiments demonstrate that the transformer can ap- proximate the finite-dimensional particle dynamics (4). In
5

---

Universal Approximation of Mcan-Field Models via Transformers
this section, we will prove that the infiniteddimensional vec- tor field F (1) can be approximated by the expected output of transformer, which we shall denote Tx (z, u)) Further, we show that the solutions to the continuity equation (2) can be approximated by the solutions to the following appraxi- mate continuity equation, defined using Tm.
Definition 4.2 (1-Wasserstein Distance). Given two proba- bility measures J4, (Pp(2)) on metric space (Ω, ll2) where d is the metric on 9t the 1-Wasserstein distance be- tween and is defined as
/ QxD
Wi (i(A =
inf yETT(N))))
(1--vl227(4,11 v), (12)
dy Ot + V (T) = 0, (0)) = Ho-
(8)
where IT v) denotes the set of all couplings (transport plans) ? on Q Q with marginals μ and v.
4.1 Lifting Transformers to the Space of Measures
We now state the main assumptions required for our analysis: Assumption 4.3 (Regularity and Growth Conditions). As- sume that the vector field F : : x Pp(2)) + R satisfies the following conditions:
Traditionally, transformers are defined on sequences of vec- tors in Rd. However, the map we wish approximate, JF, is defined on fi x (22). Therefore, we lift the standard trans- former to act on (X)(( via expectation operation. This will allow us to lift any transformer' T', - + R(=+1)xd to model Tx 92 x P(S2) + Rd.
a) (Lipschitz Continuity) There exists a constant L such that for all z, y E Q2 and p.v E Pp((2).
[F( Fly, 2) < L (1 ylla + Wi (1)
Definition 4.1 (Expected Transformer). Given a transformer T: g + (xxx1)x and prescribed sequence length n, define the expected transformer Tx 92 x P(l2) + R2 by
This condition implies the following. b) (Linear Growth) There exists a constant M > such
that for all æ E Ω and E PP)))),
Tn Tn p)) Exrpon I(T (Lr; 1)
(9)
[F(x, S M (1 ||o||2 + Mi (n)).
= ff 10
(T ([=; 21....,, 2nl)), dp(21) du(zn),
where M1 (n) = Sa ly/ (j((y) is the first moment of p.
where [x; z] denotes the concatenation of x and (Z1,..., 2n) form an input sequence of length n + 1, and (T ([x; z))); denotes the first output vector.
These assumptions are standard in the analysis of mean-field models and differential equations in general and are needed to guarantee the uniqueness of solutions.
Remark 4.4 (Lipschitz Implies Linear Growth). Since f is compact, Assumption 4.3a) implies Assumption 4.3b). However, in some cases, we only need the weaker assump- tion of Linear Growth, hence we explicitly state it.
Some prior works (Geshkovski et al., 2023; Furuya et al., 2025) have defined transformers T Q P(t) - Rd through continuous version of self-attention I. For Ω and μu E P(S)), T is defined as
Remark 4.5 (Example Models). These assumptions are sat- isfied by the Cucker-Smale model as well as the training 2-layer neural network model. Specifically, Theorem 2 of piccoli et al. (2009) shows that Cucker-Smale model satis- fies Assumption 4.3a) and Lemma 2 of Mei et al. (2019) shows that the training 2-layer neural network model (7) also satisfies Assumption 4.3a)
[[z, == Z(e, ()) J Att ([æ; yl) du(y). (10)
1
where 7(x, μ4) a normalization factor and Att is the at- tention layer (D.1). Then the transformer 7 in Geshkovski et al. (2023); Furuya et al. (2025) is defined as
To establish our approximation results for functions H. Ω x P(2)= + Rd, we define the following norm. Definition 4.6. Given a function 7. Ω x P(2)) + Rd, we define its norm by
f(x, 14) == FCC, oT, .... o FC,, oT (z),
(11)
where Io, and FC2, are attention and red-ffowwrr layers with parameters ij and E,, respectively.
|| = sup sup (M(x,,))l rEQUEP(S)
(13)
4.2 Assumptions
To state our result, we require some assumptions the map F in (1). The key assumption we make is that F is Lipschitz continuous with respect to the probability measure. To formalize this, we require a metric on the space of probability measures P(2)). A commonly used metric is the p- Wasserstein distance.
4.3 Approximating the Vector Field F
To approximate the vector field F, we define, for fixed the finite-dimensional map F. 22" - BR as
(z = [F(21, v2)
F(zn, (14)
6

---

Universal Approximation of Mcan-Field Models via Transformers
We now state our main result regarding the universal approx- imation of the mean-field vector field F by the expected transformer T
with the weak * topology. Then for all > O, there exist I and parameters (05,5))}=1 such that
(T - (1, FF
Se, V(r,)) E S2 P(R*)
Theorem 4.7 (Universal Approximation), Ler S C R be compact set containing O. Let F Q P(2) + Rdd satisfy Assumption 4.3a) for given p. Given transformer let
where the parameters 0,,5; depend on the dimension
T: gn+11 + [[(++))xd
To compare with Theorem 4.9, we state the following corol- lary to Theorem 4.7.
:= sup zEfcot
[T(z)= (2)l
(15)
Corollary 4.10. Let E >>>>>>>> 0 and 2 1. Let C Rd be a compact set containing O. Let F Z x P(2) + Rd satisfy Assumption 4.3a) for a given p. Then there exists a transformer T with depth ((1), one attention layer with width ((d) such that the expected transformer Tn satisfies (16) with & =
Then, forq > P ttheee exists a constant Cp, q, d), depending only on p, q. and d, such that for all n 2 1, the correspond- ing continuum version Tn: St P(S)) + a (9) satisfies
|TT - FFIl L diam(52)
1
+ cG(n,p.))
n2 +
While Corollary 4.10 and Theorem 4.9 are about two differ- ent models, they share notable similarities while exhibiting key differences. Both results feature (cd) width for the attention layers, independent of and n, and neither pro- vides bounds for the width of feedforward layers. However, there are key differences. Our work leverages prior results to establish a bound on network depth, which Furuya et al. (2025) does not. Moreover, we note that providing a bound on the width of the feedforward network, in our case, is straightforward, owing to recent developments that provide bounds on both width and depth (Augustine, 2024).
p 5/2, q # 2p
v vn
G(n.p.q)
log(n + 1) p d/2, q # 2p (16)
p p< d/d, q + 2
Remark 4.8. In the Theorem 4.7 above, we observe that the approximation of infinite-dimension maps F by finite- dimensional transformers T depends on two key quantities. First, it depends on how well 7 approximates the finite- dimensional map Fu. This corresponds to the & term in the bound (15). Due to universal approximation results for transformers, this term can be made arbitrarily small. For example, see Theorem 4.3 in Alberti al. (2023).
Further, we provide detailed error rates that depend on the length of the sequence n that the transformer is trained on. However, we note that Furuya et al. (2025) impose weaker assumptions for the map F.
Second, the approximation depends on the convergence rate of W1 (1(,, (,),, which, by Theorem of Fournier & Guillin (2015) is tight and depends on n. p: and d. Furthermore, we observe that the stronger the regularity of the map F (i.e., larger the constant p). the better the approximation. The best rates are obtained for -= _ + +1] In which case the rate becomes independent of the dimension d. Conse- quently, if we use longer sequences, we obtain an improved approximation of the vector field F.
4.4 Approximating the Mean Field Dynamics
We build upon our previous approximation results to prove that the solution μ T to the approximate continuity equa- tion (8) approximates the solution μ² (t) continuity equa- tion (2). To formalize this, we first introduce an appropriate notion of a solution to the continuity equation.
Definition 4.11. Let μ E ((p)) be measure- valued function, is called Lagrangian solution of the con- tinuity (2) if there exists X [0, 7] x Red - Rd. referred to as the flow map, satisfies
We empirically verify Theorem 4.7 using the transformer trained for the Cucler-Smale model in Section 3.1. We test five (a., v) E [0.1]2 values and vary (u, v) over an x 11 grid. Using Gibbs sampling, we approximate the expected transformer and compute the error with the true vector field (5). Figure 4 shows the heat maps, with a maximum error of 0.025, while the bound from Theorem 4.7 is at least 0.05.
X(t,z)= = x [ F(((s
(17)
for all x E Rd and pl() = X (1,-) pn for all E [0,
Comparison with Result From (Furuya et al., 2025): The concurrent work (Furuya et al., 2025) proves the fol- lowing approximation result for continuous maps F by the continuum version of the transformer Î (11).
Existence and Uniqueness of Solutions Under Assump- tion 4.3a), it is known that there is a unique Lagrangian solution corresponding to (2), see Proposition 4.8 in Cav- agnari et al. (2022). As stated in Remark 4.5 the numerical examples considered in Section satisfy this assumption.
Theorem 4.9. Let Ω C Rd be a compact set and F+ : Ω x ((Rd) -> dR.d be continuous, where P(R.d) is endowed
7

---

Universal Approximation of Mcan-Field Models via Transformers
300
A20
300
Figure 4. Figure shows the emrrr ||T, Fll- for the CS model. Here (2, y) is held fixed while (u, u) is varied in a 11 x 11 grid.
However, to ensure that solutions of (2) can be approxi- mated, we also need existence of unique solutions for the approximating continuity equation (8). Unfortunately, the expected transformer might not be globally Lipschitz as re- quired in Assumption 4.3a) since T(2, z in (9) might not be globally Lipschitz (Kim et al., 2021). ) Hence, the continuity equation (8) with the expected transformer as the vector field may not have a unique Lagrangian solution in general.
where uf and ym are the solutions to (2) and (8), respec- rively, and the estimate (19) independent of po-
Remark 4.15. In essence, Theorem 4. 14 states that by choos- ing a large enough ball of radius R that covers the support of /o, and selecting and appropriate K, as a function of R, final time T, and regularity constant M,, we can ensure that if Tn approximates F on the ball of radius K, then Tn can be used to simulate the dynamics (2) over the time interval [0, r]. We observe that the error bound (19) grows expo- nentially. Therefore, a small approximation error for the vector field JF implies a small approximation error (19) for the solution of the continuity equation over [0, r). However, the bound (19) also depends on the regularity of F, namely p and L. Therefore, the more regular the vector-field F, i.e., larger and smaller L, the better the bound (19).
The goal of the following theorem is to show that if Assump- tion 4.3b) holds locally on sufficiently large set, then we can prove existence of a unique solution to (8), even if the expected transformer is only locally Lipschitz. This will be used later to show the ability of the expected transformer to approximate solutions of the continuity equation (2), using its universal approximation property.
Theorem 4.12. Suppose po E PP)) is such that the sup- port of no is contained within (0) C Q. for some R >>00 O. Additionally, let K satisfy K > (R (R+2.44 2. Further- more, assume that
We can combine Theorem 4.7 and Theorem 4.14 obtain the following corollary.
Corollary 4.16. Suppose F satisfies Assumption 4.3a) for some p. Let the support of pn0 be contained within (n( C Ω2 for some R > O, and let K E R be such that K > (R +
(18) for all E (0) and all u with support in Bx(0), Then, there exists a unique Lagrangian solution to the approximate continuity equation (8) such that supp μ² (t) C Bc. (0) for all t E [0, rI. where Ct- = (R + t.
I1T 1)1l2 S A (1 (1+111011 + M1()),
2. )e S. Given a transformer' T: -2+ - (x++))xdd
let E := sup
[[T(z) - Fn+1(z))p 12. If & is small enough
2EfM
and n E E is large enough, then
Wi(H (1), (0,x"((1)) < 2(E + ((n, exp(2
Remark 4.13. In the proof of Theorem 4.14 it will be shown that the bound (18) on the expected transformer directly follows from Assumption 4.3a) on F wheneverT, is close enough to F in the uniform norm.
1 G(n. p.))
where 8(n, L(2K))
where
u n
G(n,p.g) is as per (16).
We are now ready to state our main theorem regarding the approximation of mean-field dynamics using transformers.
5 Related Works
Theorem 4.14 (Mean Field Dynamics Approximation Us-
ing Transformers). Let E > 0 be small enough and n. 2 1. Suppose that F satisfies Assumption 4.3a) for some p. As- sume that the support of po contained within (n( C. Q. for some R. O, and let K I. be such that K> (R 2. M If the transformer Tn satisfies the condition
In this section, we provide more in-depth comparison with similar recent prior work Furuya et al. (2025); Kratsios & Furuya (2025). We highlight key differences below.
Role of Measure and Approximation Target: In this work the measure f is an input argument representing the state distribution in mean-field system, directly influenc- ing the vector field F(z, p)). Specifically, this paper targets the approximation of the vector field F(z, ()) governing mean-fiel field dynamics and the subsequent approximation of the dynamical system's m's evolution (solution to the continuity
TT p) - F(x F(I,,))l++ <s.
for all 2 (0) and μ ((Bx(0)). Then we have that
(1), s etexp(229t)
(19)
8

---

Universal Approximation of Mean-Field Models via Transformers
equation). We require the target vector field F(z, (i) to be Lipschitz continuous w.r.t. spatial and measure arguments (using Wasserstein distance). This directly models physical or biological system interactions. A key strength is its ap- plicability to general Borel probability measures P(()) on a compact set f2.
sults and numerical simulations, we established that trans- formers be powerful tools for modeling and learning the collective behavior of particle systems.
A Aeknowledgment
Shiba Biswal acknowledges partial support from the Labo- ratory Directed Research and Development program of Los Alamos National Laboratory (LANL) under Projects No. 20230338ER and No. 20250638DI, and from the U.S. De- partment of Energy/Office of Science Advanced Scientific Computing Research Program.
Furuya al. (2025) also has yl as input and aim approx- imate general continuous in-context mappings A (u, They require A to be continuous w.r.t. weak-star topol- ogy (plus Lipschitz conditions on contexts for masked case). Their focus is broad representational power.
For Kratsios & Furuya (2025) the measure i1 is an input argument to the target function [(,, 2).. However, it's re- stricted to the specific class of Cermutation-Invariant Con- texts within geometrically constrained domain, aimed at analyzing general in-context function approximation.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Map Definition and Transformer: Our work defines the map 14 ~ Tn (Measure to Vector Field) using the Expected Transformer, derived by taking the expectation of a standard, finite-sequence transformer T. This provides a practical link between standard architectures and measure-theoretic inputs.
References
Adu, D. O. and Gharesifard, B. Approximate controllability of continuity equation of transformers. IEEE Control Systems Letters, 2024.
Furuya et al. (2025) define transformers directly on the space of probability measures using a measurc- formulation with continuous attention layers. Kratsios & Furuya (2025) defines the map from finite vector spaces to finite vector spaces. However, the input and outputs are interpreted measures on discrete sets.
Alberti, S., Demn, N., Thesing, L.,, and Kutyniok, G. Sum- former: Universal approximation for efficient transform- ers. In Topological, Algebraic and Geometric Learning Workshops 2023, pp. 72-86. PMLR, 2023.
Augustine, M. T. survey on universal approximation theorems. arXiv preprint arXiv:2407.12898 2024.
This work provides quantitative bounds the vector field approximation eπrоr that explicitly show convergence as the number of particles increases, linking the error to the quality of the underlying finite transformer. Furthermore, it connects this to the approximation of the system's dynamics via Gronwall's lemma.
Guarantees:
Cavagnari, G., Lisini, S., Orrieri, C., and Savaré, G. La- grangian, eulerian and kantorovich formulations of multi- agent optimal control problems: Equivalence and gamma- convergence. Journal of Differential Equations, 322: 268-364, 2022. ISSN 0022-0396.
Kratsios & Furuya (2025) provides quantitative probabilis- tic bounds on the output approximation error, dependent the target function's 's modulus of continuity and the domain's 's geometry, focusing on the network size needed for given precision. While Furuya et al. (2025) show a single trans- former architecture (with fixed dimensions/heads) works uniformly for an arbitrary number of input tokens (even infinite) for given precision E.
Cucker, F. and Smale, S. Emergent behavior in flocks. IEEE Transactions on automatic control, 52(5)) 862, 2007.
D'Orsogna, M. R., Chuang, Y. L.. Bertozzi, A. L., and Chayes, L. S. Self-propelled particles with soft-core interactions: Patterns, stability, and collapse. Phys. Rev. Lett., 104302, Mar 2006. doi: 10. 10.1 101110/yyssee..ett 104302. URL https: : /link. aps. org/ doi/ 10. 1ww.1110//hyyyeeie.
96 104302.
6 Conclusion
Elamvazhuthi, K. and Berman, S. Mean-field models in swarm robotics: A survey, Bioinspiration Biomimetics, 15(1)x11001,, 2019.
This paper demonstrated the efficacy of transformer archi- tectures in approximating the mean-field dynamics of inter- acting particle systems. We showed that fiite-dimennional transformer models can be lifted approximate the infinite- dimensional mean-ficeld dynamics. Through theoretical re-
Feng, J., Kulick, C., Ren, Y., and Tang, S. Learning parti- cle swarming models from data with gaussian processes. Mathematics of Computation, 93(349):2391-22 2024.

---

Universal Approximation of Mcan-Field Models via Transformers
Fournier, N. and Guillin, A. On the rate of convergence in Wasserstein distance of the empirical measure. Probabil- ity theory and related fields, 162(3):707-738, 2015.
go neural:
Higher-order graph neural networks. In Pro- ceedings of the AAAI conference artificial intelligence, volume 33, pp. 4602-4609, 2019.
Furuya, T., de Hoop, M. V., and Peyré, G. Transform- ers are universal in-context learners. In The Thirteenth International Conference on Learning Representations, 2025. URL https: htttps://oppeeeeeeetett wwwpeerreeee.eeetffoeeeeeeeeeeee id=6s4WQD1L22
Pham, H. and Warin, X. Mean-field neural networks: learn- ing mappings Wasserstein space. Neural Networks, 168:380-393, 2023.
Piccoli, B., Tosin, A., et al. Vehicular traffic: review of continuum mathematical models. Encyclopedia of complexity and systems science, 22:9727-9749, 2009.
Geshkovski, B., Letrouit, C., Polyanskiy, Y., and Rigollet, P. A mathematical perspective on transformers. arXiv preprint arXiv:2312.1077 2023.
Shi, Y., Huang, Z., Feng, S., Zhong, H., Wang, W., and Sun, Y. Masked label prediction: Unified message passing model for semi-supervised classification. In Proceedings of the Thirtieth International Joint Con- ference on Artificial Intelligence, IJCAI-21, pp. 1548- 1554, 2021. URL https: http:///doo.rg111..2..4. i ijca 2021 /24
Katz, Y., Kolbjørn, T., Christos C, I, Huepe, C., and lain D, C. Fish schooling data subset: Oregon state university, http::httttttt.ttttctteeee wtw..irrrrrr..oooottttttte.. /otp//oonceem//atassts//tt.
2021.
Kim, H., Papamakarios, G., and Mnih, A. The Lipschitz constant of self-attention. In International Conference on Machine Learning, pp. 5562-5571. PMLR, 2021.
Siri, S., Pasquale, C., Sacone, S., and Ferrara, A. Freeway traffic control: survey, Automatica, 130: 109655, 2021.
Kratsios, A. and Furuya, T. in-context universality enough? MLPs are also universal in-context. arXiv preprint arXiv:2502.03327, 2025.
Veličković, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., and Bengio, Graph attention networks. In International Conference on Learning Representations, /ttts:/:///ppeeeee//// http:////ppeeeiiieeeemm net net/forun? forum? id=rJXMpixC2.
2018. URL https:
Kratsios, A., Zamanlooy, B.. Liu, T., and Dokmanic, Universal approximation under constraints is possible with transformers. In The Tenth International Conference Learning Representations, 2022. URL https:// openreview net forum? id JGOBCV6559.
Vuckovic, J., Baratin, A., and des Combes, R. A math- ematical theory of attention. ArXiv, 0bbs/2007.02876 2020. URL https: http://///aa.ssssiiieeeeeeee semanticschola: org http/CooppuIDD::2333599
Lopez, U., Gautrais, J., Couzin, I. D., and Theraulaz, G. From behavioural analyses models collective motion in fish schools. Interface focus, 2(6):693-707, 2012.
Lu, F,, Zhong, M., Tang. S., and Maggioni. M. Nonparamet- ric inference interaction laws in systems of agents from trajectory data. Proceedings of the National Academy of Sciences, 116(29):14424 14433, 2019.
Mei, S., Misiakiewicz, T., and Montanari, A. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on learning the- ory, pp. 2388-2464. PMLR, 2019.
Messenger, D. A. and Bortz, D. M. Learning mean-field equations from particle data using WSINDy. Physica D: Nonlinear Phenomena, 439:133406 2022.
Miller, J., Tang, S., Zhong M., and Maggioni, M. Learning theory for inferring interaction kernels in second-order interacting agent systems. Sampling Theory, Signal Pro- cessing, and Data Analysis, 21(1):21, 2023.
Morris, C., Ritzert, M., Fey, M., Hamilton, W. L., Lenssen, E., Rattan, G., and Grohe, M. Weisfeiler and leman
10

---

Universal Approximation of Mcan-Field Models via Transformers
A Proof of Theorem 4.7
Theorem 4.7 (Universal Approximation). Let St
R be compact set containing 0. Let F 1 x P(L) - R satisfy
Assumption 4.3a) for given p. Given a transformer T -2++1 - (a++ let
= sup ZEQ
[T(z) F
(15)
Then, for > p there exists a constant C(p, q, d), depending only on p. q and d, such that for n 2 1, the corresponding
continuum version Tn P(2) R (9) satisfies
| FIl. SE+ + L diam((?))
+ cC(m.D.))
p> d/2, q 4 d/2, 2p < d/2, :5p
G(n,p.9) =
log(n + 1)
(16)
Proof.
F(x, 11) / (T(x, dp (z)
| - Tall= = sup sup 1
52"
< sup sup 1
|| 14) (T(x 2))1 du (z)
S
= sup
[[F(z, (4) F(x, 1²) + F(z, v2") - (T z))1 (z)
#
< sup
||F(x,14 - =F(x, 12) du (z))
s2m
[ [[F(x, - (T(x, 2)) d1 (z)
+ sup
(+)
The second inequality follows from the standard triangle inequality. The first integral on the RHS can be bounded from above as,
supsup
[F(x, F(x, 12) dp (z)
52
sup
sup [F p)) - F(x, 12)ll2 (2)
52"
S L11 " V,l,,, (z)
sup
S2m
L sup Exrpon [W (H, [1)]
(**)
Recall that (4) denotes the q-moment of μ i.e My(M) = So (r)(d)())) then as per Theorem (Fournier & Guillin, 2015) there exists constant C(p, q, d) (a function of p, q. d) such that, Ez on [Wr (11, x²)] from (**) can be bounded from above by (M)
(n) g). We obtain:
sup /
[F((, p) Flx, v2)ll (z) s I sup (u
< LC diam (oamiimm))) (f)))((n,p.)))))
(20)
where we have used the fact that y is probability measure on Q.

---

Universal Approximation of Mean-Field Models via Transformers
Next, obtain upper bound for the second integral in (*).
supsup
|[F(x, "n) (T (((x.2)) 2)) (z)
x
Tm
S sup
sup | 122) - (F (2, zn))1 ll2 du (z))
+ sup
sup
(I,..,.. - (T 2))1 ll2 du (2)
(***)
Consider the first term in the expression above
sup
sup
[F(x,,))) (Fn+1(I, 21...., zn))illa du (z)
m2
= =sup
sup
F(r, 12) - x. (1.2)) (z)
a 32
S sup
I LW
(12 (222((2))
(2/20) (z)
[
1
S sup
diam(12))
dp (z)
n1
1
L (iiam(11))
(21)
n +
Since we have assumed (15), the second term in (* +) evaluates to
sup / sup
- (T(x, 2))1ll dp (2) E.
(22)
Putting together (20), (21), and (22), we get that
|Ta Fll. s I diam(S)) (cC(m.p.0) ,+11 E.
B Proof of Theorem 4.12
Proposition B.1. Suppose Po PS()) such that the support of po lies in (R) for some R > O, and that there exists a Lagrangian solution 1F E C(O, r); (R.) of (2). Additionally, suppose that F satisfies Assumption 4.3b). Then the solution satisfies,
supp (=) S (0)
(23)
for all t E [0, where C₁ = (R - 2. 3.
Proof. By definition of the Lagrangian solution (17),
|((t, < |01|2 + ( [F)(X(s, ds
[' M1 |[ (u(s))) ds
< ||]||
(24)
Integrating both sides of (24) with respect μ4 and noting that X(t, 400 me. get,
(i(((t) & Mi(wo) + +.1l S. (1+ (1+2M 22222M2(20 (m(s))) ds
12

---

Universal Approximation of Mean-Field Models via Transformers
Combining this with (24) itself we get,
[txtt (M((((t) s Mi(po) + f
f(22 [1++I((s, + 3M1(((s)) ds
This implies
(x, + Mi(e(t) & (Mi(Ho) +2 + [. 3.4l (I)X(s + Mi(u(s))) ds.
We note that the above equation is in the integral form of Gronwall's lemma (u(t) S a(t) + [ B(s)u(s) ds) with a(t)
(uo) + 2 t and (s) = 3.H.. Therefore, an application of Gronwall's lemma gives us
[X(22)) Mi((()) (M1())) +
$ (AAt)
Hence supp o C (0)
Theorem 4.12. Suppose μi E PP(( is such that the support of μo is contained within (g(0) C. Q for some >> Additionally, let K satisfy K (R 2.11 r) 3 Furthermore, assume that
(T, 1)ll2 S M (1 ||=||2 ||]|| Mi(H)),
(18)
for all x E By (0) and all μ with support in B5 (0). Then, there exists a unique Lagrangian solution to the approximate continuity equation (8) such that supp μ² (1) S Bc, (0) for all E [0, where C, = (R + 2.
Proof. Let (R( be compactly support smooth function such that <xx) 1 for all R+C (0). We construct the function Ma. 2] = o(x 2z, z). Let Fa be the expected transformer corresponding to î. We check that T, satisfies Assumption 4.3a). Towards, this end we compute
(a, - (v. 1) 11 (f(x;2))) du (z) / (T( 2))) du(2)
Let P(Rdd x 1dd) be the optimal plan with marginals a and that solves the optimization problem (12). Then
1 x
[+( - (zr, x)122
(f([=;))) dy(2, 2) - S (((())))) dr(z, 2)
xm2
2
1 (f([z)))) ((y: dr 2))
Dr x
The function T is globally Lipschitz for some Lipchitz constant as a product of compactly supported smooth function & and locally Lipschitz function T. We use this conclude that
[T n) v)l12 f
|| ul12 dr(z, 2) L Σ
I Jnn
|| - 1 dz((tt, 24)
Jarxar
/ 10x00n
L|| - ul12 + L Σ
| - Eill2 dr(2s, 2.)
L||æ vll2 + LmWi
Therefore, 7 satisfies Assumption 4.3a). Then, Proposition 4.8 of Cavagnari et al. (2022) guarantees that there exists unique Lagrangian solution yu (t) of (8) corresponding to Ā. Moreover, we know from Proposition B. that the support of the solution n² (t) lies in Bc. (0) for all t [0, r]. However, note that F(z, 11) T(z, p)) for all x e Bc, (0) and all E ((2) with support in Bc, (0). Therefore, 4 (t) = fut (t), the unique Lagrangian solution of (8) for the expected transformer T.
13

---

Universal Approximation of Mcan-Field Models via Transformers
C Proof of Theorem 4.14
Theorem 4.14 (Mean Field Dynamics Approximation Using Transformers). Let > 0 be small enough and n 2 1. Suppose that F satisfies Assumption 4.3a) for some p. Assume that the support of po contained within Bn(0) 92 for some R > O, and let K E R. be such that K (R 2.ll T) If the transformer Tn satisfies the condition
[T, 1) - F(x, <E.
for all z E (0 and μ E PPBx((0) Then we have that
Wi (1) (m)))=((()== S exp((2t)
(19)
where je and n are the solutions to (2) and (8), respectively, and the estimate (19) is independent of po-
Proof. Using the uniform norm approximation [T) F(x, < and Assumption 4.3b), we can conclude that
(Tx (M + +1 + | + M2())
for all Rd and μ P(Bx(0)). Since K> (R + T. for 0 small enough, and large enough we can conclude that, K > (R + 2(M + From Theorem 4.12, there exists Lagrangian solution for (8) such that
supp (s) S (0)
for all [0, r]. Due Assumption 4.3 and Proposition 4.8 of (Cavagnari al., 2022), there exists unique Lagrangian solution for (2). Once again, using Proposition B.1 conclude that
supp f ((s) C Bk(0)
for s [0, r]. Let X, be the flow maps associated with with respect the vector fields F and Tn respectively. From the definition of Lagrangian solutions know that,
X(t,x) = x I H(
x(, S'" (T,((((, (h)))
for a t e [0, r). From this we get
()(,,)) x(1,x))22- /' (=) J F(X(s,)),1F (s))ds
Ty (s))ds / (( (s))ds
F(Y(s,)) x), 18 (s))(s / F(X(s,)), 1f (F))
14

---

Universal Approximation of Mean-Field Models via Transformers
By triangle inequality of the || 2/2 norm,
I Tn / Fl (t))
1
1./ F( (8) / F(X(s,)), (s))ds
(V( (s))dds F(Y(s,,)), (8))ds
/ 2( (
F(X(s,z). (s))
/ (Y) (s))dss
F(Y(s,x),, Tx (s))
L F(Y (s)
F(Y Hr (s))ds
1. ((, (s)) F(X(s, (s)))ds
+
Using the fact that ||
2|2 is convex, applying Jensen's incquality yields.
Ta(Y 1-
(((() F(Y(S,)),II'()) (8)) ds
ds
x)) (s))ds
+
2
+
F(X(S,)),NE (s)) ((Y (s))ds
ds
Using the fact that
supp (s) (k
we get,
( eds
(Y(t,)) x(1,x))l2 S
+ {.- 2Wi (m) I 2| -
Integrating with respect to no and noting that pi a probability measure,
(i(( (o) Set
+ S smi( 'r sevviuu (o)
$ at +2 +2 o sevi(u (z)
Now, applying Gronwall's incquality, we get,
Wi(H Wi(H=(0) (m) < etexp(2211
15

---

Universal Approximation of Mean-Field Models via Transformers
D Transformer
Definition D.1 (Multi-Headed Self- Attennttttiion.. Let X E Raxd be matrix whose rows data points in R Let Wo. WK. Wv IRddd be learnable weight matrices. Define the query, key, and value matrices by
Q X =
= VWV
Let softmax denote the softmax function applied row-wise matrix. The self-attention head function AttHead Raxd Raxd is defined as
QK
AttHead( X) softmax
V.
Vd
Let h E 24, be the number of attention heads. Let AitHead,,... AttHead, be attention heads with their own weight matrices, and let Wo E hdd be learnable weight matrix. The multi-head self-attention layer Att : Raxd - xxd defined as
Xxt [AttHead] AttHodd, (X)) Wo.
where [)] denotes concatenation along the feature dimension. Definition D.2 (Transformer Network). A transformer block Block R= paxd is defined
(lock(X)= == X + FC FC Att(X))
where FC are feed-forward layers (position-wise fully connected layers), and ReLU the rectified linear unit activation function. The addition operations represent residual connections. Let L Z+, and let Block Block, be transformer blocks. A transformer network T Roxd Ta is defined composition of transformer blocks followed by an output network:
T(X)= = FC (Block,, Block,, .... Block (X)),
where FCout : RRox -> BRan is fully connected neural network applied position-wise.
E Training
Here provide the training details.
E.1 Learning the Vector Field
Hyperparameters We consider depths in (3,4,5), widths in {128, 256, 512). and learning rates in {0.0002, 0.0001, 0.001}. For the kernel method, use polynomial basis of (1, (1.,,..., aa) and sine and cosine basis of sinne))ssn((2)))) sin and {cos(I), cos(22),........................ cos(kr). We search over E {2, 3, 4} and k E (3,4, 5}.
Training Details We used Adam with a cosine annealing decay rate for the step size. For the synthetics CS data, we used a batch size of 500 and trained the model for 1000 epochs. For the fish milling data, we used a batch size of 1, and trained the model for 10 epochs.
E.2 Simulating dynamics
For the simulating dynamics experiment, we also train a transformer to learning the training dynamics two layer neural network. We fix the transformer to have hidden dimension 512 and layers. We train the model for 250 epochs, using a learning rate 0.0002, batch size of 1000, using Adam with cosine annealing.
16


Reviews:
[{"review_id": "YH2vFqfQ4l", "reviewer": "Reviewer_vHZs", "summary": "This paper explores the application of transformers in modeling the mean-field dynamics of interacting particle systems.  The study empirically shows that transformers can effectively approximate diverse mean field models, such as the Cucker-Smale model and systems for training two-layer neural networks. It supports these empirical findings with mathematical theory, proving that the approximation error between the transformer-based and the true mean-field dynamics can be quantified and is dependent on the number of particles used in training. Finally, it establishes theoretical bounds on these errors, enhancing the understanding of transformer capabilities in complex system dynamics.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewers for their questions and comments, which help improve the paper. We thank the reviewer for finding that our paper \"supports these empirical findings with mathematical theory\" and \"establishes theoretical bounds on these errors, enhancing the understanding of transformer capabilities in complex system dynamics.\"\n\nWe hope that our response can clear the reviewers' doubts. \n\n> I am somewhat confused by Figure 1 and Table 1. Firstly, for the mean-field model, is it appropriate to solely use mean square error to assess performance? ... it should also be evaluated by the distance between distributions.\n\nThis is a great question: we have two points. \n\n1. The MSE shown in Table 1 is about approximating the vector field $\\mathcal{F}$. For a given distribution $\\mu$, $\\mathcal{F}$ is a map from $\\Omega \\times \\mathcal{D}(\\Omega) \\subset \\mathbb{R}^d \\to \\mathbb{R}^d$. Since the range of $\\mathcal{F}$ is $\\mathbb{R}^d$, using MSE is appropriate in this case. Alternatively, if we were minimizing the particle positions between the predicted model and data, it would have made sense to use 2-Wasserstein distance.\n\n2. Additionally, if we were minimizing the particle positions between the predicted model and data, bounds on MSE imply bounds on the 2-Wasserstein distance but not vice versa. According to the definition,\n\n$ \\mathcal{W}\\_2(\\mu,\\nu)^2 = \\inf\\_{\\gamma \\in \\Pi(\\mu,\\nu)} \\int \\|x-y\\|^2 d\\gamma(x,y).$\n\nThe MSE bound restricts the integrand, thereby providing an upper bound on \n$ \\mathcal{W}\\_2(\\mu, \\nu)^2$.\nHowever, the reverse implication does not hold. For instance, consider two particles, $a\\_1$ and $a\\_2$, that both start at zero. With equal probability, either $a\\_1$ moves to 1 and $a\\_2$ to $-1$ or vice versa. At the level of distributions, these scenarios yield zero distance. However, the MSE between the two scenarios is 4. \n\n> For table 1, I do not see any benefit of proposed method, regardless that I feel the distribution distance should also be reported. The proposed method seems to have small variance but the absolute mean is not competitive?\n\nWe apologize for the confusion; the $\\times 10^{-k}$ factor multiplies both the mean and the variance. For example, for the Cucker Smale model, our method has a mean of $1.9 \\times 10^{-6}$. The next best mean is the TransformerConv with $m = 20$ has a mean of $3.3 \\times 10^{-6}$ which is nearly double. For the fish milling dataset, we have a mean of $2.2 \\times 10^{-2}$ and the next best model is the TransformerConv with $m=3$ has a mean of $6.5 \\times 10^{-2}$ which is nearly three times that of the Transformer. Thus, Table 1 shows that transformers have the best MSE error. \n\n> For figure 2, it makes me even more confused. Why the comparison is between SGD (optimizer) and Transformer (Model)? Am I missing anything here?\n\nWe use the Transformer to approximate the ODE dynamics that govern the evolution of parameters in a two-layer neural network during training. That is, we are using the Transformer to train a distinct two layer neural network. This approach is analogous to the Cucker-Smale model, with the true dynamics defined by Equations (5) and (6). In this context, the true dynamics we aim to approximate are those induced by SGD. As Mei et al. (2019) demonstrated, the SGD dynamics can be expressed through a mean-field equation (Equation (7)), ensuring that our theory is directly applicable here.\n\n> I am not an expert in this area, but the benchmark experiments seem to be limited to models similar to Cucker-Smale, which I believe is not comprehensive enough.\n\nWe consider three benchmark datasets. The first consists of simulated dynamics generated using the true Cucker-Smale equations. The second dataset features real-world observations of fish milling in a pond—actual fish behavior—which, although often conjectured to follow the Cucker-Smale model, is not guaranteed. The third dataset captures the dynamics of SGD for a two-layer neural network, representing a model that is notably distinct from the other two.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "I am somewhat confused by Figure 1 and Table 1. Firstly, for the mean-field model, is it appropriate to solely use mean square error to assess performance? I was under the impression that it should also be evaluated by the distance between distributions.", "theoretical_claims": "No I did not. I am still confused with problem set up and the correctness of evaluation.", "experimental_designs_or_analyses": "I am not an expert in this area, but the benchmark experiments seem to be limited to models similar to Cucker-Smale, which I believe is not comprehensive enough.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "I am not an expert in this area so I am not sure.", "other_strengths_and_weaknesses": "First of all, I feel really confused about the experiment results.\n\n1. For table 1, I do not see any benefit of proposed method, regardless that I feel the distribution distance should also be reported. The proposed method seems to have small variance but the absolute mean is not competitive?\n\n2. For figure 2, it makes me even more confused. Why the comparison is between SGD (optimizer) and Transformer (Model)? Am I missing anything here?", "other_comments_or_suggestions": "See weakness.", "questions_for_authors": "See weakness.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kGhViaBVgh", "reviewer": "Reviewer_svzQ", "summary": "This paper shows, both empirically and with theoretical guarantees, that mean-field dynamics (\"transport-type\" dynamical systems over the space of probability measures, i.e., which take the form of a continuity equation $\\partial_t \\mu_t = -\\nabla_z \\cdot (\\mu_t \\mathcal{F}(z,\\mu_t))$) can be approximated up to any finite time horizon using transformers, provided the vector field $\\mathcal{F}$ is Lipschitz-continuous in the Wasserstein sense. This is achieved by:\n- fixing a number $n$ of particles and considering the $n$-particle dynamics corresponding to the desired mean-field dynamics, $\\frac{d}{dt} \\mathbf{z}\\_t = \\mathcal{F}(z, \\nu^n\\_{\\mathbf{z}\\_t})$ where $\\mathbf{z}\\_t \\in (R^d)^n$ and $\\nu^n\\_{\\mathbf{z}\\_t} = \\frac1n \\sum\\_{i=1}^n \\delta\\_{\\mathbf{z}\\_{ti}}$\n- choosing a transformer model $T_\\theta: \\Omega^{n+1} \\to R^{(n+1) \\times d}$, where $\\Omega \\subset R^d$ is the domain over which we consider probability measures ($\\Omega = R^d$ in the experiments and $\\Omega=$a compact set in the theory sections)\n- learning a transformer $T = T_{\\hat{\\theta}}$ that approximates the mapping $\\Omega^{n+1} \\ni \\mathbf{z} \\mapsto \\left[ \\mathcal{F}(\\mathbf{z}\\_1, \\nu^{n+1}\\_{\\mathbf{z}}, ..., \\mathcal{F}(\\mathbf{z}\\_{n+1}, \\nu^{n+1}\\_{\\mathbf{z}} \\right]$\n- considering the vector field, called expected transformer, $\\mathcal{T}\\_n: (x, \\mu) \\mapsto \\mathbb E\\_{\\mathbf{z} \\sim \\mu^{\\otimes n}} T([x, \\mathbf{z}])$, and using the associated dynamics $\\partial_t \\mu\\_t = -\\nabla\\_z \\cdot (\\mu\\_t \\mathcal{T}\\_n(z,\\mu_t))$ as an approximation of the desired mean-field dynamics.\n\nOn the experimental side, this methodology is validated on two toy examples (the Cucker-Smale and the Fish Milling models) in dimension 4, and on approximating the training dynamics of two-layer neural networks (in the mean-field parametrization).\n\nOn the theoretical side, the paper provides quantitative estimates of the approximation error for the proposed scheme. In particular the approximation error upper bounds vanish when $\\Omega$ is compact, the time horizon is fixed, and $n$ goes to infinity. These theoretical guarantees rest upon previous works' results on the approximation power of transformers as sequence-to-sequence mappings.\n\n\n## update after rebuttal\n\nSee discussion in the comments. I decided to keep my score of 3, though it's a \"strong\" 3, because there are still some natural questions raised by this paper that are not really addressed, particularly the significance of the in-expectation result; this being said this is arguably okay for a conference publication.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and feedback. We thank the reviewer for finding that \"the paper shows end-to-end guarantees\" and that it is  \"worthwhile to write those bounds down properly, which this paper does well,\" and for finding that the paper takes a novel, \"arguably simpler approach\" compared to prior work.  \n\nWe hope that this response answers the reviewer's concerns.\n\n> Last step of Thm 4.7 proof, upper bound different from theorem statement\n\nThe reviewer is right, we fixed the upper bound. It now reads:\n\n$$\\left\\|\\mathcal{T}_n - \\mathcal{F} \\right\\|\\_{*} \\le \\mathcal{E} +  \\mathscr{L} \\text{diam}(\\Omega)^{p} \\left(\\frac{1}{n+1} + CG(n,p,q) \\right) $$\n\n> Technical details for theorem 4.14\n\nThe reviewer is correct on all counts. \n\nWe were missing a factor of $2^{p-1}$ in line 704.\n\nThe term should be $\\int_0^t \\varepsilon^p ds$. Moreover, we did miss a dimension-related term : $d^{\\frac{p}{2}}$ that is due to the conversion between $\\|\\cdot\\|_\\infty$ and $\\|\\cdot\\|_2$ norms. The final result that we obtain is:\n\n$$ \\mathcal{W}^p\\_p(\\mu^{\\mathcal{F}}(t),\\mu^{\\mathcal{T}\\_n}(t))  \\leq 2^{2p-1} d^{\\frac{p}{2}} \\varepsilon^p t\\exp(\\mathscr{L}^p 2^{2p-1} d^{\\frac{p}{2}} t). $$\n\nWe don't need Young's. We use $(a+b)^p \\leq 2^{p-1}(a^p+b^p)$ with the triangle inequality to give us the needed result. \n\nWe now use the classical definition using $\\|\\cdot\\|_2$.\n\n> Corollary 4.10 proof missing\n\nWe thank the reviewer for pointing this out. Here is the sketch: The proof follows from Theorem 4.3 of Alberti et al., which states -\n\nFor each permutation equivariant function $f$ and $\\epsilon > 0$, there exists a Transformer $T$ such that \n\n$$ \\sup\\_{X \\in \\mathcal{X}^n}\\|f(X) - T(X)\\|\\_\\infty < \\epsilon $$\n\nAlthough the statement does not explicitly provide bounds on the sizes, these can be inferred from the construction. Specifically, the architecture comprises a two-layer network described by Hornik (1989), followed by a single attention layer, and then another two-layer feedforward network from Hornik (1989), resulting in constant depth. The result of Hornik et al. (1989) does not impose bounds on the widths of these networks. The single attention layer has a width of $1+2d+d'$, where $d'$ remains constant with respect to $d$.\n\n> Generalize to maps containing isotropic diffusion term\n\nWe thank the reviewer for this insightful comment. This is something we are currently looking at. While approximation of score functions using transformers is a potential application, the score function is not Lipschitz in the measure, so one requires additional work since results of our paper don't immediately extend to this case. \n\n> Uniform-in-time (resp. polynomial-in-time) approximation error\n\nThis a good point, however, this would require the notion of stability on Wasserstein spaces and robustness of stable mean-field systems to perturbations, so it would involve more work to extend ideas from approximation of classical ODE theory to approximation of mean-field ODEs.\n\n> Learning mean-field dynamics from a finite number of observed trajectories\n\nThis work is relevant to generative modeling or sampling problems in which a noise distribution is transported to a target distribution. \n\nAnother interesting application is training neural networks. We can train a smaller network, approximate the training dynamics using the Transformer, and then increase the width (i.e., increase the number of particles). Then, we could train the large-width model using the Transformer. Note that this doesn't require knowing the training data. \n\n> Approximation bounds for the dynamics using the expected transformer\n\nIn practice, the expected transformer can be approximated quickly. For example, if $x$ is a data point and $\\mu$ is the measure. Suppose we have $B$ collections $z^{(1)}, \\ldots, z^{(B)}$ of particles, where each $z^{(i)}$ is state of $n$ particles. Then, we can append $x$ to each of the $B$ collections. \n\nFor the transformer, we represent the input with a batch size $B$ and a sequence length of $n+1$. This allows the forward pass to be efficiently parallelized in any modern ML library. After obtaining the outputs, we compute their mean to approximate the expected transformer output. Instead of calculating the theoretical mean, we use the sample average over $B$ samples, concentrating around the true mean at a rate proportional to the transformer's variance divided by $B$. We expect that a reasonably sized $B$ will yield an accurate estimate.\n\nWe did this for the Cucker Smale model. The results can be seen in Figure 4. Here, we used $B = 10000$. On a single GPU, this takes a few seconds.  \n\n> Time Complexity\n\nThis is a great question. \n\nFor the fish milling dataset, on a single L4 GPU: Transformer: 6-10 minutes per model,\n\nTransformerConv: 5-10 minutes per model,\n\nCyclindrical, FNN, and Kernel were all under 2 minutes for the largest model.\nThe smaller models were on the order of seconds.", "claims_and_evidence": "Yes the claims are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Yes the proposed methods and evaluation criteria make sense.", "theoretical_claims": "I have checked Appendix A (proof of Thm 4.7) and Appendix C (proof of Thm 4.14). Both contain minor mistakes (or missing steps I wasn't able to fill) which may slightly impact the constants in the bounds:\n- The last step of the proof of Thm 4.7 concludes to an upper bound which is different from the one in the theorem statement: $\\mathcal{L} \\frac{2}{n+1}$ instead of $C \\mathcal{L} \\mathrm{diam}(\\Omega)^p n^{p/q-1}$.\n- In the proof of Thm 4.14, I don't see how line 704 is obtained: if it is again by Young's inequality then there should be an extra $2^{p-1}$ factor a priori. On line 717, I don't see how the first term can be upper-bounded by $\\int_0^t \\varepsilon ds$, it seems it should be $(\\int_0^t \\varepsilon ds)^p$, and similarly for the two other terms. Why not just consider $\\\\|Y(t,x)-X(t,x)\\\\|_p$ and use triangle inequalities? Even then, the norms used in the inequalities stated in Assumption 4.3 don't seem to be sufficient to have dimension-independent bounds at this step of the proof. By the way, this proof assumes that in the definition of the distance $W_p$, the distance $d$ is the $p$-norm, which is not necessarily standard, so it may be useful to specify it. For simplicity it might be preferable to stick to using the $2$-norm on $R^d$ and considering $W_p$ distances defined as in Definition 4.2 with $d$ being the $2$-norm, and in the proof of Thm 4.14, bound $\\\\|Y(t,x)-X(t,x)\\\\|_2$.\n\nThe proof of Corollary 4.10 appears to be missing.", "experimental_designs_or_analyses": "I have not checked the soundness of experimental designs.", "relation_to_broader_scientific_literature": "The paper shows end-to-end guarantees for the problem of approximating mean-field dynamics using transformers. The possibility of using transformers for this purpose is not surprising given the universal approximation capabilities of transformers as sequence-to-sequence maps, but this paper gives a complete rigorous analysis. The error bounds obtained in this paper are relatively straightforward consequences of classical techniques in the context of mean-field dynamics, and are likely way too pessimistic, but it is still worthwhile to write those bounds down properly, which this paper does well.\n\nCompared to related works, in particular Furuya et al. (2024), this paper takes an alternative and arguably simpler approach. Indeed, by considering the expected transformer (and using an easy triangle inequality argument, line 516), this paper's approach allows to apply previous finite-sequence-length approximation results directly \"off the shelf\". (I am not aware of other works taking this approach, but I am not familiar with the literature on transformers.)", "essential_references_not_discussed": "I am not aware of any essential references that were not discussed.", "other_strengths_and_weaknesses": "See my comment on the use of the expected transformer in \"Relation to Broader Scientific Literature\".", "other_comments_or_suggestions": "It seems to me that the approach taken in this work could be generalized to the case where the map $\\mathcal{F}(z,\\mu)$ is of the form $\\mathcal{G}(z,\\mu) + \\nabla \\log \\mu(z)$ where $\\mathcal{G}$ is Wasserstein-Lipschitz (in the sense of Assumption 4.3a). That is, it could be generalized to maps containing an isotropic diffusion term. This would connect to the use of transformers for score learning in the context of score-based diffusion models.\n\nIf the target dynamical system $\\partial_t \\mu_t =-\\nabla_z \\cdot (\\mu_t \\mathcal{F}(z, \\mu_t)$ converges (or is stable), then we can hope for a uniform-in-time (resp. polynomial-in-time) approximation error. I wonder if the proposed methodology achieves this, and if this could be shown using this paper's tools. This remark is inspired by the favorable observed behavior for the simulation on the Cucker-Smale model in dimension 4.", "questions_for_authors": "- Please address the minor technical concerns listed in \"Theoretical Claims\".\n- Do you know of any practical ML settings in which learning a mean-field dynamics from a finite number of observed trajectories is directly relevant?\n- The approximation bounds given in the paper are for the dynamics using the expected transformer. In practice, how can the expected transformer mapping be computed? Can the permutation invariance of $T$ be exploited for fast computation?\n- To complement the empirical comparison of the proposed method's performance compared to baselines in section 3.1: how do these methods compare in terms of computational cost?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mQIXmBdusQ", "reviewer": "Reviewer_6g3z", "summary": "The authors study how transformes can be used to approximate mean-field models. The analysis is both theoretical and empirical.\nEmpirically, they test the transformers' on two different mean field models. Theoretically they provide bounds in therms of $L_\\infty$ distance between the expected transformer and the mean field vector field.\n\n**Update after rebuttal**\n\nI am satisfied with the authors' answer and confirm my positive evaluation of the paper thus recommending its acceptance", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and feedback. We thank the reviewer for finding that the problem we study has \"very important implications in machine learning\" and that \"the paper looks original to me and the question the authors try to answer is fundamental\". We hope that this response answers the reviewer's concerns. \n\n> Comparison with Furaya et al 2024\n\nWe acknowledge the reviewer's comment regarding the need for a more detailed comparison with the work of Furuya et al. (2024) in the introduction.\n\nFuruya et al. (2024) establish that deep transformers are universal approximators for general continuous 'in-context' mappings defined on probability measures. Their measure-theoretic approach defines transformers directly on the space of probability distributions, leveraging a continuous version of attention. A key result is that a single deep transformer architecture, with fixed embedding dimensions and a fixed number of heads, can approximate any such continuous mapping uniformly, even when dealing with an arbitrarily large or infinite number of input tokens represented by the measure. \n\nOur work, while also connecting transformers to measure theory, takes a different approach tailored to approximating the specific structure of mean-field dynamics in interacting particle systems. Instead of defining a continuous transformer directly, we utilize standard transformers designed for finite sequences and propose a novel lifting mechanism: the 'Expected Transformer' ($\\mathcal{T}_n$). This construct maps the standard finite-particle transformer to the space of measures via an expectation over the particle distribution. Our primary goal is not general function approximation on measures, but rather to specifically approximate the vector field $\\mathcal{F}$ governing mean-field dynamics and, consequently, the evolution of the system's distribution described by the associated continuity equation. **Hence Theorems 4.14 is a major contribution of our work.**\n\nThe substantial difference lies in both the model definition and the nature of the approximation guarantee. While Furuya et al. provide an existence result for a single, deep, fixed-dimension transformer achieving a target precision $\\epsilon$ for arbitrarily many tokens, our work provides quantitative approximation bounds for the Expected Transformer $\\mathcal{T}_n$. These bounds explicitly characterize how the approximation error for the infinite-dimensional vector field $\\mathcal{F}$ depends on two factors: (i) the approximation quality ($\\mathcal{E}$) of the underlying *finite-dimensional* transformer on $n+1$ particles, and (ii) the number of particles $n$ used in the expectation, leveraging known convergence rates of empirical measures in Wasserstein distance. We further utilize these bounds to establish guarantees on approximating the *solution trajectories* of the mean-field continuity equation, linking the vector field approximation error to the error in the dynamics via stability results like Gronwall's inequality.\n\nBelow are the specific substantial differences highlighted:\n\n**Approximation Target**\n\n1. Furuya et al. Aim to approximate *general continuous in-context mappings* $\\Lambda^*(\\mu, x)$ defined on probability measures. Require target mapping $\\Lambda^*$ to be *continuous* w.r.t. weak* topology (plus Lipschitz conditions on contexts for masked case). Their focus is broad representational power.\n\n2. Our Work: Specifically targets the approximation of the *vector field* $\\mathcal{F}(z, \\mu)$ governing *mean-field dynamics* and the subsequent approximation of the *dynamical system's evolution* (solution to the continuity equation). Requires target vector field $\\mathcal{F}$ to be *Lipschitz continuous* w.r.t. spatial and measure arguments (using Wasserstein distance).\n\n**Transformer Definition:**\n\n1. Furuya et al.: Define transformers directly on the space of probability measures using a *measure-theoretic formulation* with continuous attention layers ($\\Gamma_{\\theta}(\\mu, x)$)\n\n2. Our Work: Uses standard transformers $T$ designed for *finite sequences* of length $n+1$. Introduces the \"Expected Transformer\" $\\mathcal{T}_n(x, \\mu)$ which *lifts* the finite transformer's output to the measure space via an expectation operation.\n\n**Handling Input Size (Number of Tokens/Particles):**\n\n1. Furuya et al. Show a *single* transformer architecture (with fixed dimensions/heads) works uniformly for an *arbitrary* number of input tokens (even infinite) for a given precision $\\epsilon$.\n\n2. Our Work: The approximation quality of $\\mathcal{T}_n$ explicitly *improves as n increases*, reflecting empirical measure convergence. Focus is on convergence behavior.\n\n> In the introduction the symbols are undefined\n\nWe thank the reviewer for pointing this out, we have now added the definitions.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "I checked the proof of the main theorem (Theorem 4.7) in Appendix A and looks correct to me", "experimental_designs_or_analyses": "yes. The experimental anaysis looks sound and valid to me", "relation_to_broader_scientific_literature": "Yes, there is a dedicated paragraph in the introduction where the authors mention several works where mean field games are learned from the particle level dynamics. Moreover, in section 4 the authors make more specific comparison with a work in the literature in terms of theoretical resutls obtained.", "essential_references_not_discussed": "not that I am aware of", "other_strengths_and_weaknesses": "I think that studying how transformers can be used to approximate mean field games has very important implications in machine learning.\nThe paper looks original to me and the question the authors try to answer is fundamental.", "other_comments_or_suggestions": "The authors are invited to provide a more extensive comparison with the work of Furuya 2024 in the introduction already as it seems to be very related to what they are doing.\n\nThe sentence \"Additionally, Furuya et al.\" use a continuous version of transformes and attention and provide universal approximation of measure theoretic maps is too minimalistic and does not explain what is the substantial difference between their work and your work.\n\n-In the introduction the symbols $P(\\Omega)$ and $D(\\Omega)$ are undefined", "questions_for_authors": "The authors are invited to provide a more extensive comparison with the work of Furuya 2024 in the introduction already as it seems to be very related to what they are doing.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Rzs6ZXfkkz", "reviewer": "Reviewer_NEo9", "summary": "The papers consider a mild variant of the transformer model, as part of a larger literature connecting transformers and maps to/from probability measures.  Their main result, from my perspective, is Theorem 4.14 which provides a sort of \"small time\" approximation guarantees that their version of the transformer model can efficiently approximate certain MF ODEs.  This is largely a consequence of their main technical result (Theorem 4.7).\n\nI find that, the title perhaps mis-matches the claims, as when reading I was expecting to see approximation guarantees for MFGs or interacting particle systems.  Perhaps ODEs should be **clearly** placed in the paper's title and abstract, and not \"models\" which is far too overpromissing.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and feedback, and we thank the reviewer for finding that our results show that \"transformer model can efficiently approximate certain MF ODEs.\" We hope that this response answers the reviewer's concerns.\n\n> Assumption 4.3 \n\nAssumption 4.3 a) is an assumption on $\\mathcal{F}$ map and not on the measures. If we restrict to compactly supported measures, arbitrary maps $\\mathcal{F}$ may not satisfy the Lipschitz condition for any $p$. Hence, we assume that the map satisfies the Lipschitz condition for some *finite* p. \n\nThus, for any $p$ for which we have the Lipschitz condition, Theorem 4.14 holds for that $p$. \n\n> Comparison with Kratsios and Furyaya  (KF) 2025 and Kratsios et al. (K) 2022\n\nWe thank the reviewer for pointing us to these interesting papers. In the interest of fairness, *we would like to point out that KF 25 follow-up paper was posted to arxiv on 5th February after the ICML deadline on 30th January.* \n\n**Role of the Measure**\n\n1. Our Work: The measure $\\mu$ is an input argument representing the state distribution in a mean-field system, directly influencing the vector field $\\mathcal{F}(z, \\mu)$. This directly models physical or biological system interactions. A key strength is its applicability to general Borel probability measures $\\mathcal{P}(\\Omega)$ on a compact set $\\Omega$.\n\n2. K 22: The measure $\\mathbb{P}$ is the output of the model $\\hat{F}(x)$, explicitly designed to handle constraint satisfaction by ensuring the output distribution lies within the constraint set $K$.\n\n3. KF 25: The measure $\\mu$ is an input argument to the target function $f(\\mu, x)$. However, it's restricted to the specific class of Permutation-Invariant Contexts $\\mathcal{P}\\_{C,N}(\\mathcal{X})$ within a geometrically constrained domain $\\mathcal{K}$, aimed at analyzing general in-context function approximation.\n\n**Map Definition and Transformer**\n\n1. Our Work: Defines the map $\\mu \\mapsto \\mathcal{T}_n(\\cdot, \\mu)$ (Measure to Vector Field) using the Expected Transformer $\\mathcal{T}_n$, derived by taking the expectation of a standard, finite-sequence transformer $T$. This provides a practical link between standard architectures and measure-theoretic inputs.\n\n2. K 22: Defines the map $x \\mapsto \\hat{F}(x)$ (Vector to Output Measure) using a modified transformer incorporating Probabilistic Attention, explicitly outputting a measure.\n\n3. KF 25: Defines the map from finite vector spaces to finite vector spaces. However, the input and outputs are interpreted as measures on discrete sets.\n\n4. Connection via Thm 4.7: If the map approximated in K \\& F satisfies our assumptions, then our Theorem 4.7 could be applied to the transformer $\\hat{\\mathcal{T}}$ constructed in K \\& F's Corollary 5.\n\n**Guarantees**\n\n1. Our Work: Provides quantitative $L_\\infty$ bounds on the vector field approximation error ($\\|\\mathcal{T}_n - \\mathcal{F}\\|$) that explicitly show convergence as the number of particles $n$ increases, linking the error to the quality ($\\mathcal{E}$) of the underlying finite transformer. Furthermore, it connects this to the approximation of the system's dynamics ($\\mathcal{W}_p$ bounds for $\\mu(t)$) via Gronwall's lemma.\n\n2. K 22: Guarantees exact support constraint satisfaction ($supp(\\hat{F}(x)) \\subseteq K$) combined with quantitative $W_1$ bounds on how well the output measure approximates the target.\n\n3. KF 25: Provides quantitative probabilistic $W_1$ bounds on the output approximation error, dependent on the target function's modulus of continuity $\\omega$ and the domain's geometry $q$, focusing on the network size needed for a given precision $\\epsilon$.\n\n> Interpret (2)\n\nMean field Equations of the kind in (2) are used for modeling the dynamics of a large system of interacting particles. Here, the derivative in time models the change in distribution (or measure) of the particles, and there is not necessarily any underlying variational aspect to their evolution. In this way, mean field games (MFGs) are distinct, where the evolution of the agents are dependent on the measure, through a loss function of a coupled game.\n\n> Lower Bounds\n\nThe reviewer is correct that using this, we can lower bound the quantity in Line 535. However, this is only one of the three terms that we bound. Even if we manage to individually bound the terms, we used a variety of inequalities, such as Jensen's and triangle inequality, to get our decomposition. \n\n> Title\n\nPlease note that the particle models we consider are formulated as ODEs (e.g. Eq. 6), while the mean-field models (e.g. Eqs. 2 and 7) are expressed as PDEs. In Theorem 4.14, we demonstrate that solutions of the continuity equation can be approximated by the *approximate continuity equation*, where the transformer replaces $\\mathcal{F}$. In essence, our work shows that mean-field models derived from both ODEs and PDEs can be approximated, therefore the word *models* aligns nicely with the overall theme of our study.", "claims_and_evidence": "Rigorous and correct proofs.", "methods_and_evaluation_criteria": "Good", "theoretical_claims": "Not always careful; e.g.\n- In Assumption 4.3, what range are you allowing for p?  $p>0$, or I assume $p\\ge 1$.  Can $p$ be infinite (e.g. compactly supported measures)?  The same issue is in the first line of the proof of Theorem 4.14. \n\nOtherwise the proofs seem correct", "experimental_designs_or_analyses": "Effectively NA", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "The authors compare extensively to Furuya et al. (2024); however, the same author has since provided improved, i.e. fully quantitative results without passing to a \"non-continuous limit\" of attention in:\n- \"Is In-Context Universality Enough? MLPs are Also Universal In-Context.\" (2025).\nCan you please compare to more contemporary result of the same author, and not old results?\n\nAdditionally, the authors should comment on the relationship to the probabilistic transformer of\n- Kratsios A, Zamanlooy B, Liu T, Dokmanić I. \"Universal Approximation Under Constraints is Possible with Transformers.\" International Conference on Learning Representations (Spotlight).\nwhich is measure-valued and, which also yields a universal vector-valued map, when the expectation/barycenter map is applied at the output layer.", "other_strengths_and_weaknesses": "I don't understand the point of Assumption 4.3 in proving Theorem 4.14.  If $\\Omega$ is compact then for any Borel probability measure $\\mu \\in \\mathcal{P}(\\Omega)$ and any $p>0$ we have\n$$\n\\mathbb{E}_{X\\sim \\mu}[\\|X\\|^p] \\le \\operatorname{diam}(\\Omega)^p,\n$$\nwhence $\\mu\\in \\cap_{p>0}\\mathcal{P}_p(\\Omega)$; so why the subscript.  Also, if we're in the metric setting, which we are, then not consider the weakest condition where $p=1$?", "other_comments_or_suggestions": "NA", "questions_for_authors": "- Can you shed some light on how to interpret (2) given that the second argument of $\\mathcal{F}$ on measures, and there is a derivative of $\\mu$ in time (for non-experts in MFGs).\n\n- A key step in the proof of Theorem 4.7 in the quantization results of Fournier et al.  Can the authors provide LBs by combining the LBs in [1] with those in [2] into their analysis.\n\n[1] Kloeckner, Benoit. \"Approximation by finitely supported measures.\" ESAIM: Control, Optimisation and Calculus of Variations 18.2 (2012): 343-359.\n\n[2] Ronald A DeVore, Ralph Howard, and Charles Micchelli. Optimal nonlinear approximation.\nManuscripta mathematica, 63(4):469–478, 1989\n\n\n- Can the authors comment on the aforementioned transformer results and theirs?  Especially the more recent Furuya et al. (2025) result?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "NA"}}]

Author rebuttal:

Meta-review: This paper demonstrated that transformers can approximate the dynamics of mean-field models, systems in physics, biology, and modern machine learning. Through the concept of the expected transformer, the authors provide approximation guarantees with quantitative error bounds on the distance between the true mean-field dynamics and those obtained using the transformer. These results are complemented by careful empirical validation on both synthetic and real-world datasets.

Reviewers appreciated the novelty, depth, and clarity of the contribution. The authors have addressed reviewer concerns by correcting and improving the theoretical analysis, and providing precise comparisons with recent related work (e.g., Furuya et al., Kratsios et al.). 

Overall, this is a high-quality paper with broad relevance across ML theory and applications. I recommend acceptance.
Decision: Accept (poster)