StocBench: A Benchmark for Generative Modeling of Stochastic Dynamics

Motivation and related work

• Existing SPDE benchmarks evaluate neural operators or propose surrogate evaluation criteria, but do not carefully compare transport-based methods and distilled variants across inference budgets.The cited examples are SPDEBench and the concurrent TRIE framework (Srikishan et al., 2026).Source: §1 · p. 1

Evaluation

• Evaluation measures one-step mean and standard-deviation errors and energy distance, plus rollout differences in the enstrophy spectrum as a probe of invariant-measure preservation.Small conditional biases can accumulate autoregressively, causing long-run statistics to drift from the invariant measure.Source: §1 · pp. 1–2

Evaluation scale

• The evaluation uses 5,000 simulated next states and 3,000 model samples for each of 48 one-step conditions, and 250,000 predicted transitions across 5,000 rollouts.Training data comprise 1,500 trajectories of 200 frames each; long-term rollouts span 50 autoregressive steps.Source: §1 · p. 2

Contributions

• StocBench provides a reproducible benchmark for conditional stochastic forecasting in stochastically forced Kolmogorov flow, with datasets and robust ground-truth statistics.The benchmark is designed to be comparatively easy to use and includes evaluation metrics in a ready-to-use Python package.Source: §1 · pp. 1–2

Methods and inference budgets

• StocBench compares transport-based methods at 10–400 network evaluations per sample with distillation methods using one or two evaluations.Evaluated transport-based methods include DDPM (Ho et al., 2020), DDIM (Song et al., 2021a), DPM-2 (Lu et al., 2022), flow matching (Lipman et al., 2023), and stochastic interpolation (Albergo & Vanden-Eijnden, 2023; Albergo et al., 2025; Chen et al., 2024b).Source: §1 · p. 2

Benchmark tasks

• The benchmark evaluates stochastic forcing forecasts and a deterministic variant conditioned on exact forcing to separate aleatoric from epistemic uncertainty and test whether model rankings transfer.With forcing observed, the next state is fully determined by the current vorticity and forcing; the stochastic task instead has irreducible uncertainty.Source: §1 · p. 2

Neural PDE and SPDE benchmarks

• Neural PDE benchmarks emphasize deterministic short-horizon error, while stochastic benchmarks and studies cover architectural comparisons, trajectory diversity, and partially observed systems.PDE-Refiner (Lippe et al., 2023) uses iterative denoising to improve rollout stability but remains deterministic. SPDEBench evaluates architectures across noise regimes and discretizations; Kohl et al. (2026) assess trajectory accuracy and diversity, and Harder et al. (2025) study probabilistic surrogates.Source: §2 · p. 2

Benchmark scope

• Unlike concurrent TRIE (Srikishan et al., 2026), STOCBENCH systematically compares transport-based generative methods and distillation under controlled inference budgets.TRIE evaluates invariant-measure reproduction, CRPS-calibrated uncertainty, and computational efficiency on stochastic PDEs.Source: §2 · p. 3

Transport-based generative models for physical dynamics

• Diffusion, flow matching, and stochastic interpolation transport simple priors to target distributions through ODEs or SDEs, and STOCBENCH compares their strengths under limited inference budgets.These methods converge to the target distribution as integration steps increase (Ho et al., 2020; Song et al., 2021b; Lipman et al., 2023; Liu et al., 2023; Albergo & Vanden-Eijnden, 2023; Chen et al., 2024b; Albergo et al., 2025).Source: §2 · p. 3

Few-step and single-step inference

• Consistency models and adversarial diffusion distillation compress multi-step transport for faster ensemble forecasting, and STOCBENCH tests whether they match transport methods at low inference budgets.Distillation trades theoretical convergence guarantees for speed; Chen et al. (2025) show that gains from additional consistency steps depend on training-error reduction, explaining observed quality saturation beyond two steps.Source: §2 · p. 3

Multi-frame prediction and extended temporal context

• Temporal bundling and extended input histories can reduce compounding rollout errors or uncertainty, but STOCBENCH restricts evaluation to single-step predictions to isolate generative-step quality.Diffusion Forcing (Chen et al., 2024a) combines joint trajectory generation, flexible past-frame conditioning, and diffusion-style refinement; extending STOCBENCH to richer temporal formulations is left for future work.Source: §2 · p. 3

3.1. Stochastic Kolmogorov Flow

