Robust Errant Beam Prognostics with Conditional Modeling for Particle Accelerators

SNS accelerator employs a DCM [42, 43] to protect the Super Conducting Linac (SCL) from adverse effects of beam losses [44]. Figure 1 shows a schematic diagram of the DCM. It continuously monitors upstream and downstream beam current waveforms to detect any beam loss across the SCL. When significant beam loss is detected, it aborts the beam to minimize possible damage or activation. DCM uses a Field Programmable Gate Array (FPGA) and a dedicated communication line with the Machine Protection System (MPS) to enable fast aborts. Initially, the DCM only streamed a subset of the waveforms to a server for storage. At regular intervals, normal waveforms were stored and waveforms right before, during, and right after an errant beam pulse. All waveforms of all beam pulses are streamed to an edge computer for inferences of models and to store the data.

With this study, we aim to predict errant beam pulses before they occur using normal prefault pulses. Our previous studies  [4, 3] have shown that there are precursors in the normal pulse to indicate an impending errant beam pulse. To achieve this, the DCM was programmed to archive beam current waveforms preceding errant beam pulses and regularly archive waveforms during normal operation. The prefault normal waveforms are labeled as anomalies and waveforms from the normal operation are labeled as normal (reference) for the ML study.

We also recorded a snapshot of beam configuration settings along with the beam current data for conditional modeling. Eight relevant beam tuning parameters are considered for this study, namely, Ramp-Up-Pulse-Width, Ramp-Up-Pulse-Width-Change, Ramp-Up-Width, Beam-On-Width, Beam-On-Pulse-Width, Ramp-Down-Pulse-Width, Gate-Source-Offset, and Beam-Ref-Gate-Width. The description of these parameters is included in the Appendix.

Figure 2 shows an example of the beam current waveform. The top plot shows a series of macro-pulses, a 1 ms long pulse repeated at 60 Hz. This macro-pulse consists of approximately 1000 mini-pulses. Each mini-pulse is $\sim$650 ns, followed by a gap of $\sim$350 ns. Within each mini-pulse are the micro-pulses, not shown, which are the Radio Frequency (RF) buckets filled with the beam particles spaced at 402.5 MHz. The bottom plot shows a beam pulse’s current waveform with the initial ramp-up in intensity in the beginning of the macro-pulse, as well as the different widths of the mini-pulses during the macro-pulse. This setup is typical for a production-style beam.

The waveform has a length of about 1.2 ms and is digitized at 100 MS/s. The beam pulse’s timestamp, beam configuration parameters, and cycle ID are collected and archived along with the upstream and downstream beam current waveforms.

Each macro pulse comprises $\sim$120K time steps; however, for this study, we only use 10K time steps (from time step 3K to 13K) as our research indicated that the pre-cursors for an upcoming errant beam pulse are prevalent in this region. Furthermore, a peak counter approach was employed to exclusively select beam current waveforms associated with neutron production and avoid outliers such as test beams or beam studies. This technique leverages the “find peaks” method in the SciPy library [45]. The parameters for peak identification were set with a minimum height threshold of 2 mA and a stipulated minimum distance of 75 steps, equivalent to 750 ns, between adjacent peaks. The waveforms of the macro-pulse must encompass a minimum of 900 mini-pulses and exhibit a consistent repetition frequency of 60 Hz to indicate that these pulses were produced during full power beam and not during physics studies or testing. These strict requirements effectively filter out any potential outliers in accelerator conditions.

The dataset is divided into three parts, specifically training (64%), validation (16%), and testing (20%) according to their beam setting configurations. All data collected during normal operation was considered normal data. For SNN training, since the model requires two input samples, we developed a training data generator using tensorflow data generator [46] to create combinations of normal-normal data labeled as 0 and normal-prefault data labeled as 1. The data generator has two hyper-parameters: the number of anomalies per batch and the samples per anomaly. The former controls the number of unique prefault beam current samples in a batch, and the latter controls the number of normal beam current waveforms paired with each anomaly. The number of normal-normal samples matches the number of normal-prefault samples per training batch in the data generator. It is important to note that only samples from the same beam configuration are paired with each other. The data generator maintains separate counters for both normal and anomaly waveforms for each configuration setting to iterate over them in sequential order during training and evaluation. To introduce randomness, the normal and anomaly data are shuffled at the end of each epoch.

