SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Phillip Lippe
Microsoft Research AI4Science*
phillip.lippe@googlemail.com
Paris Perdikaris
Bastiaan S. Veeling
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.
In recent years, mostly due to a 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 can 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 of them [1, 2, 18, 33, 43, 79]; (iii) neural approaches in which the learned evolution operator maps the current state to a future state of the system [4, 7, 9, 21, 49, 54, 68, 87].
*Work done during internship at Microsoft Research; on leave from University of Amsterdam.
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 in two and three dimensions [21, 39, 43, 46, 51, 60, 63, 67, 76, 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 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.e., 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, 42, 72, 75, 82], and, thus, opens the door for new families of uncertainty-aware 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 state-of-the-art neural operators such as modern U-Nets [22] and Fourier Neural Operators (FNOs) [49], and identify a 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-Refiner on solving the 1D Kuramoto-Sivashinsky 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.,  $t \in [0,T]$ , and possibly multiple spatial dimensions, i.e.,  $\pmb{x} = [x_{1},x_{2},\dots,x_{m}] \in \mathcal{X}$ . Time-dependent PDEs relate solutions  $\pmb{u}(t,\pmb{x}) : [0,T] \times \mathcal{X} \to \mathbb{R}^n$  and respective derivatives for all points in the domain, where  $\pmb{u}^{0}(\pmb{x})$  are initial conditions at time  $t = 0$  and  $B[\pmb{u}](t,\pmb{x}) = 0$  are boundary conditions with boundary operator  $B$  when  $\pmb{x}$  lies on the boundary  $\partial \mathcal{X}$  of the domain. Such PDEs can be written in the form [9]:
$$
\boldsymbol {u} _ {t} = F (t, \boldsymbol {x}, \boldsymbol {u}, \boldsymbol {u} _ {\boldsymbol {x}}, \boldsymbol {u} _ {\boldsymbol {x x}}, \dots), \tag {1}
$$
where the notation  $\pmb{u}_t$  is shorthand for the partial derivative  $\partial \pmb{u} / \partial t$ , while  $\pmb{u}_x, \pmb{u}_{xx}, \ldots$  denote the partial derivatives  $\partial \pmb{u} / \partial x, \partial^2 \pmb{u} / \partial x^2$  and so on². Operator learning [48, 49, 53-55] relates solutions  $\pmb{u}: \mathcal{X} \to \mathbb{R}^n$ ,  $\pmb{u}': \mathcal{X}' \to \mathbb{R}^{n'}$  defined on different domains  $\mathcal{X} \in \mathbb{R}^m$ ,  $\mathcal{X}' \in \mathbb{R}^{m'}$  via operators  $\mathcal{G}$ :  $\mathcal{G}: (\pmb{u} \in \mathcal{U}) \to (\pmb{u}' \in \mathcal{U}')$ , where  $\mathcal{U}$  and  $\mathcal{U}'$  are the spaces of  $\pmb{u}$  and  $\pmb{u}'$ , respectively. For time-dependent PDEs, an evolution operator can be used to compute the solution at time  $t + \Delta t$  from time  $t$  as
$$
\boldsymbol {u} (t + \Delta t) = \mathcal {G} _ {t} (\Delta t, \boldsymbol {u} (t)), \tag {2}
$$
where  $\mathcal{G}_t: \mathbb{R}_{>0} \times \mathbb{R}^n \to \mathbb{R}^n$  is the temporal update. To obtain predictions over long time horizons, a temporal operator could either be directly trained for large  $\Delta t$  or recursively applied with smaller time intervals. In practice, the predictions of learned operators deteriorate for large  $\Delta t$ , while autoregressive approaches are found to perform substantially better [22, 50, 86].
Long Rollouts for Time-Dependent PDEs. We start with showcasing the challenges of obtaining long, accurate rollouts for autoregressive neural PDE solvers on the working example of the 1D
${}^{2}\partial \mathbf{u}/\partial \mathbf{x}$  represents a  $m \times  n$  dimensional Jacobian matrix  $\mathbf{J}$  with entries  ${\mathbf{J}}_{ij} = \partial {u}_{i}/\partial {x}_{j}$  .
(a) Rollout examples
(b) Frequency spectrum
Figure 1: Challenges in achieving accurate long rollouts on the KS equation, comparing PDE-Refiner and an MSE-trained model. (a) Example trajectory with predicted rollouts. The yellow line indicates the time when the Pearson correlation between ground truth and prediction drops below 0.9. PDE-Refiner maintains an accurate rollout for longer than the MSE model. (b) Frequency spectrum over the spatial dimension of the ground truth data and one-step predictions. For PDE-Refiner, we show the average spectrum across 16 samples. (c) The spectra of the corresponding errors. The MSE model is only accurate for a small, high-amplitude frequency band, while PDE-Refiner supports a much larger frequency band, leading to longer accurate rollouts as in (a).
(c) One-step error
Kuramoto-Sivashinsky (KS) equation [44, 73]. The KS equation is a fourth-order nonlinear PDE, known for its rich dynamical characteristics and chaotic behavior [35, 40, 74]. It is defined as:
$$
u _ {t} + u u _ {x} + u _ {x x} + \nu u _ {x x x x} = 0, ^ {3} \tag {3}
$$
where  $\nu$  is a viscosity parameter which we commonly set to  $\nu = 1$ . The nonlinear term  $uu_{x}$  and the fourth-order derivative  $u_{xxxx}$  make the PDE a challenging objective for traditional solvers. We aim to solve this equation for all  $x$  and  $t$  on a domain  $[0,L]$  with periodic boundary conditions  $u(0,t) = u(L,t)$  and an initial condition  $u(x,0) = u_0(x)$ . The input space is discretized uniformly on a grid of  $N_{x}$  spatial points and  $N_{t}$  time steps. To solve this equation, a neural operator, denoted by  $\mathbb{N}0$ , is then trained to predict a solution  $u(x,t) = u(t)$  given one or multiple previous solutions  $u(t - \Delta t)$  with time step  $\Delta t$ , e.g.  $\hat{u}(t) = \mathbb{N}0(u(t - \Delta t))$ . Longer trajectory predictions are obtained by feeding the predictions back into the solver, i.e., predicting  $u(t + \Delta t)$  from the previous prediction  $\hat{u}(t)$  via  $\hat{u}(t + \Delta t) = \mathbb{N}0(\hat{u}(t))$ . We refer to 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:  $\mathcal{L}_{\mathrm{MSE}} = \| u(t) - \mathsf{NO}(u(t - \Delta t))\|^2$ . By minimizing this one-step MSE, the model learns to 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 to differ significantly from the ground truth. In Figure 1a 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 1b shows that the main dynamics of the KS equation are modeled within a frequency band of low wavenumbers (1 to 25). As a result, the primary errors in predicting a one-step solution arise from inaccuracies in modeling the dynamics of these low frequencies. This is evident in Figure 1c, where the error of the MSE-trained model is smallest for this frequency band relative to the ground truth amplitude. Nonetheless, over a long time horizon, the non-linear term  $uu_{x}$  in 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 a neural solver that models all spatial frequencies across the spectrum as accurately as possible. Essentially, our objective should give high amplitude frequencies a higher priority, since these are responsible for the main dynamics 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.
3We omit the bold notation for 1D cases where the field  $u(x,t)$  is scalar valued.
Figure 2: Refinement process of PDE-Refiner during inference. Starting from an initial prediction  $\hat{u}^1(t)$ , PDE-Refiner uses an iterative refinement process to improve its prediction. Each step represents a denoising process, where the model takes as input the previous step's prediction  $u^k(t)$  and tries to reconstruct added noise. By decreasing the noise variance  $\sigma_k^2$  over the  $K$  refinement steps, PDE-Refiner focuses on all frequencies equally, including low-amplitude information.
3 PDE-Refiner
In this section, we present PDE-Refiner, a model that allows for accurate modeling of the solution across all frequencies. The main idea of PDE-Refiner is that we allow the model to look multiple times at its prediction, and, in an iterative manner, improve the prediction. For this, we use a model  $\mathsf{NO}$  with three inputs: the previous time step(s)  $u(t - \Delta t)$ , the refinement step index  $k\in [0,\dots ,K]$ , and the model's current prediction  $\hat{u}^k (t)$ . At the first step  $k = 0$ , we mimic the common MSE objective by setting  $\hat{u}^0 (t) = 0$  and predicting  $u(t)$ :  $\mathcal{L}^0 (u,t) = \| u(t) - \mathsf{NO}\left(\hat{u}^0 (t),u(t - \Delta t),0\right)\| _2^2$ . As discussed in Section 2, this prediction will focus on only the dominating frequencies. To improve this prediction, a simple approach would be to train the model to take its own predictions as inputs and output its (normalized) error to the ground truth. However, such a training process has several drawbacks. Firstly, as seen in Figure 1, 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  $k \geq 1$ , we remove low-amplitude information of an earlier prediction by applying noise, e.g. adding Gaussian noise, to the input  $\hat{u}^k(t)$  at refinement step  $k$ :  $\tilde{u}^k(t) = \hat{u}^k(t) + \sigma_k \epsilon^k$ ,  $\epsilon^k \sim \mathcal{N}(0,1)$ . The objective of the model is to predict this noise  $\epsilon^k$  and use the prediction  $\hat{\epsilon}^k$  to denoise its input:  $\hat{u}^{k+1}(t) = \tilde{u}^k(t) - \sigma_k \hat{\epsilon}^k$ . By decreasing the noise standard deviation  $\sigma_k$  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 on low-amplitude information, typically corresponding to the non-dominant frequencies. Generally, we find that an exponential decrease, i.e.  $\sigma_k = \sigma_{\min}^{k/K}$  with  $\sigma_{\min}$  being the minimum noise standard deviation, works well. The value of  $\sigma_{\min}$  is chosen based on the frequency spectrum of the given data. For example, for the KS equation, we use  $\sigma_{\min}^2 = 2 \cdot 10^{-7}$ . We train the model by denoising ground truth data at different refinement steps:
$$
\mathcal {L} ^ {k} (u, t) = \mathbb {E} _ {\epsilon^ {k} \sim \mathcal {N} (0, 1)} \left[ \| \epsilon_ {k} - \mathbb {N} \left(u (t) + \sigma_ {k} \epsilon_ {k}, u (t - \Delta t), k\right) \| _ {2} ^ {2} \right] \tag {4}
$$
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  $\sigma_{k}$  and ignore potentially larger errors that, during inference, could have occurred in previous steps. To train all refinement steps equally well, we uniformly sample  $k$  for each training example:  $\mathcal{L}(u,t) = \mathbb{E}_{k\sim U(0,K)}\left[\mathcal{L}^k (u,t)\right]$ .
At inference time, we predict a solution  $u(t)$  from  $u(t - \Delta t)$  by performing the  $K$  refinement steps, where we sequentially use the prediction of a refinement step as 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-Refiner even improves on low frequencies with small amplitudes.
3.1 Formulating PDE-Refiner as a Diffusion Model
Denoising processes have been most famously used in diffusion models as well [12, 29-31, 61, 69, 77]. Denoising diffusion probabilistic models (DDPM) randomly sample a noise variable  $\mathbf{x}_0 \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$  and sequentially denoise it until the final prediction,  $\mathbf{x}_K$ , is distributed according to the data:
$$
p _ {\theta} \left(\mathbf {x} _ {0: K}\right) := p \left(\mathbf {x} _ {0}\right) \prod_ {k = 0} ^ {K - 1} p _ {\theta} \left(\mathbf {x} _ {k + 1} \mid \mathbf {x} _ {k}\right), \quad p _ {\theta} \left(\mathbf {x} _ {k + 1} \mid \mathbf {x} _ {k}\right) = \mathcal {N} \left(\mathbf {x} _ {k + 1}; \boldsymbol {\mu} _ {\theta} \left(\mathbf {x} _ {k}, k\right), \boldsymbol {\Sigma} _ {\theta} \left(\mathbf {x} _ {k}, k\right)\right), \tag {5}
$$
where  $K$  is the number of diffusion steps. For neural PDE solving, one would want  $p_{\theta}(\mathbf{x}_K)$  to model the distribution over solutions,  $\mathbf{x}_K = u(t)$ , while being conditioned on the previous time step  $u(t - \Delta t)$ , i.e.,  $p_{\theta}(u(t)|u(t - \Delta t))$ . For example, Lienen et al. [51] recently proposed DDPMs for modeling 3D turbulent flows due to the flows' unpredictable behavior. Despite the similar use of a denoising process, PDE-Refiner sets itself apart from standard DDPMs in several key aspects. First, diffusion models typically aim to model diverse, multi-modal 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 a very low minimum noise variance  $\sigma_{\mathrm{min}}^2$ , decreasing much faster and further than common diffusion schedulers. Second, our goal with PDE-Refiner 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 [70], which smoothly transitions from predicting the sample  $u(t)$  to the additive noise  $\epsilon$ :  $\mathrm{v}^k = \sqrt{1 - \sigma_k^2}\epsilon - \sigma_k u(t)$ . Here, the first step  $k = 0$ , yields the common MSE prediction objective by setting  $\sigma_0 = 1$ . With an exponential noise scheduler, the noise variance is commonly much smaller than 1 for  $k \geq 1$ . In these cases, the weight of the noise is almost 1 in the v-prediction, giving a diffusion process that closely resembles PDE-Refiner.
Nonetheless, the similarities between PDE-Refiner and DDPMs indicate that PDE-Refiner has a potential interpretation as a 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 to this intriguing possibility later in Section 4.1. Further, we find empirically that implementing 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 diffusion 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.
We demonstrate the effectiveness of PDE-Refiner on a diverse set of common PDE benchmarks. In 1D, we study the Kuramoto-Sivashinsky equation and compare to typical temporal rollout methods. Further, we study the models' robustness to different spatial frequency spectra by varying the visocity in the KS equation. In 2D, we compare PDE-Refiner to hybrid PDE solvers on a turbulent Kolmogorov flow, and provide a speed comparison between solvers. We make our code publicly available at https://github.com/microsoft/pdearena.
4.1 Kuramoto-Sivashinsky 1D equation
Experimental setup. We evaluate PDE-Refiner and various baselines on the Kuramoto-Sivashinsky 1D 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 \cdot 64, 1.1 \cdot 64]$  and the time step  $\Delta t \sim U(0.18, 0.22)$ . The initial conditions are sampled from a distribution over truncated Fourier series with random coefficients  $\{A_m, \ell_m, \phi_m\}_m$  as  $u_0(x) = \sum_{m=1}^{10} A_m \sin(2\pi \ell_m x / L + \phi_m)$ . We generate a training dataset with 2048 trajectories of rollout length  $140\Delta t$ , and test on 128 trajectories with a duration
High-Correlation Rollout Times on the Kuramoto-Sivashinsky equation
Figure 3: Experimental results on the Kuramoto-Sivashinsky equation. Dark and light colors indicate time for average correlation to drop below 0.9 and 0.8, respectively. Error bars represent standard deviation for 5 seeds. We distinguish four model groups: models trained with the common one-step MSE (left), alternative losses considered in previous work (center left), our proposed PDE-Refiner (center right), and denoising diffusion (center right). All models use a modern U-Net neural operator [22]. PDE-Refiner surpasses all baselines with accurate rollouts up to nearly 100 seconds.