Data-driven Approaches to Surrogate Machine Learning Model Development

The computational model Parmec [Koziara (2019)] is the underlying model which was the subject of a MLS model in [Jones et al. (2022)] and the same problem and base dataset is considered in this work. Parmec is employed to simulate the seismic response of the graphite core within the UK’s advanced gas-cooled reactor (AGR). This model consists of a simplified 3-dimensional representation of the AGR graphite core, including the positional arrangement of the graphite bricks and other components. Parmec can be used to simulate a range of different seismic scenarios with the resulting component translation, rotation etc. being calculated by the model.

In addition to the seismic configuration, another input to the Parmec model is configuration of cracked fuel bricks within the graphite core. Due to years of exposure to high temperatures and irradiation, some of the fuel bricks within the reactor are cracking, causing them to break into two pieces. The presence and configuration of these cracks has an impact on the reaction of the core to seismic loading. It is possible that up to 40% of the fuel bricks will eventually crack, although it is difficult to determine or predict where and when cracks will occur.

The relationship between crack configuration and seismic response of core components is complex, hence the Parmec model consists of many thousands of parameters and equations. In addition, there are over $10^{2500}$ possible permutations of crack configuration, assuming 40% cracking. With each configuration requiring around 2 hours to compute the seismic response via Parmec, it is clearly impractical to generate data for even a small percentage of them. Instead, industry practice is to generate random configurations of cracks, passing each through Parmec in order to build up a stochastic distribution of the seismic response.

In previous machine learning assessments of AGR graphite core seismic analysis [Jones et al. (2022)], each crack configuration is considered an individual data instance, with the encoding of cracked bricks being the input features and the response of core components to the earthquake being the output labels. The Parmec software generates a time-history of the earthquake response for all of the thousands of components within the core. For the sake of simplicity and focus, the MLS model was trained to predict the earthquake response for a single interstitial brick at a single time frame - see Figure 2.

To summarise the features of the MLS, each instance has an input size of 1988 with this being the number of fuel bricks within the AGR graphite. This input was arranged into a 3D tensor which retains physical positional relationships within the actual AGR graphite core (Figure 3). Each element is either a 1, -1 or 0 representing a cracked brick, uncracked brick or ‘empty’ position. The 3-dimensional encoding of the input features also allows the dataset to be used with a convolutional neural network [Albawi et al. (2017)] which was found to be the best performing type of machine learning model.

For the aforementioned study, a dataset of approximately 8300 instances was created using the random crack pattern generator and the Parmec software. Out of these instances, 6300 (75%) were used for training with the remaining 2000 samples retained for testing.

Data augmentation is frequently employed in classification problems within the field of machine learning [Shorten and Khoshgoftaar (2019)], where the model predicts a discrete category for each dataset instance. A classic example of classification is in computer vision, where a 2D or 3D tensor representing an image is used to predict a category that is depicted. For example, models trained on the ImageNet dataset [Deng et al. (2009)], which contains millions of images each representing one of 1000 discrete classifications, attempt to categorise the image depicted in an instance it is presented with. Classification is in contrast to regression problems, where there is a continuous, rather than discrete, output variable.

When dealing with problems such as ImageNet classification, performance is constrained by the size of the dataset. Model performance tends to improve with a larger number of training examples. A related constraint is dataset bias: regardless of the overall number of examples in the entire dataset, if one or more classes exhibits a significantly lesser or greater number of examples than the rest, model performance may be inhibited. Should there be a lack of examples for one particular class, not only will the model have difficulty identifying examples of this class, but performance for other classes will also be impacted.

Both of the aforementioned constraints tend to cause the phenomenon known as overfitting [Hawkins (2004)], where the model optimises too closely to the training data, including any noise or unrelated variability. There exist several methods to alleviate the effects of overfitting, including randomised nodal dropout [Srivastava et al. (2014)]. A commonly used solution to overfitting within image classification is data augmentation, where image manipulation techniques are used to generate additional data from existing examples.

Figure 4 shows an example of a data augmentation process on a single image instance. The image on the left of this figure would correctly be classified as a bird. This image can be used to generate five additional instances of the same class: three by rotation and two by mirroring. By applying this process to all images within a dataset, the number of available training instances can be multiplied by a factor of seven.

From the example, it can be seen that data augmentation techniques are highly suited to problems where the data is structured as a 2D or 3D tensor (such as greyscale or colour images, respectively). By extension, data augmentation is highly effective for applications in which localised or spatial patterns are of importance, for example where CNNs are employed. It should be noted at this point that the research work that we attempt to build on here employs 3D encoding of data as well as CNNs.

