Title: Parametric model reduction of mean-field and stochastic systems via higher-order action matching
Venue: NeurIPS 2024
Accepted: yes
Keywords: partial differential equations

reduced modeling

model reduction

stochastic dynamical systems

generative models
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

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

1 Introduction

Predicting the behavior of time-dependent processes X t,µ over time t and across varying physics parameters µ is a key challenge in computational science and engineering [46, 65]. The dynamics of X t,µ typically are described by systems of (stochastic) differential equations, which are derived from physics models and can be computationally expensive to simulate [40, 32]. Thus, it is desirable to learn reduced or surrogate models that can be rapidly evaluated to predict the system behavior across varying physics parameters [72, 10, 11, 45].

Reduced modeling via learning population dynamics Given a data set of samples, i.e., realizations of the random variable X t,µ on a suitable domain X ⊆ R d ,

we aim to learn a dynamical-system reduced model to rapidly predict samples that approximately follow the same law ρ t,µ as X t,µ over time t and varying physics parameter µ . We refer to the evolution of ρ t,µ in time as population dynamics. Learning the population dynamics instead of learning the dynamics of the individual trajectories t ↦→ X i t,µ for all i = 1 , . . . , N x and µ can be beneficial: There are cases where ρ t,µ does not change in time, yet every sample trajectory t ↦→ X i t,µ 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 a 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 Fokker-Planck 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 [8] as well as the so-called action-matching loss introduced in [61], we pose a variational problem to learn gradient fields ∇ s t,µ so that the continuity equation corresponding to the vector field given by ∇ s t,µ approximates the population dynamics ρ t,µ of the samples (1). In the spirit of reduced modeling [72, 10, 11, 45], we seek a vector field s t,µ that generalizes to different values of the physics parameters µ . We therefore optimize for s t,µ that minimizes the average objective of a variational problem over all parameters µ ∼ ν , where ν describes the distribution of parameters on the domain D ⊂ R p . We parametrize s t,µ with a neural network with weight modulation [39, 13] so that it can be evaluated quickly over t and µ .

Rapid sample generation in inference phase Predictions at inference time at new physics parameters µ are made by sampling based on the vector field ∇ s t,µ , which means that our approach represents ρ t,µ through the application of ∇ s t,µ on an initial condition. Importantly, time t in the inference step corresponds to the time of the physics problem so that in one inference step a whole sample trajectory is obtained, rather than a sample at one specific time point as in regular conditioning-based methods (see literature review). Thus, we can rapidly generate samples that follow the law ρ t,µ in 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 (1), we introduce an empirical loss (8) that utilizes higher-order quadrature [27] in the time direction so that the learned ∇ s t,µ accurately captures the dynamics over time t . 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 ∇ s t,µ that accurately capture the evolution in time t and that generalize across physics parameters µ .

Literature review We review relevant literature; see Figure 1 for an overview.

Non-intrusive and data-driven surrogate modeling There is a range of 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 [71, 76, 86, 12, 46, 58, 92] as well as neural network-based methods such as neural ordinary differential equations [19, 28, 48]. There also are methods for stochastic systems [51, 42, 88, 19, 28, 73, 21]. However, all of these methods ignore physics parameter dependencies and/or aim to learn the sample dynamics, whereas we focus on parametric population dynamics.

Population dynamics and trajectory inference Learning population dynamics has been considered 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 [34, 30, 93, 75, 85, 47]; however, many of these approaches [17, 84] are

· · ·

· · ·

ρ

0

ρ

t

n

ρ

t

n

+1

ρ

1

score-based

modeling

HOAM

· · ·

· · ·

· · ·

· · ·

· · ·

· · ·

flow-based modeling

· · ·

simulation-based and thus require integrating dynamics during the training or parameterizing 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, 79, 36, 41, 81, 82] and flow-based modeling [2, 54]; see [1] for a detailed review. These approaches are not taking into account time t because they learn paths between a reference and a target distribution only. There are works that condition on time t and a parameter µ such as [68, 14, 26, 37, 33, 38, 52], but this requires then generating a path for each time step at inference time, which is computationally expensive. Furthermore, the conditioning on time t means that the target distribution ρ t,µ at each time t and µ is different, and thus a separate hyper-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 [15, 78, 50] compute transport-based solutions but parametrize different quantities than our approach, require actively sampling data, and ignore physics parameters µ . We note that there also is work on forecasting with diffusion- and flow-based modeling [68, 62, 18, 20], which is a different task than our task of predicting across varying physics parameters.

Optimal transport Besides the machine learning literature, variational approaches for inferring vector fields are extensively used in optimal transport theory [5, 4]. Of particular importance to us is the formulation by Benamou and Brenier [8]. The Bennamou-Brenier formula describes a joint optimization problem over vector fields and paths in probability space and the action matching loss [61] is the restriction of this optimization problem to the case of a fixed path and the vector field parametrized by a neural network, which are core building blocks for us that we show can be used together with a parameter dependency.

Contributions We summarize our contributions:

(a) Developing a loss to learn population dynamics that remain valid across varying physics parameters by building on optimal transport literature [8] and action matching [61].

(c) Demonstrating on a range of physics problems from Vlasov-Poisson instabilities to high-dimensional chaotic systems that our approach leads to (i) accurate predictions of population dynamics and (ii) orders of magnitude speedups in inference/prediction over

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

classical methods that numerically 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://github.com/julesberman/HOAM .

2 Method

2.1 Parameter-dependent population dynamics

Continuity equation Let us consider data (1) corresponding to the probability measure ρ t,µ , which is absolutely continuous for t ∈ [0 , 1] and µ ∈ D . We use the same notation for the measure and its density. The density ν of µ is also assumed to be absolutely continuous on D . We consider population dynamics of X t,µ ∼ ρ t,µ that can be described by the continuity equation with the initial condition ρ t =0 ,µ =: ρ 0 ,µ and vector field v t,µ . Notice that in our case the continuity equation (2) depends on the physics parameter µ ∼ ν . There can be many vector fields v t,µ that lead to the same population dynamics (2). For example, if v t,µ is a vector field that describes the dynamics of ρ t,µ via (2), then another vector field is given by v ′ t,µ = v t,µ + w/ρ t,µ with any other w that satisfies ∇· w = 0 as long as ρ t,µ is positive.

