PREDICTING PHYSICS IN MESH-REDUCED SPACE WITH TEMPORAL ATTENTION
Year: 2022

Abstract:
Graph-based next-step prediction models have recently been very successful in modeling complex high-dimensional physical systems on irregular meshes. However, due to their short temporal attention span, these models suffer from error accumulation and drift. In this paper, we propose a new method that captures long-term dependencies through a transformer-style temporal attention model. We introduce an encoder-decoder structure to summarize features and create a compact mesh representation of the system state, to allow the temporal model to operate on a low-dimensional mesh representations in a memory efficient manner. Our method outperforms a competitive GNN baseline on several complex fluid dynamics prediction tasks, from sonic shocks to vascular flow. We demonstrate stable rollouts without the need for training noise and show perfectly phase-stable predictions even for very long sequences. More broadly, we believe our approach paves the way to bringing the benefits of attention-based sequence models to solving high-dimensional complex physics tasks.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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 cannot directly handle unstructured data with moving or varying-size meshes, and many of them do not consider parameter variations. For example, POD cannot 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].

heading: 3 METHODOLOGY

heading: 3.1 PROBLEM DEFINITION AND OVERVIEW
text: 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.

heading: 3.2 GRAPH REPRESENTATION AND MESH REDUCTION
text: 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.

heading: 3.3 ATTENTION-BASED SIMULATOR
text: As a dynamics model, we learn a simulator which can predict the sequence (z1, . . . ,zT ) autoregressively based on an initial state z0 = GMR(G,Y0) and the system parameters µ. The system parameters include conditions such as different Reynolds numbers, initial temperatures, or/and parameters of the spatial domain (see Table 3 for details). The model is conditioned on µ, to be able to predict flows with arbitrary system parameters at test time.We encode µ into a special parameter token z−1 of the same length as the state representation vectors zt
z−1 = mlpp(µ). (7) Then, the transformer model [43], trans(·) predicts the subsequent latent vectors in an autoregressive manner.
z̃1 = trans(z−1, z0), z̃t = trans(z−1, z0, z̃1, . . . , z̃t−1) (8) Specifically, we use a single layer of multi-head attention in our transformer model, which we found sufficiently powerful for the environments we studied. We first describe the computation of a single attention head: At step t, the prediction z̃t−1 from the previous step issues a query to compute attention a(t−1) over latents (z−1, z0, z̃1, . . . , z̃t−1).
a(t−1) = softmax
( z̃>t−1W
> 1 W2√ d′
· [ z−1, z0, z̃1, . . . , z̃t−1, ]) (9)
Here W1 and W2 are learnable parameters, and d′ denotes the length of the vector zt−1. This equation corresponds the popular inner-product attention model [43].
Next, the attention head computes the prediction vector gt = W3 [ z−1, z0, z̃1, . . . , z̃(t−1) ] at−1, (10) with the learned parameterW3.
Our multi-head attention model usesK parallel attention heads with separate parameters to compute K vectors (g1t , . . . , g K t ) as described above. These vectors are concatenated and fed into an MLP to predict the residual of z̃t over z̃t−1. z̃t = z̃t−1 + mlpm ( concat(g1t , . . . , g K t ) ) (11)
We train the autoregressive model on entire sequences by minimizing the prediction loss
Latt = T∑ t=1 ‖zt − z̃t‖22. (12)
Compared to next-step models such as MeshGraphNet [34], this model can propagate gradients through the entire sequence, and can use information from all previous steps to produce stable predictions with minimal drift and error accumulation.
The proposed model also contains next-step models as special cases: if we select all nodes as pivotal nodes and set the temporal attention model such that it predicts the next step with only physical parameters and the current step, then the proposed model becomes a next-step model. The proposed model shows advantage when it reduces the dimension of representation by using pivotal nodes and expands the input range when predicting the next step. A.2 shows how this method reduces the amount of computation and memory usage.

