PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers
Authors: Phillip Lippe
Bastiaan S. Veeling
Paris Perdikaris
Richard E Turner
Johannes Brandstetter
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Neural PDE Solvers
Neural Operators
Temporal Stability
Long-Horizon Modeling
Autoregressive Forecasting
TL;DR: PDE-Refiner enables neural operator training for accurate and stable predictions over long time horizons.

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.

PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers
Phillip Lippe
Bastiaan S. Veeling
Microsoft Research A14Science' phillip www.ppgooppoolllocco email. com
Microsoft Research A14Science
Paris Perdikaris Microsoft Research AI4Science
Richard E. Turner Microsoft Research A14Science
Johannes Brandstetter Microsoft Research Al4Science brandstetter@ll. jku at
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 their ability provide accurate, stable predictions over long time horizons, which 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 infor- mation, often associated with high frequencies in PDE solutions, the primary pit- fall 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 multistep refinement process. We validate marks of complex fluid dynamics, demonstrating stable and accurate rollouts that consistently hybrid outperform state-of-the-ar art models, including neural, numerical, and architectures. We further demonstrate that PDE- Refiner greatly enhances data efficiency, since denoising objective implicitly induces a novel form of spectral data augmentation. Finally, PDE-Refiner's connection to dif- fusion models enables an accurate and efficient assessment of the model's predic- tive uncertainty, allowing estimate when the surrogate becomes inaccurate.
on challenging bench-
Introduction
In recent years, mostly due to rapidly growing interest in modeling partial differential equations (PDEs), deep neural network based PDE surrogates have gained significant momentum as a more computationally efficient solution methodology [10 81). Recent approaches be broadly classified into three categories: (i) neural approaches that approximate the solution function of the underlying
PDE [24 66); (ii) hybrid approaches, where neural networks either augment numerical solvers or
replace parts them 1 2 18, 33 43 [9]; (iii) neural approaches in which the learned evolution
operator maps the current state future state of the system
68 87
Work done during internship at Microsoft Research; on leave from University Amsterdam.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Approaches (i) have had great success in modeling inverse and high-dimensional problems [37] whereas approaches (ii) and (iii) have started to advance fluid and weather modeling two and 89 These problems are usually described by complex time-dependent PDEs. Solving this class of PDEs over long time horizons presents fundamental challenges. Conventional numerical methods suffer accumulating approximation effects, which in the temporal solution step can be counteracted by implicit methods [23] 36.. Neural PDE
three dimensions [1 39
43
46
51
60
63
76
solvers similarly struggle with the effects of accumulating noise, an inevitable consequence of
autoregressively propagating the solutions of the underlying PDEs over time 9 43 59
79 Another
critique of neural PDE solvers is that - besides very few exceptions, e.g.. [33 they lack convergence
guarantees and predictive uncertainty modeling, i.c., estimates of how much to trust the predictions. Whereas the former is in general notoriously difficult to establish in the context of deep learning, the
latter links to recent advances in probabilistic neural modeling 17 28 30
72 75 82] and, thus,
opens the door for new families of uncertainty-awan neural PDE solvers. In summary, to the best of our understanding, the most important desiderata for current time-dependent neural PDE solvers comprise long-term accuracy, long-term stability, and the ability to quantify predictive uncertainty.
In this work, we analyze common temporal rollout strategies, including simple autoregressive unrolling with varying history input, the pushforward trick [9], invariance preservation (58), and the Markov Neural Operator [50]. We test temporal modeling by ofate-of-the-art neural operators such modern U-Nets [22] and Fourier Neural Operators (FNOs) [49]. and identify shared pitfall in all these unrolling schemes: neural solvers consistently neglect components of the spatial frequency spectrum that have low amplitude. Although these frequencies have minimal immediate impact, they still impact long-term dynamics, ultimately resulting in a noticeable decline in rollout performance.
Based on these insights, we draw inspiration from recent advances in diffusion models (30 61 70) to introduce PDE-Refiner. PDE-Refiner is a novel model class that uses an iterative refinement process to obtain accurate predictions over the whole frequency spectrum. This is achieved by an adapted Gaussian denoising step that forces the network to focus on information from all frequency components equally at different amplitude levels. We demonstrate the effectiveness of PDE-Re Refiner on solving the ID Kuramoto-Sivashinsk equation and the 2D Kolmogorov flow, a variant of the incompressible Navier-Stokes flow. On both PDEs, PDE-Refiner models the frequency spectrum much more accurately than the baselines, leading to a significant gain in accurate rollout time.
2 Challenges of Accurate Long Rollouts
Partial Differential Equations. In this work, we focus on time-dependent PDEs in one temporal dimension, i.e., E [0. 7I. and possibly multiple spatial dimensions, i.e., x = [.1 .., Fmm X. Time-dependent PDEs relate solutions at, [0, 7J x X + R" and respective derivatives for all points in the domain, where u (xx) are initial conditions at time = 0 and Blu) a) = are boundary conditions with boundary operator B when lies on the boundary AX of the domain. Such PDEs can be written in the form
U = F(t, x. us Ugy u..::....
where the notation ue is shorthand for the partial derivative 0t/Ot. while Ugz Use ... denote the partial derivatives du/ 0dx Ou/da² and so on learningg
learning 48 49 53-551 relates solutions u. : X + RB, u' X" + R=' defined on different domains X' E Rm X' E Rmmd' via operators G: G (u E (f) + (u' E ('). where U and '2' are the spaces of u and u', respectively. For time-dependent PDEs, an evolution operator can be used to compute the solution at time t + At from time t as
u(t + At) = (Att u(t))
(2)
where G. : Rss x RT + IR" is the temporal update. To obtain predictions over long time horizons, a temporal operator could either be directly trained for large At or recursively applied with smaller time intervals. In practice, the predictions of learned operators deteriorate for large At, while autoregressive approaches are found to perform substantially better 22 50 86).
Long Rollouts for Demendepnn PDEs. We start with showcasing the challenges of obtaining long, accurate rollouts for autoregressive neural PDE solvers on the working example of the 11)
2 Puu/O /xc represents n dimensional Jacobian matrix J with entries Ju = Owu (xx

---

Spectrum the One Step Prodiction
Absolute
Ground Tuala
Auglieak
To
Spatial dimessin
101
FDE-Rcizer (Gus))
MSF Thaining
PUE-Belin (Ows))
Auplinch
net
Edattr amplitel
Fance relaiivn
Amplitaai
MSI Eritnned
PDR-Reieer (Chus)
0
urt
MSE Teaain
20
50
70
100
Timer [10 ssssmnda)
(a) Rollout examples
(b) Frequency spectrum
(c) One-step error
Figure 1: Challenges in achieving accurate long rollouts the KS equation, comparing PDE-Refinn 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 small, high-amplitude frequency band, while PDE-Refiner supports much larger frequency band, leading to longer accurate rollouts as in (a).
Kuramoto- Kuramoto-Sivashinst Sivashinsky (KS) equation 44 known for its rich dynamical characteristics and chaotic behavior 35 40
The KS equation a fourth-order nonlinear PDE,
74 It is defined as:
Ut + ut, + Ugz + Vilyzaa = 0,
where v is viscosity parameter which we commonly set to = 1. The nonlinear term wu, and the fourth-order derivative Myzrr make the PDE a challenging objective for traditional solvers. We aim to solve this equation for all and t on domain [0, L.] with periodic boundary conditions u(0, 6) = u(..,t t) and an initial condition uzx 0) = un(2). The input space is discretized uniformly on grid of N₂ spatial points and N₂ time steps. To solve this equation, neural operator, denoted by NO, is then trained to predict a solution u(æ, t) = uu() given one or multiple previous solutions ult u(t-AV/ At) with time step At, e.g. i(t) = (O(( At)). Longer trajectory predictions are obtained by feeding the predictions back into the solver, i.e., predicting u(t + At) from the previous prediction a(t) via ú(t + At) = vt). We refer this process as unrolling the model or rollout. The goal is to obtain a neural solver that maintains predictions close to the ground truth for as long as possible.
The MSE training objective. The most common objective used for training neural solvers is the one-step Mean-Squared Error (MSE) loss: LMSE = [|x(t)) NO - At) By minimizing this one-step MSE, the model learns replicate the PDE's dynamics, accurately predicting the next step. However, as we roll out the model for long trajectories, the error propagates over time until the predictions start differ significantly from the ground truth. In Figure the the solver is already accurate for 70 seconds, so one might argue that minimizing the one-step MSE is sufficient for achieving long stable rollouts. Yet, the limitations of this approach become apparent when examining the frequency spectrum across the spatial dimension of the ground truth data and resulting predictions. Figure shows that the main dynamics of the KS equation are modeled within a frequency band of low wavenumbers (I to 25). As result, the primary errors in predicting a one-step solution arise from inaccuracies in modeling the dynamics of these low frequencies. This is evident Figure Ic where the error of the MSE-trained model is smallest for this frequency band relative to the ground truth amplitude. Nonetheless, over long time horizon, the non-linear term uu, the KS equation causes all frequencies to interact, leading to the propagation of high-frequency errors into lower frequencies. Hence, the accurate modeling of frequencies with lower amplitude becomes increasingly important for longer rollout lengths. In the KS equation, this primarily pertains to high frequencies, which the MSE objective significantly neglects.
Based on this analysis. we deduce that in order to obtain long stable rollouts, we need neural solver that models all spatial frequencies across the spectrum as accurately as possible. Essentially, our objective should give high amplitude frequencies higher priority, since these are responsible for the main dynamies of the PDE. However, at the same time, the neural solver should not neglect the non-dominant, low amplitude frequency contributions due to their long-term impact on the dynamics.
We omit the bold notation for ID cases where the field u(x, t) scalar valued.
3

---

Initial Prediction
Refinement process
u(t An A
u(t An) M
Neural Operator
Neural Operator
a (²(()
a++ (0)
(0)
Inpur a
af(1)
Refined prediction
Noisc ~ (0,
Repeat for k
Figure 2: Refinement process of PDE-Refiner during inference. Starting from an initial prediction (t), PDE-Refiner uses iterative refinement process to improve its prediction. Each step represents denoising process, where the model takes as input the previous step's prediction w²() and tries to reconstruct added noise. By decreasing the noise variance of over the refinement steps, PDE- Refiner focuses all frequencies equally, including low-amplitude information.
3 PDE-Refiner
In this section, we present PDE-Refiner, model that allows for accurate modeling of the solution across all frequencies. The main idea of PDE-Refiner is that we allow the model look multiple times at its prediction, and, in an iterative manner, improve the prediction. For this, we use model NO with three inputs: the previous time step(s) u(t At), the refinement step index & E [0, .... K]. and the model's current prediction n (t). At the first step k O, mimic the common MSE objective by setting (t) 0 and predicting u(t): - t) = ||u(t) NO (t), u(t At), 0) I12. As discussed in Section 2 this prediction will focus on only the dominating frequencies. To improve this prediction, simple approach would be train the model to take its own predictions inputs and output its (normalized) error to the ground truth. However, such a training process has several drawbacks. Firstly, as in Figure the dominating frequencies in the data also dominate in the error, thus forcing the model to focus on the same frequencies again. As we empirically verify in Section 4.1 this leads to considerable overfitting and the model does not generalize.
Instead, we propose to implement the refinement process as a denoising objective. At each refinement step A 2 1. we remove low-amplitude information of an earlier prediction by applying noise, e.g. adding Gaussian noise, to the input a² (²(t) at refinement step k: a ū²(t) a (t) + a&C ~ NO, 1). @*(t)
The objective the model to predict this noise * and use the prediction z to denoise its input: (1) (() o+E#_ By decreasing the noise standard deviation σą over refinement steps, the model focuses on varying amplitude levels. With the first steps ensuring that high-amplitude information is captured accurately, the later steps focus low-ttplitude information, typically corresponding to the non-dominant frequencies. Generally, we find that an exponential decrease, i.c.
σ% = uin with mmin being the minimum noise standard deviation, works well. The value of ain is chosen based on the frequency spectrum of the given data. For example, for the KS equation, we use o = 2 10-7 We train the model by denoising ground truth data at different refinement steps:
k/K
(u, = E,+~N(0.1) [l|e - NO (u(t) + okck, u(t -)),
Crucially. by using ground truth samples in the refinement process during training, the model learns to focus on only predicting information with a magnitude below the noise level a₁ and ignore potentially larger errors that, during inference, could have occurred in previous steps. To train all refinement steps equally well, we uniformly sample & for each training example: L(u, t) = E4-U((,KK)
[ (u, t)]-
At inference time, we predict solution u(t) from u(t At) by performing the K refinement steps, where we sequentially use the prediction of a refinement step the input to the next step. While the process allows for any noise distribution, independent Gaussian noise has the preferable property that it is uniform across frequencies. Therefore, it removes information equally for all frequencies, while also creating a prediction target that focuses on all frequencies equally. We empirically verify in Section 4.1 that PDE-R Reffinee even improves on low frequencies with small amplitudes.

---

3.1 Formulating PDE-Refiner as Diffusion Model
Denoising processes have been most famously used diffusion models well 12 29
69 77
Denoising diffusion probabilistic models (DDPM) randomly sample noise variable Xo
and sequentially denoise it until the final prediction, XK, is distributed according the data:
Pe (xo:k) = p(xo) II pe ( xx)), pe(xx+1 (xx)) N (x ueo (x k), Ex k)) (5)
where K is the number of diffusion steps. For neural PDE solving, one would want pa (Xk) to model the distribution over solutions, XK = u(t), while being conditioned on the previous time step u(t At), i.e., Po (o(u(() (u()))((L A()). For example, Lienen al. (51] recently proposed DDPMs for modeling 3D turbulent flows due to the flows' unpredictable behavior. Despite the similar use of denoising process, PDE-Refiner sets itself apart from standard DDPMs in several key aspects. First, diffusion models typically aim to model diverse, mulli--odaa distributions like in image generation, while the PDE solutions we consider here are deterministic. This necessitates extremely accurate predictions with only minuscule errors. PDE-Refiner accommodates this by employing an exponentially decreasing noise scheduler with very low minimum noise variance o² min' decreasing much faster and further than common diffusion schedulers. Second, our goal with PDE-Refinet is not only to model a realistic-looking solution, but also achieve high accuracy across the entire frequency spectrum. Third, we apply PDE-Refiner autoregressively to generate long trajectories. Since neural PDE solvers need to be fast to be an attractive surrogate for classical solvers in applications, PDE- Refiner uses far fewer denoising steps in both training and inferences than typical DDPMs. Lastly, PDE-Refiner directly predicts the signal u(t) at the initial step. while DDPMs usually predict the noise residual throughout the entire process. Interestingly, a similar objective to PDE-Refiner is achieved by the v-prediction 770). which smoothly transitions from predicting the sample u(t) to the additive noise e: y* v o (t. Here, the first step k = 0, yields the common MSE prediction objective by setting do 1. With an exponential noise scheduler, the noise variance is commonly much smaller than I for l: 2 1. In these cases, the weight of the noise is almost in the v-prediction, giving diffusion process that closely resembles PDE-Refiner.
Nonetheless, the similarities between PDE-Refiner and DDPMs indicate that PDE-Refiner has a potential interpretation
probabilistic latent variable model. Thus, by sampling different noises during the refinement process, PDE- Refiner may provide well-calibrated uncertainties which faithfully indicate when the model might be making errors. We return this intriguing possibility later in Section 4.1 Further, we find empirically that implementing PDE- PDE-Refiner as a diffusion model with our outlined changes in the previous paragraph, versus implementing it as an explicit denoising process, obtains similar results. The benefit of implementing PDE-Refiner as a diffusson model is the large literature on architecture and hyperparameter studies, as well as available software for diffusion models. Hence, we use a diffusion- based implementation of PDE-Refiner in our experiments.
4 Experiments
We demonstrate the effectiveness of PDE-Refiner on a diverse set of common PDE benchmarks. In 1D,, we study the Further, we study the models' robustness to different spatial frequency spectra by varying the visocisity in the KS equation. In 2D, we compare PDE-Refiner to hybrid PDE solvers on turbulent Kolmogorov flow, and provide speed comparison between solvers. We make our code publicly available at
equation and compare to typical temporal rollout methods.
4.1 Kuramoto-Sivashinsky 1D equation
Experimental setup. We evaluate PDE-Refinee and various baselines on the Sivashinsk ID equation. We follow the data generation setup of Brandstetter et al. [8] by using a mesh of length L discretized uniformly for 256 points with periodic boundaries. For each trajectory, we randomly sample the length L between [0.9 64, 1.1 64] and the time step At ~ UO0.18, 0.22). The initial conditions are sampled from a distribution over truncated Fourier series with random coefficients { em, Pm Jm as un(r) = -m= Am (2 + @m). We generate training datasct with 2048 trajectories of rollout length 140At, and test on 128 trajectories with duration
5

---

High-Correlation Rollout Times on the Kuramoto-Slvaelinskk equation
l1gh col--corriattnn (inocomnc))
MSE Training
PDF-Refiee (Ouns Dilinsion Ablacio
9825
58.36
98.5x
Alternative Losses
89.86
01.2s
73.77 70.7=
75
15 dx
74.8o
70.74
71.24
71.
70
66 95
60.6
Beseion
Histoy
ds parouutes
Exeeable
Prefitammrd
Subslrr A Sobuder A m
Eina Curretion
MNN
Extr Prefecion
stop
aduje
stoge
stops
stopa
stepa Mean
Caiin Schedale
Schedole
Figure 3: Experimental results the -iiasiiiii 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 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 modern U-Net neural operator [22] PDE-Refiner surpasses all baselines with accurate rollouts up to nearly 100 seconds.
of 64 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] [8) and diffusion modeling [2 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, detailed study with FNOs in Appendix E.T shows that even here PDE-Refiner offers significant performance gains. Finally, we also evaluate on Dilated ResNets (78) in Appendix [..] 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(!), each model takes as input the previous time step u(l 4A1) and the trajectory parameters L and At. Thereby, the models predict the residual between time steps Au()) = u(t) - u(t - 4 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' nigh-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 Appendix D leading to the same conclusions.
MSE Training. We compare PDE-Refinee to three groups of baselines in Figure 3 The first group are models trained with the one-step MSE error, i.c., predicting Au(t) from u(t 4A)). 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.c u(t 4 and u(t S&t). improves the one-step prediction but worsens rollout performance. The problem arising is that the difference between the inputs u(t 4 - u(t 8At) highly correlated with the model's target Au(t), the residual of the next time step. This leads the neural operator to focus on modeling the second-order difference Au(t) Au(t 4 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 [F.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.
The second baseline group includes alternative losses and post-pros 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 our setting, confirming previous results K. 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-frrqquncyy information, the Sobolev norm loss [50) maps the
Alternative losses.

---

Spectrum One-Step. Prodbction
Spertrum Differenee to CT
Spectrure Nois lovel
Ground truth
MSI Tralning
10'
laitial pentiction
30
lutiil prodiction
10'
Bettieeenn step
Reflocment step
Angituuk
-0-1
Relissment step
Auplituls
10
Refincasni
Amplitude
Reliieeent step
Reficcucc step
-0-3
Netor at refinmment step
10-1
10-05
10-3-
Notac at reflactonnt step
Notse refforment sto
125
100
Wiosmmmierr
Wevemmber
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 (t). i (D), i² (t). (1) of PDe-erfiner's refinement process compared the Ground Truth. Center: the spectrum of the difference between ground truth and intermediate predictions, i.e. [FFT(u(t)) - a (f)) Right: the spectrum of the noise o&c² 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 aeppltul 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 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 O, as their amplitude is below float32 precision (see Appendix D.I] 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 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- -egh-mmllttue 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 trade- off 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 standard denoising diffusion model (30] that we condition on the previous time step u(t - 4At). When using common cosine noise schedule [6]), 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 Reff-eeieeer we check whether samples single modes under potentially multi-moda. 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 PeE-Refiner has close-to uniform error pattern across frequencies. While the first refinement step shows an improvement across all frequencies, refinmment 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 PDF-Reefieer does not only improve the high frequencies, but also the lowest frequencies (wavenumber 1 1-6) with low amplitude.
7

---

KS rollout stability over input resolution
Cateh-CCrrcltiion Rollout Times uver Datasel Size
05
High-smelltion thme
100
MSF Training
91 6x
97.5s
PDE- PoE-Roieer (Oun))
53
PDE-Re MSR Training
75.34
75.45
11gh comdeliion
85-
65.24 22
80
60
52
70-
125 Spatial Resolution
236
27/6 Proportion training trajectories used
Figure 5: Left: Stable rollout time over input resolution. PDE-Refiner models the high frequencies to improve its rollout on higher resolutions. Right: Training smaller datasets. PDE-Refine consistently outperforms the MSE baseline, increasing its relative improvement to 50% for the lowest data regime.
and the MSE baseline on
Input Resolution. We demonstrate that capturing high-frequency information is crucial for PDE-Refiner's performance gains over 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 Figure 5 MSE models achieve similar rollout times for resolutions between 32 and 256, emphasizing its inability model high frequreeyy information. At resolution of 32, PDE-Refinct achieves similar performance to the MSE baseline due to the missing high- frequency information. However, resolution increases, PDE-Refine significantly outperforms the baseline, showcasing its utilization of high-freyuency information.
Spectral data augmentation. A pleasant side effect of PDE-Refiner is data augmentation, which is induced by adding varying Gaussian noise ORE", e ~ N0 1) at different stages & 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] 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 leave more rigorous testing of PDE-Refinee induced data augmentation for future work, we show results for the low training data regime in Figure 5 When training -De-Refiner and the MSE baseline on 10%, 20% and 50% of the training trajectories, PDE-Refine consistently outperforms the baseline in all settings. Moreover, with only 10% of the training data, PDE-Refiner performs on par to the MSE model 100%. Finally, the relative improvement of PDE-Refiner increases 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 tra- jectories 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 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 accu- rate, leading Rt2 coefficient of 0.86 between the times. Furthermore, the prediction for how long the rollout remains accurate depends strongly on the individual trajec- tory - PDE- P-E--DDeiinee 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 common approaches. PDE-Refiner provides more accurate estimates than input modulation 5 while only requiring trained model compared to a model ensemble [45 71).
Higlrorrdation secomde)
140
Lancar fit
Trajectories
120
100
80
60
60
80
100
120
1400
Croes-correlatin time (In seconds))
Figure 6: Uncertainty estimate of PDE- Refiner. Each point represents the esti- mated correlation time via sample cross- correlation (x-axis) and the ground truth time (y-axis) for test trajectory.

---

Frequency spectrim of KS equation
High-osredation time (In seconds)
Rollout over viscosity
10'
PDE-Refiner
120
Avergg amplftude
MSE Training
1.25
10-1
Vaccusty
1.00
80
10-1
0.75
60
10
0.50
40
Wavenmmber
Viscosily
Figure 7: Visualizing the parameter-dependen 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, have focused the KS equation with viscosity term of 1. Under varying values of v, the Suramoto-SSvvshinsky equation has been shown to develop diverse behaviors and fixed points This offers ideal benchmark for evaluating neural surrogate methods
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 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 and add the viscosity y to the conditioning set of the neural operators.
35
40 74
setup identical to Section
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 [V 0.1, 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 the given underlying data. Furthermore, both models achieve similar performance to their unconditional counterpart for v = 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-dy muid-dynamics benchmark, we apply
to the 2D Kolmogorov flow, variant of the incompressible Savier-Stokes flow. The PDE is defined as:
v - Vp
uu+ (u =
(6)
2 the solution, the tensor product, the kinematic viscosity, p the fluid density, the pressure field, and, finally, f the external forcing. Following previous work [43] [9], we set the foreing to f sin (4y) & 0.lu, the density م 1, and viscosity v 0.001, which corresponds to Reynolds number of 1000. The ground truth data is generated using finite volume-base direct numerical simulation (DNS) method of 2048 x 2048, and afterward downscaled to 64 x 64. To align our experiments with previous results, we use the same dataset 128 trajectories for training and trajectories for testing as Sun et [79
where u [0,
57 with time step of At 7.0125 x 10-3 and resolution
Experimental setup. We employ a modern U-Net 222 as the neural operator backbone. Due to the lower input resolution, we set - 10-3 and use 3 refinement steps in PDE-Refiner. For efficiency, we predict 16 steps (16At) into the futuree and use the difference Au = u(t) u(t 16At) as the output target. Besides the MSE objective,
compare PDE-Refiner with FNOs [50], classical PDE solvers (i.e., DNS) on different resolutions, and state-of-the-an hybrid machine learning solvers which estimate the convective flux u via neural networks. Learned Interpolation (LJ) (43
43 79
takes the previous solution u(t - At) as input to predict u(t). similar to PDE-Refiner In contrast, the Temporal Steneil 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 classical solver with neural networks. For evaluation, roll out the models on the 16 test trajectories and determine the Pearson correlation with the
(L.C) 43 83 which corrects the outputs

---

ground truth in terms of the scalar vorticity field = 0g - 0y Following previous work [79]. we report Table the time until which the average correlation across trajectories falls below 0.8.
we find that modern U-Nets outperform FNOs the 2D do- main for long rollouts. Our MSE-trained U-Net al- ready surpasses all classical and hybrid PDE solvers. This result highlights the strength of our baselines, and improving upon those poses significant challenge. Nonetheless, PDE- -De-Refiner manages to provide sub- stantial 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 A for this dataset Appendix [..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 spec- trum. This accounts for the smaller relative gain of PDE-Refiner the MSE baseline here.
Results. Similar to previous work 12
55,
Table 1: Duration of high correlation
0.8) on the 2D Kolmogorov flow. Results for classical PDE solvers and hybrid meth- ods taken from Sun et al. [79
Method
Corr. > 0.8 time
Classical PDE Solvers
DNS 64 64
2.805
DNS 128 x 128
3,983
DNS 256 x 256
5.386
DNS 512 x 512
6.788
DNS 1024 x 1024
8,752
Hybrid Methods
LC
43
33
CNN
6,900
LC
43
33
FNO
7.630
LI
43
CNN
7.910
TSM
79
FNNO
7.798
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 (=00.01,, with PDE-Refiner taking 4 times longer (16.53 + 0.04 seconds) due to four model calls per step. With that, PDE-Refiner still faster than the best hybrid solver, TSM, which needs 20.25 seconds (+0.05). In comparison to the ground truth solver at resolution 2048 x 2048 with 31 minute generation time GPU, all surrogates provide significant specdup.
TSM
79
CNN
8,359
TSM
79
НHPPO
9,481
ML. Surrogates
MSE training FNNO
6.451 + 0.105
MSE training
U-Net
9.663 0.117
PDE-Refiner U-Net
10.659 0.092
5 Conclusion
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 IDE-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.
Linitations. 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, applying distillation and enhanced samplers 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 interesting avenue for future work. Further architectures like Transformers [13 84 can be explored too, having been shown also suffer from spatial frequency biases for PDEs . 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 the KS dataset, 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.
10

---

References
[1] Troy Arcomano, Istvan Szunyogh, Alexander Wikner, Jaideep Pathak, Brian R Hunt, and Edward Ott. 2022. Hybrid Approach to Atmospheric Modeling That Combines Machine Learning With Physics-Basee Numerical Model. Journal of Advances in Modeling Earth Systems, 140)x2202MMS002722
[2] Yohai Bar-Sinai, Stephan Hoyer, Jason Hickey, and Michael P Brenner. 2019. Learning data- driven discretizations for partial differential equations. Proceedings of the National Academy of Sciences, 116(31):1
[3] David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap, Shuangfei Zhai, Siyuan Hu, Daniel Zheng. Walter Talbot, and Eric Gu. 2023. TRACT: Denoising Diffusion Models with Transitive Closure mint-Distillation. arXiv preprint arXiv:2303.04228.
[4] Saakaar Bhatnagar, Yaser Afshar, Shaowu Pan, Karthik Duraisamy, and Shailendra Kaushik. 2019, Prediction of aerodynamic flow fields using convolutional neural networks. Computational Mechanics, 64(2):525-544
[5] Neill E. Bowler. 2006. Comparison of error breeding, singular vectors, random perturbations and ensemble Kalman filter perturbation strategies a simple model. Tellus A: Dynamic Meteorology and Oceanography, 58(5):538-548.
[6] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milh Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs. 2018. Software URL: http://gttubb.co//ooglgeejx
[7] Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K Gupta. 2023. Clifford Neural Layers for PDE Modeling. In The Eleventh International Conference on Learning Representations.
[8] Johannes Brandstetter, Max Welling, and Daniel E Worrall. 2022. Lie Point Symmetry Data Augmentation for Neural PDE Solvers. Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 2241- 2256. PMLR.
[9] Johannes Brandstetter, Daniel E. Worrall, and Max Welling, 2022. Message Passing Neural PDE Solvers. In International Conference on Learning Representations.
[10] Steven Brunton and J Nathan Kutz. 2023. Machine Learning for Partial Differential Equations. arXiv preprini arXiv:2303.1710 17078.
[11] Ashesh Chattopadhyay and Pedram Hassanzadeh. 2023. Long-term instabilities of deep learning-based digital twins of the climate system: The cause and a solution. arXiv preprint 077002
[12] Prafulla Dhariwal and Alexander Nichol. 2021. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-877
[13] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby, 2021. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations.
[14] William Falcon and The PyTorrh Lightning team. PyTorch Lightning. 2019. Software URL:
https://gi thubb
[15) Vladimir Fanaskov, Tianchi Yu, Alexander Rudikov, and Ivan Oseledets. 2023. General
Covariance Data Augmentation for Neural PDE Solvers. arXiv preprint arXiv:2301.12730.
[16] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley. Sherjil Ozair, Aaron Courville, and Yoshua Bengio. 2014. Generative Adversarial Nets. In Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc
[17] Jonathan Gordon, Wessel P Bruinsma, Andrew YK Foong, James Requeima, Yann Dubois, and Richard Turner. 2019. Convolutional conditional neural processes. International Conference Learning Representations.

---

[18] Daniel Greenfeld, Meirav Galun, Ronen Basri, Irad Yavneh, and Ron Kimmel. 2019. Learning Optimize Multigrid PDE Solvers. In International Conference on Machine Learning (ICML),
pages 2415-2423.
[19] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. 2020. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 33:1474-1487.
[20] Albert Gu, Karan Goel, and Christopher Re. 2022. Efficiently Modeling Long Sequences with Structured State Spaces. In International Conference on Learning Representations.
[21] Xiaoxiao Guo, Wei Li, and Francesco Iorio. 2016. Convolutional neural networks for steady flow approximation. In Proceedings of the 22nd ACM SIGKDD International Conference Knowledge Discovery and Data Mining, pages 481-490.
[22] Jayesh K Gupta and Johannes Brandstetter. 2022. Towards Multi-spatiotemporaa-cceo General- ized PDE Modeling, arXiv preprint arXiv:
[23] Ernst Hairer and Gerhard Wanner. 1996. Solving ordinary differential equations. II, volume 14 of Springer Series in Computational Mathematics.
[24] Jiequn Han, Arnulf Jentzen, and Weinan E. 2018. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 15(34):8505- 8510.
[25] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778.
[26] Kaiming Hc, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. 2016. Identity Mappings in Deep Residual Networks. In Computer Vision ECCV 2016, pages 630 645, Cham. Springer International Publishing.
[27] Dan Hendrycks and Kevin Gimpel. 2016. Gaussian error linear units (gelus). arXiv preprint 1rXiv:1606.084/ 15.
[28] Philipp Hennig, Michael A Osborne, and Hans P Kersting. 2022. Probabilistic Numerics: Compulation Machine Learning. Cambridge University Press.
[29] Jonathan Ho. William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederík Kingma, Ben Poole, Mohammad Norouzi, David Fleet, et al. 2022. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303.
[30] Jonathan Ho, Ajay Jain, and Pieter Abbeel. 2020. Denoising Diffusion Probabilistic Models. In Advances in Neural Information Processing Systems, volume 33, pages 6840-6851. Curran Associates, Inc
[31] Jonathan Ho, Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. 2022. Cascaded Diffusion Models for High Fidelity Image Generation. J. Mach. Learn. Res., 23(47):1-33.
[32] Emiel Hoogeboom and Tim Salimans. 2023. Blurring Diffusion Models. In The Eleventh International Conference on Learning Representations.
[33] Jun-Ting Hsich, Shengjia Zhao, Stephan Eismann, Lucia Mirabella, and Stefano Ermon. 2019. Learning Neural PDE Solvers with Convergence Guarantees. arXiv preprint arXiv:1906.00 012000
[34] D. Hunter. 2007. Matplotlib: A 2D graphics environment. Computing in Science & Engineer- ing, 9(3):90-95 Software URL: https: github.. com//aatplooiiib/mmppooooo
[35] James M. Hyman and Basil Nicolaenko. 1986. The Kuramoto-Sivashinse equation: A bridge between PDE'S and dynamical systems. Physica D: Nonlinear Phenomena, 18(1):113-226
[36] Arieh Iserles. 2009. A first course the numerical analysis of differential equations. Number 44 in Cambridge Texts in Applied Mathematics. Cambridge university press.
[37] George Em Kariiadaiis, loannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. 2021. Physics-informed machine learning. Nature Reviews Physicx, 3(6))222-44))
[38] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine, 2022. Flucidating the Design Space of Diffusion-Basse Generative Models. In Advances in Neural Information Processing Systems.
12

---

[39] Ryan Keisler. 2022. Forecasting Global Weather with Graph Neural Networks. arXiv preprint arX.v:2202.07575.
[40] loannis G. Kevrekidis, Basil Nicolaenko, and James C. Scovel. 1990. Back in the Saddle Again: A Computer Assisted Study of the Kuramoto Sivasshiskv Equation. SIAM Journal on Applied Mathematics, 50(3):760-790.
[41] Diederik P. Kingma and Jimmy Ba. 2015. Adam: A Method for Stochastic Optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings.
[42] Diederik P. Kingma and Max Welling. 2014. Auto-Encoding Variational Bayes In 2nd International Conference on Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Conference Track Proceedings.
[43] Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan Hoyer. 2021. Machine learning-acceleren National Academy of Sciences, (18(2)):e210178411
computational fluid dynamics. Proceedings of the
[44] Yoshiki Kuramoto. 1978. Diffusion-induces chaos in reaction systems. Progress of Theoretical Physics Supplement, 64:346-367.
[45] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. 2017. Simple and Scalable Predictive Uncertainty Estimation Using Deep Ensembles. In Proceedings of the s1 Interna- tional Conference on Neural Information Processing Systems, NIPS' 17, page 6405 6416, Red Hook, NY, USA. Curran Associates Inc.
[46] Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortu- nato, Alexander Pritzel, Suman Ravuri, Timo Ewalds, Ferran Alet, Zach Eaton-Rosen, al. 2022. GraphCast: Learning skillful medium-rante range global weather forecasting. arXiv preprint arXiv:2212.12794.
[47] Sangyun Lee, Hyungjin Chung, Jaehyeon Kim, and Jong Chul Ye. 2022. Progressive deblurring of diffusion models for coarse-to-fine image synthesis. arXiv preprint arXiv:2207.1 /992.
[48] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. 2020. Neural operator: Graph kernel network for partial differential equations. arXiv preprint arXiv:2003.0344
[49] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhat- tacharya, Andrew Stuart, and Anima Anandkumar. 2021. Fourier Neural Operator for Paramet- ric Partial Differential Equations. In International Conference on Learning Representations.
[50] Zongyi Li, Miguel Schiaahiafinni, Nikola Borislavov Kovachki, Kamyar Arizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. 2022. Learning Chaotic Dynamics in Dissipative Systems. Advances in Neural Information Processing Systems.
[51] Marten Lienen, Jan Hansen-Palman Palmus, David Lüdke, and Stephan Günnemann. 2023. Generative Diffusion for 3D Turbulent Flows. arXiv preprint arXiv:2306.01770
[52] Ily Loshchilov and Frank Hutter. 2019. Decoupled Weight Decay Regularization. In Interna- tional Conference on Learning Representations.
[53] Lu Lu, Pengzhan Jin, and George Em Karniadakis. 2019. DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. arXiv preprint arXiv: 19 1910.03193.
[54] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. 2021. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218-22) 229.
[55] Lu Lu, Xuhui Meng, Shengze Cai, Zhiping Mao, Somdatta Goswami, Zhongqiang Zhang, and George Em Karniadakis. 2022. A comprehensive and fair comparison of two neural operators (with practical extensions) based on fair data. Computer Methods in Applied Mechanics and Engineering, 397:1
[56] Hao Ma, Yuxuan Zhang, Nils Thuerey, Xiangyu Hu, and Oskar J Haidn. 2021. Physics-driven learning of the steady Navier-Stokes equations using deep convolutional neural networks. arXiv preprint 21X::2200600930.
13

---

[57]
James M. McDonough. 2007. Lectures in Computational Fluid Dynamics of Incompressible Flow: Mathematics, Algorithms and Implementations. 4. Mechanical Engineering Textbook Gallery.
[58]
Nick McGreivy and Ammar Hakim. 2023. Invariant preservation in machine learned PDE solvers via error correction. arXiv preprint 261
[59] Jonas Mikhaeil, Zahra Monfared, and Daniel Durstewitz. 2022. On the difficulty of learning chaotic dynamics with RNNs. In Advances in Neural Information Processing Systems, vol- ume 35, pages 11297-1 1312. Curran Associates, Inc.
[60] Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K Gupta, and Aditya Grover. 2023. ClimaX: A foundation model for weather and climate. arXiv preprint arXiv:2301. 10343.
[61] Alexander Quinn Nichol and Prafulla Dhariwal. 2021. Improved denoising diffusion probabilis- tic models. In International Conference on Machine Learning, pages 8162-8171. PMLR.
[62] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, An- dreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chil- amkurthy. Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019. PyTorch: An Imperative Style, High-Perfornance Deep Learning Library. In Advances in Neu- ral Information Processing Systems, volume 32. Curran Associates, Inc. Software URL: https:/ /ttps::/gi// http:///gttuu..... wwt./oooomeeeeeee http:/ytttttpooor/eccccc
[63] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, Pedram Hassanzadeh, Karthik Kashinath, and Animashree Anandkumar. 2022. FourCastNet: A Global Data-driven High- figh-resolution Weather Model using Adaptive Fourier Neural Operators. arXiv preprint 2224114402...
[64] Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron Courville. 2018. FiLM: Visual Reasoning with a General Conditioning Layer. In Proceedings of the Thirty- Second AAAI Conference Artificial Intelligence and Thirlieth Innovative Applications of Artificial Intelligence Conference and Eighth AAAI Symposium on Educational Advances in Artificial Intelligence, AAATI8MAAF18EAATT 18. AAAI Press.
[65] Patrick von Platen, Suraj Patil, Anton Lozhkov, Pedro Cuenca, Nathan Lambert, Kashif Rasul, Mishig Davaadorj, and Thomas Wolf. 2022. Diffusers:: State-of-the-art diffusion models. Software URL: https: httL:htts:///ttt //github. cww http:///uugiingfffceeffuseeeeeee
[66] Maziar Raissi, Paris Perdikaris, and George Karniadakis. 2019. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707.
[67] Stephan Rasp and Nils Thuerey. 2021. Data-driven medium-ramm rrange weather prediction with a resnet pretrained on climate simulations: A new model for weatherbench. Journal of Advances in Modeling Earth Systems, 13Q))e2020MS00240..
[68] David Ruhe, Jayesh K Gupta, Steven De Keninck, Max Welling, and Johannes Brandstetter. 2023. Geometric Clifford Algebra Networks. In Proceedings of the 40th International Confer- ence on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 29306-29337. PMLR.
[69] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily I. Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. 2022. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 01136479-36494
[70] Tim Salimans and Jonathan Ho. 2022. Progressive Distillation for Fast Sampling of Diffusion Models. In International Conference on Learning Representations.
[71] Sebastian Scher and Gabriele Messori. 2021. Ensemble Methods for Neural Network-Based
Weather Forecasts. Journal of Advances in Modeling Earth Systems, 13(2).
[72] Jacob H Seidman, Georgios Kissas, George J Pappas, and Paris Perdikaris. 2023. Variational Autoencoding Neural Operators. arXiv preprint arXiv:2302.10331.
[73] G.I. Sivashinsky. 1977. Nonlinear analysis of hydrodynamic instability in laminar flames-L. Derivation of basic equations. Acta Astronautica, 4(11):1177-1266.
14

---

[74] Yiorgos S. Smyrlis and Demetrios T. Papageorgiou. 1991. Predicting Chaos for Infinite Dimen- sional Dynamical Systems: The of the National Academy of Sciences of the United States of America, 88(24): (112)):11229-1111 1292-11122
Equation, A Case Study. Proceedings
[75] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. 2015. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2266-2265. PMLR.
[76] Casper Kaae Sønderby, Lasse Espeholt, Jonathan Heek, Mostafa Dehghani, Avital Oliver, Tim Salimans, Shreya Agrawal, Jason Hickey, and Nal Kalchbrenner. 2020. Metnet: A neural weather model for precipitation forecasting. arXiv preprint arXiv:2003. 12140.
[77] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. 2021. Score- Generative Modeling through Stochastic Differential Equations. In International Conference on Learning Representations.
[78] Kim Stachenfeld, Drummond Buschman Fielding, Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Zonzalz. 2022. Learned Simulators for Turbulence. In International Conference on Learning Represen- tations.
[79] Zhiqing Sun. Yiming Yang, and Shinjae Yoo. 2023. A Neural PDE Solver with Temporal Stencil Modeling, arXiv preprint arXiv:2302.08105.
[80] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. 2022. PDEBench: An Extensive Benchmark for Scientific Machine Learning. In 36th Conference on Neural Information Processing Systems (NewrIPS 2022) Track on Datasets and Benchmarks.
[81] Nils Thuerey, Philipp Holl, Maximilian Mueller, Patrick Schnell, Felix Trost, and Kiwon Um. 2021. Physics-based Deep Learning. arXiv preprint arXiv:2109.05237.
[82] Jakub M Tomezak. 2022. Deep generative modeling. Springer.
[83] Kiwon Um, Robert Brand, Yun (Raymond) Fei, Philipp Holl, and Nils Thuerey. 2020. Solver- in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDF-Solvers. In Advances in Neural Information Processing Systems, volume 33, pages 6111-6122. Curran Associates, Inc.
[84] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, t. ukasz Kaiser, and Illia Polosukhin. 2017. Attention is All you Need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
[85] Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, An- drew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C Carey, ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Hen- riksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. 2020. SciPy 1.0: Fundamental Al- gorithms for Scientific Computing in Python. Nature Methods, 17:261-272. Software URL: https://gi thub. cww http:/ssppyyscsspp
[86] Sifan Wang and Paris Perdikaris. 2023. Long-time integration of parametric evolution equations with physics-informed deeponets. Journal of Computational Physics, 475: 175:111855.
[87] Sifan Wang, Hanwen Wang. and Paris Perdikaris. 2021. Learning the solution operator of parametric partial differential equations with physics-informed DeepONets. Science advances,
=40))eab08605.
[88] Daniel Watson, William Chan, Jonathan Ho, and Mohammad Norouzi. 2022. Learning Fast Samplers for Diffusion Models by Differentiating Through Sample Quality. In International Conference on Learning Representations.
[89] Jonathan A Weyn, Dale R Durran, and Rich Caruana. 2020. Improving data-driven global weather prediction using deep convolutional neural networks cubed sphere. Journal of Advances in Modeling Earth Systems, 1209):2020MS002109.
[90] Yuxin Wu and Kaiming He. 2018 Group Normalization. In Proceedings of the European Conference Computer Vision (ECCV).
15

---

[91] Yasin YaZıeı, Chuan-Sheng Foo, Stefan Winkler, Kim-Hui Yap, Georgios Piliouras, and Vijay Chandrasekhar. 2019. The Unusual Effectiveness of Averaging in GAN Training. In Interna- tional Conference on Learning Representations.
[92] Fisher Yu and Vladlen Koltun. 2016. Multi-Scale Context Aggregation by Dilated Convolutions. In International Conference Learning Representations.
16

---

SUPPLEMENTARY MATERIAL
PDE-REFINER: ACHIEVING ACCURATE LONG ROLLOUTS WITH NEURAL PDE SOLVERS
TABLE OF CONTENTS
Broader Impact
18
Reproducibility Statement
18
PDE-Refiner Pseudocode
19
Experimental details
21
D.I
Kuramoto-Sivashinsky ID dataset
21
D.2
Parameter-dependent KS dataset
25
D.3
Kolmogorov 2D Flow
26
Supplementary Experimental Results
29
Fourier Neural Operator
29
Dilated ResNets
30
E.3
Step Size Comparison
31
History Information]
32
E.5
Uncertainty Estimation
33
E.6
Frequency Analysis for 2D Kolmogorov Flow
34
Minimum noise variance in Refiner
35
Stability of Very Long Rollouts
36
17

---

A
Broader Impact
Neural PDE solvers hold significant potential for offering computationally cheaper approaches to modeling 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 at http://wwww..ooo.c.ooooooo////ae We report the used model architectures, hyperparameters, and dataset properties detail in Section and Appendix [ We additionally include pseudocode for our proposed method, PDE-Refine in Appendix C All experiments
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 wethon-based training framework for neural PDE solvers in PyTorch [2] 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 scipy [85] in the public code of
which implements
Brandstetter et al. (8] for the KS equation, and JAX (6 in the public code Kochkov et al. 43) Sun
et al. [79] for the 2D Kolmogorov Flow dataset. The usage of these assets is further described in Appendix
In terms of computational resources, all experiments have been performed on NVIDIA V100 GPUs with 16GB memory. For the experiments on the KS equation, each model was trained on a single NVIDIA V100 for to 2 days. We note that since the model errors are becoming close to the float32 precision limit, the results may differ on other GPU architectures (e-g. A100), especially when different precision like tensorfloat (TF32) or float16 are used in the matrix multiplications. This can artificially limit the possible accurate rollout time a model can achieve. For reproducibility, we recommend using V100 GPUs. For the 2D Kolmogorov Flow dataset, we parallelized the models across 4 GPUs, with a training time of days. The speed comparison for the 2D Kolmogorov Flow were performed required roughly 250 GPU days, with additional 400 GPU days for development, hyperparameter search, and the supplementary results in Appendix B
NVIDIA 1100 GPU with 80GB memory. Overall, the experiments in this paper
18

---

C PDE-Refiner Pseudocode
In this section, we provide pseudocode to implement PDE-Refiner in Python with common deep learning frameworks like PyTorch [62] and JAX [6] The hyperparameters PDE-Refiner the number of refinement steps K, called num_steps in the pseudocode, and the minimum noise standard deviation mins called min noise_std. std. Further, the neural operator NO can be an arbitrary network architecture, such as a U-Net as in our experiments, and is represented by MyNetwork/ self www..eeeerol_orrtttoooooo in the code.
The dynamics of PDE-Refiner can be implemented via three short functions. The train, step function takes input training example of solution s(t) (named u_t) and the previous solution u(t Af) (named u prev). We uniformly sample the refinement step we want to train, and use the classical MSE objective if & O. Otherwise, we train the model to denoise u(t). The loss can be used to calculate gradients and update the parameters with common optimizers. The operation Landn_like samples Gaussian noise of the same shape as u_t. Further, for batch-wise inputs, we sample k for each batch element independently. For inference, we implement the function seeriit_nee__sosuuue
which iterates through the refinement process of PDE-Refiner. Lastly, to generate a trajectory from an initial condition u_ iinitial, the function rollout autoregressively predicts the next solutions. This gives us the following pscudocode:
class PDERefiner:
def __init__ (self, num_steps, min_noise_std). self nue_steps num_steps std min_noise_sta std
self www.neeeuuerrrttt
MyNetworkk...)
def train stel_, u_t, u.prev): randint (0, num_steps 1)
if k == 0:
pred = self nwuran target
u_prev, k)
else:
noise_std self www_mminnnnttttttteooeo ** (k self num_ntmmes) noise
u_t_noised u_t noise noise_std
pred = self www.aal tnised, u_prev, k)
target noise
loss (pred, target)
return loss
def predict_nnxx stw.eettsolutiimeee1f,, u.prev):
u_hat_t self neural, operator neural (zeros, set___like___ like(u_prev), u_prev, 0)
for in range(1, self num_stepa 1):
noise_std self min_noise_std ** (k self num_steps)
noise randn like (u_t) t_hat_t_noised u_hat_t
noise_std
pred self nwuraluullloooo couraaloperattrr( tu u_prev, k)
u_hat_t t_has___noisedd pred noise_std return u_hat_t
def rollout u_initial, timesteps):
trajectory [u_initial]
for in range (tinestee::):::: u_hat_t self, predictreict..mm trajectory. crajectory.append( (ppend((_hat_t) return trajectory
As discussed Section 3.1 PDE-Refiner can alternatively implemented diffusion model.
To demonstrate this implementation, we use the Python library diffusers (65) (version 0.15) in the pseudocode below. We create a DDPM scheduler where we set the number diffusion steps to number of refinement steps and the prediction type to v_prediction [70). Further, for simplicity,
19

---

we set the betas to the noise variances of PDE-Refiner. We note that in diffusion models and in diffusers, the noise variance o diffusion step l calculated as:
K II K=
o = 0 = II (x-
-
(1
0
min
2
x=
2 Since we generally use few diffusion steps such that the noise variance falls quickly, i.e. 0 2(++1/KK the product in above's equation is dominated by the last term a 2 Thus, the noise 2 Further, for k 0 and k K, the variances always
min
T min
variances in diffusion of Ri min the same since the product is 0 or a single element, respectively. needed, one could correct for the product terms in the intermediate variances. However,
show in Appendix [E.7] PDE-Refine is robust to small changes in the noise variance and no performance difference was notable. With this in mind, PDE-Refiner can be implemented follows:
diffusers. schedulers import DDPMScheduler
class PDERefinerDiffusionn
def __init__ (self, , num_atepe, min.noise_std):
betas imin_noise_are for in reversed (range (num_steps 1))]
(k num_stepa
self scheducee DDPMScheduler (num_traat ( Arrai_,timestepssmmmmttt + 1,
trained_ ses__betttbbbate
prediction_type sype='vuprediccionn
saip_ample--Fase)
self. num_ steps num_steps self self neural MyNetwork(...)
def train stee_sff, u.t, u_prev): randint (0, self num_steps + The scheduler uses t-X for first step prediction, and t-0 for minimum noúse. # To be consistent with the presentation in the paper, we keep l and the also the scheduler time step
scheduler time separate. However, directly and acts conditional input the neural operator.
scheduler_t self .um_stepp
noise_ factor = self schfuucced www.eeee....aapaaaaaaaaaea
signal_factor = noise_factor
noise Land__likeu__).
u_t_noised self scheduler noiseeet noise, scheduler_t) pred self operatr_ (u, tu_t_ tuissdd u_prev, k)
target (noise (notse_factor ** 0.5) noise (signal_factor ** 0.5) u_t loss mse (pred, target)
return loss
def saalcc__ext_soluttioneeeel u_prev):
i_hat___noised randn landn_like (u_prev)
for scheduler_t in self scheduchedller. timesteps:
self nel_tnumssteps scheduler_t pred self neurat_operator (u_ ttt_ttnnssssdeeee u_prev, k)
out self www.duler.sttepeeeeeeeeeeeeeeeee u_hat___noised out. uut.prev__smple u_hat_t u_hat_t_noised return u_hat_t
acheduler_t, t_hat_t_noised]
def (self, u_initial, timesteps):
trajectory = [u_initial] for range(timestege
u_hat_t self, aoradiitteeet.oluttootteeee srajectory.appentt (u__aatt_ hat_t)
return trajectory
20

---

Training examples
Suatad drveson
Spatial droweeees
Npotal dioeeiion
20.5
au-
20.2
2872
Spatull
64.75
16.3
00
14.5
211
50
347
213
15.7
213
0.0
340
26.
0.
20
Tane (0) seasal
Tax: (a soneab
Tane (asaceib)
Tiae (arancana)
Test examples
1214
1240
wooab)
dataset. The training trajectories generated with 140 time steps, while the test trajectories consist of 640 time steps. The spatial dimension is uniformly sampled from (0.9-64, 64, 1.1 164, and the time step seconds from [0.18, 0.22].
Figure 8: Dataset examples of the
Prequemey spectrum Kuramoto-Shsschissky trajectories
Precision
Float64 Precision
Jonrae aupiirole
10
30
10"11
10"11
10
10
10
10-
175
Worcomsserr
Wovcomber
Figure Frequency spectrum of the Kuramoto- -Sivasiiii dataset under different precisions. Casting the input data to float32 precision removes the high frequency information due to adding noise with higher amplitude. Neural surrogates trained on float64 did not improve over float32,, showing that it does not affect models in practice.
D Experimental details
In this section, provide detailed description of the data generation, model architecture, and hyper- parameters used our three datasets: vwrrmmto-SSvvshinsky (KS) equation, parameter-dependent KS equation, and the 2D Kolmogorov flow. Additionally, we provide an overview of all results with cor- responding error bars in numerical table form. Lastly, we show example trajectories for each dataset.
D.1 Kuramoto-Sivashinsk 1D dataset
Data generation. We follow the data generation setup of Brandstetter er al. 8 which uses the method of lines with the spatial derivatives computed using the pseudo-spectral method. For each trajectory our dataset, the first 360 solution steps are truncated and considered as a warmup for the solver. For further details the data generation setup, refer Brandstetter et
Our dataset can be reproduced with the public code of Brandstetter et al. [8]. obtain the training data, the data generation command in the repository needs to be adjusted by setting the number of training samples to 2048, and 0 for both validation and testing. For validation and testing, we increase --end_time=2000
the rollout time by adding the arguments --o0=1000 --nt_ --nt_effedh www.ffcctieeeeAA and setting the number of samples to 128 each. We provide training and test examples in Figure
http://wwwwoooooooeeoeooooee///n datpppppaattpprouuuceoddaaaataaa datasets- for- horttttsso----k--oot---. kuramoto-
shivashinsky- ks- eusutiutionn
21

---

Table 2: Detailed list of layers the deployed modern U-Net. The parameter channels next to layer represents the number of feature channels of the layer's output. The U-Net uses the four different channel sizes C1, c2 63, c4, which are hyperparameters. The skip connection from earlier layers in residual block is implemented by concatenating the features before the first GroupNorm. For the specifics of the residual blocks, see Figure]
Index
Layer
Encoder
Conv(kernel size=3, channels=c1, stride=1)
ResiduallBoccKcchnnnls===))
3
ResidualBlock(Channeeess
Conv(kernel size=3, channels=cj, stride==)
5
7
Convekernel size=3, channels=cy, stride=2)
ResidualBlocke (chammlellc
10
Conv(kernel size=3, channels=cg, stride=2)
11
BesidualBlocK(ccannell----
Middle block
13
ResidualBloc((chamelsss..)
14
ResidualBlockKChannesssa)
Decoder
15
ReiddualBlccKChaanelsss. skip connection from Layer 12)
16
Residal/BBecKCCcanellsssss skip connection from Layer 11)
17
ResidualBBlccK(channee==zz skip connection from Layer 10)
18
TranspseConvvllliionnkenne size=4, channels=c3, stride=2)
19
ResidulBBiocKKChanneessss skip connection from Layer 9)
20
ResidualBleXKKchannelssss skip connection from Layer
21
Residul/BlecK(Cannnee==z. skip connection from Layer 7)
22
TranposeCoovvolltionnnkneeel size=4. channels=cg. stride=2)
19
ResiduulBBecK(Channnesssz. skip connection from Layer 6)
20
Reiduu/Bloo((channelsse.. skip connection from Layer
21
Residaa/BlecKKChammeezzzz skip connection from Layer 4)
22
TranspoeCCanvolliiinnkkenee size=4, channels=cg, stride=2)
23
ResidualBBeccKCchammeeszzzz skip connection from Layer 3)
24
ResidualBloccKchaannbsszzj skip connection from Layer 2)
25
ResidualBBlocK(Channee-- skip connection from Layer
26
CroupNorm(chhnnellszz groups=8)
groupsp:)
27
GELU activation
28
Convolution(kernel size=3, channels=1, stride=1)
The data is generated with float64 precision, and afterward converted to float32 precision for storing and training of the neural surrogates. Since we convert the precision in spatial domain, causes minor artifacts in the frequency spectrum as seen in Figure 9 Specifically, frequencies with wavenumber higher than 60 cannot adequately represented. Quantizing the solution values in spatial domain introduce high-frequency noise which is greater than the original amplitudes. Training the neural surrogates with float64 precision did not show any performance improvement, besides being significantly more computationally expensive.
Model architecture. For all models in Section we use the modern U-Net architecture from Gupta et al. (22), which we detail in Table The U-Net consists of an encoder and decoder, which are implemented via several pre-activation ResNet blocks 25 26 with skip connections between encoder and decoder blocks. The ResNet block is visualized in Figure 10 and consists of Group Normalization [90], GELU activations (27), and convolutions with kernel size 3. The conditioning parameters At and A are embedded into feature vector space via sinusoidal embeddings, for
22

