Parametric model reduction of mean-field and stochastic systems via higher-order action matching

Authors: jules  berman
tobias  blickhan
benjamin  peherstorfer
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 21:30:56

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.

Main paper: 
 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] .Given a data set of samples, i.e., realizations of the random variable X t,µ on a suitable domain X ⊆ R d , {X i tj ,µ k | i = 1, . . . , N x , j = 1, . . . , N t , k = 1, . . . , N µ } ⊂ X , (1) 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 µ.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.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.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] 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.• • • • • ρ 0 ρ tn ρ tn+1 ρ 1 score-based modeling HOAM • • • • • • • • • • • • • • • • • • • • • flow -based modeli ngThere 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.(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] . (b) Introducing higher-order quadrature schemes for the loss to efficiently couple the gradient fields over time. This leads to lower variance estimators of the loss that critically stabilize training. (c) Demonstrating 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 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. 
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 equationEQUATION) 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.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,µ dx = 0, leads to parametric elliptic problems inEQUATIONwith 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:EQUATIONfor each t ∈ [0, 1] and µ ∈ D. The space H 1 (ρ t,µ , X ) contains functions s with X |∇s| 2 ρ t,µ dx < ∞, 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,µ dx, 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] ,EQUATIONwhich depends on ϵ ≥ 0. The energy E ϵ t,µ is of particular interest for two reasons: One, the Euler-Lagrange equation of (5) 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] ,EQUATIONNotice 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 ϵ = 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.The energy E t,µ defined in (4) as well as the energy E ϵ 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 atE ϵ (s)= D 1 0 X 1 2 |∇s t,µ | 2 +∂ t s t,µ + ϵ 2 2 ∆s t,µ ρ t,µ dxdt- X s t,µ ρ t,µ dx t=1 t=0 dν(µ) . (7)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.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 ϵ = 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) .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 formC(x) = W x + ϕ(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 ≪ 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 Gauss-Legendre 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 Ênx∼ρ [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 byEQUATION)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.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 Xt,µ = ∇s t,µ ( Xt,µ )dt + ϵdW t , where W t are Wiener processes and ϵ is the same ϵ 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.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. 
numerical experiments: 
 Examples We consider the following parametric dynamical systems; details in Appendix B. The high variance leads to inaccurate estimates of the time integral by Monte Carlo, whereas higher-order numerical quadrature produces accurate estimates. Center right: Numerical quadrature in HOAM leads to stable estimates of the loss whereas Monte Carlo integration in AM leads to unstable behavior. Right: HOAM based on higher-order quadrature is stable and more accurate than AM.subject to a quadratic potentialV (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.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.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.EQUATION)We consider only the case where the nodes t i are uniformly distributed random variables in [a, b] . The weights are:EQUATION)The root mean-squared integration error of Monte Carlo estimation decays as O(N -1/2 ), for integrands with bounded variance.The trapezoidal rule approximates the integral of a function f over an interval [a, b] bydividing it into N subintervals of equal width h = b -a N .The approximation is given by:EQUATION)where the nodes t i are:EQUATION)The trapezoidal rule is a second-order rule, which means that the integration error decays as O(h 2 ) for sufficiently smooth functions.Composite Simpson's rule approximates the integral by fitting parabolas through intervals. ItEQUATION)with nodes:EQUATION)Composite Simpson's rule is a fourth-order rule, which means that the integration error decays as O(h 4 ) for sufficiently smooth functions.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 formEQUATION)where t i are the roots of the Legendre polynomial P n (t), and the weights are:EQUATION)For integration over [a, b], a linear transformation maps [-1, 1] to [a, b]:EQUATION)The approximation becomes:EQUATION)Gauss-Legendre quadrature exactly integrates polynomials of degree 2n -1, where n is the number of nodes.The equation of motion in four-dimensional phase-space forEQUATION)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.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 meansd dt 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.We slightly change the notation here to be consistent with the references. f : X x × R d × R × D → R, d ∈ {1, 2,EQUATION)EQUATION)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.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 toEQUATIONwith α = 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 .We consider the initial distributionEQUATION)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 andEQUATIONwith 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]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}.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 SDEdZ i = g(t, Z i )dt + M j=1 K(Z i , Z j )dt + 2γdW i , i = 1, . . . , M ,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.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 ϵ 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 asEQUATION)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 isEQUATION)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] .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] . Figure 6 : Shows the projections of other dimensions of the nine-dimensional chaotic system [69] ; see also Figure 3 . 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] .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 [13] versus FiLM [66] and MLP. CoLoRA layers achieve the lowest mean Wasserstein distance compared to FiLM and MLP. In particular, CoLoRA avoids outliers with larger errors.homgeneous Neumann boundary conditions. In its weak form, it readsEQUATION)With the choice f = log ρ, we find the useful identity ∆ ρ log ρ = ∆ρ. Next, recall the objective E ϵ from ( 5):E ε (s) = 1 2 X ∇ s - ε 2 2 log ρ 2 ρdx - X ∂ t ρsdx.Now denote by δs an arbitrary element of C ∞ (X ). Then, if s ϵ is a minimizer of the (strictly convex) objective, it holds thatEQUATIONHence,EQUATION)Furthermore, note that ( 5) is identical toEQUATIONafter integration by parts. The last term is the Fisher information of the data at t, µ and plays no role in the optimization.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 } nt i=0 a strictly increasing sequence with t 0 = 0, t nt = 1, and t i+1 -t i = δt i . Then, for fixed but arbitrary µ, we obtain n t coupled problems of the form minEQUATIONAdding these problems and shifting the indices, one can eliminate ρ ti+1,µ , explicitly coupling s ti,µ and s ti+1,µ . The continuous equivalent of this of course is an integration over t, followed by an integration by parts.We omit the dependence on the parameter µ here for the sake of simpler notation and write dρ for ρ dx 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] .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 L2 (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-ϵ, t+ϵ) such that ρ 1 τ τ =t = ρ 2 τ τ =t = ρ t .The optimal transport metric on T P(X ) at ρ t ∈ P (X ) is given byEQUATION)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.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)):EQUATIONwith W 2 the Wasserstein (or Kantorochiv-Rubinstein) distance.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.The objective defined in Equation ( 5) corresponds to the choiceEQUATIONThe associated momentum s ϵ therefore satisfies s ϵ = δL ϵ δ(∂tρ) , hence -∆ ρ s ϵ + ϵ 2 2 ∆ρ = ∂ t ρ, a Fokker-Planck equation. Furthermore, the action of the curve t → ρ t is given byEQUATION)This expression is known as the dual formulation of the Kantorovich-Schrödinger problem ( [31] , Theorem 36, except for the fact that the ϵ therein corresponds to ϵ 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 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.• 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 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 method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • 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. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.The derivation of fluid dynamics from variational principles is, of course, much older and goes back as far as Langrange's Mécanique analytique published in 1789.The variables ρ and ∂t here denote any probability density and a scalar field on X . 
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.In Figure 2 we consider the harmonic oscillator example. We learn the field s t and plot q(s)(t) = Ênµ µ∼ν Ênx x∼ρt,µ [ 1 2 |∇s t,µ | 2 +True HOAM-S HOAM-G AM NCSM CFM ∂ t s t,µ + ϵ 2 2 ∆s tn,µ ] 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.Poisson problems describe the interaction of charged particles with dynamics that depend on all other particles, which leads to mean-field dynamics for large numbers of particles N x . Thus, reduced modeling with HOAM is well suited for this problem because the natural 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.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 .We now consider the nine-dimensional dynamical system introduced in [69] , which leads to chaotic behavior. We show in Figure 3 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. 
 
 Appendix:  


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

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

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

Reviews: review: paper_summary: paper_summary: This paper develops models of population dynamics in physical systems that exhibit stochastic and mean-field effects, influenced by physics parameters. The goal is to create models that can efficiently predict system behavior as alternatives to classical numerical methods. By utilizing the Benamou-Brenier formula from optimal transport and action matching, the approach involves solving a variational problem to infer gradient fields that approximate population dynamics. These gradient fields enable the generation of sample trajectories that mimic physical system dynamics under various physics parameters. The study highlights the importance of combining Monte Carlo sampling with higher-order quadrature rules for accurate estimation and stable training. The models are demonstrated to perform well on Vlasov-Poisson instabilities and high-dimensional particle and chaotic systems, outperforming state-of-the-art diffusion-based and flow-based models that rely solely on time and physics parameters.
strength_weakness: strength_weakness: 1) The model assumes access to a dense set of time points for the Gauss-Legendre quadrature, which may not be applicable when only a few time samples are available. This was already mentioned as a limitation of the current work. As this forms the main part of the approach, the practical benefit would be limited. 
2) 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.
questions: questions: 1)	Authors 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.
2)	Parametrizing 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?
3)	Solving eq (7) from data is challenging and prone to potential numerical issues. As a remedy, a combination of higher order numerical quadrature and MC sampling strategy is proposed. However, the details of this approach would be better to provide with a stability analysis on the training as mentioned in the text.
limitations: limitations: Yes, limitations of their work is adequately addressed. No potential negative societal impact of their work is identified.
paper_summary: paper_summary: The authors focus on learning models for population dynamics of parameterized physical systems that exhibit stochastic and mean-field effects in time. To do so, the authors leverage the Benamou-Brenier formula to learn gradient fields that transport the probability density as time evolves, and that enable the generation of sample trajectories reflecting the dynamics of the population. Numerical experiments show compeling results and state-of-the-art performance in high-dimensional particle systems and in chaotic systems.
strength_weakness: strength_weakness: To this reviewer the paper does not have any evident weakness. However, when introducing a high-order quadrature in time, which improves the performance, the authors do not elaborate on why one would, or would not, expect further improvements if similar quadrature rules were used in the spatial variable or in the physical parameters.
questions: questions: I have the following minor questions:
- In (6) it seems that the integrand should be evaluated on $s_t$ instead of $s$ for consistency with (4). 
- In (7) there seems to be a parenthesis missing that factors out the density $\rho_{t,\mu}$.
limitations: limitations: The authors adequately address the limitations of their work in Section 4. These are mostly technical. For instance, they require a high sampling rate in time, which is not always available in practice. This limits the applicability of the method, but does not detract from their main contribution.
paper_summary: paper_summary: The authors present a framework for learning a time-dependent gradient vector field (parametric model) that describes how the particle evolves in the state space. The idea is based on the action-matching framework (https://arxiv.org/abs/2210.06662), which learns a time-dependent gradient field that interpolates between marginals at different times in a simulation-free manner. The authors propose using high-order quadrature rules for evaluating nested integrals in the objective. They demonstrate the forecasting ability of their framework on diverse examples in high dimensions and show that it outperforms other diffusion-based and flow-based models in inference runtime and accuracy
strength_weakness: strength_weakness: The objective used by the authors is based on the action-matching paper, now adapted to include the dependence on the parameters (\mu). From a methodological standpoint, the authors' only important contribution is using a higher-order quadrature scheme to discretize the nested integrals
questions: questions: General comments:

i) 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? 

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