VAEs [6] are commonly used for anomaly detection [47, 35] when there are few anomaly examples. VAEs are normally used for semi-supervised learning and can address imbalanced classification problems. VAEs are trained only on normal data to model the normal behavior of the system. The model identifies any anomalous data using the reconstruction error at inference time. A typical VAE consists of two Artificial Neural Networks (ANNs), an encoder $\Psi$, and a decoder $\Phi$. The $\Psi$ maps input space $x$ into a reduced representation $z$ that is forced to follow a specific well-known distribution, such as a Gaussian distribution. The latent $z$ is generated by sampling from $\mu$ and $\sigma$ parameters estimated by the $\Psi$. A new parameter $\epsilon$ is also introduced to reparametrize the sampling layer $z$ and allow the model to backpropagate through the entire network. The latent variable is then fed to the $\Phi$ to reconstruct the input data. The loss function for a VAE is motivated by variational inference [48] by minimizing the Kullback–Leibler divergence (KLD) [49] between the posterior $p(z|x)$ and the encoded prior distribution $q(z)=\mathcal{N}(0,1)$:

where the first term is the reconstruction error, the second term computes KLD, and $\eta$ is the harmonic parameter to balance the two.

The SNN model takes a different approach to anomaly detection than traditional classification methods, or VAEs, by learning the similarities between a pair of inputs. Instead of focusing on classifying individual data samples, it seeks to capture similarities (or dissimilarities for anomalies). This is achieved by encoding two separate input vectors using twin networks (sharing weights and biases) and then comparing the reduced latent space using a distance metric, as illustrated in Figure 3. The twin networks are deployed to learn and extract the relevant features for similarity prediction between normal-normal and normal-abnormal combinations. As it learns the relevant features, it produces a reduced representation of the original inputs that are then compared using a Lambda difference layer as described in Equation 2.

Here, $h_{x_{1}}$ and $h_{x_{2}}$ are reduced representations of reference and inference inputs, respectively. The output from the difference layer is fed into a multi-layer perceptron (MLP) block with dropouts to avoid over-fitting before the final output layer that produces scalar similarity score between the two inputs.

During training, predicted similarity scores are compared with the ideal scores (labels) using a modified contrastive loss function described in Equation 3.

Here, $y$ is the true label, $\hat{y}$ is the predicted similarity score, $\alpha$ (set to 5.2 for this study) is a tuning parameter used to emphasize the similar pulses, and $\beta$ is another tuning parameter used to emphasize dissimilar pulses (set to 1 for this study).

This study uses a ResNets [50] model for the twin network with three stacked convolution residual blocks. Each residual block consists of skip connections between two convolution layers to improve the information propagation in the deep model. Each ResNet block consists of a convolution layer, Rectified Linear Unit (ReLU) activation layer, a batch-normalization layer, and a dropout layer. The model architecture parameters are provided in Table 2.

Conditional models learn a relationship between input and output data under given conditions. These conditions can be explicit features or variables influencing the relationship between inputs and outputs. In other words, the model’s predictions depend on specific conditions or contextual information. We employed conditional models that take beam configurations as input to an orthogonal MLP to handle the beam current data belonging to different beam configurations. We compared both CSNN and CVAE for our application.

CSNN consists of the same architecture described in Section 4.2 with an additional conditional input. The conditional input is then passed to an embedding sub-model as shown in red in Figure 4. The configuration embedding MLP consists of 3 dense layers with 64, 32, and 16 nodes respectively along with leaky relu activations and 5% dropouts. The embedding transforms the beam configuration values into a latent space representation, which allows the model to learn the beam configuration landscape and leverage any correlations. The embedded beam configuration is concatenated to the output from the difference layer as shown in Figure 4.

CVAE [10] was proposed to make diverse predictions for different input samples. A CVAE is a type of VAE that introduces a conditional variable into both the encoder and the decoder. The objective function of the VAE can be modified by adding the variable $c$,

where we condition all of the distributions with $c$. CVAEs are an extension of VAEs by adding the conditional part in the $encoder$ and $decoder$ to associate the input samples with labels. Therefore, at inference time, we have more control to generate samples that belong to specific conditions in contrast to a VAE that does not have control over the generated samples.

