Low-Energy Electron-Track Imaging for a Liquid Argon Time-Projection-Chamber Telescope Concept using Probabilistic Deep Learning

We prepared the training data used to develop the models presented in this paper in a two-step process. We used the PENetration and Energy LOss of Positrons and Electrons (PENELOPE) code (NEA, 2019) to generate raw electron tracks, and then subsequently applied detector-response effects with a custom code. PENELOPE provides a detailed microphysical simulation of electron energy loss processes in matter, including simulations of elastic and inelastic scattering off of other particles in the media, bremsstrahlung emission, and generation of delta-rays. The custom code voxelizes the raw-track energy depositions, simulating digitized pixel readout with a specified pixel pitch determining the $X-Y$ pixel spacing. The $Z$ spacing is determined by the specified electric field, electron mobility, and digitization frequency. The electric field is set to 500 V/cm, and the electron mobility is calculated according to Equation 21 from Li et al. (2016). The digitization frequency is then set to achieve an effective $Z$ spacing equal to the $X-Y$ spacing. For the track head reconstruction, we simulated data with a 500 $\upmu$m pixel pitch at 4 different drift depths (1 cm, 5 cm, 10 cm, and 20 cm)²²2These parameters are a good starting point based on initial instrument design considerations, but may eventually change as the design evolves.. The training data for the track direction reconstruction also uses a 500 $\upmu$m pitch, and is simulated at a drift depth of 5 cm. Each sample therefore consists of a 3D array of floating-point numbers, representing a simulated pixel readout at each position in a 3D grid. The size of this grid in each training data set is given by the bounding box of the largest sample in that data set. All other simulated charge depositions are zero-padded symmetrically to match that dimension.

All simulations compute detected quanta from first principles based on raw energy depositions in a detailed manner similar to that of the Noble Element Simulation Technique (NEST) (Szydagis et al., 2011). All simulations also have a transverse diffusion coefficient of 12 cm²/s, taken from Li et al. (2016). The longitudinal diffusion coefficient is computed from Equation 23 of Li et al. (2016), and is approximately 6 cm²/s, giving $\sim$0.4 mm root-mean-squared (RMS) diffusion over a 10 cm drift. We also apply an effective electronic noise of approximately 30 electrons, and a simulated digitization threshold of 3$\sigma$ above noise. We simulated a number of electrons equal to within 6% at 5 energies (50, 300, 500, 750, and 1,000 keV), obtaining 24,472 simulated electrons in total. We then split the data randomly into training and testing sets in a 90:10 ratio for the track head and direction deterministic models, and a 92.5:7.5 ratio for the track head uncertainty quantification model.

We assess subsequent performance of the trained models not only with standard measures such as accuracy and coverage, but also with figures-of-merit for the quality of the overall Compton event reconstruction. We do this using data simulated with the Medium-Energy Gamma-ray Astronomy library, or MEGAlib (Zoglauer, 2019). After implementing the detector geometry using MEGAlib’s built-in geometry specification, we generate 1,000 keV gamma-rays originating at 64 degrees from the zenith using the MEGAlib FarFieldPointSource generator. We then process the generated energy depositions with custom code that applies the detector response, including a parameterization of the interaction locations (with uncertainties) based on the output of the trained neural networks on test electron track simulations. We present a description of this parameterization in Sections A.1 and A.2. We then feed this simulated detector output, in the form of a MEGAlib-specific file type, into the MEGAlib reconstruction code revan, which determines the time-ordering of the Compton scatters. At this point we remove any events where revan was unable to determine the first two Compton scatters correctly. Because events with an incorrect reconstruction of the scatter order tend to have significantly larger pointing error, this selection will produce a smaller pointing error than would occur in a real experiment. However, since we are primarily concerned with understanding specifically the behavior of the electron scatter position reconstruction, and not the Compton ordering reconstruction, we choose to do this so as to eliminate that possibly confounding source of error from our assessment of the performance of our models. These events are then then used by a modified version of the ComPair Python library (Moiseev et al., 2015) to compute the location of and uncertainty in the gamma-ray origin.

