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

Phillip Lippe; Bastiaan S. Veeling; Paris Perdikaris; Richard E. Turner; Johannes Brandstetter

Source-linked AI summary

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

Phillip Lippe, Bastiaan S. Veeling, Paris Perdikaris, Richard E. Turner, Johannes Brandstetter

arXiv:2308.05732v2Aug 10, 2023cs.LGcs.AI

TL;DR

Neural PDE solvers struggle to maintain accurate, stable long-term predictions because rollout strategies neglect low-amplitude spatial frequencies. PDE-Refiner uses iterative diffusion-inspired refinement to model all frequency components and achieves longer accurate rollouts across fluid-dynamics datasets.

Problem
Neural PDE solvers struggle with accumulating errors during long-horizon rollouts, partly because common strategies neglect low-amplitude spatial-frequency components.
Method
PDE-Refiner uses an iterative diffusion-inspired denoising process to model all frequency components more accurately at different amplitude levels.
Results
PDE-Refiner consistently outperforms neural, numerical, and hybrid baselines, achieving stable, accurate rollouts for longer across three fluid-dynamics datasets.
Takeaways & Limitations
PDE-Refiner offers a practical route to more reliable long-horizon neural PDE simulation while also supporting predictive-uncertainty assessment.
Takeaways & Limitations
PDE-Refiner requires increased computation time per prediction, despite remaining faster than hybrid and classical solvers.

Abstract

from arXiv · show

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.

1 Introduction

Deep neural PDE surrogates offer computational efficiency, but long-horizon prediction for complex time-dependent PDEs remains fundamentally difficult. PDE-Refiner addresses a shared rollout failure—neglect of spatial-frequency components—through iterative diffusion-inspired refinement across the whole spectrum.

Motivation: Deep neural PDE surrogates have gained momentum as a more computationally efficient methodology for modeling PDEs.Recent approaches include neural, hybrid, and numerical-replacement categories.
Motivation: Neural approaches have succeeded on inverse and high-dimensional problems, while hybrid and numerical-replacement approaches have advanced two- and three-dimensional fluid and weather modeling.These applications are commonly described by complex time-dependent PDEs.
Challenge: Long-horizon prediction for complex time-dependent PDEs presents fundamental challenges, including accumulating approximation effects in conventional numerical methods.The supplied passage introduces these difficulties but is truncated before fully describing them.
Challenge: Across autoregressive, pushforward, invariance-preserving, and Markov Neural Operator rollouts, neural solvers consistently neglect components of the spatial frequency spectrum.The analysis tests modern U-Nets and Fourier Neural Operators among state-of-the-art neural operators.
Contribution: PDE-Refiner uses diffusion-inspired iterative refinement to produce accurate predictions over the whole frequency spectrum.An adapted Gaussian denoising step makes the network focus equally on all frequency components at different amplitude levels.

2 Challenges of Accurate Long Rollouts

Accurate long rollouts for time-dependent PDEs are difficult because learned temporal updates accumulate errors, especially when recursively applied over many steps. Analysis of the KS equation identifies incomplete modeling of spatial frequencies as a key limitation, motivating methods that preserve both dominant and non-dominant frequency information.

Time-dependent PDE solvers can use large-step temporal operators or recursively apply smaller updates, but learned operators deteriorate for large ∆t and autoregressive approaches perform better.
Long Rollouts for Time-Dependent PDEs: PDE-Refiner maintains an accurate KS rollout longer than an MSE-trained model and uses iterative refinement to focus equally on all frequencies, including low-amplitude information.During inference, successive denoising steps reduce the noise variance σ2_k across K refinement steps.
Long Rollouts for Time-Dependent PDEs: The fourth-order nonlinear Kuramoto-Sivashinsky equation is challenging because its nonlinear term uux and fourth-order derivative uxxxx produce rich, chaotic dynamics.The equation is solved on a uniformly discretized spatial-temporal grid with periodic boundary conditions.
Long Rollouts for Time-Dependent PDEs: One-step MSE training accurately predicts the next step, but rollout errors propagate until predictions significantly diverge from the ground truth.In Figure 1a, the solver is reported to remain accurate for 70 se…
Long Rollouts for Time-Dependent PDEs: Stable long rollouts require accurately modeling all spatial frequencies, prioritizing high-amplitude frequencies while retaining low-amplitude contributions because of their long-term dynamical impact.