heading: 3.4 MODEL TRAINING AND TESTING
text: We train the mesh reduction and recovery modules GMR, GMUS by minimizing equation 6 and the temporal attention model by minimizing equation 12. We obtain best results by training those components separately, as this provides stable targets for the autoregressive predictive module. A second reason is memory and computation efficiency; we can train the temporal attention method on full sequences in reduced space, without needing to jointly train the more memory-intensive GMR/GMUS modules.
The transformer model is trained by incrementally including loss terms computed from different time steps. For every time step t, we train the transformer by minimizing the objective ∑t t′=1 ‖zt′− z̃t′‖22. Until the objective reaches a threshold, we add the next loss term ‖zt+1 − z̃t+1‖ to the objective and continue to train the model We found that this incremental approach leads to much more stable training than directly minimizing the overall objective.
As the transformer model can directly attend to all previous steps via equation 10, we don’t suffer as much from vanishing gradients, compared to e.g. LSTM sequence models. By predicting in meshreduced space, the memory footprint is limited, and the model can be trained without approximate gradient calculations such as teacher forcing or limiting the attention history. Details for training procedure can be found in section A.3.
Once we have trained our models, we can make predictions for a new problem. We first represent system parameters as a token z−1 by equation 7 and encode the initial state into a token z0 = GMR(G,Y0), then we predict the sequence z̃1, . . . , z̃T by equation 8, and finally we decode the system state Ỹ by
Ỹt = GMUS(z̃t), t = 1, . . . , T. (13)

