PaperlayerView exact revision on arXiv

Source-linked AI summary

StocBench: A Benchmark for Generative Modeling of Stochastic Dynamics

Sebastian Pfister, Benjamin Holzschuh, Nils Thuerey

arXiv:2608.22309v1cs.LG

TL;DR

Probabilistic forecasting of stochastic fluid flows requires models that capture conditional distributions and preserve long-run statistics under limited inference budgets. StocBench benchmarks transport and distillation methods on stochastic and forcing-observed deterministic Kolmogorov-flow tasks, finding that method rankings depend on budget and uncertainty setting. Flow matching is strongest at high budgets, DPM-2 at low budgets, while distilled and diffusion samplers exhibit setting-dependent trade-offs.

  • Problem

    Stochastic fluid forecasting needs evaluation of conditional distributions, rollout invariant measures, and uncertainty under limited inference budgets.

  • Method

    StocBench compares transport-based and few-step distillation models on stochastic Kolmogorov flow using one-step distributional metrics and rollout enstrophy spectra, plus a forcing-conditioned deterministic control task.

  • Results

    Performance depends on regime: flow matching is most accurate at high stochastic-task budgets, DPM-2 is strongest at low budgets, and stochastic-versus-deterministic rankings reverse for several samplers.

  • Takeaways & Limitations

    Inference budget and whether forcing is observed should guide sampler choice; distilled models are competitive stochastically but least accurate on the deterministic control task.

  • Takeaways & Limitations

    The benchmark focuses on accessible scenarios and leaves higher dimensions, resolutions, and other stochastic systems for future work.

Abstract

from arXiv · show

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.

1. Introduction

StocBench benchmarks probabilistic forecasting of stochastically forced fluid dynamics, evaluating both one-step conditional distributions and long-horizon invariant-measure preservation. It compares transport-based and few-step distillation methods across inference budgets and includes a deterministic control task separating aleatoric from epistemic uncertainty.

  • Evaluation: The benchmark measures one-step mean, standard deviation, and energy-distance errors, plus rollout enstrophy-spectrum differences.The long-term metric probes whether autoregressive predictions preserve the invariant measure.
  • Experimental foundation: The benchmark uses 1,500 training trajectories and large simulated ensembles for both one-step and long-term evaluation.Ground-truth statistics use 5,000 next states for each of 48 conditions, while model ensembles use 3,000 samples per condition.
  • Motivation: StocBench evaluates conditional next-state distributions rather than single best predictions for stochastic fluid dynamics.The benchmark targets probabilistic forecasting settings such as turbulent flows with unresolved stochastic forcing.
  • Methods: StocBench compares transport-based samplers from 10 to 400 network evaluations with distillation methods using one or two evaluations.The comparison is designed to assess performance under constrained inference budgets.
  • Task design: A deterministic forcing-conditioned task separates irreducible aleatoric uncertainty from epistemic uncertainty caused by imperfect learning.This setting tests whether performance transfers between stochastic and deterministic forecasting.

2. Related Work

Related work spans stochastic PDE benchmarks, transport-based physical forecasting, and few-step distillation. StocBench emphasizes controlled distributional comparison of transport and distillation methods under inference budgets while restricting the task to single-step prediction.

  • Stochastic PDE benchmarks: Existing stochastic PDE benchmarks evaluate architectures, noise regimes, discretizations, invariant measures, uncertainty calibration, or computational efficiency.The cited benchmarks include SPDEBench and the concurrent TRIE framework.
  • Transport-based forecasting: Transport-based generative models, including diffusion and flow matching, are increasingly used for physical simulation and probabilistic forecasting.Prior applications include weather prediction and surrogate modeling of physical dynamics.
  • Few-step inference: Consistency models and adversarial diffusion distillation compress multi-step transport into one or a few network evaluations.This trade-off is relevant to ensemble forecasting, which requires many independent samples.
  • Temporal extensions: Temporal bundling predicts multiple future frames jointly to reduce exposure bias and compounding rollout errors.Extended history conditioning can also provide temporal context that reduces uncertainty and improves long-term consistency.
  • Scope of StocBench: StocBench deliberately restricts evaluation to the t_i → t_i+1 setting for a controlled comparison of generative modeling and distillation strategies.Its dataset preview pairs consecutive vorticity snapshots with next-state ensemble statistics estimated from 3,000 samples.