3 PDE-Refiner

PDE-Refiner improves predictions through iterative denoising refinement, enabling modeling across dominant and non-dominant frequencies. Its training and inference procedures use noise scheduling and sequential refinement, while its diffusion-model connection supports uncertainty estimation.

Model: PDE-Refiner repeatedly refines predictions using previous time steps, a refinement index, and the model’s current prediction to model all frequencies.The first refinement step mimics the common MSE objective, while later steps improve the prediction iteratively.
Denoising objective: At each refinement step, Gaussian noise is added to the current prediction, and the model predicts that noise to denoise the input.The update is ˆu_k+1(t) = ũ_k(t) − σ_kˆϵ_k.
Noise schedule: Decreasing σ_k over refinement steps shifts attention from high-amplitude information to low-amplitude information associated with non-dominant frequencies.An exponential schedule is reported to work well, with σ_min selected from the data’s frequency spectrum.
Training: During training, ground-truth samples make each refinement step learn information below its noise level, while uniformly sampling k trains all refinement steps equally well.This prevents the model from learning to reproduce potentially larger errors from earlier inference steps.
Inference and uncertainty: At inference, PDE-Refiner sequentially performs K refinement steps, using each prediction as the next input; independent Gaussian noise treats frequencies uniformly.Its diffusion-model connection also suggests that sampling different refinement noises can provide calibrated uncertainty estimates indicating possible model errors.

4 Experiments

Experiments on 1D Kuramoto–Sivashinsky and 2D Kolmogorov flow show that PDE-Refiner improves long-rollout accuracy across frequency spectra and challenging solver baselines, while enabling spectral data augmentation and uncertainty estimation.

Kuramoto–Sivashinsky: PDE-Refiner reaches almost 100 seconds of stable rollout on the Kuramoto–Sivashinsky benchmark, significantly outperforming baseline strategies.Performance improves with more refinement steps, but additional steps require more model calls and slow the solver.
Additional analyses: PDE-Refiner induces spectral data augmentation by adding Gaussian noise at different refinement stages, forcing recovery of underlying structure across scales.The resulting inputs and objectives vary across stages, encouraging the model to fit different parts of the spectrum.
Additional analyses: Predictive uncertainty is estimated from 32 noisy rollouts per test trajectory by measuring when their cross-correlation falls below 0.8.The divergence time is investigated as an estimate of how long predictions remain accurate.
Kuramoto–Sivashinsky: PDE-Refiner consistently improves stable rollout time over the MSE model across viscosities, demonstrating robustness to different frequency spectra.The comparison evaluates trajectories grouped by viscosity and finds similar performance to the unconditional counterpart at ν = 1.0.
Kolmogorov flow: On 2D Kolmogorov flow, PDE-Refiner remains accurate 32% longer than the best single-input hybrid method and 10% longer than the best multi-input hybrid method.The MSE-trained U-Net already surpasses all classical and hybrid PDE solvers, making further improvement difficult.
Kolmogorov flow: PDE-Refiner generates the 2D test trajectories in 16.53 ± 0.04 seconds, slower than the MSE U-Net but faster than the TSM hybrid solver at 20.25 seconds (±0.05).The comparison uses 16 trajectories of 20 seconds on an NVIDIA A100 GPU; PDE-Refiner requires four model calls per step.

5 Conclusion

The paper identifies neglected low-amplitude information as a key limitation of accurate neural PDE rollouts and introduces PDE-Refiner’s iterative refinement to model all frequency components, achieving considerably longer accuracy across three fluid-dynamics datasets. Its primary limitation is increased computation time per prediction, although it remains faster than hybrid and classical solvers.

