Building high accuracy emulators for scientific simulations with deep neural architecture search

The example outputs of the trained emulators with DENSE for some cases are shown in figure 2. The outputs for other cases can be seen in figure 7. We see that the output of the emulators generally matches closely the output of the actual simulations, even in MOPS and GCM where only 410 and 39 data points are available. When only a limited number of data is available, the choice of model architectures that give the right priors become important. Complex model architectures with bad priors could still fit the sampled training data, but are likely to overfit, giving bad accuracy on the out-of-samples data.

Zoom In Zoom Out Reset image size

With DENSE, the model architecture that gives a good prior on the problem is automatically searched for by preferring models that can fit well the out-of-samples data (i.e. the validation dataset). Moreover, randomly choosing an operation in every layer acts as a regularizer in updating the weights during the training to minimize overfitting. These two advantages make DENSE suitable for learning to emulate a wide range of simulations including expensive ones where only a limited number of datasets can be generated.

While the simulations presented typically run in minutes to days, the DENSE emulators can process multiple sets of input parameters in milliseconds to a few seconds with one CPU core, or even faster when using a GPU card. For the GCM simulation which takes about 1150 CPU-hours to run, the emulator speedup is a factor of 110 million on a like-for-like basis, and over 2 billion with a GPU card. The speed up achieved by DENSE emulators for each test case is shown in figure 3(b).

Zoom In Zoom Out Reset image size

Compared with other non-deep learning techniques usually employed in building emulators [25], the models found and trained by DENSE achieved the best results in all tested cases, and in most cases by a significant margin. The DENSE model also performs better than an emulator designed specifically for ICF JAG simulation [26] (i.e. CycleGAN in figure 1 of the supplementary material (available online at stacks.iop.org/MLST/3/015013/mmedia)). As seen in figure 3(a), the emulators built by DENSE achieved a loss function up to 14 times lower than the best performing non-deep learning model.

We also compared DENSE with a manually-designed deep neural network model by an architecture from the super-architecture in figure 1, where all the convolutional layers have size 3. The use of kernel size 3 and skip connections follows the idea of ResNet [11]. As with DENSE, the hyperparameters for manually-designed deep neural network were tuned to optimize its result for OTS.

Shown in figure 3(a) is the comparison between DENSE emulators and manually-designed deep neural network. Although in most cases their performances are similar, in some cases DENSE emulators can give a considerable improvement in terms of loss function, as can be seen in Halo and MOPS. This illustrates the robustness of DENSE emulators in a wide range of cases.

The high fidelity emulators built by DENSE are sufficiently accurate to allow us to substitute simulations even for more advanced tasks such as solving the inverse problem [27]. To illustrate this, we took a simulated output signal randomly from the test dataset where the actual parameters are known. A small noise (${\sim} 1\%$) was added to the chosen signal to closely mimic a real observed signal from an experiment. Using this signal, we use an optimization algorithm [28] to retrieve the input parameters by minimizing the error between the sample signal and the output of the emulators, which we call it the retrieval error.

The results of the parameter retrieval using the emulators are compared with the retrieval using the simulations in figure 4, where we plot the relative retrieval error histograms for two cases. We observe that the relative retrieval errors from the emulators are very similar to those from the simulations.

Zoom In Zoom Out Reset image size

Without much loss in accuracy, the parameter retrieval with the emulators only takes about 800 ms with a single GPU card. This is to be compared with using the actual simulations which could take up to 2 d (XES) even when using 32 CPU cores. As the parameters can be retrieved in less than one second rather than in hours or days, one can envisage employing this technique for online diagnostics, real-time data interpretation, or even on-the-fly intelligent automation with an accuracy comparable to high-fidelity simulations that are by far too computationally expensive to be used directly. The use of DENSE emulators also enables parameter retrieval with resource-intensive simulations, such as MOPS and GCM, that were too expensive before.

