Conditionally Parameterized, Discretization-Aware Neural Networks for Mesh-Based Modeling of Physical Systems

1 Introduction . Numerical simulations of partial differential equations ( PDEs ) have become an indispensable tool in the study of complex physical systems . High-resolution simulations are , however , prohibitively expensive or intractable in many practical problems . Machine learning techniques have recently been explored to improve the efficiency and accuracy of traditional numerical methods . Successful applications include nonlinear model order reduction [ 1 , 2 , 3 ] , model augmentation [ 4 , 5 , 6 ] , and super-resolution [ 7 , 8 , 9 ] . Neural networks have also been used to replace traditional PDE-based solvers , and serve as a standalone prediction tool . Popular approaches include auto-regressive timeseries predictions [ 10 , 11 , 12 , 13 , 14 ] , Physics-Informed Neural Networks ( PINNs ) [ 15 , 16 , 17 ] . Despite promising results on canonical problems , commonly used network architectures such as autoencoders and CNNs have inherent limitations . An autoencoder generates a fixed mapping between the geometric coordinates and the encoded digits . This limits their portability for new geometries and dynamic patterns . A CNN requires interpolation of existing data to a structured , Euclidean space , introducing additional cost and error . Irregular geometry boundaries require constructs such as elliptic coordinate transformation [ 18 ] and Signed Distance Function ( SDF ) [ 1 ] . Moreover , models often ignore the hierarchical relations between heterogeneous features , and concatenate them into a single input vector , e.g . the common concatenation of the edge and node features in Graph 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Neural Networks ( GNNs ) . The learning of high-order terms remains mostly unguided – even simple quadratic terms are often fitted via a number of hidden units in a brute-force manner . With a focus on mesh-based modeling of physical systems , we use the idea of conditional parameterization ( CP ) to build the hierarchical relations between different physical quantities as well as numerical discretization information into the network architectures . The key contributions of our work are as follows 1 : 1 . We demonstrate that a drop-in CP modification can bring significant improvements for various existing models on several tasks essential to the modeling of physical systems . 2 . We propose a conditionally parameterized graph neural network ( CP-GNet ) , which effectively models complex physics such as chemical source terms , irregular mesh discretizations , and different types of boundary conditions . 3 . We conduct extensive numerical tests and demonstrate state-of-the-art performances on problems of different complexities , ranging from the basic viscous Burgers equation to a complex reacting flow . 2 Methodology . Conditional Parametrization : The idea of conditional parametrization ( CP ) is to use trainable functions of input parameters to generate the weights of a neural network . To demonstrate this , we start from a standard dense ( fully connected ) layer : h ( u ; W , b ) = σ ( Wu + b ) , ( 1 ) where u ∈ Rnx is the input feature vector , h ∈ Rnh is the output hidden state vector , W ∈ Rnh×nx and b ∈ Rnh are the trainable weights and bias , and σ is the activation function . It can be seen that in the evaluation stage , the values of W and b are fixed regardless of the inputs . Thus the performance of Eq . ( 1 ) is largely limited by the interpolation range of training data . By introducing a parameter vector p ∈ Rnp and a trainable function f ( p ) : Rnp → Rnh×nx that computes the weights W based on p , the conditionally parameterized version of Eq . ( 1 ) is given by : h ( u ; f ( p ) , b ) = σ ( f ( p ) u + b ) . ( 2 ) An easy way to incorporate the formulation into existing neural network models is by making f a single-layer MLP , the conditionally parameterized dense ( CP-Dense ) layer can be represented by : h ( u , p ; W , B , b ) = σ ( σ ( 〈W , p〉+ B ) u + b ) . ( 3 ) It should be noted that this would bring a change in the dimensions of weights and biases , which become W ∈ R ( nh×nu ) ×np , B ∈ Rnh×nu . When the layer width is kept the same , the total number of trainable parameters increases linearly with the parameter size np . In applications , p is not limited to an additionally-introduced parameter . When simply taking u as the parameter for itself , the quadratic terms will be introduced . High-order terms , which are prevalent in physical systems , can be easily modeled using multiple such layers . In Appendix B , we demonstrate how certain discretized PDE terms can be fitted exactly with simple conditionally parameterized layers . 2.1 CP-GNet for mesh-based modeling of physical systems . Graph representation of discretized systems : Consider a physical system governed by a set PDEs for a time-variant vector of variables q ( t ) . Using the popular finite volume discretization , the computational domain is divided into contiguous small cells , indexed by i . The discretized form of equation can be written as : dqi ( t ) dt = 1 Ωi ∑ j∈N ( i ) f ( qi , qj , nij ) Aij + s ( qi ) , ( 4 ) 1The source code is released to facilitate future research at https : //github.com/ davidxujiayang/cpnets where qi is the cell-centered value of cell i , Ωi is the volume ( 3D ) /area ( 2D ) of the cell , and N ( i ) is the neighborhood set of cells around i . Between a neighboring pair of cells i and j , Aij is area ( 3D ) /length ( 2D ) of the shared cell boundary , and nij = ( xi − xj ) /|xi − xj | is a vector between the cell center locations xi and xj . In the explicit numerical simulation of Eq . ( 2.1 ) , solutions are updated by computing the increment of ∆qki = q k+1 i − qki between discrete time steps indexed by k , which is determined by two terms . The flux term f computes the exchange of quantity between neighboring cells , which is a complex function involving both the cell values as well as the vector between them , e.g . [ 19 ] . The source term s computes physics that are local to the cell , such as the reaction of chemical species . In our setting , the discretized system is mapped to a graph G ( V , E ) , defined by nodes V of size |V | = nv connected by edges E ⊂ V × V of size |E| = ne . Each node vi is located at the corresponding cell center xi , and each edge ( i , j ) corresponds to a shared boundary between the finite volume cells . Denoting the sets of mapped quantities on all nodes and edges of G , Q = { qi , i ∈ V } , N = { nij , ( i , j ) ∈ E } , the target is to develop a graph neural network operator g that predicts the increment as ∆Qk = g ( Qk , N ) . CP-GNet architecture : The architecture for the proposed conditionally parameterized graph neural network , CP-GNet , can be written in an encoder-processor-decoder form . A schematic is provided in Fig . 1 . For clarity of different variables in the description of the network , we use ui for the latent variables on node i to distinguish from the physical variables qi , and use eij for latent variables on edge ( i , j ) to distinguish from the vector nij . Encoder : Numerical solution of PDEs ( e.g. , the compressible Navier–Stokes equations ) requires the processing of arbitrarily complex interactions between mesh elements . While large MLP architectures can represent this complexity , the data requirements to reliably train such networks might be large . In contrast , our proposed encoder takes two CP-Dense layers , taking the output from the previous layer as both the input and the conditional parameter . Through the encoder , high-order interactions can be easily extracted , allowing a degree of extrapolation by virtue of linearity . The CP-GNet uses two separate , but similarly constructed encoders to process the input node features qki and edge features nij , respectively . Processor : The flux term f in Eq . ( 2.1 ) can be effectively approximated by CP message-passing ( CP-MP ) between adjacent nodes on a graph . The source term s , on the other hand , can be modeled by CP-Dense layers . In CP-GNet the processor consists of multiple identical blocks with independent weights . Residual connections are added between the blocks . As shown in Fig . 1 , each block includes a CP-MP based section and a CP-Dense based section to address the two types of terms . Modified from the Edge Conditioned Convolution ( ECC ) [ 20 ] , the CP-MP computation is formulated as : Wij = σ ( 〈 W , eφij 〉 + B ) , ; hi = ∑ j∈N ( i ) wijσ ( 〈Wij , [ ui ; uj ] 〉 ) , ( 5 ) where eφij is the output from the edge encoder , ui and uj are the latent node features from the previous layer , hi is the nodal latent output , and wij = Aij/Ωi is the flux weight from Eq . ( 2.1 ) . Decoder : It is common in a PDE solver to use a Jacobian matrix J = ∂u/∂h to transform the variable increments ∆u = J∆h . The decoder in the GP-GNet serves a similar purpose – to convert hidden variables to the output on the physical space . Similar to the encoder , the decoder consists of three conditionally parameterized dense layers . The first two layers can actually be viewed as a dedicated “ encoder ” that is similar to the initial node encoder , taking qki as the input , but with independent weights . The purpose of this “ encoder ” is to extract a final conditional parameter , which is used in the third CP-Dense layer in the decoder to determine the weights for the output node feature uχi from the processor . The third decoder layer is also the final layer of the model , which outputs ∆qki ( with proper scaling ) . Except for the edge encoder and the last two layers in the decoder , all dense , CP-Dense , CP-MP layers are appended with LayerNormalization ( LN ) layers . Treatments for boundaries : The computational domain of a practical problem includes multiple types of boundaries , e.g . the case in Sec . 4.3 . In classic PDE solvers , they are treated with different boundary conditions , which define explicit formulations to compute relationships of the domain with the external world . However , these conditions and formulations are only defined for the physical quantities , thus can not be easily transferred for latent variables , especially when multiple messagepassing/convolution steps are used . To enable the GP-GNet to model different types of boundaries efficiently , special treatments are necessary . For boundaries with known inputs , such as the inlet and the outlet , their values are directly input to the corresponding nodes at every time step . For the boundaries imposing certain constraints , instead of a given physical value , such as Neumann and symmetry boundaries , ghost edges are introduced . For a cell i with a face lying on a boundary , we introduce a ghost edge vector nig , that points from the corresponding node i to the center of the boundary face . Ghost edges are processed together with the normal edges in the edge encoder . However , the CP-MP layer in the processor of the CP-GNet is slightly modified . More specifically , the concatenation [ ui ; uj ] in Eq . ( 5 ) is replaced with only ui . And for each type of boundary , the weights for the CP-MP layer are trained independently , to let the model learn different types of boundary condition for the latent variables . The effectiveness of this treatment is shown in Sec . 4.3 and further discussed in Appendix . A.3.3 .

Summary: The authors of the paper are adding a  parametrized (by input data) layer to perform a renormalisation step from a small scale PDE to larger scale system. This requires a definition of a closure term that modifies the equation at the larger scale to match the behaviour at the a more accurate, smaller scale. This is shown to significantly reduce the required computing resources, while still retaining high accuracy.  More, the system is defined on graph convolutional network that is applicable to complex mesh structures that are required for proper description of many real world problems.  The authors introduce a simplified form of the attention mechanism that is missing the value vectors and softmax portions, but gives a possibility to introduce a degree of  non-linearity to the activation, besides the activation function. The expectations is that this will help the network to reach accuracy with fewer computations. This is  validated in the  CFD cases considered in the paper.