3. Benchmark

The benchmark uses stochastic and deterministic Kolmogorov-flow variants to evaluate conditional forecasting, uncertainty handling, and invariant-measure preservation during rollouts.

  • 3.1. Stochastic Kolmogorov Flow: Stochastic Kolmogorov flow tests learning P(ωt+∆t|ωt) when forcing realizations are unobserved.The next state is uncertain because the forcing over [t, t + ∆t] is not observed.
  • 3.2. Deterministic Kolmogorov Flow: The deterministic variant conditions on fixed forcing, making the next vorticity deterministic up to numerical solver accuracy.This setting tests prediction when the applied forcing is known throughout the interval.
  • 3.2. Deterministic Kolmogorov Flow: Both variants have very similar average enstrophy spectra despite differing forcing procedures.The stochastic version uses Wiener forcing, whereas the deterministic version holds forcing fixed over each snapshot interval.
  • 3.3. Metrics: One-step evaluation measures conditional mean error, conditional standard-deviation error, and energy distance.These metrics assess errors in expected state, stochastic spread, and distributional mismatch beyond the first two moments.
  • 3.3. Metrics: Rollout evaluation measures invariant-measure preservation through enstrophy spectrum error across autoregressive steps and radial wavenumber bins.Predicted and solver ensemble-averaged radial enstrophy spectra are compared from initial states drawn from the invariant measure.

4. Baselines

The baselines represent diffusion, flow-matching, stochastic-interpolation, and distillation approaches, compared through continuous-time transport formulations and varying inference costs.

  • 4. Baselines: Inference cost is measured by network function evaluations, enabling comparison of distributional fidelity against computational budget.Changing solver steps changes the number of forward passes required to generate a sample.
  • 4.1. Diffusion Models: Diffusion baselines reverse a Gaussian-noise corruption process using either a reverse SDE or a probability-flow ODE.The score is represented through a trained noise predictor, and sampling integrates from τ = 1 to τ = 0.
  • 4.1. Diffusion Models: DDIM uses one NFE per step, DPM-2 uses two NFEs per step, and DDPM performs stochastic ancestral sampling with one NFE per step.The samplers share one trained noise-prediction model but differ in numerical integration or stochastic update procedure.
  • 4.2. Flow Matching: Flow matching learns a time-dependent vector field that transports a Gaussian prior to the target distribution.The conditional optimal-transport path uses straight trajectories, allowing fewer integration steps for accurate sampling.
  • 4.3. Stochastic Interpolation: Stochastic interpolation bridges consecutive physical states with a Wiener-driven process rather than transporting from a Gaussian prior.Its endpoints are x0 = ωt and x1 = ωt+∆t, and the corresponding SDE is integrated with Euler–Maruyama at one NFE per step.

5. Experiments

The experiments compare transport-based and distilled models on stochastic and deterministic Kolmogorov-flow forecasting, emphasizing one-step accuracy, rollout stability, and invariant-measure preservation across inference budgets. Flow matching is strongest for several high-budget stochastic metrics, DPM-2 performs well at low budgets, while distilled models' stochastic-task strengths do not transfer to deterministic control.

  • Experimental setup: The benchmark evaluates stochastic and deterministic Kolmogorov-flow baselines using one-step distributional metrics, predictive variability, deterministic-solver accuracy, and enstrophy-spectrum preservation during rollouts.The stochastic setting uses unobserved forcing, while the deterministic control setting conditions on the exact forcing over the prediction interval.
  • 5.1. Stochastic Kolmogorov Flow: Flow matching achieves the lowest mean error across all stochastic inference budgets and the lowest standard deviation error at higher budgets.Stochastic interpolation approaches flow matching for standard deviation error as inference budgets increase.
  • 5.1. Stochastic Kolmogorov Flow: DPM-2 performs particularly well at low stochastic inference budgets, reaching its lowest standard deviation error at 20 NFEs and outperforming other multi-step baselines.Its low-budget advantage is associated with coarse discretization effects compensating for systematic standard-deviation underestimation.
  • 5.1. Stochastic Kolmogorov Flow: Most stochastic models underestimate one-step predictive variability at low budgets and approach the ground truth as NFEs increase, whereas DPM-2 overestimates it at low budgets.Predictive variability is measured as the spatial L2 norm of the pointwise ensemble standard deviation.
  • 5.1. Stochastic Kolmogorov Flow: On stochastic rollouts, DDPM preserves most enstrophy, while DDIM and DPM-2 converge to invariant measures with lower enstrophy than the reference; flow matching performs best among high-budget multi-step methods.The reported pattern is attributed to stochastic sampling preserving small-scale energy through injected noise, unlike deterministic denoising.
  • 5.2. Deterministic Kolmogorov Flow: Distilled models show higher residual variability and larger deterministic mean errors than multi-step methods, with distillation trajectories deviating faster from the numerical solver.This indicates greater epistemic uncertainty in few-step deterministic prediction; CD-1 also develops excessive high-frequency enstrophy after rollout.