---

Conditioning features
Scale-and malle-and-oe Saile-and-Shiffe Shift
GroupNorm
Convolution
GroupNorm
Convolution
GELU
GELU
Feature map a.1+1
Feature map a
Figure 10: ResNet block of the modern U-Net 22 Each block consists of two convolutions with GroupNorm and GELU activations. The conditioning features, which are A,, Ar for the KS dataset and additionally v for the parameter-dependem KS dataset, influence the features via a scale-and- shift layer. Residual blocks with different input and output channels use convolution with kernel size on the residual connection.
Table 3: Hyperparameter overview for the experiments on the KS equation. Hyerparameters have been optimized for the baseline MSE-trained model on the validation dataset, which generally worked well across all models.
Hyperparameter
Value
Input Resolution
256
Number of Epochs
400
Batch size Optimizer
128
AdamW
52
Learning rate
CosineScheduler(1e- - le-6)
Weight Decay
1e-5
Time step
0.8s / /4tt
Output factor
0.3
Network
Modern U-Net [22
Hidden size
61 64, 62
128, 03- = 256, C4 1024
Padding
circular 0.995
EMA Decay
example used in Transformers 84 We combine the feature vectors via linear layers and integrate them in the U-Net via AdaGN
6 layers, which predicts scale and shift parameter for each channel applied after the second Group Normalization in each residual block. We represent it as We also experimented with adding attention layers the residual blocks, which, however, did not improve performance noticeably. The implementation of the U-Net architecture can be found in the public code of Gupta et 22
'scale-and-shiff layer in Figure
Hyperparameters. We detail the used hyperparameters for all models in Table 3 We train the models for 400 epochs batch size of 128 with an Adam optimizer (52). One epoch corresponds to iterating through all training sequences and picking 100 random initial conditions each. The learning rate is initialized with le-4, and follows cosine annealing strategy to end with final learning rate of le-6. We did not find learning rate warmup to be needed for our models. For regularization, we use weight decay of le-5. As mentioned in Section 4.J train the neural operators to predict 4 time steps ahead via predicting the residual Au u(t) u(t 4At). For better output coverage of the neural network, we normalize the residual standard deviation of about by dividing it with 0.3.
twod_ soidd_une.pp py
23

---

KS Rollout Loes
Correlation Rollont KS duraset
10'
1..9
Cumietio to Groand
0.8
80*
0.6
037*
0.4
MSE Traising Basdiue
03-1
Pushforrand
00.
Pocthewadd
Defecce 201 Betroe
0.2
PDE-Rdiie 4 stop (Ouns)
Diffusion
Codine
20
40
300
131
0.0
201
Rellor Tire Sky sourch
Rollant Time (In socords))
Figure 11: Left: Visualizing the average MSE error over rollouts on the test set for four methods: the baseline MSE-trained model (blue), the pushforward trick (green), the diffusion model with standard cosine scheduling (orange), and PDE-Refiner with 8 refinement steps. The markers indicate the time when the method's average rollout correlation falls below 0.8. The y-axis shows the logarithmic scale of the MSE error. While all models have a similar loss for the first 20 seconds, PDE-Refiner has much smaller increase of loss afterwards. Right: showing the average correlation over rollout time. Different thresholds would lead to the same findings in this paper.
Thus, the neural operators predict the next time step via ú(() ra(t 4At) + 0.3 NO(u(t 4At)). We provide an ablation study on the step size in Appendix E33 For the modern U-Net, we set the hidden sizes to 64, 128, 256, and 1024 on the different levels, following Gupta al. [22]. This gives the model a parameter count of about 55 million. Crucially, all convolutions use circular padding in the U-Net account for the periodic domain. Finally, we found that using an exponential moving average (EMA) 44 of the model parameters during validation and testing, as commonly used in diffusion models 30 38) and generative adversarial networks [16 9T], improves performance and stabilizes the validation performance progress over training iterations across all models. We set the decay rate of the moving average to 0.995, although did not appear to be sensitive hyperparameter.
Next, we discuss extra hyperparameters for each method in Figure 3 individually. The history 2 model includes earlier time steps by concatenating u(t- 8&t) with u(t 4At) over the channel dimension. We implement the model with 4x parameters by multiplying the hidden size by 2, i.e. use 128, 256, 512, and 2048. This increases the weight matrices by a factor of 4. For the pushforward trick, we follow the public implementation of Brandstetter et al. [ Ind the ground truth with a prediction over the first 10 epochs. Additionally, we found beneficial to use the EMA model weights for creating the predictions. and rolled out the model up to 3 steps. We implemented the Markov Neural Operator following the public codd of Li et al. [50]. We performed hyperparameter search over À (0.2,0.5, 0.8}. E {0.001, 0.01, 0.1), k E {0,1), for which we found o = 0.5, a 0.01, = 0 to work best. The error correction during rollout is implemented by performing an FFT on each prediction, setting the amplitude and phase for wavenumber O and above 60 to zero, and mapping back to spatial domain via an inverse FFT. For the error prediction, in which one neural operator tries to predict the error of the second operator, we scale the error back to an average standard deviation of I to allow for a better output scale of the second U-Net. The DDPM Diffusion model is implemented using the diffusers library [65]. We use a DDPM scheduler with squaredcos_cape cap_v2 v2 scheduling, a beta range of le-4 to le- 1, and 1000 train time steps. During inference, we set the number of sampling steps to 16 (cqually spaced between 0 and 1000) which we found to obtain best results while being more efficient than 1000 steps. For our schedule, we set the betas the same way shown in the pseudocode of Appendix C Lastly, we implement PDE-Refine using the diffusers library [65] as shown in Appendix C We choose the minimum noise variance of = 2e-7 based hyperparameter search the validation, and provide ablation study on it in Appendix E.7
and increase the probability of replacing
of
Results. We provide an overview of the results in Figure 3 as table in Table 4 Besides the high- correction time with thresholds 0.8 and 0.9, we also report the one-step MSE error between the prediction ii!) and the ground truth solution vu!). A general observation is that the one-step MSE is not a strong indication of the rollout performance. For example, the MSE loss of the history model is twice as low as the bascline's 's loss, but performs significantly worse in rollout. Similarly,
httpp: http://gghhhblooo///aaaatttttrer Sataauanw//---korrra----------- http://ggghhbcoooo////moppopppro
24