In the first model, we use a deterministic CNN-based model to predict the location of the electron track head. We report the architecture of the CNN used for electron track head reconstruction in Figure 3. The input to the model is a 3D image of an electron track from the detector. We include spatial dropout (Tompson et al., 2015) (as opposed to conventional dropout (Srivastava et al., 2014)) and Batch Normalization layers after each convolutional layer to ameliorate overfitting. We use a 3-dimensional global average pooling operation at the end of the feature extraction stage, as opposed to a conventional flattening operation, to reduce the parameters in the ensuing fully connected layers and improve generalization. The final output from the fully connected layers is a 3-dimensional regression prediction for the location of the electron track head. The model architecture and ancillary hyperparameters were established based on Bayesian Optimization (O’Malley et al., 2019), followed by subsequent manual fine tuning. This model was implemented with Keras (Chollet et al., 2015), using the Adam optimizer (Kingma & Ba, 2014), a mean-squared-error loss function, and was trained for 200 epochs on an NVIDIA Tesla A100 GPU at the SLAC Scientific Data Facility, which took about 10 minutes.

Since the model predicts the location of the $X$, $Y$, and $Z$ components of the electron track head independently, and treats each physical dimension identically, we can combine the predictions of all 3 dimensions and analyze them in the same distribution. The distribution of errors on the test data set in all 3 dimensions is shown in Figure 4, broken down by initial electron energy. However, from an applications perspective, the goal is a sub-mm accuracy of the electron track head location across all three spatial dimensions. Therefore, we add in quadrature the $X$, $Y$, and $Z$ error for each prediction to obtain the absolute error (i.e. the Euclidean distance between the true and predicted electron track head position) for each sample in the test data set. We show this distribution in Figure 5. Clearly the distribution of errors is wider for higher energy electron tracks. Because the beginning of an electron track tends to become straighter as the electron’s initial energy increases (this effect is illustrated in more detail in Section 3.3), one might expect that the beginning of the track should therefore become easier to identify, not harder. We hope to investigate this with a future interpretability study.

We plot the mean of each of these distributions in Figure 6, as a function of drift depth. As expected, the RMS error increases for all energies with drift depth, consistent with diffusion effects spreading out the signal and making it harder to pinpoint the track head location. However it is important to note that this increase with drift length is modest, and we readily achieve sub-mm accuracy up to at least a 20 cm drift length. There are no error bars given for the points in Figure 6 because the statistical errors are negligible, and the deterministic approach does not provide a good way of estimating a systematic uncertainty. We approach this problem with uncertainty estimation techniques in Section 3.4.

In Figure 7 we present an assessment of the deterministic model using the MEGAlib-simulated performance of the GammaTPC instrument. We apply a parameterization of the deterministic model output to the simulated data produced by MEGAlib and our custom detector-effects code, which is described in detail in Appendix Section A.1. The distribution of the Angular Resolution Measure (ARM), which is the angular difference between the true and reconstructed origin of gamma-rays, is peaked at zero degrees, and is shown fitted to a Lorentz distribution³³3We also attempted a fit to a Voigt profile, but the best fit contained no Gaussian component and was therefore also a Lorentzian.. The Full-Width-at-Half-Maximum (FWHM) of this distribution for a 5 cm drift depth (as shown) is 2.746°.

We also developed a model to reconstruct the direction of the initial electron scatter. Obtaining the direction is very useful, because, as explained in Section 2, breaking the axial degeneracy of the reconstructed gamma-ray origin can greatly improve the pointing accuracy of a well-reconstructed event.

The model is mostly identical to that used in Section 3.2. The difference is that the final 3 outputs shown Figure 3 are concatenated with the predicted electron track head position produced by the trained model from Section 3.2, and then reduced with a feed-forward layer back to 3 numbers which predict the $X$, $Y$, and $Z$ components of the scattered electron’s initial direction unit vector. In principle the initial direction vector has only two independent variables. Removing the degeneracy of the third fitted parameter may further improve the direction reconstruction in future studies. The model was implemented with Keras, using the Adam optimizer, a mean-squared-error loss function, and was trained for 20 epochs on an NVIDIA Tesla A100 GPU at the SLAC Scientific Data Facility, which took about 1 minute. In Figure 8 we show the performance of this model using a histogram of the cosine of the angle between the predicted electron direction and the true electron direction. A model that can reconstruct the direction with perfect accuracy would present as a delta function at 1 in this plot, while random guessing would produce a flat distribution. Since the distribution is skewed towards higher values of $\cos{(\delta)}$, this indicates it is able to extract useful information on the initial electron direction from the pixel data.

