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.
Figures & tables
Figure 1: (a) MSE training leads to high-frequency information loss, while stochastic refiner recovers the structures. (b) Sampling variability accumulates during stochastic autoregressive rollout, especially in the high-frequency regime. Our ReAl mitigates this accumulation of sampling variability.
Figure 2: Overview of Re fine and Al ign ( ReAl ), a model-agnostic plug-in for pretrained PDE surrogates. (a) Refine : a conditional flow matching refiner generates fine-scale residuals for the frozen deterministic prediction. (b) Align : trajectory rewards from multiple stochastic rollouts post-train the refiner for stable long-term PDE forecasting. Both stages optimize the same conditional velocity field, with trajectory rewards weighting its flow-matching objectives during alignment.
Figure 3: Qualitative results on Controlled Cylinder Real2Real for velocity components u and v . Residuals are computed against the ground truth, while variance is measured across random seeds. Base+ReAl produces more accurate and consistent predictions.
Table 1: Main experimental results on RealPDEBench ( Hu et al., 2026 ) . We report RMSE over the 10 predicted future frames on test set. lower is better. The small percentage in the bottom-right corner indicates the relative RMSE reduction over the corresponding baseline.
Figure 4: RMSE over long autoregressive rollout frames under the RealPDEBench Real2Real setting. Long-horizon error accumulation causes severe degradation for several base models, most notably DPOT on Foil and FNO on FSI. By contrast, Base+ReAl consistently limits error growth and preserves stable predictions over long rollout horizons.
Figure 5: Visualization of long-horizon Real2Real rollouts (v channel): Base+ ReAl (a) restores fine-scale details, (b) suppresses error accumulation, and (c, d) mitigates numerical dissipation to preserve flow energy and wake structures.
Figure 7
Method
Reward
hr
Evaluation R1
Evaluation R5
Evaluation R10
RMSE ↓
Rel. L2↓
fRMSE ↓
FE ↓
KE ↓
RMSE ↓
Rel. L2↓
fRMSE ↓
FE ↓
KE ↓
RMSE ↓
Rel. L2↓
fRMSE ↓
FE ↓
KE ↓
Base
–
–
0.0100
0.0750
0.0012
21.5557
3.77e-5
0.0227
0.1771
0.0040
46.0869
5.57e-5
0.0320
0.2603
0.0065
64.3384
6.28e-5
+ Re
–
–
0.0093
0.0673
0.0011
26.3517
3.39e-5
0.0243
0.1830
0.0043
49.6342
5.63e-5
0.0331
0.2492
0.0067
65.5248
6.33e-5
+ ReAl
RMSE
1
0.0082
0.0602
0.0010
20.0997
3.12e-5
0.0223
0.1694
0.0039
35.1164
5.49e-5
0.0314
0.2358
0.0063
59.6282
6.20e-5
+ ReAl
RMSE
5
0.0084
0.0619
0.0010
20.5995
3.22e-5
0.0200
0.1560
0.0032
35.4562
5.38e-5
0.0273
0.2070
0.0054
55.0176
5.91e-5
+ ReAl
RMSE
10
0.0087
0.0637
0.0011
20.1948
3.31e-5
0.0210
0.1605
0.0037
35.7731
5.45e-5
0.0269
0.2052
0.0053
55.2919
5.91e-5
Table 3: Reward ablation results on Controlled Cylinder. hr denotes the rollout block used to compute the training reward. The best results are shown in bold , and the second-best results are underlined .
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: More visualizations on the RealPDEBench Real2Real setting. Our ReAl can correct the generation trajectory of the Base model.
Figure 8: Visualizations on the RealPDEBench Sim2Real setting. Compared with Real2Real , Sim2Real involves a substantially larger distribution shift, making autoregressive rollouts more susceptible to error accumulation and long-horizon divergence. We therefore visualize predicted Frames 1, 5, 10, 15, and 20, which cover the informative stage before catastrophic rollout failure. Our ReAl can rectify the generation trajectory of the Base model.
Autoregressive neural PDE surrogates predict future states by repeatedly applying a learned one-step operator. This is a simple and widely used method, but small one-step errors can accumulate during long rollouts. The resulting drift often appears as spectral amplitude distortion, phase misalignment, and nonlinear mode-interaction error. These effects are especially important for time-dependent PDEs with clear Fourier structure. We introduce the Spectral Generator Neural Operator (SGNO), a structured autoregressive neural operator for long-horizon PDE forecasting. SGNO organizes each learned one-step map as a structured spectral evolution update. A real-valued nonpositive diagonal generator provides a gain-controlled spectral backbone, while a learned correction pathway with complex-valued spectral mixing completes the residual evolution. This design gives the autoregressive step an evolution-like structure while retaining the flexibility needed for dissipative, dispersive, transport-dominated, and nonlinear PDEs. SGNO is designed for periodic linear and semilinear evolution PDEs with Fourier multiplier linear dynamics. Across ten mechanism-matched APEBench tasks spanning this regime, SGNO consistently outperforms strong single-step autoregressive baselines in long-horizon rollout accuracy, reducing GMean100 by a median of 74.8% relative to the strongest available non-SGNO baseline, with per-task reductions ranging from 13.6% to 92.9%. The gains are strongest on dispersive and transport-dominated tasks, as well as tasks involving nonlinear closure and mode coupling. Spectral diagnostics show lower spectral energy error and improved rollout-level phase fidelity. Ablations show that the constrained generator, the structured update, and the learned correction pathway each contribute to performance. The code is available at https://github.com/cruiseresearchgroup/SGNO.
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
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.
Qi Liu, Laure Zanna, Joan Bruna
Courant Institute School of Mathematics, Computing, and Data Science, New York University, USA.