Parametric model reduction of mean-field and stochastic systems via higher-order action matching
Authors: Jules Berman
Tobias Blickhan
Benjamin Peherstorfer
Venue: NeurIPS 2024 poster (2024)
Keywords: partial differential equations
reduced modeling
model reduction
stochastic dynamical systems
generative models
Primary area: machine_learning_for_physical_sciences
TL;DR: Learning surrogate models to efficiently predict behavior of systems with stochastic and mean-field dynamics across varying physics parameters.

Abstract: The aim of this work is to learn models of population dynamics of physical systems that feature stochastic and mean-field effects and that depend on physics parameters. The learned models can act as surrogates of classical numerical models to efficiently predict the system behavior over the physics parameters. Building on the Benamou-Brenier formula from optimal transport and action matching, we use a variational problem to infer parameter- and time-dependent gradient fields that represent approximations of the population dynamics. The inferred gradient fields can then be used to rapidly generate sample trajectories that mimic the dynamics of the physical system on a population level over varying physics parameters. We show that combining Monte Carlo sampling with higher-order quadrature rules is critical for accurately estimating the training objective from sample data and for stabilizing the training process. We demonstrate on Vlasov-Poisson instabilities as well as on high-dimensional particle and chaotic systems that our approach accurately predicts population dynamics over a wide range of parameters and outperforms state-of-the-art diffusion-based and flow-based modeling that simply condition on time and physics parameters.