5 Conclusion: PDE-Refiner addresses neglected low-amplitude information with iterative refinement that accurately models all frequency components.The paper identifies this neglect as a key factor limiting accurate rollout times.
5 Conclusion: PDE-Refiner remains considerably longer accurate during rollouts on three fluid dynamic datasets, overcoming a common pitfall.
Limitations: PDE-Refiner’s primary limitation is increased computation time per prediction, despite remaining faster than hybrid and classical solvers.Future work could reduce compute for early refinement steps or use distillation and enhanced samplers to accelerate refinement.
Limitations: PDE-Refiner provides smaller gains with FNOs because it models high-frequency noise, motivating further investigation.

SUPPLEMENTARY MATERIAL · A Broader Impact

Neural PDE solvers could make simulations cheaper across natural-science applications while reducing research-related carbon emissions. Their use nevertheless requires rigorous cross-checking, continuous monitoring, and careful attention to errors outside the training distribution.

A Broader Impact: Neural PDE solvers could provide cheaper approaches to modeling natural phenomena than classical solvers.Potential applications include fluid dynamics and weather modeling.
A Broader Impact: PDE surrogates could advance research in fields such as fluid dynamics and weather modeling.
A Broader Impact: Reducing simulation compute may lower the carbon footprint of research institutes and industries relying on these models.
A Broader Impact: Neural surrogates require rigorous cross-checks and continuous monitoring because they may introduce errors outside their original training distribution.This is especially important when the surrogates were trained on simulations themselves, making their assumptions and limitations important to understand.

B Reproducibility Statement · C PDE-Refiner - Pseudocode · D Experimental details

The paper supports reproducibility through public code, detailed experimental specifications, repeated seeds, and reported standard deviations. It also provides direct and diffusion-based pseudocode for PDE-Refiner, while documenting datasets, architectures, hyperparameters, results, and example trajectories.

B Reproducibility Statement: The authors publish code and document model architectures, hyperparameters, and dataset properties, with KS experiments repeated for five seeds and Kolmogorov Flow experiments using three seeds.Plots and tables with quantitative results report standard deviations.
B Reproducibility Statement: The implementation builds on PDE-Arena, PyTorch, PyTorch Lightning, diffusers, Matplotlib, NumPy, scipy, and JAX.These assets support training, diffusion modeling, plotting, data handling, and dataset generation.
C PDE-Refiner - Pseudocode: The core PDE-Refiner pseudocode exposes the refinement-step count K and minimum noise standard deviation σmin, while allowing arbitrary neural operators such as U-Nets.The implementation represents these parameters as num_steps and min_noise_std.
C PDE-Refiner - Pseudocode: Training uniformly samples a refinement step, uses classical MSE at k = 0, and otherwise trains the neural operator to denoise a noised solution.Prediction and target construction are implemented through the train_step function.
C PDE-Refiner - Pseudocode: Prediction initializes the next solution at k = 0, iteratively removes noise across refinement steps, and rolls out by repeatedly appending predicted solutions to the trajectory.The pseudocode also provides an equivalent implementation using diffusers' DDPM scheduler with v_prediction.
D Experimental details: The experimental-details section specifies data generation, architectures, and hyperparameters for the KS, parameter-dependent KS, and 2D Kolmogorov-flow datasets, and presents tabulated error bars and example trajectories.Dataset examples show KS training trajectories with 140 time steps and test trajectories with 640 time steps.

D.1 Kuramoto-Sivashinsky 1D dataset

The 1D Kuramoto–Sivashinsky dataset uses pseudo-spectral data generation with truncated warmup steps and float32 storage, while models use a conditioned modern U-Net. PDE-Refiner maintains lower rollout MSE growth than baselines after the initial 20 seconds, at higher inference cost than the MSE baseline.

