Deep Neural Networks to Correct Sub-Precision Errors in CFD

Computational simulations of fluids and other complex physical systems have critical applications in engineering and the physical sciences. Accurate and dependable numerical simulations are essential for designing aircraft components and turbomachinery, predicting climate patterns, and analysing flow in blood vessels. The mathematical models that describe these physical systems often come in the form of Partial Differential Equations (PDEs) represented in a discretised form in computers and are solved using various numerical schemes. The complexity of real-world systems being simulated presents a trade-off between resolving the finest spatio-temporal details and keeping the computational cost to a feasible limit. Information loss in a numerical simulation can arise from various sources: the spatial and temporal discretisation of the PDEs, or the availability of finite precision in the representation of numerical values in computer memory. This work examines the benefits of low numerical precision in Computational Fluid Dynamics (CFD) and uses deep learning to correct the resulting errors.

The rise of machine learning (ML) as a powerful tool for information processing, together with the large amounts of data generated in CFD simulations, has led to the use of ML for understanding, modelling, optimising and controlling fluid flows [1]. Purely data-driven approaches to spatio-temporal forecasting learn deep neural networks in a supervised manner without enforcing constraints on the physical properties of the output [2, 3, 4]. Such methods are computationally efficient but fail to enforce constraints such as conservation of mass and momentum and are hence suitable for accelerating simulations in computer animations and games where the exact physical correctness of the flow is not essential. Loss-based approaches try to guide the neural network predictions by including prior scientific knowledge in the loss function. Physics-guided or physics-informed neural networks encode physical relationships between variables and underlying physical laws such as PDEs in the loss function, and hence are biased towards producing physically accurate outputs [5, 6, 7]. Machine learning has also been used in turbulence closure modelling [8, 9, 10].

Recent efforts to use machine learning in scientific computing have involved differentiable programming [11, 12, 13]. Software frameworks for deep learning (e.g. PyTorch and JAX) implement reverse-mode automatic differentiation to compute gradients in neural networks [14, 15]. This capability can be leveraged by re-implementing numerical solvers in these frameworks, allowing a neural network to be trained end-to-end along with the numerical solver by backpropagating gradients through the solver. Recently this method has been demonstrated to improve generalisation to new inputs and the adherence to physical laws [16, 17]. Since the gradients can be readily calculated end-to-end, a solver-ML system trained in this manner can also be used for flow control [18], inverse design and shape optimization [19, 20]. Several numerical solvers written in automatic differentiation frameworks have been recently open-sourced [17, 21, 22].

Figure 1 illustrates the double-precision (float64) and single-precision (float32) floating-point formats commonly used in scientific computing. Most modern computer systems, including Graphics Processing Units (GPUs) often used for CFD, implement the IEEE 64-bit floating-point arithmetic standard, consisting of 53 mantissa bits resulting in roughly 15-17 significant digits of precision. The IEEE 32-bit floating-point standard consists of 23 mantissa bits resulting in approximately 6-9 significant digits of precision. In recent times, the realisation that deep neural networks retain their learning and predictive power even when using reduced-precision numerics has spurred the widespread adoption of floating-point standards such as half-precision IEEE 16-bit format (float16), which consists of 10 mantissa bits resulting in roughly 4 significant digits of precision [23]. This format is currently well supported on modern GPUs, with around 2-3 times more theoretical throughput than float32 [24]. The GPU memory savings obtained from the reduced numerical precision can also enable simulations with more mesh elements than would otherwise be possible. However, with 5 exponent bits, float16 has a maximum representable value of 65504. This can lead to overflow issues in scientific computing applications, including deep neural networks [25]. The ’Brain’ floating-point format (bfloat16) was developed at Google Brain to mitigate this issue [26]. The advantages of this format for computing are that the range of values it can represent is the same as that of float32, and conversion to and from float32 is simple, making it a drop-in replacement for float32 in situations where reduced precision is desired to trade-off precision and compute cost. The use of 16-bit formats is further motivated by a trend towards new hardware containing special units optimised for low-precision/mixed-precision computations. For example, bfloat16 is implemented in modern Intel x86 CPUs, Google’s TPU accelerator and the new NVIDIA Ampere and Hopper architecture GPUs.