All the models in this study are implemented using Tensorflow 2.8 [46] with Adam optimizer [51] for training. For efficient optimization, we used a learning rate decay (callback provided by tensorflow) of 85% using the validation loss plateau with a patience value of 5 epochs. To automatically tune the number of epochs required for convergence, we used an early stopping callback at the validation loss plateau with a patience of 20 epochs.

The goal of this study is to show that ML can improve beam availability and minimize damage to the accelerator by developing advanced fault prediction models and comparing their performance. Any anomaly prediction technique we deploy should not significantly increase the levels of beam aborts or abort duration due to false alarms. Hence, our evaluation metric is built to keep a fixed minimal False Positive Rate (FPR) while achieving a significant True Positive Rate (TPR). Even a modest level of FPR, such as 2%, would noticeably reduce the beam availability and negatively affect the science research program. Based on input from the operators and subject matter experts a FPR of 0.1% is considered the upper acceptable limit.

Deep learning models can be built with a wide range of different configurations and architectures, producing varying results. It is essential to explore and choose the best hyper-parameters and model architecture for a given task.

There are multiple HPO and/or NAS toolkits that are available for researchers to use and produce an optimal architecture(s). In this study, we use the Tree-based Parzen Estimator (TPE) [52] algorithm available in the HyperOpt [53] library in python in combination with ml-flow [54] to perform HPO and NAS studies on both Siamese and VAE models in order to pick the best model architecture for the comparison. The model initialization was fixed for this process. The objective was set to minimize validation loss with 50 trials to be consistent with the HyperOpt usage in the literature.

In addition to the numerical hyper-parameters and model architecture parameters shown in Figure 5, we also explore activation functions between ReLU, and Leaky_ReLU for both the models. Based on this study, we compared the model architectures and hyper-parameters that produced the best TPR at the target FPR of 0.1% in the HPO/NAS study.

The models were trained on the above-described beam current data set from DCM with the training data generator to learn to produce a low similarity score between normal pulses collected during normal operation (named normal) and normal pulse preceding faults (named anomaly).

We trained eight individual SNN models on data sets belonging to each beam configuration and a single CSNN model on the entire data set using beam configurations as conditional input. To statistically evaluate the performance, stability, and robustness of the model architecture, we trained each model (SNN and CSNN) 10 times with the same architecture (chosen after architecture optimization as described in Section 5.2), using different weight and bias initialization. Figure 6, compares the prediction performance between SNN and CSNN models regarding TPR at FPR of 0.1%. In the plot, each blue violin represents an ensemble of SNN model inferences on the test data of the beam configuration they are trained on. The corresponding purple violins represent the ensemble of CSNN model inferences on the same test data. The beam configuration ID are marked on the x-axis. As shown in Figure 6, the CSNN outperforms SNN model trained on a single configuration for our application. We believe that this improved performance is due to the CSNN model being able to leverage the correlation between beam current waveforms from different configurations and hence can learn from an increased amount of data.

To understand the stability and robustness of our models, we trained an ensemble of 10 CVAE models with different weights and bias initialization on the same data set. Figure 7 shows the predictive results of comparing CSNN and CVAE regarding TPR at FPR of 0.1%. For all the configurations, CSNN is seen to outperform CVAE. Even though both models, in this case, are trained on the same data set and use beam configuration as a conditional input, we believe that the CSNN model, being a supervised learning method, can leverage the label information and is much more powerful than CVAE. In addition, the CVAE model needs to reproduce all the features to reconstruct the normal waveforms. In contrast, CSNN does not need to learn to reproduce the whole waveforms; instead learns to extract the relevant features to distinguish anomalies from normal.

It is important to note that the two models also vary in architecture and how they learn to distinguish abnormal inputs from normal. However, deep learning models are often compared based on the number of trainable parameters they have. The model architectures are compared in Table 2; the number of trainable parameters in the CVAE (after a light HPO and NAS) is roughly an order of magnitude higher than that in the CSNN model. This indicates that the CVAE requires much more expressivity to learn to reproduce the normal data waveforms. On the other hand, CSNN leverages the label information and does not need to learn to reproduce the entire waveforms. This requires much lower trainable parameters and still outperforms the CVAE model.