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
Temporal surrogate models are effective for predicting chaotic dynamical systems where computational cost can be prohibitive. Several deep neural network architectures can be used for such purposes. In this work, a few commonly used architectures are compared using a common training protocol. The objective is to fairly assess the impact of model architectures for long-horizon prediction stability. Experiments are carried out for three problems, the double pendulum, the Kuramoto-Sivashinsky equations, and the Kolmogorov flow. The experiments are carried out with matching model capacity. Analysis is also carried out for a scenario where each model is individually optimized. It is observed that in both scenarios, the models exhibit categorical differences in long-horizon rollouts. For a concrete quantification, stepwise error injections and perturbation amplifications are analyzed using metrics such as local jacobian, relative one-step bias, and finite-time Lyapunov growth. Additionally, an attractor analysis is also conducted to assess how well the learned models replicate the underlying system geometry. An ablation study to isolate the impact of each component of a continuous-update architecture is also carried out. It is concluded that models that having integrator-like updates show lower bias and perturbation amplification yielding stable long-horizon rollout and more accurate predictions.
Fast surrogate modeling for high-dimensional physical dynamics requires more than low short-term error: useful models must roll out efficiently while preserving the statistical structure of long trajectories. Neural operators provide inexpensive autoregressive forecasts but can drift in turbulent regimes, whereas rolling diffusion and latent generative surrogates can represent stochastic transitions at the cost of multi-step denoising, noise-schedule design, or auxiliary compression models. We propose MeanFlow Long-term Invariant Spatiotemporal Consistency Autoregressive Models (MeLISA), a latent-free autoregressive generative surrogate built on pixel-space MeanFlow. MeLISA defines a blockwise stochastic transition kernel that generates each forecast block with a single model evaluation, avoiding latent encoders and iterative diffusion solvers at inference time. To stabilize long-horizon rollouts, MeLISA combines a Window-Consistency MeanFlow objective that learns conditional spatiotemporal generation from partially observed temporal windows with a Time Increment Consistency loss that constrains multi-lag finite increments and targets temporal-correlation structure. We evaluate MeLISA with compact UNet and scalable DiT backbones on two high-resolution benchmarks, extended 2D Kolmogorov flow at 256×256 and turbulent channel-flow slice at 192×192. MeLISA outperforms neural-operator baselines on short-term forecasting accuracy and long-horizon statistical metrics, including energy spectra, turbulent kinetic energy, and mixing-rate-related dynamics, while achieving inference speeds comparable to, and in some cases faster than, neural operators. Compact 3.7-5.7M-parameter variants already deliver strong parameter efficiency, and DiT variants provide a scalable path up to 150M parameters. Overall, MeLISA benefits both rollout efficiency and long-horizon statistical accuracy.
Tianyue Yang, Xiao Xue
The Center for Computational Science University College London
Neural partial differential equation (PDE) surrogates trained with mean squared error (MSE) can capture dominant dynamics but often lose high-frequency information and fine-scale structures. Generative models (e.g., Flow Matching) can capture the distribution of fine-scale details, yet stochastic errors may accumulate during autoregressive rollouts and compromise long-term stability. To address these limitations, we introduce Refine and Align(ReAl), a model-agnostic plug-in that augments pretrained PDE surrogates to improve the reconstruction of high-frequency structures associated with complex physical phenomena. In the refinement stage, a conditional flow-matching model learns the residual distribution around the deterministic base prediction, recovering missing fine-scale structures while retaining the dominant dynamics. In the alignment stage, to mitigate the generative discrepancies amplified in the autoregressive rollout, particularly in chaotic systems, we generate multiple trajectories from the same input and reinforce the refiner with trajectory-level rewards, thereby achieving long-horizon dynamical alignment. Experiments across diverse surrogate architectures and PDE tasks demonstrate that ReAl consistently improves high-frequency fidelity, forecasting accuracy, and long-horizon stability.
Jing Xu, Zenglin Xu, Li Yi +2
Shanghai Academy of AI for Science · Fudan University