Previous studies on numerical precision in scientific computing noted that certain operations are more susceptible to error accumulation than others. In particular, large summations or accumulation operations, repeated operations in loops, large-scale parallel computation (since floating-point arithmetic is non-associative), and resolving small-scale phenomena cause the most error accumulation [27]. Simulations of turbulent fluid flows satisfy all of these conditions: resolving small-scale eddies is important; the PDEs are inherently non-linear, resulting in vastly different evolution trajectories with small changes in the initial conditions; and they are solved using large scale matrix operations on parallel processing devices such as GPUs. Climate modelling is an application where the accumulation of numerical error makes reproducibility and verification of the results difficult. He and Ding [28] showed that employing double-double precision arithmetic in long inner product loops in atmospheric simulations dramatically reduces the variability of results between subsequent runs.

To examine the benefits of reduced precision, we performed preliminary tests on the canonical test case of the Taylor Green Vortex on a cartesian grid of $2\pi^{3}$ [29]. We have used our in-house high-order CFD solver COMPSQUARE [30, 31] to carry out these direct numerical simulations. The general purpose solver can handle body-fitted curvilinear meshes and uses high order compact schemes to solve the conservative form of the compressible Navier-Stokes equations. A V100 GPU with 32 GB of memory can accommodate a grid size of upto $415^{3}$ points ($\approx 72\times 10^{6}$) using float32 and around $330^{3}$ points ($\approx 36\times 10^{6}$) using float64. For a given grid, the simulations using float32 are $\approx 1.65\times$ faster and can accommodate twice the grid size than those using float64. On hardware optimized for 16-bit computation, the speedup and the memory savings are expected to be even greater. For example, NVIDIA’s latest hardware optimized for 16-bit (and lower) precision is advertised to show between $3-12\times$ speedup across all applications ranging from High Performance Computing (HPC) to Artificial Intelligence (AI) [32].

This research is motivated by: (a) The speedup seen in low-precision simulations, (b) the reduced memory requirements, (c) the increasing hardware support for low-precision computation and (d) the demonstrated ability of deep learning to improve numerical simulations.

In Section 2, we quantify the error accumulation in a Kolmogorov forced turbulence test case. In Section 3, we explore deep learning strategies to correct this error. We present the key results from our experiments in Section 4, following which we conclude with a discussion on future directions in Section 6.

In this section, we attempt to use deep learning to learn a function that can correct the error arising in a small duration of time between the low-precision and high-precision solvers. We use a neural network that interacts with the solver through differentiable programming [16, 17].

For our experiments, we define $\Delta$ to be the total time duration (number of solver iterations) that the model sees in a single training sample. We define $N$ as the number of corrections the neural network makes to the low-precision simulation during the time duration $\Delta$. Therefore $T$, given by the relation $T=\Delta/N$, is the number of iterations the solver steps forward before passing the output to the neural network for a correction.

We experiment with training the model on data generated using $\Delta=150$ and $\Delta=300$. In each case, we experiment with four values of $N$ (the number of corrections made during $\Delta$). For each set of values $\{\Delta,N\}$, we generate a training, validation and held-out test set with dataset size $S$ = 200, 10 and 2, respectively. The validation and held-out test sets are generated as in the Pseudocode, whereas in the case of training data, the low-precision snapshots are generated on-the-fly during training.

While the models see windows of 150 and 300 timesteps during training, we test the error reduction and stability of the model outputs over a longer time duration. We generate a long-run test set consisting of simulation sequences of several thousand timesteps, allowing us to test the ability of the trained models to produce corrected sequences far longer than what they saw during training.

The metrics we use to quantify the performance of the models broadly fall into two categories: a) aggregate metrics such as kinetic energy spectra and mean divergence, which deal with statistical accuracy of the flow, and b) error metrics such as Mean Squared Error (MSE) and correlation coefficient which deal with pointwise accuracy of the flow, i.e., the exact evolution of eddy locations over time.

Tables 1 and 2 show the relative improvement in MSE during a time duration $\Delta$, which is calculated on the held-out test set for each $\{\Delta,N\}$. $t_{90\%}$ measures the number of timesteps for which the MSE between the corrected sequence and the high-precision sequence is $\leq 90\%$ of the MSE between the low-precision sequence and the high-precision sequence.

