SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1: The average relative reconstruction error of three systems, with the unit of  $1 \times 10^{-3}$
Dataset	GMR-GMUS	PbGMR-GMUS
Cylinder flow	14.3	1.9
Sonic flow	1.11	0.24
Vascular flow	10	2.8
flow case, our model achieves improvements of  $52\%$ ,  $45\%$ , and  $95\%$  for  $u$ ,  $v$ , and  $p$ . The ablation
Figure 2: For each case, the first row is one of our predicted samples for velocity. The second row is the predicted standard deviation for velocity. Our model-predicted sample accurately reflects the physical systems. Meanwhile, the predicted standard deviation also has a physical spatial-temporal pattern.
study in the Appendix section (Table 9) indicates that both the new encoder and the conditional normalizing flow are essential to improve prediction accuracy.
The error in the time axis is shown in Figure 3. The performance of our model is shown with the solid black line. The model has small initial errors in all three tasks. Then its error accumulates over time, but the accumulation is much slower than that of competing models, so our model consistently has the lowest error in the final time step.
To better understand our model, we have also done an extensive ablation study reported in the appendix. Here we only show a comparison of our model with GMR-GMUS in the encoding-decoding tasks. The results in Table. 1 show that the new model has clearly better encoding and decoding capabilities. In one encoding and decoding case, the reconstruction error of our model is less than one fifth of that of the GMR-GMUS model in the appendix. The ablation study also shows that both residual connections and encoding centers enhance the encoding-decoding accuracy (please see Table 8 in the Appendix section).
Contour analysis. As a probabilistic model, we plotted the contour of the predicted sample and the standard deviation (std) for all three datasets. Figure 2 visualizes the velocity predictions from two datasets. In each subplot, the top row is one prediction (e.g. one sample from our model), and the bottom row visualizes the standard deviation of predictions. To clearly compare different predictions, we use the same color bar within the same dataset. With visual inspection, the sample predictions are reasonable and stable for all three datasets.
The standard deviations also show reasonable spatial-temporal patterns consistent with our expectations. For cylinder flow with  $Re = 307$ , the standard deviation is larger around the vortex-shedding region, indicating the model is less certain about the fast-changing dynamics. Moreover, the standard deviation tends to grow with time, reflecting the error accumulation behavior for long-time rollouts. Moreover, the standard deviation is smaller when  $Re$  is higher, indicating that the model is more certain about the prediction. The same trends can also be seen in other datasets. Therefore, our proposed model can accurately predict the dynamics for deterministic systems while providing a spatial-temporal uncertainty estimate, potentially improving the interpretability of deep learning systems. More contours can be found in Appendix A.6
5.2 Stochastic dynamics
Dataset. We apply the proposed method to solve a stochastic dynamical system governed by the unsteady-state incompressible Navier-Stokes equations. We also compare our method to the unsteady Reynolds-averaged Navier-Stokes equations (URANS), the most applied method in the industry due to the balance of accuracy and computational efficiency. Let  $\Omega$  be the channel with a backward-facing
Table 2: The average relative rollout error of three systems, with the unit of  $1 \times 10^{-3}$ . We compare with five different baselines. Two variances of MeshGraphNet with or without noise injection (NI) and three variants of GMR-GMUS (LSTM, GRU, or Transformer).
Dataset-rollerstep Variable	Cylinder flow-400	Sonic flow-40	Vascular flow-250
u	v	p	u	v	p	T	u	v	p
MeshGraphNet [3]	NI	25	778	136	1.71	3.67	0.4	0.027	57	133	55
without NI	98	2036	673	4.12	6.13	0.24	0.020	3117	1771	601
GMR-GMUS [11]	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
decoding-only	4.9	89	38	0.95	2.8	0.43	0.39	7.3	10	22
PbGMR-GMUS	encoding-decoding	3.8	74	20	0.37	0.85	0.079	0.01	3.49	5.47	1.05
Figure 3: Averaged error over all state variables on cylinder flow (left), sonic flow (middle) and vascular flow (right), for the models MeshGraphNets(MGN)( $\cdots$ ), MGN-NI(), GMR-GMUS-GRU(), GMR-GMUS-LSTM( $\cdots$ ), GMR-GMUS-Transformer(), Ours model().
step shown in Figure 5. The goal of the problem is to generate the temporal unsteady velocity field close to the ones simulated by large eddy simulation (LES) with the inflow velocity along the inlet boundary subject to uncertainty. The inflow at every time step is uncertain and sampled from a 60-dimensional stochastic space subject to a uniform distribution as  $\{u_1^{\mathrm{inlet}},\dots,u_{60}^{\mathrm{inlet}}|u_i^{\mathrm{inlet}} = 10 + u',u'\sim \mathcal{U}(0,1)\}$ .
Analysis of fluid motion. To visually study the generated results, we plot several time steps of the stream-wise velocity  $(u)$  from a collection of model samples in Figure 4. The proposed method clearly generates diverse turbulent flow samples. Specifically, the multiscale structure of the vortex can be seen in the contour plots. To quantitatively examine the model performance on turbulent statistics, we plot and compare the mean and variance of the velocity profiles from the URANS, the proposed model, and LES results in Figure 5. The mean velocity profile predicted by the URANS model (gray lines) has significant discrepancies compared to the LES result (cyan lines), particularly in the right half of the domain where the backward-facing step's geometry is less important and small-scale turbulent features are dominant. As for the velocity variance, the URANS model underestimates the fluctuation in the whole domain. On the contrary, both the LES velocity mean and variance can be captured well by our generated samples (purple lines).
Comprehensive evaluations. To evaluate the temporal energy transport, we plot the power spectral density using Welch's method [47] in Table 3 (left). The high wave numbers reflect the rapid changes of fluid motion (e.g., instantaneous fluctuations), and the low wave numbers reflect the slow changes, such as the re-circulation introduced by the backward-facing step. By leveraging the sequence models,
Figure 4: Velocity contour at time step 40, 80, 120, 160 (left to right) from different samples.
Figure 5: Evaluation lines: Inlet (—), Outlet (—), Wall (—),  $x$ -evaluation lines (---). Quantities from LES (—), URANS (—), and our model (—)
the multiscale problem in the time domain can be automatically learned. However, URANS directly calculate an ensemble average over infinite experiments, and it can not capture stochastic behavior. Therefore, the energy  $(E(k))$  of high-wave-number  $(k)$  signals is underpredicted by URANS while the proposed model has a consistent energy spectrum pattern with the LES simulations. Moreover, the turbulent kinetic energy (TKE) is computed to evaluate the quality of generated samples. Physically, the TKE is characterized by measured root-mean-square velocity fluctuations. By applying the decoder, the point-wise instantaneous flow can be recovered from the latent vector to the physical solution reasonably well. The TKE metric can be defined as a scalar:
$$
\mathrm {T K E} := \frac {\int_ {\Omega} \int_ {[ 0 , T _ {\mathrm {e n d}} ]} [ (u - \bar {u}) ^ {2} + (v - \bar {v}) ^ {2} ] d \Omega d t}{\int_ {\Omega} \int_ {[ 0 , T _ {\mathrm {e n d}} ]} [ (u ^ {\mathrm {l e s}} - \bar {u} ^ {\mathrm {l e s}}) ^ {2} + (v ^ {\mathrm {l e s}} - \bar {v} ^ {\mathrm {l e s}}) ^ {2} ] d \Omega d t}. \tag {14}
$$
It indicates the preservation of the TKE benchmarked by the LES simulations. Our model can generate a flow field preserving  $99\%$  of the energy, whereas URANS can only do less than  $20\%$ . As for the temporal mean field, although URANS aims to solve it directly, it is inevitable to introduce the bias from the LES result due to the ergodicity assumption of RANS. Instead, we directly formulate the probabilistic problem to learn the distribution of the spatiotemporal field. We calculate the error of the temporal mean field and find that although URANS leverages the conservation laws, it still obtains a much larger error than our generated samples. Finally, we further propose to measure the quality of predicted rollouts using evaluation metrics of video generation in computer vision. In particular, we consider two metrics: the continuous ranked probability score (CRPS) [48, 49] and Frechet Video Distance (FVD) [50]. We use CRPS to assess the respective accuracy of our probabilistic forecasting and LES models. Since our model considers the distribution of the forecasts as a whole, it outperforms URANS, which only focuses on the mean of the distribution. Our model is also better evaluated in the FVD metric. This indicates the learned distribution by our model is close to the samples from LES.
We have also experimented with two deep learning methods, MeshGraphNet and GMR-GMUS, which are designed for deterministic systems. Their predictions are both far from LES simulations, not to mention that they can only make deterministic predictions. More results can be found in Appendix A.13.
We propose a new learning model to predict/generate deterministic and stochastic fluid dynamics over unstructured mesh in a uniform way. With an integration of a novel graph auto-encoder, a transformer, and a normalizing flow model, the new model decomposes temporal and spatial correlations in a dynamical system. It outperforms competitive baseline models for deterministic systems while providing a reasonable spatial-temporal pattern of forward uncertainty estimations. The samples
Table 3: Energy cascade spectrum at (2.58, 0.21)(left), and criteria of generation quality(right).
URANS (→), LES (→), Ours(→)	Quantity of interest	URANS	Ours
E(k)	10^0	CRPS (C_u,↓)	3.24	1.28
10^-7	CRPS (C_v,↓)	2.0	1.08
10^-14	FVD (d_u,↓)	228112	1262
FVD (d_v,↓)	137860	397
mean flow error (e_u,↓)	0.31	0.0176
mean flow error (e_v,↓)	0.94	0.176
turbulent energy (TKE,↑)	0.192	0.99
from the model trained on stochastic systems capture the rich physical patterns of expensive LES simulations. The current model still has one limitation: it can not accurately model stochastic variables close to boundary areas, which is shown in Appendix A.12. In future work, we will design new learning architectures to overcome this issue.
We thank all reviewers for their insightful feedback. Liu was supported by NSF CAREER 2239869. Wang's group would like to acknowledge the funds from Office of Naval Research under award numbers N00014-23-1-2071 and National Science Foundation under award numbers OAC-2047127.
[1] Fadl Moukalled, Luca Mangani, Marwan Darwish, F Moukalled, L Mangani, and M Darwish. The finite volume method. Springer, 2016.
[2] Stephen B Pope and Stephen B Pope. Turbulent flows. Cambridge university press, 2000.
[3] 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.
[4] Apostolos F Psaros, Xuhui Meng, Zongren Zou, Ling Guo, and George Em Karniadakis. Uncertainty quantification in scientific machine learning: Methods, metrics, and comparisons. Journal of Computational Physics, page 111902, 2023.
[5] Yarin Gal, Petros Koumoutsakos, Francois Lanusse, Gilles Louppe, and Costas Papadimitriou. Bayesian uncertainty quantification for machine-learned models in physics. Nature Reviews Physics, 4(9):573-577, 2022.
[6] Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. Advances in neural information processing systems, 32, 2019.
[7] Jayaraman Thiagarajan, Rushil Anirudh, Vivek Sivaraman Narayanaswamy, and Timo Bremer. Single model uncertainty estimation via stochastic data centering. Advances in Neural Information Processing Systems, 35:8662-8674, 2022.
[8] James M Dolezal, Andrew Srisuwananukorn, Dmitry Karpeyev, Siddhi Ramesh, Sara Kochanny, Brittany Cody, Aaron S Mansfield, Sagar Rakshit, Radhika Bansal, Melanie C Bois, et al. Uncertainty-informed deep learning models enable high-confidence predictions for digital histopathology. Nature communications, 13(1):6572, 2022.
[9] Sifan Wang and Paris Perdikaris. Long-time integration of parametric evolution equations with physics-informed deeponets. Journal of Computational Physics, 475:111855, 2023.
[10] Zhong Yi Wan, Leonardo Zepeda-Nuñez, Anudhyan Boral, and Fei Sha. Evolve smoothly, fit consistently: Learning smooth latent dynamics for advection-dominated systems. arXiv preprint arXiv:2301.10391, 2023.
[46] Zeqian Ju, Peiling Lu, Xu Tan, Rui Wang, Chen Zhang, Songruoyao Wu, Kejun Zhang, Xiangyang Li, Tao Qin, and Tie-Yan Liu. Telemelody: Lyric-to-melody generation with a template-based two-stage method. arXiv preprint arXiv:2109.09617, 2021.
[47] Peter Welch. The use of fast fourier transform for the estimation of power spectra: a method based on time averaging over short, modified periodograms. IEEE Transactions on audio and electroacoustics, 15(2):70-73, 1967.
[48] James E Matheson and Robert L Winkler. Scoring rules for continuous probability distributions. Management science, 22(10):1087-1096, 1976.
[49] Christopher AT Ferro, David S Richardson, and Andreas P Weigel. On the effect of ensemble size on the discrete and continuous ranked probability scores. Meteorological Applications: A journal of forecasting, practical applications, training techniques and modelling, 15(1):19-24, 2008.
[50] Thomas Unterthiner, Sjoerd van Steenkiste, Karol Kurach, Raphaël Marinier, Marcin Michalski, and Sylvain Gelly. Fvd: A new metric for video generation. 2019.
[51] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[52] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415, 2016.
[53] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):1735-1780, 1997.
[54] George Papamakarios, Theo Pavlakou, and Iain Murray. Masked autoregressive flow for density estimation. Advances in neural information processing systems, 30, 2017.
The Jacobian of the transformation in Equation (4) can be written as:
$$
\frac {\partial f ^ {i} (x)}{\partial f ^ {i - 1} (x)} = \left[ \begin{array}{c c} \mathbb {I} ^ {d ^ {\prime}} & 0 \\ \frac {\partial f ^ {i} (x) ^ {d ^ {\prime} + 1 : d}}{\partial f ^ {i - 1} (x) ^ {1 : d ^ {\prime}}} & \operatorname {d i a g} \left(\exp \left[ s \left(f ^ {i - 1} \operatorname {c o n c a t} ((x) ^ {1: d ^ {\prime}}, c)\right) \right]\right) \end{array} \right] \tag {15}
$$
where  $\mathrm{diag}\bigl (\exp [s(f^{i - 1}\mathrm{concat}((x)^{1:d'},c))]\bigr)$  is the diagonal matrix whose diagonal elements correspond to the vector  $\exp \left[s\left(f^{i - 1}\mathrm{concat}((x)^{1:d'},c)\right)\right]$ . Since the concatenation operation still keeps the matrix triangular, the determinant can be computed as:
$$
\det \left(\partial f ^ {i} (\boldsymbol {x}) / \partial f ^ {i - 1} (\boldsymbol {x})\right) | = \exp \left(\sum \left(s _ {i} \left(\operatorname {c o n c a t} \left(f ^ {i - 1} (\boldsymbol {x}) ^ {1: d ^ {\prime}}, c\right)\right)\right)\right) \tag {16}
$$
A.2 Training
As a regeneration learning method, there are two stages during the training. The first stage is to train a graph auto-encoder, PbGMR-GMUS, in this case, through self-supervised learning. This training process happens over all time steps and sequences in the training set. And for a given sequence  $\left[\mathbf{Y}_0,\dots ,\mathbf{Y}_T\right]$ , we compute  $\hat{\mathbf{Y}}_t = \mathrm{PbGMUS}(\mathcal{G},\mathrm{PbGMR}(\mathcal{G},\mathrm{Y}))$  at each step and minimize the reconstruction loss:
$$
\mathcal {L} _ {\text {g r a p h}} = \sum_ {t = 1} ^ {T} \| \boldsymbol {Y} _ {t} - \hat {\boldsymbol {Y}} _ {t} \| _ {2} ^ {2}. \tag {17}
$$
In the second stage, the parameters in PbGMR-GMUS will be fixed. And the attention-based sequence model, as well as the conditional flow model, will be trained. We maximize the log-likelihood in the latent space:
$$
\mathcal {L} _ {p} = - \frac {1}{T - 1} \sum_ {t = 2} ^ {T} \log p _ {\psi} \left(\boldsymbol {z} _ {t} \mid \boldsymbol {c} _ {t}\right) \tag {18}
$$
We don't set a time window so that the whole trajectory will be trained. The entire training process can be found in Algorithm 1. Particularly in the second stage, such a teacher forcing-like strategy allows a faster training process compared with the previous work [11], which adopts a rollout strategy during the training.
For inference we first encode the first two snapshots into latent space  $[z_0, z_1]$ . The condition vector  $c_t$  is computed from the attention-based sequence model. By sampling a noise vector  $\pmb{x}$  from an isotropic Gaussian, and putting it going backward through the flow, we can obtain a new sample from  $p_{\phi}(\pmb{z}_t | c_t)$  as latent representation at each step. We describe the procedure of training in Algorithm 2.
A.4 Latent analysis
Figure 6 plots the latent  $z_{t}$  and conditional vectors  $c_{t}$  for the stochastic BFS cases. For visualization purposes, we plot the first two modes of the PCA result. Moreover, we plot two different realizations of the stochastic process with clearly distinctive trajectories, as shown in the  $c_{t}$  plot (right). Since the cases are inherently from the same distribution, the sampled latent trajectories  $z_{t}$  plot (left) have a lot of overlaps after the few initial development steps. The plots of different cases visually conform to the physical intuition and indicate that our model can learn the relation between different cases well.
Algorithm 1 Training Process
Input: Domain graph  $\mathcal{G}$  , Node features over time  $[\mathbf{Y}_0,\dots ,\mathbf{Y}_T]$ $\mathrm{PbGMR}_{\theta}$ $\mathrm{PbGMUS}_{\phi}$  , Attention encoder MHA, Attention decoder Masked-MHAl, conditional flow model  $p_{\psi}(.|)$  , physical condition parameter  $\mu$    
Output: Learned parameters  $\theta$ $\phi$ $\omega$  and  $\psi$    
repeat for  $\mathbf{Y}_t\in [\mathbf{Y}_0,\ldots ,\mathbf{Y}_T]$  do  $\hat{\mathbf{Y}}_t = \mathrm{PbGMUS}_{\phi}(\mathrm{PbGMR}_{\theta}(\mathbf{Y}_t,\mathcal{G}))$  end for Compute  $\nabla_{\theta ,\phi}\gets \nabla_{\theta ,\phi}\mathcal{L}_{graph}(\theta ,\phi ,[\mathbf{Y}_0,\dots ,\mathbf{Y}_T],[\hat{\mathbf{Y}}_0,\dots ,\hat{\mathbf{Y}}_T])$  Update  $\phi ,\theta$  using the gradients  $\nabla_{\phi},\nabla_{\theta}$    
until convergence of the parameters  $(\theta ,\phi)$  {Self-supervised learning}  $[z_0,z_1,\dots ,z_T] = \mathrm{GMR}_{\theta}([Y_0,\dots ,Y_T])$    
repeat  $\hat{\boldsymbol{\mu}},\hat{\boldsymbol{z}}_{0} = \mathrm{MHA}_{\omega}(\boldsymbol {\mu},\boldsymbol {z}_{0})$ $[c_2,\dots ,c_t] =$  Masked-MHAs  $(z_{1},\dots ,z_{t - 1}),I = (\hat{\mu},\hat{z}_{0}))$  Compute  $\nabla_{\omega ,\lambda ,\psi}\gets \nabla_{\omega ,\lambda ,\psi}\mathcal{L}_p(\omega ,\lambda ,\psi ,[c_2,\dots ,c_t],[z_2,\dots ,z_t])$  Update  $\omega ,\lambda$ $\psi$  using the gradients  $\nabla_{\omega},\nabla_{\lambda},\nabla_{\psi}$    
until convergence of the parameters  $(\omega ,\lambda ,\psi)$
Algorithm 2 Inference Process
Input: Domain graph  $\mathcal{G}$  , The first two snapshot of node features  $[\mathbf{Y}_0,\mathbf{Y}_1]$ $\mathrm{PbGMR}_{\theta}$ $\mathrm{PbGMUS}_{\phi}$  Attention encoder  $\mathrm{MHA}_{\omega}$  , Attention decoder Masked-MHAA, conditional flow model  $p_{\psi}(|)$  physical condition parameter  $\mu$  , sample length  $T$    
Output: A trajectory sample  $[\mathbf{Y}_2,\dots ,\mathbf{Y}_T]$ $[z_0,z_1] = \mathrm{PbGMR}_{\theta}([Y_0,Y_1],\mathcal{G})$ $\hat{\boldsymbol{\mu}},\hat{\boldsymbol{z}}_0 = \mathrm{MHA}_{\omega}(\boldsymbol {\mu},\boldsymbol {z}_0)$    
for  $t\in [2,\ldots ,T]$  do   
 $c_{t} = \mathrm{Masked - MHA}_{\lambda}((z_{1},\ldots ,z_{t - 1}),I = (\hat{\boldsymbol{\mu}},\hat{\boldsymbol{z}}_{0}))[-1]$  Sample latent representation from conditional flow  $z_{t}\sim p_{\phi}(z_{t}|c_{t})$    