One of the most important model metrics to be selected during machine learning model development is the loss function. This function is used during training to calculate the difference between the ground truth and model prediction (known as the loss rate). Further, the derivative of the loss rate is used to iteratively update the model weights during the training and optimisation process.

Typically, one of only a few loss functions will be selected for regression model training. A common selection is the mean squared error [Wallach and Goffinet (1989)] or one of its derivatives such as root mean squared error. Other options include mean absolute error [Chai and Draxler (2014)] and Huber loss [Huber (1964)] which was found to be optimal in the preceding research work on this topic [Jones et al. (2022)].

An issue encountered in the aforementioned preceding research was that the data was not evenly distributed throughout the data space. The output data generated by the Parmec model tends to be distributed around a central modal value, with increasingly fewer examples towards the extremes of the data space. This in turn results in a model which tends to over-predict values in the lower part of the data space, and under-predict those in the upper part (Figure 5).

This problem can be compared with the issue of class imbalance encountered in the field of machine learning classification [Johnson and Khoshgoftaar (2019)]. Much literature has been written on the subject of correcting for data imbalance in classification, with a recent work [Sarullo and Mu (2019)] using a weighted loss function.

Conversely, for regression based problems, research attention has been scarce by comparison. Some recent works [Branco et al. (2017)][Yang et al. (2021)] note the lack of research on imbalanced regression have proposed solutions to problems caused by gaps or rarefactions in the data space. The solutions proposed involve the application of smoothing or dataset resampling. We note that the imbalance featured in these research works is not of the relatively smooth trend seen at the bottom of Figure 5.

Whereas the data distribution in the aforementioned research papers contain discontinuities and irregular gaps, our dataset follows a regular pattern. This makes the methods explored previously in this area potentially unsuitable for the research problem at hand. If we wish to counteract the data imbalance problem in the research at hand, a bespoke custom loss function will have to be developed.

Transfer learning is a machine learning technique that focuses on using already trained models to serve a purpose beyond their original intent. Often in machine learning, models are trained from scratch, a process that consumes significant resources in terms of time and computation to achieve optimal performance. Transfer learning can serve to make the process more efficient and less resource intensive by using the knowledge from the pre-trained models in the training process.

In [Jones et al. (2022)], an optimal model architecture was developed for the purpose of surrogation of a nuclear engineering model. Using this architecture and the training data, the parameters of the model were optimised from random starting weights. As training started from random weights, the process of model optimisation was repeated multiple times. The best performing models from this study, along with their pre-trained weights, can be transferred to this study as a starting point for exploitation of the methods described in subsections 2.3 & 2.4.

Recall from Figure 3 that the input features for this research problem is a 3D tensor representing the position of cracked fuel bricks within a given instance. This tensor is comparable to that of a colour image such as the one shown in Figure 4. Being a tensor of a similar encoding to that in the aforementioned figure, the same image manipulation techniques can also be applied to the tensor for this research problem.

Recall also from Figure 2 that the output labels for this research problem are continuous values representing displacement in the central brick of the core. Note also from the aforementioned figure that the overall data space from which the output variables are extracted is planar and can be expressed as a 2D tensor. For a given instance i.e. input/output pair, we can apply any of the rotational or mirroring techniques shown in Figure 4. Note that if we rotate or mirror about the vertical axis at the centre-point of the core, the encoding order of the input features will change, but the output label will not (as it represents the top brick in the central channel). Hence, data augmentation for the dataset in this study will create new instances with restructured feature tensors, but the same output label value - see Figure 6.

If we apply rotation (90°, 180°and 270°) and mirroring (vertical and horizontal) to all examples in our dataset, we can multiply the available number of instances by a factor of six.

As mentioned in subsection 2.3, the motivation behind using image manipulation techniques to create an augmented dataset is based on the conjecture that symmetries do exist in the Parmec data space. What is the justification behind the belief that rotational and symmetric manipulations of Parmec inputs would yield similarly transformed outputs? There are three lines of evidence that can be used to support the validity of this process:

1. 1.

   AGR Design: The design of the AGR (and the Parmec model that is based on it) contains four-fold symmetry [Nonbøl (1996)]. This effectively means that each quarter of the AGR (and Parmec) model is a rotation or mirror of the others. However, this is an incomplete justification as the simulated earthquake always impacts the model on one particular point on its periphery i.e. it is not symmetric.