We observe that the models making more frequent corrections ($N=30$ and $N=50$) perform better in terms of relative improvement in MSE over $\Delta$, compared to the models making infrequent corrections ($N=2)$. In terms of $t_{90\%}$, the $N=2$ models fail to keep the error under control for very long. The $N=15$, $N=30$ and $N=50$ models, on the other hand, have $t_{90\%}$ significantly greater than the time duration which they saw during training. The $T=6$, $N=50$ model shows an anomalous value of rel. improvement in MSE, despite which it has a better $t_{90\%}$. Figure 7 shows the evolution of the MSE for the low-precision and corrected sequences on the long-run test set. Except for the $T=150,N=2$ model, all the models delay the error accumulation by some amount. Looking beyond the timesteps shown in Fig. 7, however, all the flows eventually become uncorrelated with their target high-precision flows.

We measure the effective computational throughput of each of our learned solvers on our workstation using a single NVIDIA V100 GPU. For a given $\Delta$, the more frequently the neural network is invoked to make a correction, the more computational overhead is introduced. This presents a trade-off, which can be seen in Tables 1 and 2. Making more frequent corrections results in lower error but incurs greater computational cost resulting in lower throughput. The throughput can be further improved by training and evaluating the neural network using bfloat16 arithmetic, which has been shown to preserve their learning capability [26]. However, due to the lack of support for bfloat16 on the Tensor Cores on our available hardware, [37] this could not be done in the present study. For this same reason, a direct comparison of throughput with the base 16-bit solver is not meaningful.

Figure 8 visualizes the vorticity fields obtained from the experiment with $\Delta=150,N=50$. The low-precision flow appears grainy and distorted, whereas the corrected flow regains the lost visual quality. Furthermore, the locations and shapes of the vortices in the final snapshot match more closely with the corresponding high-precision snapshot, which indicates that the corrective model delays the time when the low- and high-precision flows become uncorrelated (see snapshots at $t=1980$).

Mean Squared Error, being a pointwise comparison metric, only considers local distances and may fail to capture structural similarity and physical characteristics of numerical simulations [36]. Therefore it is not a sufficient metric on its own to capture some of the differences seen in Fig. 8, which can be better characterised by inspecting flow properties such as the kinetic energy spectrum. Therefore, we compare the kinetic energy spectra $E(k)=\frac{1}{2}|u(k)|^{2}$ of the low-precision, corrected and high-precision sequences in each experiment. Figure 9 shows the average of the scaled kinetic energy spectrum $E(k)\cdot k^{5}$ over 1000 snapshots from the long-run test set. These 1000 snapshots are taken between 3000 and 150000 timesteps. From Fig. 2, it is evident that the corrected and reference flows are uncorrelated at these timesteps. Generally, both the low-precision and corrected sequences yield spectra which match closely with the high-precision spectra in the low to mid frequency range ($k=4\text{--}20$). However, we see that the graininess caused by low numerical precision results in surplus energy accumulation at high frequencies (larger $k$). The models trained with $N=2$ fail to produce accurate spectra compared to the high-precision simulation at all frequency ranges. The learned solvers with more frequent corrections ($N=30$ and $N=50$) produce significantly improved spectra. They reduce the surplus high-frequency energy accumulation, except at the highest frequencies, where all the models deviate from the high-precision spectra.

We compute the divergence of the velocity fields in the high-precision, low-precision and corrected sequences using finite difference derivatives and inspect its mean value over the domain. We observe that all three flows have a mean divergence of order 1e-11, which shows that the correction made by the neural network preserves the net-zero divergence property of the field.

Figure 10 shows the statistical correlation of the velocity magnitude between the corrected and the high-precision flows as a function of time. In this case, as well, the models making more frequent corrections perform better. With the exception of the $\Delta=300,N=15$ and all the $N=2$ models, the corrections made to the low-precision flow delay its decorrelation with the reference flow. Compared to the coarse-grid simulation of Kochkov et al. [17], the rate of decorrelation of the low-precision simulation is significantly slower. Consequently, our models achieve a smaller relative delay in decorrelation than their models.

