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

TL;DR

This work uses a variational problem to infer parameter- and time-dependent gradient fields that represent approximations of the population dynamics and shows 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.

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.

Figures (8)

• Figure 1: Parametric model reduction with our HOAM seeks to learn vector fields that represent population dynamics $\rho_{t}$ over time $t$. In contrast, parametric model reduction with score-based diffusion denoising and flow-based modeling requires conditioning on time $t$, which leads to separate, costly inference steps for each time step of a sample trajectory.
• Figure 2: Left: During training the high-variance function $q(s)(t)$ needs to be numerically integrated for estimating the loss. Center left: The high variance leads to inaccurate estimates of the time integral by Monte Carlo, whereas higher-order numerical quadrature produces accurate estimates. Center right: Numerical quadrature in HOAM leads to stable estimates of the loss whereas Monte Carlo integration in AM leads to unstable behavior. Right: HOAM based on higher-order quadrature is stable and more accurate than AM.
• Figure 3: Histograms of solution fields. Top: Bump-on-tail ($t = 20$) instability. Middle top: two-stream ($t = 20$) instability. Middle bottom: Strong Landau damping ($t = 4$) instability. HOAM with Simpson's and Gauss quadrature accurately predicts the fine scale features and multi-modality of the population density in the Vlasov problems. AM does not converge on the 6 dimensional problem. Bottom: High-dimensional chaos PeterReiterer_1998 ($t = 3.7$, dim 3 vs dim 9). HOAM accurately predicts the low probability region that connects the two high probability regions while AM does not converge.
• Figure 4: Electric energy of bump-on-tail (top) and two-stream (bottom) instability. HOAM with Simpson's and Gauss quadrature accurately predicts the energy growth in the transient regime and oscillations at later times. The ground truth is displayed in blue.
• Figure 5: Left: HOAM accurately predicts the time evolution of the mean position of a 100-dimensional particle system in an aharmonic moving trap (dim 1 vs dim 100). Right: HOAM reduced models provide about 2 orders of magnitude speedup over traditional numerical (full) models for the 6 dimensional strong Landau problem. HOAM is also 1--2 orders of magnitude faster than CFM and NCSM, which provide no speedup over the full models in our problems.

Theorems & Definitions (2)

• Remark 1
• Definition 1: otto_geometry_2001