6. Discussion

The discussion finds that sampler strengths depend on inference budget and task setting: flow matching leads at high budgets, DPM-2 at low budgets, and distillation can preserve invariant measures efficiently. Spectral preservation also depends on whether forcing is stochastic or deterministic.

  • Flow matching is most accurate at high inference budgets on the stochastic task and best preserves the enstrophy spectrum among multi-step samplers.
  • DPM-2 performs best in the low-budget regime, producing accurate conditional distributions with only a few NFEs.Its exponential integrator is especially effective for standard deviation accuracy.
  • DDPM, DDIM, and DPM-2 share one trained model but differ in rollout spectral preservation because their sampling procedures produce different invariant-measure behavior.On the stochastic task, DDPM preserves the invariant measure while DDIM and DPM-2 converge to a damped one; the ranking reverses on the deterministic task.
  • Distillation methods can be competitive at very low inference budgets, with ADD-FM showing strong invariant-measure preservation at one NFE.
  • No single method dominates all regimes: flow matching leads at high budgets, DPM-2 at low budgets, and distillation is competitive on the stochastic task but weaker in the deterministic setting.Stochastic interpolation requires a relatively large inference budget to become competitive.

7. Conclusion

STOCBENCH evaluates generative models for stochastic dynamics across inference budgets and task settings. Its results show that preferred methods differ by budget and that performance on stochastic forecasting does not transfer directly to deterministic control.

  • STOCBENCH benchmarks generative modeling of stochastic dynamics, with method choice determined partly by the available inference budget.
  • Flow matching is most accurate at high budgets, while distilled models remain competitive at one NFE on the stochastic task.
  • DPM-2 is the strongest multi-step method at low budgets and does not suffer from the unstable training reported for distilled models.
  • Performance does not translate between stochastic forecasting and deterministic control: distilled models are competitive in the former but least accurate in the latter.Among diffusion samplers, DDPM best preserves stochastic-task rollout enstrophy, whereas DDIM and DPM-2 preserve it better in deterministic control.
  • Future work could extend the benchmark to higher dimensions, higher resolutions, and other stochastic systems.

A. Dataset Generation

The dataset is generated from simulated Kolmogorov-flow trajectories and evaluated using large reference ensembles. Warm-up integration targets the invariant measure before training data and test states are collected.

  • The training dataset contains 1,500 trajectories with 200 frames each, generated by pseudo-spectral Euler–Maruyama integration.
  • Simulations start from zero vorticity and run through 100 warm-up snapshot intervals so the state distribution can approach the invariant measure.
  • The simulations use a 256 × 256 grid and are reduced to 64 × 64 through spatial averaging before normalization.
  • One-step testing uses 48 conditioning states with 5,000 next-snapshot samples per state to estimate reference conditional statistics.Another 5,000 states initialize invariant-measure autoregressive rollouts.

A.2. Deterministic Kolmogorov Flow

The deterministic variant replaces stochastic Wiener forcing with observed piecewise-constant forcing. Its one-step targets are therefore deterministic, while the prediction horizon is set by the snapshot interval.

  • The deterministic dataset replaces Wiener forcing with piecewise-constant forcing held fixed between consecutive snapshots.
  • The model receives the current vorticity field and forcing during the prediction interval, and predicts the next vorticity snapshot.
  • Because the interval forcing is observed, each one-step test case has a deterministic next state, so ensemble mean and standard deviation are compared with the target and zero.
  • The snapshot interval ∆t defines the prediction horizon, with the model mapping the state at t to the snapshot at t + ∆t.

