Recent advances in autoregressive neural surrogate models have enabled orders-of-magnitude speedups in simulating dynamical systems. However, autoregressive models are generally prone to distribution drift: compounding errors in autoregressive rollouts that severely degrade generation quality over long time horizons. Existing work attempts to address this issue by implicitly leveraging the inherent trade-off between short-time accuracy and long-time consistency through hyperparameter tuning. In this work, we introduce a unifying mathematical framework that makes this tradeoff explicit, formalizing and generalizing hyperparameter-based strategies in existing approaches. Within this framework, we propose a robust, hyperparameter-free model implemented as a conditional diffusion model that balances short-time fidelity with long-time consistency by construction. Our model, Self-refining Neural Surrogate model (SNS), can be implemented as a standalone model that refines its own autoregressive outputs or as a complementary model to existing neural surrogates to ensure long-time consistency. We also demonstrate the numerical feasibility of SNS through high-fidelity simulations of complex dynamical systems over arbitrarily long time horizons.
Figures & tables
Figure 1: Top 5 rows: Vorticity fields of the Kolmogorov flow from a numerical solver, Gaussian approximation of the transition density, ACDM with 200 denoising steps, Thermalizer and SNS refined trajectories for Kolmogorov flow over a trajectory of 15,000 steps. Bottom row: kinetic energy spectra at each timestep, averaged over 20 randomly initialized trajectories.
Figure 2: Phase space for the forward and reverse process : Moving to the right and moving to the left corresponds to the forward process and the reverse process for xt−1 . Moving up and down corresponds to the forward and the reverse process for xt , respectively. The dashed curved arrows represent a jump in the reverse process.
Figure 3: The x-axis corresponds to the diffusion time of the forward process for the conditional variable. The images below the x-axis are realizations of the forward process with the same initial condition xt−10 at different times s . The images in the top row are in direct correspondence with the images below via an estimator of the multi-noise-level denoising oracle x^t0=Dθ(xtS,xt−1s) . The y-axis in the L2 distance between the denoised image and the original image. The degradation in denoising accuracy and the smoothing of fields with higher noise injection in the conditional variable align with the theory.
Figure 4: Top 6 rows: Top and bottom layer vorticity fields of the two-layer QG system from a numerical solver, Gaussian approximation of the transition density, and SNS refined trajectories for Kolmogorov flow over a trajectory of 20,000 steps. Bottom 2 rows: kinetic energy spectra of top and bottom layer at each timestep, averaged over 5 randomly initialized trajectories, with min-max spread shaded.
Figure 5: For a trajectory xt∈Rd , we report two temporal consistency metrics. The spatio-temporal correlation at lag τ is computed as C(τ)=⟨⟨xt−xˉt,xt+τ−xˉt+τ⟩∥xt−xˉt∥2−1∥xt+τ−xˉt+τ∥2−1⟩t,IC where xˉt is the spatial mean of xt , and ⟨⋅⟩t,IC denotes averaging over time and initial conditions. The rate of change is computed from the discrete time derivative, R(t)=∥xt+1−xt/Δt∥1, and we plot the mean of R(t) over different random initial conditions. The blue curves disappear in the bottom left plot for two different reasons. From 0 to 250, the blue curve is not visible because it overlaps with the green one, indicating strong temporal accuracy of the point-wise estimate in the short rollouts. From 30000 to 30250, the blue curve is missing because the point-wise estimates are returning NaN from the blow up in scale when the distribution drifted too far away from the stationary distribution.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Different reasonable traversal strategies in the reverse SDE phase space
Figure 7: Histogram of denoising history for the numerical experiments in the main text
Figure 8: Each row in a rollout from the same initial condition using point-wise estimate surrogate and SNS refinement, but each following a different traversal strategy . The spectra are averaged over 20 snapshots at time t with independent initial conditions. On the right is the L2 distance between the temporal mean of numerical fields and the mean of each run, and similarly for the L2 distance between the densities.
Figure 9: 2 layer Quasigeostrophic Turbulence with SNS varients