---

Table 4: Results of Figure 3 in table form. All standard deviations are reported over 5 seeds excluding Ensemble, which used all 5 baseline model seeds and has thus no standard deviation. Further, we include the average one-step MSE error of each method on the test set. Notably, lower one-step MSE does not necessarily imply longer stable rollouts (e.g. History versus baseline).
Method
Corr. >
0.8 time
Corr. > 0.9 time
One-step MSE
MSE Training
Baseline
75.4 1.1
66.5 0.8
2.70e-08 1 8.52e-09
History
61.7 + 1.1
54.3 = 1.8
1.50c-08 + 1.67c-09
4x parameters Ensemble
79.7 0.7
71.7 0.7
1.02c-08 + 4.91e-10
79.7 = 0.0
72.5
5,56c-09 = 0.00c+00
Alternative Losses
Pushforward [9]
75,4 = 1.1
1.7
2.76c-08 + 5.68c-09
Sobolev norm =0 50
71,4 2.9
6.92
1.33c-07 + 8.70c-08
Sobolev norm
[50
66.9
5.5
1.04e-07 1: 3.28e-08
Sobolev norm = 2 150
8.7 = 0.9
7.3
7.84e-04 +: 9.30e-05
Markov Neural Operator
50
66.6 1.0
58.5
2.66c-07 + 1.08c-07
Error correction [58]
74.8 1.1
66.2
1.46e-08 + 1.99c-09
Error Prediction
75.7 0.5
67.3 0.6
2.96e-08 +: 2.36e-10
Diffussoo Ablations
Diffusion Standard Scheduler
30
75.2 = 1.0
66.9 = 0.7 79.7 1.1
3.06e-08 08 + 5.24e-10 2.85e-09 + 1.65e-10
Diffusion Our Scheduler
88.9 1.0
PDE Refiner
PDE-Refines step (ours)
89.8 0.4
80.6
3.14e-09 + 2.85e-10
PDE-Refines 2 steps (ours)
94.2 + 0.8
84.2
5.24e-09 1 1.54e-10
PDE-Refiner - steps (ours)
97.5 + 0.5
87.0
5.80e-09 + 1.65e-09
PDF-Refiner 4 steps (ours)
98.3 0.8
87.8 1.6
5.95e-09 + 1.95e-09
PDE-Re ADE-Refiner 8 steps (ours)
98.3 0.1
89.0
6.16e-09 + 1.48e-09
PDE-Refing 3 steps mean (ours)
98.5 - 0.8
88.6 1.16111
1.28e-09 i 6.27e-11
the Ensemble has lower one-step error than PDE-Refiner with more than refinement steps, but is almost 20 seconds behind in rollout.
As an additional metric, we visualize in Figure II the mean- mean-squared error loss between predictions and ground truth during rollout. In other words, we replace the correlation we usually measure during rollout with the MSE. While PDE-Refiner starts out with similar losses the baselines for the first 20 seconds, it has significantly smaller increase in loss afterward. This matches our frequency analysis, where only later time steps, the non-dominant, high frequencies start to impact the main dynamics. Since PDE-Refine can model these frequencies in contrast to the baselines, it maintains smaller error accumulation.
Speed comparison. We provide speed comparison of an MSE-trained baseline with PDE-Refiner on the KS equation. We time the models on generating the test trajectories (batch size 128, rollout length 640Att on an NVIDIA A100 GPU with a 24 core AMD EPYC CPU. We compile the models in PyTorch 2.0 62]. and exclude compilation and data loading time from the runtime. The MSE model requires 2.04 seconds (=0.011, while
with refinement steps takes 8.67 seconds (+0.01). In contrast, the classical solver used for data generation requires on average 47.21 seconds trajectory, showing the significant speed-up of the neural surrogates. However, it should be noted that the solver implemented on CPU and there may exist faster solvers for the 1D Kuramoto- Sivashinsky equation.
D.2 Parameter-depenneet KS dataset
Data generation. We follow the same data generation as in Appendix D.1 To integrate the viscosity v, we multiply the fourth derivative estimate Uzzza by v. For each training and test trajectory, we uniformly sample between 0.5 and 1.5. We show the effect of different viscosity terms in Figure [2
25

---

Training trajectories
701
k-11.5
63.0
P=1.00
071
P-1.44
Spotid dirciinn
Spotid dardddr
Steeed
20.1
0
18.0
210
00
13.
305
30
135
27.5
0.
12.1
28.3
00
13.
31.6
30
143
101
Time soundh)
Tix (6n wonld))
Tixe seocak)
Trae (oroone)))
Tux (6 wordd)
Test trajectories
25
25
Figure 12: Dataset examples of the parameter-dependent Muramoto-Sivashinse dataset. The viscosity noted above each trajectory. The training trajectories 140 time steps, while the test trajectories are rolled out for 1 140 time steps. Lower viscosities generally create more complex, difficult trajectories.
Table 5: Results of Figure in table form. All standard deviations are reported over seeds.
Method
Viscosity
Corr. 0.8 time
Corr. 0.9 time
0.55, 0.7)
41.8 + 0.4
35.6 = 0.6
0.7, 0.9)
57.7 + 0.6
50.7 = 1.3
MSE Training
0.9,
73.3 + 2.3
66.0 2.5
1.1,
88.0 + 1.5
76.7 + 2.2
1.5
97.0 + 2.7
85.5 2.2
0.5, 0.7)
53.1 + 0.4
46.7 + 0.4
0.77
71.4
64.3 = 0.6
PDE-Refiner
0.19
94.5 0.6
84.9 =
1.1,
112.2 + 0.9
98.5 + 1.5
1.5
130.2 + 1.5
116.6 0.7
Model architecture. We use the same modern U-Net as in Appendix D.1 The conditioning features consist of At, Az, and For better representation in the sinusoidal embedding, scale the range [0, 100] before embedding it.
Hyperparameters. We reuse the same hyperparameters of Appendix D.I number of epochs to 250. This is since the training dataset is twice as large as the original KS dataset, and the models converge after fewer epochs.
except reducing the
Results. We provide the results of Figure 7 in table form in Table 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 [79] as detailed the publicly released code For hyperparameter tuning, we additionally generate 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.
26