iii) 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? 

Specific comments

i) In Figure 5, a color bar showing the scales would help with the interpretation.

ii) In Figure 4, a label to show what the two curves (blue and orange) are is needed. What is the ground truth and what is the prediction?
limitations: limitations: -

Comments: comment: Thanks for the clarification.
Yes this is a typo. We thank the reviewer for catching this. We apologize for the confusion and will update the pdf accordingly. 

In all cases HOAM does outperform AM.
I thank the authors for their responses. I have the following additional question. In the pdf file you attached, Fig. 3 shows the relative error in mean and the caption states that "we see that HOAM remains relative stable, while AM increases in error." However, HOAM is the _blue_ line which not only is above the orange line (representing AM) but also seems to increase faster as $T$ increases. Is this a typo?
Thank you for the comments. If there is any other information we can provide, please let us know.
> I believe all of these will require additional rewrite in multiple places.

We thank the reviewer for raising these concerns. We are confident that all of the concerns are addressed in the paper, which we concisely summarize in this response. The reviewer’s comments will be helpful for guiding the final revision of this paper (if accepted) to make our points even clearer. 



> it is yet unclear to me how the proposed model improves the results and helps the stabilization of the training as claimed. 

We summarize here the reasoning for why "the proposed model improves the results and helps the stabilization of the training", which is also provided in the paper but we will make this clearer in a revision, if the paper gets accepted: 

