Physics-Embedded Neural Networks: E( n ) -Equivariant Graph Neural PDE Solvers
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Graph neural network (GNN) is a promising approach to learning and predictingphysical phenomena described in boundary value problems, such as partial differential equations (PDEs) with boundary conditions.
However, existing modelsinadequately treat boundary conditions essential for the reliable prediction of suchproblems.
In addition, because of the locally connected nature of GNNs, it isdifﬁcult to accurately predict the state after a long time, where interaction betweenvertices tends to be global.
We present our approach termed physics-embeddedneural networks that considers boundary conditions and predicts the state aftera long time using an implicit method.
It is built based on an E( n ) -equivariant
GNN, resulting in high generalization performance on various shapes.
We demonstrate that our model learns ﬂow phenomena in complex shapes and outperformsa well-optimized classical solver and a state-of-the-art machine learning modelin speed-accuracy trade-off.
Therefore, our model can be a useful standard forrealizing reliable, fast, and accurate GNN-based PDE solvers.
Introduction
Partial differential equations (PDEs) are of interest to many scientists because of their application invarious ﬁelds such as mathematics, physics, and engineering.
Numerical analysis is used to solve
PDEs because most PDE problems in real life cannot be solved analytically.
For example, predictingﬂuid behavior in complex shapes is an essential topic because it is helpful for product design, disasterreduction, weather forecasting and many others; however, it is a difﬁcult problem and takes time tosolve using classical solvers.
Machine learning is a promising approach to predicting such phenomenabecause it can utilize data similar to the state to be predicted, while classical solvers cannot utilizesuch information.
However, the main challenge in dealing with complex phenomena such as ﬂuids is to guaranteegeneralization performance because possible states in complex systems can be huge and may not becovered using a purely data-driven approach.
Therefore, we must apply appropriate inductive biasesto machine learning models.
Many approaches successfully introduced various inductive biasessuch as local connectedness using graph neural networks (GNNs) and symmetry under coordinatetransformations using equivariance.
While these methods have made great progress in solving PDEs using machine learning, there is stillroom for improvement.
First, there is need for an efﬁcient and provable way to respect boundaryconditions like Dirichlet and Neumann.
Rigorous fulﬁllment of Dirichlet boundary conditions isindispensable because they are hard constraints and different Dirichlet conditions correspond todifferent problems users would like to solve.
Second, there is need to reinforce the treatment of globalinteraction to predict the state after a long time, where interactions tend to be global.
GNNs have E( n ) [Figure] Boundary
[Figure]
Pseudoinverse
Decoder
Encoded Dirichlet boundary condition [Figure] nonlinear solver excellent generalization properties because of their locally-connected nature; however, they may missglobal interaction due to their localness. We propose physics-embedded neural networks (PENNs), a machine learning framework to addressthese issues by embedding physics in the models.
We build our model based on IsoGCN (Horie et al.,2021), a lightweight E( n ) -equivariant GNN to reﬂect physical symmetry and realize fast prediction.
Furthermore, we construct a method to consider Dirichlet and Neumann boundary conditions.
Finally,we reconsider a way to stack GNNs based on an iterative method used for nonlinear optimizationproblems, which naturally introduces the global pooling to GNNs as the global interaction with highinterpretability.
In experiments, we demonstrate that our treatment of Neumann boundary conditionsimproves the predictive performance of the model, and our method can fulﬁll Dirichlet boundaryconditions with no error.
Our method also achieves state-of-the-art performance compared to aclassical, well-optimized numerical solver and a baseline machine learning model -accuracy
Our main contributions are summarized as follows:
• We construct four models to satisfy Dirichlet and Neumann boundary conditions: theboundary encoder , Dirichlet layer , pseudoinverse decoder , and NeumannIsoGCN .
Theconsidered models show provable fulﬁllment of boundary conditions, while existing modelscannot.•
We propose neural nonlinear solvers , which realize global connections to stably predict thestate after a long time.
• We demonstrate that the proposed model shows state-of-the-art performance in speedaccuracy trade-off.
Figure 1 shows the overview of the proposed model.
Background and related work
In this section, we review the foundation of PDEs to clarify the problems we solve and introducerelated works where machine learning models are used to solve PDEs.
2.1 Partial differential equations (PDEs) with boundary conditions
A general form of the d -dimensional temporal PDEs that we consider can be expressed as follows:
[Equation]
where ⌦ is the domain, @ ⌦ is the boundary of ⌦ , and @ ⌦ Dirichlet and @ ⌦ Neumann are boundarieswith Dirichlet and Neumann boundary conditions.
ˆ · is a known function, and D is a known nonlineardifferential operator, which can be nonlinear and contain spatial differential operators (see equation 18for an example of D ).
n ( x ) is the normal vector at x 2 @ ⌦ .
Equation 3 is called the Dirichletboundary condition, where the value on @ ⌦ Dirichlet is set as a constraint.
Equation 4 corresponds tothe Neumann boundary condition, where the value of the derivative u in the direction of n is set on@ ⌦ Neumann rather than the value of u .
When u : (0 , T ) ⇥ ⌦ !
R f satisﬁes Equations 1 – 4, it iscalled the solution of the (initial-) boundary value problem.
2.1.1 Discretization
PDEs are deﬁned in a continuous space to make differentials meaningful.
Discretization can beapplied in space and time so that computers can solve PDEs easily.
In numerical analysis of complexshaped domains, we commonly use meshes (discretized representations of shapes), which can beregarded as graphs.
We denote the position of the i th vertex as x i and the value of a function f at thex i as f i .
The simplest method to discretize time is the explicit Euler method formulated as:
[Equation]
which updates u ( t, x i ) with a small increment D ( u )( t, x i )  t .
Another way to have time discretization is the implicit Euler method formulated as: [Equation]
[Equation]
which solves Equation 6 rather than simply updating variables to ensure the original PDE is satisﬁednumerically.
The equation can be viewed as a nonlinear optimization problem by formulating it as:
[Equation]
where R ( v ) is the residual vector of the discretized PDE.
The solution of Equation 8 correspondsto the state after  t , u ( t +  t, x ) .
By letting r  = R for an appropriate  , solving Equation 8corresponds to optimizing  in an ( f ⇥ n ) -dimensional space, where n is the number of vertices inthe considered mesh.
A simple way to solve such an optimization problem is by applying gradientdescent formulated as:
[Equation]
where v (0) =
u ( t ) and ↵ ( i ) 2 R is determined using line search.
However, due to the highcomputational cost of the search, ↵ can be ﬁxed to a small value, which corresponds to the explicit Euler method with the time step size ↵  t .
Barzilai & Borwein (1988) suggested another simpleyet effective way to determine the step size using a two-point approximation to the secant equationunderlying quasi-Newton methods.
2.2 Neural PDE solvers
We review machine learning models used to solve PDEs called neural PDE solvers, typically formulated as u ( t n +1 , x i ) ⇡
F NN ( u )( t n , x i ) for ( t n , x i ) 2 { t 0 , t 1 , . . . } ⇥ ⌦ , where F NN is a machine 95 learning model.
2.2.1 Physics-informed neural networks (PINNs)
Raissi et al.
(2019) made a pioneering work combining PDE information and neural networks, called
PINNs, by adding loss to monitor how much the output satisﬁes the equations.
PINNs can be usedto solve forward and inverse problems and extract physical states from measurements (Pang et al.
,2019; Mao et al., 2020;
Cai et al., 2021).
However, PINNs’ outputs should be functions of spacebecause PINNs rely on automatic differentiation to obtain loss regarding PDEs.
This design constraintsigniﬁcantly limits the model’s generalization ability because the solution of a PDE could be entirelydifferent when the shape of the domain or boundary condition changes.
Therefore, these methods arenot applicable in building models that are generalizable over shape and boundary condition variations.
2.2.2 Graph neural network based PDE solvers
As discussed in Section 2.1.1, one can regard a mesh as a graph.
GNNs can take any graphs asinputs (Gori et al., 2005;
Scarselli et al., 2008; Kipf & Welling, 2017; Gilmer et al., 2017), having thepossibility to generalize over various graphs, i.e., meshes.
Therefore, GNNs are strong candidatesfor learning mesh-structured numerical analysis data, as seen in Alet et al.
; Chang & Cheng(2020); Pfaff et al.
Brandstetter et al.
(2022) advanced these works by suggesting temporalbundling and pushforward trick for efﬁcient and stable prediction.
Their method could also considerboundary conditions by feeding them to the models as inputs.
Here, one could expect the modelto learn to satisfy boundary conditions approximately, while there is no guarantee to fulﬁll hardconstraints such as Dirichlet conditions.
2.2.3 Equivariant models
In addition to GNNs, another essential concept to help machine learning models generalize isequivariance.
Equivariance is characterized by using group action as f ( g · x ) =
g · f ( x ) forf : X ! Y and g 2 G acting on X and Y .
In particular, E( n ) -equivariance is essential to predictthe solutions of physical PDEs because it describes rigid body motion, i.e., translation, rotation, andreﬂection.
Ling et al.
(2016) and Wang et al.
introduced equivariance to a simple neuralnetwork and CNN to predict ﬂow phenomena.
Both works showed that equivariance improvedpredictive and generalization performance compared to models without equivariance.
Horie et al.(2021) proposed E( n ) -equivariant GNNs based on GCNs (Kipf & Welling, 2017), called IsoGCNs.
A form of their model is formulated as:
[Equation]
where N i is the neighborhood of the i th vertex and ⌦ is the tensor product operator.
Here, we denote 126 IsoGCN [-
-]0 !
1 an IsoGCN layer that converts the input scalar (rank-0 tensor) ﬁeld to the outputvector (rank-1 tensor) ﬁeld.
This layer corresponds to the gradient operator, which helps learn PDEsbecause spatial derivatives such as gradient play an essential role in PDEs.
They applied the model tothe heat equation problem, showing high predictive performance and fast prediction, while boundarycondition treatment was out of their scope.
Proposed method
In this section, we present our model architecture.
We adopt an encode-process-decode architecture,proposed by Battaglia et al.
(2018), which has been applied successfully in various previous works,e.g., Horie et al. ; Brandstetter et al.
Our key concept is to encode input features,including information on boundary conditions, apply a GNN-based nonlinear solver loop reﬂectingboundary conditions in the encoded space, then decode carefully to satisfy boundary conditions inthe output space.
3.1 Dirichlet boundary model
As demonstrated theoretically and experimentally in literature (Hornik, 1991; Cybenko, 1992;
Nakkiran et al., 2021), the expressive power of neural networks comes from encoding in a higherdimensional space, where the corresponding boundary conditions are not trivial.
However, if there are no boundary condition treatments in layers inside the processor, which resides in the encodedspace, the trajectory of the solution can be far from the one with boundary conditions.
Therefore,boundary condition treatments in an encoded space are essential for obtaining reliable neural PDEsolvers that fulﬁll boundary conditions.
To ensure the same encoded space between variables and boundary conditions, we use the sameencoder for variables and the corresponding Dirichlet boundary conditions, which we term theboundary encoder , as follows: [Equation]
[Equation]
One can easily apply Dirichlet boundary conditions in the aforementioned encoded space using the
Dirichlet layer deﬁned as:
[Equation]
This process is necessary to return to the state respecting the boundary conditions after someoperations in the processor, which might disrespect the conditions.
After the processor layers, we decode the hidden features using functions satisfying:
[Equation]
This condition ensures that the encoded boundary conditions correspond to the ones in the originalphysical space.
Demanding that Equation 13 holds for arbitrary ˆ u ; we obtain f decode  f encode = Id u ,resulting in f decode = f
+encode , which we call the pseudoinverse decoder .
It is pseudoinverse becausef encode , in particular encoding in a higher-dimensional space, may not be invertible.
Therefore, weconstruct f +encode using pseudoinverse matrices and inverse activation functions.
For more details,see Appendix A.
3.2 Neumann boundary model
Matsunaga et al.
(2020) proposed a wall boundary model to deal with Neumann boundary conditionsfor the least squares moving particle semi-implicit (LSMPS) method (Tamai & Koshizuka, 2014), aframework to solve PDEs using particles.
The LSMPS method is the origin of the IsoGCN’s gradientoperator, so one can imagine that the wall boundary model may introduce a sophisticated treatmentof Neumann boundary conditions into IsoGCN.
We modiﬁed the wall boundary model to adapt tothe situation where the vertices are on the Neumann boundary, which differs from the situation ofparticle simulations. For more details of the model’s derivation, see Appendix A.
Our formulation of IsoGCN with Neumann boundary conditions, which is termed NeumannIsoGCN , is expressed as:
[Equation]
where ˆ g i is the value of the Neumann boundary condition at x i , and w i > 0 is an untrainableparameter to control the strength of the Neumann constraint.
As w i ! 1 , the model strictlysatisﬁes the given Neumann condition. However, as the directional derivatives can be computedin various directions, i.e., ( x j  x i ) , we keep the value of w i moderate to respect information inthe neighborhood.
In particular, we set w i = 10 . 0 , assuming that around ten vertices may virtuallyexist "outside" the boundary on a ﬂat surface in a 3D space.
NeumannIsoGCN is a straightforwardgeneralization of the original IsoGCN by letting n i = [-
 whenx-]0 when x i / 2 @ ⌦ Neumann .