• Models forecast the stochastic next-state distribution P(ωt+∆t|ωt), because unobserved forcing makes outcomes uncertain; Figure 1 shows vorticity snapshots and ensemble statistics.The stochastic setup uses forced wavenumbers K = {(6, 0), (7, 0), (5, 5), (8, 8)}.Source: §3.1 · p. 4

3.2. Deterministic Kolmogorov Flow

• Holding sampled forcing ηt fixed over each interval makes the next vorticity deterministic; despite differing from Wiener forcing, the variants have similar average enstrophy spectra.Figure 2 presents the comparison of their average radial enstrophy spectra.Source: §3.2 · p. 4, Figure 2 · p. 4

3.3. Metrics

• Stochastic evaluation measures Conditional Mean Error (ϵµ), Conditional Standard Deviation Error (ϵσ), Energy Distance (ϵD), and Enstrophy Spectrum Error (ϵZ) for autoregressive invariant-measure preservation.The energy distance assesses distributional mismatch beyond the first two moments; enstrophy spectrum error compares rollout spectra with the solver’s.Source: §3.3 · p. 4

4. Baselines

• Baselines are compared by varying solver steps and measuring the trade-off between distributional fidelity and inference cost, counted as network forward passes (NFEs).The baselines learn P(ωt+∆t | ωt) and receive ωt as an additional network input.Source: §4 · p. 5

4.1. Diffusion Models

• Diffusion baselines share a trained noise predictor and sample by integrating either a reverse SDE or its probability-flow ODE, using DDIM, DPM-2, or DDPM.DDIM and DDPM require one NFE per step, while the second-order DPM-2 midpoint scheme requires two; DDPM uses ancestral sampling.Source: §4.1 · p. 5

4.2. Flow Matching

• Flow matching regresses a network onto conditional vector fields, and its conditional optimal-transport path yields straighter trajectories that need fewer integration steps.The benchmark integrates the learned ODE with Euler at one NFE per step.Source: §4.2 · pp. 5–6

4.3. Stochastic Interpolation

• Stochastic interpolation bridges consecutive physical states, using the current state as its base and the next snapshot as its target.The benchmark follows Chen et al. (2024b) schedules and integrates the corresponding SDE with Euler–Maruyama at one NFE per step.Source: §4.3 · p. 6

4.4. Consistency Distillation

• Consistency distillation trains a self-consistent denoiser along a pretrained diffusion teacher’s probability-flow ODE and evaluates it once or twice for sampling.CD-1 uses one evaluation at maximum noise; CD-2 adds an evaluation at a task-specific noise level.Source: §4.4 · p. 6

4.5. Adversarial Diffusion Distillation

• ADD-FM adapts adversarial diffusion distillation to flow matching with a one-step generator and a 10-step teacher, requiring one NFE.The generator combines student–teacher distillation and adversarial objectives, so it is not strictly limited by teacher sampling quality.Source: §4.5 · p. 6

5.1. Stochastic Kolmogorov Flow

• During stochastic rollouts, DDIM and DPM-2 converge to lower-enstrophy invariant measures, whereas DDPM retains most reference enstrophy and FM has the lowest high-budget spectrum error, nearly constant over 50 steps.The passage attributes stronger spectral damping to deterministic denoising and small-scale energy preservation to noise injected by stochastic sampling.Source: Figure 3 · p. 7, §5.1 · p. 8, Figure 4 · p. 8

5.2. Deterministic Kolmogorov Flow

• On the deterministic task, multi-step methods form a low-error cluster, while distilled models have higher mean error and deviate faster from the solver trajectory, although all models remain stable over 50 rollout steps.The task conditions on forcing over the next snapshot interval, so residual predictive variability reflects epistemic uncertainty rather than aleatoric uncertainty.Source: §5.2 · p. 8, Figure 6 · p. 9
• Deterministic rollouts preserve spatial structures more readily than stochastic rollouts, but CD-1 is an exception, with excessive enstrophy especially at high-frequency modes.Distillation methods do not consistently outperform multi-step methods in the deterministic setting.Source: §5.2 · pp. 8–9
• In the deterministic control task, single- and few-step distillation methods have higher residual variability than multi-step methods, and CD-1 has substantially more than ADD-FM.CD-1’s contrast with its stochastic-task performance is attributed to the sampling schedule: deterministic dynamics need additional lower-noise evaluations to eliminate residual variability.Source: §5.2 · p. 8, Figure 5 · p. 9

5.1. Stochastic Kolmogorov Flow

