M ESSAGE P ASSING N EURAL PDE S OLVERS
Anonymous authors
Paper under double-blind review
A BSTRACT
The numerical solution of partial differential equations (PDEs) is difﬁcult, 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 ﬁnite differences, ﬁnite 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 ﬂuid-like ﬂow 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.
I NTRODUCTION
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 (Leli ` evre &
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 ﬁeld of numerical methods is a splitter ﬁeld (Bartels, 2016), tending to build handcrafted solvers for each sub-problem, rather than a lumper ﬁeld, where a mentality of “one method to rule them all” reigns.
This tendency is commonly justiﬁed 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 ﬂexibility to satisfy all structural requirements of a typical PDE problem.
This design is motivated by the insight that some classical solvers (ﬁnite differences, ﬁnite 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 coefﬁcients can be input to the solver.
B ACKGROUND AND RELATED WORK
Here in Section 2.1 we brieﬂy outline deﬁnitions 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. P ARTIAL DIFFERENTIAL EQUATIONS
We focus on PDEs in one time dimension t = [0 , T ] and possibly multiple spatial dimensions x
= [ x 1 , x 2 , . .
, x D ] (cid:62) ∈ X .
These can be written down in the form
[Equation]
where u : [0 , T ] × X → R n is the solution , with initial condition u 0 ( 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 ∂ x u , ∂ xx u , . . . is shorthand for partial derivatives ∂ u /∂ x , ∂ 2 u /∂ x 2 , 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 B D [ u ] = u − b D for ﬁxed function b D and Neumann boundary conditions B N [ u ] = n (cid:62) ∂ x u − b N , 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 ﬁeld for having studied these (Bar-Sinai et al., 2019; Li et al., 2020a).
Conservation form PDEs are written as
[Equation]
where ∇· J is the divergence of J .
The quantity J : R n → R n is the ﬂux , which has the interpretation of a quantity that appears to ﬂow.
Consequently, u is a conserved quantity within a volume, only changing through the net ﬂux J ( u ) through its boundaries.
2. C LASSICAL SOLVERS
Grids and cells Numerical solvers partition X into a ﬁnite grid X
= { c i } Ni =1 of N small nonoverlapping volumes called cells c i ⊂ X .
In this work, we focus on grids of rectangular cells.
Each cell has a center at x i .
u ki is used to denote the discretized solution in cell c i and time t k .
There are two main ways to compute u ki : sampling u ki = u ( t k , x i ) and averaging u ki = (cid:82) c
i u ( t k , x ) d x .
In our notation, omitting an index implies that we use the entire slice , so u k = ( u k 1 , u k 2 , ..., u kN ) .
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 u k .
We then solve
[Equation]
where f is the form of F acting on the vectorized u t 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, ∂ x u k | x i might become ( u ki +1 − u ki ) /
( x i +1 − x i ) .
Principled ways to derive stencils can be found in Appendix A. FDM is simple and efﬁcient, 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 ﬂux into the cell.
In 1D, this leads to
[Equation]
where t k , respectively.
The problem thus boils down to estimating the ﬂux at cell boundaries x i ± 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 m th derivative is computed as IFFT { ( ιω ) m FFT ( 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.
2. N EURAL SOLVERS
We build on recent exciting developments in the ﬁeld 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 u 0 , a neural operator M : [0 , T ] × F → F , where F is a (possibly inﬁnite-dimensional) function space, is trained to satisfy
[Equation]
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.
Inﬁnite-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
[Equation]
where A : R > 0 × R N → R N is the temporal update.
In this work, since ∆ t is ﬁxed, 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.
Each paper focuses on a different class of PDE solver: ﬁnite volumes, multigrid, and iterative ﬁnite 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.
M ETHOD
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.
3. T RAINING FRAMEWORK
Autoregressive solvers A map solutions u k to causally consequent ones u k +1 .
A straightforward way to train this is with one-step training .
If p 0 ( u 0 ) is the distribution of initial conditions in the training set then let p k ( u k )
= (cid:82) p ( u k | u 0 )
p 0 ( u 0 )
d u 0 be the groundtruth solution distribution at iteration k .
We minimize
[Equation]
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
[Figure]
rollouts), and lead to divergence from the groundtruth.
This can be interpreted as overﬁtting 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 p k
(cid:55)→
A (cid:93) p k at iteration k + 1 , where A (cid:93) : 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 (cid:93) p k instead of the distribution p k +1 , and unfortunately A (cid:93) p k (cid:54) =
p k +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
[Equation]
, (9) where (cid:15) | u k is an adversarial perturbation sampled from an appropriate distribution.
For the perturbation distribution, we choose (cid:15) such that ( u k + (cid:15) ) ∼ A (cid:93) p k .
This can be easily achieved by using ( u k + (cid:15) )
=
A ( u k − 1 ) for u k − 1 one step causally preceding u k .
Our total loss is then
L one-step + L stability .
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 ﬁrst 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 (cid:15) 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 (cid:107)A ( u 0 + (cid:15) ) − u 1 (cid:107) < κ (cid:107) (cid:15) (cid:107) , 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 u 0 (cid:55)→ u 1 ; whereas, we predict K steps u 0 (cid:55)→ ( u 1 , u 2 , ..., u K ) =
u 1: 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.
3. A RCHITECTURE
We model the grid X as a graph G =
( V , E ) with nodes i ∈ V , edges ij ∈ E , and node features f i ∈ R c .
The nodes represent grid cells
c i and the edges deﬁne local neighborhoods.
Modeling
[Figure]
the domain as a graph offers ﬂexibility 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.
and Sanchez-Gonzalez et al.
(2020) , with adjustments.
We are not the ﬁrst 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 u k − K : k i , node position x i , current time t k , and equation embedding θ PDE to node embedding vector f 0 i = (cid:15) v ([ u k − K : k
i , x i , t k , θ PDE ]) .
θ PDE contains the PDE coefﬁcients 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 G 1 , G 2 , ..., G M .
The speciﬁc updates we use are edge j → i message:
[Equation]
node [Equation] 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 x j − x i can be justiﬁed by the translational symmetry of the PDEs we consider.
Solution differences u
i − u j 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 x i .
For each node i , the processor outputs a vector f Mi .
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 efﬁcient 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 d i =
( d 1 i , d 2 i , ..., d Ki ) with each element d k i corresponding to a different point in time.
We use this to update the solution as
[Equation] [Equation]
The motivation for this choice of decoder has to do with a property called consistency (Arnold, 2015), which states that lim ∆ t → 0 (cid:107)A (∆ t, u 0 ) − u (∆ t ) (cid:107)
= 0 .
That is, the prediction matches the exact solution in the inﬁnitesimal time limit.
Consistency is a requirement for zero-stability of the rollouts.
[Figure]
MLP MLP
MPNN MPNN
CNN
CNN CNN
[Figure]
Connections.
As mentioned in Bar-Sinai et al.
(2019), both FDM and FVM are linear methods, which estimate n th -order point-wise function derivatives as
[Equation]
for appropriately chosen coefﬁcients α ( n ) j , where N ( 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 ﬂux 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, andlayers 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.
E XPERIMENTS
We demonstrate the effectiveness of the MP-PDE solver on tasks of varying difﬁculty 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.
4. I NTERPOLATING BETWEEN PDE S
Data We focus on the family of PDEs
[Equation]
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 coefﬁcients sampled uniformly in A j ∈ [ − 0 . 5 , 0 . 5] , ω j ∈ [ − 0 .
4 , − 0 .
4] , (cid:96) j ∈ { 1 , 2 , 3 } , φ j ∈ [0 , 2 π ) .
This setup guarantees periodicity of the initial conditions and forcing.
Space is uniformly discretized to n x = 200 cells in [0 , 16) with periodic boundary and time is uniformly discretized to n t = 200 points in [0 , 4] .
Our training sets consist of 2096 trajectories, downsampled to resolutions ( n t , n x ) ∈ { (250 , 100) , (250 , 50) , (250 , 40) } .
Numerical groundtruth is generated using a 5 th -order WENO scheme (WENO5) (Shu, 2003) for the convection term ∂ x u 2 and 4 th -order ﬁnite 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 .
≤ α ≤ 3 . 0 , 0 . 0 ≤ β ≤ 0 .
4 and 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- (cid:8)(cid:8)(cid:8) θ 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 (T OP ) shows shock formation at different resolutions ( E1 ), a traditionally difﬁcult phenomenon to model—FDM and
PSM methods cannot model shocks.
Strikingly, shocks are preserved even at very low resolution.
4. V ALIDATING TEMPORAL BUNDLING AND THE PUSHFORWARD METHOD
We observe solver survival times on E1 , deﬁned as the time until the solution diverges from groundtruth.
A solution ˆ u ki diverges from the groundtruth u ki when its max-normalized L 1 -error 1 n
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
|
[Figure]
[Equation]
| u
[Equation]
| exceeds 0 .
The solvers are unrolled to n t = 1000 timesteps with T =s. Examples are shown in Figure 4 (B OTTOM ), where we 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 of SanchezGonzalez et al.
(2020) or Figure S9 of Bar-Sinai et al.
In a second experiment, we compare the efﬁcacy 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
[Figure]
[Table]
[Figure]
we show a model trained with Gaussian noise adversarial perturbations, similar to that proposed in Sanchez-Gonzalez et al.
We see that applying the pushforward trick leads to far higher survival times, conﬁrming 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.
4. S OLVING ON IRREGULAR GRIDS WITH DIFFERENT BOUNDARY CONDITIONS
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
[Equation]
(16) where c is wave velocity ( c = 2 in our experiments).
We consider Dirichlet B [ u ] =
u =and Neumann B [ u ] =
∂ x u = 0 boundary conditions.
This PDE is 2 nd -order in time, but can be rewritten as 1 st -order in time, by introducing the auxilliary variable v = ∂
t u and writing ∂ t [ u, v ] − [ v, c 2 ∂ xx u ] = 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 ( n t , n x ) =
(250 , 200) on a Chebyshev extremal point grid (cell edges are located at x i = cos( iπ/ ( n x + 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 (cid:8)(cid:8)(cid:8) θ PDE ) have no equation speciﬁc 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]
4. 2D EXPERIMENTS
For 2D experiments we use data from P HI F LOW 1 , an open-source ﬂuid simulation toolkit.
We look at ﬂuid simulation based on the Navier-Stokes equations, and simulate smoke inﬂow 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 inﬂow 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 inﬂow accurately over the given time period, suggesting scalability of MP-PDE solver to higher dimensions.
C ONCLUSION
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 ﬂexibility 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.
For the latter, a trained MP-PDE solver can be ﬁtted 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 ﬁrst 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 ﬁx 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 ﬁeld of research, which is especially appealing for building neural PDE solvers since every PDE is deﬁned via a unique set of symmetries (Olver, 1986).
R EPRODUCIBILITY S TATEMENT
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 ﬁelds and therefore have included a detailed description of those in our appendix.
For reproducibility, we plan to provide our code upon publication.
E THICAL S TATEMENT
The societal impact of MP-PDE solvers is difﬁcult 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.
A I NTERPOLATION
To compute classical numerical derivatives of a function u : R → R , which has been sampled on a mesh of points x 1 < x 2 < ... <
x N it is common to ﬁrst ﬁt a piecewise polynomial p : R → R on the mesh.
We assume we have function evaluations
u i = u
( x i ) at the mesh nodes for all i = 1 , ..., N and once we have ﬁtted the polynomial, we can use its derivatives at any new off-mesh point, for instance the half nodes x i + 12 .
Here we illustrate how to carry out this procedure.
A polynomial of degree N − 1 (note the highest degree polynomial we can ﬁt to N points has degree
N − 1 ) can be ﬁtted at the mesh points by solving the following linear system in a
[Equation]
To ﬁnd the polynomial interpolation at a new point x , we then do
[Equation]
where x (cid:62) =
[1 , x, x 2 , ..., x N − 1 ] .
To retrieve the m th derivative, where m < N − 1 , is also very simple.
For this we have that
[Equation]
19) where x ( m ) (cid:62) =
[(0) m x 0 − m , (1) m x
1 − m , ..., ( N − 1) m x N − 1 − m ] (20) where ( i ) m
=
i ( i − 1)( i − 2) · · · ( i − m + 1) is the m th Pochhammer symbol and ( i ) 0 :
= 1 .
Note that is typical to fold s = x (cid:62) X − 1 into a single object, which we call a stencil .
B R ECONSTRUCTION
Reconstruction is the task of ﬁtting a piecewise polynomial on a mesh of points, when instead of functions values at the grid points we have cell averages ¯ u i , where
[Equation]
where each cell I i =
[ x i − 12 , x i + 12 ] , the half nodes are deﬁned as x i + 12 :=
12 ( x i + x i +1 )
and ∆ x i = x i + 12 − x i − 12 .
The solution is to note that we can ﬁt a polynomial P ( x ) to the integral of u ( x ) , which will be exact at the half nodes, so
[Equation]
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
[Equation]
where ¯ u is the vector of cell averages and L is a lower triangular matrix performing the last sum in
Equation 22.
Thus the m th derivative of the polynomial p ( x ) = P (cid:48)
( x ) is
[Equation] [Equation]
C WENO S CHEME
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 ∈ [ x i , x i +1 ] .
We design N stencils to ﬁt the function on shifted overlapping intervals
I 1 , ..., I N , where I k
=
[ x i − N +
+ k , x i + k ] .
In the case of
WENO reconstruction these intervals are I k
=
[ x i − N +
+ k − 12 , x i + k + 12 ] .
If a discontinuity lies in
I = (cid:83) k
I k , then it is likely
that one of the substencils s ( m ) k
or ¯ s ( m ) k
(deﬁned on interval I k ) 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)
[Equation]
Here the nonlinear weights w k satisfy (cid:80) Nk
=1 w k
=
It was shown in Jiang & Shu (1996) that these nonlinear weights can be constructed as
[Equation]
[Equation]
where γ k is called the linear weight , (cid:15) is a small number, and β k is the smoothness indicator .
The linear weights are set such that the sum (cid:80) Nk
=
γ k ¯ s k over the order
N − 1 stencils matches a larger order 2 N − 2 stencil.
The smoothness indicators are computed as
[Equation]
where p k is the polynomial corresponding to substencil ¯ s k .
C. W ENO 5 SCHEME
A conservative ﬁnite difference spatial discretization approximates a derivative f ( u ) x by a conservative difference
[Equation]
[Equation]
where ˆ f i +
and ˆ f i − 12 are numerical ﬂuxes.
Since g ( u ) y is approximated in the same way, ﬁnite difference methods have the same format for more than one spatial dimensions.
The leftreconstructed (reconstruction is done from left to right) ﬁfth order ﬁnite difference WENO scheme (WENO5) (Shu, 2003) has the ˆ u − i + 12 given by:
[Equation]
In equation 29, ˆ u − ( j ) i + 12 are the left WENO reconstructions on three different stencils given by
[Equation] [Equation]
[Equation]
(32) and the non-linear weights w j given by
[Equation]
with γ { 1 , 2 , 3 } = { 110 , 35 , 310 } , and (cid:15) a tiny-valued parameter to avoid the denominator becoming 0.
The left smoothness indicators β − k are given by:
[Equation]
The right-reconstructed WENO5 scheme has the ˆ u + i − 12 given similarly to ˆ u − i
+
but with all coefﬁcients ﬂipped since the reconstruction is done from the other side (from right to left).
The right reconstruction on three different stencils are given by
[Equation]
and the right smoothness indicators β + k are given by
[Equation]
Both ˆ u − and ˆ u + are needed for full ﬂux reconstruction as explained in the next section.
C. F LUX RECONSTRUCTION
We consider ﬂux reconstruction via Godunov.
For Godunov ﬂux, ˆ f i + 12 =
ˆ f ( u i + 12 ) is reconstructed from ˆ u
+ i
+
and ˆ u − i + 12 via:
[Equation]
C. C OMPARING WENO SCHEME TO ANALYTICAL SOLUTIONS
First analytical case.
An analytical solvable case for Burgers equation arises for the boundary conditions u
[Equation]
and the initial conditions
[Equation]
where
[Equation]
The analytical solutions for this speciﬁc set of boundary and initial conditions gives
[Equation]
where
[Equation]
For this ﬁrst analytical solveable case, the analytical solution, the WENO scheme and the fourth order ﬁnite difference scheme (FDM) are compared in Fig.
6 for a diffusion term of ν = 0 .
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:
[Equation]
and the initial condition
[Equation] [Equation]
Solutions are (Basdevant et al., 1986)
[Equation]
with f ( y ) =
exp( − cos( πy 2 πν ) .
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 .
The WENO scheme models the analytical solution perfectly.
Modeling via the FDM scheme fails completely.
6 x       u ( m / s )
Analytical solution, 1D rollout t=0.0s t=0.1s t=0.1s t=0.2s t=0.2s t=0.3s t=0.4s t=0.4s t=0.5s   6 x 0.
t ( s e c o n d s )
Analytical solution, 2D rollout          6 x       u ( m / s )
WENO solution, 1D rollout t=0.0s t=0.1s t=0.1s t=0.2s t=0.2s t=0.3s t=0.4s t=0.4s t=0.5s   6 x 0.
t ( s e c o n d s )
WENO solution, 2D rollout          6 x      u ( m / s )
FDM solution, 1D rollout t=0.0s t=0.1s t=0.1s t=0.2s t=0.2s t=0.3s t=0.4s t=0.4s t=0.5s   6 x 0.
t ( s e c o n d s )
FDM solution, 2D rollout
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.
[Figure]
D P SEUDOSPECTRAL METHODS FOR WAVE PROPAGATION ON IRREGULAR GRIDS
We consider Dirichlet B [ u ] = u
= 0 and Neumann B [ u ] = ∂ x u = 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 fulﬁll 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.
[Table]
[Table] [Table]
8 x 0.
[Table] [Table]
conditions        8 0.
8 x 0.
E E XPLICIT R UNGE -K UTTA METHODS
The family of Runge-Kutta methods (Butcher, 1987) is given by:
[Equation] [Equation]
where
[Equation] [Equation]
For a particular Runge-Kutta method one needs to provide the number of stages s , and the coefﬁcients a ij (1 ≤ j < i ≤ s ) , b i ( i = 1 , 2 , . . .
, s ) and c i ( i = 1 , 2 , . . . , s ) .
These data are usually arranged in so-called Butcher tableaux (Butcher, 1963).
F E XPERIMENTS
Flux terms of equations that we study—the Heat, Burgers, Korteweg-de-Vries (KdV), and
Kuromoto-Shivashinsky (KS) equation—are summarized in Table 3.
[Table]
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, perform N forward passes, and ﬁnally 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 (cid:46) data is the complete PDE trajectory t ← DrawRandomNumber t ∈ { 1 , ..., T } (cid:46)
We draw a random starting point input ← data( t − K : t ) (cid:46)
We input the last K timesteps for n ∈ { 1 , . . .
, N } do input ← model(input) end for (cid:46)
Here we cut the gradients target ← data( t + NK : t + N ( K + 1) )
(cid:46) Actual labels after the pushforward operation output ← model(input) loss ← criterion(output, target)
Implementation details.
MP-PDE architectures, consist of three parts (sequentially applied):
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 2layer edge update network φ following Equation (10), and a 2-layer node update network ψ following Equation (11).
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 usechannels 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 inﬂow 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
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 underﬁtting, whereas a hidden size of 200 is not improve performance signiﬁcantly. 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. E XPERIMENTS E1, E2, E
Figure 9 displays exemplary 1D rollouts at different resolution for the Burgers’ equation with different diffusion terms.
Lower diffusion coefﬁcients 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.
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Figure]
1.1s t=1.3s t=1.4s t=1.6s t=1.8s t=1.9s
[Figure]
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.
[Figure]
G A RCHITECTURE ABLATION AND COMPARISON TO CNN S
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 ﬂexibility 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 coefﬁcients 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 veriﬁed:
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.
For the decoder part, how much does a 1D convolutional network with shared weights across spatial locations approve upon a standard MLP decoder?
How much does the inclusion of PDE coefﬁcients θ PDE help to generalize over e
g. different equations or different boundary conditions?
[Figure]
Table 4 shows the three ablation (MP-PDE- (cid:8)(cid:8)(cid:8) θ PDE , MP-PDE- (cid:40)(cid:40)(cid:40)(cid:40) 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- (cid:8)(cid:8)(cid:8) θ PDE ablation results are the same as reported in the main paper.
The effect gets more prominent if more equation speciﬁc 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- (cid:40)(cid:40)(cid:40)(cid:40) 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 ( n x ).
The K previous timesteps used for temporal bundling are treated as K 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 al
, 2016) non-linearities are applied.
Circular padding is implemented to reﬂect the periodic boundary conditions.
The CNN output is a new vector d i =
( d 1 i , d 2 i , ..., d Ki ) with each element d ki corresponding to a different point in time.
Analogously to the MP-PDE solver, we use the output to update the solution as
[Equation]
) 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 ( n x ) 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 reﬂect the periodic boundary conditions, for the temporal dimension zero padding is used.
The CNN output is a new vector d i =
( d 1 i , d 2 i , ..., d Ki ) with each element d ki 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.
[Table]
[Table]
H G ENERALIZATION OVER MULTIPLE SPATIAL RESOLUTIONS
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 ( n t , n x ) =
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.
[Table]