Predicting Physics in Mesh-reduced Space with Temporal Attention

1 INTRODUCTION . There has been an increasing interest in many scientific disciplines , from computational fluid dynamics [ 3 , 39 ] over graphics [ 42 , 40 ] to quantum mechanics [ 20 , 1 ] , to accelerate numerical simulation using learned models . In particular , methods based on Graph Neural Networks ( GNN ) have shown to be powerful and flexible . These methods can directly work with unstructured simulation meshes , simulate systems with complex domain boundaries , and adaptively allocate computation to the spatial regions where it is needed [ 7 , 38 , 34 , 51 ] . Most models for complex physics prediction tasks , in particular those on unstructured meshes , are next-step prediction models ; that is , they predict the next state u ( t+1 , x ) of a physical system from the current state u ( t , x ) . As next-step models suffer from error accumulation , mitigation strategies such as training noise [ 38 ] or augmented training data using a solver-in-the-loop [ 41 ] have to be used to keep rollouts stable . These remedies are not without drawbacks– training noise can be hard to tune , and ultimately place a bound on the achievable model accuracy . And worse , next-step models also tend to show drift , which is not as easily mitigated . Failure examples include failure to conserve volume or energy , shift in phase , or loss of shape information ( see e.g. , the failure case example in [ 38 ] ) . On the other hand , auto-regressive sequence models such as Recurrent neural networks ( RNNs ) , or more recently transformers , have been hugely successful in predicting sequences in NLP and image applications [ 35 , 30 ] . They can capture stochastic dynamics and work with partial observations [ 49 ] . Furthermore , their long attention span allows them to better preserve phase and conserved quantities [ 17 ] . However , as memory cost for full-sequence transformer models scales with both ∗Equal contribution . sequence length and spatial extents , it is hard to directly extend such models to predict physical systems defined on large unstructured meshes . This paper combines powerful GNNs and a transformer to model high-dimensional physical systems on meshes . The key idea is creating a locally coarsened yet more expressive graph , to limit memory consumption of the sequence model , and allow effective training . We first use a GNN to aggregate local information of solution fields of a dynamic system into node representations by performing several rounds of message passing , and then coarsen the output graph to a small set of pivotal nodes . The pivotal nodes ’ representations form a latent that encodes the system state in a low-dimensional space . We apply a transformer model on this latent , with attention over the whole sequence , and predict the latent for the next step . We then use a second GNN to recover the full-sized graph by up-sampling and message passing . This procedure of solving on a coarse scale , upsampling , and performing updates on a fine-scale is related to a V-cycle in multigrid methods [ 4 ] . We show that this approach can outperform the state-of-the-art MeshGraphNets [ 34 ] baseline on accuracy over a set of challenging fluid dynamics tasks . We obtain stable rollouts without the need to inject training noise , and unlike the baseline , do not observe strong error accumulation or drift in the phase of vortex shedding . 2 RELATED WORK . Developing and running simulations of complex , high-dimensional systems can be very timeintensive . Particular for computational fluid dynamics , there is considerable interest in using neural networks for accelerating simulations of aerodynamics [ 3 ] or turbulent flows [ 21 , 45 ] . Fast predictions and differentiable learned models can be useful for tasks from airfoil design [ 39 ] over weather prediction [ 36 ] to visualizations for graphics [ 42 , 40 ] . While many of these methods use convolutional networks on 2D grids , recently GNN-based learned simulators , in particular , have shown to be a very flexible approach , which can model a wide range of systems , from articulated dynamics [ 37 ] to dynamics of particle systems [ 23 , 38 ] . In particular , GNNs naturally enable simulating on meshes with irregular or adaptive resolution [ 34 , 7 ] . These methods are either steady-state or next-step prediction models . There are , however , strong benefits of training on the whole predicted sequence : next-step models tend to drift and accumulate error , while sequence models can use their long temporal range to detect drift and propagate gradients through time to prevent error accumulation . Models based on RNNs have been successfully applied to 1D time series and small n-body systems [ 6 , 53 ] or small 2D systems [ 49 ] . More recently , transformer models [ 43 ] , which have been hugely successful for NLP tasks [ 35 ] , are being applied to low-dimensional physics prediction tasks [ 17 ] . However , applying sequence models to predict high-dimensional systems remains a challenge due to their high memory overhead . Dimensionality reduction techniques , such as CNN autoencoders [ 33 , 32 , 26 , 22 , 29 , 16 , 11 , 27 ] , POD [ 44 , 48 , 5 , 31 , 18 , 8 , 47 , 10 ] , or Koopman operators [ 24 , 9 , 14 ] can be used to construct a lowdimensional latent space . The auto-regressive sequence model then operates on these linear ( POD modes ) or nonlinear ( CNNs ) latents . However , these methods can not directly handle unstructured data with moving or varying-size meshes , and many of them do not consider parameter variations . For example , POD can not operate on state vectors with different lengths ( e.g. , variable mesh sizes data in Fig.2 bottom ) . On the other hand , CNN auto-encoders can only be applied to rectangular domains and uniform grids ; and rasterization of complex simulation domains to uniform grids are known to be inefficient and are linked to many drawbacks as discussed in [ 12 ] . In contrast , our method reduces the input space locally by aggregating information on a coarsened graph . This plays to the strength of GNNs to learn local universal rules and is very closely related to multi-level methods [ 50 , 15 ] and graph coarsening [ 13 , 52 , 2 ] . 3 METHODOLOGY . 3.1 PROBLEM DEFINITION AND OVERVIEW . We are interested in predicting spatiotemporal dynamics of complex physical systems ( e.g. , fluid dynamics ) , usually governed by a set of nonlinear , coupled , parameterized partial differential equations ( PDEs ) as shown in its general form , ∂u ( x , t ) ∂t = F [ u ( x , t ) ; µ ] , x , t ∈ Ω× [ 0 , Te ] , ( 1 ) where u ( x , t ) ∈ Rd represents the state variables ( e.g. , velocity , pressure , or temperature ) and F is a partial differential operator parameterized by µ . Given initial and boundary conditions , unique spatiotemporal solutions of the system can be obtained . One popular method of solving these PDE systems is the finite volume method ( FVM ) [ 28 ] , which discretizes the simulation domain Ω into an unstructured mesh consisting of N cells ( Ci : i = 1 , . . . , N ) . At time t , the discrete state field Yt = { ui , t : i = 1 , . . . , N } can thus be defined by the state value ui , t at each cell center . Traditionally , solving this discretized system involves sophisticated numerical integration over time and space . This process is often computationally expensive , making it infeasible for realtime predictions or applications requiring multiple model queries , e.g. , optimization and control . Therefore , our goal is to learn a simulator , which , given the initial state Y0 and system parameters µ , can rapidly produce a rollout trajectory of states Y1 ... YT . As mentioned above , solving these systems with high spatial resolution by traditional numerical solvers can be quite expensive . In particular , propagating local updates over the fine grid , such as pressure updates in incompressible flow , can require many solver iterations . Commonly used technique to improve the efficiency include multigrid methods [ 46 ] , which perform local updates on both the fine grid , as well as one or multiple coarsened grids with fewer nodes , to accelerate the propagation of information . One building block of multigrid methods is the V-cycle , which consists of down-sampling the fine to a coarser grid , performing a solver update , up-sampling back on the fine grid , and performing an update on the fine grid . Noting that GNNs excel at local updates while the attention mechanism over temporal sequences allows long-term dependency ( see remark A.2 ) , we devise an algorithm inspired by the multigrid V-cycle . For each time step , we use a GNN to locally summarize neighborhood information on our fine simulation mesh into pivotal nodes , which form a coarse mesh ( section 3.2 ) . We use temporal attention over the entire sequence in this lower-dimensional space ( section 3.3 ) , upsample back onto the fine simulation , and perform local updates using a mesh recovery GNN ( section 3.2 ) . This combination allows us to efficiently make use of long-term temporal attention for stable , high-fidelity dynamics predictions . 3.2 GRAPH REPRESENTATION AND MESH REDUCTION . Graph Representation . We construct a graph G = ( V , E ) to represent a snapshot of a dynamic system at time step t. Here each node i ∈ V corresponds the mesh cell Ci , so the graph size is |V| = N . The set of edges E are derived from neighboring relations of cells : if two cells Ci and Cj are neighbors , then two directional edges ( i , j ) and ( j , i ) are both in E . We fix this graph for all time steps . At each step t , each node i ∈ V uses the local state vector ui , t as its attribute . Therefore , ( G , ( Y0 , . . . , YT ) ) forms a temporal graph sequence . The goal of the learning model is to predict ( Y1 , . . . , YT ) given G and Y0 . Mesh Reduction . After representing the system as a graph , we use a GNN to summarize and extract a low-dimensional representation zt from Yt for each step t. In this part of discussion , we omit the subscript t for notational simplicity . We refer to the encoder as Graph Mesh Reducer ( GMR ) since its role is coarsening the mesh graph . GMR first selects a small set S ⊆ V of pivotal graph nodes and locally encodes the information of the entire graph into representations at these nodes . By operating on rich , summarized node representations , the dynamics of the entire system is well-approximated even on this coarser graph . There are a few considerations for selection of pivotal nodes , for example , the spatial spread and the centrality of nodes in the selection . We generally use uniform sampling to select S from V . This effectively preserves the density of graph nodes over the simulation domain , i.e . pivotal nodes are more concentrated in important regions of the simulation domain . More details and visualizations on the node selection process can be found in section A.6 . GMR is implemented as a Encode-Process-Decode ( EPD ) GraphNet [ 38 ] . GMR first extracts node features and edge features from the system state using the node and edge Multi-Layer Perceptrons ( MLPs ) . v0i = mlpv ( Y [ i ] ) , e 0 ij = mlpe ( p ( i ) − p ( j ) ) . ( 2 ) Here Y [ i ] is the i-th row of Y , i.e . the state vector at each node , and p ( i ) is the spatial position of cell Ci . Then GMR uses L GraphNet processer blocks [ 37 ] to further refine node representations through message passing . In this process a node can receive and aggregate information from all neighboring nodes within graph distance L. Each processor updates the node and edge representations as e ` ij = mlp e ` ( e ` −1ij , v ` −1 i , v ` −1 j ) , v ` i = mlp v ` v ` −1i , ∑ j∈Ni e ` −1ij , ` = 1 , . . . , L. ( 3 ) Here v0i , e 0 ij are the outputs of Equation 2 , andNi denotes all neighbors of node i . The two functions mlpe ` ( · ) and mlpv ` ( · ) respectively concatenate their arguments as vectors and then apply MLPs . The calculation in equation 2 and equation 3 computes a new set of node representations V = ( vLi : i ∈ V ) for the graph G. Finally GMR applies an MLP to representations of the pivotal nodes in S only to “ summarize ” the entire graph onto a coarse graph : hi = mlpr ( v L i ) , i ∈ S ( 4 ) We concatenate these vectors into a single latent z = concat ( hi : i ∈ S ) as reduced vector representation of the entire graph . We collectively denote these three computation steps as z = GMR ( G , Y ) . The latents z can be computed independently for each time step t and will be used as the representation for the attention-based simulator . Mesh Recovery To recover the system state from the coarse vector representation z , we define a Graph Mesh Up-Sampling ( GMUS ) network . The key step of GMUS is to restore information on the full graph from representations of pivotal nodes . This procedure is the inverse of operation of GMR . We first set the representations at pivotal nodes by splitting z , that is , ri = hi , i ∈ S. We then compute representations of non-pivotal nodes by spatial interpolation [ 2 ] : for a non-pivotal node j , we choose a set N ′j of k nearest pivotal nodes in terms of spatial distance and compute the its representations rj by rj = ∑ i∈N ′j wijhi∑ i∈N ′j wij , wij = 1 d ( j , i ) 2 ( 5 ) Here d ( j , i ) is the spatial distance between cells Cj and Ci . Then every node i ∈ V has a representation ri , and all nodes ’ representations are collectively denoted asR = ( ri : i ∈ V ) . Similar to GMR , GMUS applies EPD GraphNet to the initial node representation R to restore Y on the full graph G. We denote the chain of operations so far as Ŷ = GMUS ( G , z ) . Details about information flows in GMR and GMUS can be found in section A.1 . We train GMR and GMUS as an auto-encoder over all time steps and sequences . For each time step t , we compute Ŷt = GMUS ( G , GMR ( G , Yt ) ) and minimize the reconstruction loss Lgraph = T∑ n=1 ‖Yt − Ŷt‖22 . ( 6 ) With these two modules we can encode system states ( Y1 , . . . , YT ) to latents ( z1 , . . . , zT ) as a lowdimensional representation of system dynamics . In the next section , we train a transformer model to predict the sequences of latents .

Summary: This work proposes a new algorithm combining graph-neural-network (GNN) and auto-regressive sequence models for physics prediction problems. The authors first use GNNs to compress the physical graphs, then use transformers to predict the next steps of the compressed representations, and finally use GNNs to recover the graph representations from the predicted representations. Through empirical studies, the authors show that this method outperform the previous SOTA model (MeshGraphNet) significantly, especially in the long-rollout prediction scenarios. The authors further analyze the models to understand the success and find that the ability of the transformer model to replay the earlier sequences seems to be the critical for the better performance especially in the scenarios with oscillations.