• Most stochastic-task methods underestimate one-step predictive variability at low inference budgets and approach ground truth as NFE increases; DPM-2 instead overestimates it at low budgets.The 20-NFE minimum in DPM-2’s standard-deviation error is attributed to coarse discretization compensating for systematic underestimation.Source: §5.1 · pp. 7–8, Figure 5 · p. 9

6. Discussion

• Flow matching achieves the best one-step conditional distribution at high stochastic-task budgets and, among multi-step samplers, best preserves the enstrophy spectrum.DPM-2's exponential integrator yields accurate conditional distributions at few NFEs, especially for the standard deviation.Source: §6 · p. 9
• DDIM, DPM-2, and DDPM share one trained model yet differ fivefold in stochastic rollout spectrum error: DDPM preserves the invariant measure, while DDIM and DPM-2 damp it.The ranking reverses on the deterministic task: DDPM's stochastic updates reintroduce small-scale enstrophy, beneficial stochastically but excessive deterministically; similar one-step metrics do not imply long-rollout spectral accuracy.Source: §6 · pp. 9–10, Figure 7 · p. 10
• Distillation methods can be highly competitive, with ADD-FM particularly preserving the invariant measure at one NFE.Their training often does not converge monotonically, and some predictive variability is epistemic, arising from imperfect learning rather than stochastic forcing.Source: §6 · p. 10
• Stochastic interpolation bridges nearby simulation states rather than transporting from pure noise, but requires a relatively large inference budget to become competitive.Source: §6 · p. 10
• No method dominates: flow matching leads stochastic forecasting only at high budgets, DPM-2 is strongest among multi-step methods at low budgets, and distillation is competitive at 1–2 NFEs only on the stochastic task.DPM-2's accuracy plateaus early as budgets increase; distilled methods fall behind all multi-step samplers on the deterministic task.Source: §6 · p. 10

7. Conclusion

• On the stochastic task, flow matching is most accurate at high inference budgets, while distilled models are competitive at one NFE and DPM-2 is the strongest low-budget multi-step method.Unlike the distilled models, DPM-2 does not suffer from unstable training.Source: §7 · p. 10
• Performance does not translate between stochastic and deterministic control tasks: distilled models are competitive on the former but least accurate on the latter, suggesting some variability is epistemic.DDPM best preserves the rollout enstrophy spectrum on the stochastic task, whereas DDIM and DPM-2 preserve it better on deterministic control.Source: §7 · p. 10
• STOCBENCH aims to make evaluation accessible and guide future model and distillation development; future work could add higher-dimensional, higher-resolution, and other stochastic-system scenarios.Source: §7 · p. 10

A.1. Stochastic Kolmogorov Flow

• Training uses 1,500 trajectories of 200 frames, simulated with a pseudo-spectral solver and Euler–Maruyama time stepping.Source: §A.1 · p. 14
• Each simulation starts from zero vorticity and runs 100 warm-up snapshot intervals to allow convergence toward the invariant measure µ.Source: §A.1 · p. 14
• Simulations run on a 256 × 256 grid and are reduced to 64 × 64 by spatial averaging.Source: §A.1 · p. 14
• All trajectories are normalized using the global mean and standard deviation to improve training stability.Source: §A.1 · p. 14
• One-step evaluation uses 48 conditioning states, each paired with an ensemble of 5,000 next-snapshot states to estimate ground-truth conditional statistics.Source: §A.1 · p. 14
• Invariant-measure evaluation starts autoregressive rollouts from 5,000 additional states.Source: §A.1 · p. 14

Training

• Training samples pair the current vorticity field and fixed forcing over the prediction interval with the next-snapshot target.The forcing replaces Wiener forcing and is provided to the model as conditioning information.Source: §A.2 · p. 14

Testing

• One-step testing uses 48 conditioning states and compares the ensemble mean with the deterministic target and ensemble standard deviation with zero.For invariant-measure and long-term trajectory evaluations, the study generates 5,000 trajectories, each containing 50 state–forcing pairs.Source: §A.2 · p. 14

B.1. Diffusion Models: DDIM, DPM-2, DDPM

• DDIM, DPM-2, and DDPM share a continuous-time ϵθ-prediction model with a linear VP noise schedule, while DDIM uses first-order probability-flow ODE integration and DDPM uses ancestral sampling (Song et al., 2021a; Ho et al., 2020).DDIM sets η = 0; DDPM sets η = 1.Source: §B.1 · pp. 14–15
• DPM-2 uses second-order DPM-Solver-2 with steps uniform in time; uniform-time and uniform log-SNR discretizations perform comparably, with slightly lower low-budget standard deviation errors for uniform time (Lu et al., 2022).Source: §B.1 · p. 14