Uniqueness via gradient fields and the corresponding elliptic problems Because we aim to learn a vector field from sample data (1) that describes the population dynamics (2) of the corresponding law ρ t,µ , it is helpful to remove this non-uniqueness. One way to do so is to restrict the vector field to v t,µ = ∇ s t,µ so that it is a gradient field [4, p. 45]. Plugging v t,µ = ∇ s t,µ into (2), together with the assumptions ρ t,µ > 0 and ∫ X ∂ t ρ t,µ d x = 0, leads to parametric elliptic problems in s t,µ

with coefficient function ρ t,µ , right-hand side (source term) ∂ t ρ t,µ , and homogeneous Neumann boundary conditions ρ t,µ ∇ s t,µ · ˆ n = 0 on ∂ X with normal vector ˆ n for all t ∈ [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 s t,µ via optimization:

for each t ∈ [0 , 1] and µ ∈ D . The space H 1 ( ρ t,µ , X ) contains functions s with ∫ X |∇ s | 2 ρ t,µ d x < ∞ , which is the energy (semi-)norm corresponding to the ρ t,µ -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 1 ( X ) under strong assumptions on ρ t,µ such as uniform boundedness by a positive constant for all t and µ ; see [29, Proposition 2.2] and [11, Section 3.2]. The theory of optimal transport allows treating the much more general case when ρ t,µ is not uniformly bounded away from zero and possibly atomic; we refer to [8] and [74, Section 5.3.1] for details. Among all vector fields v t,µ that are compatible to ρ t,µ in the sense of (2), gradient fields ∇ s t,µ have the smallest associated kinetic energy 1 2 ∫ X | v | 2 ρ t,µ d x , which is the objective considered in [8]. 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 t ↦→ ρ t,µ . The metric is the inner product of L 2 ( ρ t,µ , X , R d ). This is the weak Riemannian structure of P ( X ) equipped with the Kantorovich-Rubinstein metric and described in detail in [5, Chapter 8]. We give a short description in Appendix E.

Energy functional with entropy term Instead of the energy (4), we can also use other choices of the energy to 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 ( X ), the entropic optimal transport or Schrödinger bridge problem [77, 56],

which depends on /epsilon1 ≥ 0. The energy E /epsilon1 t,µ is of particular interest for two reasons: One, the Euler-Lagrange equation of (5) in strong form is the Fokker-Planck equation for s /epsilon1 t,µ :

∂ t ρ t,µ = -∇· ( ρ t,µ ∇ s /epsilon1 t,µ ) + /epsilon1 2 2 ∆ ρ t,µ , again with homogeneous Neumann boundary conditions for all t ∈ [0 , 1] and µ ∈ D ; see Appendix C. This means we can efficiently generate samples after learning s /epsilon1 t,µ via corresponding stochastic differential equations (SDEs). Two, it can be interpreted as regularizing the field s t,µ , which we discuss in Appendix C.

2.2 Loss for learning vector fields over time t and physics parameter µ

Variational formulation over t and µ So far we just carried along time t and physics parameter µ but did not address them in the variational problems, i.e., we had separate variational problems (4) for all t ∈ [0 , 1] and µ ∼ ν . We now propose to consider the average energy over t and µ to infer a map s : [0 , 1] ×D → H 1 ( ρ t,µ , X ) , ( t, µ ) ↦→ s t,µ , which is called a solution map in reduced modeling [72, 10, 11, 45],

Notice that time t and physics parameter µ have two different effects on the gradient field ∇ s t,µ : Time t couples the elliptic problems (i.e., (3) for /epsilon1 = 0) via the time derivative ∂ t ρ t,µ ; 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 µ ′ ∈ 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 an accurate approximation of the time integral to ensure the coupling between the time points is reflected in s t,µ ; see Section 2.3.

Loss for learning gradient fields from samples over t and µ The energy E t,µ defined in (4) as well as the energy E /epsilon1 t,µ defined in (5) leads to a loss that can be estimated from samples (1). The quantity ∂ t ρ t,µ appears in (4) and (5), which is typically unavailable when we have access to data samples (1) only. Integration by parts of the term involving ∂ t ρ t,µ eliminates it, see also Appendix D. We arrive at

Note that this loss is comprised only of expectation values with respect to ρ t,µ and is therefore well-defined also for empirical distributions. The choice ε > 0 assumes that the Fisher information of ρ t,µ is finite.

Remark 1. Loss functions of the form as (7) but without the parameter dependence have been used in [61] and [47, Theorem 2.1]. In fact, the case with /epsilon1 = 0 appears already in [8, Equation 35] and [64, Section 3]. We build on these results but work with population dynamics that depend on physics parameters, which leads to the loss shown in (7) .

2.3 Parameterizing the vector field, estimating the loss from data, sampling

Parametrizing s t,µ with weight modulations We parametrize the vector field s t,µ 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 ( x ) = Wx + φ ( t, µ ) ABx + b , where W is a weight matrix, A,B are low-rank matrices, b is a bias vector, and φ ( t, µ ) ∈ R is a scalar weight modulation; see Appendix B. Only the weight modulations φ ( t, µ ) depend on time t and physics parameter µ . We use a hyper-network h : [0 , 1] ×D× Ψ → R that depends on the weight vector ψ ∈ Ψ ⊆ R q to map t and µ to the modulation weights φ ( t, µ ) = h ( t, µ ; ψ ). The weights W,A,B,b , which are independent of t and µ , over all layers are collected into the weight vector θ ∈ Θ ⊆ R q ′ . Typically q /lessmuch q ′ . Using the hyper-network encourages continuity of s t,µ in time t , which is key for many physics problems [13].

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 i t,µ , time t , and physics parameter µ 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 3 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 to be estimated with particular high accuracy to ensure the coupling between the time points as well as the coupling to the boundary terms to match the path from ρ 0 ,µ at time t = 0 to ρ 1 ,µ 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 GaussLegendre quadrature [27]; see Appendix A. 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 and at the Gauss-Legendre nodes in the case of Gauss quadrature. If the data set (1) does not contain samples at these time points then we interpolate the data to the appropriate times. We note that for Simpson's quadrature, interpolation is typically unnecessary 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 ˆ E n x ∼ ρ [ f ] := ∑ n i =1 f ( X i ) where X 1 , X 2 , . . . , X n ∼ ρ . Then, the empirical loss with mini-batching of sizes n x , n µ and n t quadrature points in time is given by

where w n are numerical quadrature weights and t n are the corresponding nodes; see Appendix A for the Simpson's quadrature and Gauss-Legendre 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 ∇ s t,µ , which approximates the law ρ t,µ of X t,µ . Because we train with the loss (7), we integrate the SDE d ˆ X t,µ = ∇ s t,µ ( ˆ X t,µ )d t + /epsilon1 d W t , where W t are Wiener processes and /epsilon1 is the same /epsilon1 that is used in the training loss (7); see Appendix C. As initial condition, we use samples from ρ 0 ,µ at time t = 0. Of course other sampling schemes can be used [70].

Notice that the time t in the SDE used for generating samples is the same time as of the physics problem and thus of the sample trajectory. This means that the costs of the inference step of our HOAM for generating a trajectory of length K scales as O ( K ). In contrast, introducing a conditioning on time and physics parameter in, e.g., noise-conditioned score matching (NCSM) [80] and conditional flow matching (CFM) or stochastic interpolants [2, 54] requires inferring a separate sampling path for each t and µ pair of interest. In particular, the inference costs of CFM scale as O ( Kτ ), where τ is the number of steps 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 ( Kτσ ), 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. Harmonic oscillator: A collection of particles evolves in four-dimensional phase-space

10

0

.

0

0

.

5

1

.

0

time

-

200

0

200

q

(

s

)(

t

)

Gauss

Simp

trap

AM

10

-

8

10

-

5

10

-

2

rel.

err.

in

∫

1

0

q

(

s

)(

t

)

dt

0

10000

0

-

4

-

8

Loss

AM

HOAM

HOAM-G

HOAM-S

AM

10

-

5

10

-

3

10

-

1

10

1

1

subject to a quadratic potential V ( x ) = -1 2 ω 2 | x | 2 . In the experiments shown ω = 8. The particles are initially at rest and follow a normal Gaussian distribution in space with mean m 0 = [1 , 1]. To avoid the formation of a singularity at ωt = 1 2 π , we add white noise of strength η = 5 × 10 -2 to the momentum equation. For the case η = 0, we have analytical expressions for ρ and s , see Appendix B.1.

Two-stream instability We numerically solve the Vlasov-Poisson partial differential equations using a particle-in-cell method to generate samples (1). We consider the two-stream instability [22, 43] in a 1D1V configuration with collisions [87, Sec 2(b)(i)], with β = 10 -3 and v 0 = 1 as in [49]. These collisions lead to stochastic sample trajectories. The parameter µ ∈ [1 . 2 , 1 . 9] is a normalization constant related to the Debye length [83]. It controls the ratio between electric and inertial effects in the simulation.

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 [7, 35, 43]. The parameter varies as µ ∈ [1 . 3 , 2 . 0].

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 (sixdimensional) setup. A perturbation in the x 1 -direction leads to the formation of phase-space structures [59]. The parameter µ ∈ [0 . 5 , 1 . 5] is the mass of the charged particles.

High-dimensional chaos A Rayleigh-Bénard convection leads to a density gradient that sets a 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 µ ∈ [13 . 7 , 14 . 4] is the reduced Rayleigh number.

Particles in aharmonic trap We consider 50 particles in an aharmonic trap [16], which lead to 100-dimensional samples X i t,µ that encode the positions of the particles. The particle positions are governed by a stochastic differential equation. The parameter µ ∈ [0 . 3 , 0 . 9] controls the velocity of the trap.

Setup We compare our higher-order action matching (HOAM) to the original version of action matching (AM) [61], where we handle the parameter dependence on µ in the same way as in our approach. Additionally, we compare to noise-conditioned score matching (NCSM) where samples are generated via annealed Langevin dynamics [80] and conditional flow matching (CFM) [2, 54], for which we condition on time t and µ ; see Appendix B.

HOAM stabilizes training with higher-order quadrature In Figure 2 we consider the harmonic oscillator example. We learn the field s t and plot q ( s )( t ) = ˆ E n µ µ ∼ ν ˆ E n x x ∼ ρ t,µ [ 1 2 |∇ s t,µ | 2 +

True

HOAM-S

HOAM-G

AM

NCSM

CFM

∂ t s t,µ + /epsilon1 2 2 ∆ s t n ,µ ] over time t , which is the function that needs to be integrated in time to estimate the loss (7). As Figure 2 (left) shows, this function is far from smooth and exhibits several sharp peaks, which make estimating the loss challenging. In AM [61], the time integral is estimated by averaging samples uniformly taken in time, which is equivalent to Monte Carlo estimation. Figure 2 (center left) shows the relative error in estimating the time integral using Monte Carlo integration as done by AM versus numerical quadrature as in our HOAM. The trapezoidal rule, the composite Simpson's rule, and Gauss-Legendre quadrature all produce highly accurate estimates, whereas Monte Carlo integration yields inaccurate estimates of the integral with high variance.

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 2 (center right). In contrast, our HOAM, where training is done with higher-order quadrature (in this case Simpsons' rule), the loss curve is stable and of low variance. In Figure 2 (right) we plot the mean Wasserstein distance over time of solutions generated via a gradient field s trained with with Simpson'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 instabilities lead AM to fail. In contrast, our quadrature-based HOAM is consistently stable and provides orders of magnitude more accurate results.

Accurate predictions with speedups for Vlasov-Poisson equations Our VlasovPoisson 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 x . Thus, reduced modeling with HOAM is well suited for this problem because the natural

10

-

3

10

-

2

10

-

1

HOAM-S

HOAM-G

AM

NCSM

CFM

0

40

10

-

2

10

-

1

10

0

0

40

0

40

0

40

0

40

Electric Energy

Time

dynamics to learn from such a system are the population dynamics ρ t,µ rather than the sample dynamics; see Appendix B.2. We observe the particles computed with a particle-incell method and learn the gradient field ∇ s t,µ with the proposed HOAM approach. For a test physics parameter µ that controls the wave number, we then generate samples with ∇ s t,µ and plot a histogram in Figure 3 for the bump-on-tail (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 over the classical particle-in-cell (full) model, as NCSM and CFM lead to 1-2 orders of magnitude longer inference times than HOAM and the full models.

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 is 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 of the full model increases by almost two orders of magnitude. In contrast, the runtime of our HOAM reduced model increases only from 6 to 8 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 a quantity of interest from the generated samples over time t for the test physics parameters, which we plot in Figure 4 and its relative error averaged over time (e.e.) in Table 1 (see (25)). Our HOAM approximates the electric energy well at later times, whereas NCSM and CFM lead to poorer approximations at later times t . 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 of the electric energy with orders of magnitude speedups; see Table 1 and Figure 3 as well as Figure 7 in the appendix.

Speedups in inference step (predictions) Recall two limitations of introducing a time and physics parameter dependence in NCSM/CFM via conditioning (see page 3 and Section 2.3): (i) For each t and µ , a separate sampling path has to be computed, which leads to orders of magnitude higher inference runtimes than in HOAM; see Table 1, Section 2.3. (ii) For each t and µ pair, the target distribution ρ t,µ is different, which can require t - and µ -specific tuning of hyper-parameters of the inference step, which is impractical and thus can lead to a deterioration of accuracy compared to our HOAM approach; see Figure 3-4.

Predicting statistics of chaotic and particle dynamics in high dimensions We now consider the nine-dimensional dynamical system introduced in [69], which leads to chaotic behavior. We show in Figure 3 (bottom) the sample histogram corresponding to a test physics parameter that represents the Rayleigh number. At time t = 3 . 7 and projecting onto

-

0

.

2

-

0

.

1

0

.

0

0

.

1

-

0

.

2

0

.

0

True

HOAM

two-stream

bump

10

1

10

2

10

3

10

4

runtime (s)

HOAM

CFM

NCSM

FOM

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 AM fails to converge. Consider now the example of the particles in an aharmonic trap, which leads to 100-dimensional samples X i t,µ . For a test physics parameter, Figure 5 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 t and physics parameter µ with our variational approach helps reduce inference runtime compared to standard diffusion- and flow-based modeling that condition on t and µ and therefore have to solve a separate inference problem for each time step and physics parameter at test time. Because we learn the dynamics over time t , it is critical to accurately capture the coupling over the time steps, for which we propose to use higher-order 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 errors as state-of-the-art methods while at the same time reducing inference runtime by 1-2 orders of magnitude. Additionally, HOAM provides speedups of up to 2 orders of magnitude to classical numerical full models.

Limitations : 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 a limitation in computational biology [23, 9]. Second, we currently seek a vector field that minimizes the kinetic energy or a variant thereof. Investigating other notions of energy that might lead to vector fields with other desired properties in certain problems remains a challenge.

We do not expect that this work has negative societal impacts.

References

M. S. Albergo, N. M. Boffi, and E. Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv , 2303.08797, 2023.

M. S. Albergo and E. Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations , 2023.

L. Ambrosio and W. Gangbo. Hamiltonian ODEs in the Wasserstein space of probability measures. Communications on Pure and Applied Mathematics , 61(1):18-53, Jan. 2008.

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.

L. Ambrosio, N. Gigli, and G. Savaré. Gradient Flows . Lectures in Mathematics ETH Zürich. Birkhäuser-Verlag, Basel, 2005.

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.

J. W. Banks and J. A. F. Hittinger. A new class of nonlinear finite-volume methods for Vlasov simulation. IEEE Transactions on Plasma Science , 38(9 PART 1):2198 - 2207, 2010.

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.

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.

P. Benner, S. Gugercin, and K. Willcox. A survey of projection-based model reduction methods for parametric dynamical systems. SIAM Review , 57(4):483-531, 2015.

P. Benner, M. Ohlberger, A. Cohen, and K. Willcox, editors. Model Reduction and Approximation: Theory and Algorithms . Society for Industrial and Applied Mathematics, Philadelphia, PA, July 2017.

P. Benner and M. Redmann. Model reduction for stochastic systems. Stochastic Partial Differential Equations: Analysis and Computations , 3(3):291-338, Sep 2015.

J. Berman and B. Peherstorfer. CoLoRA: Continuous low-rank adaptation for reduced implicit neural modeling of parameterized partial differential equations. In R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, and F. Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning , volume 235 of Proceedings of Machine Learning Research , pages 3565-3583. PMLR, 21-27 Jul 2024.

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.

N. M. Boffi and E. Vanden-Eijnden. Probability flow solution of the fokker-planck equation. Machine Learning: Science and Technology , 4(3):035012, jul 2023.

J. Bruna, B. Peherstorfer, and E. Vanden-Eijnden. Neural Galerkin schemes with active learning for high-dimensional evolution equations. Journal of Computational Physics , 496:112588, 2024.

C. Bunne, L. Papaxanthos, A. Krause, and M. Cuturi. Proximal optimal transport modeling of population dynamics. In G. Camps-Valls, F. J. R. Ruiz, and I. Valera, editors, Proceedings of The 25th International Conference on Artificial Intelligence and Statistics , volume 151 of Proceedings of Machine Learning Research , pages 6511-6528. PMLR, 28-30 Mar 2022.

S. R. Cachay, B. Zhao, H. James, and R. Yu. DYffusion: A dynamics-informed diffusion model for spatiotemporal forecasting. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

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. CesaBianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

Y. Chen, M. Goldstein, M. Hua, M. S. Albergo, N. M. Boffi, and E. Vanden-Eijnden. Probabilistic forecasting with stochastic interpolants and Föllmer processes. arXiv , 2403.13724, 2024.

Y. Chen and D. Xiu. Learning stochastic dynamical system via flow map operator. Journal of Computational Physics , 508:112984, 2024.

Y. Cheng, A. J. Christlieb, and X. Zhong. Energy-conserving discontinuous Galerkin methods for the Vlasov-Ampère system. Journal of Computational Physics , 256:630-655, 2014.

S. Chewi, J. Clancy, T. Le Gouic, P. Rigollet, G. Stepaniants, and A. Stromme. Fast and Smooth Interpolation on Wasserstein Space. In A. Banerjee and K. Fukumizu, 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.

S.-N. Chow, W. Li, and H. Zhou. Wasserstein Hamiltonian flows. Journal of Differential Equations , 268(3):1205-1219, Jan. 2020.

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.

A. Davtyan, S. Sameni, and P. Favaro. Efficient video prediction via sparsely conditioned flow matching. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) , pages 23263-23274, October 2023.

P. Deuflhard and A. Hohmann. Numerical Analysis in Modern Scientific Computing . Springer, 2003.

E. Dupont, A. Doucet, and Y. W. Teh. Augmented neural ODEs. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

A. Ern and J.-L. Guermond. Theory and Practice of Finite Elements , volume 159 of Applied Mathematical Sciences . Springer New York, New York, NY, 2004.

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 , 360(6392):eaar3131, June 2018.

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 Iberoamericana , 36(4):1071-1112, Jan. 2020.

R. G. Ghanem and P. D. Spanos. Stochastic Finite Elements: A Spectral Approach . Springer, 1991.

W. Harvey, S. Naderiparizi, V. Masrani, C. D. Weilbach, and F. Wood. Flexible diffusion modeling of long videos. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems , 2022.

T. Hashimoto, D. Gifford, and T. Jaakkola. Learning population-level diffusions with generative RNNs. 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.

J. Hittinger and J. Banks. Block-structured adaptive mesh refinement algorithms for Vlasov simulation. Journal of Computational Physics , 241:118-140, 2013.

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 Neural Information Processing Systems , volume 33, pages 6840-6851. Curran Associates, Inc., 2020.

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 Neural Information Processing Systems , volume 35, pages 8633-8646. Curran Associates, Inc., 2022.

B. Holzschuh, S. Vegetti, and N. Thuerey. Solving inverse physics problems with score matching. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

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 Representations , 2022.

T. J. R. Hughes. The Finite Element Method: Linear Static and Dynamic Finite Element Analysis . Dover Publications, 2012.

A. Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(24):695-709, 2005.

P. Kidger, J. Foster, X. Li, and T. J. Lyons. Neural SDEs as infinite-dimensional 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.

K. Kormann and A. Yurova. A generalized Fourier-Hermite method for the Vlasov-Poisson system. BIT Numerical Mathematics , 61(3):881-909, Sept. 2021.

T. Koshizuka and I. Sato. Neural Lagrangian Schrödinger Bridge: Diffusion Modeling for Population Dynamics. In The Eleventh International Conference on Learning Representations , 2023.

B. Kramer, B. Peherstorfer, and K. E. Willcox. Learning nonlinear reduced models from data with operator inference. Annual Review of Fluid Mechanics , 56(Volume 56, 2024):521-548, 2024.

J. N. Kutz, S. L. Brunton, B. W. Brunton, and J. L. Proctor. Dynamic mode decomposition: data-driven modeling of complex systems . SIAM, 2016.

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.

K. Lee and E. J. Parish. Parameterized neural ordinary differential equations: applications to computational physics problems. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences , 477(2253):20210162, 2021.

A. Lenard and I. B. Bernstein. Plasma Oscillations with Diffusion in Velocity Space. Physical Review , 112(5):1456-1459, Dec. 1958.

L. Li, S. Hurault, and J. Solomon. Self-consistent velocity matching of probability flows. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

X. Li, T.-K. L. Wong, R. T. Q. Chen, and D. K. Duvenaud. Scalable gradients and variational inference for stochastic differential equations. In C. Zhang, F. Ruiz, T. Bui, A. B. Dieng, and D. Liang, editors, Proceedings of The 2nd Symposium on Advances in Approximate Bayesian Inference , volume 118 of Proceedings of Machine Learning Research , pages 1-28. PMLR, 08 Dec 2020.

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.

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.

Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow Matching for Generative Modeling. In The Eleventh International Conference on Learning Representations , 2023.

J. Lott. Some Geometric Calculations on Wasserstein Space. Communications in Mathematical Physics , 277(2):423-437, Nov. 2007.

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.

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.book.....M.

I. Mezić. Spectral properties of dynamical systems, model reduction and decompositions. Nonlinear Dynamics , 41(1-3):309-325, 2005.

C. Mouhot and C. Villani. On Landau damping. Acta Mathematica , 207(1):29-201, 2011.

A. Muntean, J. Rademacher, and A. Zagaris, editors. Macroscopic and Large Scale Phenomena: Coarse Graining, Mean Field Limits and Ergodicity , volume 3 of Lecture Notes in Applied Mathematics and Mechanics . Springer International Publishing, Cham, 2016.

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. PMLR, July 2023.

F. Noé, A. Tkatchenko, K.-R. Müller, and C. Clementi. Machine learning for molecular simulation. Annual Review of Physical Chemistry , 71(Volume 71, 2020):361-390, 2020.

F. Otto. The geometry of dissipative evolution equations: the porous medium equation. Communications in Partial Differential Equations , 26(1-2):101-174, Jan. 2001.

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.

B. Peherstorfer, K. Willcox, and M. Gunzburger. Survey of multifidelity methods in uncertainty propagation, inference, and optimization. SIAM Review , 60(3):550-591, 2018.

E. Perez, F. Strub, H. de Vries, V. Dumoulin, and A. Courville. Film: visual reasoning with a general conditioning layer. In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence and Thirtieth Innovative Applications of Artificial Intelligence Conference and Eighth AAAI Symposium on Educational Advances in Artificial Intelligence , AAAI'18/IAAI'18/EAAI'18. AAAI Press, 2018.

S. Possanner, F. Holderied, Y. Li, B. K. Na, D. Bell, S. Hadjout, and Y. Güçlü. HighOrder Structure-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 in Computer Science.

K. Rasul, C. Seward, I. Schuster, and R. Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. 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 8857-8868. PMLR, 18-24 Jul 2021.

P. Reiterer, C. Lainscsek, F. Schürrer, C. Letellier, and J. Maquet. A nine-dimensional lorenz system to study high-dimensional chaos. Journal of Physics A: Mathematical and General , 31(34):7121, aug 1998.

C. Robert and G. Casella. Monte Carlo Statistical Methods . Springer, 2004.

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.

G. Rozza, D. Huynh, and A. Patera. Reduced basis approximation and a posteriori error estimation for affinely parametrized elliptic coercive partial differential equations. Archives of Computational Methods in Engineering , 15(3):1-47, 2007.

C. Salvi, M. Lemercier, and A. Gerasimovics. Neural stochastic pdes: Resolutioninvariant learning of continuous spatiotemporal dynamics. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 1333-1344. Curran Associates, Inc., 2022.

F. Santambrogio. Optimal Transport for Applied Mathematicians , volume 87 of Progress in Nonlinear Differential Equations and Their Applications . Springer International Publishing, Cham, 2015.

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-Transport Analysis of SingleCell Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell , 176(4):928-943.e22, Feb. 2019.

P. J. Schmid. Dynamic mode decomposition of numerical and experimental data. Journal of Fluid Mechanics , 656:5-28, 2010.

E. Schrödinger. Über die Umkehrung der Naturgesetze. Technical Report 1931 IX, Akademie der Wissenschaften, Berlin, 1931.

Z. Shen, Z. Wang, S. Kale, A. Ribeiro, A. Karbasi, and H. Hassani. Self-consistency of the fokker planck equation. In P.-L. Loh and M. Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 817-841. PMLR, 02-05 Jul 2022.

J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium 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.

Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

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 Aviv, Israel, July 22-25, 2019 , page 204, 2019.

Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Scorebased generative modeling through stochastic differential equations. In International Conference on Learning Representations , 2021.

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.

A. Tong, J. Huang, G. Wolf, D. Van Dijk, and S. Krishnaswamy. TrajectoryNet: A dynamic optimal transport network for modeling cellular dynamics. 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.

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):381-386, Apr. 2014.

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-421, 2014.

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(2252):20210167, Aug. 2021.

B. Tzen and M. Raginsky. Neural stochastic differential equations: Deep latent gaussian models in the diffusion limit. arXiv , 1905.09883, 2019.

C. Villani. Optimal transport: old and new . Number 338 in Grundlehren der mathematischen Wissenschaften. Springer, Berlin, 2009.

C. Villani. Topics in optimal transportation . Number 58 in Graduate studies in mathematics. American Mathematical Society, Providence, Rhode Island, reprinted with corrections edition, 2016.

P. Vincent. A connection between score matching and denoising autoencoders. Neural Computation , 23(7):1661-1674, 2011.

M. O. Williams, I. G. Kevrekidis, and C. W. Rowley. A data-driven approximation of the Koopman operator: Extending dynamic mode decomposition. Journal of Nonlinear Science , 25(6):1307-1346, 2015.

F. A. Wolf, F. K. Hamey, M. Plass, J. Solana, J. S. Dahlin, B. Göttgens, 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.

A Quadrature rules

A.1 Monte Carlo estimation

Monte Carlo estimation approximates an integral by evaluating the integrand at randomly sampled nodes within an interval [ a, b ],

We consider only the case where the nodes t i are uniformly distributed random variables in [ a, b ]. The weights are:

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 a function f over an interval [ a, b ] by dividing it into N subintervals of equal width h = b -a N . The approximation is given by:

where the nodes t i are:

The trapezoidal rule is a second-order rule, which means that the integration error decays as O ( h 2 ) for sufficiently smooth functions.

A.3 Composite Simpson's Rule

Composite Simpson's rule approximates the integral by fitting parabolas through intervals. It divides [ a, b ] into an even number N of subintervals of width h = b -a N . The approximation is:

with nodes:

Composite Simpson's rule is a fourth-order rule, which means that the integration error decays as O ( h 4 ) for sufficiently smooth functions.

A.4 Gauss-Legendre Quadrature

Gauss-Legendre quadrature approximates the integral over [ -1 , 1] by choosing nodes t i and weights w i so that polynomials of the highest possible degree are integrated exactly. A Gauss-Legendre quadrature has the form

where t i are the roots of the Legendre polynomial P n ( t ), and the weights are:

For integration over [ a, b ], a linear transformation maps [ -1 , 1] to [ a, b ]:

Gauss-Legendre quadrature exactly integrates polynomials of degree 2 n -1, where n is the number of nodes.

B Details about numerical examples

B.1 Harmonic oscillator with background collisions

The equation of motion in four-dimensional phase-space for X = [ X 1 , X 2 , V 1 , V 2 ] is given by:

Here, η = 5 × 10 -2 and ξ denotes white noise. The initial configuration is a Gaussian centered at m 0 = [1 , 1] with covariance equal to Σ 0 = 10 -2 × Id in the spatial coordinates X 1 , X 2 and Gaussian in the velocity coordinates V 1 , V 2 centered at zero and with covariance Σ 0 = 10 -2 × Id.

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 a single particle depend on the position of all other particles. Assuming N particles in the system, this means d d t X i t,µ = v ( t, X i t,µ ; µ, X 1 t,µ , . . . , X N t,µ ). 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 a PDE of the form ∂ t ρ t,µ + ∇· ( ρ t,µ v mf ( t, · ; µ, ρ t,µ )) = 0 that describes the evolution of the collection (or population, ensemble) of particles denoted by ρ t,µ . In the specific case of the Vlasov-Poisson problem, Coulomb interactions in the mean-field limit give rise to a Poisson equation determining an electric field that is generated by the collection of particles and influences its dynamics. For completeness sake, we 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 [60] 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.

Governing equation We slightly change the notation here to be consistent with the references. f : X x × R d × R ×D → R , d ∈ { 1 , 2 , 3 } , denotes the distribution function governed by the Vlasov-Poisson system

In the notation of the rest of this work, f ( · , · , t ; µ ) = ρ t,µ , X x × R d = X . The spatial domain X x is a subset of R d , in all our examples it is of the form [0 , l 1 ] × [0 , l 2 ] × [0 , l 3 ] with periodic boundary conditions.

The approximation becomes:

Two-stream instability In this case, d = 2, so the particle positions vary in X x = [0 , l 1 ] with periodic boundary conditions and their velocity evolves in R . For the two-stream instability, we set the initial distribution to

with α = 0 . 05 , l 1 = 50 , v 0 = 3. The parameter µ varies as µ train ∈ { 1 . 2 , 1 . 3 , . . . , 1 . 9 } and µ test ∈ { 1 . 25 , 1 . 85 } . We use a particle-in-cell method for generating the data based on the repository https://github.com/pmocz/pic-python . The number of marker particles is N = 25000 and for the sake of computing the electric field, a uniform grid of N/ 8 cells is used. Integration in time is done via a Störmer-Verlet splitting over t ∈ [0 , 40] with time-step size 10 -2 .

Bump-on tail We consider the initial distribution

with α = 0 . 05 , l 1 = 50 , v b = 4 , δ = 9 / 10 , σ 1 = 1 , σ 2 = 1 / √ 2. The parameter µ varies as µ train ∈ { 1 . 3 , 1 . 4 , . . . , 2 . 0 } and µ test ∈ { 1 . 35 , 1 . 95 } . The other parameters are the same as in the two-stream case.

Strong Landau damping In this case, d = 6 and

with l 1 = 4 π and l 2 = l 3 = 1. The data is generated using the Struphy package [67], the exact specifications of the simulation are available at https://gitlab.mpcdf.mpg.de/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 · ∇ x f . This implies µ ∈ { 0 . 5 , 0 . 6 , . . . , 1 . 5 } , where µ = 1 . 0 corresponds to the default settings. This µ = 1 . 0 is also the test parameter 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 a default option of the used code. For a single MPI process, it extends to 27 minutes while for 16, it can be reduced to 4 minutes.

The high-fidelity data we generate is using a control variate approach in order to 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 our reconstructed potential in comparison to the physical one, but we observe in practice that this is only by a multiplicative constant. We save 10 5 marker particles from the high-order simulations and use N = 25000 of them as input data for our method. We integrate in time over t ∈ [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 to 2 × 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 5 × 10 -2 . We integrate 25000 particles of the system up to T = 20 using the Euler-Maruyama scheme with time step size equal to 10 -2 . The parameter µ varies as µ train ∈ { 13 . 5 , 13 . 6 , . . . , 14 . 2 } and µ test ∈ { 13 . 65 , 14 . 05 } .

B.4 Particles in aharmonic trap

We consider the evolution of interacting particles in an aharmonic trap [16]. The twodimensional particle positions Z 1 ( t, µ ) , . . . , Z M ( t, µ ) are governed by an SDE

where γ > 0 is the diffusion coefficient and W i are independent Wiener processes. The function g ( t, Z ) = ( a ( t ) -Z ) 3 describes a time-dependent one-body force, where a ( t ) = 5 / 4(sin( πt )+3 / 2))+ µ cos(2 πt ) is the position of the trap. The function K ( Z, Z ′ ) = α M ( Z ′ -Z ) describes a pairwise interaction term. We set α = -1 / 4 and γ = 10 -2 . The parameter µ is in the range D = [0 . 3 , 0 . 9] and modifies the position of the trap. A sample X i t,µ corresponds to a vector [ Z 1 ( t, µ ) , . . . , Z M ( t, µ )] T of dimension 100, because we have M = 50 particles and each position Z j ( t, µ ) as 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 2.

B.5 Other details about numerical experiments

In terms of network architecture, we follow [13] 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 linear 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 × 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 iterations for Vlasov examples and for 25 , 000 Adam iterations for all other examples.

The results were computed on NVIDIA Quadro RTX 8000 GPUs. All code was implemented in Python using the JAX library with JIT complication where possible.

Hyper-parameter /epsilon1 in the loss (7) searched over { 0 , 1 , 2 , 5 , 7 } × 10 -2 for both HOAM and AM.

The relative error in the electric energy is computed as

where e true ( t ) is the electric energy predicted by the high-fidelity numerical simulations at time t and e predict ( t ) is the electric energy computed from samples of either HOAM (ours), AM, NCSM, or CFM. The relative error in the mean is

where the expected values are estimated via Monte Carlo from the generated samples.

The Sinkhorn distance is computed with https://ott-jax.readthedocs.io/en/latest/ with threshold 10 -3 ; see also [25].

B.6 Additional numerical results

In Figure 6 we show the various projections at time t = 3 . 7 of the sample distribution corresponding to the nine-dimensional chaotic system [69].

2,7

True

HOAM-S HOAM-G

AM

NCSM

CFM

3,9

4,1

4,6

CC

0

5

10

-

3

10

-

1

Electric Energy

HOAM-S

0

5

Time

HOAM-G

0

5

AM

0

5

NCSM

0

5

CFM

In Figure 7 we show the particle histograms and the electric energy curves for the sixdimensional Vlasov-Poisson problem corresponding to strong Landau damping.

In Figure 8, for the linear oscillator example, we compare CoLoRA to two other modulation schemes: FiLM [66] and MLP. For the MLP the inputs x, t, µ are concatenated together and input directly to the model. There is no hyper-network or modulation scheme. For FiLM, we closely follow the original paper. The main network takes x as input and the hyper-network t, µ as input. The hyper-network and main network have the same parameter counts as in the CoLoRA experiments. The output of the hyper-network then directly modulates the activation of each layer of the main network as detailed in the original FiLM paper [66]. Figure 8 shows that parameterizing the vector field s t,µ with CoLoRA layers achieves the lowest mean Wasserstein distance, which motivates the use of the CoLoRA modulation scheme [13].

C Calculations regarding the entropic loss

In the following, assume that ρ ∈ P ( X ) is a smooth density bounded away from zero. We begin by showing some calculation rules of the operator -∆ ρ : s ↦→ -∇ · ( ρ ∇ s ) with

CoLoRA

FiLM

MLP

modulation Scheme

4

×

10

-

4

5

×

10

-

4

6

×

10

-

4

7

×

10

-

4

homgeneous Neumann boundary conditions. In its weak form, it reads

With the choice f = log ρ , we find the useful identity ∆ ρ log ρ = ∆ ρ . Next, recall the objective E /epsilon1 from (5):

Now denote by δs an arbitrary element of C ∞ ( X ). Then, if s /epsilon1 is a minimizer of the (strictly convex) objective, it holds that

Hence,

Furthermore, note that (5) is identical to

after integration by parts. The last term is the Fisher information of the data at t, µ and plays no 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 t are coupled through the term ∂ t ρ t,µ . This is most apparent when one discretizes the equation in time. Denote by { t i } n t i =0 a strictly increasing sequence with t 0 = 0 , t n t = 1, and t i +1 -t i = δt i . Then, for fixed but arbitrary µ , we obtain n t coupled problems of the form

Adding these problems and shifting the indices, one can eliminate ρ t i +1 ,µ , explicitly coupling s t i ,µ and s t i +1 ,µ . The continuous equivalent of this of course is an integration over t , followed by an integration by parts.

E Geometric picture of the optimization problem

We omit the dependence on the parameter µ here for the sake of simpler notation and write d ρ for ρ d x for brevity. Note that the following considerations are purely formal. They are meant to illustrate a geometric picture of the optimization problems we consider. We claim no originality of these ideas; the exposition is based on Chapter 7 of [89] 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 ρ t ∈ P ( X ), the closure of {∇ f : f ∈ C ∞ ( X ) } in L 2 ( X , ρ t , R d ), see Definition 8.4.1 in [5]), one can view P ( X ) formally as a Riemannian manifold:

Definition 1 ([63]) . Let τ ↦→ ρ 1 τ and τ ↦→ ρ 2 τ be two curves valued in P ( X ) for τ ∈ ( t -/epsilon1, t + /epsilon1 ) such that ρ 1 τ ∣ ∣ τ = t = ρ 2 τ ∣ ∣ τ = t = ρ t . The optimal transport metric on T P ( X ) at ρ t ∈ P ( X ) is given by

This formalism is commonly named after the author of [63] and is closely linked to Arnold's considerations on geometric hydrodynamics [6] 2 As both the identification of s t from ∂ t ρ t and the metric depend on ρ t , the geometry defined on P ( X ) in this way is non-trivial.

Action of a curve The optimization Equation (3) has an appealing 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 1 2 ∫ 1 0 ∫ X |∇ s t | 2 d ρ t the action of the curve t ↦→ ρ t with tangent velocity ∇ s t . We want to stress that while this procedure is reminiscent of physical action principles, in the latter a solution corresponds to a stationary point given boundary conditions at the beginning and end of the curve. The problem we consider in Equation (6) is more narrow and concerned with finding ∇ s t that matches a given curve t ↦→ ρ t . Determining curves of minimal action in P ( X ), leads to the Benamou-Brenier formula ([4], Proposition 2.30)):