heading: 4 RESULTS
text: Datasets. We tested the proposed method to three datasets of fluid dynamic systems: (1) flow over a cylinder with varying Reynolds number; (2) high-speed flow over a moving wedge with different initial temperatures; (3) flow in different vascular geometries. These datasets correspond to applications with fixed, moving, and varying unstructured mesh, respectively. A detailed description of the datasets can be found in section A.4.
Methods. The proposed methods is compared to the state-of-the-art MeshGraphNet method [34], which is a next-step model and has been shown to outperform a list of previous methods. Two versions of MeshGraphNet are considered here, with and without Noise Injection (NI). We also study variants of our model with LSTM and GRU (GMR-LSTM and GMR-GRU) instead of temporal attention. These variants also operate in mesh-reduced space, and use the same encoder (GMR) and decoder (GMUS) as our main method. Model details can be found in section A.5. And section A.6 shows pivotal nodes for each dataset.
Stable and visually accurate model predictions. We tested the proposed method on three fluid systems. Figure 2 compares our model’s predictions for velocity against the ground truth (CFD) on six test set trajectories. Our model shows stable predictions even without noise injection, and model rollouts are visually very closely to the CFD reference. On the other hand, without mitigation strategies, next-step models often struggle to predict such long sequences stably and accurately. For cylinder flow, our model can successfully learn the varying dynamics under different Reynolds numbers, and accurately predict their corresponding transition behaviors from laminar to vortex shedding. In supersonic flow over a moving wedge, our model can accurately capture the shock bouncing back and forth on a moving mesh. In particular, the small discontinuities at the leeward side of the edge are discernible in our model predictions. Lastly, in vascular flow with a variable-size circular thrombus, and thus different mesh sizes, the model is able to predict the flow accurately. In both cylinder flow and vascular flow, both frequency and phase of vortex shedding is accurately captured. The excellent agreement between model predictions and reference simulation demonstrates the capability of our model when predicting long sequences with varying-size meshes. Results for pressure can be found in section A.9.
Error behavior under long rollouts To quantitatively assess the performances of our models and baselines, we compute the relative mean square error (RMSE), defined as
RMSE(ûprediction, ûtruth) = ∑ (ûpredictioni −û truth i )
2∑ (ûpredictioni ) 2 , over the full rollout trajectory. Table 1 com-
pares average prediction errors on all three domains. Our attention-based model outperforms baselines and model variants in most scenarios. For sonic flow, which has 40 time steps, our model performs better than MeshGraphNet when predicting velocity values (u and v). The quantities p and T are relatively simple to predict, and all models show very low relative error rates (< 10−3), though our model has slightly worse performance. The examples with long rollout sequences (predicting cylinder flow and vascular flow) highlight the superior performance of our attention-based method compared with baseline next-step models, and the RMSE of our method is significantly lower than baselines.
Figure 3 shows how error accumulates in different models. Our model has higher error for the first few iterations, which can be explained by the information bottleneck of mesh reduction. However, the error only increases very slowly over time. MeshGraphNet on the other hand suffers from strong error accumulation, rendering it less accurate for long-span predictions. Noise injection partially addresses the issue, but error accumulation is still considerably higher compared to our model. We can attribute this better long-span error behavior of our sequence model, to being able to pass gradients to all previous steps, which is not possible in next-step models. We also note that the
transformer model performs much better compared to the GRU and LSTM model variants. We believe its ability to directly attend to long-term history helps to stabilize gradients in training, and also to identify recurrent features of the flow, as discussed later in this section.
In Figure 4 we illustrate the difference in the models’ abilities to predict long rollouts. For cylinder flow, our model is able to retain phase information. MeshGraphNet on the other hand drifts to an opposite vortex shedding phase at t = 250, though its prediction still looks plausible. For vascular flow, our model gives an accurate prediction at step t = 228, while the prediction of MeshGraphNet is significantly different from the ground-truth at this step: the left inflow diminishes, making the prediction less physically plausible.
Diagnosis of state latents. The low-dimensional latent vectors, computed by the encoder-decoder (GMR and GMUS), can be viewed as lossy compression of high-dimensional physical data on irregular unstructured meshes. In our experiment, we observe the compression ratio is ≤ 4% while the loss of the information measured by RMSE is below 1.5%. More details are in section A.7.
We also visualize state latents z-s from the cylinder flow dataset and show that the model can distinguish rollouts from different Reynolds numbers. We first use Principal Component Analysis (PCA) to reduce the dimension of z-s to two and plot them in Figure 5 (left). Rollouts with two different Reynolds numbers have two clearly distinctive trajectories. They eventually enters their respective periodic phase. We also apply PCA directly to the system states (Y) and get 2-d latent vectors directly from the original data. We plot these vectors in Figure 5 (right). Latent vectors from PCA start in the center and form a circle. The two trajectories from two different Reynolds numbers are hard to distinguish, which makes it challenging for the learned model to capture parameter variations. Although CNN-based embedding methods [14, 50, 29] are also nonlinear, they cannot handle irregular geometries with unstructured meshes due to the limitation of classic convolution operations. Using pivotal nodes combined with GNN learning, the proposed model is both flexible in dealing data with irregular meshes and effective in capturing state transitions in the system.
Diagnosis of attention weights. The attention vector at in equation 9 plays an important role in predicting the state of the dynamic system. For example, if the system enters a perfect oscillating
Re = 307 ( ), Re = 993 ( ) start from and respectively.
stage, the model could directly look up and re-use the encoding from the last cycle as the prediction, to make predictions with little deterioration or drift.
We can observe the way our model makes predictions via the attention magnitude over time. Figure 6 shows attention weights for the initial parameter token z−1, which encodes the system parameters µ, over one cylinder flow trajectory. In the initial transition phase, we see high attention to the parameter token, as the initial velocity fields of different Reynold numbers are hard to distinguish. As soon as vortex shedding begins around t = 120, and the flow enters an oscillatory phase, and attention shift away from the parameter token, and locks onto these cycles. This can be seen in Figure 7 (left): we perform a Fourier analysis of the attention weights, and notice that the peak of the attention frequency (blue) coincides with the vortex shedding frequency in the data (purple). That is, our dynamics model learned to attend to previous cycles in vortex shedding, and can use this information to make accurate predictions that keep both frequency and phase stable. Figure 7 (right) shows that this observation holds for the whole range of Reynolds numbers, and their corresponding vortex shedding frequencies.
For sonic flow, the system state during the first few steps is much easier to discern even without knowing the system parameters µ, and hence, the attention to z−1 quickly decays. This shows that the transformer can adjust the attention distribution to the most physically useful signals, whether that is parameters or system state. Analysis on this, as well as the full attention weights maps can be found in section A.10.

heading: 5 CONCLUSION
text: In this paper we introduced a graph-based mesh reduction method, together with a temporal attention module, to efficiently perform autoregressive prediction of complex physical dynamics. By attending to whole simulation trajectories, our method can more easily identify conserved properties or fundamental frequencies of the system, allowing prediction with higher accuracy and less drift compared to next-step methods.

heading: ACKNOWLEDGEMENT
text: Han Gao and Jianxun Wang are supported by the National Science Foundation under award numbers CMMI-1934300 and OAC-2047127. Li-Ping Liu are supported by NSF 1908617. Xu Han was also supported by NSF 1934553.

