SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Luning Sun*
Lawrence Livermore National Lab
Livermore, CA 94550
sun42@11nl.gov
Xu Han
xu.han@tufts.edu
Han Gao
hgao1@seas.harvard.edu
Jian-Xun Wang
Notre Dame, IN 46556
jwang33@nd.edu
Accurate prediction of dynamical systems in unstructured meshes has recently shown successes in scientific simulations. Many dynamical systems have a nonnegligible level of stochasticity introduced by various factors (e.g. chaoticity), so there is a need for a unified framework that captures both deterministic and stochastic components in the rollouts of these systems. Inspired by regeneration learning, we propose a new model that combines generative and sequential networks to model dynamical systems. Specifically, we use an autoencoder to learn compact representations of full-space physical variables in a low-dimensional space. We then integrate a transformer with a conditional normalizing flow model to model the temporal sequence of latent representations. We evaluate the new model in both deterministic and stochastic systems. The model outperforms several competitive baseline models and makes more accurate predictions of deterministic systems. Its own prediction error is also reflected in its uncertainty estimations. When predicting stochastic systems, the proposed model generates high-quality rollout samples. The mean and variance of these samples well match the statistics of samples computed from expensive numerical simulations.
Accurate prediction of long-term dynamics of physical systems is of great interest in many science and engineering fields. The classical simulation of a system heavily relies on a spatial/temporal discretization of the space and the numerical solution to a finite-dimensional algebraic system derived from the governing equation [1]. However, due to the multi-scale, stochastic nature of the complex physics and the complexity of the geometry, the simulation of large-scale real-time applications is extremely expensive in computation [2]. In recent years, deep learning models have been applied to predict rollouts of large and complex physical systems thanks to their flexibility and scalability [3]. Moreover, they can handle uncertainties and non-linearity in physical problems [4-8] more effectively than traditional methods.
Previous work has studied two types of complex dynamical systems: deterministic ones and stochastic ones. A deterministic system is often considered under a perfectly controlled experiment with exactly
known PDE terms and initial conditions (IC) or boundary conditions (BC) [9, 10, 3, 11-13]. This type of system can be modeled by autoregressive predicting models. A stochastic system such as quantum mechanics [14] and statistical physics [15] has a stochastic rollouts. When there are stochastic forcing terms or IC/BC terms in the governing equations, researchers have developed various models for predicting the stochastic state variables, such as turbulence velocity and stock prices [16-20]. Given that there is not a clear boundary between the two types of systems, it is highly desirable that a unified model can model either a deterministic or a stochastic system automatically. However, such models so far are limited to classical numerical models such as OpenFOAM [21]. There is an urgent need to develop a unified deep-learning model that can model both types of systems.
In this work, we propose a unified framework based on deep generative models to predict deterministic and stochastic systems. The new model is based on graph-structured state representations [11], which can handle irregular spatial areas commonly seen in dynamical systems. It uses an autoencoder to encode a state representation to a low-dimensional latent vector. Accurate encoding and decoding is critical for recovering the state of the dynamical system. This work makes several innovations to enhance the autoencoder's ability to preserve the information in system states. We provide a new approach to encoding the graph-structured representation with a vector with a fixed length. We also get some inspiration from regeneration learning [22] and train our autoencoder with self-supervised learning.
To describe stochasticity in the system, we use a sequential probabilistic model for the latent representations. We integrate a transformer and a normalizing flow to construct a step-wise predictive neural network in the latent space. When there is stochasticity, the model will learn the conditional distribution of the next latent state; when there is no stochasticity, it can place probabilities to correct deterministic predictions and still can minimize the predictive error. By including the ability to simulate both deterministic and stochastic systems in a uniform framework, it reduces the effort of developing separate models for different problems.
We evaluate our proposed framework in an extensive empirical study. The results indicate the proposed model outperforms the SOTA baselines on deterministic datasets regarding accuracy. More importantly, for the first time, we introduce several alternative evaluation metrics other than normalized RMSE for stochastic fluid dynamics, which help improve comparisons between different methods in this domain. The proposed framework can produce high-quality samples for stochastic systems in the mesh space.
Let's consider a general partial differential equation (PDE) defined on the  $d$ -dimensional space and one-dimensional time domain,
$$
\frac {\partial \boldsymbol {u}}{\partial t} = j (\boldsymbol {u}, \boldsymbol {\mu}, \iota) \quad \text {i n} \Omega \times [ 0, T _ {\text {e n d}} ] \tag {1}
$$
where  $\Omega \subset \mathbb{R}^d$  is the spatial domain,  $T_{\mathrm{end}}$  is the endpoint of time,  $\iota : \partial \Omega \times [0,T]$  is the random parameter over time for stochastic systems (e.g., boundary conditions), and  $\pmb{u}:\Omega \times [0,T]$  is the primary solution variable (e.g., velocity and pressure of fluid flow). Here  $\pmb{\mu}:\mathcal{D}$  is the global physical system parameter and is time-invariant (e.g., Re number) for a given rollout.  $j:\Omega \times [0,T]\times \mathcal{D}\to \Omega \times [0,T]$  is an aggregation of the spatial terms of conservation laws (e.g., source and flux). To numerically solve the conservation law in Equation 1, let  $\mathcal{C}_h$  be a mesh of  $\Omega$ , that is,  $\mathcal{C}_h = \{C_i\in \Omega :i = 1,\dots ,N\}$  is a collection of non-overlapping cells that cover  $\Omega$ . We further use  $\pmb{u}_{i,t}$  to denote the evaluation of  $\pmb{u}$  at the cell center of  $C_i$  at time step  $t$ .
Using the mesh described above, we apply a finite volume discretization to yield the parametrized, nonlinear dynamical system. Here, the dynamical system can be computationally intensive because a fine-level mesh will lead to discretization with a large degree of freedom. For stability issues, the numerical time step also needs to be very small. Therefore, numerical simulations with traditional methods require extremely expensive computation and storage.
In this paper, we are interested in two different scenarios: deterministic dynamics with invariant parameter  $\iota$  over time (e.g., fixed boundary condition), and stochastic dynamics with random
parameter  $\iota$  over time (e.g., perturbations in boundary conditions of turbulent flow). Both cases will also include the global physical parameters  $\mu$ . The scope of our paper is to build a unified data-driven surrogate model for generating/predicting parametric PDE solutions.
2.2 Normalizing flow models
A normalizing flow model constructs a flexible probabilistic distribution by applying a learnable bijective mapping to a simple random variable (e.g. Gaussian distributed). Suppose the mapping is  $z = f(x)$  with  $x \in \mathbb{R}^d$  being the simple input variable, then we have the probability  $p(z)$  of  $z$  as follows:
$$
p (\boldsymbol {z}) = p (\boldsymbol {x}) \left| \det  \left(\frac {\partial (f (\boldsymbol {x}))}{\partial (\boldsymbol {x})}\right) \right| ^ {- 1}. \tag {2}
$$
Here  $\frac{\partial(f(x))}{\partial(x)}$  is the Jacobian of  $f$  at  $x$ . The function  $f$  is usually a neural network with a layered structure, with each layer being a bijective mapping. Then the determinant of  $\frac{\partial(f(x))}{\partial(x)}$  is the product of determinants of these layers' Jacobian matrices. With a special design of layer structures, their determinants can be efficiently computed. For example, a RealNVP model [23] is constructed by stacking several coupling layers, each of which has a lower triangular Jacobian matrix. With  $h \in \mathbb{R}^d$  as the input, a coupling layer  $f_\ell$  runs the following calculation:
$$
f _ {\ell} (\boldsymbol {h}) = \operatorname {c o n c a t} \left(\boldsymbol {h} \left[ 1: d ^ {\prime} \right], \boldsymbol {h} \left[ \left(d ^ {\prime} + 1\right): d \right] \odot \exp \left(s \left(\boldsymbol {h} \left[ 1: d ^ {\prime} \right]\right)\right) + t \left(\boldsymbol {h} \left[ 1: d ^ {\prime} \right]\right)\right) \tag {3}
$$
Here  $\text{concat}$  concatenates its arguments as one vector, and  $d'$  is usually about one-half of  $d$ . In this calculation, the first half of the vector is directly copied to the output. The second half goes through an entry-wise linear transformation: the operation  $\odot$  is the Hadamard product, the neural network  $s$  provides scaling coefficients, and the neural network  $t$  provides biases. With  $K$  such coupling layers, we have a transformation that defines a flexible distribution  $p(z)$ .
$$
p (\boldsymbol {z}) = p (\boldsymbol {x}) \cdot \prod_ {\ell = 1} ^ {K} | \det  \left(\partial f _ {\ell} \left(\boldsymbol {h} ^ {\ell - 1}\right) / \partial \boldsymbol {h} ^ {\ell - 1}\right) | ^ {- 1} \tag {4}
$$
Here  $\pmb{h}^0 = \pmb{x}$ , and  $\pmb{h}^\ell = f_\ell(\pmb{h}_{\ell-1})$ . At the same time,  $p(z)$  also has an efficient sampling procedure given by  $f$ .
Our new predicting model has three components: an encoder that compresses a graph representation of the state into a fixed-length vector, a sequential model that predicts next-step representations in the latent space, and a decoder that decodes spatial states from graph representations. The encoder is an improved version of the GMR-GMUS encoder [11]. The sequential model is a conditional flow model. The encoder and decoder are trained via self-supervised learning as in regeneration learning.
3.1 Graph representation learning
Following GMR-GMUS, we use a graph  $\mathcal{G} = (\mathcal{V},\mathcal{E})$  to encode a snapshot of a deterministic or stochastic dynamical system at time step  $t$ . Here each node  $i\in \mathcal{V}$  corresponds to the mesh cell  $C_i$ , and each edge  $(i,j)\in E$  represents a neighboring relationship between two cells. At each step  $t$ , the solution  $\pmb{u}_{i,t}$  at cell  $i$  becomes the feature vector of the node  $i\in \mathcal{V}$  at time  $t$ . Because the mesh is pre-defined and fixed, the graph representation  $(G,(\pmb{u}_{i,t},i\in V))$  preserves the full information of the mesh  $\mathcal{C}_h$  at time  $t$ . Here we use  $\pmb{Y}_t = (\pmb{u}_{i,t}:i\in V)$  to denote a snapshot at time  $t$ .
The key task for the encoder is to encode the graph into a low-dimensional vector  $\mathbf{z}_t$ . For this task, we use an improved version of the encoder of GMR-GMUS. The GMR-GMUS encoder runs a graph neural network over the graph representation to learn node representations. Then it takes node vectors of a selection of "pivotal" nodes and concatenates them to get  $\mathbf{z}_t$ . In our new encoder, we select a set of locations in the spatial space instead of graph nodes to aggregate spatial information: each selected location aggregates vectors of nearby nodes to get its representation. It decouples the graph representation and the aggregation operation so that a select location can encode nodes within an arbitrary distance. To improve the training stability, we also improve the graph neural network's
Figure 1: The diagram of the proposed model, which first compresses the whole graph into the latent representation  $\mathbf{z}_t$  by PbGMR using selected positions. During the generation process, a transformer encodes physical parameters and previous latent representations into a condition vector  $c_{t+1}$ , from which a normalizing flow model describes the conditional probability of  $\mathbf{z}_{t+1}$ . Finally, the decoder PbGMUS decodes  $\mathbf{z}_{t+1}$  through to obtain the next-step prediction  $\mathbf{Y}_{t+1}$ .
architecture by adding residual connections between its message-passing layers. We call this new encoder Position-based Graph Mesh Reducer (PbGMR). These two modifications clearly improve the encoder's ability to preserve the state information, which will be empirically shown in the experiment section.
The architecture of PbGMR is specified as follows. For notational convenience, we omit the time from the subscript. For each graph node  $i \in V$ , we first extract node and edge features as follows.
$$
\boldsymbol {v} _ {i} ^ {0} = \operatorname {m l p} _ {v} (\boldsymbol {u} _ {i}), \quad \boldsymbol {e} _ {i j} ^ {0} = \operatorname {m l p} _ {e} (\operatorname {p o s} (i) - \operatorname {p o s} (j)). \tag {5}
$$
Here  $\mathrm{pos}(i)$  is the spatial location of the cell center of  $C_i$ . After that, we apply  $L$  message-passing layers.
$$
\boldsymbol {e} _ {i j} ^ {\ell} = \boldsymbol {e} _ {i j} ^ {\ell - 1} + \text {l a y e r n o r m} \left(\operatorname {m l p} _ {\ell} ^ {e} \left(\boldsymbol {e} _ {i j} ^ {\ell - 1}, \boldsymbol {v} _ {i} ^ {\ell - 1}, \boldsymbol {v} _ {j} ^ {\ell - 1}\right)\right) \tag {6}
$$
$$
\boldsymbol {v} _ {i} ^ {\ell} = \boldsymbol {v} _ {i} ^ {\ell - 1} + \text {l a y e r n o r m} \left(\mathrm {m l p} _ {\ell} ^ {v} \left(\boldsymbol {v} _ {i} ^ {\ell - 1}, \sum_ {j \in \mathcal {N} _ {i}} \boldsymbol {e} _ {i j} ^ {\ell - 1}\right)\right), \quad \ell = 1, \dots , L. \tag {7}
$$
Here  $\mathcal{N}_i$  denotes all neighbors of node  $i$ ,  $\mathrm{mlp}_{\ell}^{e}$  and  $\mathrm{mlp}_{\ell}^{v}$  are two separate multi-layer perceptrons, and layernorm is the layer-normalization operation. When a perceptron accepts multiple arguments, its arguments are first concatenated into a single vector. Each layer here is similar to a GraphNet block [24] but with slight differences. The prominent ones are residual connections and layer normalization in the update of node and edge representations: they help to stabilize the training procedure [25-27].
After we have learned node representations, we need to aggregate them into a single vector to represent the entire state. We randomly select a small set  $S$  of centers in the spatial area. For each center  $c \in S$ , we select  $k$  nearest mesh cells  $\mathrm{knn}(c)$  based on spatial distance and then compute the position representation  $h_c$  by interpolation [28]:
$$
\boldsymbol {h} _ {c} = \sum_ {j \in \operatorname {k n n} (c)} \frac {w _ {c j} \boldsymbol {v} _ {j} ^ {L}}{\sum_ {j \in \operatorname {k n n} (c)} w _ {c j}}, \quad w _ {c j} = \frac {1}{d (c , j) ^ {2}},, c \in \mathcal {S} \tag {8}
$$
Here  $d(c, j)$  is the spatial distance between cell  $C_j$  and the center  $c$ . Given that the calculation is scaled by the sum of  $w_{cj}$ -s, the unit of the spatial distance does not change the calculation here. The number  $k$  of neighbors is a hyper-parameter, and we fix it to 10. Finally, we concatenate all representations of centers into a single vector  $z = \operatorname{concat}(h_c : c \in S)$  as the latent for the entire graph. Note that the centers in  $S$  are fixed for a problem.
As a comparison, the encoder in GMR-GMUS uses graph nodes as centers and only considers connected neighbors of a center. Then the number of neighbors in the interpolation operation is
limited by the graph structure. As a result, each center vector from GMR-GMUS can only represent information in a short range. Our new encoder overcomes this issue by decoupling the neighbors in interpolation and the neighbors in the graph representation and then gives the interpolation operation more freedom to represent a state of the system.
3.2 Decoding and self-supervised training
We devise a decoder PbMUS to recover node features on the graph from latent representation  $z$ :  $\hat{Y} = \mathrm{PbGMUS}(z)$ . Here we also consider the computation at a single time step and omit time indices. We first split  $z$  and get vectors at interpolation centers:  $(h_c : c \in S) = z$  and then compute the initial node representation  $r_i^0$  by spatial interpolation from centers:
$$
\boldsymbol {r} _ {i} ^ {0} = \sum_ {j \in \operatorname {k n n} ^ {\prime} (i)} \frac {w _ {i c} \boldsymbol {h} _ {c}}{\sum_ {c \in \operatorname {k n n} ^ {\prime} (i)} w _ {i c}}, \quad w _ {i c} = \frac {1}{d (i , c) ^ {2}},, i \in \mathcal {V}, c \in \mathcal {S} \tag {9}
$$
Here  $\mathrm{knn}^{\prime}(i)$  are  $k$  centers that are nearest to the cell center  $i$ . Then we apply  $L$  message-passing layers to compute  $\hat{\mathbf{Y}} = \mathrm{gnn}(G, (\mathbf{r}_i^0 : i \in V))$ , with gnn representing the  $L$  network layers. These layers have the same architecture as PbGMR but use different learnable parameters.
Self-supervised training. Without considering the sequential property of the data, we first train the encoder and decoder on single steps with self-supervised training. This training method shares the same spirit of regeneration learning and improves encoder and decoder's abilities to capture spatial patterns in the data. In particular, we use the reconstruction error as the minimization objective to train the encoder-decoder pair.
$$
\min  \sum_ {t = 0} ^ {T} \left| \left| \boldsymbol {Y} _ {t} - \operatorname {P b G M U S} (\operatorname {P b G M R} (\boldsymbol {Y} _ {t})) \right| \right| _ {2} ^ {2} \tag {10}
$$
3.3 Attention-based temporal conditioned generative model
Now we consider the sequence of latent representations from the encoder and devise a sequential generative model to model the sequence of latent representations. The low-dimensional latent space reduces the modeling difficulty, and the vector form of latent representations  $\mathbf{z}_t$  avoids the graph structure and allows more model choices. We first decompose the sequence as follows.
$$
P \left(\boldsymbol {z} _ {1: T} \mid \boldsymbol {\mu}, \boldsymbol {z} _ {0}\right) = p \left(\boldsymbol {z} _ {1} \mid \boldsymbol {\mu}, \boldsymbol {z} _ {0}\right) \prod_ {t = 2} ^ {T} p \left(\boldsymbol {z} _ {t} \mid \boldsymbol {\mu}, \boldsymbol {z} _ {0}, \boldsymbol {z} _ {1: t - 1}\right). \tag {11}
$$
The key is to devise a model for the conditional  $p(\pmb{z}_t|\pmb{\mu},\pmb{z}_0,\pmb{z}_{1:t - 1})$ . We take two steps to construct this conditional: we first represent the condition  $(\pmb{\mu},\pmb{z}_0,\pmb{z}_{1:t - 1})$  with a single vector  $\pmb{c}_t$  and then adapt RealNVP [23], a normalizing flow model, to describe the conditional  $p(\pmb{z}_t|\pmb{c}_t)$ .
Calculate the conditional vector with a transformer. Because physical parameters  $\mu$  control the entire system, and the initial condition  $z_0$  contains information about the initial condition, we consider them special and use them in the prediction of  $z_{t}$  for each  $t$ . We structure the problem as a "translation" problem and use the transformer [29] to run the calculation. In our case, the input "sentence" is  $(\mu ,z_0)$ , the first few "tokens" in the target sentence are  $(z_{1},\ldots ,z_{t - 1})$ , and the "next token" to be predicted is  $c_{t}$ .
$$
\boldsymbol {c} _ {t} = \operatorname {t r a n s f o r m e r} \left(\left(\boldsymbol {\mu}, \boldsymbol {z} _ {0}\right), \left(\boldsymbol {z} _ {1}, \dots , \boldsymbol {z} _ {t - 1}\right)\right) \tag {12}
$$
While the computation is exactly the same as one step in a translation task, the rationale is very different. First, the translation task directly gets the  $z_{t}$  from the transformer, but our model needs to send  $c_{t}$  to a conditional flow model to get the prediction  $z_{t}$ . Second, the translation task has an informative input sequence, but our model uses a less informative input and depends on the transformer to get a reasonable output sequence.
Predict  $z_{t}$  with a conditional flow model. Once we have a vector  $c_{t}$  representation of the condition, we can construct a conditional flow model from RealMVP. Specifically, we append the condition vector  $c_{t}$  to the input of the scaling function  $s$  and the bias function  $t$  in equation 3 in each layer.
$$
f _ {\ell} (\boldsymbol {h}) = \left[ \boldsymbol {h} \left[ 1: d ^ {\prime} \right], \boldsymbol {h} \left[ d ^ {\prime} + 1: d \right] \odot \exp \left(s \left(\boldsymbol {h} \left[ 1: d ^ {\prime} \right], \boldsymbol {c} _ {t}\right)\right) + t \left(\boldsymbol {h} \left[ 1: d ^ {\prime} \right], \boldsymbol {c} _ {t}\right) \right] \tag {13}
$$
From these layers, we have the flow model  $p(\boldsymbol{z}_t|\boldsymbol{c}_t)$ . By chaining up the transformer and the conditional flow model, we have our sequential model for conditional probability  $p(\boldsymbol{z}_t|\boldsymbol{\mu},\boldsymbol{z}_0,\boldsymbol{z}_{1:t - 1})$ .
During training, we train the sequential model on  $(z_0,\dots ,z_T)$  that are computed from our PbGMR encoder. During inference, we can efficiently sample  $z_{t}$  from the sequential model. The whole training and inference processes can be found in Appendix A.2 and Appendix A.3.
Deep learning for physics. For a deterministic system, one-step learning simply takes the current state as the input and outputs the next-step prediction [3, 30-32]. To further improve the accuracy of long-term forecasts, dimension reduction is applied together with sequence nets, aiming to solve long dynamical systems on a regular domain [33-35]. Such works often adopt CNNs as encoders so that can not be applied to irregular mesh space. To further work on mesh data directly with dimension reduction, the GNN encoder/decoder was initially introduced by [11]. Various GNN architectures were also proposed to facilitate the learning of physics. [13, 15, 36-39]. For physical stochastic systems, CNN is mainly used for probabilistic prediction [18, 40]. However, it lacks researches solving stochastic system in the graph space.
Regeneration learning and generative modeling. Regeneration learning is a new learning paradigm for data generation, which first generates a latent representation for the data. Then the generation process happens on the latent space. There are many recently popular generative models for images[41-43], videos[44], speeches[45], and music [46] are built on this paradigm. We notice that it should be further investigated for graph generation tasks.
5.1 Deterministic dynamics
We benchmark our proposed model with three datasets from three flows: flows over the cylinder, high-speed flows over the moving edge, and vascular flows [11]. Note that for test cases, we feed the model physical parameters  $\mu$  and the snapshot at the beginning time to generate the whole trajectory. Detailed description can be found in Appendix A.5.
We compare our new method against five SOTA models for fluid dynamics, including two variants of MeshGraphNet[3] and three variants of the GMR-GMUS[11]. We use relative mean square error (RMSE) as the evaluation metric:  $\mathrm{RMSE} = \frac{\sum(\hat{u}_i^{\mathrm{prediction}} - \hat{u}_i^{\mathrm{truth}})^2}{\sum(\hat{u}_i^{\mathrm{prediction}})^2}$ . More details can be found in the Appendix A.2. For cylinder and vascular flows, we calculate the RMSE with respect to velocity variables  $u$ ,  $v$ , and the pressure variable  $p$ . For high-speed flows over the moving edge, we also calculated the RMSE for the temperature  $T$ .
The results on prediction errors. Table 2 shows the comparison of different models. Our model outperforms the baseline models in all the scenarios. For the cylinder flow case, our model has improvements of  $22\%$ ,  $17\%$ , and  $47\%$  for  $u$ ,  $v$ , and  $p$ . For the sonic flow case, our model has improvements of  $61\%$ ,  $70\%$  for the  $u$  and  $v$  variables, and improvements of  $82\%$  and  $50\%$  respectively for the  $p$  and  $T$  variables. For the vascular