In addition to interpreting signals and parameter retrieval, the emulators can also be used to significantly speed up Bayesian uncertainty quantification [27]. Bayesian uncertainty quantification is usually done by constructing Bayesian posterior distribution using Markov Chain Monte Carlo (MCMC) algorithms. However, the cost of running MCMC to collect sufficient samples from the Bayesian posterior distribution is typically much larger than the cost for parameter retrieval, and is often intractable in practice. Here we perform the Bayesian posterior sampling using an ensemble MCMC algorithm [29] with the same conditions as in [27]. In short, we collect all parameters sets that produce spectra that lie in a certain band around a central spectrum.

Figure 5 compares the results of sampling the posterior distribution using simulations and emulators in two cases to interpret scattering and spectroscopy data. The posterior distribution sampled by the emulators are very similar to those by actual simulations, and we see that the emulators are well-suited to capture the correlations between parameters. However, note that while collecting 200 000 XES samples via simulations takes over 22 d, the sampling process with the emulators was completed in just a few seconds. Interestingly, building the emulator for XES from scratch only needs some 14 000 samples plus 8 h for training, so the whole pipeline to build the emulator and use it for MCMC is still considerably faster than directly collecting 200 000 samples using the original simulation.

Zoom In Zoom Out Reset image size

A final important advantage of building emulators with DENSE is the availability of an intrinsic estimator of the emulator's prediction uncertainty. The randomization of network architectures from the super-architecture can be seen as a special case of dropout [30]. Thus, by adapting the theory of prediction uncertainty with Monte Carlo (MC) dropout by [31], we can show that DENSE emulators can produce the uncertainty of their outputs. The expected value and variance of a DENSE emulator prediction can be obtained by

$$\begin{align} \mathbb{E}(\mathbf{y}|\mathbf{x}) & = \mathbb{E}_{a\sim\mathcal{A}(\mathbf{b})}\left(\mathbf{y} | \mathbf{x}, a\right) \nonumber\\ \mathrm{Var}(\mathbf{y}|\mathbf{x}) & = \mathrm{Var}_{a\sim\mathcal{A}(\mathbf{b})}\left(\mathbf{y} | \mathbf{x}, a\right), \end{align} \tag{ 3 }$$

where a is the architecture sampled from the super-architecture $\mathcal{A}$ based on the final values of the network parameters, b. Figure 6 shows the prediction uncertainty of the DENSE emulators for some cases, illustrating regions where they are either uncertain or confident in their predictions. The prediction uncertainty for other cases can be seen in figure 8.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

As also observed in MC dropout [32], the prediction uncertainty in this case can be smaller than the difference between the predicted and simulated output, indicating an overconfident prediction. This problem of overconfidence can be resolved by stopping the training of network variables early. The investigation of prediction uncertainty tuning will be the subject of future work.

Although DENSE has the capability of emulating a wide range of simulations, it is still limited to simulations with a few scalar inputs. The DENSE algorithm has not been tested on building emulators with one, two, or three-dimensional direct inputs. One way to fit a simulation with multi-dimensional inputs to DENSE is by parameterizing the inputs using several scalar parameters (as done in ELMs) or employing a dimensionality reduction techniques [33].

Another limitation observed in DENSE is that it does not learn very well in regions where output variability is high, i.e. where changing the input parameters slightly changes the outputs significantly. This limitation is also observed in other cases of deep learning [34]. Due to the difficulty in learning, regions with high variability tend to require more samples than regions with low variability. This problem can thus be overcome by sampling the parameter space intelligently [35].

Our DENSE approach opens up numerous applications that require fast calculations. One of the main applications of DENSE emulators is real-time diagnostics of complex systems. For research in some fields, such as in plasma physics, diagnostics are usually done by solving an inverse problem using simulations, which involves running the simulations hundreds of times or more. By using the fast emulators instead of simulations, solving the inverse problem could be significantly accelerated without sacrificing the quality of the solutions, as shown in section 2.2. Real-time diagnostics are the key in automating operations of some machines with complex systems, such as tokamaks [36] and particle accelerators [37].

Another application is the optimization of a very expensive simulations where the simulations can only be executed a few times. The emulators can be used to make high-quality guesses about the optimum parameters which can then be tested using the expensive simulations. This idea of utilizing a cheap model for expensive simulation optimization has been used in surrogate-model optimization [38] and Bayesian optimization [39]. By using a more accurate cheap model, such as DENSE emulators, the number of simulations to be executed can be much lower than in previous approaches. This is a potential avenue for future works.