It can be seen in Figure 8 that the quality of the direction reconstruction increases with electron energy, as the $\cos{(\delta)}$ distributions peak more towards 1. Specifically, we find that the median angle between the true and predicted initial electron direction is 87°, 71°, 59°, 50°, and 48° for 50, 300, 500, 750, and 1,000 keV electrons respectively.⁴⁴4A median angle of 90° would indicate no ability to reconstruct the direction, while a median angle of 0° would indicate perfect reconstruction. This could be because high energy electrons travel in a straighter line initially, making it easier to establish the initial direction the electron was traveling. This effect is visible in Figure 9, which shows the cosine of the angle between each electron’s current and initial direction of travel, vs. the distance between its current and original positions. By comparing the 1,000 keV electrons (purple) and the 300 keV electrons (green), we see that the 1,000 keV electrons travel further before starting to deviate from their initial direction (i.e. moving away from a $Y$-value of 1). This trend of higher-energy electrons traveling straight initially for longer than low-energy electrons holds up as a general trend in this energy regime. This runs counter to the effect of energy on position reconstruction, where lower-energy scatters create more localized charge clouds, which facilitate better position reconstruction. The fact that these effects offset is encouraging, as it suggests that superior electron direction reconstruction for high-energy electrons could counteract lower-quality interaction position reconstruction.

Deep learning models are often considered to be high-fidelity models with accurate predictive powers. However, even predictions from a well-trained neural network may contain substantial errors and uncertainties due to bias, noise and complexity inherent in the data, the volume of the training data, the error minimum chosen during optimization, etc. In this light, it is important to be able to construct reliable prediction intervals in addition to point predictions. This can enable otherwise unavailable techniques when the data-driven model is to be used in scientific discovery. Reliable uncertainty quantification for deep learning has been identified as a priority for scientific machine learning (SciML) applications (Baker et al., 2019; Ihme et al., 2022). In our case, since electron tracks are highly random, they will have great diversity in shape, particularly at higher energies. It is reasonable to assume that the accuracy of a machine learning model may vary as a function of electron energy and/or track shape. The ability to select only tracks where the model has high confidence in its reconstruction should improve the overall pointing accuracy of the instrument.

There exist different algorithms for uncertainty quantification for deep learning models. These include purely Bayesian approaches, non-Bayesian approaches and hybrid methodologies. A priori, one is typically unable to determine the best algorithm for a specific application. As of yet, there have been very few studies focusing on evaluating, comparing and contrasting these approaches for scientific problems. This is exacerbated by the fact that most computer science studies comparing these algorithms have focused on classification problems, as opposed to problems requiring regression as is the case in many science applications, including this one. To this end, we evaluate different deep learning uncertainty quantification approaches on the data sets for electron track head reconstruction. These include Bootstrapped Ensembles of neural networks, Monte Carlo Dropout, Probabilistic Neural Networks (PNNs), and Evidential Deep Learning, a brief description of which follows. Additionally, we also outline the hyperparameters of the algorithm and their selection using the validation dataset.

Bootstrapped ensembles use bootstrap aggregation for uncertainty estimation. Bootstrap data sets of $N$ samples each are produced for training each individual model in the ensemble by sampling uniformly with replacement from the original training data set of $N$ unique samples. Individual neural networks in the model ensemble are trained separately on each of these bootstrap data sets. The predictions of the individual trained models in the ensemble are averaged to get a mean prediction. The variance in predictions provides a measure of the predictive uncertainty. Boostrapped ensembles are very commonly used in uncertainty quantification for scientific problems (Manly, 2018; Jeong & Maddala, 1993). Herein, the key hyperparameter is the number of constituent neural network models in the ensemble. In this investigation, the number of neural network models in the ensemble was increased until the coverage and MPIW became constant over the validation dataset, beyond a value of $15$ individual neural network models in the ensemble. Thus, the comparison and the comparative results correspond to a bootstrapped ensemble with $15$ neural network models.