Parametric model reduction of mean-field and stochastic systems via higher-order action matching
Jules Berman"
Tobias Blickhan"
Courant Institute of Mathematical Sciences New York University New York, NY 10012 jmb11740nyu.edi
Courant Institute of Mathematical Sciences New York University New York, NY 10012 tobias, bli cwhan@nyu.
Benjamin Peherstorfer
Courant Institute of Mathematical Sciences New York University New York, NY 10012
pehersto@cims.nyu.co
Abstract
The aim of this work is to learn models of population dynamies of physical systems that feature stochastic and mean-field effects and that depend on physics parameters. The learned models can act as surrogates of classical mumerical models efficiently predict the system behavior over the physics parameters. Building on the Cemamou-Brenie formula from optimal trans- port and action matching, we use a variational problem infer parameter- and time-dependent gradient fields that represent approximations of the population dynamics. The inferred gradient fields can then be used rapidly generate sample trajectories that mimic the dynamics of the physical system on population level over varying physics parameters. We show that combining Monte Carlo sampling with higher-order quadrature rules is critical for accurately estimating the training objective from sample data and for stabilizing the training process. We demonstrate Plasov-Poisson instabilities as well as on high-dimensional particle and chaotic systems that approach accurately predicts population dynamics parameters and outperforms state-of-the diffusion- based and flow-bases modeling that simply condition time and physics parameters.
wide range of
Introduction
Predicting the bchavior of time-dependent processes Xt, time t and across varying physics paramcters y is a koy challenge in computational science and enginecring 40 [ The dynamics of Xt,, typically are described by systems of (stochastic) differential equations, which are derived from physics models and can be computationally expensive to simulate Thus, it is desirable to learn reduced surrogate models that. can be rapidly evaluated to predict the system behavior across varying physics parameters 72
40 32.
"Equal contribution
38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

Reduced modeling via learning population dynamics Given a data set of samples, ie., realizations of the random variable Xt,, on a suitable domain & S R
(Xi i N2. j j= i.....,,NN No, k= k= C
(1)
we aim to learn a dynamical-system reduced model to rapidly predict samples that ap- proximately follow the same law Pt. as Xt. over time i and varying physics parameter μ. We refer to the evolution of Pt, in time as population dynamics. Learning the population dynamics instead of learning the dynamics of the individual trajectories H> Xi for all i= -I..... N₂ and l can be beneficial: There are cases where Pr.u does not change in time, yet every sample trajectory t Xi,, follows complicated dynamics. For example, consider incompressible fluid dynamics with constant density. Samples corresponding to particles that comprise the fluid can have complicated trajectories, whereas on distribution level, the density of the fluid is constant and so are the population dynamics. Furthermore, learning population dynamics seamlessly treats deterministic and stochastic systems because the stochastic models that we consider can be expressed as deterministic Pokker-Plane equations on the population level.
Our approach: Learning parametric minimal energy vector fields that represent population dynamics Building on standard literature on optimal transport theory [ as well as the so-called action-matching loss introduced in [1 we pose a variational problem to learn gradient fields St, so that the continuity equation corresponding to the vector field given by Vst, approximates the population dynamics Pt.. of the samples [L. In the spirit of reduced modeling 722, 10 11 45 we seek a vector field St, that generalizes to different values of the physics parameters js. We therefore optimize for 8t, that minimizes the average objective of variational problem over all parameters the distribution of parameters on the domain D C R". We paramctrize St, with neural so that it can be evaluated quickly over & and μ.
V, where v describes
network with weight modulation 39 [3]
Rapid sample generation in inference phase Predictions at inference time at new physics parameters I are made by sampling based on the vector field Vs,, which means that our approach represents Pt..u through the application of Vst, on an initial condition. Importantly, time in the inference step corresponds to the time of the physics problem so that in one inference step a whole sample trajectory obtained, rather than a sample at one specific time point as in regular conditioning- based methods (sce literature review). Thus, we can rapidly generate samples that follow the law Pt. the inference phase.
Stabilizing training with higher-order quadrature An important part of our contribution is stabilizing the training procedure by accurately estimating the objective of the variational problems from few data samples. In particular, instead of uniformly sampling over the data [ we introduce an empirical loss (8) that utilizes higher-order quadrature Pi in the time direction so that the learned s accurately captures the dynamics over time Consequently, we refer to our approach as higher-order action matching (HOAM). Our numerical experiments show that the higher-order quadrature in the empirical loss is key for learning gradient fields Vst,u that accurately capture the evolution in time and that generalize across physics parameters μ.
Literature review We review relevant literature; see Figure 1 for overview.
Non-inlrusive and dala-driven surrogale modeling There a range surrogate and latent modeling methods that aim to learn or reduce the sample dynamics of the realizations rather than the population dynamics, such as dynamic mode decomposition, Koopman- based methods, and others [7T as neural ordinary difierential equations
as well as neural network- baased methods such AS There also are methods for stochastic 73 21. However, all of these methods ignore physics parameter dependencics and or aim to learn the sample dynamics, whereas we focus on paramctric population dynamics.
76, 86 12
46
58 92
systems 51 22 88
[9
28
Population dynamies and trujectory inference Learning population dynamics has been con- sidered extensively in computational biology in the context of gene expression, where the focus is on learning from independent samples at selected time points rather than from
sample trajectories
RZ
30
[3]
75
N5
0r1; however, many of these approaches [i 84 are
2

---

score-basee modeling
HOAM
po
Pla
Pla-1
P1
flows based modeling
Figure 1: Parametric model reduction with our HOAM seeks to learn vector fields that represent population dynamics At over time t. In contrast, parametric model reduction with score-based diffusion denoising and flow-based modeling requires conditioning on time t, which leads to separate, costly inference steps for each time step of sample trajectory.
simulation-basee and thus require integrating dynamics during the training parameteriz- ing the density additionally to the vector field. These works also are not concerned with generalizing over a range of physics parameters in many cases.
Diffusion- and flow-based modeling There is a large body of work on diffusion-based 91 II [2 and flow-based modeling b4 see [I] for detailed review. These approaches are not taking into account time f because they learn paths between a reference and a target distribution only. There works that condition
36 41
time t and parameter μ such as [68] [14] 26 37 33 38, 52), but this requires then generating a path for each time step al inference time, which is computationally expensive. Furthermore, the conditioning on time I means that the target distribution Pt.u at each time f and is different, and thus a separate paper-parameter tuning can be required, which is impractical over many time steps and physics parameters as in our physics problems; see our numerical experiments. The works compute transport-based solutions but parametrize different quantities than our approach, require actively sampling data, and ignore physics parameters We note that there also work on forecasting with diffusion- and Bow-basod modcling which is different task than our task of predicting across varying physics parameters.
15 78 50
18
Oplimal transport Besides the machine learning literature, variational approaches for inferring vector fields are extensively used in optimal transport theory [ 4 Of particular importance to is the formulation by Benamou and Brenier [ The Soiiaammu-Breniic formula. describes a joint optimization problem over vector fields and paths in probability space and the action matching loss [I| is the restriction of this optimization problem to the case of fixed path and the vector field parametrized by a neural network, which are core building blocks for that we show can be used together with parameter dependency.
Contributions We summarize contributions:
(a) Developing a loss to learn population dynamics that remain valid across varying physies parameters by building on optimal transport literature g and action matching [1] (b) Introducing higher-order quadrature schemes for the loss to efficiently couple the gradient fields over time. This leads to lower variance estimators of the loss that critically stabilize training.
(c) Demonstrating
range of physics problems from Vlasov-Poisson instabilities to high-dinennsioall chaotic systems that our approach leads to (i) accurate predictions of population dynamics and (ii) orders of magnitude speedups in naference/predicttnnn over
3

---

classical methods that mumerically solve the underlying partial differential equations as well as standard diffusion- and flow-based models that condition on physical time.
We provide an implementation of our method at https://githhut com/
Method
2.1
Paramcter-dependent population dynamics
Continuity equation Let us consider data ! corresponding to the probability measure Pt.,, which is absolutely continuous for & E [0. 1] and μ e D. We use the same notation for the measure and its density. The density of μ is also assumed to be absolutely continuous
D. We consider population dynamics of Xt, ~ Pt,, that can be described by the continuity equation
(2) with the initial condition Pt=0, =: Po.a and vector field Vt.u- Notice that in our case the contimuity equation (2) depends
p. = - (r,,,)
for all t [0.1], D,
the physics parameter / V. There can be many vector fields Vt that lead to the same population dynamics (2). For example, if Ue, a vector field that describes the dynamics of Pt,a via [2] then another vector field is given by 2k. = Ur + w/Pt. with any other w that satisfies
w = 0 as long as Pt,, is positive.
Uniqueness via gradient fields and the corresponding elliptic problems Because we aim learn a vector field from sample data (1) that describes the population dynamics (2) of the corresponding law Pt.u, it is helpful to remove this non-uniqueness. One way to do so to restrict the vector field to ft Vst so that is a gradient field 4 p. 45). Plugging Vt,, = st, into ( together with the assumptions Pt.p 0 and Jx de 0, leads to parametric clliptic problems in St,u
-V- Vst,u) = jePt.a.
with coefficient function Pt.y, right-hand side (source term) a.pr, and homogeneous Neumann boundary conditions pt,, S = on òx' with normal vector î for all t E [0, 1] and D. The weak forms of the elliptic problems (3) lead to energy minimization problems that can be used to learn the gradient field $t,p via optimization:
1 min - min Exg(()= (E (P)
V S Oepr,ode
(4)
2
for each 1 E [0, and I E D. The space HP X) contains functions with Sx Vs|"Pr,,dz < co, which is the energy (semi-)nnrmn corresponding to the Pt. weighted inner product 29 Sec. 2.3.2).
Optimal transport Standard elliptic theory guarantees unique solutions up to constants of (4) in the Sobolev space H' (X') under strong assumptions on Pt.n. such as uniform boundedness by a positive constant for all and μи see [2] Proposition 2.2] and [I] Section 3.2). The theory of optimal transport allows treating the much more general case when Pt. not uniformly bounded away from zero and possibly atomic; we refer to [ and [74] Section Pt, in the sense of (2. gradient fields Vst, have the smallest associated kinetic energy Ix pt, which is the objective considered in N In the language of optimal transport and in particular the formalism of [63], vector fields with minimal kinetic energy describe tangent vectors to the curve + Pt.u- The motric is the inner product of [' .. X, R. This is the weak Riemannian structure of P(X) equipped with the Kantorovich-Rubiann Rubinstein metric and described in detail in 5 Chapter 8). We give a short description in Appendix
5.3.1] for details. Among all vector fields Dt that are compatible
we can also use other Energy functional with entropy term Instead of the energy 4
choices of the energy select gradient fields, as long as energy functions are convex to maintain uniqueness. We consider an energy that is based on a different notion of discrepancy on P()), the entropic optimal transport or Schrödinger bridge problem 77
1 / [v log pt.w)) Pe - - d,pt, (s) 2
2

---

which depends on e 2 O. The energy Ei, is of particular interest for two reasons: One, the Shule-LLgraaee
equation of in strong form is the Fokker-Planck equation for " Pept. = -V (Pt.pu V + S Pp.,. again with homogeneous Neumann boundary conditions for all t E [0, 1] and E D; see Appendix C This mcans we can efficiently generate samples after learning si., via corresponding stochastic differential equations (SDEs). Two, it can be interpreted as regularizing the field $t,a, which we discuss in Appendix C
2.2 Loss for learning vector ficlds over time t and physics paramcter
Variational formulation over and So far we just carried along time and physics parameter z but did not address them in the variational problems, i.c., we had separate variational problems 4 for all E [0, 1] and μ v. We now propose to consider the average energy over t and μ to infer a map s: [0, 1] x D + H (Pr.uu 2²,, (t, p) + St,0 which is called solution map in reduced modeling 72
151,
E (8) = min / ('
S...'' Ei (opp)) dl de()).
min = (pC.
(6)
Notice that time t and physics parameter have two different effects on the gradient field VoL.i: Time couples the elliptic problems (i.e., (3) for e 0) via the time derivative O.pt.pi see Appendix D In contrast, the elliptic problems are uncoupled over μ and can be considered separately. This means that to compute the solution to an elliptic problem for one value of μ D, one does not need to consider any other μ E D. This will allow us to sample the physics parameters over D independently from each other when estimating the corresponding loss, whereas we will use higher-order quadrature to obtain accurate approximation of the time integral to ensure the coupling between the time points is reflected in St,ui see Section 2.3
Loss for learning gradient fields from samples over I and μ The energy E defined leads to a loss that can be estimated from samples (1]. The quantity 0ipt, appears in 44 and (5), which is typically unavailable when we have access to data samples only. Integration by parts of the term involving O.Pe, eliminates it, see also Appendix
in 4 as well as the energy Fi.,, defined in
D
We arrive al.
F3 (5)) 5 (/.
/ V VoL. 2 + Ast. Penalade / st
t=1 du(u). (7)
t-
Note that this loss is comprised only of expectation values with respect to Pt.p and is therefore well-defined also for empirical distributions. The choice e > 0 assumes that the Fisher information of pt.p is finite.
Remark 1. Loss functions of the form 7 but without the parameter dependence have been used in bT and 177, Theorem 2.1]. In fact, the case with e = appears already in R. Equalion 35) and 64, Section 3/. We bwild on these resulls but work with population dynamics that depend on physics parameters, which leads to the loss shown in
2.3 Parameterizing the vector field, estimating the loss from data, sampling
Parametrizing St.u with weight modulations We parametrize the vector field st via a neural network with continuous versions of low-rank adaptation (CoLoRA) layers, which have been successfully used for parametric model reduction of deterministic time-dependent dynamical systems 13: see also 39. The layers have the form C(r) = Wx+ + Φ(t, (1) ABx + b, where W is weight matrix, A, B are low-rank matrices, b is a bias vector, and @(t, E TR. is a scalar weight modulation; see Appendix B Only the weight modulations ø(t, 4)) depend on time and physics parameter μ. We use hyper-network h [0, 1] D x depends on the weight vector vp E V C Rs to map t and to the modulation weights H),, hit, μc; v). The weights W, A, B, b, which are independent of t and μ, over all layers are collected into the weight vector 0 C RRY. Typically q < q'- Using the hyper-network encourages continuity of st in time t, which is key for many physics problems B
R that
5

---

Combining higher-order quadrature and Monte Carlo sampling for estimating the loss from sample data Estimating the loss 7 from data can be challenging because the three nested integrals (expectations) over the samples X; time 4, and physics parameter Iu can have different properties and correspondingly need different numerical treatment. Our numerical results show that it is critical to accurately estimate the loss to avoid instabilities in the training; see Section 8 and Figure 2
We propose a combination of higher-order numerical quadrature and Monte Carlo sampling to estimate the loss (7). In particular, we propose to use a higher-order quadrature rule for the time t integral. Because it is a one-dimensional integral, standard higher-order quadrature rules from numerical analysis are applicable [27]. The time integral needs be estimated with particular high accuracy to ensure the coupling between the time points as to
well as the coupling to the boundary terms to match the path from Po., at time Pi, at time t = 1. Our numerical results will show that estimating the time integral to high accuracy is essential for stabilizing the training. In contrast to the one-dimensional integral over time, the integrals over X' and the parameter domain D can be high dimensional and thus we estimate them via Monte Carlo estimation.
We consider two high-order quadrature rules, composite Simpson's quadrature and Gauss- Legendre quadrature [77]; see Appendix [ We refer to our method as HOAM-S and HOAM-G when using either quadrature, respectively. Importantly, these quadrature rules require samples on specifically spaced time points, equidistant in the case of Simpson's i's and at the Gauss-Tegendre nodes in the case of Gauss quadrature. Tf the data set. (1) does not contain samples these time points then we interpolate the data the appropriate times. We note that for Simpson's 's quadrature, interpolation is typically unnccessary as data simulated with numerical methods often come at equispaced points in time.
We denote a Monte Carlo estimate of an expectation value obtained from a mini-batch as 8 B = Ei= f(X*) where x2,x2..... X ~ p. Then, the empirical loss with mini- batchinthhng of sizes ng By and ny quadrature points in time given by
The
ES (5) pau
Wn Baz Pegppe..
Vet..nel +
A A
Eng Exppsn
[sc,e]
(8)
where Wn are numerical quadrature weights and in are the corresponding nodes; see Ap- pendix A for the Simpson's 's quadrature and Gauss-Leendre weights and nodes.
Rapid predictions (inference) with learned reduced models Making predictions in the inference step means drawing samples that follow the law represented by the learned gradient field St,,u which approximates the law Pt,ys of Xt,u- Because we train with the loss 77, we integrate the SDE d&t,,. V V (Xt, )dt + edW,, where W, are Wiener processes and is the same e that is used in the training loss (7); see Appendix C As initial condition, we use samples from Po.a at time
O. Of course other sampling schemes can be used [70].
Notice that the time in the SDE used for generating samples is the same time as the physics problem and thus of the sample trajectory. This means that the costs of the inference stop of our HOAM for generating a trajectory of length K scales as O(K). In contrast, introducing conditioning on time and physics parameter in, e.g., noise-conditioned score matching (NCSM) [ and conditional flow matching (CFM) or stochastic interpolants 2 5 requires inferring separate sampling path for each and μ pair of interest. In particular, the inference costs of CFM scale as O(KT), where 7 is the number of sleps taken in the differential equation for generating one sample at one time point. For NCSM with annealed Langevin sampling, the inference costs scale as O(KTr), where is the number of annealing steps. Contrasting this to the scaling of O(K) of our HOAM approach shows that HOAM is well suited for fast predictions over t and as required in parametric model reduction.
3 Numerical experiments
Examples We consider the following parametric dynamical systems; details in Appendix B 1. Hormonic oscillator: A collection of particles evolves in dour-dimensional phase-space

---

rel. err. q())
200
mean Wasserstein
10
MMM
10-2
((s)())
10
L
10
10
000
AM
-00
HOAM
10
10-8
HOAM-G HOAM-S
AM
0.0
0.5 time
1.0
Gauss
Simp
trap
AM
10000
Adam Iteration
Figure 2: Left: During training the high-variance function (())(t) needs to be numerically integrated for estimating the loss. Center left: The high variance leads to inaccurate estimates of the time integral by Monte Carlo, whereas higher-order numerical quadrature produces accurate estimates. Center right: Numerical quadrature in HOAM leads to stable estimates of the loss whereas Monte Carlo integration in AM leads to unstable behavior. Right: HOAM based on higher-order quadrature is stable and more accurate than AM.
subject to a quadratic potential - -
the experiments shown w 8. The particles are initially at rest and follow a. normal Gaussian distribution in space with mean mo = [1, To avoid the formation of a singularity at wt _ we add white noise of strength 17 = x 10-2 to the momentum equation. For the case n = O, we have analytical expressions for p and s, see Appendix B.1
Two-stream instability We numerically solve the Tlasov-Poissie partial differential equa- tions using a particle-in-cell method to generate samples (I). We consider the two-stream
instability 22
a 1D1V configuration with collisions [87] See 2(b) with 3 10-3
and 0o = as in 49. These collisions lead to stochastic sample trajectories. The parameter
y E [1.2, ratio between electric and inertial effects in the simulation.
normalization constant related to the Debye length
S3
It controls the
Bump-on-tail instability Using the same numerical setup of the Vlasov-Poisson equation as for the two-stream instability, we also consider the the bump-on-tail instability The parameter varies as μ E 2.0..
43
4. Strong Landau damping We consider the strong Landau damping phenomenon that is governed by Vlasov-Poisson partial differential equations again but now in a 3D3V (six- dimensional) setup. A perturbation in the æ -d-direction leads to the formation of phase-space structures 50 The parameter ja E [0.5, 1.5) is the mass of the charged particles.
dimessiimenio chaos A Rayleigh Bénard convection leads to a density gradient that sets fluid in motion. We consider a nine-dimensional dynamical system that is derived from such a flow, which exhibits cascades that lead to chaos 69 The parameter μ E [13.7, 14.4] is the reduced Rayleigh number.
Particles in aharmoníc trap We consider 50 particles in an aharmonic trap 16 which lead to 100-dimensional samples X that encode the positions of the particles. The particle positions are governed by stochastic differential equation. The parameter [0.3, 0.9]
controls the velocity of the trap.
We compare our higher-order action matching (HOAM) to the original version of action matching (AM) 6I. where we handle the parameter dependence on μ in the same way as in our approach. Additionally, we compare to noise-conditioned score matching
Sctup
(NCSM) where samples are generated via annealed Langevin dynamics [X and conditional
flo matching (CFM) 2 54 for which we condition on lime i and μc see Appendix B
HOAM stabilizes training with higher-order quadrature In Figure E
consider the [
harmonic oscillator example. We learn the field s, and plot (())(t)
st, 2

---

True
HOAM-S
HOAM-G
AM
NCSM
CFM
Figure 3: Histograms of solution fields. Top: Bump-on-tail (t 20) instability. Middle top: two-stream (t = 20) instability. Middle bottom: Strong Landau damping (t 4) instability. HOAM with Simpson's and Gauss quadrature accurately predicts the fine scale features and multi-modality of the population density in the Vlasov problems. AM does not converge on the 6 dimensional problem. Bottom: High-dimensional chaos 60 (t = 3.7, dim 3 vs dim 9). HOAM accurately predicts the low probability region that connects the two high probability regions while AM does not converge.
D,st.n + A.....u over time t, which is the function that needs to be integrated time to estimate the loss 77. As Figure 2 (left) shows, this function is far from smooth and exhibits several sharp peaks, which make estimating the loss challenging. In AM [1], the time integral is estimated by averaging samples uniformly taken in time, which is equivalent (center left) shows the relative error in estimating the time integral using Monte Carlo integration done by AM versus mumerical quadrature as in our HOAM. The trapezoidal rule, the composite Simpson's rule, and Gauss Legende quadrature all produce highly accurate estimates. whereas Monte Carlo integration yields inaccurate estimates the integral with high variance.
to Monte Carlo cstimation. Figure 2
Poor numerical estimates by Monte Carlo lead to unstable and inaccurate estimates of the loss function in AM, which eventually causes the optimization to diverge as shown in Figure (center right). In contrast, our HOAM, where training is done with higher-order quadrature (in this case Simpsons' rule), the loss curve stable and of low variance. Figure 2 (right) plot the mean Wasserstein distance over time of solutions generated via gradient field trained with with Simpson's 's (HOAM-S), Gauss quadrature (HOAM-G), and with Monte Carlo (AM) for five seeds, which determine the random initialization of the neural network. For some seeds, AM yields reasonable solutions while for others numerically instabilit ies lead AM to fail. In contrast, quadrature-based HOAM is consistently stable and provides orders magnitude more accurate results.
Accurate predictions with speedups for Vlasov-Poisson equations Our Vlasov- Poisson problems describe the interaction of charged particles with dynamics that depend on all other particles, which leads to mean-field dynamics for large numbers of particles N₂ Thus, reduced modeling with HOAM is well suited for this problem because the natural
8

---

HOAM-S
HOAM-G
AM
NCSM
CFM
10
myma
Electric Energy
10
10
10
10-
10-2
40
40
Time
40 0
40 0
40
Figure 4: Electric energy of bump-on-tail (top) and two-stream (bottom) instability. HOAM with Simpson's and Gauss quadrature accurately predicts the energy growth in the transient regime and oscillations at later times. The ground truth is displayed in blue.
dynamics to learn from such a system are the population dynamics Pr, rather than the sample dynamics; see Appendix B.2 We observe the particles computed with particle-in- cell method and learn the gradient field Vsr, with the proposed HOAM approach. For a test physics parameter μ that controls the wave number, VSt. and plot a histogram in Figure for the bump-on-tern onil (top) and two-stream (middle top) instability. Our approach approximates well the histogram obtained with the classical particle-in-cell method. Figure 5 (right) shows that HOAM is the only method which provides speedup the classical particlo-in-cell (full) model, as NCSM and CFM lead to 1-2 orders of magnitude longer inference times than HOAM and the full models.
then generate samples with
For the strong Landau damping problem in six dimensions (three spatial and three velocity), our HOAM approach achieves about 2 orders of magnitude speedup. This because the runtime of the full model based on the traditional particle-in-cell method to compute the mean-field dynamics scales poorly with the dimension. In this example, the runtime the full model increases by almost two orders of magnitude. In contrast, the runtime of our HOAM reduced model increases only from to seconds. This importantly shows that the computational costs of the inference step of reduced models built with HOAM avoid exponential scaling with the dimension in this example.
We now compute the electric energy as quantity of interest from the generated samples over time for the test physics parameters, which we plot in Figure and its relative error averaged over time (e.e.) in Table I (see (25)). Our HOAM approximates the electric energy well at later times, whereas NCSM and CFM lead lo poorer approximations
later times 1. This is relevant because this non-linear regime is where numerical solvers become important; the initial (linear) growth regime can be approximated well by analytical perturbation theory. Also for the six-dimensional strong Landau damping problem, our HOAM approach provides accurate predictions the electric energy with orders of magnitude speedups; see Table and Figure as well Figure in the appendix.
Speedups in inference step (predictions) Recall two limitations of introducing time and physics parameter dependence in NCSM/CFM via conditioning (see page 3 and Section 2.3 : (i) For each and μ, separate sampling path has be computed, which leads to orders of magnitude higher inference runtimes than in HOAM; see Table [ Section 2.3 (ii) For each and a pair, the target distribution Pt., is different, which can require t- and j-specific tuning of hyper-parameters of the inference step, which is impractical and thus can lead a deterioration of accuracy compared to our HOAM approach; see Figure
Predicting statistics of chaotic and particle dynamics high dimensions We now consider the nine-dimensional d dynamical system introduced in G4 which leads to chaotic behavior. We show Figure (bottom) the sample histogram corresponding to a test physics parameter that represents the Rayleigh number. At time = 3.7 and projecting onto

---

10a
True
HOAM
10
0.0
HOAM
CFM
runtime
10a
NOSM
-0.2
10
FOM
0.2
001
0.0
0.1
two-stream bump 8. Landau
Figure Left: HOAM accurately predicts the time evolution of the mean position of 100- dimensional particle system in an aharmonic moving trap (dim vs dim 100). Right: HOAM reduced models provide about 2 orders of magnitude speedup over traditional numerical (full) models for the 6 dimensional strong Landan problem. HOAM is also 1-2 orders of magnitude faster than CFM and NCSM, which provide no speedup over the full models in our problems.
example:
two-stream
bump-on-tail
strong Landau
9D chaos
metric:
c.c.
r.t. [s]
c.c.
r.t. [s]
c.c.
r.t. [s]
sinkhorn
r.t. [s]
CFM 2 54
1.44
139
5.52
141
0.629
161
0.259
NCSM 80
0.245
1142
0.626
1133
4.06
4531
0.869
1109
AM 61
0.275
0.892
6
NaN
80.1
HOAM-S (ours)
0.078
0.427
0.641
7
0.214
HOAM-- (ours)
0.208
0.429
6
0.447
7
0.217
7
Table 1: HOAM with Simpson's and Gauss quadrature outperforms state-of-the-art methods w.r.t. inference runtime (r.t.) with comparable errors when applied to various physics problems for parametric model reduction. Metrios: e.e. the relative error in electric energy, see 25
for the Sinkhorn divergence, see Appendix B.5
dimension three and nine, the histograms show that the proposed HOAM accurately matches the low probability region that connects the two high probability regions, whereas A. fails to converge. Consider now the example of the particles in aharmonic trap, which leads to 100-dimensional samples Xi For a test physics parameter, Figure shows that HOAM accurately predicts the mean particle positions even for this high dimensional system.
4 Conclusions, limitations, and future work
For parametric model reduction, learning population dynamics via minimal-energy vector fields over time / and physics parameter with our variational approach helps reduce inference runtime compared to standard diffusion- and boseb modeling that condition on and and therefore have to solve separate inference problem for each time step and physics parameter at test time. Because we learn the dynamics over time t, it critical to accurately capture the coupling quadrature schemes when estimating time integrals in the training loss. The higher-order quadrature of the time integrals considerably improves training stability. Our approach achieves comparable crrors as stato-of-the-art methods while at the same time reducing inference runtime by 2 orders of magnilude. Additionally, HOAM provides speedups of up to 2 orders of magnitude to classical numerical full models.
the time steps, for which we propose to higher-order
First, if there are only very few samples in time, even numerical quadrature cannot provide an accurate enough estimation of the loss, which could be limitation in computational biology 23 g. Second, we currently seek vector field that minimizes the kinetic energy or a variant thercof. Investigating other notions of energy that might lead to vector fields with other desired properties in certain problems remains challenge.
Limitations:
We do expect that this work has negative socictal impacts.
10

---

Acknowledgements
Berman and Peherstorfer were partially supported by the Air Force Office of Scientific Research under award FA9550-21-1-0222, Peherstorfer and Blickhan were partially supported We thank
by the Office of Naval Research, United States under award N00014-22---279 Stefan Possanner and Dominik Bell (Max Planck Institute for Plasma Physics) for their support with using the high-fidelity code used in the axx-dimensionl Vlasov experiments.
References
[ M. S. Albergo, N. M. Boffi, and E. Vanden-Eijnded Stochastic interpolants: A unifying framework for flows and diffusions. arXiv, 2303.08797, 2023.
2] M. S. Albergo and E. Vanden-Eijuden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.
3 L. Ambrosio and W. Gangbo. Hamiltonian ODEs in the Wasscrstein space of probability measures. Communications on Pure and Applied Mathematics, (1(1):18-5 (11):11853,, 53, Jan. 2008.
[4] L. Ambrosio and N. Gigli. A User's Guide to Optimal Transport In Modelling and Optimisation of Flows on Networks, volume 2062, pages 1-155. Springer Berlin Heidelberg, Berlin, Heidelberg, 2013. Series Title: Lecture Notes in Mathematics.
5] L. Ambrosio, N. Gigli, and G. Savaré. Gradient Flows. Lectures in Mathematics ETH Zürich. Birkhäuser- Verla,, Basel, 2005.
66] V. Arnold. Sur la géométrie différentielle des groupes de Lie de dimension infinie et ses applications à l'hydrodynamique des fluides parfaits. Annales de l'institut Fourier, 16(1):319 361, 1966.
[7] J. W. Banks and J. A. F. Hittinger. A new class of nonlinear finite-volume methods for Vlasov simulation. IEEE Transactions Plasma Science, 38(9 PART 1):2198 2207, 2010.
[8] J.-D. Benamou and Y. Brenier. A computational fluid mechanics solution to the Monge-Kann Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, Jan. 2000.
19] J.-D. Benamou, T. O. Gallouêt, and F.-X. Vialard. Second-Order Models for Optimal Transport and Cubic Splines on the Wasserstein Space. Foundations of Computational Mathematics, 19(5):1113 1143, Oct. 2019.
[10] Benner, S. Gugercin, and K. Willcox. A survey projection- based model reduction methods for parametric dynamical systems. STAM Review, 57(4))483 101-133331,, 2015.
[11] Benner, M. Ohlberger, A. Cohen, and K. Willcox, editors. Model Reduction and Approcimation: Theory and Algorillmss. Society for Industrial and Applied Mathematics, Philadelphia, PA, July 2017.
[12] P. Benner and M. Redmann. Model reduction for stochastic ic systems. Slochastic Partial Differential Equations: Analysis and Computations, 3(3):291-33) 293-338, Sep 2015.
[13] J. Berman and B. Peherstorfer. CoLoRA: Continuons low-rank adaptation for reduced implicit neural modeling of parameterized partial differential equations. In R. Salakhut- dinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, and F. Berkenkamp, editors, Proceedings of the st International Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 3565 3583. PMI.R, 21 21-27 Jul 2024.
[14] A. Blattmann, T. Dockhorn, S. Kulal, D. Mendelevitch, M. Kilian, D. Lorenz, Y. Levi, Z. English, V. Voleti, A. Letts, V. Jampani, and R. Rombach. Stable video diffusion: Scaling latent video diffusion models to large datasets, 2023.
11

---

[15] N. M. Boffi and E. Vanden-Eijnden. Probability flow solution of the fokker-plance equation. Machine Learning: Science and Technology, (33)033501,, jul 2023.
[16] J. Bruna, B. Peherstorfer, and E. Vanden-Eijnden. Neural Galerkin schomes with active learning for high-dimensional evolution equations. Journal of Computational Physics, 496:112588, 2024.
[17] C. Bunne, L. Papaxanthos, A. Krause, and M. Cuturi. Proximal optimal transport modcling of population dynamics. In G. Camps- Valls, F. J. R. Ruiz, and Valera, editors, Proceedings of The 25th International Conference on Artificial Intelligence and Statislies, volume 151 of Proceedings of Machine Learning Research, pages 1-12..... PMI.R, 28 30 Mar 2022.
[18] S. R. Cachay, B. Zhao, H. James, and R. Yu. DYffusion: A dynamics-informed difTusion model for spatiolemporal forecasting. In Phirly-sevenlt Conference on Newral Information Processing Systems, 2023.
[19] R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. Neural ordinary differential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[20] Y. Chen, M. Goldstein, M. Hua, M. S. Albergo, N. M. Bolli, and E. Vanden-Eijnden. Probabilistic forecasting with stochastic interpolants and Föllmer processes. arXiv, 2403.13724 2024.
[21] Y. Chen and D. Xiu. Learning stochastic dynamical system via flow map operator. Journal of Computational Physics, 508: 508:112984, 2024.
[22] Y. Cheng, A. J. Christlieb, and X. Zhong. Energy-conserving discontinuous Galerkin methods for the Alpér system. Journal of Computational Physics, 256:630-655, 2014.
[23] S. Chewi, J. Clancy, T. Le Gouic, P. Rigollet, G. Stepaniants, and A. Stromme. Fast and Smooth Interpolation editors, Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pages 3061-3069. PMLR, Apr. 2021.
Wasserstein Space. In A. Banerjee and K. Fukumizu,
[24] S.-N. Chow, W. Li, and H. Zhou. Wasscrstein Hamiltonian flows. Journal of Differential Equations, 26313:11055-, Jan. 2020.
[25] M. Cuturi. Sinkhorn distances:
Lightspeed computation of optimal transport. In C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Weinberger, editors, Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013.
[26] A. Davtyan, S. Sameni, and P. Favaro. Efficient video prediction via sparsely conditioned flow matching. In Proceedings of the IEEE/CVF /CFF International Conference on Computer Vision (ICCV), pages 23263 23274, October 2023.
[27] P. Deufthard and A. Hohmann. Numerical Analysis in Modern Scientific Computing. Springer, 2003.
[28] F. Dupont, A. Doucet, and Y. W. Teh. Augmented neural ODEs. In H. Wallach, H. Larochelle, A. Beygelzimer, F. a' A'Ach- Buc, E. Fox, and R. Garnett, editors, Advances in Newral Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[29] A. Ern and J.-T.. Guermond. Theory and Practice of Finile Flements, volume 159 of Applied Mathematical Sciences. Springer New York, New York, NY, 2004.
[30] J. A. Farrell, Y. Wang, S. J. Riesenfeld, K. Shekhar, A. Regev, and A. F. Schier. Single-cell reconstruction of developmental trajectories during zebrafish embryogenesis. Science, (60(633))), June 2018.
12

---

[31] I. Gentil, C. Léonard, and L. Ripani. Dynamical aspects of the generalized Schrödinger problem via Otto calculus A heuristic point of view. Revista Matemática Iberoomeri- cana, 3614),107111112 Jan. 2020.
[32] R. G. Ghanem and P. D. Spanos. Stochastic Pinile Blements: A Spectral Approuch. Springer, 1991.
[33] W. Harvey, S. Naderiparizi, V. Masrani, C. D. Weilbach, and F. Wood. Flexible diffusion modcling of long vidcos. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022.
(34] T. Hashimoto, D. Gifford, and T. Jaakkola. Learning population-level diffusions with generative RNNss In M. F. Balcan and K. Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 2417 2426, New York, New York, USA, 20 22 Jun 2016. PMLR.
[35] J. Hittinger and J. Banks. Block-structured adaptive mesh refinement algorithms for Vlasov simulation. Journal of Computational Physics, 241:118 140, 2013.
[36] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Newral Information Processing Systems, volume 33, pages 6840-6851. Curran Associates, Inc., 2020.
(37] J. Ho, T. Salimans, A. Gritsenko, W. Chan, M. Norouzi, and D. J. Fleet. Video diffusion models. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Newral Information Processing Systems, volume 35, pages 8633 8646. Curran Associates, Inc., 2022.
[38] B. Holzschuh, S. Vegetti, and N. Thuerey. Solving inverse physics problems with score matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[39] E. J. Hu, yelong shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. LoRA; Low-rank adaptation of large language models. In International Conference on Learning Represenlations, 2022.
[40] T. J. R. Hughes. The Finite Element Method: Linear Static and Dynamic Finite Element Analysis. Dover Publications, 2012.
[41] A. Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(24):695-704 2005.
(42) P. Kidger, J. Foster, X. Li, and T. Lyons. Neural SDEs as infiniteedimmensional GANs. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 5453 5463. PMLR, 18 24 Jul 2021.
[43] K. Kormann and A. Yurova. A generalized Fourier-Hermite method for the Vlasov Poisson system. BIT Numerical Mathematics, 01(3):881-900 Sept. 2021.
[44] T. Koshizuka and I. Sato. Neural Lagrangian Schrödinger Bridge: Diffussion Modeling for Population Dynamics. In The Eleventh International Conference on Learning Representations, 2023.
[45] B. Kramer, B. Peherstorfer, and K. E. Willeox. Learning nonlinear reduced models from data with operator inference. Annual Review of Pluid Mechanics, 56( Volume 56, 9024):521-5488 2024.
[46] J. N. Kutz, S. I. Brunton, B. W. Brunton, and J. L. Proctor. Dynamic mode decompo- sition: data-driven modeling of complez systems. SIAM, 2016.
[47] H. Lavenant, S. Zhang, Y.H. Kim, and G. Schiebinger. Toward a mathematical theory of trajectory inference. The Annals of Applied Probability, 34(1A):428 500, 2024.
13

---

[48] K. Lee and E. J. Parish. Parameterized neural ordinary differential equations: ap- plications to computational physics problems. Proceedings of the Royal Society A: Mathemalicnl, Physical and Engineering Sciences, 477((2533:202111622 2021.
[49] A. Lenard and I. B. Bernstein. Plasma Oscillations with Diffusion in Velocity Space. Physical Review, 112(5):1456 1459, Dec. 1958.
[50] L. Li, S. Hurault, and J. Solomon. Self-consistent velocity matching of probability flows. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[51] X. Li, T.-K. L. Wong, R. T. Q. Chen, and D. K. Duvenaud. Scalable gradients and variational inference for stochastic differential equalions. In C. Zhang, F. Ruiz, T. Bui, A. B. Dieng, and D. Liang, editors, Procerdings of The 2nd Symposium on Advances in Approzimate Bayesion Inference, volume 118 of Proceedings of Machine Learning Research, pages 1-28. PMLR, 08 Dec 2020.
[52] M. Lienen, D. Lüdke, J. Hansen-Palmus, and S. Günnemann. From zero to turbulence: Generative modeling for 3d flow simulation. In The Twelfth International Conference on Learning Representations, 2024.
[53] E. M. Lifshitz and L. P. Pitaevski. Chapter III Collisionless Plasmas. In E. M. Lifshitz and L. P. Pitaevski, editors, Physical Kinetics, volume 10 of Course of Theoretical Physics, pages 115-167. Pergamon, Amsterdam, Jan. 1981.
(54) Y. Lipman, R. T. Q- Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow Matching for Generative Modeling. In The Fleventh International Conference on Learning Represen- tations, 2023.
[55] J. Lott. Some Geometric Calculations on Wasserstein Space. Communications in Mathematical Physics, 277( 437, Nov. 2007.
[56] C. Léonard. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems A, 34(4):1533-1574, 2014.
[57] D. B. Melrose. Instabilities in Space and Laboratory Plasmas. Cambridge University Press, Aug. 1986. Publication Title: Instabilities in Space and Laboratory Plasmas ADS Bibcode: 1986islp. bo.......M.
[58] I. Mezić. Spectral properties of dynamical systems, model reduction and decompositions. Nonlinear Dynamics, 41 (1-3))3300 -25,, 2005.
[59] C. Mouhot and C. Villani. On Landau damping. Acta Matherattca, 207(1):29 201, 2011.
[60] A. Muntean, J. Rademacher, and A. Zagaris, editors. Macroscopic and Large Scale Phenomena: Coorse Graining, Mean Field Limils and Frgodicity, volume of Lecture Notes in Applied Mathematics and Mechanics. Springer International Publishing, Cham, 2016.
[61] K. Neklyudov, R. Brekelmans, D. Severo, and A. Makhzani. Action Matching: Learning Stochastic Dynamics from Samples. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 25858 25889. PMER, July 2023.
[62] F. Noé, A. Tkatchenko, K.RR. Müller, and C. Clementi. Machine learning for molecular simulation. Anruual Review of Physical Chemistry, (Voolue 71, 2020):361-33 -900,, 2020.
[63] F. Otto. The geometry of dissipative evolution equations: the porous medium equation. Communications in Partial Differential Equations, 26(1-2):101-174, Jan. 2001.
[64] F. Outo and C. Villani. Generalization of an Inequality by Talagrand and Links with the Logarithmic Soboley Inequality. Journal of Functional Analysis, 173 June 2000.
14

---

[65] B. Peherstorfer, K. Willcox, and M. Gunzburger. Survey of multifidelity methods in uncertainty propagation, inference, and optimization. SIAM Review, 60(3):550-594) 2018.
[66] E. Perez, F. Strub, H. de Vries, V. Dumoulin, and A. Courville. Film: visual rea- soning with general conditioning layer. In Proceedings of the Thirty- Second AAAI Conference on Artificial Intelligence and Thirtieth Innovative Applications of Artifi- ciol Intelligence Conference and Eighth AAAI Symposium on Educational Advances in Artificial Intelligence, AAAP18/IAATI AAAAAA/////AATS /EAAI1.. AAAI Press, 2018.
[67] S. Possanner, F. Holderied, Y. Li, B. K. Na, D. Bell, S. Hadjout, and Y. Güçlü. High- Order Structure-Prestt Preserving Algorithms for Plasma Hybrid Models. In F. Nielsen and F. Barbaresco, editors, Geometric Science of Information, volume 14072, pages 263-271. Springer Nature Switzerland, Cham, 2023. Series Title: Lecture Notes Computer Science.
[68] K. Rasul, C. Seward, L. Schuster, and R. Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. In M. Mcila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8857-8868. PMLR, 18-24 Jul 2021.
[69] P. Reiterer, C. Lainscsek, F. Schürrer, C. Letellier, and J. Maquet. A nine-dimensional lorenz system to study sigh-dimensional chaos. Journal of Physics A: Mathematical and General, 31 (34):7121, aug 1998.
[70] C. Robert and G. Casella. Monte Carlo Statistical Methods, Springer, 2004.
[71] C. W. Rowley, I. Mezié, S. Bagheri, P. Schlatter, and D. S. Henningson. Spectral analysis of nonlinear flows. Journal of Fluid Mechanics, 641:115 127, 2009.
[72] G. Rozza, D. Huynh, and A. Patera. Reduced basis approximation and posteriori error estimation for affinely parametrized elliptic coercive partial differential equations. Archives of Compulational Melhods in Engineering, 15(3):1 47, 2007.
[3 C. Salvi, M. Lemercier, and A. Gerasimovics. Neural stochastic pdes: Resolution- invariant learning of continuous spatiotemporal dynamics. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advonces in Neural Information Processing Systems, volume 35, pages 1333-1344 Curran Associates, Inc., 2022.
[74] F. Santambrogio. Optimal Transport for Applied Mathematicians, volume 87 of Progress in Nonlinear Differential Equations and Their Applications. Springer International Publishing, Cham, 2015.
[75] G. Schiebinger, J. Shu, M. Tabaka, B. Cleary, V. Subramanian, A. Solomon, J. Gould, S. Liu, S. Lin, P. Berube, L. Lee, J. Chen, J. Brumbaugh, P. Rigollet, K. Hochedlinger, R. Jaenisch, A. Regev, and E. S. Lander. Optimal-Tramne Transport Analysis Single- Cell Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell, (76(4):928-94)) 943.e22, Feb. 2019.
(76] J. Schmid. Dynamic mode decomposition of numerical and experimental data. Journal of Fluid Mechanics, 656:5 28, 2010.
[77] E. Schrödinger. Über die Umkehrung der Naturgesetze. Technical Report 1931 IX, Akademie der Wissenschaften, Berlin, 1931.
(78] Shen, Wang, Kale, Ribeiro, Karbasi, and H. Hassani. Self-consistency of the lokker planck equation. In P.-L. Loh and M. Raginsky, editors, Proceedings of Thirly Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages PMLR, 02 05 Jul 2022.
15

---

[79] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequllibrium thermodynamics. In F. Bach and D. Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 2256 2265, Lille, France, 07-09 Jul 2015. PMLR.
[80] Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distributton. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' AlAAhhé--Bu, E. Fox, and R. Garnett, editors, Advances in Neural Informalion Processing Systems, volume 32. Curran Associates, Inc., 2019.
[81] Y. Song.
S. Garg, J. Shi, and S. Ermon. Sliced score matching: A scalable approach to density and score estimation. In Proceedings of the Thirty-Fifth Conference on Uncertainty in Artificial Intelligence, UAI 2019, Tel Aviu, Israel, July 22-25, 2019, page 204, 2019.
[82] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score- basod generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[83] E. Sonnendrücker, A. Wacher, R. Hatzky, and R. Kleiber, A split control variate scheme for PIC simulations with collisions. Journal of Computational Physics, 295:402 419, Aug. 2015.
[84] A. Tong, J. Huang, G. Wolf, D. Van Dijk, and S. Krishnaswamy, TrajectoryNeet: A dynamic optimal transport network for modeling cellular dynamies. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 9526 9536. PMLR, 13-18 Jul 2020.
[85] C. Trapnell, D. Cacchiarelli,, J. Grimsby, P. Pokharel, S. Li, M. Morse, N. J. Lennon, K. J. Livak, T. S. Mikkelsen, and J. L. Rinn. The dynamics and regulators of cell fate decisions are revealed by pseudotemporal ordering of single cells. Nature Biotechnology, 32(4) 32(4):381-380 Apr. 2014.
[86] J. H. Tu, C. W. Rowley, D. M. Luchtenburg, S. L. Brunton, and J. N. Kutz. On dynamic mode decomposition: Theory and applications. Journal of Computational Dynamics, 1(2):391-42) 421, 2014.
[87] T. M. Tyranowski. Stochastic variational principles for the collisional Vlasov Maxwell and Vlasov Poisson equations. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, 477(22252:20210067,, Aug. 2021.
[88] B. Tzen and M. Raginsky. Neural stochastic differential equations: Deep latent gaussian models in the diffusion limit. arviv, 1905.09883 2019.
[89] C. Villani. Optimal transport: old and new. Number 338 in Grundlehren der mathema- tischen Wissenschaften. Springer, Berlin, 2009.
[90] C. Villani. Topics in optimal transportation. Number 58 in Graduate studies in mathematics. American Mathematical Society, Providence, Rhode Island, reprinted with corrections edition, 2016.
[91] P. Vincent. A connection between score matching and denoising autoencoders. Newral Computation, 23(7):1661 1674, 2011.
[92] M. O. Williams, I. G. Kevrekidis, and C. W. Rowley. data-driven approximation of the Koopman operator: Extending dynamic mode decomposition. Journal of Nonlinear Science, 25(:) -344-13633 2015.
[93] F. A. Wolf, F. K. Hamey, M. Plass, J. Solana, J. S. Dahlin, B. Göltgens, N. Rajewsky, L. Simon, and F. J. Theis. PAGA: graph abstraction reconciles clustering with trajectory inference through a topology preserving map of single cells. Genome Biology, 20(1):59, Dec. 2019.
16

---

A Quadrature rules
A.1
Monte Carlo estimation
Monte Carlo estimation approximales an integral by evaluating the integrand randomly sampled nodes within an interval [a,b],
[_ (()) di
b b- E S(A).
(9)
N
We consider only the case where the nodes t, are uniformly distributed random variables in (a, bj. The weights are:
a N 1= N.
(10)
The root mean-squared integration error of Monte Carlo estimation decays as O(N-1/2), for integrands with bounded variance.
A.2 Trapezoidal Rule
The trapezoidal rule approximates the integral of function f over interval [a, b] by The approximation is given by:
dividing it into N subintervals of equal width
N
[ ((1) di
N-1
1 f(1,) + |
5 +
((4)
(11)
where the nodes t, arc:
4 = a ih, 0, O......,NN N.
(12)
The trapezoidal rulc second-order rule, which means that the intogration crror decays as ((h)) for sufficiently smooth functions.
A.3 Composite Simpson's Rule
Composite Simpson's 's rule approximates the integral by fitting parabolas through intervals. Tt b a
divides |a, b] into an even mumber subintervals width
The approximation
N
is:
S f(t) dt
N/2 f(to) +2 f(t2) + Σ f(12-11) +f(tN)
N/2-1
Σ
(13)
with nodes:
4 ih, i O...,,..,
(14)
Composite Simpson's rule is a fourth-order rule, which means that the integration error decays as Ohh*) for sufficiently smooth functions.
A.4 Gauss-Legendre Quadrature
Causs-Legendrs quadrature approximates the integral over [-1, 1] by choosing nodes t, and weights w; so that polynomials the highest possible degree are integrated exactly. A Gauss-Tegendre quadrature has the form
[ f(t) dt = Σ wo (tz),
(15)
17

---

where E; are the roots of the Legendre polynomial (t()), and the weights are:
2
W1 (1 (1-
(16)
For integration over [a, b), linear transformation maps [-1, to [a, bj:
f b- t,
a+b
b b-6
W
wi.
(17)
The approximation becomes:
f(t) dt ~
ff..)
(18)
Gauss-Legendre quadrature exactly integrates polynomials of degree 2n 1, where is the number of nodes.
B Details about numerical examples
B.1 Harmonic oscillator with background collisions
The equation of motion four-dimensional phase-space for = X2, V/1, V22 is given by:
Xi
Vi
0
d
X2
V22 W
dt
Vi
(t) 22
=
X1
(f)
+
"
E(t).
(19)
X2
Here, = 5 x 10-2 and & denotes white noise. The initial configuration is a Gaussian
centered at mo = [1, 1] with covariance equal to Eg 10-2 x Id in the spatial coordinates X2 and Gaussian in the velocity coordinates V1, 1/2 centered at zero and with covariance No 10 x ld.
B.2 Vlasov-Poisson problems
Mean field approximations The Vlasov-Poisson system describes the interaction of charged particles. Due to the presence of the Coulomb force, the dynamics of single particle depend on the position of all other particles. Assuming N particles in the system, this means d di X = =tt X; Xigilh 1t, X1 t.o.... X.N.). Given the fact that N is in practice extremely large, it is natural to pass to the mean-field limit. Assuming the particles are indistinguishable the result is PDE of the form pept. + V Pmt Pt. 0 that describes the evolution of the collection (or population, ensemble) particles denoted by Pt. In the specific case of the Vlasov-Poisson problem, Coulomb interactions in the mean-field limit give rise to Poisson equation determining an eloctric field that is generated by the collection particles and inftuences its dynamics. For completeness sake, mention that the singularity of the Coulomb interaction poses a considerable technical challenge when passing to this limit. We refer to 53 59 for the derivation of the Vlasov-Poisson equation and 600 for more examples of mean-field systems. The theory behind the test-cases we run in this work can be found in 57 Chapter 3.
We slightly change the notation here to be consistent with the references. f: X. x Rd xxx + I, d E 21, denotes the distribution function governed by the Vlasov-Poisse system
Governing equation
0.f(z, 1; - (x,(x, Vo(2,)) V. 1, -
(20)
=1 [ 1;
(21)
In the notation of the rest of this work, /1.,,,444 = Pt.,, X, Rd X. The spatial domain
X, a subset of Rd, all our examples of the form [0,4] x [0, 12] x [0, lg] with periodic boundary conditions.
18

---

Two-stream instability In this case, d d= 2, so the particle positions vary in A₂ = [0 with periodic boundary conditions and their velocity evolves in R. For the two-stream instability, we set the initial distribution to
1 cos
(v (V-210) v10)²) exp
( =
fo(2, 2/2x
2T 1
exp
(22)
with a = 0.05, = 50, up = 3. The parameter l varies as Prrain E 1.3.. 1.9) and Mrest {1.25, 1.85). We use particle-in-cell method for generating the data based on the repository https://githu:
The number of marker particles is N 25000 and for the sake of computing the electric field, uniform grid of N/8 cells is used. Integration in time is done via a Störmer- Verlet splitting over t E [0, 40) with time-step size 10-2
com/ con//poocz/pic-yyhhor
Bump-on tail We consider the initial distribution
fo(2, =
1 I
(
exp
(2o7
+
1
exp
(_(v Ub) )
2
20
20
(23)
The parameter l varies as Mtrain < (1.3, 2.0} and Preat (1.35, 1.95). The other parameters are the same as in the two-stream case.
with a 0.05, 1 50, 1
4,6 9/10, 01
1,02 = 1/V2..
Strong Landau damping In this case, d and
1 + c cos 22T
2T I1 h
|u|2
fo(x,
exp
(24)
with I 4z and 1 1. The data is generated using the Struphy package dr the exact
specifications of the simulation are available at betps://gitLan wttpb/.pdaabb.........mm. mpg /el.dd///seeeee struphy as an example problem. The physics parameter we vary is the mass of the charged particles, which has the effect of changing the strength of the inertial term accelerating the particles relative to the advection term v V.f. This implies μ E {0.5, 0.6, .., 1,5), where μ = 1.0 corresponds to the default settings. This μ 1.0 is also the test paramcter and is excluded from the training set. The timing for the full order method has been obtained on a computing cluster with AMD EPYC Genoa 9554 CPUs using 8 MPI processes, which is default option of the used code. For a single MPI process, extends to 27 minutes while for 16, it can be reduced to 4 minutes.
The high-fidelity data we generate using a control variate approach order reduce numerical noise introduced by the finite number of marker particles. Since we require the particles to be identical for our method, we assume they are all weighted equally when re-constructing the electric potential. This biases reconstructed potential in comparison to the physical one, but we observe in practice that this is only by multiplicative constant. We save 10a marker particles from the high-order simulations and use N 25000 of them as input data for our method. We integrate in time over E [0, 8.75)
B.3 High-dimensional chaos
We consider the dynamical system introduced in 69 We generate samples by initializing a 9 dimensional Gaussian centered at the origin with width equal x 10-2, We then integrate these samples forward as an SDE whose drift is given by the 9-dimensional system of ODEs described in 69 and whose diffusion term is given as diagonal noise equal to x 10-2, We integrate 25000 particles of the system up to 7 20 using the Euler-Maruyana, scheme with time step size equal to 10-2, The parameter μ varies as Ptrain E (13.5, 13.6, 14.2) and Mtest (13.65, 14.05).
19

---

B.4 Particles in aharmonic trap
We consider the evolution of interacting particles in an aharmonic trap 16 The two- dimensional partiele positions (t, p)..... (t. are governed by an SDE
dz, = dt
Σ K(z, V2dW,, i= M,
where y > 0 is the diffusion coefficient and W are independent Wiener processes. The function gtt, Z) (a(t) Z)' describes time-dependen one-body force, where a(t) 5/ (((3im(Tt+++/2) + cos((ttt) is the position of the trap. The function (Z, Z') = M (Z'
describes a pairwise interaction term. We set a = 1/4 and y = 10-2. The parameter μ is in the range D = 00.3, 0.9] and modifics the position of the trap. A sample X corresponds to a vector (1. :.)..... Zas (t, n0)² of dimension 100, because we have M 50 particles and each position (t, (4) two dimensions. We generate samples via Monte Carlo by using the Euler-Maruyama scheme. The time step size is öt = 10-3 and we integrate up to final time
B.5 Other details about numerical experiments
In terms of network architecture, we follow [3 closely because we use their network architecture. We use MLPs to parameterize both the main network and the hyper-network with swish activation functions. The main network is depth 7 and width 64 lincar layers while the hyper-network is depth 3 with width 15 linear layers. The rank of the CoLoRA modulations is set to 3. Identical CoLoRA architectures are used for all HOAM experiments as well as the comparisons with AM, NCSM, and CFM. The only difference is the size of the output layer for NCSM and CFM whose outputs must be the same dimensionality as their inputs.
For all experiments we use an Adam optimizer at a 2 x 10-3 learning rate with a cosine learning rate scheduler. For all experiments unless otherwise noted, we take a batch size of 256 particles over 256 time points. We optimize for 50, 000 Adam itcrations for Vlasov examples and for 25, 000 Adam iterations for all other examples.
The results were computed on NVIDIA Quadro RIX 8000 GPUs. All code was implemented in Python using the JAX library with JIT complication ion where possible.
in the loss searched over {0, 5, x 10-2 for both HOAM and
AM.
The relative error in the electric energy is computed
Ctrue (f) (tr
(25)
[Ctrue (1)
where (t) is the electric energy predicted by the high-fidelity numerical simulations at time and epredict (t) is the electric energy computed from samples of either HOAM (ours), AM, NCSM, or CFM. The relative error the mcan
1 T
Elptrue (t) E Ppredict (1)]
(26)
E(ptrue
where the expected values are estimated Monte Carlo from the generated samples.
The Sinkhorn distance is computed with https://ott- wtw.aaaeeeateeeco..o/eee///////e
with threshold 10-3; see also [25].
B.6 Additional numerical results
In Figure f
we show the various projections at time /. = 3.7 of the sample distribution
corresponding to the nine-dimensions chaotic system 69
20

---

True HOAM-S HOAM-G
AM
NCSM
CFM
2,7
3,9
C
C
4,1
V
M
4,6
Figure 6: Shows the projections of other dimensions of the nine-dimensions chaotic system
69
see also Figure 3
HOAM-S
HOAM-G
AM
NOSM
CFM
Electric Energy
10-1
10-3
5
0
5
Time
Figure Electric energy and solution field at time 4 for the dimensional strong Landau example.
In Figure we show the particle histograms and the electric energy curves for the six- dimensional Vlasov-Poisson problem corresponding to strong Landau damping.
In Figure for the linear oscillator example, compare CoLoRA to two other modulation
schemes: FiLM [6 and MLP. For the MLP the inputs æ, tt, are concatenated together and modulation scheme. For FILM, closely follow the original paper. The main network takes x as input and the hyper-network 1 input. The hyper-network and main network have the same parameter counts in the CoLoRA experiments. The output of the hyper-networ then directly modulates the activation of each layer of the main network detailed in the original FiLM paper [6.. Figure S shows that parameterizing the vector field stw, with CoLoRA layers achieves the lowest mean Wasserstein distance, which motivates the of the CoLoRA modulation
imput directly to the model. There is hyper-networls
scheme 13
C Calculations regarding the entropic loss
In the following, assume that p E P(2') a smooth density bounded away from zero. We begin by showing some calculation rules of the operator -Ap + -V (pVs) with
21

---

T 10
mesn Wassersein
6 6 10"
10
10-4
CoLoRA
FILM
MLP
modulation Schome
Figure & Comparison of CoLoRA modulation scheme CoLoRA layers achieve the lowest mean Wasserstein distance compared to FiLM and MLP. In particular, CoLoRA avoids outliers with larger errors.
18
versus FiLM 66 and MLP.
homgeneous Neumann boundary conditions. In its weak form, it reads
- fA,sdr = / Vf Vs pdæ e
(27)
With the choice f log P, we find the usefuu identity Aploge = Ap. Next, recall the objective E from
22
)) pda -J O.psdx.
(s)
loge)
2
Now denote by ós an arbitrary element of C (X). Then, if sf is minimizer of the (strictly convex) objective, holds that
- dr (= + rōs)
=- SsL, (5
8² log dr / Ospösdr Wōs.
(28)
2
==
Hence,
8² log
82 Ap Ap+ O.p.
Ap 5r
+ + V (pVs'
(29)
Furthermore, note that 5 is identical to
(s) ((((/vv
V + As)
Pe,ude J Supe + 2 / V log pe. Pr (30)
after integration by parts. The last term is the Fisher information of the data at t, p and plays role in the optimization.
D Motivating the partial integration in time in the loss
Note that the problems from Equation (3) corresponding to different values of L are coupled through the lerm Sept,p- This is most apparent when one discretizes the equation in time. Denote by H4 strictly increasing sequence with to 0, ln, = 1, and l4t1 4 = 644- Then, for fixed but arbitrary μ, we obtain ne coupled problems of the form
2 / [V ( po
1
5t. J (Pt. Pta. pp Vi,
min N.
(31)
Adding these problems and shifting the indices, one can eliminate Pt explicitly coupling 8t. and sa The continuous equivalent of this course an integration over i, followed by integration by parts.
22

---

E Geometric picture of the optimization problem
We omit the dependence on the parameter μ here for the sake of simpler notation and write dp for pdz for brevity. Note that the following considerations are purely formal. They are mcant illustrate gcometric picture of the optimization problems we consider. We claim no originality of these idcas; the exposition is based on Chapter 7 of [9] as well as 24 55
Otto calculus Based on the identification of the tangent space of P(X) with the space of gradients (more rigorously, at point Ps E P(X), the closure of {Vf : ff E (²)))) in (X,, Pt. Pt, RR,, see Definition 8.4.1 in ) one can view P(X) formally as a Ricmannian manifold:
Definition 1 (66) Let + p. and + o² be two curves valued in P(X) for E (t (t-e,t+e) Prlr=t = Pr. The optimal transport metric on TP(X) at P. E P(X) is such that pt P=l -
given by
(pP)( O a.p 0 / (Vs. Vs 0.p + v Vp (
+ (p: Vs 0. (32)
This formalism is commonly named after the author of 68 and is closely linked to Arnold's considerations geometric hydrodynamics b 2 As both the identification of 8, from pep. and the metric depend on Pt; the geometry defined on P((2) in this way is non-trivial.
Action of curve The optimization Equation (3) has appcaling physical interpretation: The vector field we define as tangent to the curve is, among all compatible ones, the one with the smallest integrated kinetic energy. In analogy with the physical literature, we call So Sx Vsj/ddp the action of the curve +> Pt with tangent velocity Vst. We want to stress that while this procedure is reminiscent of physical action principles, in the latter solution corresponds to a stationary point given boundary conditions at the beginming and end of the curve. The problem we consider in Equation (6 is more narrow and concerned with finding Vs, that matches a given curve + Pr. Determining curves of minimal action in P(X), leads to the Benamou-Brenier formula (α. Proposition 2.30)):
(2 01) inf (1 (2// 2
/ /e di Oept + V Vs) = Pt-o = Po, Po. = p1
p.a
(33)
with W2 the Wasserstein (or Kantorochiv-Rubiisseen distance.
Lagrangian functions The selection criterion based on kinetic energy alone is not without alternatives. In 24. the relation 0,p- - interpreted as a form of Legendre trans- form, hence s plays the role of momentum and O.pe.)) t) fx V Lagrangian. Here, we introducod the notation A to denote the pscudo inverse operator. Note that, formally, it is sensible to consider 0,p as an element of the tangent space of P(X). After all, p + Tupp P(X) for p strictly positive and small enough. In this picture, is element of the cotangent. space. The introduction of 63 addresses the two concepts and how they relate.
that of
Any function L: (p, + L(p, dep, t), strictly convex and superlinear in its second argument, can be chosen to define the minimization objective @ Details can be found in Chapter of [59] which also features a comprehensive discussion of the history and applications of this problem. In recent years, this formulation has been applied for modeling purposes, e.g. in 44. To give an example, the choice L(p, (p.d,p,tt t) = 2 fx /Aja, fx V dp for a potential V : X' + R can be used to model obstacles in the path of the samples.
"The derivation of fluid dynamics from variational principles is, of course, much older and goes back as far as Langrange's Mécanique analytique published in 1789. The variables p and O, here denote any probability density and scalar field on X.
23

---

There exist a number of partial differential equations whose solutions Pt can be described as curves of stationary action with respect to such Lagrangians, described in well as 89] Chapter 23, and 90], Chapter 8.
Schrödinger Bridge
The objective defined in Equation
5 corresponds to the choice
(pp. 1 2
/ v -App
) dp.
2
log
(34)
The associated momentum sf therefore satisfies s 8(0,p) hence - -Aps* A dep, Fokker-Planck equation. Furthermore, the action of the curve ++ P given by
6L
S (pc, Dep, 1)dt L. (2// VD vajanaPpp <² VVlogi pil dppdp))
2 + 2 ?
(J log pedpi log
Sx log pidpi
(35)
t=1
This expression known as the dual formulation of the 81 Theorem 36, except for the fact that the therein corresponds to e²/2 herc). While concerned with the path connecting Po and pr minimizing the time integral of the kinetic energy (which coincides with the transport cost), the Schrödinger-Bridee Bridge problem is concerned with finding the most likely configuration al intermediate times, subject to the information that the configuration is given al times and and assuming that the particles X, undergo Brownian motion with diffusivity e²/2. Unless P1 the result of convolution of Po with Gaussian kernel of width e, the evolution of the system towards is rare event and the most likely solution is to be understood conditional on the observation of this event.
-erooooviicc-Schrrinnngeee problem
the classical optimal transport problem
Rigorous results can be found in Section 5 of 31. Another derivation of the loss function from Equation (6) starting from the static formulation and linking to the dynamical picture presented here, can also be found in 47 Theorem 2.1. In their notation, v -s.
24

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Section 3.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations.
No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect
how much the results can be expected to generalize other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section 4.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate 'Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested results often depend on implicit assumptions, which should be articulated.
few datasets or with few runs. In general, empirical
The aut hors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algo- rithms and how they scale with dataset size.
applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honcsty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assump- tions and a complete (and correct) proof?
25

---

Answer:
[Yes]
Justification: Section 2.1, Appendix C E.
Guidelines:
The answer NA means that the paper does not include thcorctical results. All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the complemented by formal proofs provided in appendix or supplemental material.
of the paper should be
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: Section 3, Appendix A B, code implementation link in Section 1 (retracted for review).
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
If the contribution
Depending on the contribution, reproducibility ways. For example, the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to cither make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of model checkpoint, other means that are appropriate to the research performed.
be accomplished in various
While NeurlPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend the nature of the contribution. For example
(a) If the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, may be that access to the model is limited in some way (e.g., .,, to registered users), but it should be
26

---

possible for other researchers to have some path to reproducing or verifying the results.
5. Open access data and code
Question:
Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Section provides link to code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (nttps://nips
for more details.
While we encourage the release code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g. for new open-souree benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submis- sion guidelines for more details.
The authors should provide instructions data access and preparation in- cluding how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Botting/Detaile
Question:
Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
J Justification: Section 3, Appendix A B, code publication discussed in Section 1.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Figure shows replicates, results reported in Figure 6, Table are based on thousands of samples.
Guidelines:
The answer NA means that the paper does not include experiments.
27

---

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/tees split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumpliions made should be given (e-g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be carcful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negalive error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: Table 1, Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relovant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well as estimate the total compute.
The paper should disclosc whether the full rescarch project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make inlo the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips. http///wwrrp..ccoooo/ol//sccccaa
Answer:
[Yes]
Justification: Only data from numerical simulations used. We do not expect that this work has negative socictal impacts.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question:
Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
28

---

Justification: Section 4.
Guidelines:
The answer NA means that there no socictal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g, disinformation, generating fake profiles, surveillance), fairness consid- erations (e-g., deployment of lechnologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational rescarch and not tied particular applications, Ict alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it legitimate to point out that an improvement of generalive models could be used to generate deepfakes for disinformation. On the other hand, not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
the quality
The authors should consider possible harms that could arise when the technology being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional
unintentional misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategics (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mochanisms to monitor how system learns from feedback over time, improving the elficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
Appendix B.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
29

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset,
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are relcasod, the license, copyright information, and terms of use in the package should be provided. For popular datascts, paperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out. to the asset's 8 creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
Section and Appendix B.
Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details of the dataset/con hatt///o////aoaa modell as part of their submissions via structured templates. This includes details about training, license, limitations, elc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Question:
Answer:
[NA]
Justification The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor rescarch with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (TRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
30

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects rescarch. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between insti- tutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
31


Reviews:
[{"review_id": "obqoGRtU7J", "reviewer": "Reviewer_zRLH", "summary": "This paper develops models of population dynamics in physical systems that exhibit stochastic and mean-field effects, influenced by physics parameters. The goal is to create models that can efficiently predict system behavior as alternatives to classical numerical methods. By utilizing the Benamou-Brenier formula from optimal transport and action matching, the approach involves solving a variational problem to infer gradient fields that approximate population dynamics. These gradient fields enable the generation of sample trajectories that mimic physical system dynamics under various physics parameters. The study highlights the importance of combining Monte Carlo sampling with higher-order quadrature rules for accurate estimation and stable training. The models are demonstrated to perform well on Vlasov-Poisson instabilities and high-dimensional particle and chaotic systems, outperforming state-of-the-art diffusion-based and flow-based models that rely solely on time and physics parameters.", "questions": "1)\tAuthors aims to learn population dynamics $\\rho_{t,\\mu}$ instead of learning the dynamics of individual trajectories $t\\to X_{t,\\mu}^i$ for all $i$/. What are the assumptions on model so that it admits density? If not, would this work apply to the situation under weak convergence where no density is assumed. Moreover, it is not clear why the equations (5) to (7) should admit solutions in the strong forms.\n2)\tParametrizing the $s_{t,\\mu}$ with weight modulation as in CoLoRA [38] only applies to deterministic time-dependent dynamical systems. However, one goal of the work was stated as learning the population dynamics should allow for seamless treatment on deterministic and stochastic systems. How the latter can be handled with the corresponding parametrization is unclear. The forms of the layers assumes some low-rank structure, and only the weight modulations $\\phi$ is assumed to depend on time and parameters.  As such, are stochastic systems omitted from the problem definition unlike the initial motivation?\n3)\tSolving eq (7) from data is challenging and prone to potential numerical issues. As a remedy, a combination of higher order numerical quadrature and MC sampling strategy is proposed. However, the details of this approach would be better to provide with a stability analysis on the training as mentioned in the text.", "limitations": "Yes, limitations of their work is adequately addressed. No potential negative societal impact of their work is identified.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1)\tEfficient parametric model reduction: The model is demonstrated to reduce inference runtime significantly compared to standard diffusion- and flow-based models by leveraging minimal-energy vector fields.\n2)\tAccurate and stable dynamics learning: It captures the coupling over time steps accurately, using higher-order quadrature schemes for estimating time integrals, which enhances training stability.\n3)\tHigh accuracy and reduced runtime: Achieves error rates comparable to state-of-the-art methods while reducing inference runtime by 1-2 orders of magnitude.", "weaknesses": "1) The model assumes access to a dense set of time points for the Gauss-Legendre quadrature, which may not be applicable when only a few time samples are available. This was already mentioned as a limitation of the current work. As this forms the main part of the approach, the practical benefit would be limited. \n2) Regarding the vector field complexity, the model seeks a vector field that minimizes kinetic energy, but in some cases, this may be more complicated than other vector fields that produce the same population dynamics. Examples include situations where the minimal-energy field varies with time, making it challenging to determine the appropriate energies to use for different problems.", "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 respond to the following points:\n\n\"The model assumes access to a dense set of time points for the Gauss-Legendre quadrature [...] the practical benefit would be limited\"\n- Having data at a good amount of time points is a common situation in the setting of surrogate modeling and model reduction, which we consider here. The data are generated with a high-fidelity numerical model, which has to take small time steps for numerical stability reasons. With HOAM we show that not only we can leverage the dense sampling in time, but that it is absolutely critical to exploit it with appropriate numerical quadrature as existing methods fail due to inaccurate estimates of the time integral (see Figure 2 in paper as well as new results to provide further evidence in the uploaded PDF). \n- We stress that the problems that we examine here are real problems from the physics literature (see references) that are intractable with current surrogate modeling techniques (see references in paper). With HOAM we construct fast, predictive surrogate models for such problems for the first time. Additionally we show that other reasonable approaches like NCSM or CFM perform worse while incurring orders of magnitude higher inference costs. Thus, this shows that there is a large practical benefit to HOAM—creating fast surrogate models of stochastic and mean-field systems.\n- If the paper gets accepted, we will emphasize more clearly in the Introduction that in our problem setting it can be reasonably expected that data are available at many time points.\n\n\\\n\"[the learned vector fields] may be more complicated than other vector fields that produce the same population dynamics.\"\n- This is an interesting point, which we briefly discuss in the Conclusion section.\n- The main challenge is mathematically formalizing what “complicated” and “easy” means in the context of neural network approximations, which will also depend on the problem at hand.\n- We build on the kinetic energy because our results show it is a reasonable energy for many different problems that also links nicely with the theory developed in optimal transport theory (see our appendix). One desirable feature of the kinetic energy is that $\\nabla s$ of the minimum kinetic energy vector field is identically zero when the population dynamics are stationary - this means that if the population dynamics are stationary, then sampling with the minimum kinetic energy field is trivial as the samples don’t move. This is not necessarily the case for other compatible vector fields.\n- If this paper gets accepted, we will expand the comments in the Conclusion section where we briefly discuss this point.\n\n\\\n\"Authors aim to learn population dynamics [...] What are the assumptions on the model so that it admits density? If not, would this work apply to the situation under weak convergence where no density is assumed. Moreover, it is not clear why the equations (5) to (7) should admit solutions in the strong forms.\"\n- Objective (7) (and (5), the special case where $\\varepsilon = 0$) is expressed entirely in the form of expectations values w.r.t. $\\rho$ - both are meaningful if no density is admitted and can be evaluated on empirical distributions (sums of dirac masses). We will add a comment about this if the paper gets accepted. \n- We use (5) and (7) to formally derive an optimality criterion for the learned gradient field $\\nabla s^*(\\theta)$. This field is smooth as it is parametrized by a neural network with smooth activation functions. For any empirical distribution (samples from $\\rho_{t, \\mu}$), we therefore obtain a vector field that is suited for inference.\n\n\\\n\"Parameterizing the with weight modulation as in CoLoRA [38] only applies to deterministic time-dependent dynamical systems. [...] As such, are stochastic systems omitted from the problem definition unlike the initial motivation?\"\n-  We do consider stochastic systems in the numerical experiments (e.g., particles in aharmonic trap, high dimensional chaos). What we parametrize with CoLoRA is the gradient field, which is deterministic even if the system is stochastic (see line 37 in paper, as well as Section 2.1). Thus, the CoLoRA parametrization applies independent of whether the system is deterministic or stochastic. This allows a seamless treatment of deterministic and stochastic systems, which we consider a major advantage of our approach.\n\n\\\n\"Solving eq (7) from data is challenging and prone to potential numerical issues. [...] would be better to provide with a stability analysis on the training as mentioned in the text.\"\n- An empirical stability analysis is provided in Figure 2, where we show that the higher-order quadrature is essential for stabilizing the training. We stress that our results show that higher-order quadrature is not just a nice add-on to improve the accuracy a bit but that it is a core component to make the approach work in practice. In fact, when just using MC, the error blows up during training (see also the new results in the uploaded PDF). \n- A formal stability analysis is work in progress. As an outlook, we will include the following if this paper gets accepted: The difficulty is that introduction of a numerical quadrature breaks the exact correspondence between the objective\n$$\n\\int_0^1 \\left(\\frac{1}{2}\\mathbb E_{\\rho_{t, \\mu}}\\left[|\\nabla s_{t, \\mu}|^2\\right] + \\mathbb E_{\\rho_{t, \\mu}} \\left[ \\partial_t s_{t, \\mu} \\right] \\right) \\mathrm{d} t - \\mathbb E_{\\rho_{t, \\mu}} \\left[ s_{t, \\mu} \\right] \\big |_0^1\n$$\nand the continuity equation. Assume $\\rho$ admits a density, the derivative $\\partial_t$ is approximated exactly and denote by $(\\dots)^n$ the value at time $t_n$ and by $w^n$ quadrature weights. The mismatch\n$$\n\\left| \\sum_n w^n \\int s^n \\partial_t \\rho^n dx - \\left( - \\sum_n w^n \\int \\rho^n \\partial_t s^n dx + \\int \\rho s \\big |_0^1 dx \\right) \\right|\n$$\nequals the numerical integration error of $\\frac{d}{dt}\\int s\\rho \\\\, dx$ in time.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "mZcu777dvi", "reviewer": "Reviewer_gej1", "summary": "The authors focus on learning models for population dynamics of parameterized physical systems that exhibit stochastic and mean-field effects in time. To do so, the authors leverage the Benamou-Brenier formula to learn gradient fields that transport the probability density as time evolves, and that enable the generation of sample trajectories reflecting the dynamics of the population. Numerical experiments show compeling results and state-of-the-art performance in high-dimensional particle systems and in chaotic systems.", "questions": "I have the following minor questions:\n- In (6) it seems that the integrand should be evaluated on $s_t$ instead of $s$ for consistency with (4). \n- In (7) there seems to be a parenthesis missing that factors out the density $\\rho_{t,\\mu}$.", "limitations": "The authors adequately address the limitations of their work in Section 4. These are mostly technical. For instance, they require a high sampling rate in time, which is not always available in practice. This limits the applicability of the method, but does not detract from their main contribution.", "rating": 9, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The contributions of the paper are novel, and the presentation is very clear. I commend the author's attempts at making the paper reasonably self-contained by adding context in Appendices C through E. Each of the theoretical developments are clearly motivated. Finding a vector field as proposed enables the interpolation of probability measures at different times. This notion has broad impacts beyond those stated in the paper. It is also interesting that the authors illustrate the importance of using a proper quadrature in time as opposed to random sampling. The experiments are also compelling, in particular the ability to resolve the low probability connection between two high probability regions as shown in Fig. 5.", "weaknesses": "To this reviewer the paper does not have any evident weakness. However, when introducing a high-order quadrature in time, which improves the performance, the authors do not elaborate on why one would, or would not, expect further improvements if similar quadrature rules were used in the spatial variable or in the physical parameters.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive and thorough review. We will address all the points brought up.\n\n\"In (6) it seems that the integrand should be evaluated on  instead of for consistency with (4).\"\n- Thank you, yes, this should be $s_{t, \\mu}$.\n\n\"In (7) there seems to be a parenthesis missing that factors out the density rho\"\n- Thank you, we will fix this in the final version, if the paper gets accepted.\n\n\"the authors do not elaborate on why one would, or would not, expect further improvements if similar quadrature rules were used in the spatial variable or in the physical parameters.\"\n- Typically the dimension of $x$ is too high for quadrature rules to be effective, which is why we use Monte Carlo for it. We also specifically assume that our data is in the form of samples, so evaluating $\\rho$ at quadrature points would require some form of density estimation, a non-trivial task in high dimensions.\n- For the parameter $\\mu$, higher-order quadrature rules can be used but in surrogate modeling we typically have data points only at very few $\\mu$ training parameters, which makes higher-order quadrature rules difficult to apply.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "S1yqpeQ4yn", "reviewer": "Reviewer_rAtu", "summary": "The authors present a framework for learning a time-dependent gradient vector field (parametric model) that describes how the particle evolves in the state space. The idea is based on the action-matching framework (https://arxiv.org/abs/2210.06662), which learns a time-dependent gradient field that interpolates between marginals at different times in a simulation-free manner. The authors propose using high-order quadrature rules for evaluating nested integrals in the objective. They demonstrate the forecasting ability of their framework on diverse examples in high dimensions and show that it outperforms other diffusion-based and flow-based models in inference runtime and accuracy", "questions": "General comments:\n\ni) Maybe, I am not understanding things here.. Can the authors elaborate a bit more on the O(K\\tau) complexity? If for instance, I train a velocity field using continuous normalizing Flow (simulation-based) using empirical marginals at different times. I can then integrate the learned velocity in time to get samples, in the same way, the authors draw samples from the gradient field through Langevin dynamics. Where is the extra \\tau in complexity coming from? \n\nii) Or, is this framework a simulation-free strategy to train a Continuous Normalizing Flow but adapted for multiple snapshots? This was not so clear in the paper.\n\niii) The authors learn a gradient field as a function of parameters and time. Can I authors comment a bit on the stability of the learned velocity during inference time? \n\nSpecific comments\n\ni) In Figure 5, a color bar showing the scales would help with the interpretation.\n\nii) In Figure 4, a label to show what the two curves (blue and orange) are is needed. What is the ground truth and what is the prediction?", "limitations": "-", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Developing efficient and high-fidelity surrogate models is of great importance to the computational science and engineering community. The authors take a step forward in this direction by learning the effective population dynamics of the underlying physical process. The paper is well-written, and the related works are adequately referenced.", "weaknesses": "The objective used by the authors is based on the action-matching paper, now adapted to include the dependence on the parameters (\\mu). From a methodological standpoint, the authors' only important contribution is using a higher-order quadrature scheme to discretize the nested integrals", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough and insightful comments:\n\n\"The objective used by the authors is based on the action-matching paper, now adapted to include the dependence on the parameters (\\mu) [...] the authors' only important contribution is using a higher-order quadrature scheme to discretize the nested integrals.\"\n\n- We stress that the higher-order quadrature scheme is not just a nice add-on to improve the accuracy a bit but the **core component to make the approach work in practice at all**. In fact, we show that action matching is unstable even on simple problems (Figure 2, see HOAM vs rand). That AM is unstable on presumably simple problems is also documented in community implementations such as the TorchCFM library. Moreover we show that AM fails on systems with complex dynamics and high dimensions (Strong Landau and 9D chaos, see Figure 8 and 9 in appendix). By contrast HOAM succeeds, allowing us to derive predictive surrogate models with low inference costs of such challenging problems for the first time.\n- We further uploaded a PDF that shows more detailed results showing that HOAM is essential for stable training as the method proposed in the action-matching framework fails on the problems we consider (see global response).\n- In terms of novelty of the objective function, variants of it have been widely used in other works far earlier than the action-matching paper (this is also stated in the AM paper). For example, there is [Reich 2011, page 240], with the earliest form appearing in [Otto & Villani 2000, Section 3, and Bennamou & Brenier 2000, eq. 35]. Thus, rather than inventing a new objective, we made the objective computationally tractable, which we consider an important contribution. \n\n\\\n\"Can the authors elaborate a bit more on the O(K\\tau) complexity? [...] Where is the extra \\tau in complexity coming from?\"\n- The standard way of handling a time $t$ and parameter $\\mu$ dependency with NCSM and CFM is conditioning (see literature review in paper). As we point out, this means that for each of the K time steps, a separate inference problem has to be solved, which is expensive: $\\tau$ refers to the number of steps taken in solving the SDE/ODE in one inference step (for one out of the K time steps) in conditioned CFM/NCSM. Thus with conditioned CFM and NCSM a separate SDE must be solved at every $t$ and $\\mu$ for which one wants samples.\n- By contrast, $\\tau$ does not appear in the complexity for HOAM because $\\nabla s$ evolves particles such that they match $\\rho$ at each time. Thus physical time $t$ and the SDE time $\\tau$ are aligned.\n\n\\\n\"[...] is this framework a simulation-free strategy to train a Continuous Normalizing Flow but adapted for multiple snapshots?\"\n- As we discuss above, simple conditioning is not competitive for the surrogate modeling task that we consider. It is absolutely possible that other methods for training CNFs can be extended and modified to be more efficient for surrogate modeling, but instead of tweaking another method we opted to go with our approach as it is simulation-free and naturally couples the physical time t with the sampling time so that in one inference step a whole sample trajectory is obtained. This is what ultimately provides speedups, which are key for surrogate modeling.\n\n\\\n\"The authors learn a gradient field as a function of parameters and time. Can I authors comment a bit on the stability of the learned velocity during inference time?\"\n- It is sufficient that the drift and diffusion are uniformly Lipschitz continuous such that the inference SDE (or ODE when $\\epsilon = 0$) is well-posed. See, for example, [Ambrosio, Gigli & Savaré 2005, Lemma 8.1.4]. The diffusion is a constant in this work and because we parametrize with a CoLoRA neural network, the gradient $\\nabla s$ is smooth in $t, x$, which is sufficient for the well-posedness. We will add comments about this in the paper if it gets accepted.\n- Regarding the analytical vector field we hope to converge to, its properties depend on the data. When $t \\mapsto \\rho_{t, \\mu}$ describes a regular curve in Wasserstein space [Gigli 2012, Definition 2.7], then the inference ODE is well-posed [Gigli 2012, Theorem 2.6].\n\n\\\nSpecific comments about improving figures:\n- We will address these, if the paper gets accepted.\n- The ground truth in Figure 4 is blue.\n\nReferences:\n- [Ambrosio, Gigli & Savaré 2005] Luigi Ambrosio, Nicola Gigli, and Guiseppe Savaré. Gradient Flows. Lectures in Mathematics ETH Zürich. Birkhäuser-Verlag, Basel, 2005. doi:10.1007/978-3-7643-8722-8\n- [Bennamou & Brenier 2000] J.-D. Benamou and Y. Brenier. A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375–393, Jan. 2000. doi:10.1007/s002110050002 \n- [Gigli 2012] Nicola Gigli. Second Order Analysis on (P2(M),W2). Memoirs of the American Mathematical Society, Volume 216; 2012\n- [Otto & Villani 2000] F. Otto and C. Villani. Generalization of an Inequality by Talagrand and Links with the Logarithmic Sobolev Inequality. Journal of Functional Analysis, 173(2):361–400, June 2000. doi:10.1006/jfan.1999.3557 \n- [Reich 2011] Reich, S. A dynamical systems framework for intermittent data assimilation. Bit Numer Math 51, 235–249 (2011). doi:10.1007/s10543-010-0302-4", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": {}}]

Author rebuttal:
We thank the reviewers for taking the time to read the paper and very much appreciate their detailed comments. Below we provide a detailed response. We summarize here some of the main points and how we addressed them: 
- We now stress that the higher-order quadrature scheme that is introduced by our approach HOAM is not just a nice add-on to improve the accuracy a bit but the core component to make the approach work in practice at all. While variants of the objective that we consider have been around in the literature for a long time, we show that higher-order quadrature in time makes the objective computationally tractable and so avoids the instabilities in training that we observe in our experiments and that agree with the instability results shown in the literature (see detailed comments below). **We further uploaded a PDF that shows more detailed results showing that HOAM is essential for stable training.**
- The problems that we examine here are intractable with current surrogate modeling and model reduction techniques (see references in paper). The improvements made with HOAM allow us to construct fast, predictive surrogate models for such problems for the first time. 
- We address concerns about densely sampling in time: We consider the setting of surrogate modeling and model reduction where data are typically generated with a high-fidelity numerical model, which has to take small time steps due to numerical stability constraints. In fact, we show that HOAM performs well on a range of real problems from the physics literature (see references) for which data are generated with standard tools from the physics community. Thus, the data requirements align with what is typically available for the problems that we aim to address with our HOAM approach.

Meta-review: This work is concerned with inference in dynamic systems, where the authors propose an approach that combines Monte Carlo sampling with higher-order quadrature rules. This work was reviewed by three reviewers, where one reviewer argued strongly for acceptance, one gave a borderline accept, and the third a borderline reject.

The reviewers largely agree on the strengths related to practical benefits and efficient parametric model reduction. The work shows that the model captures the temporal coupling of physical systems accurately by utilizing higher-order quadrature schemes to estimate time integrals. The idea itself can be seen as rather straightforward, but the idea is well presented.

The large spread in the reviewer scores is perhaps best explained by the presentation in the paper that slightly deviates from what one would typically expect from a NeurIPS paper. After reading the paper, I lean towards accepting this work despite the slightly mixed reviews. I trust that the authors will fix the issues (which are minor but many) that came up during the review/discussion phase.
Decision: Accept (poster)