Source-linked AI summary
Model reduction in time: Learning flow maps for reducing stochastic dynamical systems in time
Shreya Jha, Benjamin Peherstorfer
TL;DR
Ensemble prediction can be costly because reduced models may still require many sequential time steps. FM-FTM learns a deterministic representation of stochastic probability-current dynamics and approximates its evolution with a flow map. Experiments report accurate ensemble quantities over long horizons with substantial speedups.
Problem
Reduced models can lower the cost per time step, but large numbers of sequential steps can still make ensemble prediction costly.
Source: §1 · p. 1, §2.2 · p. 4Method
FM-FTM uses FTM to learn deterministic probability-current dynamics, then learns a flow map that predicts their evolution over long horizons without sequential integration.
Source: §1 · pp. 1, 3Results
FM-FTM accurately captures time marginals, probability-current quantities, and physical ensemble statistics over horizons orders of magnitude larger than the underlying model's time-integration step.
Source: §5 · p. 19Takeaways & Limitations
The experiments show that accurate ensemble prediction can reduce time-integration steps through flow-map evaluations rather than relying only on state-dimension reduction or cheaper steps.
Source: §5 · p. 19The formulation assumes that the stochastic process has a probability density ρ(t).
Source: §2.1 · p. 4
Abstract
from arXiv · showhide
Predictions with dynamical systems can be computationally expensive---even with reduced models---because reaching future prediction times can require many sequential time integration steps. We propose Flow-Map First-Order Trajectory Matching (FM-FTM), a nonintrusive model reduction approach that accelerates ensemble predictions of stochastic dynamical systems by learning flow maps that replace multiple time-integration steps. Our method uses first-order trajectory matching (FTM) to learn a deterministic representation of the stochastic probability-current dynamics and then reduces its evolution over time into a flow map whose evaluation directly yields predictions over long time horizons. Numerical experiments on evolution problems with stochastic forcing demonstrate that FM-FTM accelerates ensemble predictions by up to two orders of magnitude while accurately predicting time marginals and ensemble statistics.
1 Introduction
FM-FTM reduces ensemble-prediction costs by learning flow maps that replace sequential integration with long-horizon evaluations. It builds on FTM to represent stochastic probability-current dynamics, then approximates their evolution; experiments report faster predictions with accurate statistics and quantities of interest.
1 Introduction
- Prior reduced models lower per-step costs, but large numbers of sequential steps can remain expensive for ensemble prediction.State-space reduction does not necessarily reduce the number of time steps.Source: §1 · p. 1
- Ensemble predictions support uncertainty quantification and probabilistic forecasting, but large ensembles can make accurate numerical simulation computationally costly.Reducing state dimension lowers cost per time step, but does not necessarily reduce the number of steps.Source: §1 · p. 1
- Deterministic flow-map methods reduce sequential evaluations, but do not directly address preserving stochastic probability distributions and statistical quantities of interest.Generative-model flow maps reduce integration in artificial sampling time, whereas this work targets physical time [71] [48] [8] [31] [9].Source: §1 · p. 2
- Unlike approaches that repeatedly apply learned stochastic maps, FM-FTM uses one neural-network evaluation per long-horizon prediction and transports ensembles between arbitrary physical times without modeling the full transition law.Other approaches either repeat learned stochastic transitions [19] [50] or evaluate time-marginal dynamics per time step [55] [6] [67].Source: §1 · pp. 2–3
- The method first uses FTM to learn a deterministic velocity field representing probability-current dynamics, then learns a flow map to predict states over far horizons without sequential integration.This two-stage design separates the representation of stochastic dynamics from reducing computational cost in time.Source: §1 · pp. 1, 3
- FM-FTM accelerates ensemble predictions by at least one order of magnitude while accurately predicting ensemble statistics and probability-current-dependent quantities of interest.The approach requires trajectory data but not the governing equations or discretized operators.Source: §1 · p. 4
2 Preliminaries and problem formulation
The section formulates ensemble prediction from trajectory data and targets accurate forecasts over horizons much larger than the integration step, without computing every intermediate step. It states the stochastic-process setting and the associated data and cost assumptions.
2.1 Trajectory data and stochastic dynamics
- The problem uses N observed trajectories at K + 1 equally spaced times over [0, T] as training data.The process is continuous-time and has state space R^d.Source: §2.1 · p. 4
- The formulation assumes the stochastic process has a probability density ρ(t), whose distribution at a single time is called the time marginal.The initial density is ρ(0).Source: §2.1 · p. 4
- The considered stochastic dynamics follow an SDE with smooth drift b and diffusion coefficient A, but learning uses trajectory data without requiring either coefficient.The drift specifies the deterministic contribution, while A determines how random forcing enters the state.Source: §2.1 · p. 4
2.2 Problem formulation
- The objective is to generate ensembles at selected prediction times for estimating quantities of interest and statistics of the distributions ρ(t)..Source: §2.2 · p. 4
- Fixed-step rollout to time T requires at least T/h sequential steps, even when predictions are requested only at intervals τ much larger than h.A reduced model or surrogate can lower per-step cost but typically does not reduce the step count.Source: §2.2 · p. 4
- The target is to advance ensembles over horizons τ ≫ h while retaining accurate statistical estimates, reducing sequential steps rather than only per-step cost..Source: §2.2 · p. 4
3 Flow maps for first-order trajectory matching
FM-FTM first learns deterministic probability-current dynamics from stochastic trajectory data, then learns a flow map to predict their evolution across long time horizons. The map retains the dynamics’ marginal and probability-current properties up to approximation errors while avoiding intermediate integration steps.
3 Flow maps for first-order trajectory matching
- FM-FTM separates stochastic-dynamics representation from time-cost reduction: FTM learns a deterministic model, and a learned flow map approximates its evolution.The construction uses trajectory data to learn the probability-current velocity and then learns the corresponding flow map.Source: §3 · p. 5
3.1 A deterministic model of ensemble evolution with first-order trajectory matching
- The probability-current velocity captures both time marginals and the direction and rate of average local motion, which marginals alone do not determine.The paper illustrates the distinction with clockwise and counterclockwise rotations that preserve the same rotationally symmetric distribution.Source: §3.1 · p. 5
3.2 Flow maps
- Applying the probability-current flow map to a state distributed as ρ(t) produces the distribution ρ(s) at the destination time.The map accumulates velocity along the evolving trajectory from the initial state at t to the state at s.Source: §3.2 · p. 6
3.3 Learning the flow map
- The learned flow map is trained to approximate the integrated learned velocity without numerically integrating that velocity over the jump horizon at evaluation.Its objective averages prediction error over starting states and time pairs, so one map is trained across starting times and jump horizons.Source: §3.3 · p. 6
3.4 Ensemble prediction with the learned flow map
- For requested prediction times, ensemble rollouts apply the map sequentially, requiring M evaluations rather than approximately T/h FTM integration steps.Each predicted state becomes the initial state for the next jump, and horizons can vary between requested times.Source: §3.4 · pp. 6–7
3.5 Quantities of interest
- FM-FTM aims to retain time-marginal and probability-current quantities of interest, including fluxes and circulations, subject to FTM and flow-map approximation errors.The paper distinguishes ensemble averages depending only on time marginals from current-like observables of local probability transport.Source: §3.5 · p. 7
4 Numerical experiments
Across four stochastic evolution experiments, FM-FTM preserves time marginals and statistics over tested jump horizons while gaining runtime speedups. It avoids integrating the FTM velocity through intermediate steps.
4 Numerical experiments
- FM-FTM removes the need to integrate the FTM velocity through each intermediate time step, enabling predictions over large jump horizons.Source: §4 · p. 7
- FM-FTM preserves time marginals and statistics across tested jump horizons while gaining runtime speedups.The experiments cover the Duffing oscillator, Rayleigh-Bénard convection, stochastic Burgers, and stochastic Navier-Stokes.Source: §4 · p. 7
- The experiments evaluate FM-FTM on the Duffing oscillator, Rayleigh-Bénard convection, stochastic Burgers, and stochastic Navier-Stokes.Source: §4 · p. 7
4.1 Training
Training first learns the FTM velocity field, then fits a flow map to predict states across sampled time intervals. This lets FM-FTM replace multiple sequential integration steps with a direct long-horizon prediction.
4.1 Training
- FM-FTM replaces multiple FTM integration steps with a learned flow map that predicts the state at a later time.The map approximates evolution under the FTM velocity field.Source: §4.1 · pp. 7, 9
- Training pairs combine a trajectory state at time t with a later sampled time s, and the FTM velocity is integrated from t to s to produce the target state.Source: §4.1 · p. 10
- The flow map is trained in minibatches by minimizing empirical mean-squared error between its prediction and the target state.Training times are sampled from the trajectory time grid, with later times chosen from available grid points.Source: §4.1 · p. 10
4.2 Evaluation Metrics
Evaluation measures both agreement in time marginals and trajectory-dependent probability-current quantities, since matching marginals alone does not establish agreement in probability transport.
4.2 Evaluation Metrics
- Time-marginal accuracy is evaluated using the mean sliced Wasserstein-2 distance across rollout times.Source: §4.2 · p. 10
- Trajectory-dependent accuracy is assessed through time-local probability-current quantities and their time-weighted average over rollout intervals.The FM-FTM estimate uses a midpoint rule and an additional short flow-map evaluation of length δQoI = 0.01.Source: §4.2 · pp. 10–11
- Marginal agreement alone is insufficient because it does not determine how probability mass moves between states.Source: §4.2 · p. 10
4.3 Duffing oscillator with stochastic forcing
The Duffing experiment tests FM-FTM on a stochastically forced oscillator using trajectory data and comparisons with FTM. FM-FTM preserves marginal and probability-current accuracy across tested horizons, including jumps up to 400 times the FTM step size.
4.3.1 Duffing oscillator: Setup, training data, and evaluation metrics
- The oscillator uses 5,000 training trajectories over [0, 12] with reference step size 0.01, and evaluates 5,000 test initial conditions.The stochastic system is integrated using Euler-Maruyama; tested jump horizons are τ ∈ {0.5, 1.0, 2.0, 4.0}.Source: §4.3.1 · p. 11
4.3.2 Duffing oscillator: Results
- FM-FTM captures the main distributional structure at jump horizons up to 400 times the FTM step size, while coarse-step FTM loses accuracy at h = 0.25.The comparison includes predictions at t = 4, 8, and 12.Source: §4.3.2 · pp. 11–12
- FM-FTM achieves time-marginal errors comparable to FTM at h = 0.01, with little sensitivity across τ ∈ {0.5, 1.0, 2.0, 4.0}.Source: §4.3.2 · p. 12
- FM-FTM accurately predicts the probability-current quantity across tested horizons, retaining information about local probability transport beyond the marginals.Source: §4.3.2 · p. 12
4.4 Rayleigh-Bénard convection with stochastic forcing
The Rayleigh-Bénard experiment evaluates FM-FTM on a nine-dimensional stochastically forced system with chaotic deterministic dynamics. FM-FTM preserves multimodal distributions and probability-current quantities at large jump horizons, while marginal errors remain moderate.
4.4.1 Rayleigh-Bénard: Setup, training data, and evaluation metrics
- The nine-dimensional system uses 20,000 training trajectories over [0, 20] with reference step size 0.01, and 20,000 test initial conditions.The setting uses µ = 13.65 and σ = 0.05; tested jump horizons range from 0.5 to 4.0.Source: §4.4.1 · p. 12
4.4.2 Rayleigh-Bénard: Results
- At τ = 4, FM-FTM captures both modes of the Rayleigh-Bénard distribution, whereas coarse-step FTM fails to capture its distributional structure.The jump horizon is 400 times the FTM integration step size.Source: Figure 5 · p. 13, §4.4.2 · p. 14
- FM-FTM maintains moderate sliced Wasserstein-2 errors through τ = 4, with no systematic deterioration as the horizon increases, though errors exceed FTM at h = 0.01.Source: §4.4.2 · p. 14, Figure 6 · p. 14
- FM-FTM accurately predicts the probability-current quantity across tested horizons, retaining information about rotational probability transport beyond time marginals.Source: §4.4.2 · p. 14
4.5 Stochastically forced Burgers equation
The Burgers experiment tests whether FM-FTM predicts physical ensemble statistics across large jump horizons. It reports accurate energy and enstrophy predictions while replacing multiple integration steps with flow-map evaluations.
4.5.1 Burgers: Setup and training data
- The experiment evaluates a stochastically forced Burgers equation with periodic boundaries on [0, 1), viscosity ν = 0.007, and a 64-point spatial grid.Reference trajectories use Euler–Maruyama with step size 0.0005 over [0, 4].Source: §4.5 · p. 14, §4.5.1 · pp. 14–15
- Training data comprises 4,096 trajectories with 800 snapshots each, downsampled from 8,000 time steps to a snapshot spacing of 0.005.The initial conditions are Gaussian bumps with noise structure matching the stochastic forcing.Source: §4.5.1 · p. 15
4.5.2 Burgers: Evaluation metrics.
- The evaluation averages energy and enstrophy estimates over 4,096 test trajectories and compares their relative errors over time.The tested jump horizons are τ ∈ {0.25, 0.5, 1.0}.Source: §4.5.2 · pp. 15–16
4.5.3 Burgers: Results
- Energy relative errors remain below 2% across τ ∈ {0.25, 0.5, 1.0}; enstrophy is also captured, though its errors are generally larger and exceed FTM’s at later times.Figure 7 compares relative errors for energy and enstrophy across the tested horizons.Source: Figure 7 · p. 16, §4.5.3 · p. 16
4.6 Forced turbulence
The forced-turbulence experiment evaluates FM-FTM on stochastic two-dimensional Navier–Stokes dynamics. It reports accurate ensemble statistics and energy spectra alongside substantial reductions in trajectory-generation time.
4.6.1 Forced turbulence: Setup and training data
- The experiment uses stochastic forcing in two-dimensional incompressible Navier–Stokes dynamics and trains on 2,048 trajectories over [0, 25], subsampled to 1,000 snapshots and a 64 × 64 grid.Reference dynamics use a 256 × 256 grid; training trajectories use reference step size 0.001.Source: §4.6 · p. 16, §4.6.1 · pp. 16–17
4.6.3 Forced turbulence: Results
- Approximately 26× and 49× faster than FTM at τ = 0.625 and τ = 1.25, respectively; at τ = 1.25, FM-FTM is more than two orders of magnitude faster than the physics model.Figure 9 presents the trajectory-generation timing comparison.Source: Figure 9 · p. 18, §4.6.3 · p. 19
4.6.2 Turbulence: Evaluation metrics
- Evaluation uses 1,024 test initial conditions and compares relative errors in ensemble-mean kinetic energy and enstrophy, kinetic-energy spectra, and trajectory-generation wall-clock time.The physics model, FTM, and FM-FTM are compared for computational efficiency.Source: §4.6.2 · p. 18
4.6.3 Forced turbulence: Results
- At τ = 1.25, FM-FTM requires 20 flow-map evaluations over [0, 25], versus 1,000 FTM integration steps; ensemble-mean kinetic energy and enstrophy errors are similar in magnitude over much of the interval.The tested horizons are τ ∈ {0.625, 1.25}.Source: Figure 8 · p. 17, §4.6.3 · p. 18
- With τ = 1.25, the predicted kinetic-energy spectrum at t = 25 closely matches the reference across plotted wavenumbers, capturing energy distribution across spatial scales.The experiment assesses spectra alongside ensemble-mean kinetic energy and enstrophy.Source: §4.6.3 · pp. 18–19, §4.6.2 · p. 18
5 Conclusion
The conclusion argues that reducing dynamics in time can accelerate stochastic ensemble predictions. FM-FTM replaces sequential integration with learned flow-map evaluations while retaining accurate predictions over long horizons.
5 Conclusion
- FM-FTM reduces stochastic ensemble-prediction cost by reducing dynamics in time, complementing approaches that reduce state dimension or make individual steps cheaper.The method approximates the evolution of a learned probability-current velocity with a flow map.Source: §5 · p. 19
- The experiments report accurate time marginals, probability-current quantities, and physical ensemble statistics over horizons orders of magnitude larger than the underlying integration step.Source: §5 · p. 19
- When predictions are needed only at selected times, the integration time resolution need not determine their computational cost.The conclusion presents flow-map approximations as complementary to state-space model reduction.Source: §5 · p. 19