B.2. Flow Matching

• Flow matching trains velocity model vθ on a conditional OT path with minimum noise level σmin and samples using Euler steps uniform in τ.Source: §B.2 · p. 15

B.3. Stochastic Interpolation

• Stochastic interpolation uses ατ = στ = 1 − τ and βτ = τ^2; it initializes from the conditioning state, learns conditional drift, and uses Euler–Maruyama steps over τ ∈ [0, 0.999].The integration excludes τ = 1 to avoid the degenerate noise-free endpoint.Source: §B.3 · p. 15

B.4. Consistency Distillation

• Consistency distillation trains an EMA-targeted student initialized from a pretrained EDM teacher to map neighboring probability-flow ODE trajectory points to the same clean sample across 100 Karras noise levels (Karras et al., 2022).Source: §B.4 · p. 15
• CD-1 samples once at σ(99), while CD-2 adds a task-specific refinement evaluation at σ(97) for stochastic forecasting or σ(50) for deterministic forecasting.Source: §B.4 · p. 15

B.5. Adversarial Diffusion Distillation

• ADD combines a frozen flow-matching teacher, a trainable generator, and a discriminator: the generator matches the teacher’s 10-step prediction in one step while producing realistic samples.Its objective balances teacher distillation against adversarial scoring; λ = 1 is pure distillation and λ = 0 is a standard GAN. The discriminator uses teacher U-Net encoder and bottleneck features with pooled-logit heads and an R1 penalty.Source: §B.5 · p. 15
• ADD updates the generator once per five discriminator updates to keep the discriminator ahead and stabilize training, with additional balancing similar to Berthelot et al. (2017).Source: §B.5 · p. 15

C. Consistency Distillation Sampling Schedules

• Consistency sampling is not derived from discretizing an SDE or probability-flow ODE, so adding steps need not improve sample quality and can degrade it.This motivates empirical schedule analysis in both stochastic and deterministic settings.Source: C. Consistency Distillation Sampling… · p. 15
• In the stochastic case, one consistency step at the highest noise level performs near optimally; a second high-noise step slightly improves mean error, but further steps yield no benefit.Figure 8 reports mean and standard deviation errors after each denoising step.Source: C. Consistency Distillation Sampling… · p. 16, Figure 8 · p. 16
• Consistency schedules that spread many steps across the full noise range perform worse on both mean and standard deviation errors in the stochastic case.Figure 8 compares the different schedules for stochastic and deterministic Kolmogorov flow.Source: C. Consistency Distillation Sampling… · p. 16, Figure 8 · p. 16

D. DPM-2 Discretization Schedule

• DPM-Solver uses uniform log-SNR steps concentrated at low noise, whereas the DPM-2 baseline uses uniform time steps; their mean errors are nearly identical on both tasks.Figures 9a and 9b compare these discretizations.Source: D. DPM-2 Discretization Schedule · p. 16, Figure 9 · p. 17
• For stochastic DPM-2, standard-deviation errors converge at high inference budgets, while log-SNR discretization amplifies their overestimation at low budgets.The low-budget overestimation is also noted in Figure 5, left.Source: D. DPM-2 Discretization Schedule · p. 16, Figure 9 · p. 17
• For deterministic DPM-2, uniform-time discretization has lower residual variability than log-SNR discretization across all inference budgets.The comparison is shown in Figures 9a and 9b.Source: D. DPM-2 Discretization Schedule · p. 16, Figure 9 · p. 17

• Figure 10 shows ensemble-averaged radial enstrophy spectra after 50 autoregressive rollout steps on stochastic Kolmogorov flow, with line styles denoting 1, 2, and 400 NFEs.Source: Figure 10 · p. 18, E. Supplementary Evaluation Figures · p. 18
• Figure 11 shows ensemble-averaged radial enstrophy spectra after 50 autoregressive rollout steps on deterministic Kolmogorov flow, with line styles denoting 1, 2, and 400 NFEs.Source: Figure 11 · p. 18, E. Supplementary Evaluation Figures · p. 18
• Figure 12 presents ground-truth conditional means and standard deviations for three stochastic Kolmogorov flow conditions, alongside model errors relative to ground truth.Color scales are shared within each column.Source: Figure 12 · p. 19