Message Passing Neural PDE Solvers
Authors: Johannes Brandstetter
Daniel E. Worrall
Max Welling
Venue: ICLR 2022 Spotlight (2022)
Keywords: neural PDE solvers
message passing
autoregressive models
zero-stability

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 backprop-optimized 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.

Published as a conference paper at ICLR 2022
MESSAGE PASSING NEURAL PDE SOLVERS
Johannes Brandstetter" University of Amsterdam Johannes Kepler University Linz brandstetter jku.at at
Daniel E. Worrall
Qualcomm Al Research
dworrall quwllguann..n.... com
Max Welling
University of Amsterdam m. we wingguin.a........
ABSTRACT
The numerical solution of partial differential equations (PDEs) is difficult, hav- ing led to a century of research so far. Recently, there have been pushes to build neural-numeries 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 which a generic solver would be faced, including: reso- lution, topology, geometry, boundary conditions, domain discretization regular- ity, dimensionality, etc. In this work, we build a solver, satisfying these proper- ties, where all the components are based on neural message passing, replacing all heuristically designed components in the computation graph with backprop- optimized neural function approximators. We show that neural message passing solvers representationally contain some classical methods, such as finite differ- ences, finite volumes, and WENO schemes. In order to encourage stability in training autoregressive models, we put forward a method that is based on the prin- ciple of zero-stability, posing stability idate our method on various fluid-like flow problems, demonstrating fast, stable, and accurate performance across different domain topologies, discretization, etc. in ID and 2D. Our model outperforms state-of-the-al numerical solvers in the low resolution regime in terms of speed and accuracy.
domain adaptation problem. We val-
INTRODUCTION
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 temporal partial differential equations (PDE). Solving these differential equations is of huge importance for problems in all numerate disciplines such as weather forecast- ing (Lynch] 2008), astronomical simulations (Courant et al.) 1967), molecular modeling (Lelièvre & Stoltz 2016) or jet engine design (Athanasopoulos et al.) 2009p. Solving most equations of impor- tance 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 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 compu- tational 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 maintaining stability over long rollouts and preserving invariants. It is precisely because of this considerable list of requirements that the field of numerical methods a splitter field (Bartels) 2016), tending build handcrafted solvers for each sub-problem, rather than lumper field, where a mentality of "one method to rule them all" reigns. This tendency is commonly justified with
Equua contribution
QQulcoomm AI Research is initiative of Qualcomm Technologies, Inc.

---

Published as a conference paper at ICLR 2022
reference to free lunch theorems. We propose to numerically solve PDEs with an end-to-end, neural solver. Our contributions can be broken down into three main parts: (i) An end-to-end fully neural PDE solver, based on neural message passing, which offers flexibility to satisfy all structural requirements 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 special cases of mes- sage passing. (ii) Temporal bundling and the pushforward trick, which are methods to encourage zero-stability in training autoregressive models. (iii) Generalization across multiple PDEs within given class. At test time, new PDE coefficients can be input to the solver.
2 BACKGROUND AND RELATED WORK
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.
2.1
PARTIAL DIFFERENTIAL EQUATIONS
We focus on PDEs in one time dimension t = [o, and possibly multiple spatial dimensions E X. These can be written down in the form
[21, 22, ................................
Tp) Tp]
0,u F(4, x, u, Uxu, Oxxu,...)
(1,x) E 00, x X
(1)
u(0,x) = (x),
B(ul(t,x)=0
X X, (t,x) E [0, x 0X
(2)
where u 10. X + R2t is the solution, with initial condition u (x) at time t = 0 and bound- ary conditions [tu] = 0 when x is on the boundary 0% of the domain X. The notation 0zu, uxxu, is shorthand for partial derivatives 0u/Ox, 0u/Ox², and so forth. Most notably, 0u/8x represents a n x D dimensional Jacobian matrix, where each row is the transpose of the gradient of the corresponding component of u. We consider Dirichlet boundary conditions, where the boundary operator B [u] = u by for fixed function by and Neumann boundary conditions, where Bulu] n 0zu by for scalar-valued u, where is an outward facing normal on O%.
Conservation form Among all PDEs, we hone in on solving those that can be written down in conservation form, because there already precedent in the field for having studied these (Bar-Sinal et al. 2019 Li et al. 2020a). Conservation form PDEs are written as
8, + J(u) = =0,
where V is the divergence of J. The quantity J R + R 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 Ju) through its boundaries.
2.2 CLASSICAL SOLVERS
[c)) of N small non- overlapping volumes called cells C X. In this work, we focus on grids of rectangular cells. Each cell has a center at x; u² used denote the discretized solution in cell e and time tg. There are two main ways compute ut: sampling uf u(te, x,) and averaging u Sc, u(te, x) dx. In our notation, omitting an index implies that we use the entire slice, so u = (u) .u.,.. )
Grids and cells Numerical solvers partition into finite grid X
Method of lines A common technique to solve temporal PDEs is the method of lines (Schiesser 2012), discretizing domain X and solution u into grid X and a vector u². We then solve 0,u lt, f(t, )²) for t [0, 2'] where f the form of F acting on the vectorized u² instead of the function ut.,x). The only derivative operator is now in time, making it an ordinary differential equation (ODE), which can be solved with off-the-shelf ODE solvers (Butcher, 1987; [verhart] 1985). f can be formed by approximating spatial derivatives the grid. Below are three classical techniques.
Finite difference method (FDM) In FDM, spatial derivative operators (e.g.. 0x) are replaced with difference operators, called stencils. For instance, 0_ur1z lz, might become (uf (u ///// /2(111111 xz). 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.
2

---

Published as a conference paper at ICLR 2022
Finite volume method (FVM) FVM works for equations in conservation form. It can be shown via the divergence theorem that the integral of u over cell increases only by the net flux into the J 1/2) where Ar; is the cell width, and JK-1/2- 1211/2 the flux at the left and right cell boundary at time (g. respectively. The problem thus boils down to estimating the flux at cell boundaries 1/ The beauty of this technique is that the integral of u is exactly conserved. FVM generally more stable and accurate than FDM, but can only be applied to conservation form equations.
cell. In ID, this leads to f(t, uf) X², (-/21/2 1/2
Pseudospectral method (PSM) PSM computes derivatives in Fourier space. In practical terms, the m derivative is computed IFFT (ww) (u))""FF(" for V-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.
2.3 NEURAL SOLVERS
We build on recent exciting developments in the field to learn PDE solvers. These neural PDE solvers, refer to them, laying the foundations of what is becoming both 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 Ta
Neural operator methods Neural operator methods treat the mapping from initial conditions to solutions at time as an input-output mapping learnable via supervised learning. For given PDE and given initial conditions un. neural operator M [0, x F + F, where J.' is a (possibly infiniteddimensionat) function space, trained satisfy
M(I. u²) = un))
Finite-diman dimensional operator methods (Raissi 2018
Sirignano & Spiliopoulos) 2018 Bhatnagar Khoo et al. 2020), where dim(F) < 0o are infiite-dimensiona operator methods (Li et al.] 2020c a Bhattacharya et al. [2021] Patel et al. 2021 by contrast resolve this is- sue. Each network is 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 L.
et al. 2019; Guo et al. 2016, Zhu & Zabaras 2018
gridddependent,
cannot generalize over geometry and sampling.
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 + At is computed as
u(t + At) A(At,u(t))
(5)
where A R50 x RN + RN is the temporal update. In this work, since At is fixed, we just write A(u(L)). 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 different class of PDE solver: finite volumes, multigrid, and iterative finite elements, respectively. Crucially, they all use a hybrid approach (Garcia Satorras al., 2019), where the solver computational graph is preserved and heuristically-chosse are predicted with neural network. Hsieh et al. 2019) even have convergence guarantees for their method, something rare 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 3 we explore why this is and seek to remedy it.
parameters
train. In Section
3 METHOD
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 message passing neural network.
3.1 TRAINING FRAMEWORK
Autoregressive solvers map solutions u² to causally consequent ones k+1 A straightforward way of training is one-step training. If (u is the distribution of initial conditions in the training set, and pa (u") p(u"|uu' Iu )o (u)) du is the groundtruth distribution at iteration k we minimize
Trone-step = EyE,, 11/uut [C(A(u*). n, x1)]
(
3

---

Published as a conference paper at ICLR 2022
Initial conditions
AR solver
AR solver
time
Neural operator
Neural operator Mapping free initial
Autoregressive model Kapping between terporally consecutive time steps
Temporal bundling calls solver reduces arror propagation speed
conditions to output
(a) Neural solver frameworks
(b) Temporal bundling
Figure 1: (a) LEFT: Neural operators perform a direct mapping from initial conditions to solution time i. RIGHT: Autoregressive models on the other hand compute the solution time based on the computed solution at fixed time offset before. (b) Our autoregressive solver outputs multiple time slices on every call, to reduce number of solver calls and therefore error propagation speed.
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 (which the vast majority of rollouts), and lead to divergence from the groundtruth. This can be interpreted 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 few rollout steps.
The pushforward trick We approach the problem in probabilistic terms. The solver maps + Arpk iteration k where A₁ : P(X) + P(X) ) is the pushforward operator for A and P(X) is the space of distributions on X. After single test time iteration, the solver sees samples from Appk instead of the distribution 1 and unfortunately A:pk * Pk++1 because errors always survive training. The test time distribution is thus shifted, which we refer to the distribution shift problem. This is domain adaptation problem. We mitigate the distribution shift problem by adding stability loss term, accounting for the distribution shift. A natural candidate is adversarial loss
Liaaiility = Ex uu [Ejju* [C(A(u* + +c), )]]]]
where e/u is adversarial perturbation sampled from appropriate distribution. For the per- turbation distribution, we choose such that (u* + E) ~ Azpk. This can be easily achieved by using (u* + E) = A(1 for = one step causally preceding u Our total loss then Lone-stop 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 backpropagate through the first unroll step. This is not only faster, also seems to more stable. Exactly why, we are not sure, but think it may be ensure the perturbations are large enough. Training the ad- versarial 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 e and there is some similarity Noisy Nodes (Godwin al.). where noise injection found 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 the condition that perturbations in the input conditions are damped out sublinearly in time, that is [|A(uu +) u < xi|El|, I. for appropriate norm and small k. The pushforward trick can be seen to minimize x directly.
The temporal bundling trick The second trick found be effective for stability and reducing rollout time is to predict multiple timesteps into the future synchronously. A typical temporal solver only predicts n r u²: whereas, we predict K steps u - (u², u² ,..., u²) = 11: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. schematic of this setup can be seen Figure Tb
3.2 ARCHITECTURE
We model the grid X as graph G (V,E) with nodes V, edges ij E E, and node features f, E Rs. The nodes represent grid cells e; and the edges define local neighborhoods. Modeling the domain graph offers flexibility over grid sampling regularity, spatial/temporal resolution,
4

---

Published as a conference paper at ICLR 2022
AR
AR
AR
AR
AR
AR
AR
One-step training Bradients flow back one time step only
Unrolled training Bradients flow back through all time steps
Pushforward training Bradients flow only through last time atep
One-step training only predicts solutions one step into the future. MIDDLE: Unrolled training predicts steps into the future. RIGHT: Adversarial training predicts N steps into the future, but only backprops on the last step.
Figure 2: Different training strategies. LEFT:
domain topology and geometry, boundary modeling and dimensionality. The solver a graph neural network (GNN) (Scarselli et al. 2009 Kipt & Welling, 2017 Defferrard et al. 2016 Gilmer et al.,
2017 Battaglia et al. 2018). 8). representationally containing the function class of several classical Peeee-Proess--eeee
solvers, see Section 3.2 We follow the
framework Ba et al. (2018) and Sanchez-Gonzan ez et al. (020) with adjustments. We are not the first to use GNNs as PDE solvers eLi all]a De Avila Belbute- Perlutee-Pereseeee 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 it maps the last K solution values u k- K-A node position X,, current time tg. and equation embedding θPDE to node embedding vector f c" ((u Kik: X,, lx. OPDE]): ОPDE contains the PDE coefficients and other attributes such as boundary conditions. An exact description is found in Section The inclusion of OPDE allows us train the solver multiple different PDEs.
The processor computes M steps of learned message passing, with intermediate graph representations G G ,... G.M The specifie updates we use are
Processor
edge j message:
m
(f,
k Kikk
#
Kik
Xj. OPDE)
(8)
node update:
f. pm+
cm
OPDE
JEN(()
where Mii) holds the neighbors of node i, and 6 and t are multilayer perceptrons (MLPs). Using relative positions x; - x; can be justified by the translational symmetry of the PDEs we consider. Solution differences u, uj make sense by thinking of the message passing local difference operator, like a numerical derivative operator. Parameters θYDE are inserted into the message passing similar to Brandstetter et (2021)
Decoder After message passing, we use a shallow ID convolutional network with shared weights across spatial locations to output the K next timestep predictions at grid point x,. For each node i, the processor outputs a vector f.M. We treat this vector as a temporally contiguous signal, which feed into CNN time. The CNN helps 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 new vector d; = (d) df d with each element d corresponding to a different point in time. We use this to update the solution as
b+t u = u. (la+e - Lx) d
1 TSESK.
(10)
The motivation for this choice of decoder has to do with property called consistency (Amold, 2015), which states that lim A[At, u²) u(At)] = 0, i.e the prediction matches the exact solution in the infinitesimal time limit. Consistency is requirement for zero-stability of the rollouts.
5

