P REDICTING PHYSICS IN MESH - REDUCED SPACE WITH TEMPORAL ATTENTION
A BSTRACT
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 efﬁcient manner.
Our method outperforms a competitive GNN baseline on several complex ﬂuid dynamics prediction tasks, from sonic shocks to vascular ﬂow.
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 beneﬁts of attention-based sequence models to solving high-dimensional complex physics tasks.
I NTRODUCTION
There has been an increasing interest in many scientiﬁc disciplines, from computational ﬂuid 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 ﬂexible.
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 deﬁned 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 ﬁrst use a GNN to aggregate local information of solution ﬁelds 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 ﬁne-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 ﬂuid 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.
R ELATED W ORK
Developing and running simulations of complex, high-dimensional systems can be very timeintensive.
Particular for computational ﬂuid dynamics, there is considerable interest in using neural networks for accelerating simulations of aerodynamics [3] or turbulent ﬂows [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 ﬂexible 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 beneﬁts 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 inefﬁcient 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].
M ETHODOLOGY 3. P ROBLEM D EFINITION AND O VERVIEW
We are interested in predicting spatiotemporal dynamics of complex physical systems (e.g., ﬂuid dynamics), usually governed by a set of nonlinear, coupled, parameterized partial differential equations [Figure] (PDEs) as shown in its general form,
[Equation]
where u ( x , t ) ∈ R d 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 ﬁnite volume method (FVM) [28], which discretizes the simulation domain Ω into an unstructured mesh consisting of N cells ( C i : i = 1 , . . .
, N ) .
At time t , the discrete state ﬁeld Y t = { u i,t : i = 1 , . . .
, N } can thus be deﬁned by the state value u i,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 Y 0 and system parameters µ , can rapidly produce a rollout trajectory of states Y 1 ...
Y T .
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 ﬁne grid, such as pressure updates in incompressible ﬂow, can require many solver iterations.
Commonly used technique to improve the efﬁciency include multigrid methods [46], which perform local updates on both the ﬁne 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 ﬁne to a coarser grid, performing a solver update, up-sampling back on the ﬁne grid, and performing an update on the ﬁne 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 ﬁne 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 ﬁne simulation, and perform local updates using a mesh recovery GNN (section 3.2).
This combination allows us to efﬁciently make use of long-term temporal attention for stable, high-ﬁdelity dynamics predictions.
3. G RAPH R EPRESENTATION AND M ESH R EDUCTION
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 C i , so the graph size is |V| = N .
The set of edges E are derived from neighboring relations of cells: if two cells C i and C j are neighbors, then two directional edges ( i, j ) and ( j, i ) are both in E .
We ﬁx this graph for all time steps.
At each step t , each node i ∈ V uses the local state vector u i,t as its attribute.
Therefore, ( G , ( Y 0 , . . .
, Y T )) forms a temporal graph sequence.
The goal of the learning model is to predict ( Y 1 , . . .
, Y T ) given G and Y 0 .
Mesh Reduction.
After representing the system as a graph, we use a GNN to summarize and extract a low-dimensional representation z t from Y t 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 ﬁrst 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 ﬁrst extracts node features and edge features from the system state using the node and edge Multi-Layer Perceptrons (MLPs).
[Equation]
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 C i .
Then GMR uses L GraphNet processer blocks [37] to further reﬁne 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
[Equation]
Here v 0 i , e 0 ij are the outputs of Equation 2, and N i denotes all neighbors of node i .
The two functions mlp e(cid:96) ( · ) and mlp v(cid:96) ( · ) 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
=
( v Li : 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:
[Equation]
We concatenate these vectors into a single latent z = concat( h i : 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 deﬁne 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 ﬁrst set the representations at pivotal nodes by splitting z , that is, r i
= h
i , 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 (cid:48)
j of k nearest pivotal nodes in terms of spatial distance and compute the its representations r j by
[Equation]
Here d ( j, i ) is the spatial distance between cells C j and C i .
Then every node i ∈ V has a representation r i , and all nodes’ representations are collectively denoted as R =
( r i : 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 ﬂows 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 , Y t )) and minimize the reconstruction loss
[Equation]
With these two modules we can encode system states ( Y 1 , . . .
, Y T ) to latents ( z 1 , . . .
, z T ) as a lowdimensional representation of system dynamics.
In the next section, we train a transformer model to predict the sequences of latents.
3. A TTENTION - BASED S IMULATOR
As a dynamics model, we learn a simulator which can predict the sequence ( z 1 , . .
, z T ) autoregressively based on an initial state z 0 =
GMR( G , Y 0 ) 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 ﬂows 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 z t
[Equation]
[Equation] Then, the transformer model [43], trans( · ) predicts the subsequent latent vectors in an autoregressive manner.
[Equation]
Speciﬁcally, we use a single layer of multi-head attention in our transformer model, which we found sufﬁciently powerful for the environments we studied.
We ﬁrst 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 , z 0 , ˜ z 1 , . . .
˜
z t − 1 ) .
[Equation]
Here W 1 and W 2 are learnable parameters, and d (cid:48) denotes the length of the vector z t − 1 .
This equation corresponds the popular inner-product attention model [43].
Next, the attention head computes the prediction vector
[Equation]
with the learned parameter W 3 .
Our multi-head attention model uses K parallel attention heads with separate parameters to compute
K vectors ( g 1 t , . . .
, g Kt ) as described above.
These vectors are concatenated and fed into an MLP to predict the residual of ˜ z t over ˜
z t − 1 .
[Equation]
[Equation]
We train the autoregressive model on entire sequences by minimizing the prediction loss
[Equation]
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.
3. M ODEL TRAINING AND TESTING
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 efﬁciency; 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 (cid:80) tt (cid:48)
=1 (cid:107) z t (cid:48)
− ˜ z t (cid:48) (cid:107) 22 .
Until the objective reaches a threshold, we add the next loss term (cid:107) z t +1 − ˜
z t +1 (cid:107) 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 ﬁrst represent system parameters as a token z − 1 by equation 7 and encode the initial state into a token z 0 =
GMR( G , Y 0 ) , then we predict the sequence ˜ z 1 , . . .
, ˜ z T by equation 8, and ﬁnally we decode the system state ˜ Y by
[Equation]
RESULTS
Datasets.
We tested the proposed method to three datasets of ﬂuid dynamic systems: (1) ﬂow over a cylinder with varying Reynolds number; (2) high-speed ﬂow over a moving wedge with different initial temperatures; (3) ﬂow in different vascular geometries.
These datasets correspond to applications with ﬁxed, 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. shows pivotal nodes for each dataset.
Stable and visually accurate model predictions.
We tested the proposed method on three ﬂuid 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 ﬂow, 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 ﬂow 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 ﬂow with a variable-size circular thrombus, and thus different mesh sizes, the model is able to predict the ﬂow accurately.
In both cylinder ﬂow and vascular ﬂow, 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.
Re = 307 (Cylinder Flow)
Re = 993 (Cylinder Flow) t= t= t=t= t= t= t=t=
T r u t h
O u r s
T (0) = 201 (Sonic Flow)
T (0) = 299 (Sonic Flow) t= t= t=t= t= t= t=t=
T r u t h
O u r s
R = 0 .
31 (Vascular ﬂow)
R = 0 .
49 (Vascular ﬂow) t= t= t=t= t= t= t=t=
T r u t h
O u r s
[Table]
Error behavior under long rollouts To quantitatively assess the performances of our models and baselines, we compute the relative mean square error (RMSE), deﬁned as
RMSE( ˆu
[Equation]
[Equation]
) =
[Equation] ), over the full rollout trajectory. Table 1 compares average prediction errors on all three domains. Our attention-based model outperforms baselines and model variants in most scenarios.
For sonic ﬂow, 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 ﬂow and vascular ﬂow) highlight the superior performance of our attention-based method compared with baseline next-step models, and the RMSE of our method is signiﬁcantly lower than baselines.
Figure 3 shows how error accumulates in different models.
Our model has higher error for the ﬁrst 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 [Figure]
[Figure] [Equation] [Figure] 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 ﬂow, as discussed later in this section.
In Figure 4 we illustrate the difference in the models’ abilities to predict long rollouts.
For cylinder ﬂow, 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 ﬂow, our model gives an accurate prediction at step t = 228 , while the prediction of MeshGraphNet is signiﬁcantly different from the ground-truth at this step: the left inﬂow 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 ﬂow dataset and show that the model can distinguish rollouts from different Reynolds numbers.
We ﬁrst 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 ﬂexible in dealing data with irregular meshes and effective in capturing state transitions in the system.
Diagnosis of attention weights.
The attention vector a t in equation 9 plays an important role in predicting the state of the dynamic system.
For example, if the system enters a perfect oscillating
−− 0 .− 0 .− 0 .− 0 .
C
C−− 0 . 0 .
−− 0 . 0 .
C
Time Step
P a r a m e t e r T
o k e n
A tt e n t i o n
=
[Figure]
[Figure]
M a gn it ud e
[Figure]
[Figure]
Frequency  1 ,0 .0 .0 .
Reynolds Number
[Figure]
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.
Figureshows attention weights for the initial parameter token z − 1 , which encodes the system parameters µ , over one cylinder ﬂow trajectory.
In the initial transition phase, we see high attention to the parameter token, as the initial velocity ﬁelds of different Reynold numbers are hard to distinguish.
As soon as vortex shedding begins around t = 120 , and the ﬂow 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 ﬂow, the system state during the ﬁrst 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.
C ONCLUSION
In this paper we introduced a graph-based mesh reduction method, together with a temporal attention module, to efﬁciently 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. A CKNOWLEDGEMENT
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.
A. C OMPUTATIONAL COMPLEXITY
The computation of the proposed model is from the three components: GMR, GMUS, and the transformer.
The graph G contains N nodes and also has O ( 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 (cid:48) ) 2 ) to make a prediction: a prediction needs to consider O ( T ) previous steps, and the computation over a time step is O
(( d (cid:48) ) 2 ) .
Note that d (cid:48) 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 (cid:48) 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 (cid:48) would be Nd , and the total complexity would be O ( TN 2 d 2 ) , which would not be feasible for typical computation servers. [Figure]
[Figure]
M ODEL T RAINING
The GMR/GMUS and temporal attention model are trained separately.
We ﬁrst train GMR/GMUS by minimizing the reconstruction loss L graph deﬁned 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 L att (Eqn. 12) with early stopping.
The parameter token encoder mlp p is trained jointly with temporal attention.
The training algorithm is detailed in Algorithm 1.
Algorithm
Input: Domain graph G , Node features over time
[ Y 0 , . . .
, Y T ] , GMR θ , GMUS φ , Attention ω ,
RMSE threshold R (cid:48) ,
MLP ψ , condition parameter µ
Output: Learned parameters θ , φ , ω and ψ repeatfor Y t ∈
[ Y 0 , . . .
, Y T ] do ˆ Y t =
GMUS φ (GMR θ
( Y t , G )) end for
Compute ∇ θ,φ ← ∇ θ,φ L graph ( θ, φ, [ Y 0 , . . .
, Y T ] , [ ˆ Y 0 , . . .
, ˆ Y T ]) Update φ , θ using the gradients ∇ φ , ∇ θ until convergence of the parameters ( θ , φ )
[ z 0 , z 1 , . . .
, z T ] =
GMR θ ([ Y 0 , . . .
, Y T ]) , z − 1 =
MLP ψ ( µ ) for t ∈ [1 , 2 , . .
, T ] do repeat [ˆ z 1 , . . .
, ˆ z t ] =
Attention ω
([ z − 1 , z 0 ])
Compute ∇ ω,ψ ← ∇ ω,ψ L att ( ω, ψ, [ z 1 , . . . , z t ] , [ˆ z 1 , . . .
, ˆ z t ]) Update ω , ψ using the gradients ∇
A. D ATASET DETAILS
Three ﬂow datasets, cylinder ﬂow, sonic ﬂow and vascular ﬂow, are used in our numerical experiments, and the schematics of ﬂow conﬁgurations are shown in Figure 9.
Cylinder and vascular ﬂow are governed by incompressible Navier-Stokes (NS) equations, whereas compressible NS equations govern sonic ﬂow.
The ground truth datasets are generated by solving the incompressible/compressible NS equations based on the ﬁnite volume method (FVM).
We use the open-source
FVM library OpenFOAM [19] to conduct all CFD simulations.
All ﬂow domains are discretized by triangle (tri) mesh cells and the mesh resolutions are adaptive to the richness of ﬂow features.
The case setup details can be found in Table 2.
Both cylinder ﬂow and sonic ﬂow datasets consist oftraining and 50 test trajectories, while vascular ﬂow consists of 11 training and 10 test trajectories.
[Table]
[Table]
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
[Table]
while GMR decodes latents back to the high-dimensional physical state (Figure 8).
The parameter token encoder takes the physical parameter µ as input
[Table]
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 ﬂuid dynamic cases can be summarized as follows,
[Equation]
where v is the velocity vector, µ is the viscosity, k is the thermal conductivity, c p is the speciﬁc heat of the ﬂuid and the deﬁnitions of the stream scalar Ψ , the potential scalar Φ are referred to the chapter 3 in [28].
Only the sonic ﬂow (compressible ﬂow) involves solving the third equation.
For cylinder ﬂow, the simulation mesh topology remains the same for all trajectories, and the
Reynolds number varies.
In sonic ﬂow, topology also remains ﬁxed; 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 ﬂow, 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.
A. M ODEL DETAILS
Node and edge representations in our GraphNets are vectors of width 128.
The node and edge functions (mlp v , mlp e , mlp r ) are MLPs with two hidden layers of size 128, and ReLU activation.
The input dimension of mlp v is based on the number of input features in each node (see details in table 3), and the input dimension of mlp e 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 mlp e(cid:96) and mlp v(cid:96) used in the GraphNet blocks are 3-Layer, ReLU-activated MLPs with hidden size of 128.
The output
[Figure] [Figure]
size is 128 for the all MLPs, except for the ﬁnal layer, which is of size 4 for Cylinder, Sonic andfor Vascular.
We use a single layer and four attention heads in our transformer model.
The embedding sizes of z for each dataset (cylinder ﬂow, sonic ﬂow, and vascular ﬂow) are 1024 , 1024 , and 800 , respectively.
This is calculated by # pivotal nodes × n out , in which n out is the output size of the last GraphNet block.
The mlp p 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.
[Table]
A. P IVOTAL POINTS SELECTION
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 ﬂow physics observed in training data.
For our datasets, we select 256 pivotal nodes out of 1699 cells for cylinder ﬂow, 256 pivotal nodes out of 1900 cells for sonic ﬂow, and 400 pivotal nodes out of 7561 cells for vascular ﬂow.
The pivotal nodes in vascular ﬂow are manually reduced in the aneurysm region, where ﬂow features are not rich.
The pivotal nodes distributions for the three ﬂow systems are shown in Figures 10, 11, and 12, respectively.
[Figure]
[Table]
D ATA C OMPRESSION
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 , Y 0 , . . .
, Y T ) can be compressed into a low-dimensional representation ( G , z 0 , . . .
, z T ) , which signiﬁcantly 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 signiﬁcantly 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
[Table]
A. I NFERENCE TIME
Our ground truth solver uses the ﬁnite volume method to solve the governing PDEs.
Speciﬁcally, 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 ﬂows, and a PISO-based compressible solver [25] is used to simulate the high-speed sonic ﬂow.
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 signiﬁcantly faster than running the ground-truth simulator.
The online evaluation of our learned model at inference time only involves two steps,
[Equation]
which are completely decoupled.
We compare the evaluation cost of the learned model with the
FV-based numerical models in Table 6, and observe signiﬁcant speedups for all three datasets.
[Table]
A. O THER VARIABLE CONTOUR
Figure 13 shows the pressure contours.
Re = 307 (Cylinder Flow)
Re = 993 (Cylinder Flow) t= t= t=t= t= t= t=t=
T r u t h
O u r s
T (0) = 201 (Sonic Flow)
T (0) = 299 (Sonic Flow) t= t= t=t= t= t= t=t=
T r u t h
O u r s
R = 0 .
21 (Vascular ﬂow)
R = 0 .
39 (Vascular ﬂow) t= t= t=t= t= t= t=t=
[Figure]
With different condition settings, the model accurately predict long sequences of rollouts.
A. F URTHER ANALYSIS OF LATENT ENCODINGS AND ATTENTION
Figure 14 analyzes the latent encoding and attention for the sonic ﬂows 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.
[Figure] [Figure]
[Figure]
Time Step
[Figure]
r
[Figure]
m e t e r
T ok e n
A tt e n ti on (c)
Figure 15 shows the overview of attention value distribution in a time-time diagram for all three ﬂuid 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).
[Figure]
[Table]
A. L ONGER ROLLOUTS
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 ﬂow setup ( n =400 steps), and evaluate it for lengths of 2 n .
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 −as the ﬁrst 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 ( 2 n = 2 × 400 steps).
We ﬁnd that the proposed model still outperforms the baseline for all u, v and p trajectories.
[Equation]
[Figure]
A. P REDICTION ON CONVECTION - DIFFUSION FLOW
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 ﬂow, governed by
[Equation]
(16) where c is the concentration of the transport scalar, v = [1 , 1] is the convection velocity, and Pe is the
Peclet number, which controls the ratio of the advective transport rate to its diffusive transport rate.
[Figure]
[Equation]
[Figure]
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 ).
A. C ARDIAC FLOW PREDICTION
We test the proposed model on cardiac ﬂows (pulsatile ﬂows) with varying viscosity in multiple cardiac cycles, which is more realistic for cardiovascular systems.
An idealized spatial-temporal inﬂow condition is deﬁned as
[Equation]
(17) to simulate the pulsating inﬂow from the heart.
The model can accurately predict 10 cardiac cycles (Figure 18).
The comparison with MeshGraphNet is listed in Table 8, showing signiﬁcantly better performance for our model.
[Table] [Table] =
[Table]
=
[Table] =
T r u t h O u r s