MESSAGE PASSING NEURAL PDE SOLVERS
Year: 2021

Abstract:
The numerical solution of partial differential equations (PDEs) is difficult, having led to a century of research so far. Recently, there have been pushes to build neural–numerical hybrid solvers, which piggy-backs the modern trend towards fully end-to-end learned systems. Most works so far can only generalize over a subset of properties to which a generic solver would be faced, including: resolution, topology, geometry, boundary conditions, domain discretization regularity, dimensionality, etc. In this work, we build a solver, satisfying these properties, where all the components are based on neural message passing, replacing all heuristically designed components in the computation graph with backpropoptimized neural function approximators. We show that neural message passing solvers representationally contain some classical methods, such as finite differences, finite volumes, and WENO schemes. In order to encourage stability in training autoregressive models, we put forward a method that is based on the principle of zero-stability, posing stability as a domain adaptation problem. We validate our method on various fluid-like flow problems, demonstrating fast, stable, and accurate performance across different domain topologies, discretization, etc. in 1D and 2D. Our model outperforms state-of-the-art numerical solvers in the low resolution regime in terms of speed and accuracy.

heading: 1 INTRODUCTION
text: In the sciences, years of work have yielded extremely detailed mathematical models of physical phenomena. Many of these models are expressed naturally in differential equation form (Olver, 2014), most of the time as temporal partial differential equations (PDE). Solving these differential equations is of huge importance for problems in all numerate disciplines such as weather forecasting (Lynch, 2008), astronomical simulations (Courant et al., 1967), molecular modeling (Lelièvre & Stoltz, 2016) , or jet engine design (Athanasopoulos et al., 2009). Solving most equations of importance is analytically intractable and necessitates falling back on numerical approximation schemes. Obtaining accurate solutions of bounded error with minimal computational overhead requires the need for handcrafted solvers, always tailored to the equation at hand (Hairer et al., 1993).
The design of “good” PDE solvers is no mean feat. The perfect solver should satisfy an almost endless list of conditions. There are user requirements, such as being fast, using minimal computational overhead, being accurate, providing uncertainty estimates, generalizing across PDEs, and being easy to use. Then there are structural requirements of the problem, such as spatial resolution and timescale, domain sampling regularity, domain topology and geometry, boundary conditions, dimensionality, and solution space smoothness. And then there are implementational requirements, such as maintaining stability over long rollouts and preserving invariants. It is precisely because of this considerable list of requirements that the field of numerical methods is a splitter field (Bartels, 2016), tending to build handcrafted solvers for each sub-problem, rather than a lumper field, where a mentality of “one method to rule them all” reigns. This tendency is commonly justified with reference to no free lunch theorems.
In this paper, we propose to numerically solve PDEs with an end-to-end, neural solver. Our contributions can be broken down into three main parts:
• An end-to-end fully neural PDE solver, based on neural message passing, which offers flexibility to satisfy all structural requirements of a typical PDE problem. This design is motivated by the insight that some classical solvers (finite differences, finite volumes, and WENO scheme) can be posed as special cases of message passing.
• Temporal bundling and the pushforward trick, which are methods to encourage zerostability in training autoregressive models.
• Generalization across multiple PDEs within a given class. At test time, new PDE coefficients can be input to the solver.

heading: 2 BACKGROUND AND RELATED WORK
text: Here in Section 2.1 we briefly outline definitions and notation. We then outline some classical solving techniques in Section 2.2. Lastly, in Section 2.3, we list some recent neural solvers and split them into the two main neural solving paradigms for temporal PDEs.

heading: 2.1 PARTIAL DIFFERENTIAL EQUATIONS
text: We focus on PDEs in one time dimension t = [0, T ] and possibly multiple spatial dimensions x = [x1, x2, . . . , xD]
> ∈ X. These can be written down in the form ∂tu = F (t,x,u, ∂xu, ∂xxu, ...) (t,x) ∈ [0, T ]× X (1) u(0,x) = u0(x), B[u](t, x) = 0 x ∈ X, (t,x) ∈ [0, T ]× ∂X (2)
where u : [0, T ] × X → Rn is the solution, with initial condition u0(x) at time t = 0 and boundary conditions B[u](t, x) = 0 when x is on the boundary ∂X of the domain X. The notation ∂xu, ∂xxu, . . . is shorthand for partial derivatives ∂u/∂x, ∂2u/∂x2, and so forth. Most notably, ∂u/∂x represents a gradient vector. The boundary operator B we consider in this paper are Dirichlet boundary conditions, where BD[u] = u − bD for fixed function bD and Neumann boundary conditions BN [u] = n>∂xu− bN , for scalar-valued u where n is an outward facing normal on ∂X.
Conservation form Among all PDEs, we hone in on solving those that can be written down in conservation form, because there is already precedent in the field for having studied these (Bar-Sinai et al., 2019; Li et al., 2020a). Conservation form PDEs are written as
∂tu +∇ · J(u) = 0 , (3) where∇·J is the divergence of J. The quantity J : Rn → Rn is the flux, which has the interpretation of a quantity that appears to flow. Consequently, u is a conserved quantity within a volume, only changing through the net flux J(u) through its boundaries.

heading: 2.2 CLASSICAL SOLVERS
text: Grids and cells Numerical solvers partition X into a finite grid X = {ci}Ni=1 of N small nonoverlapping volumes called cells ci ⊂ X. In this work, we focus on grids of rectangular cells. Each cell has a center at xi. uki is used to denote the discretized solution in cell ci and time tk. There are two main ways to compute uki : sampling u k i = u(tk,xi) and averaging u k i = ∫ ci u(tk,x) dx. In our notation, omitting an index implies that we use the entire slice, so uk = (uk1 ,u k 2 , ...,u k N ).
Method of lines A common technique to solve temporal PDEs is the method of lines (Schiesser, 2012). The domain X is discretized into a grid X and the solution as a vector uk. We then solve
∂tu t ∣∣ tk = f(t,uk) t ∈ [0, T ] , (4)
where f is the form of F acting on the vectorized ut instead of the function u(t,x). The only derivative operator is now in time, making it an ordinary differential equation (ODE). This ODE can be solved with off-the-shelf ODE solvers (Butcher, 1987; Everhart, 1985). f can be formed in different ways by approximating spatial derivatives on the grid. Below are three classical techniques.
Finite difference method (FDM) In FDM, spatial derivative operators (e.g., ∂x) are replaced with difference operators, called stencils. For instance, ∂xuk|xi might become (uki+1− uki )/(xi+1− xi). Principled ways to derive stencils can be found in Appendix A. FDM is simple and efficient, but suffers poor stability unless the spatial and temporal discretizations are carefully controlled.
Finite volume method (FVM) FVM works for equations in conservation form. It can be shown using the divergence theorem that the integral of u over cell i increases only by the net flux into the cell. In 1D, this leads to
f(t, uki ) = 1
∆xi (Jki−1/2 − J k i+1/2) , (5)
where ∆xi is the cell width, and Jki−1/2, J k i+1/2 the flux at the left and right cell boundary at time tk, respectively. The problem thus boils down to estimating the flux at cell boundaries xi±1/2. The beauty of this technique is that the integral of u is exactly conserved. FVM is generally more stable and accurate than FDM, but can only be applied to conservation form equations.
Pseudospectral method (PSM) PSM computes derivatives in Fourier space. In practical terms, the mth derivative is computed as IFFT{(ιω)mFFT(u)} for ι = √ −1. These derivatives obtain exponential accuracy (Tadmor, 1986), for smooth solutions on periodic domains and regular grids. For non-periodic domains, analogues using other polynomial transforms exist, but for non-smooth solution this technique cannot be applied.