1. **Gauss-Legendre quadrature results in a more accurate estimate of the time integral when compared to Monte Carlo.** This is because it is a higher-order numerical scheme (more precisely, it allows integrating higher-degree polynomials exactly) that leads to lower quadrature errors than Monte Carlo. Besides this theoretic argument, the difference in the quadrature error can be seen empirically in our experiments in, e.g., Figure 2. (See Section 2.3.)

2. **A more accurate estimate of the time integral results in a more stable and accurate estimate of the loss.** This is supported numerically by Figure 2 (left) which shows extremely high variability in the estimates of the loss for AM and low variance estimates of loss for our proposed HOAM. This also agrees with standard results from statistical learning theory where the deviation of the empirical risk (the estimate of the loss function) from the true risk (the true loss function) directly enters in bounds of generalization errors. 

3. **A more accurate estimate of the loss makes solving the optimization problem tractable.** This is supported by the numerical results in the paper, explicitly in Figure 2 (middle) and shown to be essential in global response,

All this is also described in detail in the main text of our paper in Section 2.3 and numerically supported by Figure 2 and the global response. For example see line 200: "Plain uniform sampling over the data set for mini-batching can lead to poor estimates of the loss." 



> If I understand correctly, the HOAM combines higher order numerical quadrature with MC sampling to solve (7) in the direction of time by employing a Gauss-Legendre quadrature. How can this be achieved numerically? 