with W 2 the Wasserstein (or Kantorochiv-Rubinstein) distance.

Lagrangian functions The selection criterion based on kinetic energy alone is not without alternatives. In [24], the relation ∂ t ρ = -∆ ρ s is interpreted as a form of Legendre transform, hence s plays the role of a momentum and L ( ρ t , ∂ t ρ t , t ) = ∫ X |∇ ∆ † ρ ∂ t ρ | 2 d ρ that of a Lagrangian. Here, we introduced the notation ∆ † ρ to denote the pseudo inverse operator. Note that, formally, it is sensible to consider ∂ t ρ as an element of the tangent space of P ( X ). After all, ρ + τ∂ t ρ ∈ P ( X ) for ρ strictly positive and τ small enough. In this picture, s is an element of the cotangent space. The introduction of [63] addresses the two concepts and how they relate.

Any function L : ( ρ, ∂ t ρ, t ) ↦→ L ( ρ, ∂ t ρ, t ), strictly convex and superlinear in its second argument, can be chosen to define the minimization objective. 3 Details can be found in Chapter 7 of [89], 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 ( ρ, ∂ t ρ, t ) = 1 2 ∫ X |∇ ∆ † ρ ∂ t ρ | 2 d ρ -∫ X V dρ for a potential V : X → R can be used to model obstacles in the path of the samples.