heading: A APPENDIX
text: A.1 INFORMATION FLOW IN GMR AND GMUS
Graph Mesh Reducer Graph Mesh Up-Sampling
Figure 8 shows how the information flows in Graph Mesh Reducer (GMR) and Graph Mesh UpSampling (GMUS). This is analogous to a V-cycle in multigrid methods: Consider a graph with seven nodes (|V| = 7) to be reduced and recovered by GMR and GMSU with two message-passing layers. The increased size of pivotal nodes after a message-passing layer represents the information aggregation from surrounding nodes.

heading: A.2 COMPUTATIONAL COMPLEXITY
text: The computation of the proposed model is from the three components: GMR, GMUS, and the transformer. The graph G containsN nodes and also hasO(N) edges because of its special structure. Then GMR and GMUS as L-layer graph neural networks have computation time O(NLd∗), with d∗ being the maximum number of hidden units.
The transformer takes time O(T · (d′)2) to make a prediction: a prediction needs to consider O(T ) previous steps, and the computation over a time step is O((d′)2). Note that d′ is the length of the latent vector at each step.
In our actual calculation, the transformer takes much less time than GMR or GMUS since the latent dimension d′ is small.
However, it would be very expensive if the model did not use GMR and directly applied a transformer to the original system states. In this case, d′ would be Nd, and the total complexity would be O(TN2d2), which would not be feasible for typical computation servers.

heading: A.3 MODEL TRAINING
text: The GMR/GMUS and temporal attention model are trained separately. We first train GMR/GMUS by minimizing the reconstruction loss Lgraph defined by Eqn. 6. Then, the parameters of the trained encoder-decoder are frozen, and only the parameters of the temporal attention model will be trained by minimizing the prediction loss Latt (Eqn. 12) with early stopping. The parameter token encoder mlpp is trained jointly with temporal attention. The training algorithm is detailed in Algorithm 1.
Algorithm 1 Training Process Input: Domain graph G, Node features over time [Y0, . . . ,YT ], GMRθ, GMUSφ,Attentionω , RMSE threshold R′, MLPψ , condition parameter µ Output: Learned parameters θ, φ, ω and ψ repeat
for Yt ∈ [Y0, . . . ,YT ] do Ŷt = GMUSφ(GMRθ(Yt,G)) end for Compute∇θ,φ ← ∇θ,φLgraph(θ, φ, [Y0, . . . ,YT ], [Ŷ0, . . . , ŶT ]) Update φ, θ using the gradients∇φ,∇θ
until convergence of the parameters (θ, φ) [z0, z1, . . . , zT ] = GMRθ([Y0, . . . ,YT ]), z−1 = MLPψ(µ) for t ∈ [1, 2, . . . , T ] do
repeat [ẑ1, . . . , ẑt] = Attentionω([z−1, z0]) Compute ∇ω,ψ ← ∇ω,ψLatt(ω, ψ, [z1, . . . ,zt], [ẑ1, . . . , ẑt]) Update ω, ψ using the gradients∇ω ,∇ψ
until RMSE([z1, . . . ,zt], [ẑ1, . . . , ẑt]) < R′ end for

heading: A.4 DATASET DETAILS
text: Three flow datasets, cylinder flow, sonic flow and vascular flow, are used in our numerical experiments, and the schematics of flow configurations are shown in Figure 9. Cylinder and vascular flow are governed by incompressible Navier-Stokes (NS) equations, whereas compressible NS equations govern sonic flow. The ground truth datasets are generated by solving the incompressible/compressible NS equations based on the finite volume method (FVM). We use the open-source FVM library OpenFOAM [19] to conduct all CFD simulations. All flow domains are discretized by triangle (tri) mesh cells and the mesh resolutions are adaptive to the richness of flow features. The case setup details can be found in Table 2. Both cylinder flow and sonic flow datasets consist of 51 training and 50 test trajectories, while vascular flow consists of 11 training and 10 test trajectories.
Here we summarize all input and outputs of the GMR/GMUS models, as well as the attention simulator for each dataset. GMR encodes the high-dimensional system state onto the node presentations on pivotal nodes, while GMR decodes latents back to the high-dimensional physical state (Figure 8). The parameter token encoder takes the physical parameter µ as input, and outputs a parameter token. The attention model takes the parameter token, together with the latent vectors of states at all previous time steps z, to predict the latents of the next time step. All input and output variables for each dataset are detailed in Table 3.
The governing equations for all fluid dynamic cases can be summarized as follows,
∂ρ ∂t +∇ · (ρv) = 0,
∂(ρv) ∂t +∇ · (ρvv) = ∇ · (µ∇v)−∇p+∇ · (µ(∇v)T )− 2 3 ∇(µ∇ · v),
∂ ∂t (ρcpT ) +∇ · (ρcpvT ) = ∇ · (k∇T ) + ρT Dcp Dt + Dp Dt − 2 3 µΨ + µΦ,
(14)
where v is the velocity vector, µ is the viscosity, k is the thermal conductivity, cp is the specific heat of the fluid and the definitions of the stream scalar Ψ, the potential scalar Φ are referred to the chapter 3 in [28]. Only the sonic flow (compressible flow) involves solving the third equation.
For cylinder flow, the simulation mesh topology remains the same for all trajectories, and the Reynolds number varies. In sonic flow, topology also remains fixed; however, the mesh node coordinates change over the course of the simulation to simulate the moving ramp. Here, we vary the initial temperature. Finally, for vascular flow, we vary the radius of the thrombus, which means that each trajectory will have a different mesh topology, and the simulation model should be able to work with variable-sized inputs in this case.