end for   
Recover physical parameters on the original space  $[\mathbf{Y}_2,\dots ,\mathbf{Y}_T] = \mathrm{PbGMUS}_{\phi}([z_2,\dots ,z_T])$
(b)
Figure 6: 2-D principle subspace of the (a) latent vectors and (b) conditional latent vectors for the stochastic BFS cases: trajectory  $1(\longrightarrow)$ , trajectory  $2(\longrightarrow)$  start from  $\bullet$  and  $\bullet$ , respectively.
A.5 Dataset
Four fluid simulation datasets, cylinder flow, sonic flow, vascular flow and stochastic backward-facing step (BFS) are used in the experiments. The labeled data are generated by solving the Navier-Stokes equation by OpenFOAM [21], a commonly used finite volume method-based open-source solver for fluid simulation. Cylinder and sonic flow consist of 50 cases in the training dataset and 50 cases in the test dataset. The vascular flow cases have 10 training and 10 test cases. The stochastic BFS case has 5 training cases. Unlike the first three deterministic cases, the stochastic BFS case is a stochastic case. The simulation details can be found in Table 4. Moreover, the details input and output of the PbGMR-GMUS are listed in Table 5.
The governing equation for the first three cases is the same as the one listed in [11]. For the last stochastic BFS case, the governing equation is listed in Equation (19). One significant difference in this case is that it introduces stochasticity by perturbing the inlet  $u$  velocity at every time step.
Table 4: Simulation details of datasets
Dataset	Meshing	# nodes	# steps	type
Cylinder flow	Fixed	1699	400	deterministic
Sonic flow	Moving	1900	40	deterministic
Vascular flow	Small-varying	7561 (avg.)	250	deterministic
Stochastic backward-facing step	Fixed	22500	240	stochastic
Table 5: Input-output information of PbGMR-GMUS:  $u$ : X-axis velocity,  $v$ : Y-axis velocity,  $p$ : pressure,  $T$ : temperature,  $\rho$ : density,  $m$ : cell volume.  $T_{0}, Re, r, u_{Bi}$  denote the initial temperature, Reynolds number, radius of the thrombus, and boundary velocities, respectively. BFS denotes backward-facing steps.
Dataset	PbGMR node input	PbGMUS node output	Nodal embed dim	# centers
Cylinder flow	ui,vi,pi,mi,Re	ui,vi,pi	4	256
Sonic flow	ui,vi,pi,Ti,ρi	ui,vi,pi	4	256
mi,T,0	Ti,ρi
Vascular flow	ui,vi,pi,mi,r	ui,vi,pi	4	400
Stochastic BFS	ui,vi,mi,uBi	ui,vi	4	256
$$
\frac {\partial u}{\partial x} + \frac {\partial v}{\partial y} = 0,
$$
$$
\frac {\partial u}{\partial t} + u \frac {\partial u}{\partial x} + v \frac {\partial u}{\partial y} = - \frac {1}{\rho} \frac {\partial p}{\partial x} + \nu \left(\frac {\partial^ {2} u}{\partial x ^ {2}} + \frac {\partial^ {2} u}{\partial y ^ {2}}\right), \tag {19}
$$
$$
\frac {\partial v}{\partial t} + u \frac {\partial v}{\partial x} + v \frac {\partial v}{\partial y} = - \frac {1}{\rho} \frac {\partial p}{\partial y} + \nu \left(\frac {\partial^ {2} v}{\partial x ^ {2}} + \frac {\partial^ {2} v}{\partial y ^ {2}}\right),
$$
subject to  $u_{i}^{\mathrm{inlet}} = 10 + u^{\prime}, u^{\prime} \sim \mathcal{U}(0,1)$  at every time step.
A.6 Additional visualization results
Apart from the contour shown in Figure 2, additional contour plots for the cylinder, sonic flow and vascular flow are shown in this section. Figure 7 shows the sampled velocity and std of the velocity for vascular flow cases. Figure 8 presents the sampled pressure for cylinder, sonic and vascular flow. Additionally, Figure 9 shows the sampled temperature variable in sonic flow.
Figure 7: Velocity contour for the vascular flow case, the first row is one of our predicted samples. The second row is the predicted standard deviation.
Figure 8: Pressure contour for benchmark cases
Figure 9: Temperature contour for sonic flow case.
A.7 Multi-head attention