Monte Carlo Dropout is a Bayesian technique introduced in Gal & Ghahramani (2016), where one approximates the network’s posterior distribution of class predictions by collecting samples obtained from multiple forward passes of dropout regularized networks. Dropout regularization (Srivastava et al., 2014) involves random omissions of feature vector dimensions during training time, which is equivalent to masking rows of weight matrices. Inclusion of dropout layers mitigates model overfitting and is empirically known to improve model accuracy (Srivastava et al., 2014). A key observation of Gal & Ghahramani (2016) is that under suitable assumptions on the Bayesian neural network prior and training procedure, sampling $N$ predictions from the BNN’s posterior is equivalent to performing $N$ stochastic forward passes with dropout layers fully activated. In this manner, the full posterior distribution may be approximated by Monte Carlo integration of the posterior predicted probability vector. Monte Carlo Dropout networks may also be interpreted as a form of ensemble learning (Srivastava et al., 2014), where each stochastic forward pass corresponds to a different realization of a trained neural network. In the usage of MC Dropout, key hyperparameters include the placement and number of dropout layers, the strength of the dropout rate (Kendall et al., 2015; DeVries & Taylor, 2018), the number of forward passes sampled during prediction, etc. During model selection, the best performance over the validation set was engendered by the architecture utilizing a dropout layer after each convolutional block and a single dropout layer after the global average pooling operation. Explicitly, in this architecture, no dropout was performed between the input and the first convolutional layer and neither between the last dense layer and the output. Additionally, in between the convolutional layers the best dropout rate was $0.25$ and in the dense layers the best dropout rate was $0.1$. All the hyperparameters were selected using cross-validation. While using higher dropout rates led to larger prediction intervals, they also had a deleterious effect on the model accuracy leading to poorer calibration of the uncertainty estimates. Beyond 50 samples for prediction, we did not observe any improvement in the metrics and thus, 50 samples in the prediction were used.

Probabilistic Neural Networks (PNNs) use a proper scoring rule (Gneiting & Raftery, 2007) to evaluate their performance. For regression tasks, this scoring rule is the negative log likelihood (NLL). Consequently, the training of the model ascertains network parameters that give an optimal mapping from the space of raw features to the mean and standard deviations of the predictions, so as to minimize the NLL over the training data set. Strictly speaking, this approach entails no additional hyperparameters and only requires a change to the cardinality of the output and the usage of the NLL as the loss function to be optimized. However, the learning rate of the Adam optimizer was adjusted to ensure optimal convergence for this algorithm as well.

Evidential Deep Learning (EDL) (Amini et al., 2020), refers to a class of deep neural networks that exploit conjugate-prior relationships to model the posterior distribution analytically. EDL methods have the immediate advantage of requiring only one single pass to access the full posterior distribution, at the price of restricting the space of possible posterior functions to that containing only those which have the appropriate conjugate prior forms. They also only require one to modify the loss function and the final layer of its deterministic baseline, which allows flexible integration with complex, hierarchical deep neural architectures. Under the assumption that the input data are drawn from a Gaussian distribution of unknown mean and standard deviation, the conjugate prior is then a normal-inverse-gamma distribution. In our case, we place an independent prior on each Cartesian coordinate of the electron track head position. For EDL-based regression, the hyperparameter to be determined is the coefficient for the regularization loss defined by Equation 9 of Amini et al. (2020) (thus, for EDL, total loss is the sum of the NLL loss and the product of the regularization coefficient with the regularization loss). While comparing the EDL model to the other UQ models, the regularization coefficient was utilized at a value of 1.