heading: 2.3 NEURAL SOLVERS
text: We build on recent exciting developments in the field to learn PDE solvers. These neural PDE solvers, as we refer to them, are laying the foundations of what is becoming both a rapidly growing and impactful area of research. Neural PDE solvers for temporal PDEs fall into two broad categories, autoregressive methods and neural operator methods, (see Figure 1a).
Neural operator methods Neural operator methods treat the mapping from initial conditions to solutions at time t as an input–output mapping learnable via supervised learning. For a given PDE and given initial conditions u0, a neural operator M : [0, T ] × F → F , where F is a (possibly infinite-dimensional) function space, is trained to satisfy
M(t,u0) = u(t). (6)
Finite-dimensional operator methods (Raissi, 2018; Sirignano & Spiliopoulos, 2018; Bhatnagar et al., 2019; Guo et al., 2016; Zhu & Zabaras, 2018; Khoo et al., 2020), where dim(F) < ∞ are grid-dependent, so cannot generalize over geometry and sampling. Infinite-dimensional operator methods (Li et al., 2020c;a; Bhattacharya et al., 2021; Patel et al., 2021) by contrast resolve this issue. Each network is be trained on example solutions of the equation of interest and is therefore locked to that equation. These models are not designed to generalize to dynamics for out-of-distribution t.
Autogressive methods An orthogonal approach, which we take, is autoregressive methods. These solve the PDE iteratively. For time-independent PDEs, the solution at time t+ ∆t is computed as
u(t+ ∆t) = A(∆t,u(t)), (7)
where A : R>0 × RN → RN is the temporal update. In this work, since ∆t is fixed, we just write A(u(t)). Three important works in this area are Bar-Sinai et al. (2019), Greenfeld et al. (2019), and Hsieh et al. (2019). Each paper focuses on a different class of PDE solver: finite volumes, multigrid, and iterative finite elements, respectively. Crucially, they all use a hybrid approach (Garcia Satorras et al., 2019), where the solver computational graph is preserved and heuristically-chosen parameters are predicted with a neural network. Hsieh et al. (2019) even have convergence guarantees for their method, something rare in deep learning. Hybrid methods are desirable for sharing structure with classical solvers. So far in the literature, however, it appears that autoregressive methods are more the exception than the norm, and for those methods published, it is reported that they are hard to train. In Section 3 we explore why this is and seek to remedy it.

heading: 3 METHOD
text: In this section we detail our method in two parts: training framework and architecture. The training framework tackles the distribution shift problem in autoregressive solvers, which leads to instability. We then outline the network architecture, which is a message passing neural network.

heading: 3.1 TRAINING FRAMEWORK
text: Autoregressive solvers A map solutions uk to causally consequent ones uk+1. A straightforward way to train this is with one-step training. If p0(u0) is the distribution of initial conditions in the training set then let pk(uk) = ∫ p(uk|u0)p0(u0) du0 be the groundtruth solution distribution at iteration k. We minimize
Lone-step = EkEuk+1|uk,uk∼pk [ L(A(uk),uk+1) ] , (8)
where L is an appropriate loss function. At test time, this method has a key failure mode, instability: small errors in A accumulate over rollouts greater in length than 1 (which is the vast majority of
rollouts), and lead to divergence from the groundtruth. This can be interpreted as overfitting to the one-step training distribution, and thus being prone to generalize poorly if the input shifts from this, which is usually the case after a few rollout steps.
The pushforward trick We approach the problem in probabilistic terms. The solver maps pk 7→ A]pk at iteration k + 1, where A] : P(X) → P(X) is the pushforward operator for A and P(X) is the space of distributions on X . The problem is that after a single test time iteration, the solver sees samples from A]pk instead of the distribution pk+1, and unfortunately A]pk 6= pk+1 because errors always survive training. The test time distribution is thus shifted, which we refer to as the distribution shift problem. This is a domain adaptation problem. We mitigate the distribution shift problem by adding a stability loss term, accounting for the distribution shift. A natural candidate is an adversarial-style loss
Lstability = EkEuk+1|uk,uk∼pk [ E |uk [ L(A(uk + ),uk+1) ]] , (9)
where |uk is an adversarial perturbation sampled from an appropriate distribution. For the perturbation distribution, we choose such that (uk + ) ∼ A]pk. This can be easily achieved by using (uk + ) = A(uk−1) for uk−1 one step causally preceding uk. Our total loss is then Lone-step + Lstability. We call this the pushforward trick. We implement this by unrolling the solver for 2 steps but only backpropagating errors on the last unroll step, as shown in Figure 2. This is also outlined algorithmically in the appendix. We found it important not to backpropagate through the first unroll step. This is not only faster, it also seems to be more stable. Exactly why, we are not sure, but we think it may be to ensure the perturbations are large enough. Training the adversarial distribution itself to minimize the error, defeats the purpose of using it as an adversarial distribution.
Adversarial losses were also introduced in Sanchez-Gonzalez et al. (2020) and later used in Mayr et al. (2021), where Brownian motion noise is used for and there is some similarity to Noisy Nodes (Godwin et al.), where noise injection is found to stabilize training of deep graph neural networks. There are also connections with zero-stability (Hairer et al., 1993) from the ODE solver literature. Zero-stability is the condition that perturbations in the input conditions are damped out sublinearly in time, that is ‖A(u0 + ) − u1‖ < κ‖ ‖, for appropriate norm and small κ. The pushforward trick can be seen to minimize κ directly. Typically, this property is designed directly into ODE solvers.
The temporal bundling trick The second trick we found to be effective for stability and reducing rollout time is to predict multiple timesteps into the future synchronously. A typical temporal solver only predicts u0 7→ u1; whereas, we predict K steps u0 7→ (u1,u2, ...,uK) = u1:K together. This reduces the number of solver calls by a factor of K and so reduces the number of times the solution distribution undergoes distribution shifts. A schematic of this setup can be seen in Figure 1b.

