SlowGuess/NeurIPS_2023_OCR huggingface dataset:
of  $640\Delta t$ . As the network architecture, we use the modern U-Net of Gupta et al. [22] with hidden size 64 and 3 downsampling layers. U-Nets have demonstrated strong performance in both neural PDE solving [22, 56, 80] and diffusion modeling [29, 30, 61], making it an ideal candidate for PDE-Refiner. A common alternative is the Fourier Neural Operator (FNO) [49]. Since FNO layers cut away high frequencies, we find them to perform suboptimally on predicting the residual noise in PDE-Refiner and DDPMs. Yet, our detailed study with FNOs in Appendix E.1 shows that even here PDE-Refiner offers significant performance gains. Finally, we also evaluate on Dilated ResNets [78] in Appendix E.2, showing very similar results to the U-Net. Since neural surrogates can operate on larger time steps, we directly predict the solution at every 4th time step. In other words, to predict  $u(t)$ , each model takes as input the previous time step  $u(t - 4\Delta t)$  and the trajectory parameters  $L$  and  $\Delta t$ . Thereby, the models predict the residual between time steps  $\Delta u(t) = u(t) - u(t - 4\Delta t)$  instead of  $u(t)$  directly, which has shown superior performance at this timescale [50]. Ablations on the time step size can be found in Appendix E.3. As evaluation criteria, we report the model rollouts' high-correlation time [43, 79]. For this, we autoregressively rollout the models on the test set and measure the Pearson correlation between the ground truth and the prediction. We then report the time when the average correlation drops below 0.8 and 0.9, respectively, to quantify the time horizon for which the predicted rollouts remain accurate. We investigate other evaluation criteria such as mean-squared error and varying threshold values in Appendix D, leading to the same conclusions.
MSE Training. We compare PDE-Refiner to three groups of baselines in Figure 3. The first group are models trained with the one-step MSE error, i.e., predicting  $\Delta u(t)$  from  $u(t - 4\Delta t)$ . The baseline U-Net obtains a high-correlation rollout time of 75 seconds, which corresponds to 94 autoregressive steps. We find that incorporating more history information as input, i.e.  $u(t - 4\Delta t)$  and  $u(t - 8\Delta t)$ , improves the one-step prediction but worsens rollout performance. The problem arising is that the difference between the inputs  $u(t - 4\Delta t) - u(t - 8\Delta t)$  is highly correlated with the model's target  $\Delta u(t)$ , the residual of the next time step. This leads the neural operator to focus on modeling the second-order difference  $\Delta u(t) - \Delta u(t - 4\Delta t)$ . As observed in classical solvers [36], using higher-order differences within an explicit autoregressive scheme is known to deteriorate the rollout stability and introduce exponentially increasing errors over time. We include further analysis of this behavior in Appendix E.4. Finally, we verify that PDE-Refiner's benefit is not just because of having an increased model complexity by training a model with 4 times the parameter count and observe a performance increase performance by only  $5\%$ . Similarly, averaging the predictions of an ensemble of 5 MSE-trained models cannot exceed 80 seconds of accurate rollouts.
Alternative losses. The second baseline group includes alternative losses and post-processing steps proposed by previous work to improve rollout stability. The pushforward trick [9] rolls out the model during training and randomly replaces ground truth inputs with model predictions. This trick does not improve performance in our setting, confirming previous results [8]. While addressing potential input distribution shift, the pushforward trick cannot learn to include the low-amplitude information for accurate long-term predictions, as no gradients are backpropagated through the predicted input for stability reasons. Focusing more on high-frequency information, the Sobolev norm loss [50] maps the
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.  $|\mathrm{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.
prediction error into the frequency domain and weighs all frequencies equally for  $k = 0$  and higher frequencies more for  $k = 1$ . However, focusing on high-frequency information leads to a decreased one-step prediction accuracy for the high-amplitude frequencies, such that the rollout time shortens. The Markov Neural Operator (MNO) [50] additionally encourages dissipativity via regularization, but does not improve over the common Sobolev norm losses. Inspired by McGreivy et al. [58], we report the rollout time when we correct the predictions of the MSE models for known invariances in the equation. We ensure mass conservation by zeroing the mean and set any frequency above 60 to 0, as their amplitude is below f1oat32 precision (see Appendix D.1). This does not improve over the original MSE baselines, showing that the problem is not just an overestimate of the high frequencies, but the accurate modeling of a broader spectrum of frequencies. Finally, to highlight the advantages of the denoising process in PDE-Refiner, we train a second model to predict another MSE-trained model's errors (Error Prediction). This model quickly overfits on the training dataset and cannot provide gains for unseen trajectories, since it again focuses on the same high-amplitude frequencies.
PDE-Refiner - Number of refinement steps. Figure 3 shows that PDE-Refiner significantly outperforms the baselines and reaches almost 100 seconds of stable rollout. Thereby, we have a tradeoff between number of refinement steps and performance. When training PDE-Refiner with 1 to 8 refinement steps, we see that the performance improves with more refinement steps, but more steps require more model calls and thus slows down the solver. However, already using a single refinement step improves the rollout performance by  $20\%$  over the best baseline, and the gains start to flatten at 3 to 4 steps. Thus, for the remainder of the analysis, we will focus on using 3 refinement steps.
Diffusion Ablations. In an ablation study of PDE-Refiner, we evaluate a standard denoising diffusion model [30] that we condition on the previous time step  $u(t - 4\Delta t)$ . When using a common cosine noise schedule [61], the model performs similar to the MSE baselines. However, with our exponential noise decrease and lower minimum noise level, the diffusion models improve by more than 10 seconds. Using the prediction objective of PDE-Refiner gains yet another performance improvement while reducing the number of sampling steps significantly. Furthermore, to investigate the probabilistic nature of PDE-Refiner, we check whether it samples single modes under potentially multi-modal uncertainty. For this, we average 16 samples at each rollout time step (3 steps - Mean in Figure 3) and find slight performance improvements, indicating that PDE-Refiner mostly predicts single modes.
Modeling the Frequency Spectrum. We analyse the performance difference between the MSE training and PDE-Refiner by comparing their one-step prediction in the frequency domain in Figure 4. Similar to the MSE training, the initial prediction of PDE-Refiner has a close-to uniform error pattern across frequencies. While the first refinement step shows an improvement across all frequencies, refinement steps 2 and 3 focus on the low-amplitude frequencies and ignore higher amplitude errors. This can be seen by the error for wavenumber 7, i.e., the frequency with the highest input amplitude, not improving beyond the first refinement step. Moreover, the MSE training obtains almost the identical error rate for this frequency, emphasizing the importance of low-amplitude information. For all other frequencies, PDE-Refiner obtains a much lower loss, showing its improved accuracy on low-amplitude information over the MSE training. We highlight that PDE-Refiner does not only improve the high frequencies, but also the lowest frequencies (wavenumber 1-6) with low amplitude.
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.
Input Resolution. We demonstrate that capturing high-frequency information is crucial for PDE-Refiner's performance gains over the MSE baselines by training both models on datasets of subsampled spatial resolution. With lower resolution, fewer frequencies are present in the data and can be modeled. As seen in Figure 5, MSE models achieve similar rollout times for resolutions between 32 and 256, emphasizing its inability to model high-frequency information. At a resolution of 32, PDE-Refiner achieves similar performance to the MSE baseline due to the missing high-frequency information. However, as resolution increases, PDE-Refiner significantly outperforms the baseline, showcasing its utilization of high-frequency information.
Spectral data augmentation. A pleasant side effect of PDE-Refiner is data augmentation, which is induced by adding varying Gaussian noise  $\sigma_k\epsilon^k,\epsilon^k\sim \mathcal{N}(0,1)$  at different stages  $k$  of the refinement process. Effectively, data augmentation is achieved by randomly distorting the input at different scales, and forcing the model to recover the underlying structure. This gives an ever-changing input and objective, forces the model to fit different parts of the spectrum, and thus making it more difficult for the model to overfit. Compared to previous works such as Lie Point Symmetry data augmentation [8] or general covariance and random coordinate transformations [15], the data augmentation in PDE-Refiner is purely achieved by adding noise, and thus very simple and applicable to any PDE. While we leave more rigorous testing of PDE-Refiner induced data augmentation for future work, we show results for the low training data regime in Figure 5. When training PDE-Refiner and the MSE baseline on  $10\%$ ,  $20\%$  and  $50\%$  of the training trajectories, PDE-Refiner consistently outperforms the baseline in all settings. Moreover, with only  $10\%$  of the training data, PDE-Refiner performs on par to the MSE model at  $100\%$ . Finally, the relative improvement of PDE-Refiner increases to  $50\%$  for this low data regime, showing its objective acting as data augmentation and benefiting generalization.
Uncertainty estimation. When applying neural PDE solvers in practice, knowing how long the predicted trajectories remain accurate is crucial. To estimate PDE-Refiner's predictive uncertainty, we sample 32 rollouts for each test trajectory by generating different Gaussian noise during the refinement process. We compute the time when the samples diverge from one another, i.e. their cross-correlation goes below 0.8, and investigate whether this can be used to accurately estimate how long the model's rollouts remain close to the ground truth. Figure 6 shows that the cross-correlation time between samples closely aligns with the time over which the rollout remains accurate, leading to a  $R^2$  coefficient of 0.86 between the two times. Furthermore, the prediction for how long the rollout remains accurate depends strongly on the individual trajectory - PDE-Refiner reliably identifies trajectories that are easy or challenging to roll out from. In Appendix E.5, we compare PDE-Refiner's uncertainty estimate to two other
Figure 6: Uncertainty estimate of PDE-Refiner. Each point represents the estimated correlation time via sample cross-correlation (x-axis) and the ground truth time (y-axis) for a test trajectory.
common approaches. PDE-Refiner provides more accurate estimates than input modulation [5, 71], while only requiring one trained model compared to a model ensemble [45, 71].
Figure 7: Visualizing the parameter-dependent KS equation. Left: Frequency spectrum of ground truth trajectories over viscosities. Right: Accurate rollout time over viscosities (error bars neglected if smaller than marker size). PDE-Refiner obtains improvements across all viscosities.
4.2 Parameter-dependent KS equation
So far, we have focused on the KS equation with a viscosity term of  $\nu = 1$ . Under varying values of  $\nu$ , the Kuramoto-Sivashinsky equation has been shown to develop diverse behaviors and fixed points [35, 40, 74]. This offers an ideal benchmark for evaluating neural surrogate methods on a diverse set of frequency spectra. We generate 4096 training and 512 test trajectories with the same data generation process as before, except that for each trajectory, we sample  $\nu$  uniformly between 0.5 and 1.5. This results in the spatial frequency spectrum of Figure 7, where high frequencies are damped for larger viscosities but amplified for lower viscosities. Thus, an optimal neural PDE solver for this dataset needs to work well across a variety of frequency spectra. We keep the remaining experimental setup identical to Section 4.1, and add the viscosity  $\nu$  to the conditioning set of the neural operators.
We compare PDE-Refiner to an MSE-trained model by plotting the stable rollout time over viscosities in Figure 7. Each marker represents between for trajectories in  $[\nu -0.1,\nu +0.1]$ . PDE-Refiner is able to get a consistent significant improvement over the MSE model across viscosities, verifying that PDE-Refiner works across various frequency spectra and adapts to the given underlying data. Furthermore, both models achieve similar performance to their unconditional counterpart for  $\nu = 1.0$ . This again highlights the strength of the U-Net architecture and baselines we consider here.
4.3 Kolmogorov 2D Flow
Simulated data. As another common fluid-dynamics benchmark, we apply PDE-Refiner to the 2D Kolmogorov flow, a variant of the incompressible Navier-Stokes flow. The PDE is defined as:
$$
\partial_ {t} \mathbf {u} + \nabla \cdot (\mathbf {u} \otimes \mathbf {u}) = \nu \nabla^ {2} \mathbf {u} - \frac {1}{\rho} \nabla p + \mathbf {f} \tag {6}
$$
where  $\mathbf{u}:[0,T]\times \mathcal{X}\to \mathbb{R}^2$  is the solution,  $\otimes$  the tensor product,  $\nu$  the kinematic viscosity,  $\rho$  the fluid density,  $p$  the pressure field, and, finally,  $\mathbf{f}$  the external forcing. Following previous work [43, 79], we set the forcing to  $\mathbf{f} = \sin (4y)\hat{\mathbf{x}} -0.1\mathbf{u}$ , the density  $\rho = 1$ , and viscosity  $\nu = 0.001$ , which corresponds to a Reynolds number of 1000. The ground truth data is generated using a finite volume-based direct numerical simulation (DNS) method [43, 57] with a time step of  $\Delta t = 7.0125\times 10^{-3}$  and resolution of  $2048\times 2048$ , and afterward downscaled to  $64\times 64$ . To align our experiments with previous results, we use the same dataset of 128 trajectories for training and 16 trajectories for testing as Sun et al. [79].
Experimental setup. We employ a modern U-Net [22] as the neural operator backbone. Due to the lower input resolution, we set  $\sigma_{\mathrm{min}}^2 = 10^{-3}$  and use 3 refinement steps in PDE-Refiner. For efficiency, we predict 16 steps  $(16\Delta t)$  into the future and use the difference  $\Delta \mathbf{u} = \mathbf{u}(t) - \mathbf{u}(t - 16\Delta t)$  as the output target. Besides the MSE objective, we compare PDE-Refiner with FNOs [50], classical PDE solvers (i.e., DNS) on different resolutions, and state-of-the-art hybrid machine learning solvers [43, 79], which estimate the convective flux  $\mathbf{u}\otimes \mathbf{u}$  via neural networks. Learned Interpolation (LI) [43] takes the previous solution  $\mathbf{u}(t - \Delta t)$  as input to predict  $\mathbf{u}(t)$ , similar to PDE-Refiner. In contrast, the Temporal Stencil Method (TSM) Sun et al. [79] combines information from multiple previous time steps using HiPPO features [19, 20]. We also compare PDE-Refiner to a Learned Correction model (LC) [43, 83], which corrects the outputs of a classical solver with neural networks. For evaluation, we roll out the models on the 16 test trajectories and determine the Pearson correlation with the
ground truth in terms of the scalar vorticity field  $\omega = \partial_{x}u_{y} - \partial_{y}u_{x}$ . Following previous work [79], we report in Table 1 the time until which the average correlation across trajectories falls below 0.8.
Results. Similar to previous work [22, 55], we find that modern U-Nets outperform FNOs on the 2D domain for long rollouts. Our MSE-trained U-Net already surpasses all classical and hybrid PDE solvers. This result highlights the strength of our baselines, and improving upon those poses a significant challenge. Nonetheless, PDE-Refiner manages to provide a substantial gain in performance, remaining accurate  $32\%$  longer than the best single-input hybrid method and  $10\%$  longer than the best multi-input hybrid methods and MSE model. We reproduce the frequency plots of Figure 4 for this dataset in Appendix E.6. The plots exhibit a similar behavior of both models. Compared to the KS equation, the Kolmogorov flow has a shorter (due to the resolution) and flatter spatial frequency spectrum. This accounts for the smaller relative gain of PDE-Refiner on the MSE baseline here.
Speed comparison. We evaluate the speed of the rollout generation for the test set (16 trajectories of 20 seconds) of three best solvers on an NVIDIA A100 GPU. The MSE U-Net generates the trajectories in 4.04 seconds  $(\pm 0.01)$ , with PDE-Refiner taking 4 times longer  $(16.53 \pm 0.04$  seconds) due to four model calls per step. With that, PDE-Refiner is still faster than the best hybrid solver, TSM, which needs 20.25 seconds
$(\pm 0.05)$ . In comparison to the ground truth solver at resolution  $2048 \times 2048$  with 31 minute generation time on GPU, all surrogates provide a significant speedup.
Table 1: Duration of high correlation (>0.8) on the 2D Kolmogorov flow. Results for classical PDE solvers and hybrid methods taken from Sun et al. [79].
Method	Corr. > 0.8 time
Classical PDE Solvers
DNS - 64 × 64	2.805
DNS - 128 × 128	3.983
DNS - 256 × 256	5.386
DNS - 512 × 512	6.788
DNS - 1024 × 1024	8.752
Hybrid Methods
LC [43, 83] - CNN	6.900
LC [43, 83] - FNO	7.630
LI [43] - CNN	7.910
TSM [79] - FNO	7.798
TSM [79] - CNN	8.359
TSM [79] - HiPPO	9.481
ML Surrogates
MSE training - FNO	6.451 ± 0.105
MSE training - U-Net	9.663 ± 0.117
PDE-Refiner - U-Net	10.659 ± 0.092
In this paper, we conduct a large-scale analysis of temporal rollout strategies for neural PDE solvers, identifying that the neglect of low-amplitude information often limits accurate rollout times. To address this issue, we introduce PDE-Refiner, which employs an iterative refinement process to accurately model all frequency components. This approach remains considerably longer accurate during rollouts on three fluid dynamic datasets, effectively overcoming the common pitfall.
Limitations. The primary limitation of PDE-Refiner is its increased computation time per prediction. Although still faster than hybrid and classical solvers, future work could investigate reducing compute for early refinement steps, or applying distillation and enhanced samplers to accelerate refinement, as seen in diffusion models [3, 38, 70, 88]. Another challenge is the smaller gain of PDE-Refiner with FNOs due to the modeling of high-frequency noise, which thus presents an interesting avenue for future work. Further architectures like Transformers [13, 84] can be explored too, having been shown to also suffer from spatial frequency biases for PDEs [11]. Additionally, most of our study focused on datasets where the test trajectories come from a similar domain as the training. Evaluating the effects on rollout in inter- and extrapolation regimes, e.g. on the viscosity of the KS dataset, is left for future work. Lastly, we have only investigated additive Gaussian noise. Recent blurring diffusion models [32, 47] focus on different spatial frequencies over the sampling process, making them a potentially suitable option for PDE solving as well.
[91] Yasin Yazici, Chuan-Sheng Foo, Stefan Winkler, Kim-Hui Yap, Georgios Piliouras, and Vijay Chandrasekhar. 2019. The Unusual Effectiveness of Averaging in GAN Training. In International Conference on Learning Representations.
[92] Fisher Yu and Vladlen Koltun. 2016. Multi-Scale Context Aggregation by Dilated Convolutions. In International Conference on Learning Representations.
SUPPLEMENTARY MATERIAL PDE-REFINER: ACHIEVING ACCURATE LONG ROLLOUTS WITH NEURAL PDE SOLVERS
A Broader Impact 18
B Reproducibility Statement 18
C PDE-Refiner - Pseudocode 19
D Experimental details 21
D.1 Kuramoto-Sivashinsky 1D dataset 21
D.2 Parameter-dependent KS dataset 25
D.3 Kolmogorov 2D Flow 26
E Supplementary Experimental Results 29
E.1 Fourier Neural Operator 29
E.2 Dilated ResNets 30
E.3 Step Size Comparison 31
E.4 History Information 32
E.5 Uncertainty Estimation 33
E.6 Frequency Analysis for 2D Kolmogorov Flow 34
E.7 Minimum noise variance in PDE-Refiner 35
E.8 Stability of Very Long Rollouts 36
Neural PDE solvers hold significant potential for offering computationally cheaper approaches to modeling a wide range of natural phenomena than classical solvers. As a result, PDE surrogates could potentially contribute to advancements in various research fields, particularly within the natural sciences, such as fluid dynamics and weather modeling. Further, reducing the compute needed for simulations may reduce the carbon footprint of research institutes and industries that rely on such models. Our proposed method, PDE-Refiner, can thereby help in improving the accuracy of these neural solvers, particularly for long-horizon predictions, making their application more viable.
However, it is crucial to note that reliance on simulations necessitates rigorous cross-checks and continuous monitoring. This is particularly true for neural surrogates, which may have been trained on simulations themselves and could introduce additional errors when applied to data outside its original training distribution. Hence, it is crucial for the underlying assumptions and limitations of these surrogates to be well-understood in applications.
B Reproducibility Statement
To ensure reproducibility, we publish our code at https://github.com/microsoft/pdearena. We report the used model architectures, hyperparameters, and dataset properties in detail in Section 4 and Appendix D. We additionally include pseudocode for our proposed method, PDE-Refiner, in Appendix C. All experiments on the KS datasets have been repeated for five seeds, and three seeds have been used for the Kolmogorov Flow dataset. Plots and tables with quantitative results show the standard deviation across these seeds.
As existing software assets, we base our implementation on the PDE-Arena [22], which implements a Python-based training framework for neural PDE solvers in PyTorch [62] and PyTorch Lightning [14]. For the diffusion models, we use the library diffusers [65]. We use Matplotlib [34] for plotting and NumPy [89] for data handling. For data generation, we use scipy [85] in the public code of Brandstetter et al. [8] for the KS equation, and JAX [6] in the public code of Kochkov et al. [43], Sun et al. [79] for the 2D Kolmogorov Flow dataset. The usage of these assets is further described in Appendix D.