# Predicting physics in mesh-reduced space
with temporal attention

Xu Han
Tufts University
<EMAIL_ADDRESS>
&Han Gao11footnotemark: 1
University of Notre Dame
<EMAIL_ADDRESS>
&Tobias Pfaff11footnotemark: 1
DeepMind
<EMAIL_ADDRESS>
Jian-Xun Wang
University of Notre Dame
<EMAIL_ADDRESS>
&Li-Ping Liu
Tufts University
<EMAIL_ADDRESS>
Equal contribution.

###### 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.

##  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 $\bm{u}(t+1,\bm{x})$ of a physical system from the current
state $\bm{u}(t,\bm{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 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 time-intensive. 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 low-dimensional 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].

##  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,

$\displaystyle\frac{\partial\bm{u}(\bm{x},t)}{\partial
t}=\mathcal{F}\big{[}\bm{u}(\bm{x},t);\bm{\mu}\big{]},\quad\bm{x},t\in\Omega\times[0,T_{e}],$
(1)

where $\bm{u}(\bm{x},t)\in\mathbb{R}^{d}$ represents the state variables
(e.g., velocity, pressure, or temperature) and $\mathcal{F}$ is a partial
differential operator parameterized by $\bm{\mu}$. 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 $\Omega$ into an unstructured
mesh consisting of $N$ cells $(C_{i}:i=1,\ldots,N)$. At time $t$, the discrete
state field $\bm{Y}_{t}=\\{\bm{u}_{i,t}:i=1,\ldots,N\\}$ can thus be defined
by the state value $\bm{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 real-time 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 $\bm{Y}_{0}$ and system parameters
$\bm{\mu}$, can rapidly produce a rollout trajectory of states
$\bm{Y}_{1}...\bm{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 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.

Figure 1: The diagram of the proposed model, GMR-Transformer-GMUS. We first
represent the domain as a graph and then select pivotal nodes
(red/green/yellow) to encode information over the entire graph. The encoder
GMR runs _Message passing_ along graph edges so that the pivotal nodes collect
information from nearby nodes. The latent vector $\bm{z}_{t}$ summarizes
information at the pivotal nodes, and represents the whole domain at the
current step. The transformer will predict $\bm{z}_{t+1}$ based on _all_
previous state latent vectors. Finally, we decode $\bm{z}_{t+1}$ through
message passing to obtain the next-step prediction $\mathbf{Y}_{t+1}$.

###  3.2 Graph Representation and Mesh Reduction

Graph Representation. We construct a graph
${\mathcal{G}}=(\mathcal{V},\mathcal{E})$ to represent a snapshot of a dynamic
system at time step $t$. Here each node $i\in\mathcal{V}$ corresponds the mesh
cell $C_{i}$, so the graph size is $|\mathcal{V}|=N$. The set of edges
$\mathcal{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 $\mathcal{E}$. We fix this graph for all time steps. At
each step $t$, each node $i\in V$ uses the local state vector $\bm{u}_{i,t}$
as its attribute. Therefore, $({\mathcal{G}},(\bm{Y}_{0},\ldots,\bm{Y}_{T}))$
forms a temporal graph sequence. The goal of the learning model is to predict
$(\bm{Y}_{1},\ldots,\bm{Y}_{T})$ given $\mathcal{G}$ and $\bm{Y}_{0}$.

Mesh Reduction. After representing the system as a graph, we use a GNN to
summarize and extract a low-dimensional representation $\bm{z}_{t}$ from
$\bm{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 first selects a
small set $\mathcal{S}\subseteq\mathcal{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 $\mathcal{S}$ from $\mathcal{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).

$\displaystyle\bm{v}_{i}^{0}=\mathrm{mlp}_{v}(\bm{Y}[i]),\quad\bm{e}_{ij}^{0}=\mathrm{mlp}_{e}(p(i)-p(j)).$
(2)

Here $\bm{Y}[i]$ is the $i$-th row of $\bm{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 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

$\displaystyle\bm{e}^{\ell}_{ij}=\mathrm{mlp}_{\ell}^{e}\left(\bm{e}^{\ell-1}_{ij},\bm{v}_{i}^{\ell-1},\bm{v}_{j}^{\ell-1}\right),\quad\bm{v}_{i}^{\ell}=\mathrm{mlp}^{v}_{\ell}\left(\bm{v}_{i}^{\ell-1},\sum_{j\in\mathcal{N}_{i}}\bm{e}_{ij}^{\ell-1}\right),\quad\ell=1,\ldots,L.$
(3)

Here $\bm{v}_{i}^{0},\bm{e}^{0}_{ij}$ are the outputs of Equation 2, and
$\mathcal{N}_{i}$ denotes all neighbors of node $i$. The two functions
$\mathrm{mlp}_{\ell}^{e}(\cdot)$ and $\mathrm{mlp}_{\ell}^{v}(\cdot)$
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 $\bm{V}=(\bm{v}^{L}_{i}:i\in\mathcal{V})$ for the graph
$\mathcal{G}$.

Finally GMR applies an MLP to representations of the pivotal nodes in
$\mathcal{S}$ _only_ to “summarize” the entire graph onto a coarse graph:

$\displaystyle\bm{h}_{i}=\mathrm{mlp}_{r}(\bm{v}^{L}_{i}),\quad i\in S$ (4)

We concatenate these vectors into a single latent
$\bm{z}=\mathrm{concat}(\bm{h}_{i}:i\in S)$ as reduced vector representation
of the entire graph. We collectively denote these three computation steps as
$\bm{z}=\mathrm{GMR}({\mathcal{G}},\bm{Y})$. The latents $\bm{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 $\bm{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
$\bm{z}$, that is, $\bm{r}_{i}=\bm{h}_{i},i\in\mathcal{S}$. We then compute
representations of non-pivotal nodes by spatial interpolation [2]: for a non-
pivotal node $j$, we choose a set $\mathcal{N}^{\prime}_{j}$ of $k$ nearest
pivotal nodes in terms of spatial distance and compute the its representations
$\bm{r}_{j}$ by

$\displaystyle\bm{r}_{j}=\sum_{i\in\mathcal{N}^{\prime}_{j}}\frac{w_{ij}\bm{h}_{i}}{\sum_{i\in\mathcal{N}^{\prime}_{j}}w_{ij}},\quad
w_{ij}=\frac{1}{d(j,i)^{2}}$ (5)

Here $d(j,i)$ is the spatial distance between cells $C_{j}$ and $C_{i}$. Then
every node $i\in\mathcal{V}$ has a representation $\bm{r}_{i}$, and all nodes’
representations are collectively denoted as
$\bm{R}=(\bm{r}_{i}:i\in\mathcal{V})$.

Similar to GMR, GMUS applies EPD GraphNet to the initial node representation
$\bm{R}$ to restore $\bm{Y}$ on the full graph ${\mathcal{G}}$. We denote the
chain of operations so far as
$\hat{\bm{Y}}=\mathrm{GMUS}(\mathcal{G},\bm{z})$. Details about information
flows in GMR and GMUS can be found in section 8.

We train GMR and GMUS as an auto-encoder over all time steps and sequences.
For each time step $t$, we compute
$\hat{\bm{Y}}_{t}=\mathrm{GMUS}(\mathcal{G},\mathrm{GMR}(\mathcal{G},\bm{Y}_{t}))$
and minimize the reconstruction loss

$\displaystyle\mathcal{L}_{graph}=\quad\sum_{n=1}^{T}\|\bm{Y}_{t}-\hat{\bm{Y}}_{t}\|_{2}^{2}\quad.$
(6)

With these two modules we can encode system states
$(\bm{Y}_{1},\ldots,\bm{Y}_{T})$ to latents $(\bm{z}_{1},\ldots,\bm{z}_{T})$
as a low-dimensional representation of system dynamics. In the next section,
we train a transformer model to predict the sequences of latents.

###  3.3 Attention-based Simulator

As a dynamics model, we learn a simulator which can predict the sequence
$(\bm{z}_{1},\ldots,\bm{z}_{T})$ autoregressively based on an initial state
$\bm{z}_{0}=\mathrm{GMR}(\mathcal{G},\bm{Y}_{0})$ and the system parameters
$\bm{\mu}$. 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 $\bm{\mu}$, to
be able to predict flows with arbitrary system parameters at test time.We
encode $\bm{\mu}$ into a special parameter token $\bm{z}_{-1}$ of the same
length as the state representation vectors $\bm{z}_{t}$

$\displaystyle\bm{z}_{-1}=\mathrm{mlp}_{p}(\bm{\mu}).$ (7)

Then, the transformer model [43], $\mathrm{trans}(\cdot)$ predicts the
subsequent latent vectors in an autoregressive manner.

$\displaystyle\tilde{\bm{z}}_{1}=\mathrm{trans}(\bm{z}_{-1},\bm{z}_{0}),\quad\tilde{\bm{z}}_{t}=\mathrm{trans}(\bm{z}_{-1},\bm{z}_{0},\tilde{\bm{z}}_{1},\ldots,\tilde{\bm{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 $\tilde{\bm{z}}_{t-1}$ from the previous step issues a query to
compute attention $\bm{a}_{(t-1)}$ over latents
$(\bm{z}_{-1},\bm{z}_{0},\tilde{\bm{z}}_{1},\ldots,\tilde{\bm{z}}_{t-1})$.

$\displaystyle\bm{a}_{(t-1)}=\mathrm{softmax}\left(\frac{\tilde{\bm{z}}_{t-1}^{\top}\bm{W}_{1}^{\top}\bm{W}_{2}}{\sqrt{d^{\prime}}}\cdot\big{[}\bm{z}_{-1},\bm{z}_{0},\tilde{\bm{z}}_{1},\ldots,\tilde{\bm{z}}_{t-1},\big{]}\right)$
(9)

Here $\bm{W}_{1}$ and $\bm{W}_{2}$ are learnable parameters, and $d^{\prime}$
denotes the length of the vector $\bm{z}_{t-1}$. This equation corresponds the
popular inner-product attention model [43].

Next, the attention head computes the prediction vector

$\displaystyle\bm{g}_{t}=\bm{W}_{3}\big{[}\bm{z}_{-1},\bm{z}_{0},\tilde{\bm{z}}_{1},\ldots,\tilde{\bm{z}}_{(t-1)}\big{]}\bm{a}_{t-1},$
(10)

with the learned parameter $\bm{W}_{3}$.

Our multi-head attention model uses $K$ parallel attention heads with separate
parameters to compute $K$ vectors $(\bm{g}^{1}_{t},\ldots,\bm{g}^{K}_{t})$ as
described above. These vectors are concatenated and fed into an MLP to predict
the residual of $\tilde{\bm{z}}_{t}$ over $\tilde{\bm{z}}_{t-1}$.

$\displaystyle\tilde{\bm{z}}_{t}=\tilde{\bm{z}}_{t-1}+\mathrm{mlp}_{m}\big{(}\mathrm{concat}(\bm{g}^{1}_{t},\ldots,\bm{g}^{K}_{t})\big{)}$
(11)

We train the autoregressive model on entire sequences by minimizing the
prediction loss

$\displaystyle\mathcal{L}_{att}=\sum_{t=1}^{T}\|\bm{z}_{t}-\tilde{\bm{z}}_{t}\|_{2}^{2}.$
(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.

###  3.4 Model 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 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
$\sum_{t^{\prime}=1}^{t}\|\bm{z}_{t^{\prime}}-\tilde{\bm{z}}_{t^{\prime}}\|_{2}^{2}$.
Until the objective reaches a threshold, we add the next loss term
$\|\bm{z}_{t+1}-\tilde{\bm{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 mesh-reduced 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 $\bm{z}_{-1}$ by equation 7 and
encode the initial state into a token
$\bm{z}_{0}=\mathrm{GMR}({\mathcal{G}},\bm{Y}_{0})$, then we predict the
sequence $\tilde{\bm{z}}_{1},\ldots,\tilde{\bm{z}}_{T}$ by equation 8, and
finally we decode the system state $\tilde{\bm{Y}}$ by

$\displaystyle\tilde{\bm{Y}}_{t}=\mathrm{GMUS}(\tilde{\bm{z}}_{t}),\quad
t=1,\ldots,T.$ (13)

##  4 results

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 9.

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.

|  $Re=307$ (Cylinder Flow) |  $Re=993$ (Cylinder Flow)
---|---|---
| t=0 | t=85 | t= 215 | t=298 | t=0 | t=85 | t= 215 | t=298
Truth |  |  |  |  |  |  |  |
Ours |  |  |  |  |  |  |  |
|  $T(0)=201$ (Sonic Flow) |  $T(0)=299$ (Sonic Flow)
| t=0 | t=15 | t= 30 | t=40 | t=0 | t=15 | t= 30 | t=40
Truth |  |  |  |  |  |  |  |
Ours |  |  |  |  |  |  |  |
|  $R=0.31$ (Vascular flow) |  $R=0.49$ (Vascular flow)
| t=40 | t=80 | t= 160 | t=250 | t=40 | t=80 | t= 160 | t=250
Truth |  |  |  |  |  |  |  |
Ours |  |  |  |  |  |  |  |
Figure 2: Contours of the velocity field, as predicted by our model versus the
ground truth (CFD). Our model accurately predicts long rollout sequences under
varying system parameters.

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.

Table 1: The average relative rollout error of three systems, with unit of $\times 10^{-3}$. We compare MeshGraphNet with or without noise injection (NI), to three variants of the our model (LSTM, GRU, or Transformer), which all shared the same GMR/GMUS mesh reduction models. Our model significantly outperform MeshGraphNet on datasets with long rollouts. Dataset-rollout step |  |  Cylinder flow-400 |  Sonic flow-40 |  Vascular flow-250
---|---|---|---|---
Variable |  | $u$ | $v$ | $p$ | $u$ | $v$ | $p$ | $T$ | $u$ | $v$ | $p$
MeshGraphNet | NI | 25 | 778 | 136 | 1.71 | 3.67 | $0.4$ | $0.027$ | $57$ | 133 | 55
without NI | 98 | 2036 | 673 | 4.12 | 6.13 | $\bm{0.24}$ | $\bm{0.020}$ | 3117 | 1771 | 601
Ours | GRU | 114 | 1491 | 1340 | 1.34 | 4.59 | 0.59 | $0.37$ | 8.2 | 11.2 | 23.6
LSTM | 124 | 1537 | 1574 | 1.57 | 5.8 | 0.69 | 0.45 | 8.4 | 11.1 | 23.3
Transformer | $\bm{4.9}$ | $\bm{89}$ | $\bm{38}$ | $\bm{0.95}$ | $\bm{2.8}$ | $0.43$ | $0.39$ | $\bm{7.3}$ | $\bm{10}$ | $\bm{22}$

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
$\mathrm{RMSE}(\mathbf{\hat{u}}^{\mathrm{prediction}},\mathbf{\hat{u}}^{\mathrm{truth}})=\frac{\sum(\hat{u}_{i}^{\mathrm{prediction}}-\hat{u}_{i}^{\mathrm{truth}})^{2}}{\sum(\hat{u}_{i}^{\mathrm{prediction}})^{2}}$,
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 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: Averaged error over all state variables on cylinder flow (left), sonic flow (middle) and vascular flow (right), for the models MeshGraphNets(MGN)(LABEL:line:case0:GMNpure), MGN-NI(LABEL:line:case0:GMN_noiseInjection), Ours-GRU (LABEL:line:case0:GMRGRU), Ours-LSTM (LABEL:line:case0:GMRLSTM), Ours-Transformer (LABEL:line:case0:ourModal). Our model, particularly the transformer, show much less error accumulation compared to the next-step model. $Re=405,t=250$ (Cylinder flow) |  $R=0.31,t=228$ (Vascular flow)
---|---
Next-Step | Ours | Truth | Next-Step | Ours | Truth
|  |  |  |  |
Figure 4: Predictions of the next-step MeshGraphNet model and our model,
compared to ground truth. The next-step model fails to keep the shedding
frequency and show drifts on cylinder flow. On vascular flow, we notice the
left inflow diminishing over the time, while our model remains close to the
reference simulation.

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 $\leq 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 $\mathbf{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 $\bm{z}$-s to two and plot them in Figure 6 (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 ($\mathbf{Y}$) and get 2-d latent vectors directly from
the original data. We plot these vectors in Figure 6 (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.

Figure 5: 2-D principle subspace of the latent vectors from GMR (left) and PCA
(middle) for flow past cylinder system: $Re=307$ (LABEL:ls:LowReTraj),
$Re=993$ (LABEL:ls:HighReTraj) start from LABEL:ls:LowReTrajStart and
LABEL:ls:HighReTrajStart respectively.

Figure 6:  Attention values of the parameter tokens versus time, $Re=307$
(LABEL:ls:LowReFirstAttention), and $Re=993$ (LABEL:ls:HighReFirstAttention).

Diagnosis of attention weights. The attention vector $\bm{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 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
$\bm{z}_{-1}$, which encodes the system parameters $\bm{\mu}$, 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.

Figure 7: Attention (LABEL:ls:AttentionFreq), CFD data (LABEL:ls:CFDFreq).
(left) $Re$ versus frequency for attention and CFD data. Fourier transform for
both attention and CFD data at (middle) $Re=307$ and (right) $Re=993$.

For sonic flow, the system state during the first few steps is much easier to
discern even without knowing the system parameters $\mu$, and hence, the
attention to $\mathbf{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.

##  5 Conclusion

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.

## Acknowledgement

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.

## References

  * Albergo et al. [2019] MS Albergo, G Kanwar, and PE Shanahan.  Flow-based generative models for markov chain monte carlo in lattice field theory.  _Physical Review D_ , 100(3):034515, 2019.
  * Alet et al. [2019] Ferran Alet, Adarsh Keshav Jeewajee, Maria Bauza Villalonga, Alberto Rodriguez, Tomas Lozano-Perez, and Leslie Kaelbling.  Graph element networks: adaptive, structured computation and memory.  In _International Conference on Machine Learning_ , pp. 212–222. PMLR, 2019.
  * Bhatnagar et al. [2019] Saakaar Bhatnagar, Yaser Afshar, Shaowu Pan, Karthik Duraisamy, and Shailendra Kaushik.  Prediction of aerodynamic flow fields using convolutional neural networks.  _Computational Mechanics_ , 64(2):525–545, 2019\.
  * Braess & Hackbusch [1983] Dietrich Braess and Wolfgang Hackbusch.  A new convergence proof for the multigrid method including the v-cycle.  _SIAM journal on numerical analysis_ , 20(5):967–975, 1983.
  * Chattopadhyay et al. [2020] Ashesh Chattopadhyay, Pedram Hassanzadeh, and Devika Subramanian.  Data-driven predictions of a multiscale lorenz 96 chaotic system using machine-learning methods: reservoir computing, artificial neural network, and long short-term memory network.  _Nonlinear Processes in Geophysics_ , 27(3):373–389, 2020.
  * Chen et al. [2018] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud.  Neural ordinary differential equations.  _arXiv preprint arXiv:1806.07366_ , 2018.
  * de Avila Belbute-Peres et al. [2018] Filipe de Avila Belbute-Peres, Kevin Smith, Kelsey Allen, Josh Tenenbaum, and J Zico Kolter.  End-to-end differentiable physics for learning and control.  In _Advances in Neural Information Processing Systems_ , pp. 7178–7189, 2018.
  * Eivazi et al. [2020] Hamidreza Eivazi, Hadi Veisi, Mohammad Hossein Naderi, and Vahid Esfahanian.  Deep neural networks for nonlinear model order reduction of unsteady flows.  _Physics of Fluids_ , 32(10):105104, 2020.
  * Eivazi et al. [2021] Hamidreza Eivazi, Luca Guastoni, Philipp Schlatter, Hossein Azizpour, and Ricardo Vinuesa.  Recurrent neural networks and koopman-based frameworks for temporal predictions in a low-order model of turbulence.  _International Journal of Heat and Fluid Flow_ , 90:108816, 2021.
  * Fresca & Manzoni [2022] Stefania Fresca and Andrea Manzoni.  Pod-dl-rom: enhancing deep learning-based reduced order models for nonlinear parametrized pdes by proper orthogonal decomposition.  _Computer Methods in Applied Mechanics and Engineering_ , 388:114181, 2022.
  * Fukami et al. [2020] Kai Fukami, Koji Fukagata, and Kunihiko Taira.  Assessment of supervised machine learning methods for fluid flows.  _Theoretical and Computational Fluid Dynamics_ , 34(4):497–519, 2020.
  * Gao et al. [2021] Han Gao, Luning Sun, and Jian-Xun Wang.  Phygeonet: physics-informed geometry-adaptive convolutional neural networks for solving parameterized steady-state pdes on irregular domain.  _Journal of Computational Physics_ , 428:110079, 2021.
  * Gao & Ji [2019] Hongyang Gao and Shuiwang Ji.  Graph u-nets.  In _international conference on machine learning_ , pp. 2083–2092. PMLR, 2019.
  * Geneva & Zabaras [2020] Nicholas Geneva and Nicholas Zabaras.  Transformers for modeling physical systems.  _arXiv preprint arXiv:2010.03957_ , 2020.
  * Hartmann et al. [2020] Dirk Hartmann, Christian Lessig, Nils Margenberg, and Thomas Richter.  A neural network multigrid solver for the navier-stokes equations.  _arXiv preprint arXiv:2008.11520_ , 2020.
  * Hasegawa et al. [2020] Kazuto Hasegawa, Kai Fukami, Takaaki Murata, and Koji Fukagata.  Machine-learning-based reduced-order modeling for unsteady flows around bluff bodies of various shapes.  _Theoretical and Computational Fluid Dynamics_ , 34(4):367–383, 2020.
  * Hutchinson et al. [2021] Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim.  Lietransformer: equivariant self-attention for lie groups.  In _International Conference on Machine Learning_ , pp. 4533–4543. PMLR, 2021.
  * Jacquier et al. [2021] Pierre Jacquier, Azzedine Abdedou, Vincent Delmas, and Azzeddine Soulaïmani.  Non-intrusive reduced-order modeling using uncertainty-aware deep neural networks and proper orthogonal decomposition: Application to flood modeling.  _Journal of Computational Physics_ , 424:109854, 2021.
  * Jasak et al. [2007] Hrvoje Jasak, Aleksandar Jemcov, Zeljko Tukovic, et al.  Openfoam: A c++ library for complex physics simulations.  In _International workshop on coupled methods in numerical dynamics_ , volume 1000, pp. 1–20. IUC Dubrovnik Croatia, 2007.
  * Kanwar et al. [2020] Gurtej Kanwar, Michael S. Albergo, Denis Boyda, Kyle Cranmer, Daniel C. Hackett, Sébastien Racanière, Danilo Jimenez Rezende, and Phiala E. Shanahan.  Equivariant flow-based sampling for lattice gauge theory.  _Phys. Rev. Lett._ , 125:121601, Sep 2020.  doi: 10.1103/PhysRevLett.125.121601.
  * Kochkov et al. [2021] Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang, Michael P. Brenner, and Stephan Hoyer.  Machine learning–accelerated computational fluid dynamics.  _Proceedings of the National Academy of Sciences_ , 118(21), 2021.
  * Li et al. [2020] Angran Li, Ruijia Chen, Amir Barati Farimani, and Yongjie Jessica Zhang.  Reaction diffusion system prediction based on convolutional neural network.  _Scientific reports_ , 10(1):1–9, 2020.
  * Li et al. [2019] Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B. Tenenbaum, and Antonio Torralba.  Learning particle dynamics for manipulating rigid bodies, deformable objects, and fluids.  In _International Conference on Learning Representations_ , 2019.
  * Lusch et al. [2018] Bethany Lusch, J Nathan Kutz, and Steven L Brunton.  Deep learning for universal linear embeddings of nonlinear dynamics.  _Nature communications_ , 9(1):1–10, 2018.
  * Marcantoni et al. [2012] Luis F Gutiérrez Marcantoni, José P Tamagno, and Sergio A Elaskar.  High speed flow simulation using openfoam.  _Mecánica Computacional_ , 31(16):2939–2959, 2012.
  * Maulik et al. [2021] Romit Maulik, Bethany Lusch, and Prasanna Balaprakash.  Reduced-order modeling of advection-dominated systems with recurrent neural networks and convolutional autoencoders.  _Physics of Fluids_ , 33(3):037106, 2021.
  * Morimoto et al. [2021] Masaki Morimoto, Kai Fukami, Kai Zhang, Aditya G Nair, and Koji Fukagata.  Convolutional neural networks for fluid flow analysis: toward effective metamodeling and low-dimensionalization.  _arXiv preprint arXiv:2101.02535_ , 2021.
  * Moukalled et al. [2016] Fadl Moukalled, L Mangani, Marwan Darwish, et al.  _The finite volume method in computational fluid dynamics_ , volume 113.  Springer, 2016.
  * Murata et al. [2020] Takaaki Murata, Kai Fukami, and Koji Fukagata.  Nonlinear mode decomposition with convolutional neural networks for fluid dynamics.  _Journal of Fluid Mechanics_ , 882, 2020.
  * Parmar et al. [2018] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran.  Image transformer.  In _International Conference on Machine Learning_ , pp. 4055–4064. PMLR, 2018.
  * Pawar et al. [2019] Suraj Pawar, SM Rahman, H Vaddireddy, Omer San, Adil Rasheed, and Prakash Vedula.  A deep learning enabler for nonintrusive reduced order modeling of fluid flows.  _Physics of Fluids_ , 31(8):085101, 2019.
  * Peng et al. [2020a] Jiang-Zhou Peng, Siheng Chen, Nadine Aubry, Zhi-Hua Chen, and Wei-Tao Wu.  Time-variant prediction of flow over an airfoil using deep neural network.  _Physics of Fluids_ , 32(12):123602, 2020a.
  * Peng et al. [2020b] Jiang-Zhou Peng, Siheng Chen, Nadine Aubry, Zhihua Chen, and Wei-Tao Wu.  Unsteady reduced-order model of flow over cylinders based on convolutional and deconvolutional neural network structure.  _Physics of Fluids_ , 32(12):123609, 2020b.
  * Pfaff et al. [2021] Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter W Battaglia.  Learning mesh-based simulation with graph networks.  In _International Conference on Learning Representations_ , 2021.
  * Radford et al. [2019] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.  Language models are unsupervised multitask learners.  2019\.
  * Rasp & Thuerey [2021] Stephan Rasp and Nils Thuerey.  Data-driven medium-range weather prediction with a resnet pretrained on climate simulations: A new model for weatherbench.  _Journal of Advances in Modeling Earth Systems_ , 13(2):e2020MS002405, 2021.
  * Sanchez-Gonzalez et al. [2018] Alvaro Sanchez-Gonzalez, Nicolas Heess, Jost Tobias Springenberg, Josh Merel, Martin Riedmiller, Raia Hadsell, and Peter Battaglia.  Graph networks as learnable physics engines for inference and control.  In _International Conference on Machine Learning_ , pp. 4470–4479. PMLR, 2018.
  * Sanchez-Gonzalez et al. [2020] Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia.  Learning to simulate complex physics with graph networks.  In _International Conference on Machine Learning_ , pp. 8459–8468. PMLR, 2020.
  * Thuerey et al. [2020] Nils Thuerey, Konstantin Weißenow, Lukas Prantl, and Xiangyu Hu.  Deep learning methods for reynolds-averaged navier–stokes simulations of airfoil flows.  _AIAA Journal_ , 58(1):25–36, 2020.
  * Um et al. [2018] Kiwon Um, Xiangyu Hu, and Nils Thuerey.  Liquid splash modeling with neural networks.  In _Computer Graphics Forum_ , volume 37, pp. 171–182. Wiley Online Library, 2018.
  * Um et al. [2020] Kiwon Um, Robert Brand, Philipp Holl, Nils Thuerey, et al.  Solver-in-the-loop: Learning from differentiable physics to interact with iterative pde-solvers.  _arXiv preprint arXiv:2007.00016_ , 2020.
  * Ummenhofer et al. [2020] Benjamin Ummenhofer, Lukas Prantl, Nils Thürey, and Vladlen Koltun.  Lagrangian fluid simulation with continuous convolutions.  In _International Conference on Learning Representations_ , 2020.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In _Advances in neural information processing systems_ , pp. 5998–6008, 2017.
  * Vlachas et al. [2018] Pantelis R Vlachas, Wonmin Byeon, Zhong Y Wan, Themistoklis P Sapsis, and Petros Koumoutsakos.  Data-driven forecasting of high-dimensional chaotic systems with long short-term memory networks.  _Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences_ , 474(2213):20170844, 2018.
  * Wang et al. [2020] Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu.  Towards physics-informed deep learning for turbulent flow prediction, 2020\.
  * Wesseling [1995] Pieter Wesseling.  Introduction to multigrid methods.  Technical report, INSTITUTE FOR COMPUTER APPLICATIONS IN SCIENCE AND ENGINEERING HAMPTON VA, 1995.
  * Wu et al. [2020] Pin Wu, Junwu Sun, Xuting Chang, Wenjie Zhang, Rossella Arcucci, Yike Guo, and Christopher C Pain.  Data-driven reduced order model with temporal convolutional neural network.  _Computer Methods in Applied Mechanics and Engineering_ , 360:112766, 2020.
  * Xie et al. [2019] Xuping Xie, Guannan Zhang, and Clayton G Webster.  Non-intrusive inference reduced order model for fluids using deep multistep neural network.  _Mathematics_ , 7(8):757, 2019.
  * Xingjian et al. [2015] SHI Xingjian, Zhourong Chen, Hao Wang, Dit-Yan Yeung, Wai-Kin Wong, and Wang-chun Woo.  Convolutional lstm network: A machine learning approach for precipitation nowcasting.  In _Advances in neural information processing systems_ , pp. 802–810, 2015.
  * Xu & Duraisamy [2020] Jiayang Xu and Karthik Duraisamy.  Multi-level convolutional autoencoder networks for parametric prediction of spatio-temporal dynamics.  _Computer Methods in Applied Mechanics and Engineering_ , 372:113379, 2020.
  * Xu et al. [2021] Jiayang Xu, Aniruddhe Pradhan, and Karthik Duraisamy.  Conditionally parameterized, discretization-aware neural networks for mesh-based modeling of physical systems.  _arXiv preprint arXiv:2109.09510_ , 2021.
  * Ying et al. [2018] Rex Ying, Jiaxuan You, Christopher Morris, Xiang Ren, William L Hamilton, and Jure Leskovec.  Hierarchical graph representation learning with differentiable pooling.  _arXiv preprint arXiv:1806.08804_ , 2018.
  * Zhang et al. [2020] Ruiyang Zhang, Yang Liu, and Hao Sun.  Physics-informed multi-lstm networks for metamodeling of nonlinear structures.  _Computer Methods in Applied Mechanics and Engineering_ , 369:113226, 2020.

##  Appendix A Appendix

###  A.1 Information flow in GMR and GMUS

Figure 8: Schematic of information flow in the Graph Mesh Reducer (GMR) and
Graph Mesh Up-Sampling (GMUS). Red and blue nodes are pivotal nodes. The node
size increase represents information aggregation from surroundings.

Figure 8 shows how the information flows in Graph Mesh Reducer (GMR) and Graph
Mesh Up-Sampling (GMUS). This is analogous to a V-cycle in multigrid methods:
Consider a graph with seven nodes ($|\mathcal{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.

###  A.2 Computational complexity

The computation of the proposed model is from the three components: GMR, GMUS,
and the transformer. The graph ${\mathcal{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\cdot(d^{\prime})^{2})$ to make a prediction: a
prediction needs to consider $O(T)$ previous steps, and the computation over a
time step is $O((d^{\prime})^{2})$. Note that $d^{\prime}$ 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^{\prime}$ 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^{\prime}$ would be $Nd$, and the total complexity would be
$O(TN^{2}d^{2})$, which would not be feasible for typical computation servers.

###  A.3 Model Training

The GMR/GMUS and temporal attention model are trained separately. We first
train GMR/GMUS by minimizing the reconstruction loss $\mathcal{L}_{graph}$
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 $\mathcal{L}_{att}$ (Eqn. 12) with
early stopping. The parameter token encoder $\text{mlp}_{p}$ is trained
jointly with temporal attention. The training algorithm is detailed in
Algorithm 1.

Algorithm 1 Training Process

Input: Domain graph $\mathcal{G}$, Node features over time
$[\bm{Y}_{0},\ldots,\bm{Y_{T}}]$, $\mathrm{GMR}_{\theta}$,
$\mathrm{GMUS}_{\phi}$,$\mathrm{Attention}_{\omega}$, RMSE threshold
$R^{\prime}$, $\mathrm{MLP}_{\psi}$, condition parameter $\mu$

Output: Learned parameters $\theta$, $\phi$, $\omega$ and $\psi$

repeat

for $\bm{Y}_{t}\in[\bm{Y}_{0},\ldots,\bm{Y}_{T}]$ do

$\hat{\bm{Y}}_{t}=\mathrm{GMUS}_{\phi}(\mathrm{GMR}_{\theta}(\bm{Y}_{t},\mathcal{G}))$

end for

Compute
$\nabla_{\theta,\phi}\leftarrow\nabla_{\theta,\phi}\mathcal{L}_{graph}(\theta,\phi,[\bm{Y}_{0},\ldots,\bm{Y}_{T}],[\hat{\bm{Y}}_{0},\ldots,\hat{\bm{Y}}_{T}])$

Update $\phi$, $\theta$ using the gradients $\nabla_{\phi}$, $\nabla_{\theta}$

until convergence of the parameters ($\theta$, $\phi$)

$[z_{0},z_{1},\ldots,z_{T}]=\mathrm{GMR}_{\theta}([\bm{Y}_{0},\ldots,\bm{Y}_{T}])$,
$z_{-1}=\text{MLP}_{\psi}(\mu)$

for $t\in[1,2,\ldots,T]$ do

repeat

$[\hat{z}_{1},\ldots,\hat{z}_{t}]=\mathrm{Attention}_{\omega}([z_{-1},z_{0}])$

Compute
$\nabla_{\omega,\psi}\leftarrow\nabla_{\omega,\psi}\mathcal{L}_{att}(\omega,\psi,[\bm{z}_{1},\ldots,\bm{z}_{t}],[\hat{\bm{z}}_{1},\ldots,\hat{\bm{z}}_{t}])$

Update $\omega$, $\psi$ using the gradients $\nabla_{\omega}$, $\nabla_{\psi}$

until
$\mathrm{RMSE}([\bm{z}_{1},\ldots,\bm{z}_{t}],[\hat{\bm{z}}_{1},\ldots,\hat{\bm{z}}_{t}])<R^{\prime}$

end for

###  A.4 Dataset details

Figure 9: Overview of datasets of (a) cylinder flow, (b) sonic flow, and (c)
vascular flow.

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.

Table 2: Simulation details of datasets Dataset | PDEs | Cell type | Meshing | # nodes | # steps |  $\delta t$ (sec)
---|---|---|---|---|---|---
Cylinder flow | Incompr. NS | tri | Fixed | $1699$ | $400$ | $0.5$
Sonic flow | Compr. NS | tri | Moving | $1900$ | $40$ | $5\times 10^{-7}$
Vascular flow | Incompr. NS | tri | Small-varying |  $7561$ (avg.) | $250$ | $0.5$
Table 3: Input-output information of our networks: $u,v,p,T,\rho,m$ denote X-axis velocity, Y-axis velocity, pressure, temperature, density and cell volume, respectively. $\mathbf{x}$ denote the cell coordinates. $T_{,0},Re,r$ denote the initial temperature (sonic flow), Reynolds number (cylinder flow) and radius of the thrombus (vascular flow).  Dataset |  GMR node input |  GMR edge input |  GMUS node output |  Parameter token model input |  Nodal embed dim |  # selected node
---|---|---|---|---|---|---
Cylinder flow | $u_{i},v_{i},p_{i},m_{i},Re$ |  $\bm{x}_{i}-\bm{x}_{j}$, $|\bm{x}_{i}-\bm{x}_{j}|$ | $u_{i},v_{i},p_{i}$ | $Re$ | $4$ | $256$
Sonic flow |  $u_{i},v_{i},p_{i},T_{i},\rho_{i}$ $m_{i},T_{,0}$ |  $\bm{x}_{i}-\bm{x}_{j}$, $|\bm{x}_{i}-\bm{x}_{j}|$ |  $u_{i},v_{i},p_{i}$ $T_{i},\rho_{i}$ | $T_{,0}$ | $4$ | $256$
Vascular flow | $u_{i},v_{i},p_{i},m_{i},r$ |  $\bm{x}_{i}-\bm{x}_{j}$, $|\bm{x}_{i}-\bm{x}_{j}|$ | $u_{i},v_{i},p_{i}$ | $r$ | $2$ | $400$

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 $\mu$ 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 $\mathbf{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,

$\begin{split}\frac{\partial\rho}{\partial t}+\nabla\cdot(\rho\bm{v})=0,\\\
\frac{\partial(\rho\bm{v})}{\partial
t}+\nabla\cdot(\rho\bm{v}\bm{v})=\nabla\cdot(\mu\nabla\bm{v})-\nabla
p+\nabla\cdot(\mu(\nabla\bm{v})^{T})-\frac{2}{3}\nabla(\mu\nabla\cdot\bm{v}),\\\
\frac{\partial}{\partial t}(\rho c_{p}T)+\nabla\cdot(\rho
c_{p}\bm{v}T)=\nabla\cdot(k\nabla T)+\rho
T\frac{Dc_{p}}{Dt}+\frac{Dp}{Dt}-\frac{2}{3}\mu\Psi+\mu\Phi,\end{split}$ (14)

where $\bm{v}$ is the velocity vector, $\mu$ is the viscosity, $k$ is the
thermal conductivity, $c_{p}$ is the specific heat of the fluid and the
definitions of the stream scalar $\Psi$, the potential scalar $\Phi$ 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.

###  A.5 Model details

Node and edge representations in our GraphNets are vectors of width 128. The
node and edge functions ($\text{mlp}_{v}$, $\text{mlp}_{e}$, $\text{mlp}_{r}$)
are MLPs with two hidden layers of size 128, and ReLU activation. The input
dimension of $\text{mlp}_{v}$ is based on the number of input features in each
node (see details in table 3), and the input dimension of $\text{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 $\text{mlp}_{\ell}^{e}$ and
$\text{mlp}_{\ell}^{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 $\bm{z}$ for each dataset (cylinder flow, sonic flow, and
vascular flow) are $1024$, $1024$, and $800$, respectively. This is calculated
by $\\#\text{pivotal nodes}\times n_{out}$, in which $n_{out}$ is the output
size of the last GraphNet block. The $\text{mlp}_{p}$ used to encode the
system parameter $\mu$ is a MLP with 2 hidden layers of size 100, and with
output dimensions that match $\mathbf{z}$, i.e. 1024, 1024, and 800 for
Cylinder Flow, Sonic Flow, and Vascular Flow, respectively.

Table 4: Number of parameters for each model  Dataset | MeshGraphNet | GRU | LSTM | Transformer | GMR-GMU
---|---|---|---|---|---
Cylinder flow | 2.2M | 9.4M | 11.5M | 14.1M | 1.2M
Sonic flow | 2.2M | 9.4M | 11.5M | 14.1M | 1.2M
Vascular flow | 2.2M | 5.8M | 7.0M | 8.5M | 1.2M

###  A.6 Pivotal 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 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.

(a)

(b)

Figure 10: Flow past cylinder flow system: (a) all cells in the FV mesh; (b)
pivotal nodes

(a)

(b)

(c)

(d)

(e)

(f)

Figure 11: Sonic flow system: (a-c) all cells in the FV mesh for 3 different
time steps; (d-f) pivotal nodes for 3 different time steps.

(a)

(b)

(c)

(d)

(e)

(f)

Figure 12: Vascular flow system: all cells in the FV mesh (left), pivotal
nodes (right), $r=0.3$ (a-b), $r=0.4$ (c-d), $r=0.5$ (e-f).

###  A.7 Data Compression

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 $(\mathcal{G},\bm{Y})$,
we can encode it into a latent vector $\bm{z}=\text{GMR}(\mathcal{G},\bm{Y})$
through GMR. Therefore, a high-dimensional spatiotemporal trajectory
$(\mathcal{G},\bm{Y}_{0},\ldots,\bm{Y}_{T})$ can be compressed into a low-
dimensional representation $(\mathcal{G},\bm{z}_{0},\ldots,\bm{z}_{T})$, 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.

Table 5: Report of data compressing: QoI data represent original data of the variables of interest; Embedding represents the saved embedding vectors from the GMR (encoder); Decoder is the saved GMUS model which is used to decode the embeddings; Compression Ratio is the ratio of QoI data size to the sum of the embedding and decoder sizes. The RMSE is the relative mean square error of the reconstructed data from embedding. Dataset | QoI Data | Embedding | Decoder | Compression Ratio | RMSE ($\times 10^{-3}$)
---|---|---|---|---|---
Cylinder flow | 3570 MB |  $82.1$MB |  $12.8$MB | $\bm{37}$ | $14.3$
Sonic flow |  $651$MB |  $8.5$MB |  $18$ MB | $\bm{25}$ | $1.11$
Vascular flow |  $1951$MB |  $8.5$MB |  $49.3$ MB | $\bm{33}$ | $10$

###  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 sub-timestepping. 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,

$\begin{split}&\text{Step
1:}\quad\hat{\bm{z}}_{1}=\mathrm{trans}(\bm{z}_{-1},\bm{z}_{0}),\quad\hat{\bm{z}}_{t}=\mathrm{trans}(\bm{z}_{-1},\bm{z}_{0},\hat{\bm{z}}_{1},\ldots,\hat{\bm{z}}_{t-1})\\\
&\text{Step 2:}\quad\tilde{\bm{Y}}_{t}=\mathrm{GMUS}(\tilde{\bm{z}}_{t}),\quad
t=1,\ldots,T\end{split}$ (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.

Table 6: Wall-clock time of finite volume (FV) model and attention-GMUS to simulate a trajectory for three flow cases. Dataset | GMUS | Transformer | GT Model | Speed-up
---|---|---|---|---
Cylinder flow (for 400 time steps) |  $2$sec |  $0.4758$sec |  $1688.59$sec | $\bm{682}$
Sonic flow (for 40 time steps) |  $0.2$sec |  $0.047$sec |  $25.026$sec | $\bm{100}$
Vascular flow (for 250 time steps) |  $2.63$sec |  $0.31$sec |  $2451$sec | $\bm{800}$

###  A.9 Other variable contour

Figure 13 shows the pressure contours.

|  $Re=307$ (Cylinder Flow) |  $Re=993$ (Cylinder Flow)
---|---|---
| t=0 | t=85 | t= 215 | t=298 | t=0 | t=85 | t= 215 | t=298
Truth |  |  |  |  |  |  |  |
Ours |  |  |  |  |  |  |  |
|  $T(0)=201$ (Sonic Flow) |  $T(0)=299$ (Sonic Flow)
| t=0 | t=15 | t= 30 | t=40 | t=0 | t=15 | t= 30 | t=40
Truth |  |  |  |  |  |  |  |
Ours |  |  |  |  |  |  |  |
|  $R=0.31$ (Vascular flow) |  $R=0.49$ (Vascular flow)
| t=40 | t=80 | t= 160 | t=250 | t=40 | t=80 | t= 160 | t=250
Truth |  |  |  |  |  |  |  |
Ours |  |  |  |  |  |  |  |
Figure 13: The pressure contours of rollouts predicted by our model versus the
ground truth (CFD). With different condition settings, the model accurately
predict long sequences of rollouts.

###  A.10 Further analysis of latent encodings and attention

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
$\mathbf{z}_{-1}$ also rapidly decay, demonstrating that the transformer can
adjust the attention distribution to capture most physically useful signals.

(a)  (b)  (c)

Figure 14: 2-D principle subspace of the embedded vectors from GMR (a) and PCA
(b) for the high-speed sonic flow system: $T(0)=201$ trajectory
(LABEL:ls:LowReTraj), $T(0)=299$ trajectory (LABEL:ls:HighReTraj) start from
their parameter tokens LABEL:ls:LowReTrajStart and LABEL:ls:HighReTrajStart,
respectively. (c) Attention values of the parameter tokens versus time,
$T(0)=201$ (LABEL:ls:LowReFirstAttention), and $T(0)=299$
(LABEL:ls:HighReFirstAttention).

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).

(a)  (b)  (c)

Figure 15: Examples of attention value (log value contour) for (a) flow past
cylinder, (b) sonic flow and (c) vascular flow.

###  A.11 Longer 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 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
$\tilde{\bm{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\times 400$ steps). We find that the proposed
model still outperforms the baseline for all $u,v$ and $p$ trajectories.

Figure 16: Sliding window for attention mechanism on latent representations. The parameter embedding token $\tilde{\bm{z}}_{-1}$ is always included in the first entry. Table 7: The average relative rollout error for the cylinder flow, with unit of $\times 10^{-3}$. We compare MeshGraphNet with noise injection (NI) to our model with transformer. We train the model on the training trajectories with 400 steps, and test it on the unseen trajectories with 800 steps. rollout step | 800
---|---
Variable | $u$ | $v$ | $p$
MeshGraphNet-NI | 43 | 1274 | 258
Ours-Transformer | 6 | 158 | 48

###  A.12 Prediction 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 flow, governed by

$\frac{\partial c}{\partial
t}+\nabla\cdot(\bm{v}c)=\nabla^{2}(\frac{|\bm{v}|}{Pe}c),$ (16)

where $c$ is the concentration of the transport scalar, $\bm{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. As shown in
Figure 17, the model can accurately predict the non-oscillatory sequences of
rollouts Quantitatively, our model ($\text{RMSE}=1.03\times 10^{-3}$) still
outperforms MeshGraphNet ($\text{RMSE}=2.79\times 10^{-3}$).

| Convection Flow | Diffusion Flow
---|---|---
| t=0 | t=25 | t= 50 | t=100 | t=0 | t=2 | t= 10 | t=100
Truth |  |  |  |  |  |  |  |
Ours |  |  |  |  |  |  |  |
Figure 17: The concentration contours of rollouts predicted by our model
versus the ground truth (CFD).

###  A.13 Cardiac flow prediction

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

$\begin{split}v(x,t)=\frac{0.6(x-x_{\text{min}})(x-x_{\text{max}})}{(x_{\text{max}}-x_{\text{min}})^{2}}(\sin(\frac{2\pi
t}{1.8}))^{2}+0.15\end{split}$ (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 $\times 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
Truth |  |  |  |
Ours |  |  |  |
Figure 18: The velocity contours of rollouts predicted by our model versus the
ground truth (CFD) for a cardiac flow.