heading: 3.2 ARCHITECTURE
text: We model the grid X as a graph G = (V, E) with nodes i ∈ V , edges ij ∈ E , and node features fi ∈ Rc. The nodes represent grid cells ci and the edges define local neighborhoods. Modeling
the domain as a graph offers flexibility over grid sampling regularity, spatial/temporal resolution, domain topology and geometry, boundary modeling and dimensionality. The solver is a graph neural network (GNN) (Scarselli et al., 2009; Kipf & Welling, 2017; Defferrard et al., 2016; Gilmer et al., 2017; Battaglia et al., 2018), representationally containing the function class of several classical solvers, see Section 3.2. We follow the Encode-Process-Decode framework of Battaglia et al. (2018) and Sanchez-Gonzalez et al. (2020) , with adjustments. We are not the first to use GNNs as PDE solvers (Li et al., 2020b; De Avila Belbute-Peres et al., 2020), but ours have several notable features. Different aspects of the chosen architecture are ablated in Appendix G.
Encoder The encoder computes node embeddings. For each node i it maps the last K solution values uk−K:ki , node position xi, current time tk, and equation embedding θPDE to node embedding vector f0i =
v([uk−K:ki ,xi, tk,θPDE]). θPDE contains the PDE coefficients and other attributes such as boundary conditions. An exact description is found in Section 4. The inclusion of θPDE allows us to train the solver on multiple different PDEs.
Processor The processor computes M steps of learned message passing, with intermediate graph representations G1,G2, ...,GM . The specific updates we use are
edge j → i message: mmij = φ ( fmi − fmj ,uk−K:ki − u k−K:k j ,xi − xj ,θPDE ) , (10)
node i update: fm+1i = ψ fmi , ∑ j∈N (i) mmij , θPDE  , (11) where N (i) holds the neighbors of node i, and φ and ψ are multilayer perceptrons (MLPs). We found that using node feature differences at the input of φ increases performance. Using relative positions xj − xi can be justified by the translational symmetry of the PDEs we consider. Solution differences ui − uj make sense by thinking of the messaging passing as a local difference operator, like a numerical derivative operator.
Decoder After message passing, we use a shallow 1D convolutional network with shared weights across spatial locations to output the K next timestep predictions at grid point xi. For each node i, the processor outputs a vector fMi . We treat this vector as a temporally contiguous signal, which we feed into a CNN over time. The CNN helps to smooth the signal over time and is reminiscent of linear multistep methods (Butcher, 1987), which are very efficient but generally not used because of stability concerns. We seem to have avoided these stability issues, by making the time solver nonlinear and adaptive to its input. The result is a new vector di = (d1i ,d 2 i , ...,d K i ) with each element dki corresponding to a different point in time. We use this to update the solution as
uk+`i = u k i + (tk+` − tk)d`i , 1 ≤ ` ≤ K . (12)
The motivation for this choice of decoder has to do with a property called consistency (Arnold, 2015), which states that lim∆t→0 ‖A(∆t,u0) − u(∆t)‖ = 0. That is, the prediction matches the exact solution in the infinitesimal time limit. Consistency is a requirement for zero-stability of the rollouts.
Connections. As mentioned in Bar-Sinai et al. (2019), both FDM and FVM are linear methods, which estimate nth-order point-wise function derivatives as
[∂(n)x u]i ' ∑
j∈N (i)
α (n) j uj (13)
for appropriately chosen coefficients α(n)j , whereN (i) is the neighborhood of cell i. FDM computes this at cell centers, and FVM computes this at cell boundaries. These estimates are plugged into flux equations, such as in Table 3 in the appendix, followed by an optional FVM update step, shown in Equation 5, to compute time derivative estimates for the ODE solver. The WENO5 scheme computes derivative estimates by taking an adaptively-weighted average over multiple FVM estimates, computed using different neighborhoods of cell i (see Equation 25). Equations 13, 5, and 25 are just message passing schemes with weighted aggregation (1 layer for FDM, 2 layers for FVM, and 3 layers for WENO). It is through this connection, that we see that message-passing neural networks representationally contain these classical schemes, and are thus a well-motivated architecture.

heading: 4 EXPERIMENTS
text: We demonstrate the effectiveness of the MP-PDE solver on tasks of varying difficulty to showcase its qualities. In 1D, we study its ability to generalize to unseen equations within a given family; we study boundary handling for periodic, Dirichlet, and Neumann boundary conditions; we study both regular and irregular grids; and we study the ability to model shock waves. We then show that the MP-PDE is able to solve equations in 2D. We also run ablations over the pushforward trick and variations, to demonstrate its utility. As baselines, we compare against standard classical PDE solvers, namely; FDM, pseudospectral methods, and a WENO5 solver, and we compare against the Fourier Neural Operator of Li et al. (2020a) as an example of a state-of-the-art neural operator method. The MP-PDE solver architecture is detailed in Appendix F.