heading: A.5 MODEL DETAILS
text: Node and edge representations in our GraphNets are vectors of width 128. The node and edge functions (mlpv , mlpe, mlpr) are MLPs with two hidden layers of size 128, and ReLU activation. The input dimension of mlpv is based on the number of input features in each node (see details in table 3), and the input dimension of mlpe is three due to the one-layer neighboring edge. We use L = 3 GraphNet blocks for both GMR and GMUS. The node and edge functions mlpe` and mlp v ` used in the GraphNet blocks are 3-Layer, ReLU-activated MLPs with hidden size of 128. The output
size is 128 for the all MLPs, except for the final layer, which is of size 4 for Cylinder, Sonic and 2 for Vascular.
We use a single layer and four attention heads in our transformer model. The embedding sizes of z for each dataset (cylinder flow, sonic flow, and vascular flow) are 1024, 1024, and 800, respectively. This is calculated by #pivotal nodes × nout, in which nout is the output size of the last GraphNet block. The mlpp used to encode the system parameter µ is a MLP with 2 hidden layers of size 100, and with output dimensions that match z, i.e. 1024, 1024, and 800 for Cylinder Flow, Sonic Flow, and Vascular Flow, respectively.

heading: A.6 PIVOTAL POINTS SELECTION
text: In general, pivotal nodes are selected by uniformly sampling from the entire mesh, as it preserves the mesh density distribution, which is designed based on the flow physics observed in training data. For our datasets, we select 256 pivotal nodes out of 1699 cells for cylinder flow, 256 pivotal nodes out of 1900 cells for sonic flow, and 400 pivotal nodes out of 7561 cells for vascular flow. The pivotal nodes in vascular flow are manually reduced in the aneurysm region, where flow features are not rich. The pivotal nodes distributions for the three flow systems are shown in Figures 10, 11, and 12, respectively.

heading: A.7 DATA COMPRESSION
text: The Graph Mesh Reducer (GMR) and Graph Mesh Up-Sampling (GMUS) can be treated as an encoder and decoder for compressing high-dimensional physical data on irregular unstructured meshes. For any given mesh data (G,Y ), we can encode it into a latent vector z = GMR(G,Y ) through GMR. Therefore, a high-dimensional spatiotemporal trajectory (G,Y0, . . . ,YT ) can be compressed into a low-dimensional representation (G, z0, . . . ,zT ), which significantly reduces memory for storage. Similar to other compressed sensing techniques, the original data can be restored by GMUS from latent vectors. Table 5 shows that the decoder can significantly reduce the size of the original data to about 3% with a slight loss of accuracy after being restored by the decoder. Note that only the GMR/GMUS are included in the data compressing process, the transformer dynamics model is not involved.
A.8 INFERENCE TIME
Our ground truth solver uses the finite volume method to solve the governing PDEs. Specifically, the ”pressure implicit with splitting of operators” (PISO) algorithm and ”semi-implicit method for pressure linked equations” (SIMPLE) [28] are used to simulate the unsteady cylinder and vascular flows, and a PISO-based compressible solver [25] is used to simulate the high-speed sonic flow. All these traditional numerical methods involve a large number of iterations and time integration steps. In contrast, the trained neural network model is able to make predictions without the need of subtimestepping. Hence, these predictions can be are significantly faster than running the ground-truth simulator. The online evaluation of our learned model at inference time only involves two steps,
Step 1: ẑ1 = trans(z−1, z0), ẑt = trans(z−1, z0, ẑ1, . . . , ẑt−1) Step 2: Ỹt = GMUS(z̃t), t = 1, . . . , T (15)
which are completely decoupled. We compare the evaluation cost of the learned model with the FV-based numerical models in Table 6, and observe significant speedups for all three datasets.

