SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Results. We provide the results of Figure 7 in table form in Table 5. Overall, PDE-Refiner outperforms the MSE-trained baseline by  $25 - 35\%$  across viscosities.
D.3 Kolmogorov 2D Flow
Data generation. We followed the data generation of Sun et al. [79] as detailed in the publicly released code<sup>8</sup>. For hyperparameter tuning, we additionally generate a validation set of the same size as the test data with initial seed 123. Afterward, we remove trajectories where the ground truth solver had NaN outputs, and split the trajectories into sub-sequences of 50 frames for efficient training. An epoch consists of iterating over all sub-sequences and sampling 5 random initial conditions from each. All data are stored in float32 precision.
<sup>8</sup>https://github.com/Edward-Sun/TSM-PDE/blob/main/data_generation.md
Table 6: Hyperparameter overview for the experiments on the Kolmogorov 2D flow.
Hyperparameter	Value
Input Resolution	64×64
Number of Epochs	100
Batch size	32
Optimizer	AdamW [52]
Learning rate	CosineScheduler(1e-4 → 1e-6)
Weight Decay	1e-5
Time step	0.112s / 16Δt
Output factor	0.16
Network	Modern U-Net [22]
Hidden size	[128, 128, 256, 1024]
Padding	circular
EMA Decay	0.995
Model architecture. We again use the modern U-Net [22] for PDE-Refiner and an MSE-trained baseline, where, in comparison to the model for the KS equation, we replace 1D convolutions with 2D convolutions. Due to the low input resolution, we experienced that the model lacked complexity on the highest feature resolution. Thus, we increased the initial hidden size to 128, and use 4 ResNet blocks instead of 2 on this level. All other levels remain the same as for the KS equation. This model has 157 million parameters.
The Fourier Neural Operator [49] consists of 8 layers, where each layer consists of a spectral convolution with a skip connection of a  $1 \times 1$  convolution and GELU activation [27]. We performed a hyperparameter search over the number of modes and hidden size, for which we found 32 modes with hidden size 64 to perform best. This models has 134 million parameters, roughly matching the parameter count of a U-Net. Models with larger parameter count, e.g. hidden size 128 with 32 modes, did not show any improvements.
Hyperparameters. We summarize the chosen hyperparameters in Table 6, which were selected based on the performance on the validation dataset. We train the models for 100 epochs with a batch size of 32. Due to the increased memory usage, we parallelize the model over 4 GPUs with batch size 8 each. We predict every 16th time step, which showed similar performance to models with a time step of 1, 2, 4, and 8 while being faster to roll out. All models use as objective the residual  $\Delta u = u(t) - u(t - 16\Delta t)$ , which we normalize by dividing with its training standard deviation of 0.16. Thus, we predict the next solution via  $\hat{u}(t) = u(t - 16\Delta t) + 0.16 \cdot \mathsf{N0}(\ldots)$ . Each model is trained for 3 seeds, and the standard deviation is reported in Table 1.
Results. We include example trajectories and corresponding predictions by PDE-Refiner in Figure 13. PDE-Refiner is able to maintain accurate predictions for more than 11 seconds for many trajectories.
Speed comparison. All models are run on the same hardware, namely an NVIDIA A100 GPU with 80GB memory and an 24 core AMD EPYC CPU. For the hybrid solvers, we use the public implementation in JAX [6] by Kochkov et al. [43], Sun et al. [79]. For the U-Nets, we use PyTorch 2.0 [14]. All models are compiled in their respective frameworks, and we exclude the compilation and time to load the data from the runtime. We measure the speed of each model 5 times, and report the mean and standard deviation in Section 4.3.
Figure 13: Visualizing the vorticity of three example test trajectories of the 2D Kolmogorov flow, with corresponding predictions of PDE-Refiner. PDE-Refiner remains stable for more than 10 seconds, making on minor errors at 10.66 seconds. Moreover, many structures at 14 seconds are still similar to the ground truth.
High-Correlation Rollout Times of FNOs on the Kuramoto-Sivashinsky equation
Figure 14: Experimental results of Fourier Neural Operators on the Kuramoto-Sivashinsky equation. All methods from Figure 3 are included here. FNOs achieve similar results as the U-Nets for the baselines. For PDE-Refiner and Diffusion, FNOs still outperforms the baselines, but with a smaller gain than the U-Nets due to the noise objective.
E Supplementary Experimental Results
In this section, we provide additional experimental results on the Kuramoto-Sivashinsky equation and the 2D Kolmogorov flow. Specifically, we experiment with Fourier Neural Operators and Dilated ResNets as an alternative to our deployed U-Nets. We provide ablation studies on the predicted step size, the history information, and the minimum noise variance in PDE-Refiner on the KS equation. For the Kolmogorov flow, we provide the same frequency analysis as done for the KS equation in the main paper. Finally, we investigate the stability of the neural surrogates for very long rollouts of 800 seconds.
E.1 Fourier Neural Operator
Fourier Neural Operators (FNOs) [49] are a popular alternative to U-Nets for neural operator architectures. To show that the general trend of our results in Section 4.1 are architecture-invariant, we repeat all experiments of Figure 3 with FNOs. The FNO consists of 8 layers, where each layer consists of a spectral convolution with a skip connection of a  $1 \times 1$  convolution and a GELU activation [27]. Each spectral convolution uses the first 32 modes, and we provide closer discussion on the impact of modes in Figure 15. We use a hidden size of 256, which leads to the model having about 40 million parameters, roughly matching the parameter count of the used U-Nets.
MSE Training. We show the results for all methods in Figure 14. The MSE-trained FNO baseline achieves with 73.6s a similar rollout time as the U-Net (75.4s). Again, using more history information decreases rollout performance. Giving the model more complexity by increasing the parameter count to 160 million did not show any improvement. Still, the ensemble of 5 MSE-trained models obtains a 7-second gain over the individual models, slightly outperforming the U-Nets for this case.
Alternative losses. The pushforward trick, the error correction and the error predictions again cannot improve over the baseline. While using the Sobolev norm losses decrease performance also for FNOs, using the regularizers of the Markov Neural Operator is able to provide small gains. This is in line with the experiments of Li et al. [50], in which the MNO was originally proposed for Fourier Neural Operators. Still, the gain is limited to  $3\%$ .
PDE-Refiner. With FNOs, PDE-Refiner again outperforms all baselines when using more than 1 refinement step. The gains again flatten for more than 3 steps. However, in comparisons to the U-Nets with up to 98.5s accurate rollout time, the performance increase is significantly smaller. In general, we find that FNOs obtain higher training losses for smaller noise values than U-Nets, indicating the modeling of high-frequent noise in PDE-Refiner's refinement objective to be the main issue. U-Nets are more flexible in that regard, since they use spatial convolutions. Still, the results show that PDE-Refiner is applicable to a multitude of neural operator architectures.
Diffusion ablations. Confirming the issue of the noise objective for FNOs, the diffusion models with standard cosine scheduling obtain slightly worse results than the baseline. Using our exponential noise scheduler again improves performance to the level of the one-step PDE-Refiner.
High-Correlation Rollout Times over Number of Fourier Modes
Figure 15: Investigating the impact of the choosing the number of modes in FNOs. Similar to our analysis on the resolution in the U-Nets (Figure 3), we only see minor improvements of using higher frequencies above 16 in the MSE training. Removing dominant frequencies above 8 significantly decreases performance. Similarly, increasing the modes of FNOs in PDE-Refiner has minor impact.
High-Correlation Rollout Times of Dilated ResNets on the Kuramoto-Sivashinsky equation
Figure 16: Experimental results of Dilated ResNets on the Kuramoto-Sivashinsky equation. All methods from Figure 3 are included here, with single seeds provided for all methods except the MSE baseline and PDE-Refiner, for which three seeds are shown. The results are overall similar to U-Nets, slightly outperforming the U-Net overall. Yet, PDE-Refiner again outperforms all baselines with a significant margin.
Number of Fourier Modes. A hyperparameter in Fourier Neural Operators is the number of Fourier modes that are considered in the spectral convolutions. Any higher frequency is ignored and must be modeled via the residual  $1 \times 1$  convolutions. To investigate the impact of the number of Fourier modes, we repeat the baseline experiments of MSE-trained FNOs with 8, 16, 32, 64, and 128 modes in Figure 15. To ensure a fair comparison, we adjust the hidden size to maintain equal number of parameters across models. In general, we find that the high-correlation time is relatively stable for 32 to 128 modes. Using 16 modes slightly decreases performance, while limiting the layers to 8 modes results in significantly worse rollouts. This is also in line with our input resolution analysis of Figure 5, where the MSE-trained baseline does not improve for high resolutions. Similarly, we also apply a 64 mode FNOs for PDE-Refiner. Again, the performance does not increase for higher number of modes.
E.2 Dilated ResNets
As another strong neural operator, Stachenfeld et al. [78] found Dilated ResNets to perform well on a variety of fluid dynamics. Instead of reducing resolution as in U-Nets, Dilated ResNets make use of dilated convolutions [92] to increase the receptive field of the network and take into account the whole spatial dimension. A Dilated ResNets consists of a ResNet [25] with 4 blocks, each containing 7 convolutional layers with dilation factors  $[1, 2, 4, 8, 4, 2, 1]$ . Since the convolutions do not use any stride, all operations are performed on the original spatial resolution, increasing computational cost over e.g. a U-Net with similar parameter count.
For our setup on the KS dataset, we use a Dilated ResNet with channel size 256, group normalization between convolutions and a shift-and-scale conditioning as in the U-Nets. Additionally, we change the default ResNet architecture to using pre-activations [26], allowing for a better gradient flow and
High-Correlation Rollout Times over Predicted Step Size
Figure 17: Comparing the accurate rollout times over the step size at which the neural operator predicts. This is a multiple of the time step  $\Delta t$  used for data generation (for KS on average 0.2s). For both the MSE Training and PDE-Refiner, lower step size provides longer stable rollouts, where very large time steps show a significant loss in accuracy. This motivates the need for autoregressive neural PDE solvers over direct, long-horizon predictions.
operations on the input directly. While for the MSE baseline, both ResNet version performs equally, we see a considerable improvement for PDE-Refiner, in particular for the refinement steps where early activations can augment the input noise. Overall, the model has around 22 million parameters. More parameters by increasing the channel size did not show to benefit the model.
We show all results in Figure 16. Due to the experiments being computationally expensive, we show most results for a single seed only here, but the trends are generally constant and have not shown to be significantly impacted by standard deviations for other neural operators. Dilated ResNet slightly outperform U-Nets for the MSE baseline, but generally have the same trends across baselines. Furthermore, PDE-Refiner is again able to obtain the longest accurate rollouts, similar to the U-Nets. This verifies the strength of PDE-Refiner and its applicability across various state-of-the-art neural operator architectures.
E.3 Step Size Comparison
A key advantage of Neural PDE solvers is their flexibility to be applied to various step sizes of the PDEs. The larger the step size is, the faster the solver will be. At the same time, larger step sizes may be harder to predict. To compare the effect of error propagation in an autoregressive solver with training a model to predict large time steps, we repeat the baseline experiments of the U-Net neural operator on the KS equation with different step sizes. The default step size that was used in Figure 3 is 4-times the original solver step, being on average 0.8s. For any step size below 2s, we model the residual objective  $\Delta u = u(t) - u(t - \Delta t)$ , which we found to generally work better in this range. For any step size above, we directly predict the solution  $u(t)$ .
High-correlation time. We plot the results step sizes between 0.2s and 12.8s in Figure 17. We find that the smaller the step size, the longer the model remains accurate. The performance also decreases faster for very large time steps. This is because the models start to overfit on the training data and have difficulties learning the actual dynamics of the PDE. Meanwhile, very small time steps do not suffer from autoregressive error propagation any more than slightly larger time steps, while generalizing well. This highlights again the strength of autoregressive neural PDE solvers. We confirm this trend by training PDE-Refiner with different step sizes while using 3 refinement steps. We again find that smaller time steps achieve higher performance, and we obtain worse rollout times for larger time steps.
MSE loss over rollout. To further gain insights of the impact of different step sizes, we plot in Figure 18 the MSE loss to the ground truth when rolling out the MSE-trained models over time. Models with larger time steps require fewer autoregressive steps to predict long-term into the future, preventing any autoregressive error accumulation for the first step. Intuitively, the error increases over time for all models, since the errors accumulate over time and cause the model to diverge. The models with step sizes 0.2s, 0.4s and 0.8s all achieve very similar losses across the whole time horizon. This motivates our choice for 0.8s as default time step, since it provides a 4 times speedup in comparison to the 0.2s model. Meanwhile, already a model trained with step size 1.6s performs considerable worse in its one-step prediction than a model with step size 0.2s rolled out 8 times. The gap increases
Figure 18: Visualizing the MSE error of MSE-trained models with varying step sizes over the rollout. The models with a step size of  $1\Delta t$ ,  $2\Delta t$ , and  $4\Delta t$  all obtain similar performance. For  $8\Delta t$ , the one-step MSE loss is already considerably higher than, e.g. rolling out the step size  $1\Delta t$  model 8 times. For larger time steps, this gap increases further, again highlighting the strengths of autoregressive solvers.
Figure 19: Investigating the impact of using more history / past time steps in the neural operators, i.e.,  $\hat{u}(t) = \mathbb{N}0(u(t - \Delta t), u(t - 2\Delta t), \ldots)$ , for  $\Delta t = 0.8$  and  $\Delta t = 0.2$ . Longer histories decrease the model's accurate rollout time. This drop in performance is even more significant for smaller time steps.
further the larger the time step becomes. Therefore, directly predicting large time steps in neural PDE solvers is not practical and autoregressive solvers provide significant advantages.
E.4 History Information
In our experiments on the KS equation, we have observed that using more history information as input decreases the rollout performance. Specifically, we have used a neural operator that took as input the past two time steps,  $u(t - \Delta t)$  and  $u(t - 2\Delta t)$ . To confirm this trend, we repeat the experiments with a longer history of 4 past time steps and for models with a smaller step size of 0.2s in Figure 19. Again, we find that the more history information we use as input, the worse the rollouts become. Furthermore, the impact becomes larger for small time steps, indicating that the autoregressive error propagation becomes a larger issue when using history information. The problem arising is that the difference between the inputs  $u(t - \Delta t) - u(t - 2\Delta t)$  is highly correlated with the model's target  $\Delta u(t)$ , the residual of the next time step. The smaller the time step, the larger the correlation. This leads the neural operator to focus on modeling the second-order difference  $\Delta u(t) - \Delta u(t - 2\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 also confirm this exponential increase of error by plotting the MSE error over rollouts in Figure 20. While the history information improves the one-step prediction by a factor of 10, the error of the history 2 and 4 models quickly surpasses the error of the history 1 model. After that, the error of the models continue to increase quickly, leading to an earlier divergence.
Figure 20: Comparing models conditioned on different number of past time steps on their MSE loss over rollouts. Note the log-scale on the y-axis. The markers indicate the time when the average correlation of the respective model drops below 0.8. The left plot shows a zoomed-in version of the first 4 seconds of the whole 100 second rollout on the right. While using more history information gives an advantage for the first  $\sim 5$  steps, the error propagates significantly faster through the models. This leads to a significantly higher loss over rollout.
Table 7: Comparing the uncertainty estimate of PDE-Refiner to Input Modulation [5, 71] and Model Ensemble [45, 71] on the MSE-trained models. The metrics show the correlation between the estimated and actual accurate rollout time in terms of the  $R^2$  coefficient of determination and the Pearson correlation. PDE-Refiner provides more accurate uncertainty estimates than Input Modulation while being more efficient than an Model Ensemble.
Method	R2 coefficient	Pearson correlation
PDE-Refiner	0.857 ± 0.027	0.934 ± 0.014
Input Modulation [5, 71]	0.820 ± 0.081	0.912 ± 0.021
Model Ensemble [45, 71]	0.887 ± 0.012	0.965 ± 0.007
(a) PDE-Refiner (Ours)
(b) Input Modulation
Figure 21: Qualitative comparison between the uncertainty estimates of PDE-Refiner, Input Modulation, and the Model Ensemble. Both PDE-Refiner and the Model Ensemble achieve an accurate match between the estimated and ground truth rollout times.
(c) Model Ensemble
E.5 Uncertainty Estimation
We extend our discussion on the uncertainty estimation of Section 4.1 by comparing PDE-Refiner to two common baselines for uncertainty estimation of temporal forecasting: Input Modulation [5, 71] and Model Ensemble [45, 71]. Input Modulation adds small random Gaussian noise to the initial condition  $u(0)$ , and rolls out the model on several samples. Similar to PDE-Refiner, one can determine the uncertainty by measuring the cross-correlation between the rollouts. A Model Ensemble compares the predicted trajectories of several independently trained models. For the case here, we use 4 trained models. For both baselines, we estimate the uncertainty of MSE-trained models as usually applied.
We evaluate the  $R^2$  coefficient of determination and the Pearson correlation between the estimated stable rollout times and the ground truth rollout times in Table 7. We additionally show qualitative
Figure 22: Investigating the spread of samples of PDE-Refiner. The left plot shows the frequency spectrum of 16 samples (each line represents a different sample), with the right plot showing the average difference to the ground truth and to the mean of the samples. The deviation of the samples closely matches the average error, showing that PDE-Refiner adapts its samples to the learned error over frequencies.
Frequency Spectrum on the Kolmogorov Flow
Figure 23: Frequency spectrum on the Kolmogorov Flow. The two plots show the two channels of the Kolmogorov flow. Since the data has a much more uniform support over frequencies than the KS equation, both the MSE-trained model and PDE-Refiner model the ground truth very accurately. Thus, the Ground Truth (blue), PDE-Refiner's prediction (orange) and the MSE-trained prediction (red) overlap in both plots. Plotting the error reveals that PDE-Refiner provides small gains across all frequencies.
results in Figure 21. PDE-Refiner's uncertainty estimate outperforms the Input Modulation approach, showing that Gaussian noise does not fully capture the uncertainty distribution. While performing slightly worse than using a full Model Ensemble, PDE-Refiner has the major advantage that it only needs to be trained, which is particularly relevant in large-scale experiments like weather modeling where training a model can be very costly.
To investigate the improvement of PDE-Refiner over Input Modulation, we plot the standard deviation over samples in PDE-Refiner in Figure 22. The samples of PDE-Refiner closely differs in the same distribution as the actual loss to the ground truth, showing that PDE-Refiner accurately models its predictive uncertainty.
E.6 Frequency Analysis for 2D Kolmogorov Flow
We repeat the frequency analysis that we have performed on the KS equation in the main paper, e.g. Figure 4, on the Kolmogorov dataset here. Note that we apply a 2D Discrete Fourier Transform and show the average frequency spectrum. We perform this over the two channels of  $u(t)$  independently. Figure 23 shows the frequency spectrum for the ground truth data, as well as the predictions of PDE-Refiner and the MSE-trained U-Net. In contrast to the KS equation, the spectrum is much flatter, having an amplitude of still almost 1 at wavenumber 32. In comparison, the KS equation has a more than 10 times as small amplitude for this wavenumber. Further, since the resolution is only  $64 \times 64$ , higher modes cannot be modeled, which, as seen on the KS equation, would increase the benefit of PDE-Refiner. This leads to both PDE-Refiner and the MSE-trained baseline to model all frequencies accurately. The slightly higher loss for higher frequencies on channel 0 is likely due to missing high
Frequency Spectrum of Intermediate Samples on the Kolmogorov Flow
Figure 24: Frequency spectrum of intermediate samples in the refinement process of PDE-Refiner, similar to Figure 4 for the KS equation. The refinement process improves the prediction of the model step-by-step. For the last refinement step, we actually see minor improvements for the lowest frequencies of channel 0. However, due to flatter frequency spectrum, the high frequencies do not improve as much as on the KS equation.
Figure 25: Plotting performance of PDE-Refiner over different values of the minimum noise variance  $\sigma_{\mathrm{min}}^2$ . Each PDE-Refiner is robust to small changes of  $\sigma_{\mathrm{min}}^2$ , showing an equal performance in the range of  $[10^{-7}, 10^{-6}]$ . Higher standard deviations start to decrease the performance, confirming our analysis of later refinement steps focusing on low-amplitude information. For the experiments in Section 4.1, we have selected  $\sigma_{\mathrm{min}}^2 = 2\mathrm{e} - 7$  based on the validation dataset.
frequency information, i.e., larger resolution, that would be needed to estimate the frequencies more accurately. Still, we find that PDE-Refiner improves upon the MSE-trained model on all frequencies.
In Figure 24, we additionally plot the predictions of PDE-Refiner at different refinement steps. Similar to the KS equation, PDE-Refiner improves its prediction step by step. However, it is apparent that no clear bias towards the high frequencies occur in the last time step, since the error is rather uniform across all frequencies. Finally, the last refinement step only provides minor gains, indicating that PDE-Refiner with 2 refinement steps would have likely been sufficient.
E.7 Minimum noise variance in PDE-Refiner