---

Table Hyperparameter overview for the experiments on the Kolmogorov 2D flow.
Hyperparameter
Value
Input Resolution
64x64
Number of Epochs
100
Batch size
32
Optimizer
Adam W [22
Learning rate
CosineScheduler( 1e-4 le-6)
Weight Decay
1e-5
Time step
0.112s 16At
Output factor
0.16
Network
Modern U-Net [22]
Hidden size
[128, 128, 256, 1024] circular
Padding
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, 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 128, and use ResNet blocks instead of on this level. All other levels remain the same as for the KS equation. This model has 157 million parameters.
The Fourier Neural Operator convolution with skip connection of a 1 x 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 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.
49 consists of layers, where each layer consists of a spectral
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 time step of 1, 2, 4, and 8 while being faster to roll out. All models use as objective the residual Au= utf u(t- 16At), which normalize by dividing with its training standard deviation of 0.16. Thus, we predict the next solution via ú(f u(t 16 0.16 NO(...). Each model is trained for 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. 4K Sun al. (79]. For the U-Nets, we PyTorch 2.0 [1]. All models are compiled in their respective frameworks, and we exclude the compilation and time 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
27

---

Ground Truth
0.00
3.48g
7,07s
10.66s
14.258
10
16
10
10
0
10
-10
10
10
-10
PDE-Refiner
1= 0.00s
3.488
$ 7,07s
1= 10:08s
14.26
10
10
0
Ground Truth
14.26e
10
-10
-10
PDE-Refiner
0.00s
3.486
7.0%
10.68s
14.25
10
10
10
O
0
o
10
-10
10
-10
Ground Truth
14.20a
PDE-Refiner 7.07%
0.003
3.48s
10.68s
14.25e
10
10
10
-0
10
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 10.66 seconds. Moreover, many structures 14 seconds are still similar the ground truth.
28

---

High-Correlation Rollout Times FNOs on the Kuramnto--Swwahiinsky equation
l[gh col--corriattnn (nnocomcc))
Training Alternative Losses
PDF-Refieer (Ouns
Dilinsion Ablaiiu
83 7s
85
85.3
85.3
85.4x
0s
70.5+
75.0x
73
70.06
7ads 09:0c
70
67.26
Bexeion
Histoy
as parmuates
Exeeable
Prefitamard
Subdlr A Soluder A F
Eina Curretion Extr Prefecion
MNO
stop
adule
stuge
stops
stopa Moan
Comion Schedale
Schedole
Figure 14: Experimental results of Fourier Neural Operators on the Sarramto-SSvashinee equation. All methods from Figure 3 are included here. FNOs achieve similar results as the U-Nets for the baselines. For PDE-Reliner and Diffusion, FNOs still outperforms the baselines, but with 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- Sivassii-iv 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 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 architec- tures. To show that the general trend of our results in Section 4.1 are iisiitectur-invariian we repeat all experiments of Figure 3 with FNOs. The FNO consists of layers, where each layer consists of spectral convolution with skip connection of a 1 x 1 convolution and a GELU activation [27]. Each spectral convolution uses the first 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 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 160 million did not show any improvement. Still, the ensemble of MSE-trained models obtains 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-frequen noise in PDE-Refiner's refinement objective to be the main issue. U-Nets more flexible in that regard, since they 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 Ref-Refiner.
29

---

High-Correlatton Rollout Times over Number of Fourier Modes
High-coreilttie timne (in socands)
MSP: Training
PDD-RRffier (Ouns)
100
84.4s
73.06
74.
70.04
I
60
16
128
32
Numbbrr of Modies in FNO
Figure 15: Investigating the impact of the choosing the number of modes in FNOs. Similar to analysis on the resolution in the U-Nets (FigureB). we only see minor improvements using higher frequencies above 16 in the MSE training. Removing dominant frequencies above significantly decreases performance. Similarly, increasing the modes of FNOs in PDE- Refiner has minor impact.
High- Correlation
Dilated on Kuramoto-Sivashinsk, equation
llgh ectelerreetee secuds))
110
Traising
PDE-Rdizar (thus)
988
90.Sw
90.9
90.5
Altornation Tooors
Diffusko Ablations
90
60.3+
Sl,ls
79.7
78.1
74.55
77.3k
N7. Hs
70
66.1 Is-
Baseltm
Hissar
parmetaa
Eneemble
Solaker w
Ernu Curretion
MNO
Exror Prediction
stop
atape
stepe
stape
a stops Moun
stupp
Gosine Schodale
Our Schethile
Sobaler
Figure 16: Experimental results of Dilated ResNets on 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 significant margin.
Number of Fourier Modes. 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
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 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 modes results in significantly worse rollouts. This is also in line with our input resolution analysis of Figure where the MSE-tran 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 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, 1]. Since the convolutions do not use any stride, all operations performed on the original spatial resolution, increasing computational cost over e.g. a U-Net with similar parameter count.
For our setup the KS dataset, we a Dilated ResNet with channel size 256, group normalization between convolutions and shift-and-scal conditioning as in the U-Nets. Additionally, we change the default ResNet architecture using pre-activations (26), allowing for better gradient flow and
30

---

High-Correlation Rollout Times over Predicted Step Size
High corrdation seconds
MSE Troinnn
PDE Refure (Oun)
87.64
95.8t
77.5n
75.4s
71.4s
66.1s
60
98.74
40-
0.224
0.4:
0.86
1.6s
3.2g
6.46
12.8s
0.2:
0.44
0.86
1.66
Step Nar
Figure 17: Comparing the accurate rollout times over the step size at which the neural operator predicts. This is multiple of the time step AL 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 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 perform 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.
Due the experiments being computationally expensive, show most results for single seed only here, but the trends are generally constant and have not shown 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 obtain the longest accurate rollouts, similar the U-Nets. This verifies the strength of FDE-Refined and its applicability across various trat--ofthhe-- the-art neural operator architectures.
We show all results in Figure T6
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 eтror propagation in an autoregressive solver with training a model to predict large time steps, we repeat operator the KS equation with different step sizes. The default step size that was used in Figure 4-times the original solver step, being on average 0.8s. For any step size below 2s, we model residual objective Au u(t) u(t At), which found to generally work better in this range. For any step size above, we directly predict the solution u(t).
baseline experiments of the U-Net neural
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 diffferntt step sizes, we plot in Figure 18 the MSE loss 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 achieve very similar losses across whole time horizon. This motivates our choice for 0.8s as default time step, since it provides 4 times speedup in comparison to the 0.2s model. Meanwhile, already a model trained with step size 1.68 performs considerable worse in its one-step prediction than model with step size 0.2s rolled out times. The gap increases
31

---

KS Rollout Loss over Predicted Step Size
10"2
1At 0.2s
164r
10-1
2A3 = 0.4x
32M2 6.4x
MSC Loss Ground Trath
dAt 0.8s
6tå: 12.84
10-1
8åt 1.6
10"3
10-4
10
10
Rollout Time Step in scconds
Figure 18: Visualizing the MSE error of MSE-trained models with varying step sizes over the rollout. The models with step size of AI, 221, and 4&/ obtain similar performance. For SÂ1, the one- step MSE loss is already considerably higher than, e.g. rolling out the step size 1At model 8 times. For larger time steps, this gap increases further, again highlighting the strengths of autoregressive solvers.
Iligh-Correlation Rollout Times over Input History
90
Bigh eurrelation time [in socoted))
Time stop
0.84
Time stop At
80
75.
70
61.74
60-
I
56 74
I
I
47.9k
50-
I
40
History
History
Hiskory
History
History
History 4
Number conditioning time stege history
Figure 19: Investigating the impact of using more history / past time steps in the neural operators, i.e., â(t) = NO(u(t At), u(t 2 for At = 0.8 and At 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 neural operator that took as input the past two time steps, u(t At) and u(t 2.t).. To confirm this trend, we repeat the experiments with longer history of 4 past time steps and for models with 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 larger issue when using history information. The problem arising is that the difference between the inputs u(t At) utt 2 is highly correlated with the model's target Au(t), the residual of the next time step. The smaller the time step, the larger the correlation. This leads the neural operator focus modeling the second-order difference Au(t) Au(t 2At). 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 error by plotting the MSE error over rollouts in Figure 20 While the history information improves the one-step prediction by factor of 10, the enror of the history 2 and models quickly surpasses the error of the history I model. After that, the error of the models continue to increase quickly, leading to earlier divergence.
32

---

KS Rollout Loss over Input History
Rollout Loss over Input History
Histoey
MSE: Lowe to Cound Thal
10-1
History
MSB Lose to Genaal Truth
10-7
-0--
10-1
-0
10"1
-0"
11
Thstory
10-12
0.0
0.5
1.0
1.5
2.0
2.5
3.11
3..5
4.0
20
50
100
Rollout Time Step accondss
Ralkott Time Step in sooonds
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 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 ~5 steps, the error propagates significantly faster through the models. This leads significantly higher loss over rollout.
Table 7: Comparing the uncertainty estimate of PDE-Refiner to Input Modulation S 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² coefficient of determination and the Pearson correlation. PDE-Refiner provides more accurate uncertainty estimates than Input Modulation while being more efficient than Model Ensemble.
Method
R coefficient
Pearson correlation
PDE-Refiner
0.857 + 0.027
0.934 + 0.014
Input Modulation
71
0.820 0.081
0.912 + 0.021
Model Ensemble
45
0.887 0.012
0.965 + 0.007
High-coecaatie line (is secuna))
100
Litser fit Trayeturtes
243
Theecores
100
60
Orreecomlbtion thrw (In secondee
50
330
(a) PDE- DDf-Refnner (Ours)
(b) Input Modulation
(c) Model Ensemble
Figure 21: Qualitative comparison between the uncertainty estimates of PDE-Refiner, Input Modu- lation, and the Model Ensemble. Both PDE-Refiner and the Model Ensemble achieve an accurate match between the estimated and ground truth rollout times.
E.5 Uncertainty Estimation
We extend our discussion the uncertainty estimation of Section 4.1 by comparing PDE-Refiner to two common baselines for uncertainty estimation of temporal forecasting: Input Modulation [T] 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-Refines one can determine the uncertainty by measuring the correlaton 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, estimate the uncertainty of MSE-trained models usually applied.
evaluate the R² 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
33

---

Samples
Differences
10'
10
Std between samples
Defferencc to ground truth
Ampittude
10
Amplitude
20
10-3
10:43
10
10-9
10-7
10-1-
25
50
75
100
125
25
50
75
100
125
Wavemmbeer
Wavermmbe
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 the mean of the samples. The deviation of the samples closely matches the average eπror, showing that over frequencies.
adapts its samples to the learned error
Frequency Spectrum on the Kolmogorov Flow
Channel 0
Channel
10F
10"
Ground Truth
PDE-Rofaa -Sapps
10
10'
PNE-Rofiur- Freers
100
Laining Sumples
MSF Training - Frrors
10
10:4
10
16-
10
10-1
10
10
30
10
30
Wavemmber
Waveembber
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 DDE-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 [2] The samples of PDE-RRfiner closely differs in the same distribution the actual loss the ground truth, showing that PDE-Refinet 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 on the Kolmogorov dataset here. Note that we apply 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 at wavenumber 32. In comparison, the KS equation has more than times as small amplitude for this wavenumber. Further, since the resolution is only 64 x 64, higher modes cannot be modeled, which, as seen on the KS equation, would increase the benefit of RDF-Refiner. This leads both PDE-Refiner and the MSF-trained baseline to model all frequencies accurately. The slightly higher loss for higher frequencies channel 0 is likely due to missing high-
34

---

Frequency Spectrum Intermediate Samples on the Kohmogurov Flow
Channel
Channel
10P
10
Ground Truth
Poediction
10
10
Redacmeezt
Anplitude
Auplitede
10P
Reénement step 2
Refirement step
10-1
10r1
MSE Training Erroes
10-3
10:2
10
10
10
Waveuunlee
Wavenmmber
Figure 24: Frequency spectrum of intermediate samples in the refinement process of PDE-Refiner, similar to Figure [ for the KS equation. The refinement process improves the prediction of the model ttsppb-step. For the last refinement step, we actually see minor improvements for the lowest frequencies of channel 0. However, due flatter frequency spectrum, the high frequencies do not improve as much as on the KS equation.
Higb-Correlation Rollout Times Minimum Noise Variance
Higl-morhuatin tine [In wooods)
120
PDE-Reffee (Ours)
110
100
Ds
07.5e
98.0s
98.5s
50.2s
88.96 9s
68.3s
80
40-7
46-6
1o-5
16-4
Minimum Noise Vartance dii
Figure 25: Plotting performance of PDE-Refiner over different values of the minimum noise variance minn- Each PDE-Refiner robust to small changes of om... showing an equal performance in the range of [10-7 10-6]. Higher standard deviations start decrease the performance, confirming our analysis of later refinement steps focusing low-amplitude information. For the experiments in Section we have selected 2² =2e-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 MSS--trained model on all frequencies.
Figure 24] we additionally plot the predictions PDE-Refine 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 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
Besides the number of refinement step, PDE-Refiner has as second hyperparameter the minimum noise variance o2 i.c., the variance of the added noise in the last refinement step. The noise variance determines the different amplitude levels at which PDE-Refiner improves the prediction. To show how sensitive PDE-Refing is to different values of 0 we repeat the experiments of PDE-Refiner on the KS equation while varying 22 The results in Figure 25 show that PDE-Refiner is robust to small changes of 2 and there exist larger band of values where it performs equally well. When increasing the variance further, the performance starts to decrease since the noise too high model the lowest amplitude information. Note that the results Figure 25 show the performance on the test set, while the hyperparameter selection, in which we selected 2 min 2e-7, was done on validation set.
35

---

Spettee Pridation a
Spetteen Prokerm at 0.
Specteen Treckt ww M8D5
Oheme Mual
Coom
Croad Pert
Paki Mual)
202 Sutu
(a) step
(b) 125 steps
(c) 500 steps
(d) 1000 steps
Opecttmaa L'rdtttiinn 0,d0
Cpectrea Prodectzze
Spectmen of Prodkctse 06.00
Spettmen of Trickt.n DRION
Sober 17.00
Rive
(c) step
(f) 125 steps
(g) 500 steps
(h) 1000 steps
Figure 26: Evaluating PDE solver stability over very long rollouts (800 seconds, corresponding to 1000 autoregressive prediction steps). (a-d) The frequency spectrum of predictions of an ESSE-triined model and PDF-Refiner. Over time, the MSE baseline's overestimation of the high frequencies accumulates. In comparison, PDF-Refiner shows to have an increase of extremely high frequencies, which is likely caused by the continuous adding of Gaussian noise. (e-h) When we apply the error correction (58) our models by setting all frequencies above 60 to zero, PDE-Refiner remains stable even for 1000 steps and does not diverge from the ground truth frequency spectrum.
In combination with the hyperparameter of the number of refinement steps, which PDE-Refiner showed to also be robust if more than 3 steps is chosen, PDE-Refiner is not very sensitive to the newly introduced hyperparameters and values in larger range can be considered.
E.8 Stability of Very Long Rollouts
Besides accurate rollouts, another important aspect of neural PDE solvers is their stability. This refers to the solvers staying in the solution domain and not generating physically unrealistic results. To evaluate whether our solvers remain stable for a long time, we roll out an MSE-trained baseline and PDE-Refiner for 1000 autoregressive prediction steps, which corresponds 800 seconds simulation time. We then perform a frequency analysis and plot the spectra in Figure 26 We compare the
spectra the ground truth initial condition, have reference point of common frequency spectra of solutions on the KS equation.
For the MSE-trained baseline, we find that the high frequencies, that are generally overestimated by the model, accumulate over time. Still, the model maintains a frequency spectrum close to the ground truth for wavenumbers below 40. DDE-Refiner maintains accurate frequency spectrum for more than 500 steps, but suffers from overestimating the very high frequencies in very long rollouts. This is likely due to the iterative adding of Gaussian noise, that accumulates high-frequency errors. Further, the Ü-Net has a limited receptive field such that the model cannot estimate the highest frequencies properly. With larger architectures, this may be preventable.
However, a simpler alternative is to correct the predictions for known invariances, as done in McGreivy et al. [58]. We use the same setup for Figure 3 by setting the highest frequencies to zero. This stabilizes PDE-Refiner, maintaining a very accurate estimation of the frequency spectrum even at 800 seconds. The MSSE-traine model yet suffers from overestimation of the nigh-frequencies.
In summary, the models we consider here are stable for much longer than they remain accurate the ground truth. Further, with a simple error correction, PDE--Refine can keep up stable predictions for more than 1000 autoregressive rollout steps.
36


Reviews:
[{"review_id": "oppXRFgJR2", "reviewer": "Reviewer_K6r8", "summary": "Recently, there has been active research on using deep neural networks as surrogates to predict the temporal evolution of solutions in time-dependent partial differential equations (PDEs). Improving the accuracy of long-term predictions for surrogate models that take previous time steps' solutions as inputs and predict the solution at the next time step is a critical challenge. In this paper, the authors propose a PDE-Refiner to address this issue and enable accurate prediction of solution over long time periods.", "questions": "* Do you think that PDE-Refiner is related to the neural operator method mentioned in [2]? If so, what are the similarities and differences between the autoregressive and temporal bundling approaches in [2] and the PDE-Refiner? Can these two models be applied to PDE-Refiner? What is the relationship between them?\n\n* Are wavenumber and frequency ultimately the same thing?\n\n* Looking at Figure 2, it seems that the model should involve the index $k$. How is this index incorporated? Since $\\sigma_k$ will have different effects on the input as $k$ changes, is this index necessary as an input?\n\n* In Equation (4), shouldn't the '$u(t)$' after 'NO' be $\\hat{u}^k$ instead?\n\n* As mentioned in line 203, the experiment was conducted with predictions made at intervals of $4\\Delta t$. Would the results differ if the interval was shorter? How would they differ if it was longer? What range of variations is permissible?\n\n* If more than two histories are used for History2 in Figure 3, how would the results change? Additionally, if predictions are made for multiple time steps, how would the results differ?\n\n* The graph on the far right in Figure 4 is unclear in its meaning.", "limitations": "I am curious about the reason for presenting all experimental results in terms of correlation time rather than accuracy. It is possible that I may have misunderstood, but addressing the previous questions would provide a clearer understanding of the purpose and direction of the paper.", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "In addressing the issues arising from using MSE loss to predict the solution $u(t)$ at the next time step based on the previous solution $u(t-\\Delta t)$, the authors conducted an analysis of the amplitude at different frequencies (Figure 1). To overcome these challenges, they introduced the denoising idea and developed the PDE-Refiner, enabling iterative prediction of the long-term evolution of solutions, $u(t)$. The various experiments in Section 4 confirmed the effectiveness of the PDE-Refiner in capturing high-correlation time over extended periods.", "weaknesses": "I appreciate the thorough review of the paper. It seems that there are some confusing aspects that could benefit from further clarification. Specifically, the relationship between denoising and the diverse amplitudes of frequencies in the context of enabling long rollouts for PDE solution prediction is not well-explained. Sections 3 and 3.1 touch on this topic, but it would be helpful to bridge the gap and provide a clearer explanation of how the introduced Gaussian noise, when added and subtracted, enables denoising across various sigma values, thus facilitating long rollouts for different frequency amplitudes. This clarification would greatly assist in understanding the interrelation between these concepts.\n\nFurthermore, I am curious to know how the proposed PDE-Refiner, in combination with other models such as FNO and Graph Neural Networks (GNN) that predict the next time step based on previous time steps, improves accuracy for long-term predictions compared to existing models. Although it is mentioned that combining FNO poses challenges due to the elimination of high frequencies, I believe that using the learned parameters from all frequency domains in FNO could mitigate this issue. It would be valuable to see experimental results showcasing the improved accuracy for longer time predictions when applying PDE-Refiner to models like FNO and graph-based approaches [1, 2]. Additionally, it would be beneficial to have a discussion on the reasons for focusing solely on high-correlation time instead of explicitly addressing prediction accuracy. Furthermore, given the frequent mention of Pearson correlation, a more detailed explanation of its relevance and its interpretation in all experimental results would be beneficial.\n\n\nWhile it is mentioned that the U-net, which serves as the backbone for the current experiments, exhibited a smaller correlation time when using the pushforward trick from [2] in Figure 3, it is crucial to assess the PDE-Refiner method's performance in predicting long-term solutions for time-dependent PDEs across various models proposed in the field. In order to substantiate the claim that the PDE-Refiner method is superior, it is necessary to compare the accuracy of long roll-out predictions for PDE solutions using the proposed PDE-Refiner in this paper and other models employing different tricks in the context of time-dependent PDEs. Alternatively, if we assume that each network model has its own suitable trick and consider that the proposed PDE-Refiner aligns well with the U-net model, it would be appropriate to compare the accuracy of time-dependent PDE solution predictions between the U-net+PDE-Refiner approach in this paper and other models+tricks discussed in separate papers.\n\nLastly, I am curious about the necessity of associating PDE-Refiner with the Diffusion model. While both involve denoising, I wonder if it is necessary to establish a connection between the two (Section 3.1). While they share similarities in terms of denoising, it would be interesting to understand if linking them together is crucial to the novelty of the paper.\n\nOverall, addressing these points and providing additional explanations would greatly enhance the clarity and comprehensibility of the paper.\n\n[1] Boussif, Oussama, et al. \"MAgnet: Mesh agnostic neural PDE solver.\" Advances in Neural Information Processing Systems 35 (2022): 31972-31985.\n\n[2] Brandstetter, Johannes, Daniel Worrall, and Max Welling. \"Message passing neural PDE solvers.\" arXiv preprint arXiv:2202.03376 (2022).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive review and feedback. We provide below clarifications to your questions and concerns raised.\n\n**Clearer explanation of how Gaussian noise enables denoising across sigma values / Equation (4), shouldn't the $u(t)$ after 'NO' be $\\hat{u}^k$ instead?**\n\nThe refinement process of PDE-Refiner is trained by adding noise to ground truth data, $u(t)$ (Eq. (4)). This effectively removes information of frequencies below the amplitude of the noise, $\\sigma_k$, which the model learns to recover. Frequencies above the amplitude are only minorly changed, such that the model learns to leave them mainly untouched. Thus, PDE-Refiner is trained to focus on different amplitude levels across the sigma values in the refinement process. \n\nDuring inference, we replace the ground truth data by the model's own prediction, $\\hat{u}^k$, as depicted in Figure 2. Based on its training, PDE-Refiner will again focus on information below $\\sigma_k$ and ignore potential errors in higher amplitudes. This is since the model was trained on ground-truth data, and learned to assume higher amplitudes to be correct.\n\nYour question raises a good point, since the mismatch of Figure 2 showing the inference and Equation (4) the training is slightly confusing. We will emphasize and clarify this difference.\n\n**Experimental results when applying PDE-Refiner to models like FNO and graph-based approaches**\n\nWe agree and have provided FNO studies in Fig 13 and FNO ablations in Fig 14 in the submitted Appendix E.1. Further, in the rebuttal, we have added Dilated ResNet results (Figure 1). U-Nets, FNOs, and Dilated ResNets are the best performing models in both comprehensive overviews of Gupta \\& Brandstetter [2022] and Stachenfeld et al. [2022]. On all methods, PDE-Refiner obtains a significant improvement. Regarding graph-based approaches: The proposed datasets are for regular gridded data. Non-regular data with graph-based approaches is left for future work.\n\n**Discussing reasons for focusing on high-correlation time instead of prediction accuracy**\n\nWe are interested in how long the model stays accurate. This can be measured by high-correlation time, as in previous works, with thresholds of 0.8 and 0.9 indicating when the prediction starts to visually slightly differ from the ground truth. Alternative, we could check when the prediction exceeds a certain MSE loss, as shown in Fig 10 in Appendix D.1. Both measures come to a similar conclusion, while correlation is (1) invariant to scale and thus better comparable across PDEs, and (2) limited to [-1,1], such that it cannot be dominated by a single trajectory.\n\n**Comparison of PDE-Refiner and other models/tricks**\n\nWe have done comparisons for different MSE training strategies (different history inputs, larger models, MSE ensembles), and for different alternative losses (pushforward, Sobolev, MNO, error correction, and error prediction) - for both U-Nets and FNOs, and added Dilated ResNets for the rebuttal. While some methods benefit more from FNOs than U-Nets (e.g. MNOs) or vice versa, they still stay all in the same performance range and are outperformed by PDE-Refiner. We are happy to add losses beyond those studied if suggested by the reviewer.\n\n**Associating PDE-Refiner with Diffusion crucial for novelty.**\n\nIt is not strictly necessary, since PDE-Refiner has a different motivation than Diffusion models. We point out the connection for readers familiar with Diffusion models and clarify the differences.\n\n**PDE-Refiner is related to autoregressive modeling and temporal bundling of Brandstetter et al.? / If predictions are made for multiple steps, how would the results differ?**\n\nTemporal bundling (TB) can be easily applied to PDE-Refiner, we therefore denoise all outputs at the same time. Hence, we see TB as an orthogonal trick for obtaining stable rollouts. A possible combination of TB and PDE-Refiner would be to treat time as an extra dimension and thus refine over spatial and temporal frequency components, but be computationally expensive for long inputs. We add results for predicting up to 4 time steps in TB in rebuttal Figure 3. The performance slightly worsens due to the larger predicted time step.\n\n**Are wavenumber and frequency the same thing?**\n\nCorrect! Wavenumber and frequency are defined up to some factor as the number of wavelengths per (unit) distance.\n\n**How is the k index incorporated?**\n\nThe index $k$ is a conditional parameter passed to the model. We follow the standard implementation in Diffusion models of Ho et al. [2020] and the conditioning parameters in Gupta et al. [2022] by adding $k$ as input to each residual block in U-Nets, and condition the parameters in FNOs. A sketch of the conditioning on $k$ is provided in Fig 9 in the appendix.\n\n**Would results differ if the interval was shorter/longer than $4\\Delta t$?**\n\nSee Appendix E.2 for a detailed ablation study. Smaller time steps generally give better results, but are more expensive to roll out. Larger time steps lead to suboptimal results and more difficult generalization.\n\n**If more than two histories are used, how would the results change?**\n\nSee Appendix E.3. We found more history gets even worse, especially for small time steps.\n\n**Graph on the right in Fig 4.**\n\nIt shows the average frequency band of white noise that we sample at the different refinement steps, which has different orders of magnitude. In other words, it is FFT$(\\sigma\\_k\\cdot\\epsilon\\_k)$. It gives a relation between the error and noise per step in the refinement process. We will clarify it in the paper.\n\n---\n\n[Gupta \\& Brandstetter, 2022] Gupta, Jayesh K., and Johannes Brandstetter. Towards multi-spatiotemporal-scale generalized pde modeling. arXiv preprint arXiv:2209.15616 (2022).\n\n[Ho et al., 2020] Ho, Jonathan, et al. Denoising diffusion probabilistic models. NeurIPS (2020).\n\n[Stachenfeld et al., 2022] Stachenfeld, Kimberly, et al. Learned coarse models for efficient turbulence simulation. ICLR (2022)."}, {"review_id": "OLalBzQdNB", "reviewer": "Reviewer_YYxA", "summary": "This work introduces a PDE-Refiner model to address the challenge of accurately predicting the process of time-dependent partial differential equations (PDEs) over long rollouts. The model can be divided in two parts: prediction (based on FNO or UNet) and multi-step refinement (based on the diffusion model). By refining the prediction that includes non-dominant spatial frequency information, the model outputs stability and accuracy during rollouts. The model is validated on 1D Kuramoto-Sivashinsky (KS) and 2D Kolmogorov equations.", "questions": "1. As introduced, there are three categories of PDE surrogates. The U-Net and FNO are neural approaches. Can your method be applied to hybrid methods as well?\n2. Why is the duration of the high correlation of DNS methods shorter than the neural model, even in the same resolution? ", "limitations": "Major limitations were covered in the paper.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Good presentation: the categories of recent works are clear, and the challenges of accurate long rollouts description are easy to follow. The methodology is constructed in an easy-read way.\n2. Analysis of the long-term prediction error based on the frequency spectrum is interesting and can inspire future research.\n3. Using the diffusion model to deal with low-amplitude error in PDE settings is natural and reasonable and novel to the best of my knowledge.", "weaknesses": "1. The statement “the primary errors in predicting a one-step solution arise from inaccuracies in modeling the dynamics of these low frequencies.” comes from the observation of simulation results in the 1D KS equation only. However, no further analytical support was provided for other settings, such as in the 2D cases.\n2. Even though the evaluation speed is acknowledged by the author as a limitation, the duration of high correlation improvements ($10.659$ of the proposed model v.s. $9.663$ for the MSE training with UNet) shown in Table 1 could be considered as minor, given the inference time is $4 \\times$ longer. It is not clear whether the proposed approach is Pareto efficient in terms of the speed / accuracy tradeoff.  ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for helpful review and thoughtful feedback. Please find below our response to your questions and concerns raised.\n\n**Analytical support for 2D case for the statement “the primary errors in predicting a one-step solution arise from inaccuracies in modeling the dynamics of these low frequencies.”**\n\nWe have provided a similar analysis to the main paper for the 2D Kolmogorov flow in Appendix E.5. The statement generally holds for moderate/high-resolution inputs, and typically biased frequency spectra. The 2D Kolmogorov dataset is low resolution and a bit less biased, so effects are not as strong as for the 1D Kuramoto-Sivashinsky equation, but still low frequencies are dominant over high frequencies.\n\n**It is not clear whether the proposed approach is Pareto efficient in terms of the speed / accuracy tradeoff - Table 1.**\n\nThis is a fair point. We argue that PDE-Refiner allows for a flexible selection of speed vs accuracy tradeoff. For the current state-of-the-art, it provides a closer to Pareto efficient solution, since we showed in Figure 3 that increasing the parameter size is not helping for standard MSE models (which would effectively also require 4x more compute). And thus, even with 4x longer inference, it is faster than current hybrid approaches and more than 100x faster than the classical solver. We finally argue that the presented 2D Kolmogorov dataset is already quite exhaustively studied and the obtained correlation improvements are significant.\n  \n**Can PDE-Refiner be applied to hybrid methods as well?**\n\nThe refinement process of PDE-Refiner could be applied to learned correction (LC) [Kochkov et al., 2021] as well, since the correction model will suffer from the same biases. Other hybrid methods aim to predict parameters within classical solver, which do not really fit in the framework and may not suffer from the biases as much as plain neural approaches.\n\n**Why is the duration of the high correlation of DNS methods shorter than the neural model, even in the same resolution?**\n\nThe ground truth data is DNS generated at a resolution of 2048x2048. Running DNS on smaller resolution will naturally miss data which is difficult to interpolate for these methods, and thus diverge from the ground truth earlier. Neural methods are able to predict the missing high-resolution information much better than simple interpolation, as shown in Sun et al. [2023] on this dataset. This is since certain elements of the PDE such as the equation, the Reynolds number and structure of the initial conditions (Gaussian blobs) are fixed. Thus, neural models can outperform DNS on various resolutions.\n\n---\n\n[Kochkov et al., 2021] Kochkov, Dmitrii, et al. Machine learning–accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences 118.21 (2021).\n\n[Sun et al., 2023] Sun, Zhiqing, et al. A Neural PDE Solver with Temporal Stencil Modeling. ICML (2023)."}, {"review_id": "UkXSfwmECx", "reviewer": "Reviewer_JgV1", "summary": "This paper focuses on the long-term simulation of time-dependent PDEs. The authors find that high frequencies in solutions are always neglected, impeding the long-term rollout performance. This paper presents the PDE-Refiner by introducing the diffusion model to refine the model prediction. PDE-Refiner can boost the model performance consistently in representative fluid dynamics benchmarks.", "questions": "Since the computation efficiency is the major problem of this method, is there some potential solution to speed up the PDE-Refiner efficiency?", "limitations": "They have discussed limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. This paper presents an interesting finding that the key problem in fluid simulation as the prediction of high frequencies.\n\n2. The proposed methods can consistently boost the model prediction.\n\n3. Detailed ablations are provided.", "weaknesses": "1. About the model performance.\n\nFrom Figure 1(b), we can find that the PDE-Refiner will smooth the amplitude of high frequencies. This should be discussed as a limitation.\n\n2. Compare with a simple baseline.\n\nThere is a simple baseline to enhance the model performance in high frequencies. That is calculating MSE loss in the frequency domain. Concretely, calculating MSE between the spectrum of ground truth and model prediction.\n\n3. The diffusion model design can make the method data-hungry.\n\nIt is common knowledge in deep learning community that diffusion models are data-hungry. Thus, it is necessary to compare the performance of UNet and PDE-Refiner-UNet under the limited data setting (ranging from 10% data to 100% data).\n\n4. Apply PDE-Refiner to more advanced neural operators.\n\nMany recent models perform well in long-term rollout with only MSE loss. Given the authors only apply PDE-Refiner to U-Net, it is necessary to check whether the proposed method can further boost advanced non-FNO models, such as LSM [1].\n\n[1] Solving High-Dimensional PDEs with Latent Spectral Models, ICML 2023\n\n5. Concerns about the model efficiency. \n\nAs the authors pointed out, the PDE-Refiner is time-consumption. Besides, the authors should also compare the GPU memory. I think the poor efficiency can seriously affect the model practicability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable review and helpful suggestions, especially the question on the low data regime was very interesting. We respond to your questions and concerns below.\n\n**1. From Figure 1(b), we can find that the PDE-Refiner will smooth the amplitude of high frequencies. This should be discussed as a limitation.**\n\nThis is a good catch, and actually an imprecision in the figure. We will clarify in the paper that we show the *average* over samples for PDE-Refiner, while showing a single example for the ground truth. We visualize individual samples in Figure 20 (Appendix E.4), showing that applying PDE-Refiner does not oversmooth these frequencies. We also point out that the frequencies above $\\sim 60$ are only looking so chaotic in the GT data due to float32 precision limitation. See Figure 8 (Appendix D.1) for how the amplitude continues to decrease for float64.\n\n**2. There is a simple baseline to enhance the model performance in high frequencies. That is calculating MSE loss in the frequency domain. Concretely, calculating MSE between the spectrum of ground truth and model prediction.**\n\nThis corresponds to the Sobolev norm $k=0$ baseline in our experiments (e.g. Figure 1, main paper), calculating the MSE between the predicted and GT spectrum. We do not observe any improvements over MSE baseline when using Sobolev norm, since the amplitudes of the high frequencies are still by several magnitudes smaller than the dominant frequencies. Thus, the loss focuses on the same frequencies as the regular MSE loss. Similarly, the Sobolev norm $k=1$ setup gives the high frequencies a higher weight in the loss, but results in worse modeling of the lower frequencies and hence again shorter rollout time.\n\n**3. It is common knowledge in deep learning community that diffusion models are data-hungry. Thus, it is necessary to compare the performance of UNet and PDE-Refiner-UNet under the limited data setting (ranging from 10$\\%$ to 100$\\%$ data).**\n\nThis has become a valuable input to us. We first point out that we are not training the diffusion model on a complex multi-modal distribution like unconditional image generation, but rather on a deterministic target. Hence, we don't expect PDE-Refiner to be more data hungry. Based on your suggestion, we verify this by training PDE-Refiner and an MSE baseline on 10$\\%$, 20$\\%$ and 50$\\%$ of the training trajectories in the KS dataset. The results in the rebuttal Figure 4 show that PDE-Refiner consistently outperforms the baseline in all settings, with PDE-Refiner at 10\\% of the data performing on par to the MSE model at 100\\%. Moreover, the relative improvement of PDE-Refiner increases to 50$\\%$ for the low data regime, since denoising data has an ever-changing objective and input, effectively acting as a *data augmentation*. Thank you for this input, which highlights another aspect of PDE-Refiner, which we are happy to discuss in the main paper.\n\n**4. Many recent models perform well in long-term rollout with only MSE loss. Given the authors only apply PDE-Refiner to U-Net, it is necessary to check whether the proposed method can further boost advanced non-FNO models, such as LSM [1].**\n\nFully agreed. We have added FNO studies (matching the U-Net ones in the main paper) in Fig 13, and FNO ablations in Fig 14 in the submitted appendix. Additionally, in the rebuttal Figure 1, we have added results for Dilated ResNets, in which PDE-Refiner shows similar gains as on U-Nets. U-Nets, FNOs, and Dilated ResNets are the best performing models in both comprehensive overviews of Gupta \\& Brandstetter [2022] and Stachenfeld et al. [2022], highlighting PDE-Refiners applicability across state-of-the-art neural operators. We are happy to discuss more settings for the final paper, such as the newly proposed LSM, which was unfortunately not feasible in the short time frame of the rebuttal.\n\n**5. As the authors pointed out, the PDE-Refiner is time-consumption. Besides, the authors should also compare the GPU memory.**\n\nIn terms of memory, we only add a tensor of the input size for the noise and minor computations for the conditioning on refinement step index $k$. Note that during training, we sample the index $k$ and do not require the whole refinement process to be performed. As listed in the rebuttal Table 1, the differences in terms of GPU memory between PDE-Refiner and the common MSE objective are minor (below 1\\%) in both training and testing.\n\n**Since the computation efficiency is the major problem of this method, is there some potential solution to speed up the PDE-Refiner efficiency?**\n\nAs shown in experiments, PDE-Refiner can gain higher accuracy than models with 4x the parameters, therefore allowing PDE-Refiner to run smaller, more efficient networks than standard neural operators at the same accuracy. Furthermore, early steps require lower accuracy, such that future work could investigate smaller networks and adaptive computation for different steps. Finally, for training efficiency, we share the same neural operator between refinement steps. Training a separate network per step would allow for smaller networks and fewer steps, increasing inference efficiency.\n\n---\n\n[Gupta \\& Brandstetter, 2022] Gupta, Jayesh K., and Johannes Brandstetter. Towards multi-spatiotemporal-scale generalized pde modeling. arXiv preprint arXiv:2209.15616 (2022).\n\n[Stachenfeld et al., 2022] Stachenfeld, Kimberly, et al. Learned coarse models for efficient turbulence simulation. ICLR (2022)."}, {"review_id": "fdKc6ffYkJ", "reviewer": "Reviewer_xf2d", "summary": "This paper targets the prediction of future PDE states from given initial conditions with deep-learning based surrogate simulators over long rollouts. The authors propose a method titled PDE-Refiner that combines three aspects: 1) autoregressive solution rollouts similar to numerical PDE solvers, 2) direct next-step predictions in a fully data-driven manner (referred to as MSE Training in the paper), and 3) an refinement process to improve the accuracy of predictions based on an iterative denoising objective inspired by diffusion models. Their methodology also links the rollout stability of learned surrogate PDE solvers to the neglect of non-dominant, high spatial frequencies.\n\nThe proposed method is evaluated on the 1D Kuramoto-Sivashinsky equation and a 2D Kolmogorov Flow with low spatial resolution. Especially on the former, a large range of different rollout techniques and ablations are provided. The paper also includes a detailed supplemental material that discusses reproducibility, implementation, experimental details, as well as a range of additional evaluations and results.", "questions": "1. How are the correlation thresholds of 0.8 / 0.9 (L200) chosen? Is it expected that choosing lower or higher values would substantially change the findings in terms of rollout stability?\n2. In the main methodology explanation (around L135) it could be more clear how the cases for k=0 and k!=0 are balanced. Furthermore, it is not directly obvious that the same model is used to learn k=0 and k!=0. However, this is very clear with the explanations and pseudocode in the appendix.\n", "limitations": "While there is a short discussion on interpreting different diffusion samples in terms of uncertainty estimation (L272), the aspect of posterior sampling is not discussed. Traditional diffusion models allow for generating substantially different samples from the target distribution, but the proposed method is significantly limited in that regard. As the base of the PDE-Refiner is a standard one-step prediction, the diffusion-based refinement process is not able to generate meaningful differences in the samples anymore. In my opinion, this aspect should be discussed in the paper. \n\nI think it would be important to clearly mention the number of rollout steps for each setup directly. Delta t from the simulation is given first, and that the model takes a lot fewer rollout steps is not immediately obvious (L187 and L307). Especially so for the 2D case where the overall number of network rollout steps is not given and can only roughly be inferred via 16 delta t and the t=14.25 step in Figure 12 in the appendix.\n\nOtherwise, limitations of the proposed approach are discussed with sufficient detail, but it could be mentioned that the findings only hold for data from the training domain as mentioned above. \n", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper targets the problem of the prediction of PDE solutions which are highly relevant across disciplines like medicine, weather and climate predictions as well as engineering. It explores the benefits of the recently emerging diffusion-like architectures, which have been barely studied in this context and have the potential to improve the speed and accuracy of traditional numerical PDE solvers. To the best of my knowledge the combination of a one-step prediction with an iterative diffusion-based refinement process is a novel approach. Furthermore, the proposed methodology also has the potential to impact other domains, for instance video prediction approaches could benefit from a similar method of next-step predictions followed by iterative refinements.\n\nClarity and presentation are a clear strength of this work. The methodology and experiment descriptions given in the main paper as well as the appendix are well-written and detailed. Even readers that are not directly familiar with the subtleties of the discussed topics should have no problems following the argumentation of the paper. The usage of mathematical notation is concise without being confusing, and does not interrupt the line of argumentation. Related work is cited appropriately throughout the paper and the usage of existing data sets, architectures and implementations is well documented.\n\nThe evaluations in this paper shed light on various design aspects of the proposed architecture and are performed thoroughly, especially so for the 1D cases. A large range of ablations and comparisons to related approaches strengthen the line of argumentation in the text. The appendix contains further results to highlight less crucial aspects. Results are evaluated and discussed thoroughly, and include error bars over multiple random seeds in almost every experiment. In addition, the combination of the provided training details, argumentations for hyperparameter values via ablations and search ranges, and pseudocode (with a promise to release their source code with datasets upon publication) leads to outstanding reproducibility of this work.\n\nOverall, I think this is a high-quality paper that is well-polished and has substantial strengths in terms of contributions, novelty, quality and presentation.\n", "weaknesses": "To me, the result comparison for the 2D Kolmogorov Flow in Table 1 seems problematic. It is clearly mentioned that the values for classical PDE solvers and the hybrid methods are taken from previous work. However, it appears that the correlation computation shown for the ML surrogates is performed differently; thus making it not comparable. Intuitively, I would not expect single-step predictions without refinement on a 64x64 resolution to outperform DNS results computed on a 1024x1024 grid. Furthermore, the comparison with LI-CNN seems to indicate that a major part of the benefit of the proposed method comes from the underlying U-Net architecture instead of the refinement. It would be interesting to see further evaluations to investigate and clarify this behavior.\n\nAnother weakness is the choice of generalization tests in the paper. While the different data generation setups describe the test sets clearly, they always come from the exact same domain as the training data (apart from holding out some trajectories with a different random seed). The models do have to generalize to longer trajectories during inference, but are not trained with the full training sequence length as optimization only considers single-step predictions. Especially for simulation data from PDEs, it is very valuable to test model interpolation or extrapolation abilities by holding out certain parameter ranges during training for testing lateron, e.g. different length, viscosity, or forcing values. It would be a highly interesting and promising direction to explore if the findings in terms of rollout stability of this paper hold for more challenging generalization scenarios.\n\nTo conclude, this paper has some weaknesses, however addressing them is not immediately crucial as the paper in its current state already provides a large range of solid contributions detailed above.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your encouraging review and thoughtful feedback. We provide answers to your questions and concerns below.\n\n**2D Kolmogorov Flow - the correlation computation shown for the ML surrogates is performed differently**\n\nThe evaluation for the 2D Kolmogorov Flow (lines 328-333) follows the same setup as Sun et al. [2023], from which we take the related results of DNS and hybrid solver. We use the same evaluation code and datasets to ensure fair comparisons. If there is any unclarity in whether we evaluated differently, we would be happy to clarify it.\n\n**I would not expect single-step predictions without refinement on a 64x64 resolution to outperform DNS results computed on a 1024x1024 grid.**\n\nDNS is a general PDE solver, while the neural solver have been specifically trained for this dataset. Certain elements of the PDE such as the equation, the Reynolds number and structure of the initial conditions (Gaussian blobs) are fixed. Hence, neural models can specialize on this data and interpolate the data very well from low resolution, as shown in Sun et al. [2023]. Similarly, the best hybrid approach (TSM) at a resolution of 64x64 can outperform DNS at 1024x1024 as well. We use this dataset since it is a common benchmark, especially for hybrid and neural methods, but a larger, diverse 2D study across hybrid and neural solvers would be interesting for future work.\n\n**The comparison with LI-CNN seems to indicate that a major part of the benefit of the proposed method comes from the underlying U-Net architecture instead of the refinement.**\n\nWe agree that scaling up the networks for hybrid methods, especially learned correction, could improve their performance closer to the neural methods. However, this increases computation time, since we need to execute both the classical solver *and* the network. Already for the smaller network used in the hybrid solvers, PDE-Refiner and the classical U-Net operator are faster than the hybrid solvers. Yet, adding the refinement to the U-Net significantly improves the rollout performance, highlighting the benefit of PDE-Refiner also in this setting.\n\n**Test sets always come from the exact same domain as the training data. [...] It would be a highly interesting and promising direction to explore if the findings in terms of rollout stability of this paper hold for more challenging generalization scenarios.**\n\nThis would be indeed interesting to investigate. In the current setup of the KS dataset, the models need to generalize to unseen time domains, because all training trajectories are sampled between 65s and 110s after the initial condition, while test trajectories reach until 220s. Given our parameter-dependent dataset, we plan to run a study on training models on viscosities between 0.7 and 1.3 to evaluate their extrapolation to 0.5-0.7 and 1.3-1.5, and training on the ranges 0.5-0.8 and 1.2-1.5 for testing interpolation to 0.8-1.2. Due to the limited rebuttal time, we can't report the results yet and will add them for the final paper version. Further, based on your suggestion, we will mention in the limitations that most of our study focuses on the same train-test domains.\n\n**How are the correlation thresholds of 0.8 / 0.9 (L200) chosen? Is it expected that choosing lower or higher values would substantially change the findings in terms of rollout stability?**\n\nThe thresholds were chosen to align with prior work (Sun et al. [2023], Kochkov et al. [2021]) and represent time points at which the solution starts to visually slightly differ from the ground truth. In the rebuttal Figure 2, we plot the correlation over rollout time for different methods. For any reasonable threshold, the findings remain the same. The metric closely aligns with when the prediction exceeds a certain MSE loss, as shown in Fig 10 in Appendix D.1. \n\n**It could be more clear how the cases for $k=0$ and $k\\neq 0$ are balanced. Furthermore, it is not directly obvious that the same model is used to learn $k=0$ and $k\\neq 0$. However, this is very clear with the explanations and pseudocode in the appendix.**\n\nThank you for raising this point of unclarity. For efficiency, we indeed use the same model for all steps, including $k=0$ and $k\\neq 0$, and uniformly sample them during training. Experiments with weighing the cases differently did not improve results. We will clarify this more in the main text.\n\n**[...] the aspect of posterior sampling is not discussed. Traditional diffusion models allow for generating substantially different samples from the target distribution, but the proposed method is significantly limited in that regard.**\n\nAgreed, the main goal of PDE-Refiner is not to generate highly diverse samples, but enable accurate modeling of low-amplitude frequencies. In Appendix E.4, we analyze the sample distribution of PDE-Refiner, and find that samples differ between each other in the same spectrum structure as to the ground truth, which makes them useful for estimating uncertainty. Still, we will highlight this limitation better in the final version of the paper.\n\n**I think it would be important to clearly mention the number of rollout steps for each setup directly.**\n\nThat's a good point. We initially didn't report the rollout steps since the time step itself (i.e. seconds) had a much greater impact. Smaller time steps require many more autoregressive steps, but yet perform better as shown in Appendix E.2. For the presented results on KS dataset, the MSE baseline stays accurate for 94 rollout steps, and PDE-Refiner for 122 steps. For the 2D dataset, 85 for the baseline and 95 for PDE-Refiner. We will add this to the paper.\n\n---\n\n[Kochkov et al., 2021] Kochkov, Dmitrii, et al. Machine learning–accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences 118.21 (2021).\n\n[Sun et al., 2023] Sun, Zhiqing, et al. A Neural PDE Solver with Temporal Stencil Modeling. ICML (2023)."}]

Author rebuttal:
We would like to thank all reviewers for their constructive and valuable feedback. 

We are encouraged to see that reviewers highlight this study as interesting (YYxA, JgV1, xf2d), the good presentation of the paper (YYxA, xf2d), and the thorough ablation studies (K6r8, JgV1, xf2d). For example, reviewer xf2d describes the paper as a "high-quality paper that is well-polished and has substantial strengths in terms of contributions, novelty, quality and presentation". Reviewer JgV1 mentions "this paper presents an interesting finding", and reviewer YYxA that our analysis "can inspire future research". Finally, reviewer K6r8 "appreciate[s] the thorough review of the paper".

We welcome the opportunity to clarify all concerns and questions raised. We would first like to point out that our submitted appendix contains a rich set of ablation studies and further experimental results, for which the references in the main paper may have been easily missed in its current version. We will point out the appendix references more clearly in the final paper where needed, and link here the sections that address reviewer's concerns, such as:

- applications to other models (Appendix E.1), where we have added FNO studies (matching the presented ones in the main paper) in Fig 13 and FNO ablations in Fig 14 (reviewer K6r8, JgV1);
- alternative study to high-correlation time (Appendix D.1), in which we check when the prediction exceeds certain loss as shown in Fig 10 (reviewer K6r8);
- ablations on different time intervals and different input histories in Appendix E.2 and Appendix E.3 (reviewer K6r8);
- and repeating the frequency analysis for the 2D Kolmogorov flow in Appendix E.5 (reviewer YYxA).

Based on the reviewer's suggestion, we provide additional empirical results in the submitted 1-page rebuttal PDF:

- Figure 1: results on another SOTA neural operator, Dilated ResNets [Stachenfeld et al., 2022]. PDE-Refiner significantly improves performance on this operator as well (reviwer JgV1, K6r8);
- Figure 2: showing the consistency of our findings across threshold values on the high-correlation time by plotting the correlation of different methods over rollout time. The trend stays the same for various thresholds (reviewer xf2d);
- Figure 3: predicting multiple time steps at once. The performance is limited by the same generalization challenges as predicting at large time steps (reviewer K6r8);
- Figure 4: applying PDE-Refiner to smaller datasets. The experiments highlight PDE-Refiner's data efficiency over standard neural operators, in which the denoising objective acts as augmentation (reviewer JgV1);
- Table 1: the GPU memory required by PDE-Refiner only slightly differs from standard operator training, increasing it by less than 1\% for both training and inference (reviewer JgV1).

Due to the limited rebuttal time, we show some results for a single seed only, whereby the general trends were always well outside the standard deviation ranges. For the final paper, we will repeat the experiments over multiple seeds, as done for the other experiments in the main paper.

For further clarification, please refer to our individual responses to each reviewer. If any concerns, unclarities or questions of the reviewers remain, we would be happy to answer them and elaborate on our previous responses during the discussion phase.

---

[Stachenfeld et al., 2022] Stachenfeld, Kimberly, et al. Learned coarse models for efficient turbulence simulation. ICLR (2022).

Meta-review: The paper connects a multi-step prediction for PDEs with diffusion models. This is a fresh view on the subject, with interesting algorithmic ideas and convincing numerical results. The reviewers have valued the results as well.
Decision: Accept (spotlight)