Identification of Mean-Field Dynamics using Transformers

Unfortunately, the Reviewer can not recommend the paper in its current form for publication due to some significant weaknesses:

• The paper's overall positioning is unclear. Although submitted in the learning theory category and featuring training dynamics of two-layer networks as an example, the introduction and motivation focus extensively on "identifying dynamical system models for physical processes." This approach suggests an application-focused paper aimed at modeling systems with many particles or agents. However, the related work section then shifts to reviewing literature on PINNs, Neural Operators, and methods for approximating differential equations. To improve clarity, the authors could either position the paper as primarily application-driven, aiming to model particle/agent systems, or, if the focus is indeed on learning theory, adjust the introduction and related work to align with this emphasis.

• From the reviewers perspective, the use of the theoretical results based on the universal approximation theorem in practice is very limited as they do not provide error rates. The authors also do not refer to their own theoretical results at all in the numerical illustration section. The reviewer appreciates that the authors acknowledge such limitations at the beginning of section 5, but this leads to no real connection between the theoretical and experimental parts of the paper.

• The "Numerical Simulations" section requires more precision and clarity. Besides the ambiguous title, the results and their interpretations are inconclusive. For example, Figure 1 shows that the training and test losses for the SGD approach are significantly lower, contradicting the authors' claim that "the transformer-trained models exhibit favorable training and test loss performance," which seems questionable. Furthermore, the figure on page 9 is partially cut off due to formatting issues. Descriptions should also be refined, replacing vague terms like "quite small" and "quite well" with precise, quantitative evaluations.

• The reviewer would suggest to report the details of the architecture used and conduct ablation studies to offer valuable insights into the model's performance.

• Small technical contribution: Straightforward modification of the transformer model, standard approximation results

• In the present form, rather narrow field of application (though more applications could be found, since mean field dynamics are ubiquitous)

• Insufficient experiments, both more (and more detailed) experiments as well as a better investigation of the outcomes are necessary. For example

  • More settings (e.g., what if a Cucker-Smale system does not reach consensus? Is the proposed method able to deal with these different asymptotic behaviours?)
  • More models (e.g., what about Cucker-Dong or Hegelsmann-Krause models?)
  • Higher dimensions (e.g., high-dimensional Cucker-Smale systems)

• Similarly, more details could be provided regarding the outcomes of the experiment. How reliable are the learned models? What about generalization (e.g., to different asymptotic behaviors)? What about identifiability of the underlying systems?

• Presentation needs to be improved, multiple typos, language issues and imprecisions. Some examples are given below.

L029: Language

Related work L042-44: Feng et al '21, Lu et al '19 work on the microscopic level, not the mean field level

L055: Ref should appear earlier

L061: A reference for the continuity equation should be added here, since this is a field usually not known in machine learning

L094: Why is this reference relevant?

L149: What about existence and uniqueness?

Equ (5): There is no ReLU here

L190: Output shape of $T$ is ambiguous

L200: Language

Equ (8): For consistent notation, parameters should be included here

Section title of 3.2: "with" instead of "of"

L228: Language

L234: Wrong inequality direction

L261: Language

Equ (12): Formatting

L326: Typo

L390: Is this a condition or a consequence?

L423: Language

L471: Superfluous "."

Section, Cucker Smale: How are the initial velocities computed?

Section 5.1: What does "trained with SGD" means here? And how is the transformer model then trained if not with a variant of SGD?

L505: "Note that the transformer model does not compute any gradients." This is at odds with the description before Section 5.1

Figure 1: Formatting problem

L534: Language

While the idea is interesting, the theoretical contribution appears limited and unsound, and the simulations are questionable and difficult to reproduce.

To be more specific, here is a list of some of my concerns:

1. The theoretical contributions of this paper are limited and poorly presented. Theorem 3 and Corollary 1 are straightforward consequences of Alberti et al. and a standard result on the approximation of measures. Additionally, it is unclear whether (16) should hold for all $n$ (as $n$ is refered before being defined). I have two suggestions here: a) $n$ should be defined earlier so that the right-hand side in line 317 is small, with the transformer defined afterwards; b) Theorem 3 and Corollary 1 should be merged, as the only difference between them is whether the result by Alberti et al. is cited.
2. The paper aims to propose a method for approximating solutions to mean-field systems, yet no practical method is actually studied. Indeed, $\mathcal{T}_n$ cannot be computed as it involves an expectation, and $\mu^{\mathcal{T}_n}$ cannot be computed either, as it requires applying $\mathcal{T}_n$ multiple times. This raises important questions not addressed in the paper, such as: 'In practice, should the same trajectories be used at each time step when computing $\mu^{\mathcal{T}_n}$?". The paper would benefit from performing the theoretical analysis on the practical method proposed in Remark 1 rather than on $\mathcal{T}_n$. Furthermore, a practical method for approximating $\mu^{\mathcal{F}}$ should be proposed and studied in place of $\mu^{\mathcal{T}_n}$.
3. Theorem 4 assumes that $\mathcal{F}$ satisfies Assumption 2 while it should be $\mathcal{T}_n$, which is not defined at that point. Consequently, both Assumption 2 and Theorem 4 lack clarity as they stand. In my view, Assumption 2 should be replaced with: "for any $n$ and for $\mathcal{T}_n$ constructed from a transformer $T$, (18) should have a unique Lagrangian solution". Additionally, the contribution of this paper would be strengthened if such a property were proven rather than assumed. At the very least, the authors should explain the difficulties in proving such results; otherwise, it leaves the impression that they did not attempt to do so.
4. Many of the norms used in the paper are imprecise or incorrect in terms of space, including those on lines 235, 243, 269, 271, 284, 317, 350, etc.
5. Figure 3 does not appear in the correct position.
6. The description of the simulations is too brief. As it is now, the simulations are not reproducible to me.
7. Some notations are used multiple times for different objects, making the article more difficult to read.
8. There are a large number of typos.
9. A few words on the computational complexities of the proposed methods would strengthen the analysis.

In my opinion, I think the novelty and technical contribution of the present work is not very significant. It is mainly based on the existing transformer approximation theorem. The assumptions of Theorems (eq. 16) and theorem 4 (the inequality above (19)) seem too strong, which should be further validated.