Data generation: Each trajectory is generated with method-of-lines pseudo-spectral derivatives, with the first 360 solution steps discarded as solver warmup.The dataset follows Brandstetter et al.’s setup.
Data generation: Data is generated in float64 and converted to float32, making wavenumbers higher than 60 inadequately represented and introducing high-frequency quantization noise.The conversion is performed in the spatial domain and causes minor frequency-spectrum artifacts.
Model architecture: All Section 4.1 models use a modern U-Net with pre-activation ResNet blocks, GroupNorm, GELU activations, skip connections, and conditioning on ∆t and ∆x.The U-Net uses scale-and-shift conditioning through AdaGN layers; circular padding accounts for the periodic domain.
Results: The one-step MSE is not a strong indicator of rollout quality: History 2 achieves twice the baseline’s lower loss yet performs significantly worse in rollout.The Ensemble also has lower one-step error than PDE-Refiner with more than 3 refinement steps but is almost 20 seconds behind in rollout.
Results: PDE-Refiner’s rollout MSE increases significantly less than baselines after the first 20 seconds, despite starting with similar losses.This later improvement is consistent with non-dominant high frequencies becoming influential at later rollout times.
Speed comparison: 2.04 seconds (±0.01) is the MSE baseline runtime, versus 8.67 seconds (±0.01) for PDE-Refiner with 3 refinement steps on the specified rollout benchmark.The comparison uses batch size 128 and rollout length 640∆t on an NVIDIA A100 GPU with a 24-core AMD EPYC CPU.

D.2 Parameter-dependent KS dataset · D.3 Kolmogorov 2D Flow · E Supplementary Experimental Results

The supplementary experiments extend PDE-Refiner to parameter-dependent Kuramoto–Sivashinsky data and 2D Kolmogorov flow, where it maintains long-horizon accuracy and improves over MSE baselines. Additional experiments examine alternative architectures and ablations across both PDE settings.

D.2 Parameter-dependent KS dataset: Viscosity ν is uniformly sampled from 0.5 to 1.5 for training and test trajectories in the parameter-dependent Kuramoto–Sivashinsky dataset.Training trajectories contain 140 time steps, whereas test trajectories roll out for 1140 time steps; lower viscosities generally yield more complex trajectories.
D.2 Parameter-dependent KS dataset: PDE-Refiner outperforms the MSE-trained baseline by 25-35% across viscosities.The results are reported in Table 5, with standard deviations computed over 5 seeds.
D.3 Kolmogorov 2D Flow: For 2D Kolmogorov flow, PDE-Refiner uses a modified modern U-Net with 2D convolutions and increased capacity at the highest-resolution feature level.The initial hidden size is 128, with 4 ResNet blocks at that level, while other levels match the KS model.
D.3 Kolmogorov 2D Flow: Kolmogorov-flow models predict every 16th time step, which performs similarly to shorter step sizes while enabling faster rollouts.Models use the residual Δu = u(t) − u(t − 16Δt) and are trained with 3 seeds.
D.3 Kolmogorov 2D Flow: PDE-Refiner maintains accurate Kolmogorov-flow predictions for more than 11 seconds across many trajectories.Example visualizations show stability beyond 10 seconds, with only minor errors at 10.66 seconds and similar structures still visible at 14 seconds.
E Supplementary Experimental Results: Fourier Neural Operators achieve baseline results similar to U-Nets, while PDE-Refiner and Diffusion retain gains with smaller improvements than U-Nets.The reduced gain is attributed to the noise objective in the FNO experiments.
E Supplementary Experimental Results: Supplementary experiments evaluate Fourier Neural Operators and Dilated ResNets as alternatives to the deployed U-Nets.They also ablate predicted step size, history information, and minimum noise variance on KS, and perform frequency analysis on Kolmogorov flow.

E.1 Fourier Neural Operator · E.2 Dilated ResNets

Across Fourier Neural Operators and Dilated ResNets, PDE-Refiner generally achieves the strongest long-rollout performance, while its gains depend on refinement, noise scheduling, and architecture-specific frequency modeling. These experiments also show that increasing model complexity or retained Fourier modes alone provides limited benefit.