2. 2.

   Dataset Observation: Observe Figure 7 which shows the average output value for each brick in the top layer of the Parmec model. It can be seen that there is a symmetry across both the vertical and horizontal centre-lines. This pattern is of the same four-fold symmetry as mentioned in the first bullet-point.

3. 3.

   Augmented Equivalent Data: Through the Parmec software package, we have the benefit of creating ground truth equivalents of augmented data. For a given ground truth example, we can manipulate the inputs according to one of the transformations shown in Figure 6 and then feed them through the Parmec model. Then we can compare the outputs generated by Parmec and our non-Parmec data augmentation technique.

We performed the process described in bullet point 3 above for two base instances, applying each of the five manipulation techniques on the inputs and then using these to generate labelled examples using Parmec. Simultaneously, we apply all five of our non-Parmec data augmentation techniques to the inputs and outputs of both examples. Comparing the outputs of both techniques, we notice agreement when applying rotation by 180°and when mirroring about the horizontal axis.

The validity of the augmentation technique discussed here will ultimately be tested through an experimental machine learning process. We can train two machine learning models of the exact same architecture and parameters: one with a dataset augmented with and one with the base dataset. A separate testing set of only non-augmented instances will be retained for testing both models which both models can best tested against. The performance of each image manipulation method can be evaluated in this way.

In [Jones et al. (2022)], the effectiveness of using three alternative loss functions were compared. These were the mean squared error (MSE), mean absolute error (MAE) and the Huber loss [Huber (1964)]. A model trained using the Huber loss function was found to produce the best performance. However, the other two loss functions produced a similar, albeit poorer, performance. Each of the three loss functions (Figure 8) each have their own strengths and weaknesses, with MSE heavily weighting outlying values, MAE proportionately weighting outliers and Huber being somewhere in between. As mentioned in subsection 2.4, our base dataset is highly centred around a central value with a double tailed distribution. Regardless of the loss function used, the resulting model is biased towards the central region, resulting in difficulty predicting at the upper and lower extremes (Figure 5).

We propose a loss function tailored to this dataset which applies an adjustment factor that is a function of model prediction distance from a central value.

As can be seen from (1), the regular loss, which takes the mean of the square difference between the ground truth $y_{i}$ and the model prediction $\tilde{y}_{i}$. This is then adjusted by the factor $Z_{i}$, as given by (2) and is calculated for each instance $i$. This term is based on the probability density function which in this case describes a Gaussian distribution. This Gaussian distribution is fit to that of our data distribution (Figure 9) by use of the mode $\mu$ and standard deviation $\sigma$. This summation is scaled through division by the square of $\beta$, which denotes the maximum such $Z_{i}$. Also included is the magnitude coefficient $\alpha$ which will be optimised experimentally.

This function results in a strong adjustment for predictions made near the mode, quickly dropping off to one as we move away in either direction. This adjusted loss function penalises predictions made near the region where there is a large concentration of training data and instead pushes it towards the extremes. This equation is designed with the intention of counteracting the bias in the data distribution.

In [Jones et al. (2022)], multiple machine learning models were developed and trained on the dataset. In the best performing case, a convolutional neural network was utilised with a refined architecture. Optimal performance was achieved when including inputs representing the cracking status of the top 3 levels of the AGR core. The model was used to predict displacements in a single brick on the top level of the core at a single time point during an earthquake - the distribution of these outputs for the dataset can be seen in Figure 9.

Out of all of the models produced in the aforementioned paper, the best performing five models were obtained as well as their optimised weights. We evaluated each of these models, numbered M1 to M5, against a testing dataset with the results shown in Table 1.

These models will form the basis of transfer learning experiments through refinement using the methods described in subsections 3.1 & 3.2. As mentioned in subsection 2.5, transfer learning can reduce the time and computational resources required compared to starting from randomised weights. Transferring models and weights in this case will not only reduce resource requirements but also be interesting from a research perspective.

We began by evaluating the effectiveness of each image manipulation technique mentioned in subsection 2.3. To do this, we augmented the base dataset using each of the five image manipulation techniques. This resulted in six available data sets: three using rotation, two using mirroring and the original unaugmented set. A summary of the datasets is shown in Table 2.

For the purposes of comparison, a model design was selected which is highly simplified compared to that utilised in [Jones et al. (2022)]. This simplification was made in order to reduce computational demands and to allow obtainment of results quickly. As our intention at this point is to compare the effectiveness of using different datasets and not overall optimisation, a simplified model is acceptable for this purpose.