---

Published as a conference paper at ICLR 2022
Encoder Node-wil napping the hidden sgace
Processor Message passing
Decoder shallow 10 convolution
MLP
KPWM
CHN
KPNN
CHN
KPNN
CNN
Figure 3: Schematic sketch of our MP-PDE Solver.
Connections. As mentioned in Bar-Sinai al. (2019), both FDM and FVM are linear methods, which estimate n -order point-wise function derivatives as
[0(m) u/:
Σ JEN(()
( Uy
(11)
for appropriately chosen coefficients aj where where.N(i) is the neighborhood of cell FDM computes this cell centers, and FVM computes this cell boundaries. These estimates are plugged into flux equations, see Table 3 in the appendix, followed by an optional FVM update step, to compute time derivative estimates for the ODE solver. The WENO5 scheme computes derivative estimates by taking an adaptively-weighte average over multiple FVM estimates, computed using different neighborhoods of cell (see Equation 23). The FVM update, Equation [1] and Equation 23 are just message passing schemes with weighted aggregation (I layer for FDM, 2 layers for FVM, and layers for WENO). It is through this connection, that we see that message- message-passing neural networks representationally contain these classical schemes, and are thus well-motivated architecture.
4 EXPERIMENTS
We demonstrate the effectiveness of the MP-PDE solver on tasks of varying difficulty to showcase its qualities. In ID, we study its ability generalize 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, demonstrate its utility. As baselines, we compare against standard classical PDE solvers, namely; FDM, pscudospectral methods, and a WENOS 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
4.1 INTERPOLATING BETWEEN PDEs
Data
We focus on the family of PDEs
(0,u Oz (on" 30,.u 70. u)] = 5(1,,),
(12)
u(0, x) = 2(0, 8(1,z) = (ma)=<on Aú sin(wit + 2nlrr/La Φ;).
(13)
j=1
Writing OPDE (a, B, 7). corner cases are the heat equation OPDE = (0, n,..0 0). Burgers' equation OPDe (0.5, 1/, 0), and the KdV equation OPDE (3, 0, 1). The term is forcing term, fol- 5, L = 16 and coefficients sampled uniformly in A, E [-0.5, 0.5]. wy E 0.4, 0.4.6, E (1. o, [0. 2x). This setup guarantees periodicity of the initial conditions and forcing. Space is uniformly discretized to n, = 200 cells in [0, 16) with peri- odic boundary and time is uniformly discretized to n, 200 points in [0, 4). Our training sets consist of 2096 trajectories, downsampled to resolutions (n,, nl.)) E {(250, 100), (250. 50). (250. 40)). }. Nu- merical groundtruth is generated using a 5odder WENO scheme (WENO5) (Shu 2003) for the convection term u, and -orde finite difference stencils for the remaining terms. The temporal solver is an explicit Hange-Kutta 4 solver (Rungel 1895; Kutta 1901) with adaptive timestepping. Detailed methods and implementation are in Appendix C All methods 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 establish utility in generating groundtruth.
lowing Bar-Sinai et al. 2019), with

---

Published as a conference paper at ICLR 2022
Experiments and results We consider three scenarios: E1 Burgers' equation without diffusion
OPDE = (1. (1.0.0) for shock modeling; E2 Burgers' equation with variable diffusion PDEE = (1. 17. 3.0, where 0 - 7 0.2; and E3 mixed scenario with ÓPDE (a, (4, 7) where 0.0
0.0 S 3 s 0.4 and 0.0 S S 1.0. E2 and E3 test the generalization capability. We compare against downsampled groundtruth (WENO5) and variation of the Fourier Neural Operator with autore- gressive 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-O-me OPDE), without ODDE features, and the Fourier Neural Operator method trained using our temporal bundling and pushforward tricks (FNO-PF). Errors and runtimes for experiments are in Table [ We see that the MP-PDE solver outperforms WENO5 and FNO-RNN in accuracy. Temporal bundling and pushforward trick improve FNO dramatically, to the point where beats MP-PDE on E1. But MP-PDE outperforms FNO-PF on E2 and E3 indicating that FNO is best for single equation mod- cling, but MP-PDE is better at generalization. MP-PDE predictions are best equation parameters OPDE 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 [.1 Figure +
(TOP) shows shock formation different resolutions (E1), traditionally difficult phenomenon to model-FDD and PSM methods cannot model shocks. Strikingly, shocks are preserved even at very low resolution.
4.2 VALIDATING TEMPORAL BUNDLING AND THE PUSHFORWARD METHOD
We observe solver survival times on E1, defined as the time until the solution diverges from groundtruth. solution ú diverges from the groundtruth u when its max-normalized L -ererr
exceeds 0.1. The solvers unrolled n, 1000 timesteps with T 16 i =u Ta maxy ju
Examples shown in Figure (BOTTOM), where observe increasing divergence after 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 Sanchez.
Gonzalez et al. (0200) 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 we plot the survival ratios for models trained with and without the pushforward trick. As a third comparison show model trained with Gaussian noise adversarial perturbations, similar to that proposed in Sanchez-Gonzalle et al. (2220). We see that applying the pushforward trick leads to far higher survival times, confirming model that instability can be addressed with adversarial training.
Exemplary 1D rollout of shock formation at different resolutions
0.
1.0
twb
1.34
8.0
0-)
1-2
-1.0
1-2
ground
predction
prddciion
predction
2-3.2
b-335
(-3.05
Exemplary 2D long rollout of shock formation at different resolutions
15.0
1.5
12.5
1.0
(seconds)
10.0
0.5
7.5
0.0
5.0
005
2.5
nx=200
nx=100 prediction
rx=50 prediction
nx=40 prediction
0.0
ground truth
-1.5
8
Figure 4: TOP: Exemplary ID rollout of shock formation at different resolutions. The different colors represent PDE solutions at different timepoints. Both the small and the large shock are neatly captured and preserved even for low resolutions; boundary conditions perfectly modeled. Bor- TOM: Exemplary long 2D rollout of shock formations over 1000 timesteps. Different colors repre- sent PDE solutions at different space-time pace-timee points.

---

Published as a conference paper at ICLR 2022
Table Error and runtime experiments targeting shock wave formation modeling and generalization to unseen equations. Runtimes are for one full unrolling over 250 timesteps on a GeForce RTX 2080 Ti GPU. FNO-PF, MP-PDE PEDD-ODDE, and MP-PDE are all ours. Accumulated error n2 Σx, MSE.
Accumulated Error 4
Runtime
(Rt, R2)
WENOS
FNO-RNN
FNO-PF
MP-PDE-Or6E
MP-PDE
WENO5
MP-PDE
(250, 100)
2.02
11.93
0.54
1.55
1.9
0.09
(250,50)
6.23
29.98
0.51
1.67
1.8
0.08
(250, 40)
9.63
10.44
0.57
1.47
1.7
0.08
E2
(250, 100)
1.19
17.09
2.53
1.62
1.58
1.9
0.09
E2
(250, 50)
5.35
3.57
2.27
1.71
1.63
1.8
0.09
E2
(250, 40)
8.05
3.26
2.38
1.49
1.45
1.7
0.08
E3
(250, 100)
4.71
10.16
5.69
4.71
4.26
4.8
0.09
E3
(250,50)
11.71
14.49
5.39
10.90
3.74
4.5
0.09
E3
(250, 40)
15.94
20.90
5.98
7.78
3.70
4.4
0.09
Accurate rollouts shock formation
Accurate rollouts formation ablation
1.0
1.0
0.8
0.8
ratio
0.6
0.6
0.4
no pf
0.4
nx-100
gaussian
nx=50
0.2
pf gradients
0.2
nx=40
pf
0.0
0
100
200
300
400
500
50
100
150
200
250
nt
(a)
(b)
Figure 5: Survival times at E1. Rollout for long trajectories of 8 (left), pushforward (pf) ablation
(right). The ablation compares survival times resolutions nt, = 100 (solid) and nl, = 50 (dashed) against survival times using pushforward (no pf), pushforward but putting Gaussian noise (a 0.01), pushforward but without cutting the gradients (pf gradients).
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.
4.3 SOLVING IRREGULAR GRIDS WITH DIFFERENT BOUNDARY CONDITIONS
investigate (i) how well 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.
The underlying motives for designing this experiment
Data We consider simple ID wave equation
du 20
E [-8,8
(14)
where c is wave velocity (c = in our experiments). We consider Dirichlet Blu] = u and Neumann B[u]] O,u 0 boundary conditions. This PDE is order in time, but can be rewritten as -orrrder in time, by introducing the auxilliary variable v 0,u and writing 0,(u, v] [v, 0) 0 The initial condition is Gaussian pulse with peak random loca- tion. Numerical groundtruth is generated using FVM and Chebyshev spectral derivatives, integrated time with an implicit Runge-Kutte Kutta method of Radau IIA family, order (Hairer et al. 1993). We solve for groundtruth at resolution (nt, n2z) (250, 200) on Chebyshev extremal point grid (cell edges are located z; = cos(ist (n, 1))-
Experiments and results We consider three scenarios: WE1 Wave equation with Dirichlet bound- ary conditions; WE2 Wave equation with Neumann boundary conditions and, WE3 Arbitrary com- binations of these two boundary conditions, testing generalization capability of the MP-PDE solver.
8

---

Published as a conference paper at ICLR 2022
Ablation studies (marked with (rDE) 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.
Table 2: Error and runtime comparison on tasks with non-periodic boundaries and irregular grids. Runtimes measure one full 250 timesteps unrolling on a GeForceRTX 2080 Ti GPU for MP-PDE solvers, and on CPU for our pseudospectral (PS) solver implementation based on scipy.
MSE 1 (WEI)
Σ MSE 4 (WE2)
MSE + (WE3)
Runtime [s]
(Rt. nly)
PS
MP-PDE
PS
MP-PDE
PS
MP-PDE OPgE
MP-PDE
PS
MP-PDE
(250, 100)
0.004
0.137
0.004
0.111
0.004
38.775
0.097
0.60
0.009
(250, 50)
0,450
0.035
0.681
0.034
0.610
20.445
0.106
0.35
0.09
(250, 40)
194.622
0.042
217.300
0.003
204.298
16.859
0.219
0.25
0.09
(250, 20)
breaks
0.059
breaks
0.007
breaks
17.591
0.379
0.20
0.07
4.4
2D EXPERIMENTS
We finally test the scalability of our MP-PDE solver to higher number of spatial dimensions, more specifically to 2D experiments. We use data from PHIFLOWW[ an open-source fluid simula- tion toolkit. We look fluid simulation based 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 corre- sponds to one second. Smoke inflow locations sampled randomly. Architectural details are in Appendix[F 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.
5 CONCLUSION
We have introduced a fully neural MP-PDE solver, which representationally contains classical meth- ods, 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 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 predict the solution of PDEs, but e.g. also be interpreted to optimize the integration grid and the parameters of the PDE. For the former, simply position updates need be included in the processor, similar to Satorras al. (202T). 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 common criticism of such learned numerical methods. potential fix would be to fuse this work with that in probabilistic numerics (Hennig et al., 2015). has been done for RK4 solvers (Schober et al., 2014). Another promising follow-up direction is to research alternative adversarial-stylc losses as introduced in Equation 7 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 unique set of symmetries (Olver 1986).

---

Published as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
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 imple- mented WENO scheme against two analytical solutions of the Burgers equation from literature. For experiments WI, W2, W3, we cross-checked if Gaussian wave packages keep their form through- propagation phase. Furthermore, the wave packages should change sign for
out the whole 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 pushfor- ward 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 er al. (2200a). We have not introduced new mathematical results. However, we have used data gener- ation concepts from different mathematical fields and therefore have included a detailed description of those in our appendix.
For reproducibility, we provide our at https//gttttcccommmmmttoccccccee Solvers
ETHICAL STATEMENT
The societal impact MP-PDE solvers difficult to predict. However, as stated in the introduc- tion, 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 pave the way towards shortcuts for computationally expensive simulations. Most notably, drastical computational shortcut is always somehow related to reducing the carbon foot- print. However, this regard, it also important to remind ourselves that relying on simulations or now even shortcuts of those always requires monitoring and thorough quality checks.
ACKNOWLEDGMENTS
Johannes Brandstetter thanks the Institute of Advanced Research Artificial Intelligence (IARAI) and the Federal State Upper Austria for the support. The authors thank Markus Holzleitner for helpful comments on this work.
REFERENCES
Douglas N Amoll. Stability, consistency, and convergence numerical discretizations. 2015.
Michael Athanasopoulos, Hassan Ugail, and Gabriela Gonzalez. Parametric design of aircraft ge- ometry using partial differential equations. Advances in Engineering Software, 40, 2009.
Yohai Bar-Sinai, Stephan Hoyer, Jason Hickey, and Michael P. Brenner. Learning data-driven dis- cretizations for partial differential equations. Proceedings of the National Academy of Sciences, 116(31):15344 15349, Jul 2019. ISSN 1091-6490. doi: 10.1073/pnas. 18 141581166 16.
Sören Bartels. Numerical Approximation of Partial Differential Equations. Springer, 2016.
Claude Basdevant, Michel Deville, Pierre Haldenwang, Jean Lacroix, Jalil Ouazzani, R. Peyret, Paolo Orlandi, and Patera A.T. Spectral and finite difference solutions of the burgers equation. Computers and Fluids, 14(1):23-41, 1986.
Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzaler, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Caglar Gulcehre, Francis Song, Andrew Ballard, Justin Gilmer, George Dahl, Ashish Vaswani, Kelsey Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matt Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 006.0 2018.
10

---

Published as a conference paper at ICLR 2022
Saakaar Bhatnagar, Yaser Afshar, Shaowu Pan, Karthik Duraisamy. and Shailendra Kaushik. Predic- tion of aerodynamic flow fields using convolutional neural networks. Computational Mechanics, 64(2):525-544 Jun 2019.
Kaushik Bhattacharya, Bamdad Hosseini, Nikola B. Kovachki, and Andrew M. Stuart. Model re- duction and neural networks for parametric pdes, 2021.
Johannes Brandstetter, Rob Hesselink, Hlise van der Pol, Erik Bekkers, and Max Welling. Ge- ometric and physical quantities improve c(3) equivariant message passing. arXiv preprint arXiv:2110.02905. 2021.
John Charles Butcher.
Coefficients for the study of runge-kutta integration processes. Journal of the Australian Mathematical Society, 3(2):185-201, 1963.
John Charles Butcher. The numerical analysis of ordinary differential equations: Runge-Kutta and general linear methods. Wiley-Interscience, 1987.
Djórk-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). In International Conference on Learning Representa- fions (ICLR). 2016.
Richard Courant, Kurt Otto Friedrichs, and Hans Lewy. On the partial difference equations of mathematical physics. IBM J. Res. Dev. 11(2):11--24, 1967. ISSN 0018-8646.
Filipe De Avila Belbute-Pere Thomas Economon, and Zico Kolter. Combining differentiable PDE solvers and graph neural networks for fluid flow prediction. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 2402-2411. PMLR, 13-18 Jul 2020.
Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), ), Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
Edgar Everhart. An efficient integrator that uses gauss-radau spacings. International Astronomical Union Colloquium, 83:185-202, 1985.
Victor Garcia Satorras, Zeynep Akata, and Max Welling. Combining generative and discriminative models for hybrid inference. In H. Wallach, H. Larochelle, A. Beygelzimer, d' Alché Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 13802- 13812. Curran Associates, Inc., 2019.
Justin Gilmer, Samuel Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning, pp- 1263-1272. PMLR, 2017.
Jonathan Godwin, Michael Schaarschmidt, Alexander Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Velickovic, James Kirkpatrick, and Peter W. Battaglia. Very deep graph neu- ral networks via noise regularisation. arXiv preprint 1rXiv:2106.07977.
Daniel Greenfeld, Meirav Galun, Ronen Basri, Irad Yavneh, and Ron Kimmel. Learning to opti- mize multigrid PDE solvers. In Proceedings of the 36th International Conference Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, pp. 2415-2423, 2019.
Xiaoxiao Guo, Wei Li, and Francesco Iorio. Convolutional neural networks for steady flow ap- proximation. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD '16, pp. 481-490. Association for Computing Machinery, 2016.
Emsst Hairer, Syvert Nørsett, and Gerhard Wanner. Solving Ordinary Differential Equations I (2nd Revised. Ed.): Nonstiff Problems. Springer-Verlle Berlin, Heidelberg, 1993. ISBN 0387566708.
11

---

Published as a conference paper at ICLR 2022
Philipp Hennig. Michael A. Osborne, and Mark Girolami. Probabilistic numerics and uncertainty
in computations. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, 471(2179):015011222 Jul 2015. ISSN 1471-2946. doi: 00.1098/rsaa.2015.1122 URL
http::/dxx doi http://xx..oi...r///..././/////. www./oog//...... 1098 rspa. 2015. 0142]
Jun-Ting Hsieh, Shengjia Zhao, Stephan Eismann, Lucia Mirabella, and Stefano Ermon. Learning neural PDE solvers with convergence guarantees. arXiv preprint urXiv:1906.012000 2019.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456. PMLR, R, 2015.
Guang-Shan Shan Jiang and Chi-Wang Shu. Efficient implementation weighted eno schemes. Journal of Computational Physics, 126(1):202 228, 1996. ISSN 0021-9991.
Yuehaw Khoo, Jianfeng Lu, and Lexing Ying. Solving parametric pde problems with artificial neural networks. European Journal of Applied Mathematics, 32(3):421-435, Jul 2020. ISSN 1469-4425.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional net-
works. In International Conference on Learning Representations (ICLR), 2017.
Wilhelm Kutta. Beitrag zur näherungsweisen integration totaler differentialgleichunge.. Zeitschrift für Mathematik und Physik, 46:435 453, 1901.
Tony Lelièvre and Gabriel Stoltz. Partial differential equations and stochastic methods in molecular dynamics. Acta Numerica, 25:681 13,0,,880, 2016
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08892 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Multipole graph neural operator for parametric partial differential equations. arXiv preprint 2Xiv:2006.09535, 2020b.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differ- ential equations. arXiv preprint arXiv:2003.03488 2020c.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 10000017,,.0 2017.
Peter Lynch. The origins of computer weather prediction and climate modeling. Journal of compu- tational physics, 227(7):3431 2008
Andreas Mayr, Sebastian Lehner, Arno Mayrhofer, Christoph Kloss, Sepp Hochreiter, and Jo- hannes Brandstetter. Boundary graph neural networks for 3d simulations. arXiv preprint arXiv:2106.1 12299, 2021.
Peter Olver. Introduction to partial differential equations. Springer, 2014.
P.J. Olver. Symmetry groups of differential equations. Applications of Lie Groups Differential Equations, pp. 77-185. Springer, 1986.
Ravi G. Patel, Nathaniel A. Trask, Mitchell A. Wood, and Eric C. Cyr. A physics-informed oper- ator regression framework for extracting data-driven continuum models. Computer Methods in Applied Mechanics and Engineering, 373:1 13500, Jan 2021. ISSN 0045-7825.
Maziar Raissi. Deep hidden physics models: Deep learning of nonlinear partial differential equa- tions. J. Mach. Learn. Res., 19:25:1-25:24, 2018.
Prajit Ramachandran, Barret Zoph, and Quoe V. Le. Searching for activation functions. arXiv preprint arXiv: arXiv: 1710.0544, 2017.
12

---

Published as a conference paper at ICLR 2022
Carl Runge. Über die numerische auflösung von differentialgleichungen. Mathematische Annalen, 46:167 178, 1895.
Alvaro Sanzhez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Pe- ter W. Battaglia. Learning simulate complex physics with graph networks. arXiv preprint arXiv:2002.09405. 2020.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling, E (n) equivariant graph neural net- works. arXiv preprint arXiv:2 122.09844, 2021.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions Neural Networks, 20(1):61-80, 2009.
William E Schiesser. The numerical method of lines: integration of partial differential equations. Elsevier, 2012.
Michael Schober, David Duvenaud, and Philipp Hennig. Probabilistic ode solvers with runge-kutta means, 2014.
C.-W. Shu. High-order finite difference and finite volume weno schemes and discontinuous galerkin methods for cfd. International Journal of Computational Fluid Dynamics, 17(2):107-118, 2003.
Justin Sirignano and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial differential equations. Journal of Computational Physics, 375:1339 - 1364, 2018. ISSN 0021- 9991. doi: https://doi.org/100 1016// .pp.0.....200
Eitan Tadmor. The exponential accuracy fourier and chebyshev differencing methods. SIAM Journal on Numerical Analysis, 23:1-10, 1986.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing in- gredient for fast stylization. arXiv preprint arXiv:1607.08022, 2016.
Yinhao Zhu and Nicholas Zabaras. Bayesian deep convolutional encoder-decoder networks for sur- rogate modeling and uncertainty quantification. Journal of Computational Physics, 366:415-447, Aug 2018. ISSN 0021 19911
13

---

Published as a conference paper at ICLR 2022
A INTERPOLATION
To compute classical numerical derivatives of a function u R + R. which has been sampled on mesh of points ø1 < ø2 ... < zN it is common first piecewise polynomial p : R on the mesh. We assume we have function evaluations and once we have fitted the polynomial, we can use its derivatives any new off-mesh point, for instance the half nodes 4 Here we illustrate how to carry out this procedure.
u(x;) at the mesh nodes for all i = 1,..,
A polynomial of degree N 1 (note the highest degree polynomial we can fit to N points has degree 1) can fitted the mesh points by solving the following linear system in a
N
u1
p(a1)
i==
a
(15)
UN
p(xxx)
aN
==
To find the polynomial interpolation at new point r, we then do
p(z) = x = x
(16)
where x"
11 x²
N-1].
To retrieve the mh derivative, where < N 1, is also very
simple. For this we have that
dmpp p dr
N-1
d
N-1
Σ
(1)m xf (m)T a = (m) u, (17)
a dr
aj
1=0
1=0
where x (1a)
m00)
0-m (1)mxt-m (N 1)mx
(18)
where (i)m i(i 1) ... (i m + 1) the mh Pochhammer symbol and (t)o :
Note that is typical to fold = x x into single object, which we call stencil.
B RECONSTRUCTION
Reconstruction is the task of fitting piecewise polynomial on mesh of points, when instead of functions values at the grid points we have cell averages 01 where
ū: - / u(x) da,
1
(19)
where each cell I; [a,_ J. Fe4 g. the half nodes are defined as Ti+ 2 [ax + x:41) and
Art = 4++g - 2 The solution is to note that we can polynomial P(r) to the integral of u(x), which will be exact at the half nodes,
E/
Σ F u(x) do Σ Ur Azg.
p( = ((1+4))
dr
(20)
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 (m). The system of equations we need to solve is thus
Ari
(21)
Ar1
ATN
UN
N
ON-1
X
where û is the vector of cell averages and L is a lower triangular matrix performing the last sum in Equation 20 Thus the thth derivative of the polynomial p(x) P'(x)) is
dmp
am+1p
x
(m+1)T
1)T
X Lu = {mm) u.
(22)
dx
dr
14

---

Published as a conference paper at ICLR 2022
C WENO SCHEME
The essentially non-oscillating (ENO) scheme an interpolation or reconstruction scheme to esti- mate functiin values and derivatives of a discontinuous function. The main idea is use multiple overlapping stencils to estimate a derivative at point E [Tt T++1]. We design N stencils to fit the function on shifted overlapping intervals 4 ..... IN. where 42 (ii-N+1+ke Tz++l. In the case of WENO reconstruction these intervals are 4 = -=i-N+++ ++
1 - If a discontinuity lies in
I=U; JR, then it is likely that one of the substencils s(.T) (m)) (defined interval I&) will not contain the discontinuity. We can thus use the substencil (m) or s (Tm) from the relatively smooth region estimate the function derivatives z. In the following, we focus on WENO reconstruction.
The weighted essentially non-oscillating (WENO) scheme goes one step further and takes convex combination of the substencils to create larger stencil I, where (dropping the superscript for brevity)
N
SkSS-
(23)
k=1
Here the nonlinear weights wy satisfy nonlinear weights can be constructed as
Ek=1 wk
N
1. It was shown Jiang & Shu 1996) that these
7kk (e + 3x)2²
Wk
Wk
(24)
Σj= w/
where "Yk is called the linear weight, is a small number, and 3h is the smoothness indicator. The linear weights are set such that the sum E 7A8& over the order N - 1 stencils matches a larger order 2N stencil. The smoothness indicators are computed as
dd' Pkc drm
2
An
dx,
(25)
where p& is the polynomial corresponding to substencil Sg.
C.I WENO5 SCHEME
A conservative finite difference spatial discretization approximates derivative f(u), by conser- vative difference
1 (u)alamm; Ar (1+t--1-9)
(26)
where fit3 and f. are numerical fluxes. Since g(u) is approximated in the same way, fi- nite difference methods have the same format for more than one spatial dimensions. The left- reconstructed (reconstruction is done from left to right) fifth order finite difference WENO scheme (WENO5) Shu 2003) has the ú i given by:
wiū wia + H22 143 W/ (3)
(2)
(27)
(j)
In equation 27 à
are the left WENO reconstructions on three different stencils given by
#+
(1) =
1
11
+
Ui-2
Uj-1
04
(28)
3
6
6
(2)
=
U;-1
+
U4
+
U4
(29)
6
6
3
(3)
1
u,
+
5
U;+1
41+2
(30)
6
and the non-linear weights wy given by
7
Wi
Wh
(31)
Σ
3
Jm
wk
(e + 8x)3
15

---

Published as a conference paper at ICLR 2022
10 5 ). and tiny-valued parameter to avoid the denominator becoming 0. The left smoothness indicators By are given by:
with 7(1,2,3)
13
1
BT 24-2- -14-1 + uz)²
[1i-2 4u;-1 +.34))
(32)
12
4
13
8 = [4-1 20; + )i+1)²
(ui-1 - -1++1
(33)
12
4
B3
13
2 14 2u,+1 + u(+2)
(us 4us+1 + 324+22
(34)
12
4
(35)
The right-reconstructed WENO5 scheme has the at= given similarly to "7 but with all coef-
ficients flipped since the reconstruction is done from the other side (from right to left). The right reconstruction three different stencils are given by
+11)
=
U142
Us+1
+
u4
(36)
2 3
6
+()
[- 6
U1+1
+
U4
+
U1
(37)
(3) 1
5
+
U4-2
(38)
3
and the right smoothness indicators 3 are given by
B
13
[1+2 - 2u++1 + uz)
2
(u++2 - -44;+1 + =us))
(39)
12
4
13
1
[1+11
2u; + =2=11)²
(ui+1 -w;-1)²
(40)
13
1
8 12
-21 + 22)
(u: 44g-1 + 32)
(41)
4
(42)
Both â- and tt needed for full flux reconstruction explained in the next section.
C.2 FLUX RECONSTRUCTION
We consider flux reconstruction via Godunov. For Godunov flux, f.+ = f(u,+ 3) is reconstructed
from at and â 1+1 ++1
via:
min SuSu
f(u), if uf S S4
f(u++))
(43)
f(u), if duu u
SuSu
C.3 COMPARING WENO SCHEME TO ANALYTICAL SOLUTIONS
First analytical case. An analytical solvable case for Burgers equation arises for the boundary conditions
a(t, 0) u(t, 2TT)
and the initial conditions
u(0,z)) -2v ath/Or
+ 4,
(44)
16