E.1 Fourier Neural Operator: E.1 Fourier Neural Operator: The eight-layer FNO uses spectral convolutions with 1 × 1 skip convolutions and GELU activations, retaining the first 32 Fourier modes.The experiments repeat the Section 4.1 analyses to test whether the observed trends are architecture-invariant.
E.1 Fourier Neural Operator: E.1 Fourier Neural Operator: 73.6s rollout time for the MSE-trained FNO is similar to the U-Net’s 75.4s, while five-model ensembling gains 7 seconds.More history reduces rollout performance, and increasing the parameter count to 160 million does not improve results.
E.1 Fourier Neural Operator: E.1 Fourier Neural Operator: Pushforward, error-correction, error-prediction, and Sobolev losses fail to improve the baseline, whereas Markov Neural Operator regularizers provide gains limited to 3%.The Sobolev losses decrease performance, while the regularizers produce only small improvements.
E.1 Fourier Neural Operator: E.1 Fourier Neural Operator: PDE-Refiner outperforms all baselines beyond one refinement step, with gains flattening after three steps and smaller improvements than for U-Nets.FNOs have higher training losses at smaller noise values, indicating difficulty modeling high-frequency noise in the refinement objective; exponential scheduling restores one-step PDE-Refiner-level performance.
E.1 Fourier Neural Operator: E.1 Fourier Neural Operator: Increasing retained modes above 16 yields only minor MSE improvements, while removing dominant frequencies above 8 significantly reduces performance.Increasing FNO modes also has minor impact for PDE-Refiner, and higher frequencies are otherwise modeled through residual 1 × 1 convolutions.
E.2 Dilated ResNets: E.2 Dilated ResNets: Dilated convolutions expand the receptive field without reducing resolution; the evaluated model uses four seven-layer blocks with dilation factors [8].The KS setup uses channel size 256, group normalization, shift-and-scale conditioning, and pre-activations.
E.2 Dilated ResNets: E.2 Dilated ResNets: PDE-Refiner benefits considerably from pre-activations, especially during refinement, whereas increasing the approximately 22 million parameters does not help.The MSE baseline performs similarly across ResNet versions.
E.2 Dilated ResNets: E.2 Dilated ResNets: Dilated ResNets slightly outperform U-Nets for MSE training, while PDE-Refiner achieves the longest accurate rollouts and significantly outperforms all baselines.Across both MSE Training and PDE-Refiner, smaller prediction steps support longer stable rollouts, whereas very large steps substantially reduce accuracy.

E.3 Step Size Comparison

On the Kuramoto–Sivashinsky equation, smaller prediction steps keep neural PDE solvers accurate longer, while directly predicting large steps becomes impractical because models overfit and struggle to learn the dynamics. Autoregressive solvers therefore offer significant advantages, despite error accumulation over rollout.

Step size effects: Smaller step sizes keep the model accurate longer, whereas performance declines faster at very large steps as models overfit and struggle to learn PDE dynamics.The comparison covers step sizes between 0.2s and 12.8s.
Step size effects: Models with larger time steps require fewer autoregressive steps, initially limiting error accumulation, but rollout errors still increase over time as they accumulate and cause divergence.This describes the MSE loss behavior of MSE-trained models rolled out over time.
Step size effects: The 1∆t, 2∆t, and 4∆t models obtain similar performance, while the 8∆t model has considerably higher one-step MSE than eight rollouts of the 1∆t model.For larger time steps, the performance gap increases further, highlighting autoregressive solvers’ strengths.
Conclusion: Directly predicting large time steps is not practical, so autoregressive neural PDE solvers provide significant advantages.The conclusion follows from the observed degradation at increasingly large prediction steps.

E.4 History Information

For the KS equation, conditioning neural operators on more past time steps consistently worsens long-horizon rollout performance. Although longer histories improve one-step prediction by 10× and initially help for roughly five steps, their errors propagate faster and cause earlier divergence.