The neural network architecture that was selected for this part of the research can be seen in Figure 10. Between the input and output layers, there is one convolutional layer followed by two dense layers. Activation functions and the use of dropout was utilised based on previous experience. The Huber loss function was used for back propagation and optimisation during training.

Six experiments were performed which involved training the model shown in Figure 10 individually with the datasets listed in Table 3. A 10% sample of the unagumented dataset (D0) was retained for validation and the model was trained until reaching convergence in terms of validation loss. For each experiment, the training process was repeated 32 times, each time initialising with randomised starting weights and dropout nodes. Each model was then evaluated using a separate testing dataset with the results given in subsection 5.1.1.

The datasets from the first phase were combined incrementally in the order of effectiveness as per Table 4. Again, each experiment is repeated 32 times. The results of this experiment are given in subsection 5.1.2.

The purpose of this experiment is to evaluate the bespoke loss function defined in subsection 3.2 and shown in (1)  & (2). The simplified architecture used in the previous section and shown in Figure 10 was again used in the experiments in this section. Initially, the unaugmented dataset (D0 from Table 2 was used for training and the same testing set as used in subsection 4.1 for evaluation.

In addition, we will make adjustments to the (1) & (2) and train models using the same parameters. This will not only allow us to refine the function but also understand the impact of each component of it. In this part of the experiment, the alpha coefficient ($\alpha$) is set to unity.

The next phase of research involved the use of our custom loss function in combination with the augmented datasets as detailed in subsection 4.1. All experiments outlined in Table 6 use a training set which includes all augmented datasets as per E2.4. Further, we adjust the alpha coefficient in the remaining parts of this experiment, with $\alpha$ having a value of unity in E4.0 (as it was in experiment 3). E4.1, E4.2 & E4.3 each increment $\alpha$ by unity in turn.

As per experiments 1 & 2, we will repeat each experiment 32 times to account for the stochastic nature of weight initialisation. A summary of the experiments performed in this section is given in Table 5. The results are reported in subsection 5.2.

Subsection 3.3 discusses five pre-trained models obtained using the methodology of a previous research work in this field [Jones et al. (2022)]. The performance of these models against a testing dataset is summarised in Table 1.

The intention of this experiment is to further train these models using the methods detailed in the previous two subsections (4.1 & 4.2). We begin by further training models M1 to M5 using our dataset enlarged by all augmentation methods as per E2.4. We then combine both augmentation and the use of our custom loss function as defined in Equations1 & 2. Finally, we perform the same further training of the transferred models using the conditions of the previous experiment but with an $\alpha$ coefficient of two.

This experiment is summarised in Table 7. We repeat the training process six times for each part of the experiment as opposed to the 32 times performed in earlier experiments. This is as only the dropout nodes are randomly selected and we are not initialising he weights. The results of this experiment will be presented in subsection 5.3.

As mentioned in subsection 4.1, two augmentation experimental approaches are attempted. The first tests each image manipulation technique individually, the second exploits combinations of these techniques. The following two subsections report the results of these approaches, respectively.

An experimental approach to the development of a loss function customised to the needs of the data problem at hand was discussed in subsection 4.2 with a summary of proposed experiments shown in Tables 5 & 6. The results are summarised in Tables 10 & 11 with a visual summary of E4.1 shown in Figure 12. All parts of experiment 3 show similar mean squared error to that of the baseline case (E1.0) with little variation. The addition of all augmented datasets to the training set (E4.0) yields similar model performance to that of experiment E2.4 which uses the baseline loss function. Increasing the $\alpha$ coefficient above unity appears to increase mean squared error.

We discussed in subsection 4.3 that our five existing models with their pre-trained weights (M1 to M5) were further trained using methods developed in this research work. The performance of these models further trained on our augmented training set is shown in Table 12. Comparing these results with the original model performance (Table 1) it can be seen that a significant improvement has been achieved.

Looking next at transfer models trained with the aforementioned augmented dataset and also using our custom loss function, very similar performance results are seen (Table 13).

Finally, performing the same process as the aforementioned experiment but with the $\alpha$ set to two, we see the performance of the models summarised in Table 14 with a visual comparison with the original M3 model seen in Figure 13 . Comparing E5.1 & E5.2 with E5.3, it initially appears to perform slightly worse in terms of MSE test loss. From the comparison figure, multiple improvements over the original model performance can be seen, including a closer fit with fewer examples falling outside of both the 10 and 20 point margins. Also, the model from E5.3 performs considerably better than its original counterpart at the upper and lower bounds of the data space.

Looking first at the augmentation experiments (see subsection 5.1), we can see that all dataset augmentation has a positive effect on model performance. The results of experiment 1 (Table 8) show that the addition of some augmented datasets have more of an effect than others. For example, experiments E1.4 & E1.2 (rotation by 180°and vertical mirroring) both see about a one third reduction in mean squared error in comparison to the baseline (E1.0). It is interesting that the two best performing augmentations both make adjustments about the vertical axis. Conversely, the least performing augmented dataset (D5, included in E1.5) makes its adjustment about the horizontal axis.

It is clear from experiment 2 (Table 9) that the combination of all augmented datasets to the training set yields the best result (E2.4). A deeper analysis of the performance of the model produced in experiment E2.4 can be made by comparing the left and right parts of Figure 11. It can be seen that a model trained on all augmented datasets produces an overall better fit than when using no augmentation at all. Key indicators include the percentage of data within 20 points of the ground truth - 98% for the best performing model from E2.4 compared with 94% for a similar model from E1.0. The inclusion of all augmented datasets also appears to improve model prediction in the upper part of the data space - compare the second section from the top in Figure 11. However, the best model from E2.4 also performs worse on predictions in the lower part of the data space (compare bottom parts of the aforementioned figure).

Recall that the model architecture used in these experiments is highly simplified compared to previous works in this field. We should note at this point that the results of experiment 2 are an improvement even over those of highly complex models from [Jones et al. (2022)] - compare with the performance of M1 to M5, for example (Table 1). This suggests that dataset size is far more important in this problem space than model complexity and refinement.

From an initial assessment of experiment 3 (Table 10), it appears that the use a custom loss function as defined in Equations 1 & 2 has no advantage compared to the baseline (E1.0). All of the reported mean squared error test values are within a few percent of one another.

Experiment 4 combined our custom loss function with all augmented dataset from experiment 2. A model trained with the combined custom loss function and full augmented data set (E4.0) has a similar performance to that of a model when using the augmented set alone (E2.4). Increasing the value of the coefficient $\alpha$ produces an increasing optimal and mean test performance.

So is there any value in the method evaluated in experiments 3 & 4? To answer this question, we must return to the motivation behind this method as detailed in subsections 2.4. We discuss the fact that our data space is concentrated in a central region with increasing rarefaction as we move away from it. Consequently, our models trained on this data performs poorly near the extremes as can be seen from the left part of Figure 11. Comparing the aforementioned figure with a visualisation of the optimal model from experiment E4.1 ($\alpha$ = 2), we can see significant improvements in model performance at the upper and lower extremes of the data space (Figure 12). This advantage comes at the expense of minor reductions in overall performance metrics, for example percentage of examples that are predicted outside of 20 percentage points of the ground truth (rising from 2.0% in E2.4 to 2.6% in E4.1). Whether or not this trade-off is of value will dependent on any practical application of the model. Nonetheless, it is likely that analysis using any such model would value good performance near the extremes as events in these regions are are most likely to impact safety. Therefore, we highlight the potential of this method and carry it forward into experiment 5.

We can see from the results of experiment 5.1 (Table 12) that further training our pre-trained models with an augmented dataset results in a significant performance uplift. For example, M1 sees its test mean squared error drop from 9.28e-3 to a minimum of 5.70e-3 - a performance improvement of about 38.5%. Combining the augmented dataset with our custom loss function and further training our base models appears to have little effect beyond what we saw in E5.1 (Table 13). Repeating E5.3 with an alpha coefficient of 2 initially appears to worsen performance, with test MSE rising from 5.70E-3 to between 6E-3 and 6.40E-3. However, on inspection of the visual results of E5.3 and comparison with those of E5.1 & E5.2, we see significantly improved performance at the upper and lower boundaries (above 0.6 and below 0.2 on the normalised scale of ground truth). Hence, we present the optimal results from E5.3 in Figure 13. Some notable improvements include a halving in the number of predictions that fall outside of the 20 percentage point margin (4.0% in the original M3 against 1.8% post E5.3), an increasing percentage of samples correctly placed over the 0.6 boundary (40.9% in the original M3 against 54.8% post E5.3) and finally a similar improvement below the 0.2 boundary (16.1% against 35.1%). As touched on earlier, in the field of machine learning for nuclear energy, this represents noteworthy progress as safety decisions are likely to concern events at the extremes i.e. very high or very low values.