Universal Approximation of Mean-Field Models via Transformers

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.

Assumption 4.3

Assumption 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.

Thus, for any $p$ for which we have the Lipschitz condition, Theorem 4.14 holds for that $p$.

Comparison with Kratsios and Furyaya (KF) 2025 and Kratsios et al. (K) 2022

We 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.

Role of the Measure

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

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

3. 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.

Map Definition and Transformer

1. 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.

2. 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.

3. 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.

4. 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.

Guarantees

1. 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.

2. 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.

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

Interpret (2)

Mean 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.

Lower Bounds

The 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.

Title

Please 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.

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.

Comparison with Furaya et al 2024

We acknowledge the reviewer's comment regarding the need for a more detailed comparison with the work of Furuya et al. (2024) in the introduction.

Furuya 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.

Our 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.

The 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.

Below are the specific substantial differences highlighted:

Approximation Target

1. 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.

2. 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).

Transformer Definition:

1. 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)$)

2. 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.

Handling Input Size (Number of Tokens/Particles):

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

2. Our Work: The approximation quality of $\mathcal{T}_n$ explicitly improves as n increases, reflecting empirical measure convergence. Focus is on convergence behavior.

In the introduction the symbols are undefined

We thank the reviewer for pointing this out, we have now added the definitions.

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.

We hope that this response answers the reviewer's concerns.

Last step of Thm 4.7 proof, upper bound different from theorem statement

The reviewer is right, we fixed the upper bound. It now reads:

$\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)$

Technical details for theorem 4.14

The reviewer is correct on all counts.

We were missing a factor of $2^{p-1}$ in line 704.

The 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:

$\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).$

We 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.

We now use the classical definition using $\|\cdot\|_2$.

Corollary 4.10 proof missing

We thank the reviewer for pointing this out. Here is the sketch: The proof follows from Theorem 4.3 of Alberti et al., which states -

For each permutation equivariant function $f$ and $\epsilon > 0$, there exists a Transformer $T$ such that

$\sup\_{X \in \mathcal{X}^n}\|f(X) - T(X)\|\_\infty < \epsilon$

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

Generalize to maps containing isotropic diffusion term

We 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.

Uniform-in-time (resp. polynomial-in-time) approximation error

This 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.

Learning mean-field dynamics from a finite number of observed trajectories

This work is relevant to generative modeling or sampling problems in which a noise distribution is transported to a target distribution.

Another 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.

Approximation bounds for the dynamics using the expected transformer

In 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.

For 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.

We 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.

Time Complexity

This is a great question.

For the fish milling dataset, on a single L4 GPU: Transformer: 6-10 minutes per model,

TransformerConv: 5-10 minutes per model,

Cyclindrical, FNN, and Kernel were all under 2 minutes for the largest model. The smaller models were on the order of seconds.

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."

We hope that our response can clear the reviewers' doubts.

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.

This is a great question: we have two points.

1. 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.

2. 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,

$\mathcal{W}\_2(\mu,\nu)^2 = \inf\_{\gamma \in \Pi(\mu,\nu)} \int \|x-y\|^2 d\gamma(x,y).$

The MSE bound restricts the integrand, thereby providing an upper bound on $\mathcal{W}\_2(\mu, \nu)^2$. However, 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.

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?

We 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.

For figure 2, it makes me even more confused. Why the comparison is between SGD (optimizer) and Transformer (Model)? Am I missing anything here?

We 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.

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.

We 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.