This model canalso be generalized to vectors or higher rank tensors, similarly to the original IsoGCN’s construction.
Therefore, NeumannIsoGCN can express any spatial differential operator, constituting D in PDEs.
3.3 Neural nonlinear solver
As reviewed in Section 2.1, one can regard solving PDEs as optimization.
Here, we adopt the
Barzilai–Borwein method (Barzilai & Borwein, 1988) to solve Equation 8 in the encoded space. In [Table]
our case, the step size ↵ ( i ) of gradient descent is approximated as:
[Equation]
wherer R ( h ) is the residual vector in the encoded space and h f , g i ⌦ :=
P [Equation]
g
[Equation] x i ) 183 denotes the inner product over the mesh.
Because the inner product is taken all over the mesh (graph),computing ↵ ( i )BB corresponds to global pooling.
With that view, one can ﬁnd similarities between
Equation 9 and deep sets (Zaheer et al., 2017), which is a successful method to learn point cloud dataand has a strong background regarding permutation equivariance.
Our aim is to use Equation 16, approximating the nonlinear differential operator D in Equation 7with NeumannIsoGCNs.
By doing this, we expect the processor to consider both local and globalinformation, which may have an advantage over simply stacking GNNs corresponding to the explicitmethod as discussed in Section 2.1.1.
Combinations of solvers and neural networks are alreadysuggested in, e.g., NeuralODE
(Chen et al., 2018).
The novelty of our study is the extension ofexisting methods for solving PDEs with spatial structure and the incorporation of global pooling intothe solver, enabling us to capture global interaction, which we refer to as the neural nonlinear solver .
Finally, the update step in the neural nonlinear solver is expressed as:
[Equation]
where h (0) is the output of the encoder and D NeumannIsoGCN is an E( n ) -equivariant GNN reﬂectingthe structure of D using differential operators provided by NeumannIsoGCN.
Here, Equation 17enforces hidden features to satisfy the encoded PDE, including boundary conditions, motivating us tocall our model physics-embedded neural networks because it embeds physics (PDEs) in the modelrather than in the loss.
Experiments
Using numerical experiments, we demonstrate the proposed model’s validity, expressibility, andcomputational efﬁciency.
We use two types of datasets: 1) the gradient dataset to verify the correctness of NeumannIsoGCN and 2) the incompressible ﬂow dataset to demonstrate the speed andaccuracy of the model.
We also present ablation study results to corroborate the effectiveness of the [Table] proposed method.
The implementation of our model is based on the original IsoGCN’s code.
2 Ourimplementation is in the supplemental material, which will be online upon publication.
All the detailsof the experiments can be found in Appendix B and C.
4.1 Gradient dataset
As done in Horie et al.
(2021), we conducted experiments to predict the gradient ﬁeld from a givenscalar ﬁeld to verify the expressive power of NeumannIsoGCN.
We generated cuboid-shaped meshesrandomly with 10 to 20 cells in the X, Y, and Z directions.
We then generated random scalar ﬁeldsover these meshes using polynomials of degree 10 and computed their gradient ﬁelds analytically.
Our training, validation, and test datasets consisted of 100 samples.
Table 1 and Figure 2 show thatthe proposed NeumannIsoGCN improves gradient prediction, especially near the boundary, showingthat our model successfully considers Neumann boundary conditions.
4.2 Incompressible ﬂow dataset
We tested the expressive power our model by learning incompressible ﬂow in complex shapes.
Thecorresponding nonlinear differential operator is denoted as:
[Equation]
[Equation] with the incompressible condition r · u = 0 , where, in the present case, u is the ﬂow velocity ﬁeld,p is the pressure ﬁeld, and Re is the Reynolds number.
4.2.1 Data
To generate the dataset, we ﬁrst generated pseudo-2D shapes, with one cell in the Z direction, bychanging design parameters, starting from three template shapes.
Thereafter, we performed numericalanalysis using a classical solver, OpenFOAM, 3 with  t = 10  3 , and the initial conditions were thesolutions of potential ﬂow, which can be computed quickly and stably using the classical solver.
The
Reynolds number Re was around 10 3 .
Template shapes, design parameters, and boundary conditionsused can be found in Appendix C.
To conﬁrm the expressive power of the proposed model, we used coarse input meshes for machinelearning models.
We generated these coarse meshes by setting cell sizes roughly four times largerthan the original numerical analysis.
We obtained ground truth variables using interpolation.
Thetask was to predict ﬂow velocity and pressure ﬁelds at t = 4 .
using information available beforenumerical analysis, e.g., initial conditions and the geometries of the meshes.
Training, validation, andtest datasets consisted of 203, 25, and 25 samples, respectively.
We generated the dataset by randomlyrotating and translating test samples to monitor the generalization ability of machine learning models. [Figure]
PENN prediction is consistent under rotation and translation due to the E( n ) -equivariance nature of the model, while MP-PDE’s predictive performance degrades under transformations.
4.2.2 Machine learning models
We constructed the PENN model corresponding to the incompressible Navier–Stokes equation.
Inparticular, we adopted the fractional step method, where the pressure ﬁeld was also obtained as a
PDE solution along with the velocity ﬁeld.
For the state-of-the-art baseline model, we selected MP-PDE (Brandstetter et al., 2022) as it alsoprovides a way to deal with boundary conditions.
We used the authors’ code 4 with minimummodiﬁcation to adapt to the task.
We tested various time window sizes such as 2, 4, 10, and 20, whereone step corresponds to time step size  t = 0 .
With changes in time window size, we changedthe number of hops considered in one operation of the GNN of the baseline to have almost the samenumber of hops visible from the model when predicting the state at t = 4 . 0 .
Both models weretrained for up to 24 hours using one GPU (NVIDIA A100 for NVLink 40GiB HBM2).
4.2.3 Results
Table 2 and Figure 3 show the comparison between MP-PDE and PENN.
The predictive performancesof both models are at almost the same level when evaluated on the original test dataset.
The resultsshow the great expressive power of the MP-PDE model because we kept most settings at defaultas much as possible and applied no task-speciﬁc tuning.
However, when evaluating them on thetransformed dataset, the predictive performance of MP-PDE signiﬁcantly degrades.
Nevertheless,
PENN shows the same loss value up to the numerical error, conﬁrming E( n ) -equivariance.
Inaddition, PENN exhibits no error on the Dirichlet boundaries, showing that our treatment of Dirichletboundary conditions is rigorous.
Figure 4 shows the speed-accuracy trade-off for OpenFOAM, MP-PDE, and PENN.
We varied meshcell size and the time step size for OpenFOAM to have different computation speeds and accuracy.
[Table] [Table]
[Table] [Table]
[Table]
[Table] [Table]
[Figure]
The error bar represents the standard error of the mean.
All computation was done using one core of Intel Xeon CPU E5v2@2.40GHz.
The proposed model achieved the best performancein speed-accuracy trade-off between all the testedmethods under fair comparison conditions.
Table 3 presents the results of the ablation study.
Comparison between models with and without theproposed components shows that the proposed components, i.e., the boundary encoder, Dirichlet layer,pseudoinverse decoder, and neural nonlinear solver,signiﬁcantly improve the models.
The neural nonlinear solver in the encoded space turned out to have thebiggest impact on the performance, while the Dirichlet layer ensured reliable models that strictly respect
Dirichlet boundary conditions.
Conclusion
We have presented an E( n ) -equivariant, GNN-basedneural PDE solver, PENN, which can fulﬁll boundary conditions required for reliable predictions.
Themodel has superiority in embedding the information of PDEs (i.e., physics) in the model and speedaccuracy trade-off.
Therefore, our model can be auseful standard for realizing reliable, fast, and accurate GNN-based PDE solvers.
Although theproperty of our model is preferable, it also limits the applicable domain of the model.
For instance,the proposed model cannot exploit its potential to solve inverse problems where explicit forms of thegoverning PDE are not available for such tasks.
Therefore, combining PINNs and our model, PENNs,could be the next direction of the research community.
Potential negative societal impacts
We have built a foundation to learn PDEs in a steerable manner rather than focusing on a speciﬁcapplication.
Because of that, we envisage minimal risk of direct abusing our present work.
However,as mentioned in the introduction, solving PDEs has many impacts on various domains, from bothpositive and negative aspects.
Thus, our work and possible successive ones may be abused, aiming toharm lives and the environment.
Therefore, the research community, including us, must be careful inusing them and control the research direction to prevent abusing machine learning and simulationtechnologies.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contribu-
tions and scope?
[Yes] See the abstract and Section 1.
(b) Did you describe the limitations of your work?
[Yes] See Section 5.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See Section 6.(d)
Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes]
Our paper, in particular Section 6, is written in accordance with the guideline
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimentalresults (either in the supplemental material or as a URL)?
[Yes] They are included in thesupplemental material.
(b) Did you specify all the training details (e
g., data splits, hyperparameters, how they were chosen)?
[Yes] They are included in the supplemental material.
(c) Did you report error bars (e
g., with respect to the random seed after running experiments
multiple times)?
[Yes] See Section 4, Table 1, 2,
and Figure 4.
(d) Did you include the total amount of compute and the type of resources used (e g., type of GPUs,
internal cluster, or cloud provider)?
[Yes] See Section 4.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] See Footnote
(b) Did you mention the license of the assets?
[Yes] See Footnote
(c) Did you include any new assets either in the supplemental material or as a URL?
[-(19) where W + is the pseudoinverse matrix of W and σ is an invertible activation function whose Dom( σ ) = 
-]
[Yes] It is
included in the supplemental material
(d) Did you discuss whether and how consent was obtained from people whose data you’re us-
ing/curating?
[N/A]
The dataset we used in the paper is generated by ourselves.
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
[-For each design parameter, wevaried from 0 to 1.0 with a step size of 0.1, yielding 11 shapes for type A and 121 shapes for type B and C. 
-]
[-5 d , e − 1 . 
-]
[-d − 2 . 
 d-]infor-
[-The pressure ﬁeld at t = 4 . 
-]
mation or offensive content?
[N/A]
[-Our initial learning rate was 5 . 
-]
The dataset is based on numerical analysis, so there is no
concern about personally identiﬁable information nor offensive content.
If you used crowdsourcing or
conducted research with human subjects...
[Table]
[Table]
[-= D = 0 . 
-]
[-5 d , e − 1 . 
-]
[-d , e − 2 . 
-]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]