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.