Uncertainty quantification for deep learning in particle accelerator applications

With the advent of increased computational resources and improved algorithms, machine learning-based models are being increasingly applied to complex problems in particle accelerators. However, such data-driven models may provide overly confident predictions with unknown errors and uncertainties. For reliable deployment of machine learning models in high-regret and safety-critical systems such as particle accelerators, estimates of prediction uncertainty are needed along with accurate point predictions. In this investigation, we evaluate Bayesian neural networks (BNN) as an approach that can provide accurate predictions along with reliably quantified uncertainties for particle accelerator problems, and compare their performance with bootstrapped ensembles of neural networks. We select three accelerator setups for this evaluation: a storage ring, a photoinjector, and a linac. The problems span different data volumes and dimensionalities (e.g., scalar predictions as well as image outputs). It is found that BNN provide accurate predictions of the mean along with reliable estimates of predictive uncertainty across the test cases. In this vein, BNN may offer an attractive alternative to deterministic deep learning tools to generate accurate predictions with quantified uncertainties in particle accelerator applications.

Figure 1

Test cases for utilized in this investigation: (a) Schematic of the SPEAR3 storage ring, (b) the LCLS-II Injector frontend, and (c) the LCLS electron beamline.

Figure 2

Deterministic neural network, ensemble and Bayesian neural network model mean predictions and predicted standard error for SPEAR3 measured data.

Figure 3

Deterministic neural network, ensemble and Bayesian neural network model mean predictions and predicted standard error for SPEAR3 simulation data.

Figure 4

Deterministic neural network predictions; Ensemble and Bayesian neural network model mean predictions and predicted standard error for LCLS-II injector front end simulation data.

Figure 5

Histogram outlining the distribution of log emittance in the learning dataset.

Figure 6

Comparison of (a) Deterministic neural network, (b) Bootstrapped ensembles, (c) Bayesian neural network predictions for the $x$-emittance for the LCLS electron beamline.

Figure 7

Histogram outlining the distribution of log emittance in the learning dataset.

Figure 8

Outlining the clusters of data using PCA.

Figure 9

Comparison of deterministic NN predictions on (a) in-sample test dataset samples and (b) OOD dataset samples.

Figure 10

Comparison of Bayesian neural network predictions on (a) in-sample test dataset samples and (b) OOD dataset samples.

Figure 11

Schematic for the encoder-decoder architecture for the LCLS-II injector longitudinal phase space image prediction.

Figure 12

Simulation test images, Bayesian neural network predictions and predicted standard error bounds for the LCLS-II injector longitudinal phase space.

Figure 13

Simulation test images, neural network ensemble predictions and predicted standard error bounds over the ensemble of deterministic neural networks for the LCLS-II injector longitudinal phase space. (a), (b), (c) and (d) are representative instances, randomly chosen.