heading: 4.1 INTERPOLATING BETWEEN PDES
text: Data We focus on the family of PDEs
[∂tu+ ∂x(αu 2 − β∂xu+ γ∂xxu)](t, x) = δ(t, x), (14)
u(0, x) = δ(0, x), δ(t, x) = J∑ j=1 Aj sin(ωjt+ 2π`jx/L+ φj). (15)
Writing θPDE = (α, β, γ), corner cases are the heat equation θPDE = (0, η, 0), Burgers’ equation θPDE = (1, η, 0), and the KdV equation θPDE = (6, 0, 1). The term δ is a forcing term, following Bar-Sinai et al. (2019), with J = 5, L = 16 and coefficients sampled uniformly in Aj ∈ [−0.5, 0.5], ωj ∈ [−0.4,−0.4], `j ∈ {1, 2, 3}, φj ∈ [0, 2π). This setup guarantees periodicity of the initial conditions and forcing. Space is uniformly discretized to nx = 200 cells in [0, 16) with periodic boundary and time is uniformly discretized to nt = 200 points in [0, 4]. Our training sets consist of 2096 trajectories, downsampled to resolutions (nt, nx) ∈ {(250, 100), (250, 50), (250, 40)}. Numerical groundtruth is generated using a 5th-order WENO scheme (WENO5) (Shu, 2003) for the convection term ∂xu2 and 4th-order finite difference stencils for the remaining terms. The temporal solver is an explicit Runge-Kutta 4 solver (Runge, 1895; Kutta, 1901) with adaptive timestepping. Detailed methods and implementation are in Appendix C. All methods are implemented for GPU,
so runtime comparisons are fair. For the interested reader comparison of WENO and FDM schemes against analytical solutions can be found in Appendix C to establish utility in generating groundtruth.
Experiments and results We consider three scenarios: E1 Burgers’ equation without diffusion θPDE = (1, 0, 0) for shock modeling; E2 Burgers’ equation with variable diffusion θPDE = (1, η, 0) where 0 ≤ η ≤ 0.2; and E3 a mixed scenario with θPDE = (α, β, γ) where 0.0 ≤ α ≤ 3.0, 0.0 ≤ β ≤ 0.4 and 0.0 ≤ γ ≤ 1.0. E2 and E3 test the generalization capability. We compare against downsampled groundtruth (WENO5) and a variation of the Fourier Neural Operator with an autoregressive structure (FNO-RNN) used in Section 5.3 of their paper, and trained with unrolled training (see Figure 2). For our models we run the MP-PDE solver, an ablated version (MP-PDE- θPDE), without θPDE features, and the Fourier Neural Operator method trained using our temporal bundling and pushforward tricks (FNO-PF). Errors and runtimes for all experiments are in Table 1. We see that the MP-PDE solver outperforms WENO5 and FNO-RNN in accuracy. Temporal bundling and the pushforward trick improve FNO dramatically, to the point where it beats MP-PDE on E1. But MP-PDE outperforms FNO-PF on E2 and E3 indicating that FNO is best for single equation modeling, but MP-PDE is better at generalization. MP-PDE predictions are best if equation parameters θPDE are used, evidenced in E2 and E3. This effect is most pronounced for E3, where all parameters are varied. Exemplary rollout plots for E2 and E3 are in Appendix F.1. Figure 4 (TOP) shows shock formation at different resolutions (E1), a traditionally difficult phenomenon to model—FDM and PSM methods cannot model shocks. Strikingly, shocks are preserved even at very low resolution.

heading: 4.2 VALIDATING TEMPORAL BUNDLING AND THE PUSHFORWARD METHOD
text: We observe solver survival times on E1, defined as the time until the solution diverges from groundtruth. A solution ûki diverges from the groundtruth u k i when its max-normalized L1-error
1 nx ∑nx i=1 |ûki−u k i | maxj |ukj | exceeds 0.1. The solvers are unrolled to nt = 1000 timesteps with T = 16 s. Examples are shown in Figure 4 (BOTTOM), where we observe increasing divergence after ∼ 8 s. This is corroborated by Figure 5a where we see survival ratio against timestep. This is in line with observed problems with autoregressive models from the literature—see Figure C.3 of SanchezGonzalez et al. (2020) or Figure S9 of Bar-Sinai et al. (2019)
In a second experiment, we compare the efficacy of the pushforward trick. Already, we saw that, coupled with temporal bundling, it improved FNO for our autoregressive tasks. In Figure 5b, we plot the survival ratios for models trained with and without the pushforward trick. As a third comparison
we show a model trained with Gaussian noise adversarial perturbations, similar to that proposed in Sanchez-Gonzalez et al. (2020). We see that applying the pushforward trick leads to far higher survival times, confirming our model that instability can be addressed with adversarial training. Interestingly, injecting Gaussian perturbations appears worse than using none. Closer inspection of rollouts shows that although Gaussian perturbations improve stability, they lead to lower accuracy, by nature of injecting noise into the system.

heading: 4.3 SOLVING ON IRREGULAR GRIDS WITH DIFFERENT BOUNDARY CONDITIONS
text: The underlying motives for designing this experiment are to investigate (i) how well our MP-PDE solver can operate on irregular grids and (ii) how well our MP-PDE solver can generalize over different boundary conditions. Non-periodic domains and grid sampling irregularity go hand in hand, since pseudo-spectral methods designed for closed intervals operate on non-uniform grids.
Data We consider a simple 1D wave equation
∂ttu− c2∂xxu = 0, x ∈ [−8, 8] (16)
where c is wave velocity (c = 2 in our experiments). We consider Dirichlet B[u] = u = 0 and Neumann B[u] = ∂xu = 0 boundary conditions. This PDE is 2nd-order in time, but can be rewritten as 1st-order in time, by introducing the auxilliary variable v = ∂tu and writing ∂t[u, v] − [v, c2∂xxu] = 0 . The initial condition is a Gaussian pulse with peak at random location. Numerical groundtruth is generated using FVM and Chebyshev spectral derivatives, integrated in time with an implicit Runge-Kutta method of Radau IIA family, order 5 (Hairer et al., 1993). We solve for groundtruth at resolution (nt, nx) = (250, 200) on a Chebyshev extremal point grid (cell edges are located at xi = cos(iπ/(nx + 1)).
Experiments and results We consider three scenarios: WE1 Wave equation with Dirichlet boundary conditions; WE2 Wave equation with Neumann boundary conditions and, WE3 Arbitrary combinations of these two boundary conditions, testing generalization capability of the MP-PDE solver. Ablation studies (marked with θPDE) have no equation specific parameters input to the MP-PDE solver. Table 2 compares our MP-PDE solver against state-of-the art numerical pseudospectral solvers. MP-PDE solvers obtain accurate results for low resolutions where pseudospectral solvers break. Interestingly, MP-PDE solvers can generalize over different boundary conditions, which gets more pronounced if boundary conditions are injected into the equation via θPDE features.

heading: 4.4 2D EXPERIMENTS
text: For 2D experiments we use data from PHIFLOW1, an open-source fluid simulation toolkit. We look at fluid simulation based on the Navier-Stokes equations, and simulate smoke inflow into a 32×32 grid, adding more smoke after every time step which follows the buoyancy force. Dynamics can be described by semi-Lagrangian advection for the velocity and MacCormack advection for the smoke distribution. Simulations run for 100 timesteps where one timestep corresponds to one second. Smoke inflow locations are sampled randomly. Architectural details are in Appendix F. Figure 12 in the appendix shows results of the MP-PDE solver and comparisons to the groundtruth simulation. The MP-PDE solver is able to capture the smoke inflow accurately over the given time period, suggesting scalability of MP-PDE solver to higher dimensions.

heading: 5 CONCLUSION
text: We have introduced a fully neural MP-PDE solver, which representationally contains classical methods, such as the FDM, FVM, and WENO schemes. We have diagnosed the distribution shift problem and introduced the pushforward trick combined with the idea of temporal bundling trying to alleviate it. We showed that these tricks reduce error explosion observed in training autoregressive models, including a SOTA neural operator method (Li et al., 2020a). We also demonstrated that MP-PDE solvers offer flexibility when generalizing across spatial resolution, timescale, domain sampling regularity, domain topology and geometry, boundary conditions, dimensionality, and solution space smoothness. In doing so, MP-PDE solvers are much faster than SOTA numerical solvers.
MP-PDE solvers cannot only be used to predict the solution of PDEs, but can e.g. also be reinterpreted to optimize the integration grid and the parameters of the PDE. For the former, simply position updates need to be included in the processor, similar to Satorras et al. (2021). For the latter, a trained MP-PDE solver can be fitted to new data where only θPDE features are adjusted. A limitation of our model is that we require high quality groundtruth data to train. Indeed generating this data in the first place was actually the toughest part of the whole project. However, this is a limitation of most neural PDE solvers in the literature. Another limitation is the lack of accuracy guarantees typical solvers have been designed to output. This is a common criticism of such learned numerical methods. A potential fix would be to fuse this work with that in probabilistic numerics (Hennig et al., 2015), as has been done for RK4 solvers (Schober et al., 2014). Another promising follow-up direction is to research alternative adversarial-style losses as introduced in Equation 9. Finally, we remark that leveraging symmetries and thus fostering generalization is a very active field of research, which is especially appealing for building neural PDE solvers since every PDE is defined via a unique set of symmetries (Olver, 1986).
1https://github.com/tum-pbs/PhiFlow

heading: 6 REPRODUCIBILITY STATEMENT
text: All data used in this work is generated by ourselves. It is thus of great importance to make sure that the produced datasets are correct. We therefore spend an extensive amount of time cross-checking our produced datasets. For experiments E1, E2, E3, this is done by comparing the WENO scheme to analytical solutions as discussed in detail in Appendix Section C.3, where we compare our implemented WENO scheme against two analytical solutions of the Burgers equation from literature. For experiments W1, W2, W3, we cross-checked if Gaussian wave packages keep their form throughout the whole wave propagation phase. Furthermore, the wave packages should change sign for Dirichlet boundary conditions and keep the sign for Neumann boundary conditions. Examples can be found in Appendix Section D.
We have described our architecture in Section 3.2 and provided further implementation details in Appendix Section F. We have introduced new concepts, namely temporal bundling and the pushforward method. We have described these concepts at length in our paper. We have validated temporal bundling and pushforward methods on both our and the Fourier Neural Operator (FNO) method Li et al. (2020a). We have not introduced new mathematical results. However, we have used data generation concepts from different mathematical fields and therefore have included a detailed description of those in our appendix.
For reproducibility, we plan to provide our code upon publication.

heading: 7 ETHICAL STATEMENT
text: The societal impact of MP-PDE solvers is difficult to predict. However, as stated in the introduction, solving differential equations is of huge importance for problems in many disciplines such as weather forecasting, astronomical simulations, or molecular modeling. As such, MP-PDE solvers potentially help to pave the way towards shortcuts for computationally expensive simulations. Most notably, a drastical computational shortcut is always somehow related to reducing the carbon footprint. However, in this regard, is also important to remind ourselves that relying on simulations or now even shortcuts of those always requires monitoring and thorough quality checks.

heading: B RECONSTRUCTION
text: Reconstruction is the task of fitting a piecewise polynomial on a mesh of points, when instead of functions values at the grid points we have cell averages ūi, where
ūi = 1
∆xi ∫ Ii u(x) dx, (21)
where each cell Ii = [xi− 12 , xi+ 12 ], the half nodes are defined as xi+ 12 := 1 2 (xi + xi+1) and ∆xi = xi+ 12 − xi− 12 . The solution is to note that we can fit a polynomial P (x) to the integral of u(x), which will be exact at the half nodes, so
P (xi+ 12 ) = U(xi+ 1 2 ) =
∫ x i+1
2
x 1− 1
2
u(x) dx = i∑ k=1 ∫ x k+1 2 x k− 1
2
u(x) dx = i∑ k=1 ūk∆xk. (22)
We can then differentiate this polynomial to retrieve an estimate for the u at off-mesh locations. Recall that polynomial differentiation is easy with the interpolant differentiation operators x(m). The system of equations we need to solve is thus∆x1 · · · 0... . . . ...
∆x1 · · · ∆xN  ︸ ︷︷ ︸
L
 ū1... ūN  ︸ ︷︷ ︸
ū
= x 0 1 · · · xN−11 ... . . . ...
x0N · · · x N−1 N  ︸ ︷︷ ︸
X
 a0... aN−1  ︸ ︷︷ ︸
a
. (23)
where ū is the vector of cell averages and L is a lower triangular matrix performing the last sum in Equation 22. Thus the mth derivative of the polynomial p(x) = P ′(x) is
dmp dxm = dm+1P dxm+1 = x(m+1)>X−1Lū = s̄(m)>ū. (24)

heading: C WENO SCHEME
text: The essentially non-oscillating (ENO) scheme is an interpolation or reconstruction scheme to estimate function values and derivatives of a discontinuous function. The main idea is to use multiple overlapping stencils to estimate a derivative at point x ∈ [xi, xi+1]. We design N stencils to fit the function on shifted overlapping intervals I1, ..., IN , where Ik = [xi−N+1+k, xi+k]. In the case of WENO reconstruction these intervals are Ik = [xi−N+1+k− 12 , xi+k+ 12 ]. If a discontinuity lies in
I = ⋃ k Ik, then it is likely that one of the substencils s (m) k or s̄ (m) k (defined on interval Ik) will not contain the discontinuity. We can thus use the substencil s(m)k or s̄ (m) k from the relatively smooth region to estimate the function derivatives at x. In the following, we focus on WENO reconstruction.
The weighted essentially non-oscillating (WENO) scheme goes one step further and takes a convex combination of the substencils to create a larger stencil s̄ on I , where (dropping the superscript for brevity)
s̄ = N∑ k=1 wks̄k. (25)
Here the nonlinear weights wk satisfy ∑N k=1 wk = 1. It was shown in Jiang & Shu (1996) that these nonlinear weights can be constructed as
wk = w̃k∑N j=1 w̃j
w̃k = γk
( + βk)2 , (26)
where γk is called the linear weight, is a small number, and βk is the smoothness indicator. The linear weights are set such that the sum ∑N k=1 γks̄k over the order N − 1 stencils matches a larger order 2N − 2 stencil. The smoothness indicators are computed as
βk = N−1∑ m=1 ∆x2m−1k ∫ x i+1 2 x i− 1
2
( dmpk dxm )2 dx, (27)
where pk is the polynomial corresponding to substencil s̄k.
C.1 WENO5 SCHEME A conservative finite difference spatial discretization approximates a derivative f(u)x by a conservative difference
f(u)x|x=xi ≈ 1
∆x
( f̂i+ 12 − f̂i− 12 ) , (28)
where f̂i+ 12 and f̂i− 12 are numerical fluxes. Since g(u)y is approximated in the same way, finite difference methods have the same format for more than one spatial dimensions. The leftreconstructed (reconstruction is done from left to right) fifth order finite difference WENO scheme (WENO5) (Shu, 2003) has the û−
i+ 12 given by:
û− i+ 12 = w1û −(1) i+ 12 + w2û −(2) i+ 12 + w3û −(3) i+ 12 . (29)
In equation 29, û−(j) i+ 12 are the left WENO reconstructions on three different stencils given by
û −(1) i+ 12 = + 1 3 ui−2 − 7 6 ui−1 + 11 6 ui , (30)
û −(2) i+ 12 = −1 6 ui−1 + 5 6 ui + 1 3 ui+1 , (31)
û −(3) i+ 12 = + 1 3 ui + 5 6 ui+1 − 1 6 ui+2 , (32)
and the non-linear weights wj given by
wj = w̃j∑3 k=1 wk , w̃k = γk ( + βk)2 , (33)
with γ{1,2,3} = { 110 , 3 5 , 3 10}, and a tiny-valued parameter to avoid the denominator becoming 0. The left smoothness indicators β−k are given by:
β−1 = 13
12
( ui−2 − 2ui−1 + ui )2 + 1
4
( ui−2 − 4ui−1 + 3ui )2 , (34)
β−2 = 13
12
( ui−1 − 2ui + ui+1 )2 + 1
4
( ui−1 − ui+1 )2 , (35)
β−3 = 13
12
( ui − 2ui+1 + ui+2 )2 + 1
4
( ui − 4ui+1 + 3ui+2 )2 . (36)
(37)
The right-reconstructed WENO5 scheme has the û+ i− 12 given similarly to û− i+ 12 but with all coefficients flipped since the reconstruction is done from the other side (from right to left). The right reconstruction on three different stencils are given by
û +(1)
i− 12 = +
1 3 ui+2 − 7 6 ui+1 + 11 6 ui , (38)
û +(2) i− 12 = −1 6 ui+1 + 5 6 ui + 1 3 ui−1 , (39)
û +(3)
i− 12 = +
1 3 ui + 5 6 ui−1 − 1 6 ui−2 , (40)
and the right smoothness indicators β+k are given by
β+1 = 13
12
( ui+2 − 2ui+1 + ui )2 + 1
4
( ui+2 − 4ui+1 + 3ui )2 , (41)
β+2 = 13
12
( ui+1 − 2ui + ui−1 )2 + 1
4
( ui+1 − ui−1 )2 , (42)
β+3 = 13
12
( ui − 2ui−1 + ui−2 )2 + 1
4
( ui − 4ui−1 + 3ui−2 )2 . (43)
(44)
Both û− and û+ are needed for full flux reconstruction as explained in the next section.
C.2 FLUX RECONSTRUCTION
We consider flux reconstruction via Godunov. For Godunov flux, f̂i+ 12 = f̂(ui+ 12 ) is reconstructed from û+
i+ 12 and û− i+ 12 via:
f̂(ui+ 12 ) =  min u− i+1 2 ≤u≤u+ i+1 2 f(u), if u− i+ 12 ≤ u+ i+ 12
max u− i+1
2
≤u≤u+ i+1
2
f(u), if u− i+ 12 > u+ i+ 12
(45)
C.3 COMPARING WENO SCHEME TO ANALYTICAL SOLUTIONS
First analytical case. An analytical solvable case for Burgers equation arises for the boundary conditions
u(t, 0) = u(t, 2π) ,
and the initial conditions
u(0, x) = −2ν ∂φ/∂x φ + 4 , (46)
where
φ = exp
( −x2
4ν
) + exp [ −(x− 2π)2
4ν(t+ 1)
] (47)
∂φ ∂x = −2x 4ν exp
( −x2
4ν
) − 2(x− 2π)
4ν + exp
[ −(x− 2π)2
4ν
] (48)
= −0.5x ν exp
( −x2
4ν
) − 0.5(x− 2π)
ν exp
[ −(x− 2π)2
4ν
] . (49)
The analytical solutions for this specific set of boundary and initial conditions gives
u(t, x) = −2ν ∂φ/∂x φ + 4 , (50)
where
φ = exp
( −(x− 4t)2
4ν(t+ 1)
) + exp [ −(x− 4t− 2π)2
4ν(t+ 1)
] (51)
∂φ ∂x = −2(x− 4t) 4ν(t+ 1) exp
( −(x− 4t)2
4ν(t+ 1)
) − 2(x− 4t− 2π)
4ν(t+ 1) + exp
[ −(x− 4t− 2π)2
4ν(t+ 1)
] (52)
= −0.5(x− 4t) ν(t+ 1) exp
( −(x− 4t)2
4ν(t+ 1)
) − 0.5(x− 4t− 2π)
ν(t+ 1) exp
[ −(x− 4t− 2π)2
4ν(t+ 1)
] . (53)
For this first analytical solveable case, the analytical solution, the WENO scheme and the fourth order finite difference scheme (FDM) are compared in Fig. 6 for a diffusion term of ν = 0.005. The WENO scheme models the analytical solution perfectly, whereas the FDM scheme fails to capture the shock accurately. For lower values of ν the effect gets even stronger.
Second analytical case. Another analytical solvable case for the Burgers equation arises for the boundary condition:
u(t,±1) = 0 , (54)
and the initial condition
u(0, x) = − sin(πx) . (55)
Solutions are (Basdevant et al., 1986) u(t, x) = − ∫∞ −∞ sinπ(x− η)f(x− η) exp(−η
2/4νt)dη∫∞ −∞ f(x− η) exp(−η2/4νt)dη , (56)
with f(y) = exp(− cos( πy2πν ) . Using Hermite integration allows the computation of accurate results up to t = 3/π.
For this second analytical solveable case, the analytical solution, and the WENO scheme are compared in Fig. 7 for a diffusion term of ν = 0.002. The WENO scheme models the analytical solution perfectly. Modeling via the FDM scheme fails completely.

heading: D PSEUDOSPECTRAL METHODS FOR WAVE PROPAGATION ON IRREGULAR GRIDS
text: We consider Dirichlet B[u] = u = 0 and Neumann B[u] = ∂xu = 0 boundary conditions. Numerical groundtruth is generated using FVM and Chebyshev spectral derivatives, integrated in time with an implicit Runge-Kutta method of Radau IIA family, order 5 (Hairer et al., 1993). To properly fulfill the boundary conditions, wave packages have to travel between the boundaries and are bounced back with same and different sign for Neumann and Dirichlet boundary condition, respectively. Exemplary wave propagation for both boundary conditions is shown in Figure 8.

heading: E EXPLICIT RUNGE-KUTTA METHODS
text: The family of Runge-Kutta methods (Butcher, 1987) is given by:
utn+1 = utn + ∆t s∑ i=1 biki , (57)
where
k1 = f(tn, utn) , (58) k2 = f(tn + c2∆t, utn + h(a21k1)) , (59) k3 = f(tn + c3∆t, utn + h(a31k1 + a32k2)) , (60)
... (61) ks = f(tn + cs∆t, utn + h(as1k1 + as2k2, . . . as,s−1ks−1)) . (62)
For a particular Runge-Kutta method one needs to provide the number of stages s, and the coefficients aij(1 ≤ j < i ≤ s), bi(i = 1, 2, . . . , s) and ci(i = 1, 2, . . . , s). These data are usually arranged in so-called Butcher tableaux (Butcher, 1963).

heading: F EXPERIMENTS
text: Flux terms of equations that we study—the Heat, Burgers, Korteweg-de-Vries (KdV), and Kuromoto-Shivashinsky (KS) equation—are summarized in Table 3.
Pushforward trick and temporal bundling. Pseudocode for one training step using the pushforward trick and temporal bundling is sketch in Algorithm 1.
Algorithm 1 Pushforward trick and temporal bundling. For a given batched input data trajectory and a model, we draw a random timepoint t, get our input data trajectory, performN forward passes, and finally perform the supervised learning task with the according labels. K is the number of steps we predict into the future using the temporal bundling trick, N is number of unrolling steps in order to apply the pushforward trick, T is the number of available timesteps in the training set.
Require: data, model, N , K, T . data is the complete PDE trajectory t← DrawRandomNumber t ∈ {1, ..., T} . We draw a random starting point input← data(t−K:t) . We input the last K timesteps for n ∈ {1, . . . , N} do
input← model(input) end for . Here we cut the gradients target← data(t+NK:t+N(K + 1)) . Actual labels after the pushforward operation output← model(input) loss← criterion(output, target)
Implementation details. MP-PDE architectures, consist of three parts (sequentially applied):
1. Encoder: Input→ {fully-connected layer→ activation→ fully-connected layer→ activation }, where fully connected layers are applied node-wise
2. Processor: 6 message passing layers as described in Sec. 3.2. Each layer consists of a 2- layer edge update network φ following Equation (10), and a 2-layer node update network ψ following Equation (11).
3. Decoder: 1D convolutional network with shard weights across spatial locations → {1D CNN layer→ activation→ 1D CNN layer }
We optimize models using the AdamW optimizer (Loshchilov & Hutter, 2017) with learning rate 1e-4, weight decay 1e-8 for 20 epochs and minimize the root mean squared error (RMSE). We use batch size 16 for experiments E1-E3 and WE1-WE3 and batch size of 4 for 2D experiments. For all experiments we use a hidden size of 128, and for enforcing zero-stability during training we unroll the solver for a maximum of 2 steps (see Sec. 3.1).
Message and update network in the processor consist of→ {fully-connected layer→ activation→ fully-connected layer → activation }. We use skip-connections in the message passing layers and apply instance normalization (Ulyanov et al., 2016) for experiments E1-E3 and WE1-WE3, and batch normalization (Ioffe & Szegedy, 2015) for the 2D experiments. For the decoder, we use 8 channels between the two CNN layers (1 input channel, 1 output channel) across all experiments. We use Swish (Ramachandran et al., 2017) activation functions for experiments E1-E3 and WE1WE3 and ReLU activation for the 2D experiments. ReLU activation proved most effective for 2D experiments since a characteristic of the smoke inflow dynamics we studied is that values are zero for all positions which are untouched by smoke buoyancy at a given timepoint.
Training details. The overall used architectures consist of roughly 1 million parameters and training for the different experiments takes between 12 and 24 hours on average on a GeForceRTX 2080 Ti GPU. 6 message passing layers and a hidden size of 128 for the 2-layer edge update network and
the 2-layer node update network is a robust choice. A hidden size of 64 shows signs of underfitting, whereas a hidden size of 200 is not improve performance significantly. For the overall performance, more important than the number of parameters is the choice of the output 1D CNN, the choice of inputs to the edge update network φ following Equation (10) and the node update network ψ following Equation (11). We ablate these choices in Appendix G.
Another interesting hyperparameter is the number of neighbors used for message passing. We construct our graphs by restricting the neighbors (edges) via a cutoff radius based on positional coordinates for experiments E1-E3 and the 2D experiments. We effectively use 6 neighbors for experiments E1-E3 and 8 neighbors for the 2D experiments. For experiments WE1-WE3, cutoff radii for selecting neighbors are not a robust choice since the grids are irregular and relative distances are much lower close to the boundaries. We therefore construct our graphs via a k-NN criterion, and effectively use between 20 neighbors (highest spatial resolution) and 6 neighbors (lowest spatial resolution).
F.1 EXPERIMENTS E1, E2, E3 Figure 9 displays exemplary 1D rollouts at different resolution for the Burgers’ equation with different diffusion terms. Lower diffusion coefficients result in faster shock formation. Figure 10 displays exemplary 1D rollouts for different parameter sets. Large α parameters result in fast and large shock formations. The wiggles arising due to the dispersive term γ and cannot be captured by numerical solvers at low resolution. Our MP-PDE solver is able to capture these wiggles and reproduce them even at very low resolution.
F.2 EXPERIMENTS WE1, WE2, WE3 Figure 11 displays exemplary 2D rollouts at different resolutions for the wave equation with Dirichlet and Neumann boundary conditions. Waves bounce back and forth between boundaries. MP-PDE solvers give accurate solutions on the irregular grids and are stable over time.

heading: G ARCHITECTURE ABLATION AND COMPARISON TO CNNS
text: A schematic sketch of our MP-PDE solver is displayed in Figure 13 (sketch taken from the main paper). A GNN based architecture was chosen since GNNs have the potential to offer flexibility when generalizing across spatial resolution, timescale, domain sampling regularity, domain topology and geometry, boundary conditions, dimensionality, and solution space smoothness. The chosen architecture representationally contains classical methods, such as FDM, FVM, and WENO schemes. The architectures follows the Encode-Process-Decode framework of Battaglia et al. (2018), with adjustments. Most notably, PDE coefficients and other attributes such as boundary conditions denoted with θPDE are included in the processor. For the decoder, a shallow 2-layer 1D convolutional network with shared weights across spatial locations is applied, motivated by linear multistep methods (Butcher, 1987).
For ablating the architecture, three design choices are verified:
1. Does a GNN have the representational power of a vanilla convolutional network on a regular grid? We test against 1D and 2D baseline CNN architectures.
2. For the decoder part, how much does a 1D convolutional network with shared weights across spatial locations approve upon a standard MLP decoder?
3. How much does the inclusion of PDE coefficients θPDE help to generalize over e.g. different equations or different boundary conditions?
Table 4 shows the three ablation (MP-PDE- θPDE, MP-PDE-(( ((1D-CNN, Baseline CNN) tested on shock wave formation modeling and generalization to unseen experiments (experiments E1, E2, and E3), as described in Section 4.1 in the main paper. For the training of the different architectures the optimized training strategy consisting of temporal bundling and pushforward trick is used.
The MP-PDE- θPDE ablation results are the same as reported in the main paper. The effect gets more prominent if more equation specific parameters are available (θPDE features), as it is the case for experiment E3. We also refer the reader to the experiments presented in Table 2, where MPPDE solvers are shown to be able to generalize over different boundary conditions, which gets much stronger pronounced if boundary conditions are injected into the equation via θPDE features.
The MP-PDE-((((1D-CNN ablation replaces the shallow 2-layer 1D convolutional network with in the decoder with a standard 2-layer MLP. Performance slightly degrades for the MLP decoder which is most likely due to the better temporal modeling introduced by the shared weights of the 1D CNN network.
A Baseline 1D-CNN is built up of 8 1D-CNN layers, where the input consists of the spatial resolution (nx). TheK previous timesteps used for temporal bundling are treated asK input channels. The output consequently predicts the next K timesteps for the same spatial resolution (K output channels). Using this format, again both temporal bundling and the pushforward trick can be effectively applied. The implemented 1D-CNN layers are:
• Input layer with K input channel, 40 output channels, kernel of size 3.
• 3 layers with 40 input channels, 40 output channels, kernel of size 5.
• 3 layers with 40 input channels, 40 output channels, kernel of size 7.
• 1 output layer with 40 input channels, K output channel, kernel of size 7.
Residual connections are used between the layers, and ELU (Clevert et al., 2016) non-linearities are applied. Circular padding is implemented to reflect the periodic boundary conditions. The CNN output is a new vector di = (d1i ,d 2 i , ...,d K i ) with each element d k i corresponding to a different point in time. Analogously to the MP-PDE solver, we use the output to update the solution as
uk+`i = u k i + (tk+` − tk)d`i , 1 ≤ ` ≤ K , (63)
where K is the output (and input) dimension. This baseline 1D-CNN is conceptually very similar to our MP-PDE solver.
A Baseline 2D-CNN is built up of 6 2D-CNN layers, where the input consists of the spatial resolution (nx) and the K previous timesteps used for temporal bundling. The output consequently predicts the next K timesteps for the same spatial resolution. Using this format, both temporal bundling and the pushforward trick can be effectively applied. The implemented 2D-CNN layers are:
• Input layer with 1 input channel, 16 output channels, 3× 3 kernel. • 4 intermediate layers with 16 input channels, 16 output channels, 5× 5 kernel. • 1 output layer with 16 input channels, 1 output channel, 7× 7 kernel.
Residual connections are used between the layers, and ELU (Clevert et al., 2016) non-linearities are applied. For the spatial dimension, circular padding is implemented to reflect the periodic boundary conditions, for the temporal dimension zero padding is used. The CNN output is a new vector di = (d 1 i ,d 2 i , ...,d K i ) with each element d k i corresponding to a different point in time. Analogously to the MP-PDE solver and analogously to the 1D-CNN, we use the output to update the solution as shown in Equation (63).
The 1D-CNN baseline which is conceptually very close to our MP-PDE solver performs much better than the 2D-CNN. However, we see already when looking at the results of E3 that generalization for the 1D-CNN across different PDEs becomes harder.

heading: H GENERALIZATION OVER MULTIPLE SPATIAL RESOLUTIONS
text: In Table 5, the generalization ability of MP-PDE solvers to unseen resolutions are tested. The columns mark the trained resolution, e.g. MP-PDE(250,100) is trained on a grid with (nt, nx) = (250, 100). The rows show the different tasks E1-E3 at different resolutions. MP-PDE solvers exhibit a strong potential for generalizing over multiple spatial resolutions, especially when the tested resolution is lower than the trained one. Most notably, training was only done on one resolution, and thus the network was not forced to exploit the full information injected by relative positions into the messages as done in Equation (11). Generalization from lower to higher spatial resolutions is tricky due to the different behavior when modeling shocks.
As detailed in Appendix F, for experiments WE1-WE3, cutoff radii for selecting neighbors are not a robust choice since the grids are irregular and relative distances are much lower close to the boundaries. We therefore construct our graphs via a k-NN criterion, and effectively use between 20 neighbors (highest spatial resolution) and 6 neighbors (lowest spatial resolution). Out of the box generalization to multiple spatial resolutions is in that setting therefore a bit tricky. However, this can be easily alleviated by including different spatial resolutions into the training.

Review summaries:
This paper proposes a message passing neural network to solve PDEs. As a graph neural network for solving PDEs, it proposes a few notable features, e.g. embedding in the encoder for equation-wise representation, a difference of inputs in the message function, and so on. The proposed method also adopts notable training methods to reduce error propagation. 

The paper addresses the stability of training neural PDE solvers. It uses the "pushforward trick": the adversarial-style loss is used by perturbing the input at step $k$ by a certain noise (note, that this idea looks similar to denoising autoencoders). Another trick is to predict K steps simulaneously in time. The architecture itself is graph-based: the nodes are distributed through the domain, and it is an encoder-processsor-decoder architecture. This helps to be closer to classical numerical solvers. The experiments show that the pushforward trick helps to increase "survival time" (which is defined as a time when a certain error exceed the threshold) clearly helps.  

The paper proposes a framework to solve generic PDEs written in conservation form, in an autoregressive fashion. The grid domain is considered as a graph, whose nodes contain information about solutions to the PDE until a given times steps, node positions, time steps, and hyperparameters information about the PDE. The node’s labels are first embedded in a latent space. Then message passing is applied and the output is decoded to predict the next n time steps solutions together. The network is trained with an additional loss which promotes zero stability of the solver, by constraining the network to be invariant to small perturbations in the input solutions of previous time steps. The approach is validated on a parameterized family of 1d and 2d PDEs, showing accuracy and generalization results that are above the state of the art of neural solvers and classical solver schemes. In addition, the paper show promising results dealing with irregular domains with different types of boundary conditions.