There exist a number of partial differential equations whose solutions ρ t can be described as curves of stationary action with respect to such Lagrangians, described in [3, 24], as well as [89], Chapter 23, and [90], Chapter 8.

Schrödinger Bridge The objective defined in Equation (5) corresponds to the choice

The associated momentum s /epsilon1 therefore satisfies s /epsilon1 = δL /epsilon1 δ ( ∂ t ρ ) , hence -∆ ρ s /epsilon1 + /epsilon1 2 2 ∆ ρ = ∂ t ρ , a Fokker-Planck equation. Furthermore, the action of the curve t ↦→ ρ t is given by

This expression is known as the dual formulation of the Kantorovich-Schrödinger problem ([31], Theorem 36, except for the fact that the /epsilon1 therein corresponds to /epsilon1 2 / 2 here). While the classical optimal transport problem is concerned with the path connecting ρ 0 and ρ 1 minimizing the time integral of the kinetic energy (which coincides with the transport cost), the Schrödinger-Bridge problem is concerned with finding the most likely configuration at intermediate times, subject to the information that the configuration is given at times 0 and 1 and assuming that the particles X t undergo Brownian motion with diffusivity ε 2 / 2. Unless ρ 1 is the result of a convolution of ρ 0 with a Gaussian kernel of width ε , the evolution of the system towards ρ 1 is a rare event and the most likely solution is to be understood conditional on the observation of this event.