History Information: Using four past time steps worsens KS-equation rollouts compared with using two, confirming that more history information degrades performance.The same trend holds for models using a smaller 0.2s step size.
History Information: 10×: History information improves one-step prediction, but history-2 and history-4 models soon exceed history-1 error and diverge earlier.The rollout MSE increases rapidly after the initial one-step advantage.
History Information: ~5 steps: More history provides an initial rollout advantage, but error then propagates significantly faster through the models.Figure 20 compares MSE over a 100-second rollout and marks when average correlation falls below 0.8.

E.5 Uncertainty Estimation

PDE-Refiner estimates predictive uncertainty more accurately than input modulation and nearly as well as a full model ensemble, while requiring only one trained model. Its sample spread closely matches the actual error distribution across frequencies.

Comparison with baselines: PDE-Refiner’s uncertainty estimate outperforms input modulation but performs slightly worse than a full model ensemble.Input modulation perturbs the initial condition with small random Gaussian noise, whereas model ensembles compare predictions from multiple trained models.
Comparison with baselines: PDE-Refiner requires only one trained model, unlike model ensembles, making it advantageous when training is costly.The passage highlights large-scale experiments such as weather modeling as a particularly relevant setting.
Sample-based uncertainty: PDE-Refiner’s sample deviations closely match the average error to ground truth, indicating accurate predictive-uncertainty modeling across frequencies.Figure 22 compares the frequency spectra of 16 samples with their average difference to the ground truth and sample mean.

E.6 Frequency Analysis for 2D Kolmogorov Flow · E.7 Minimum noise variance in PDE-Refiner · E.8 Stability of Very Long Rollouts

Across 2D Kolmogorov flow and very long rollouts, PDE-Refiner improves frequency predictions through iterative refinement, remains robust to small minimum-noise-variance changes, and can be stabilized beyond 1000 autoregressive steps with error correction. Very long runs nevertheless reveal overestimated very-high-frequency content caused by accumulated noise.

E.7 Minimum noise variance in PDE-Refiner: PDE-Refiner is robust to small changes in minimum noise variance, with equal performance across the range 10^-7 to 10^-6.Higher standard deviations reduce performance because later refinement steps target low-amplitude information.
E.6 Frequency Analysis for 2D Kolmogorov Flow: PDE-Refiner improves upon the MSE-trained model across all frequencies in 2D Kolmogorov flow despite its flatter frequency spectrum.The analysis uses a 2D discrete Fourier transform averaged independently over the two channels of u(t).
E.6 Frequency Analysis for 2D Kolmogorov Flow: Each refinement step improves PDE-Refiner’s prediction, while the final step provides only minor gains and no clear high-frequency bias.The error is described as relatively uniform across frequencies, suggesting that two refinement steps would likely have been sufficient.
E.7 Minimum noise variance in PDE-Refiner: Increasing the minimum noise variance too far decreases performance because the added noise becomes too high to model the lowest-amplitude information.The reported test-set results use a hyperparameter selected on validation data, with σ2_min = 2e-7.
E.8 Stability of Very Long Rollouts: Over 1000 autoregressive prediction steps, corresponding to 800 seconds, the MSE baseline accumulates overestimated high frequencies, whereas PDE-Refiner remains accurate for more than 500 steps.PDE-Refiner later overestimates very high frequencies, likely because Gaussian noise is added iteratively.
E.8 Stability of Very Long Rollouts: Setting the highest frequencies to zero stabilizes PDE-Refiner and preserves a very accurate frequency spectrum at 800 seconds, while the MSE model still overestimates high frequencies.This correction uses known invariances following the setup applied for Figure 3.
E.8 Stability of Very Long Rollouts: The evaluated models remain stable longer than they remain accurate to ground truth, and error correction keeps PDE-Refiner stable beyond 1000 autoregressive rollout steps.Stability is defined as staying within the solution domain without generating physically unrealistic results.