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.
Figures & tables
Figure 1: While FTM quickly loses accuracy at coarse time steps, FM-FTM maintains accurate predictions by taking far-horizon time jumps, replacing many sequential integration steps with a single flow-map evaluation and thereby accelerating ensemble prediction.
Figure 2: Duffing oscillator: FM-FTM accurately captures the evolving distributions over large time intervals, whereas FTM loses accuracy even with moderately increased time steps from the original h=0.01 to h=0.25 .
Figure 3: Duffing oscillator: The Wasserstein-2 errors of FM-FTM remain comparable to those of FTM and show little sensitivity to the jump horizon τ , demonstrating that FM-FTM preserves the accuracy of ensemble predictions while reducing the number of time-integration steps by a factor of up to 400× in this example.
Figure 4: The relative errors in predicting quantities of interest remain small across increasing jump horizons τ for both the Duffing oscillator (left) and Rayleigh-Bénard convection (right), demonstrating that FM-FTM preserves trajectory-dependent quantities beyond time-marginal statistics.
Figure 5: Rayleigh-Bénard: FTM fails to capture the multimodal distribution even already with moderately increased time-step sizes, whereas FM-FTM accurately captures the evolution of both modes at a jump horizon of τ=4 , which is 400 times larger than the time-step size used by FTM.
Figure 6: Rayleigh-Bénard: The Wasserstein-2 errors remain moderate across jump horizons up to τ=4 , corresponding to 400 times the time-step size of FTM.
Figure 7: Burgers: The relative errors in energy (left) and enstrophy (right) exhibit similar behavior for FTM and FM-FTM across different jump horizons τ , demonstrating that FM-FTM retains the accuracy of these physical ensemble statistics while reducing the number of time-integration steps.
Figure 8: Stochastically forced turbulence: Over the time interval, FM-FTM requires only 20 ( τ=1.25 ) flow-map evaluations, compared with 1000 sequential integration steps for FTM ( h=0.025 ), while achieving comparable relative errors in ensemble-mean energy (left) and enstrophy (right).
Figure 9: Stochastically forced turbulence: With a jump horizon of τ=1.25 , FM-FTM generates trajectories approximately 50× faster than FTM and about 500× faster than the physics model (left), while predicting solution fields that accurately match the decay of the kinetic-energy spectrum across wavenumbers (right).
We introduce First-Order Trajectory Matching (FTM), a surrogate-modeling method that learns the first-order local transport of probability mass from trajectories of stochastic systems. By matching the symmetric first-order motion of trajectories, FTM learns the probability current velocity, whose flow preserves time marginals to match ensemble averages, while also capturing current-like trajectory quantities such as fluxes, circulations, and barrier-crossing currents. FTM learns the current velocity directly from trajectories, avoiding drift, diffusion, and score estimation. Our stability analysis separates discretization error from sampling variance and shows that the one-step simulation-free FTM loss is stable when temporal resolution and sample size are properly balanced. Across stochastic dynamical systems and PDE examples, we empirically demonstrate that FTM provides trajectory-aware ensemble predictions at low, deterministic-rollout cost.
Shreya Jha, Timo Schorlepp, Nicholas Geissler +2
Courant Institute of Mathematical Sciences, New York University
Probabilistic forecasting is important for predicting complex dynamical systems because intrinsic randomness and incomplete observations can cause the same observed state to evolve into multiple plausible futures. While flow matching is a flexible approach for probabilistic forecasting, it is computationally expensive. Streaming flow (SF) reformulates this approach to model temporal evolution efficiently by learning a continuous velocity field directly in physical time. However, SF learns a deterministic velocity field. Thus, it provides only a single future trajectory for a given fixed initial state and observation history. To overcome this limitation, we introduce Variational Streaming Flow (VSF). Our approach learns a latent distribution that is conditioned on the dynamics of interest. In turn, this enables probabilistic forecasting. Importantly, we retain the computational efficiency of SF by generating in physical time. Across deterministic and stochastic dynamical systems, VSF demonstrates superior predictive accuracy and distributional fidelity. We demonstrate the advantage for both long-horizon rollouts exceeding 1,000 steps, and settings with bifurcating dynamics. Moreover, VSF can be integrated into existing Joint-Embedding Predictive Architecture (JEPA)-based world models as a plug-and-play predictor to improve temporal dynamics and goal-directed success rate in navigation, motion planning, and manipulation.
Hans Hao-Hsun Hsu, Minseon Gwak, Soon Hoe Lim +2
Georgia Institute of Technology · Lawrence Berkeley National Lab · International Computer Science Institute +2
We benchmark transport-based generative models as well as distillation-based few-step methods for the probabilistic forecasting of stochastic fluid flows, with a particular focus on performance under limited inference budgets. All methods are evaluated on a two-dimensional Kolmogorov flow with stochastic forcing. We measure one-step distributional accuracy against large simulated reference ensembles and assess whether the invariant measure is preserved during autoregressive rollouts via the enstrophy spectrum. On the stochastic task, flow matching achieves the most accurate one-step conditional distribution at high inference budgets, while the second-order exponential integrator DPM-2 is strongest at very low NFE. Few-step distillation methods are competitive with the multi-step methods and preserve the enstrophy spectrum particularly well. A deterministic control task, in which the forcing over the prediction interval is observed, separates aleatoric from epistemic uncertainty. Model performance does not translate between the two settings: the distilled models are competitive on the stochastic task but least accurate on the control task. While stochastic diffusion samplers such as DDPM better preserve the enstrophy spectrum during rollouts in the stochastic setting, deterministic samplers such as DDIM and DPM-2 show better spectral preservation in the deterministic setting.
Sebastian Pfister, Benjamin Holzschuh, Nils Thuerey
School of Computation, Information and Technology, Technical University of Munich, Germany.