---

Published as a conference paper at ICLR 2022
where
x²²
(xc 2n)²
Φ
exp
+ exp
(45)
4v/
4 + 1)
2x
2
2(x 2
(x 2T)
exp
+ exp
(46)
4v
4v
4v
0.5z
x2²
0.5(% 2
(x: )x)²
exp
exp
(47)
4v
V
4v
The analytical solutions for this specific set of boundary and initial conditions gives
u(t,x) = -2v 00/0r
+ 4
(48)
Φ
where
(x 4
41 2
0 exp
exp
(49)
Av(1 + 1)
Au(L + 1)
06
2(x 4t)
(x 4t)
2(x - 2
(x 4t 2
exp
+ exp
(50)
Ox
4u(t + 1)
4u(t + 1)
4u(t + 1)
4v(t 1)
0.5(I 4f)
-(x 4t)²
0.5(x 4t 2 v(1 +
4t )
exp
exp
(51)
11
4v(1+
Av(1 + 1)
For this first analytical solveable case, the analytical solution, the WENO scheme and the fourth order finite difference scheme (FDM) are compared in Fig.6foor a diffusion term of v = 0.005. The WENO scheme models the analytical solution perfectly, whereas the FDM scheme fails capture the shock accurately. For lower values of v the effect gets even stronger.
Second analytical case. boundary condition:
Another analytical solvable case for the Burgers equation arises for
u(t, +1) 0,
(52)
and the initial condition
u(0, x) - sin(.
(53)
Solutions (Basdevant al. 1986)
u((,, x:) = sin ( 1)f(2 17) exp(-1p /4rt)dn
(54)
1 f(r n) exp( 1) 44tt)n
with f(y) = exp(- cos( 2mn
Using Hermite integration allows the computation accurate results
up to t-3T.
For this second analytical solveable case, the analytical solution, and the WENO scheme are com- pared in Fig. for diffusion term of v 0.002. The WENO scheme models the analytical solution perfectly. Modeling via the FDM scheme fails completely.
17

---

Published as a conference paper at ICLR 2022
Analytical solution, 1D rollout
Analytical solution, 2D rollout
0.5
t=0.0s
t=0.15
0.4
t=0.1s
(m/s)
t=0.25
(seconds)
0.3
5
t=0.25
t=0.3s
0.2
t=0.45
3
t=0.45
0.1
t=0.5s
0.0
WENO solution, 1D rollout
WENO solution, 2D rollout
7
t=0.0s
t=0.1s
0.4
t=0.1s
(mis)
t=0.25
(seconds)
0.3
t=0.25
t=0.3s
0.2
t=0.4s
t=0.4s
t=0.5s
0.1
2
1
0.0
4
2
x
X
FDM solution, 1D rollout
FDM solution, 2D rollout
12
0.5
t=0.05
10
t=0.1s
0.4
10
t=0.1s
(mis)
t=0.25
(seconds)
0.3
8
t=0.25
t=0.3s
0.2
t=0,4s
t=0.4s
t=0.5s
2-
0.0
Figure 6: 1D and 2D rollouts for the first analytical case setting the diffusion term 0.005. Analytical solution (top), WENO scheme (middle) and Finite Difference scheme (FDM, bottom). The WENO scheme models the analytical solution perfectly, whereas the FDM scheme fails to capture the shock accurately.
18

---

Published as a conference paper at ICLR 2022
Analytical solution, 1D rollout
Analytical solution, 2D rollout
1.0
t=0.0s
t=0.15
t=0.25
0.5
t=0.4s
(seconds)
0.6
(m/s)
t=0.5s
0.0
t=0.6s
0.4
t=0.7s
-0.5
t=0.85
0.5
t=1.0s
-1.0
0.0
-1.0
-1.0
-0.5
0.0
1.0
-0.5
0.0
0.5
WENO solution, 1D rollout
WENO solution, 2D rollout
1.0
1.0
t=0.0s
t=0.1s
0.8
0.5
t=0.25
0.5
t=0.4s
(seconds)
0.6
(m/s)
0.0
t=0.5s
0..
t=0.6s
0.4
t=0.7s
-0.5
t=0.85
0.2
t=1.0s
-1.0
0.0
-1.0
-1.0
-0.5
0.0
0.5
1.0
-1.0
-0.5
0.0
0.5
1.0
x
Figure 7: ID and 2D rollouts for the second analytical case setting the diffusion term v 0.005. Analytical solution (top), and WENO scheme solution(bottom). The WENO scheme models the an- alytical solution perfectly.
19

---

Published as a conference paper at ICLR 2022
D PSEUDOSPECTRAL METHODS FOR WAVE PROPAGATION ON IRREGULAR GRIDS
We consider Dirichlet B[u] - 0 and Neumann B[u] 0r u 0 boundary conditions. Numeri- cal 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.] [993). To properly ful- fill 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
Wave equation, Dirichlet boundary conditions
Wave equation, Neumann boundary conditions
6.-
06-
4
-10
-65
04
Figure 8: Exemplary wave propagation data for Dirichlet boundary conditions (left) and Neumann boundary conditions (right). Solutions are obtained on irregular grids using pseudospectral solvers.
20

---

Published as a conference paper at ICLR 2022
E EXPLICIT RUNGE- METHODS
The family of Runge- Kutta methods
[utchhr 1987) is given by:
Ut+11
Mt A
b,ki
(55)
where
h1 f(nn ut,)
(56)
22 f(n, + og un + h(a21k1))
(57)
h3 f(tn + ct ue, + h(aziki a32k22)
(58)
(59)
Kg f(tn + c, ut, + h(asiki + a,2k2, (kia-1
(60)
For a particular Runge-Kutta method one needs to provide the number of stages 8, and the coef-
ficients (1 S < i S s), b, s.2.....s) and c 1.2...., s). These data are usually arranged in so-called Butcher tableaux (Butcher 1963).
21

---

Published as a conference paper at ICLR 2022
F EXPERIMENTS
Flux terms of equations that we the Heat, Burgers. -rrieweg-de-Vries (KdV), and Kuromoto-S Kuromoto-Shivashinsyy Shivashinsky (KS) equation are summarized in Table
Table 3: ID flux terms J(u) of the Heat, Burgers, Korteweg-de-Vries (KdV), and Kuramoto-
Shivashinsky (KS) equation.
Heat
Burgers
Kdv
KS
((u)
- nu
1 u nd2u
3u + 0rpu
1 + 02u + Ozzz u
2
Pushforward trick and temporal bundling. ward trick and temporal bundling sketch Algorithm
Pscudocode for one training step using the pushfor-
Algorithm Pushforward trick and temporal bundling. For given batched input data trajectory and a model, we draw a random timepoint I, get our input data trajectory, perform N 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 apply the pushforward trick, T' is the number of available timesteps in the training set.
Require:
data, model, N, K, T 1 + DrawRandomNumber & E (1...., T}
D data is the complete PDE trajectory We draw a random starting point We input the last K timesteps
input + data(t K:t)
for n E fl...., .,,....,N)....... do
input +- model(input)
end for
D Here we cut the gradients > Actual labels after the pushforward operation
target + data(t + NK: + N(K + 1))
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 activa- tion }. where fully connected layers are applied node-wise
2. Processor:
message passing layers as described in Sec. 3.2 Each layer consists of a 2-
layer edge update network @ following Equation 8 and a 2-layer node update network vi following Equation
3. Decoder: ID convolutional network with shard weights across spatial locations + {ID CNN layer >> activation > 1D CNN layer
We optimize models using the AdamW optimizer (Loshchilov & Hutter 2017) with learning rate le-4, weight decay le-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 experiments E1-E3 we use a hidden size of 164, and for experiments WE1-WE3 we use a hidden size of 128. order to enforce zero-stability during training we unroll the solver for maximum of 2 steps (see Sec. 3.1).
Message and update network in the processor consist - (fully-connected layer + activation > fully-connected layer activation J. 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 (loffe & Szegedy 2015) for the 2D experiments. For the decoder, use 8 channels between the two CNN layers (1 input channel, output channel) across all experiments. We use Swish (Ramachandran et al.) 2017) activation functions for experiments E1-E3 and WE1- WE3 and ReLU activation for the 2D experiments. ReLU activation proved most effective for 2D experiments since characteristic 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 million parameters and train- ing for the different experiments takes between 12 and 24 hours average GeForceRTXX 2080
22

---

Published as a conference paper at ICLR 2022
Ti GPU. 6 message passing layers and hidden size of 128 for the 2-layer edge update network and the 2-layer node update network is robust choice. A hidden size of 64 shows signs underfitting, whereas hidden size of 256 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 in- puts the edge update network & following Equation (8) and the node update network & following Equation (9. We ablate these choices in Appendix G
Another interesting hyperparameter is the number of neighbors used for message passing. We con- struct our graphs by restricting the neighbors (edges) via cutoff radius based on positional coor- dinates for experiments E1-E3 and the 2D experiments. We effectively use neighbors for exper- iments E1-E3 and 8 neighbors for the 2D experiments. For experiments WE1-W WE1-WEE, cutoff radii for selecting neighbors are not robust choice since the grids are irregular and relative distances are much lower close to the boundaries. We therefore construct our graphs via k--N 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 ID rollouts at different resolution for the Burgers' equation with differ- diffusion terms. Lower diffusion coefficients result in faster shock formation. Figure TO displays exemplary ID rollouts for different parameter sets. Large parameters result in fast and large shock formations. The wiggles arising due to the dispersive term y and cannot be captured by numerical solvers low resolution. Our MP-PDE solver is able to capture these wiggles and reproduce them even at very low resolution.
Exemplary 1D rollout of Burgers' equation 0.14) different resolutions
t=
1-0.36
t 6s
1. 1.05
t= 3.
ubx.t)
t= 6s
t= 9.
1= 2s
t= 65
-1:
nx=200 ground
nx=100 predictio
50=50 prediction
nx=10 predi
1-
t= 2:
tm 5s
3.9.
Exemplary 1D rollout of Burgers' equation 0.08) at different resolutions
1-0.0t
tw0. 35
1-0
tw1 0s
u(x,t)
tw 66
t=1
tw2 25
no- 200 ground TUtb
nx=100 prediction
nx=50 prediction
nx-40 predi ction
tw2 95
tw3 3.55
8
Figure 9: Exemplary ID rollout of the Burgers' equation at different resolutions. The different colors represent PDE solutions at different timepoints. Diffusion coefficients of v) = 0.14 (top) and 0.08 (bottom) are compared for the same initial conditions. Lower diffusion coefficients result in faster shock formation.
23

---

Published as a conference paper at ICLR 2022
Exemplary 1D rollout unseen equation (a 0.05, 8 0.14, y 0.55)
t-0.06
t-0.2s
t=0.3
t=555
u(k:t)
t- t=0.8c
t=1.15
ground
prediction
prediction
prediction
t-1.46
1=1.8
t=1.99
Exemplary 1D rollout equation 2.90. 8 0.20, y= 0.15)
1.5
1m0.26
0. 5s
t= 6s
u0x,t)
0.5
0. 8s
t= 0.
0.0
(=1.1 1s
t= 3.
200 ground truth
100 num baseline
hxx 50 num baseline
Tx= 40 num baseline
t= 4.
t= 6s
t- 8s
1=1.5 1.9s
Exemplary 1D rollout of an unseen equation (a 2.90, 8 0,20, 0.15)
tw0. 25
1.0
tw 5s
1-0
u(x())
tw 85
t-
tw1 15
t-
-0.5
Tx= 200 ground
nx= 100 prediction
1x- 50 prodiction
Tix 40 prediction
t-
tw1 46
tw 85
Figure 10: Exemplary 1D rollout an unseen equation with different equation parameters. The differ- ent colors represent PDE solutions at different timepoints. Low a parameters (top) result in diffusion like behavior. Large parameters (middle, bottom) result in fast and large shock formations. The wiggles arising due to the dispersive term 7. Numerical solvers cannot capture the wiggles low resolution (middle), MP-PDE solvers reconstruct them much better (bottom).
24

---

Published as a conference paper at ICLR 2022
F.2 EXPERIMENTS WEI, WE2, WE3
Figure I displays exemplary 2D rollouts at different resolutions for the wave equation with Dirich- let 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.
rollout the wave equation with Dirichlet boundary conditions
100
80
(seconds)
0.0
0.5
20
nx 100
nx-50 prediction
nx-40 prediction
20
predic tion
prediction
-10
2D rollout the equation with Neumann boundary conditions
100
(seconds)
1.0
0.5
20
nx=100 prediction
nx-50 prediction
nx-40 prediction
nk- 20
prediction
0.0
Figure 11: Exemplary 2D rollouts for the equation with Dirichlet boundary conditions (top) and Neumann boundary conditions (hottom). The different colors for the Dirichlet boundary condi- tion comes from the fact that wave propagation changes the sign each boundary.
Ground
Ground
Ground
Ground
20
20
20
20
10
>
10
10
>
10
40
40
10
10
40
60
20
40
60
40
60
20
40
60
Prediction
Prediction 108s
10
10
10
10
20
20
20
Figure 12: Exemplary 2D smoke inflow simulation. Ground truth data (top) compared to MP- PDE solvers (bottom). Simulations run for 100 timesteps corresponding 100 seconds. The MP- PDE solver is able capture the smoke inflow acccurately over the given time period.
25

---

Published as a conference paper at ICLR 2022
G ARCHITECTURE ABLATION AND COMPARISON TO CNNS
A schematic sketch of our MP-PDE solver is displayed in Figure 1 (sketch taken from the main pa- per). 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 archi- tecture representationally contains classical methods, such as FDM, FVM, and WENO schemes. The architectures follows the Encode-Process Deccess-Decode framework of Battaglia et al. (2018), with adjustments. Most notably, PDE coefficients and other attributes such boundary conditions de- noted with PPDE are included in the processor. For the decoder, a shallow 2-layer ID convolutional network with shared weights across spatial locations is applied, motivated by linear multistep meth- ods [Butcher 1987).
For ablating the architecture, three design choices are verified:
1. Does GNN have the representational power of a vanilla convolutional network on reg- ular grid? We test against 1D and 2D baseline CNN architectures.
2. For the decoder part, how much does ID convolutional network with shared weights across spatial locations approve upon a standard MLP decoder?
3. How much does the inclusion of PDE coefficients OPDE help to generalize over different equations or different boundary conditions?
Encoder
Processor
Decoder Node-wise shallow
Node-wise napping to the hidden space
Message passing
convolution
Irput 1ime
outget tine
CHN
NLP
MLP
Figure 13: Schematic sketch of our MP-PDE Solver, sketch taken from the main paper.
Table shows the three ablation (MP-PDE-0-6s MP-PDE.ID-ENN, Baseline CNN) tested on shock wave formation modeling and generalization to unseen experiments (experiments E1, E2, and E3), as described in Section 4. 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-QPOR ablation results are the same as reported in the main paper. The effect gets more prominent if more equation specific parameters available (OPDE features), as it is the case for experiment E3. We also refer the reader to the experiments presented in Table [ where MP- PDE solvers are shown to 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-ID-ENT ablation replaces the shallow 2-layer ID convolutional network with in the decoder with 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 the ID CNN network.
Baseline ID-CNN is built up of 8 ID-CNN layers, where the input consists of the spatial resolu- tion (n.)). The K previous timesteps used for temporal bundling are treated K input channels. The output consequently predicts the next K timesteps for the same spatial resolution (K output chan- nels). Using this format, again both temporal bundling and the pushforward trick can be effectively applied. The implemented ID-CNN layers are:
Input layer with K input channel, 40 output channels, kernel of size
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.
output layer with 40 input channels, K output channel, kernel of size 7.
26

---

Published as a conference paper at ICLR 2022
Residual connections used between the layers, and ELU (Clevert et al. 2016) linerrirssss applied. Circular padding is implemented to reflect the periodic boundary conditions. The CNN output is new vector d, (d. d.. ..... d with each element d corresponding
different point in time. Analogously to the MP-PDE solver, we use the output to update the solution as
i k+L u + (tx+e fx)d
1 ise5 E K
(61)
where is the output (and input) dimension. This baseline ID-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 res- olution (n,,) 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 input channel, 16 output channels, kernel.
4 intermediate layers with 16 input channels, 16 output channels, 5 x 5 kernl.
output layer with 16 imput channels, output channel, kernel.
Residual connections are used between the layers, and ELU (Clevert et al. 2016) non- con-lincarities are applied. For the spatial dimension, circular padding implemented reflect the periodic boundary conditions, for the temporal dimension zero padding is used. The CNN output is a new vector d; = (d!, d? .... d with each element d corresponding to different point time, Analogously to the MP-PDE solver and analogously to the ID-CNN, we use the output to update the solution as shown in Equation (6T).
The ID-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 ID-CNN across different PDEs becomes harder.
Table Ablation study comparing MP-PDE results on experiments on shock wave formation mod- eling and generaliziation to unseen equations (E1, E3, and E3) to an MP-PDE-0roE ablation, an MP-PDE-ID-ENT ablation, baseline ID-CNN and a baseline 2D-CNN architecture. Runtimes for one full unrolling over 250 timesteps on GeForce RTX 2080 Ti GPU. Accumulated error is 2,1 MSE.
1
Accumulated Error
Rantime [s] J
(it, Rx)
MP-PDE PDE
MP PR-PPE--ORR Orni
MP PDE. LD-ENN
ID-CNN
2D-CNN
MP-PDE
ID-CNN
2D-CNN
El
(250, 100)
1.55
2.41
3.45
25.70
0.09
0.02
0.16
E1
(250,50) 50)
1.67
2.69
3.88
32.42
0.08
0.02
0.15
E1
(250, 40)
1.47
2.50
3.07
37.13
0.008
0.02
0.14
(250, 100)
1.58
0.02
E2
(250, 50)
1.63
1.71
2.31
2.89
30.87
0.08
0.02
0.15
(250, 40)
1.45
E3
(250, 100)
4.26
4.71
6.26
9.15
42,37
0.09
0.02
0.16
E3
(250, 50)
3.74
10.90
5.15
45.41
0.09
0.02
0.15
E3
(250,40)
3,70
7.78
7.27
6.77
53.87
0.09
0.02
0.15
27


Reviews:
[{"review_id": "_7Habv975rW", "reviewer": "Reviewer_6VsU", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The ideas of the paper are interesting: the pushforward trick is appealing. The architecture, however, looks like a separate idea, which is not connected to the loss. I.e., the paper proposes two separate innovations, puts one into the title, but experiments confirm the effectiveness of only one. But I like the comparison with the standard numerical solver. However, the study of the architecture lacks the analysis of the systematic convergence of the solution, given by the new architecture, with the number of parameters. Classical solvers do have this property: increase the grid size, get better accuracy. Does the proposed model help with this respect? \n\nComments:\n- There is no ablation study (the loss or the architecture?)\n- The dependence of the properties of the method on the number of parameters in the model is not studied (at least, comments are welcome).\n", "summary_of_the_paper": "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.\n\n", "summary_of_the_review": "Reasonable idea, which proposes a loss that improves stability; however, it has two disconnected contributions and is not so easy to read. Some of the important points are not discussed/analyzed experimentally.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ZElPulyKksg", "reviewer": "Reviewer_wxXT", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "##########################################################################\n\nSummary:\n\nThe 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.\nThe 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.\nThe 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.\n\n\n##########################################################################\n\nReasons for score:\n\n \n\nIn my opinion, the paper brings significant advances to the field of neural PDE solvers. The proposed push forward trick, together with the temporal bundling trick, are valuable contributions as they impose favorably condition on autoregressive solvers (zero stability), improve the accuracy of other existing approaches as [Li et al 2020 a] and the inference time of the solver.\nThe approach is well evaluated experimentally,  showing accuracy and generalization results that are above the state of the art of neural solvers and beat classical solver schemes in inference time.\nFor these reasons I consider the manuscript to be accepted at ICLR 2022.\n\n\n \n\n##########################################################################Pr\n\n Pros:\n\n1. The paper tackles the important task of solving many families of PDEs, obtaining accuracy and generalization results upon the state-of-the-art of neural solvers.\n\n2. The approach is well motivated as classical solver schemes (FDM, FVM, WENO) are instances of message passing, as the authors themselves stated.\n\n3. The pushforward trick is by itself a valuable contribution as it enhances the performance of other existing neural solvers ([Li et al 2020 a]) as the authors show in Table 2.\n\n4. The paper is clear and well written, providing sufficient background for people not experts in the field.\n\n\n##########################################################################\n\nCons and questions:\n\n1. It is not entirely clear to me how the network can generalize to multiple spatial resolutions: more specifically how the encoder and the message passing mechanism handle a varying number of nodes (i.e. cells) in the graph? The number of cells stays fixed but the input positions vary? If this is the case, please clarify this in the experimental section.\n\n\n2. In the experimental section it would be good to provide more details regarding the time needed for training the model and the number of parameters of the network (also in comparison with [Li et al 2020 a])\n\n\n3. The paper should fit within the 9-page limit. I suggest that the background section could be shortened a bit, moving some of its parts to the Appendix.\n\n\n##########################################################################\n\nQuestions during rebuttal period:\n\n \n\nPlease address and clarify the cons above\n\n \n\n#########################################################################\n\nI spotted some typos:\n\n(1) Figure 1(right): reduces -> reduce\n\n(2) Page 5 Decoder section: an temporally -> a temporally\n\n(3) Page 9 Experiments and result section: MP-PDE solvers obtains -> MP-PDE solvers obtain\n\n(4) Page 9 section 4.4:  acccurately -> accurately\n\n(5) Page 10 section 7: many diciplines-> many disciplines\n\n(6) Page 10 section 7: paving a way -> pave the way\n\n", "summary_of_the_paper": "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.\nThe 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.\nThe 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.\n", "summary_of_the_review": "In my opinion, the paper brings significant advances to the field of neural PDE solvers. The proposed push forward trick, together with the temporal bundling trick, are valuable contributions as they impose favorably condition on autoregressive solvers (zero stability), improve the accuracy of other existing approaches as [Li et al 2020 a] and the inference time of the solver.\nThe approach is well evaluated experimentally,  showing accuracy and generalization results that are above the state of the art of neural solvers and beat classical solver schemes in inference time.\nFor these reasons I consider the manuscript to be accepted at ICLR 2022.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "GJoCUx0EBKS", "reviewer": "Reviewer_5QFf", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think the paper has a clarity issues and still in its early stages from that perspective. \n\nIn particular, the METHOD section is very vague and not clearly explained. Notation are introduced without definition and no explanation for the background needed to understand the notation. I suggest you lead with an example and motivate from there instead of going the general route.\n\nAlso, it is not clear what is the difference between using GNN and other regular networks in solving PDEs--the authors did not mention that in the paper.\n ", "summary_of_the_paper": "The paper utilizes neural message passing to propose another neural solver for PDEs. \n ", "summary_of_the_review": "The paper has a clarity issues and many notations are not introduced or defined.\n\n--Edit-- The authors either explained or addressed the issues in the paper with other reviewers. I am happy to recommend the paper for publication.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-tf2fm62z-f", "reviewer": "Reviewer_mh3h", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is very well written with: (1) thorough and detailed literature review; (2) intuitive explanation of motivation; (3) clear methodology and approach used in designing the architecture; (4) extensive details in empirical study and design of experiment. Overall, it is a good paper  with strong motivation, clear objective and solid empirical results. It is worthwhile to mention that authors design the training approach, i.e. pushforward trick, to reduce the instability due to error propagation and allow the network to solve versatile tasks. \n", "summary_of_the_paper": "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. ", "summary_of_the_review": "Overall, it is a good paper with a complete analysis of the design of architecture, training/testing approach and extensive experiments. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a message passing neural network to solve PDEs. The paper has sound motivation, clear methodology, and extensive empirical study. However, on the other hand, some reviewers also raised their concerns, especially regarding the lack of clear notations and sufficient discussions on the difference between the proposed method and previous works. Furthermore, there is no ablation study and the generalization to multiple spatial resolution is not clearly explained. The authors did a very good job  during the rebuttal period: many concerns/doubts/questions from the reviewers were successfully addressed and additional experiments have been performed to support the authors' answers. As a result, several reviewers decided to raise their scores, and the overall assessment on the paper turned to be quite positive.
Decision: Accept (Spotlight)