1. As referenced in the main text of the paper the empirical loss is given in the Appendix in section A. This is written down as a discrete summation which can be easily implemented numerically. 
2. For the Gauss-Legendre quadrature, the key algorithmic step is determining the roots of the Legendre polynomials, which is implemented (for example) in scipy (scipy.special.roots_legendre()). We also provide the code of our implementation, but the link is redacted for now per submission policy.



> It would be also useful to motivate the choice of the dynamical systems in the numerical experiments? Why the examples are relevant in this context is unclear. 

1. From a surrogate modeling perspective, the systems that we consider are exceedingly challenging because they are high dimensional, chaotic and/or stochastic. Surrogate modeling for such systems is in its infancy because the point-wise approximations of traditional surrogate modeling techniques are meaningless. We reference these attempts in detail in the introduction and literature review.
2. We make careful efforts to choose problem setups which are experimentally meaningful to practitioners (see for example Appendix B.2 that describes that the data are obtained from code that are used by plasma physicists). The physical relevance of our problem setups are detailed in the numerous sources we cite, see [7, 21, 34, 42, 87 Sec 2(b)(i), 58].

  

> A clear framing of the problem setup and all the assumptions on the model would help the reader to better understand the results. 

1. The problem setup is described in the Introduction section on page 1: “Given a data set of samples [...] we aim to learn a dynamical-system reduced model to rapidly predict samples that approximately follow the same law [...]”
2. When we make mathematical statements, we provide assumptions. For example, for the statement of ‘uniqueness’ on page 4 we assume that the density is positive and that the source term integrates to zero. 
3. We provide an extensive appendix C-E on the connection to optimal transport and additional literature.

We hope that we have been able to demonstrate that most of these concerns are addressed in the main text of the paper. If the reviewer does not have any additional concerns, we would greatly appreciate it if they would consider revising their score.
Thank you for the rebuttal and the additional results. After reading your responses, I still have same concerns regarding the contribution of the paper.  My main issue is that while the authors claim that higher-order action matching (HOAM) is better than the baseline action matching (AM) on some parametric dynamical systems, it is yet unclear to me  how the  proposed model improves the results and help the stabilization of the training as claimed.  A clear framing of the problem setup and all the assumptions on the model would help the reader to better understand the results. I also find the implementation a bit unclear. If I understand correctly, the HOAM combines higher order numerical quadrature with MC sampling to solve (7) in the direction of time by employing a Gauss-Legendre quadrature. How can this be achieved numerically? It would be also useful to motivate the choice of the dynamical systems in the numerical experiments? Why the examples are relevant in this context is unclear. Hence, I believe all of these will require additional rewrite in multiple places. Therefore, I will keep my initial rating for now.
I thank the authors for their detailed responses and clarification. I will retain my rating.