As a test of generalization, we take the neural network obtained using the best-performing set of hyperparameters ($T=3,N=50$) and use it to generate corrections to a simulation performed at a 10 times higher Reynolds number (Re=40000 instead of 4000). Figure 11 compares the KE spectra with and without NN correction. With no further training, the corrections made by the neural network to the Re=40000 low-precision solver’s output achieved similar improvements in all the metrics (Rel. improvement in MSE over $\Delta$, $t_{90\%}$, KE spectrum, etc.). This suggests that in addition to identifying trends in the model’s performance, we have also identified a set of hyperparameters that can be used as a useful starting point when applying our method to different flows.

There are broadly two kinds of metrics used in this study to evaluate the performance of NN corrections: a) aggregate metrics such as KE spectra and mean divergence which deal with the statistical accuracy of the flow, and b) error metrics such as Mean Squared Error and correlation coefficient which deal with pointwise accuracy of the flow, i.e, the exact evolution of eddy locations over time. {update} From Figure 9, the role of the proposed NN correction appears to be identical to that of a filter that damps energy content over a range of frequencies. An experiment has been performed (refer to A) in these lines where we have compared the spectra and vortical structures replacing the NN correction with a spectral filter. The physics-informed NN correction function learned from data is observed to be superior when compared to a simple filtering operation. Further investigations can explore the efficacy of a combination of NN correction (to control mid-high frequencies) with a standard filtering operation (to control high-frequency spectra), considering the computational overheads.

While our method achieves improvements in error metrics and physical quantities, the hybrid solver does not beat the reference high-precision solver in computational throughput. This is due to limitations in hardware available to carry out the study, as mentioned in Sections 2.2 and 4. The computational efficiency can be further optimized by running on hardware designed for low-precision computation, training the neural network itself using bfloat16 [26], applying convolutions in the Fourier domain [41], and applying model quantization techniques [42]. Improving the throughput of the high-precision solver is an important direction for future work.

The motivation for this method stems from related works by Kochkov et al. [17] and Um et al. [16]. In both these studies and in the current study, snapshots from a Navier-Stokes solver are used as the targets during training, which means that the neural network is implicitly trained to make corrections that satisfy the governing equations. A limitation to this strategy is that the corrections themselves need not explicitly satisfy the governing equations of the flow. Instead, by using snapshots from a real solver as the targets during training, the neural network is implicitly trained to predict a flow field that satisfies the governing equations. A direction for future work that addresses this limitation is to combine our method with the method proposed by Raissi et al. [6], where the network’s output is explicitly regularised to satisfy the governing equations by including an additional physics-based term in the loss function.

In this work, we explored the error accumulation due to low precision in a Kolmogorov forced turbulence test case and used deep learning together with a differentiable solver to improve the simulation. The learned ML-CFD hybrid reduces the Mean Squared Error accumulation over time durations longer than it saw during training, improves the kinetic energy spectra and preserves the mean divergence of the flow. We show that making more frequent corrections to the simulation achieves a greater reduction in error, although it comes with an increased computational cost. The neural network’s corrections improve the visual fidelity of the low-precision flow, making it look less grainy and distorted. Furthermore, the correction is dependent on only 2 parameters (the frequency of corrections and total time window seen during training), apart from the usual parameters involved in neural network training. We quantified the dependence of the model’s performance on these parameters and identified trends which will help implement the method in reality. The GPU memory savings obtained from the reduced numerical precision can enable simulations with more mesh elements than would otherwise be possible. Furthermore, since the hybrid solver is fully differentiable, it is suitable for use in control and inverse problems. While established industrial solvers written in FORTRAN or C are not directly compatible with our deep learning method implemented in JAX, new high-order accurate differentiable solvers are under development using the JAX framework [43, 44], which solidifies the long-term viability of our method in real-world uses. We discussed the limitations of the method and identified opportunities to extend and improve it in future work. {update} Although the current study focuses on non-reacting turbulent flows, the methods and the neural network framework discussed here can be readily extended to reacting flows.

N.R.V acknowledges the Department of Science and Technology-Science and Engineering Research Board (DST-SERB) for funding the project. The support and the computing resources provided by PARAM Sanganak under the National Supercomputing Mission, Government of India, and P G Senapathy Centre, IIT Madras are gratefully acknowledged. Y.M. acknowledges the support of Japan Science and Technology Agency, PRESTO, Grant Number JPMJPR21OB.