Here we provide a description of the test cases employed in the paper. A summary of the test case parameters is given in table 1. All 1D simulation outputs are sampled to 250 points for convenience.

5.1.1. X-ray Thomson scattering (XRTS)

5.1.2. Optical Thomson scattering (OTS)

5.1.3. X-ray emission spectroscopy (XES)

5.1.4. Edge-localized modes (ELMs) diagnostics

5.1.5. Galaxy halo occupation distribution modelling (Halo)

5.1.6. JAG model for inertial confinement fusion (ICF JAG)

5.1.7. Shatsky Rise seismic tomography (SeisTomo)

5.1.8. Global aerosol-climate modelling (GCM)

5.1.9. The model of oceanic pelagic stoichiometry (MOPS)

The super-architecture employed for most cases is shown in figure 1. It consists of two fully connected layers at the beginning, followed by combinations of different types of convolutional layers. Rectified Linear Units are used as the nonlinearity. Most of the nodes contain 64 channels with a growing size of the signal from 4 to 250 at the end. For ICF JAG that has output signal of size $64 \times 64$, the size written in figure 1 is capped at 64.

After the first two fully connected layers, each pair of adjacent nodes are connected by an identity layer and a selection of multiple convolutional layers. The identity layers serve as the skip connection for each layer which is always present in every sampled architecture. The member j in group i of layers is assigned a network parameters, b_ij , and the probability of member j being selected among the group is determined by the softmax function,

$$\begin{align} p_{ij} = \frac{\exp{b_{ij}}}{\sum_k \exp{b_{ik}}}.\end{align} \tag{ 4 }$$

Each group of layers consist of one zero layer which multiplies all the inputs to zero. This provides an option for DENSE to remove the layer and make the neural network shallower.

For two nodes with different sizes, we added modified transposed convolutional layers in the group. In the modified transposed convolutional layers, the signal is expanded just as in the normal transposed convolutional layer, but the 'holes' are filled with a trainable constant instead of zeros. The convolutional layers and identity layers between two nodes of different sizes operate by upsampling the previous node to match the size of the target node using the nearest neighbor.

The list of hyperparameters used in the algorithm are given in table 2. The hyperparameters were chosen to minimize the validation loss function for OTS using CMA-ES [14, 24]. The same hyperparameters were used for all cases.

Table 2. List of hyperparameters used in training the emulators.

Hyperparameters	For DENSE	For manual design DNN	Notes
$n_\mathrm{epochs}$	3000	3000	How many epochs
α1	$3.06 \times 10^{-4}$	$4.34 \times 10^{-3}$	Learning rate for weight update
m1	35	72	Size of minibatch in weight update
γ1	0.757	0.9913	Decaying multiplier for α1
s1	513	7	Apply the decay multiplier for α1 after this many update steps
α2	$4.88 \times 10^{-3}$	—	Learning rate for architectural update
m2	142	—	Size of minibatch in architectural update
γ2	0.701	—	Decaying multiplier for α2
s2	918	—	Apply the decay multiplier for α2 after this many epochs
$p_\mathrm{val}$	2	1	Going through the validation dataset this many times in one epoch

In training the emulators using non-deep learning methods, we employed the scikit-learn library [25] in Python. We use the default parameters suggested in the library to build the emulators for all cases.

For models that can only predict a single output, an ensemble of models are trained to predict different outputs in one simulation. For CycleGAN in the ICF JAG and ICF JAG Scalars cases, the model was trained and obtained according to [23, 26].

The parameter retrieval processes for XRTS and Halo to produce figure 4 were done using the CMA-ES [14] algorithm with population size 32 and 1200 maximum function evaluations. Default parameters suggested in [14]. were used. To give a fair results comparison, we used the same algorithm parameters and conditions in parameter retrievals via simulations and emulators.