To ascertain the quality of the uncertainty estimates from the different algorithms, we compute three metrics: accuracy, coverage and sharpness. Accuracy denotes the fidelity of the mean prediction, evaluated via the RMS error on the test data set. Coverage (Romano et al., 2019) measures the calibration of the uncertainty estimates. The prediction interval is the estimate provided by a model of the range in which the target value will lie with a certain probability, given the input features. If the actual value lies in the predicted interval, it is said to be covered by the interval. Coverage measures the ratio of test samples that lie in the prediction interval of a specific quantile. A critical consideration for prediction intervals is their reliability at reflecting the quantiles in the distribution of the target quantity. As an illustration, a $95\%$ prediction interval that subsumes less than $95\%$ of the target samples will lead to predictions with higher error than expected, which could have significant consequences, depending on the application. Prediction intervals with incorrect coverage indicate poorly calibrated models, where calibration here indicates the degree to which the predicted uncertainty matches the true underlying uncertainty in the data. In our comparison between algorithms in Table 1, we utilize a uniform prediction interval of $95\%$ for all the algorithms and report the coverage at this prediction interval. Finally, sharpness measures the size of the prediction interval, at a fixed value of coverage. We measure the sharpness of the prediction intervals using the Mean Prediction Interval Width (MPIW) (Pearce et al., 2018), which is the average size of the prediction intervals over the test data set. We require that the prediction intervals provide a minimum value of the coverage, while maintaining the MPIW to low values. To ensure a fair comparison, we compute the MPIW at a fixed coverage, by increasing the prediction interval until a set coverage is reached before making the calculation. Specifically, we measure MPIW values at prediction intervals that ensure $95\%$ coverage over the test data set.

The evaluation of deep learning uncertainty quantification algorithms for the electron head reconstruction is reported in Table 1. We observe that the ensemble-based approaches lead to a decrease in prediction error over the single deterministic model. This is due to the aggregation over decorrelated individual models and forms the basis of bagging-based meta-models (Zhou, 2021). However, the variance between the predictions of individual models in the ensemble is not a reliable indicator of predictive uncertainty. As an illustration, we find that the Bootstrapped ensembles are leading to a highly overconfident estimate of the predictive uncertainty with a coverage of only $0.22$ for a $95\%$ prediction interval. To this end, we would caution against reliance on uncertainties from Bootstrapped ensembles in use cases similar to this one. The Evidential Deep Learning approach gives well calibrated uncertainties, where the coverage is $0.79$ for a $95\%$ prediction interval. The mean prediction interval width generated by the Evidential Deep Learning models is moderated as well. This indicates that the EDL models provide sharpness of the predicted intervals. Thus, based on this empirical study, we utilize the Evidential Deep Learning approach for uncertainty estimation in this investigation. After selection of the EDL approach, we tuned the regularization coefficient to 0.02 to enforce correct coverage. The model uses the Evidential Deep Learning for Regression framework (Amini et al., 2020). Herein the model architecture is identical as the prior model in Section 3.2. However, the final layer outputs the parameters of the posterior distribution on the position of the head, as opposed to the actual estimate for the position of the head. The final model was trained using the Adam optimizer with a learning rate of $5\cdot 10^{-5}$ over 1500 epochs. It was trained at the SLAC Scientific Data Facility on an NVIDIA Tesla A100 GPU, which took about 90 minutes.

Figure 10 shows several examples of simulated 1,000 keV electron tracks. The amount of charge in each pixel is given by the colored scatter plot of circles, where the size and color of each circle is proportional to the charge signal. The true and reconstructed origin of the electron track are given by the orange and green dots respectively, with an ellipsoid around the reconstructed origin representing the predicted uncertainty. The examples have been chosen to illustrate different possible outcomes of the reconstruction. Panels 10a, 10b, and 10c show three different views of the event with the largest absolute reconstruction error. Here the model appears to have identified the wrong end of the track as the head. However, it has compensated for this mistake somewhat by correctly estimating a large error on the $Y$ dimension. While perhaps not obvious from these examples, it is the case that the uncertainty ellipsoids have no diagonal components. This is because the 3 cardinal directions are treated independently by the model, and so it lacks the ability to include non-diagonal components in the covariance matrix for the uncertainty, leading to ellipsoids with axes always aligned with the cardinal axes. Figure 10d shows the 20 cm drift track with the smallest reconstruction error. Here we clearly see the Bragg peak in the lower center of the image (larger/lighter circles), and the model has correctly identified that the opposite end of the track is the true track head.