B. Baseline Specifications

The baselines use shared diffusion parameterization where applicable but differ in sampling procedures, while flow matching, conditional flow matching, consistency distillation, and adversarial diffusion distillation use distinct training and sampling designs.

  • Diffusion baselines: DDIM, DPM-2, and DDPM share a continuous-time ϵθ-prediction model with a linear VP noise schedule but differ in sampling procedure.DDIM uses η = 0, DPM-2 uses a second-order exponential integrator, and DDPM uses η = 1 ancestral sampling.
  • Diffusion baselines: DPM-2 discretizes integration uniformly in time, with comparable performance to uniform log-SNR discretization and slightly lower low-budget standard deviation errors.The comparison is reported for both stochastic and deterministic tasks.
  • Flow matching: The velocity model is trained on a conditional OT path and sampled with Euler steps uniformly in τ.Its training uses interpolation schedules ατ = στ = 1 − τ and βτ = τ^2.
  • Conditional flow matching: Conditional flow matching initializes sampling from the conditioning state and learns conditional drift, using Euler–Maruyama steps over τ ∈ [0, 0.999].The endpoint restriction avoids the degenerate noise-free endpoint.
  • Distillation methods: Consistency distillation trains a student from a pretrained EDM teacher so neighboring points along the teacher probability-flow ODE map to the same clean sample.The teacher noise schedule is discretized into T = 100 levels, with exponential-moving-average stabilization.
  • Distillation methods: ADD combines a frozen flow-matching teacher, a trainable one-step generator, and a discriminator, with generator training balancing distillation and adversarial objectives.The discriminator distinguishes real from generated samples, while λ controls the distillation–adversarial trade-off.

C. Consistency Distillation Sampling Schedules

The section evaluates how consistency-distillation schedules and DPM-2 discretization choices affect stochastic and deterministic Kolmogorov-flow forecasting, especially under limited inference budgets.

  • Schedule rationale: Consistency models do not derive multi-step sampling from an underlying SDE or probability-flow ODE, so additional steps may fail to improve quality.Too many consistency-model sampling steps can degrade sample quality, motivating empirical schedule analysis.
  • Configuration: The consistency-distillation configuration uses a Karras schedule with ρ = 7, σmin = 2 × 10^-3, σmax = 80, T = 100 noise levels, and EMA decay 0.999.The listed hyperparameters define the student’s noise schedule and training stabilization.
  • Configuration: The ADD configuration uses a 10-step teacher, one-step student, update ratio rgen = 5, distillation weight λ = 0.9, and R1 coefficient γ = 1.0.These are the model-specific settings listed for adversarial diffusion distillation.
  • Consistency distillation: In the stochastic case, one highest-noise denoising step is near-optimal, while a second high-noise step slightly improves mean error.Further steps provide no benefit, and schedules spanning the full noise range perform worse on mean and standard deviation errors.
  • Consistency distillation: Figure 8 compares stochastic and deterministic consistency-distillation schedules using mean and standard deviation errors measured after each denoising step.Each curve represents a sampling schedule, with markers accumulating denoising steps through the indicated noise level.
  • DPM-2 discretization: DPM-2’s uniform-time and uniform-log-SNR discretizations have nearly identical mean errors on both tasks.For stochastic forecasting, standard-deviation errors converge at high budgets, whereas log-SNR discretization amplifies low-budget standard-deviation overestimation.

E. Supplementary Evaluation Figures

The supplementary figures visualize long-rollout enstrophy spectra and conditional distribution statistics for stochastic and deterministic Kolmogorov flow.

  • Stochastic rollout evaluation: Figure 10 presents ensemble-averaged radial enstrophy spectra after 50 autoregressive rollout steps on stochastic Kolmogorov flow.
  • Figure conventions: The rollout figures use dotted, dashed, and solid lines to distinguish models using 1, 2, and 400 NFEs, respectively.
  • Deterministic rollout evaluation: Figure 11 presents ensemble-averaged radial enstrophy spectra after 50 autoregressive rollout steps on deterministic Kolmogorov flow.
  • Conditional statistics: Figure 12 compares conditional means and standard deviations across three stochastic Kolmogorov-flow conditions.The top row shows ground-truth statistics, while remaining rows show model errors relative to ground truth, with shared color scales within columns.
Loading 2608.22309v1…