Rigorous results can be found in Section 5 of [31]. Another derivation of the loss function from Equation (5), 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, Ψ = -s .

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 claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

It is fine to include aspirational goals as motivation as long as 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 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 on how these assumptions might be violated in practice and what the implications would be.

The authors are encouraged to create a separate "Limitations" section in their paper.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

The authors should reflect on the factors that influence the performance of the approach. For example, a 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.

If 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 an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Guidelines:

The answer NA means that the paper does not include theoretical results.

All assumptions should be clearly stated or referenced in the statement of any theorems.

All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Theorems and Lemmas that the proof relies upon should be properly referenced.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

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 or 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 contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if 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 either make it 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 to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be

(c) If the contribution is a new model (e.g., a 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).

possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to 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 1 provides link to code.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

Please see the NeurIPS code and data submission guidelines ( https://nips. cc/public/guides/CodeSubmissionPolicy ) for more details.

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https://nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

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 Setting/Details

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]

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 full details can be provided either with the code, in appendix, or as supplemental material.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

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 2 shows replicates, results reported in Figure 3-6, Table 1 are based on thousands of samples.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The assumptions made should be given (e.g., Normally distributed errors).

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or 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]

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into 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.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Only data from numerical simulations are used. We do not expect that this work has negative societal impacts.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Section 4.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

The authors should consider possible harms that could arise when the technology is being used as 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 or unintentional) misuse of the technology.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

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 to usage guidelines or restrictions to access the model or implementing safety filters.

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]

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.

The authors should state which version of the asset is used and, if possible, include a URL.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

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.

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 1 and Appendix B.

Guidelines:

The answer NA means that the paper does not release new assets.

The paper should discuss whether and how consent was obtained from people whose asset is used.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: 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)?

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 research with human subjects.

According to the NeurIPS 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.

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.

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 (IRB) 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.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\n\nSTRENGTHS:\nDeveloping 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.\n\nWEAKNESSES:\nThe 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\n\nQUESTIONS:\nGeneral 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?\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nTo 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.\n\nQUESTIONS:\nI 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}$.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 9, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1)\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.\n\nWEAKNESSES:\n1) 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.\n\nQUESTIONS:\n1)\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.\n\nLIMITATIONS:\nYes, limitations of their work is adequately addressed. No potential negative societal impact of their work is identified.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}