PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers

TL;DR

PDE-Refiner addresses the core challenge of maintaining accurate long-horizon predictions with neural PDE surrogates by enforcing uniform attention to all frequency components through an iterative, diffusion-inspired refinement process. By training with a denoising objective and using multiple refinement steps, it overcomes the tendency of one-step MSE objectives to neglect low-amplitude frequencies that seed long-term error growth. The approach yields longer, more stable rollouts on KS and Kolmogorov-flow benchmarks, improves data efficiency via spectral augmentation, and provides a practical framework for uncertainty estimation. While it incurs higher per-step compute than simple baselines, PDE-Refiner surpasses state-of-the-art neural, numerical, and hybrid solvers in rollout quality and reliability, marking a significant advancement for physics-informed neural surrogates in long-range forecasting.

Abstract

Time-dependent partial differential equations (PDEs) are ubiquitous in science and engineering. Recently, mostly due to the high computational cost of traditional solution techniques, deep neural network based surrogates have gained increased interest. The practical utility of such neural PDE solvers relies on their ability to provide accurate, stable predictions over long time horizons, which is a notoriously hard problem. In this work, we present a large-scale analysis of common temporal rollout strategies, identifying the neglect of non-dominant spatial frequency information, often associated with high frequencies in PDE solutions, as the primary pitfall limiting stable, accurate rollout performance. Based on these insights, we draw inspiration from recent advances in diffusion models to introduce PDE-Refiner; a novel model class that enables more accurate modeling of all frequency components via a multistep refinement process. We validate PDE-Refiner on challenging benchmarks of complex fluid dynamics, demonstrating stable and accurate rollouts that consistently outperform state-of-the-art models, including neural, numerical, and hybrid neural-numerical architectures. We further demonstrate that PDE-Refiner greatly enhances data efficiency, since the denoising objective implicitly induces a novel form of spectral data augmentation. Finally, PDE-Refiner's connection to diffusion models enables an accurate and efficient assessment of the model's predictive uncertainty, allowing us to estimate when the surrogate becomes inaccurate.

Figures (26)

• Figure 1: Challenges in achieving accurate long rollouts on the KS equation, comparing PDE-Refiner and an MSE-trained model. (a) Example trajectory with predicted rollouts. The yellow line indicates the time when the Pearson correlation between ground truth and prediction drops below 0.9. PDE-Refiner maintains an accurate rollout for longer than the MSE model. (b) Frequency spectrum over the spatial dimension of the ground truth data and one-step predictions. For PDE-Refiner, we show the average spectrum across 16 samples. (c) The spectra of the corresponding errors. The MSE model is only accurate for a small, high-amplitude frequency band, while PDE-Refiner supports a much larger frequency band, leading to longer accurate rollouts as in (a).
• Figure 2: Refinement process of PDE-Refiner during inference. Starting from an initial prediction $\hat{u}^1(t)$, PDE-Refiner uses an iterative refinement process to improve its prediction. Each step represents a denoising process, where the model takes as input the previous step's prediction $u^k(t)$ and tries to reconstruct added noise. By decreasing the noise variance $\sigma_k^2$ over the $K$ refinement steps, PDE-Refiner focuses on all frequencies equally, including low-amplitude information.
• Figure 3: Experimental results on the Kuramoto-Sivashinsky equation. Dark and light colors indicate time for average correlation to drop below 0.9 and 0.8, respectively. Error bars represent standard deviation for 5 seeds. We distinguish four model groups: models trained with the common one-step MSE (left), alternative losses considered in previous work (center left), our proposed PDE-Refiner (center right), and denoising diffusion (center right). All models use a modern U-Net neural operator gupta2022towards. PDE-Refiner surpasses all baselines with accurate rollouts up to nearly 100 seconds.
• Figure 4: Analyzing the prediction errors of PDE-Refiner and the MSE training in frequency space over the spatial dimension. Left: the spectrum of intermediate predictions $\hat{u}^{0}(t),\hat{u}^{1}(t),\hat{u}^{2}(t),\hat{u}^{3}(t)$ of PDE-Refiner's refinement process compared to the Ground Truth. Center: the spectrum of the difference between ground truth and intermediate predictions, i.e. $|\text{FFT}(u(t)-\hat{u}^{k}(t))|$. Right: the spectrum of the noise $\sigma_k\epsilon^k$ added at different steps of the refinement process. Any error with lower amplitude will be significantly augmented during denoising.
• Figure 5: Left: Stable rollout time over input resolution. PDE-Refiner models the high frequencies to improve its rollout on higher resolutions. Right: Training PDE-Refiner and the MSE baseline on smaller datasets. PDE-Refiner consistently outperforms the MSE baseline, increasing its relative improvement to 50% for the lowest data regime.