For the Bayesian posterior sampling process, we employed the affine-invariant ensemble MCMC algorithm [29] with 256 walkers to collect 200 000 samples for XRTS and XES cases. The likelihood is uniform when the generated spectrum lies in a given band and it is zero when it lies outside the band. The band is 0.035 J keV⁻¹ ster⁻¹ in XES and 3.5% in XRTS as used in [27]. The justification of this form of likelihood is also provided in the supplementary materials of [27].

The randomization of the network architecture can be seen as a special case of dropout. Therefore, the derivation of the prediction uncertainty follows the derivation in Monte Carlo (MC) dropout very closely [31].

Denote the input and output from the training dataset as X and Y respectively and write ω as the active weights in the neural network. Given the training data, X and Y, the posterior distribution of the weights in the neural network can be written as

$$\begin{align} \mathbb{P}(\boldsymbol{\omega} | \mathbf{X}, \mathbf{Y}) = \frac{\mathbb{P}(\mathbf{Y}|\mathbf{X}, \boldsymbol{\omega}) \mathbb{P}(\boldsymbol{\omega})}{Z} \end{align} \tag{ 5 }$$

where Z is the normalization constant, $\mathbb{P}(\mathbf{Y}|\mathbf{X}, \boldsymbol{\omega})$ is the likelihood of observing Y with weights ω , and $\mathbb{P}(\boldsymbol{\omega})$ is the prior distribution of the weights.

The posterior distribution of the weights are intractable, so we need to use variational inference in making the approximation to the posterior distribution. Let the variational distribution takes the form of $\mathbb{Q}(\boldsymbol{\omega})$ where

$$\begin{align} \boldsymbol{\omega_{ij}} & = \mathbf{w_{ij}} z_{ij},\ \mathrm{and} \end{align} \tag{ 6 }$$

$$\begin{align} z_{ij} &\sim \mathrm{Bernoulli}[p_{ij}(b_{ij})] \end{align} \tag{ 7 }$$

where $\mathbf{w_{ij}}$ is the weights of layers j in group i, p_ij is the probability of being selected as a function of the network variable, b_ij , as written in equation (4).

In order to get the best approximation of the posterior distribution $\mathbb{P}(\boldsymbol{\omega}|\mathbf{X}, \mathbf{Y})$ with $\mathbb{Q}(\boldsymbol{\omega})$, the Kullback–Leibler (KL) divergence should be minimized. The KL divergence to be minimized can be expressed as

$$\begin{align} \mathrm{KL} = -\int \mathbb{Q}(\boldsymbol{\omega}) \log\left[\mathbb{P}(\mathbf{Y}|\mathbf{X},\boldsymbol{\omega})\right]\mathrm{d}\boldsymbol{\omega} + \mathrm{KL}\left[\mathbb{Q}(\boldsymbol{\omega})||\mathbb{P}(\boldsymbol{\omega})\right]. \end{align} \tag{ 8 }$$

The integral on the first term on the right hand side can be approximated by drawing samples from $\boldsymbol{\omega_n} \sim \mathbb{Q}(\boldsymbol{\omega})$ and performing the Monte Carlo integration on the negative log-likelihood, $-\log\left[\mathbb{P}(\mathbf{Y}|\mathbf{X},\boldsymbol{\omega_n})\right]$. The second term on the right hand side is approximated to be $\sum_{ij} \frac{p_{ij}l}{2} ||\mathbf{w_{ij}}||^2$ where l is the prior assumption of the length scale of the distribution. We can take the prior assumption of small length scale to be able to capture high variability region better and therefore making the second term small.

With various approximation above, the KL divergence in equation (8) to be minimized can be expressed as

$$\begin{align} \mathrm{KL}\approx -\frac{1}{N}\sum_n \log\left[\mathbb{P}(\mathbf{Y}|\mathbf{X},\boldsymbol{\omega_n})\right]. \end{align} \tag{ 9 }$$

By defining the negative log likelihood as the Huber loss function, we obtain that minimizing the KL divergence in the equation (9) is equivalent to minimizing the loss function in equations (1) and (2). Therefore, the optimized parameters after the training can be used to approximate the posterior distribution of the weights in the form of $\mathbb{Q}(\boldsymbol{\omega})$.