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 paper proposes to use auto-regressive sequence models with attention mechanisms to predict the evolution of physical simulations. Mesh-based discretizations are targeted, and hence the employed networks take the form of GNNs. The paper focuses on encode-process-decode structures, and a transformer is proposed as main method to predict future states in the latent space. This approach is evaluated with three CFD settings: an unsteady wake flow, a sonic backwards facing step, and a medical flow scenario. The results are additionally analyzed in terms of latent space and attention weight content.