heading: A.9 OTHER VARIABLE CONTOUR
text: Figure 13 shows the pressure contours.

heading: A.10 FURTHER ANALYSIS OF LATENT ENCODINGS AND ATTENTION
text: Figure 14 analyzes the latent encoding and attention for the sonic flows in a similar vein to section 4. Different initial temperatures (as the system parameter) lead to fast diverging of state dynamics, and thus the parameter tokens overlap for different parameters. However, the attentions to z−1 also rapidly decay, demonstrating that the transformer can adjust the attention distribution to capture most physically useful signals.
speed sonic flow system: T (0) = 201 trajectory ( ), T (0) = 299 trajectory ( ) start from their parameter tokens and , respectively. (c) Attention values of the parameter tokens versus time, T (0) = 201 ( ), and T (0) = 299 ( ).
Figure 15 shows the overview of attention value distribution in a time-time diagram for all three fluid systems. The attentions of the state to that at different time step is plotted as a log-scale contour. Note that the sum of attention values in the same row equals to 1 due to the softmax layer (Eqn. 9).

heading: A.11 LONGER ROLLOUTS
text: We perform an additional experiment to test how the model performs when rolling out of the training range. We train a model for the cylinder flow setup (n=400 steps), and evaluate it for lengths of 2n.
To limit the memory footprint for very long rollouts, we use a sliding window of n for the attention, i.e., we always attend to the last n time steps. This way allow us to retain the same complexity for the training. Fig. 16 illustrates this procedure: we always use the parameter embedding token z̃−1 as the first entry of the attention sequence, followed by a moving window of the n last timesteps.
Table 7 compares the performance of the proposed model with MeshGraphNet on a longer rollout length (2n = 2× 400 steps). We find that the proposed model still outperforms the baseline for all u, v and p trajectories.

heading: A.12 PREDICTION ON CONVECTION-DIFFUSION FLOW
text: Many of the systems we study in the main paper have oscillatory behaviors, which are hard to capture by next-step models stably. To further demonstrate the capability of the model on predicting non-oscillatory dynamics, we add an additional example: a convection-diffusion flow, governed by
∂c ∂t +∇ · (vc) = ∇2( |v| Pe c), (16)
where c is the concentration of the transport scalar, v = [1, 1] is the convection velocity, andPe is the Peclet number, which controls the ratio of the advective transport rate to its diffusive transport rate.
As shown in Figure 17, the model can accurately predict the non-oscillatory sequences of rollouts Quantitatively, our model (RMSE = 1.03 × 10−3) still outperforms MeshGraphNet (RMSE = 2.79× 10−3).

heading: A.13 CARDIAC FLOW PREDICTION
text: We test the proposed model on cardiac flows (pulsatile flows) with varying viscosity in multiple cardiac cycles, which is more realistic for cardiovascular systems. An idealized spatial-temporal inflow condition is defined as
v(x, t) = 0.6(x− xmin)(x− xmax)
(xmax − xmin)2 (sin(
2πt 1.8 ))2 + 0.15 (17)
to simulate the pulsating inflow from the heart. The model can accurately predict 10 cardiac cycles (Figure 18). The comparison with MeshGraphNet is listed in Table 8, showing significantly better performance for our model.
Table 8: The average relative rollout error for cardiac flow, with unit of ×10−3.
rollout step 300 Variable u v p
MeshGraphNet-NI 194 71 944 Ours-Transformer 4.7 2.3 106
Cardiac Flow t=0 t=90 t= 200 t=270
Tr ut h O ur s
Figure 18: The velocity contours of rollouts predicted by our model versus the ground truth (CFD) for a cardiac flow.

Review summaries:
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.

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.