To verify the model’s uncertainty calibration, it is instructive to plot the output of the model in a 2D space: the model-predicted squared error on the electron track head position against the true prediction error (i.e. the difference between the true and predicted position of the electron track head). As can be seen in Figure 11, these two parameters have a parabolic relationship, and achieve a small non-zero minimum value for the predicted squared error at a true error of zero. This is expected, since the posterior distribution of the predicted squared error is inverse-gamma, which approaches zero exponentially as its variable approaches zero. This can be seen in the marginal Y-axes of Figure 11. As can also be seen in the marginal X-axes of Figure 11, the 1D distribution of the true error is normally distributed, as expected given the normal-inverse-gamma posterior distribution of the EDL loss function.

Similar to Figure 5, Figure 12 shows the distribution of absolute errors for the trained EDL model at a 5 cm drift depth, broken down by initial electron energy. We show the dependence on drift depth of the RMS error for each initial electron energy in Figure 13. Since the EDL approach provides us with a modeled uncertainty, we can turn that into an estimate of the systematic error on the points in Figure 13. The error bars on the points in Figure 13 are an estimate of the variation in electron track shapes, leading to incorrect location predictions, that is not modeled by the EDL approach. Unmodeled track shape variation can arise not only from imperfect design/training of the model, but also from variation at a scale too small to be distinguishable by the model, once the simulated drift and detector response has been applied. While there does appear to be a clear increasing trend of the RMSE as the drift length increases, the fluctuations about that trend mostly seem to be within the estimated systematic uncertainties. For each point, its error bars are constructed by 1) fitting inverse-gamma distributions to slices of the version of Figure 11 with the appropriate initial electron energy and drift depth, 2) computing the standard deviation of those inverse-gamma distributions, 3) computing the mean standard deviation across the slices, weighted by the number of samples in each slice, 4) using that mean standard deviation to compute the uncertainty on the mean squared error by dividing it by $\sqrt{n}$ where $n$ is the number of samples, and finally 5) propagating the uncertainty on the mean squared error through to the root mean squared error. The statistical uncertainty on the points in Figure 13 is negligible. The RMS error increases as a function of drift depth for all electron energies.

Figure 14 shows the RMS error as function of an applied threshold in the predicted uncertainty, meaning that electron tracks with an uncertainty on the reconstructed track head position greater than the threshold are excluded from the calculation of the RMS error. The top panel shows the effect of excluding these tracks on the RMS error, while the bottom panel shows how many events are retained as a function of the uncertainty threshold. For reference (as seen in Figure 6), in the deterministic model the RMS error for 1,000 keV electron tracks at a 5 cm drift depth is slightly more than 0.6 mm. According to Figure 14, an RMS error similar to that of the deterministic model can be obtained by placing a threshold of approximately 0.6 mm on the predicted uncertainty. However, this requires rejecting approximately 75% of all 1,000 keV electron tracks.

Having verified the model’s calibration, we know that removing events with high predicted uncertainty from the analysis should improve the pointing accuracy of the instrument. This effect can be seen in Figure 15. Figure 15a shows the true error of the reconstructed gamma-ray initial direction against the component of the pointing uncertainty that is attributable to the position uncertainty predicted by our EDL model⁵⁵5This is computed using Equations 6-8 of Boggs, S. E. & Jean, P. (2000), for the simulated gamma source described in Section 3.1 using a parameterization of the EDL model output described in detail in Appendix Section A.2. The solid blue curve in Figure 15b plots the FWHM of a Lorentz distribution fitted to the projection of Figure 15a onto the $Y$-axis, with a cut on the position-related pointing uncertainty given by the $X$-axis value. The dashed green curve gives the fraction of total events remaining after applying such a cut. With no uncertainty cut (i.e. the $Y$-value of the solid blue curve at the right edge of Figure 15b), the FWHM of this distribution is 2.788 degrees, very close to the 2.746 degrees for the deterministic model (Figure 7). We intend to study the extent to which a cut on this parameter can be used in concert with other standard acceptance criteria (requiring sufficiently large separation between Compton scatters, choosing low-angle Compton scatters, etc.) to improve the overall pointing performance of the instrument.