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.
Figures & tables
Figure 1: Dataset preview for stochastic Kolmogorov flow. First row: consecutive vorticity snapshots. Second row: next-state ensemble statistics for a fixed initial condition, with the mean and standard deviation estimated from 3000 samples.
Figure 2: Both Kolmogorov-flow variants have very similar average radial enstrophy spectra.
Figure 3: Metrics for stochastic Kolmogorov flow that measure one-step distributional fit (a–c) as well as preservation of the invariant measure μ during autoregressive rollout, assessed in the enstrophy spectrum (d). All metrics are defined in Section 3.3 . Lower values indicate closer agreement with the reference distribution. Markers denote the inference budget while dotted and dashed lines indicate models using 1 and 2 NFEs, respectively.
Figure 4: Evaluation of model invariant measures for stochastic Kolmogorov flow, obtained from autoregressive rollouts and assessed by the enstrophy spectrum. Left : enstrophy spectrum error at each step of a 50-step rollout. Right : ratio of predicted to reference enstrophy after 50 rollout steps; values above the black line indicate overestimation, while values below indicate underestimation. In both subplots, dotted, dashed, and solid lines denote models using 1, 2, and 400 NFEs, respectively.
Figure 5: Variability of one-step predictions for the stochastic (left) and deterministic (right) Kolmogorov flow as a function of NFE. Markers denote the inference budget while dotted and dashed lines indicate models using 1 and 2 NFEs, respectively.
Figure 6: Accuracy on deterministic Kolmogorov flow relative to the numerical solver. Left : For each inference budget, we report the error between the ensemble mean of the one-step predictions and the deterministic ground-truth next snapshot generated by the numerical solver. Markers denote the inference budget, with dotted and dashed lines corresponding to 1 and 2 NFEs, respectively. Right : RMSE at each rollout step, averaged over autoregressive trajectories initialized from randomly sampled states. Markers denote the inference budget, while dotted, dashed, and solid lines correspond to evaluations using 1, 2, and 400 NFEs, respectively.
Figure 7: Evaluation of model invariant measures for deterministic Kolmogorov flow, obtained from autoregressive rollouts and assessed by the enstrophy spectrum. Left : enstrophy spectrum error at each rollout step. Right : ratio of predicted to reference enstrophy after 50 rollout steps; values above the black line indicate overestimation, while values below indicate underestimation. In both subplots, dotted, dashed, and solid lines denote models using 1, 2, and 400 NFEs, respectively.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Parameter
Symbol
Value
Solver time step
h
10−4
Snapshot interval
Δt
0.5
Solver grid resolution
N
256
Model grid resolution
N
64
Viscosity
ν
10−3
Linear drag coefficient
α
0.1
Appendix
Table 1: Physical constants and simulation parameters for the stochastic and deterministic Kolmogorov-flow datasets. The SDE is integrated with solver time step h , while snapshots are recorded at intervals of Δt . The interval Δt also defines the model prediction horizon, i.e., the model predicts the snapshot at t+Δt from the state at t .
Parameter
Symbol
Value
Diffusion
Variance schedule
β(τ)
VP linear
Schedule limits
[βmin,βmax]
[0.1,20]
Sampler
–
DPM-2 / DDIM / DDPM
Flow Matching
Minimum noise level
σmin
10−3
Appendix
Table 2: Model-specific hyperparameters. The diffusion models share the same continuous-time ϵθ -prediction parameterization and differ only in the sampler used at inference time.
Figure 8: Evaluation of different Consistency Distillation sampling schedules for stochastic and deterministic Kolmogorov flow. Each curve corresponds to a sampling schedule. The mean and standard deviation errors at each marker are obtained by performing denoising steps at all marked noise levels up to and including that marker.
Figure 9: Uniform vs. log-SNR discretization for DPM-2 on stochastic and deterministic Kolmogorov flow.
Figure 10: Ensemble-averaged radial enstrophy spectra after 50 autoregressive rollout steps on stochastic Kolmogorov flow. Dotted, dashed, and solid lines indicate models using 1, 2, and 400 NFEs, respectively.
Figure 11: Ensemble-averaged radial enstrophy spectra after 50 autoregressive rollout steps on deterministic Kolmogorov flow. Dotted, dashed, and solid lines indicate models using 1, 2, and 400 NFEs, respectively.
Figure 12: Conditional means and standard deviations for three stochastic Kolmogorov flow conditions. The top row presents the ground-truth statistics, and the remaining rows show model errors relative to the ground truth. Color scales are shared within each column.
Many scientific systems exhibit uncertainty from stochastic forcing, unresolved degrees of freedom, or imperfect observations, making reliable surrogate forecasting fundamentally distributional rather than pointwise. For such systems, deterministic neural surrogates fail to capture statistical measures and forecast uncertainty. We introduce TRIE, an evaluation framework for stochastic PDE surrogates that asks whether models reproduce invariant measures, provide trustworthy predictive uncertainty, and scale to efficient probabilistic generation. We demonstrate TRIE on two stationary chaotic spatially extended SPDEs, stochastic Kuramoto--Sivashinsky and stochastic Kolmogorov flow, across 11 parameter values. Our evaluation shows that standard pointwise-trained neural surrogates can produce plausible short rollouts while failing to match long-time statistical structure. Approximate uncertainty methods such as Monte Carlo dropout and heteroscedastic Gaussian likelihoods produce stochastic forecasts, but are often miscalibrated and overconfident under temporal and spatial uncertainty diagnostics. Across these criteria, generative models provide the most consistent performance, accurately capturing invariant measure statistics and achieving the lowest CRPS in all reported probabilistic settings. Finally, we show that latent generative models with automatic dimension discovery retain much of this statistical fidelity while reducing Kolmogorov inference time by roughly 12×. We release our code and data at https://github.com/scailab/TRIE-SPDE-Bench to support reproducible evaluation of stochastic PDE forecasting models.
Bharat Srikishan, Javier E. Santos, Nikhil Muralidhar +1
Stevens Institute of Technology · Los Alamos National Laboratory
Modern generative modeling is dominated by transport from a noise prior to data. We propose an alternative paradigm in which generation is performed by a discrete stochastic dynamics that leaves the data distribution invariant, initialized from data-supported states rather than from noise. The framework can utilize any pretrained flow model. We develop two probability-preserving sampling mechanisms, a corrected Langevin dynamics with a Metropolis adjustment and a predictor-corrector flow, that operate directly on existing checkpoints. We validate the framework on a synthetic Swiss-roll target, ImageNet-256 and Oxford Flowers-102, where our samplers consistently improve over the original generation procedures.
Eshed Gal, Md Shahriar Rahim Siddiqui, Moshe Eliasof +1
The University of British Columbia Vancouver, Canada · University of Cambridge Cambridge, United Kingdom
Analyzing unsteady fluid flows often requires access to the full distribution of possible temporal states, yet conventional PDE solvers are computationally prohibitive and learned time-stepping surrogates quickly accumulate error over long rollouts. Generative models avoid compounding error by sampling states independently, but diffusion and flow-matching methods, while accurate, are limited by the cost of many evaluations over the entire mesh. We introduce scale-autoregressive modeling (SAR) for sampling flows on unstructured meshes hierarchically from coarse to fine: it first generates a low-resolution field, then refines it by progressively sampling higher resolutions conditioned on coarser predictions. This coarse-to-fine factorization improves efficiency by concentrating computation at coarser scales, where uncertainty is greatest, while requiring fewer steps at finer scales. Across unsteady-flow benchmarks of varying complexity, SAR attains substantially lower distributional error and higher per-sample accuracy than state-of-the-art diffusion models based on multi-scale GNNs, while matching or surpassing a flow-matching Transolver (a linear-time transformer) yet running 2-7x faster than this depending on the task. Overall, SAR provides a practical tool for fast and accurate estimation of statistical flow quantities (e.g., turbulent kinetic energy and two-point correlations) in real-world settings.