A machine learning photon detection algorithm for coherent X-ray ultrafast fluctuation analysis

One key issue in the development of supervised machine learning algorithms is a robust simulator which can adequately describe the data. To describe the simulator here, we denote an input XPFS frame as $x_{i}\in\mathbb{R}^{90x90}$ and the corresponding output photon map as $p_{i}\in\mathbb{R}^{30x30}$. The 3x3 reduction in dimensionality between $x_{i}$ and $p_{i}$ is used to mimic the speckle oversampling factor that is typically used in LCLS experiments. The final calculations are performed on a 30x30 image to allow for the proper photon events to be expressed per speckle.

The detector images and corresponding photon maps are simulated according to the exact parameters described in Burdet et al. (2021a) which were tuned to mimic a previous experiment by matching the overall pixel and droplet histogram. To simulate ground truth photon maps, the following ranges were used: $\bar{k}$ $\in$ [0.025, 2.0] and $C\left(q,t\right)$ $\in$ [0.1, 1.0]. The relevant detector parameters are the probabilities ($w_{i}$) and sizes ($\sigma_{G}$) of the photon charge clouds, the variance of the zero-mean Gaussian background detector noise ($\sigma_{N}$) and the total number of analog-to-digital units (ADUs) per photon. These parameter values are reproduced below in Table 1 and an example of a detector image / photon map pair is shown in Figure 1. For comparison, we used the Gaussian Greedy Guess (GGG) algorithm with relevant parameters which were optimized for these specific simulation parameters described above Burdet et al. (2021a).

In this problem, we seek a supervised machine learning model which learns the functional mapping $f:x_{i}\mapsto p_{i}$, from $N_{f}$ paired simulated data points ($x_{i=1:N_{f}},p_{i=1:N_{f}}$). In this case, the functional mapping was chosen to be a U-Net neural network (Figure 2), a fully convolutional autoencoder architecture, which is characterized by having skip-connections between different layers of resolution and has been shown to perform well on various image segmentation tasksRonneberger, Fischer, and Brox (2015). In the schematic in Figure 2, the architecture is outlined via successive "convolutional blocks". Each such convolutional block consists of two convolutional layers sequentially applied to the input. After each convolution, we utilize batch normalization Ioffe and Szegedy (2015) to ensure robust optimization, followed by a Rectified Linear Unit (ReLU) activation.

To train the model, we use the Frobenius norm between the predicted photon maps ($\hat{P}_{i}$) and the true photon maps ($P_{i}$). This loss function measures the average squared deviation between the predicted photon map and the true photon map, where the average is taken both within a given frame (which contains $N_{p}$ pixels) and between frames ($N_{f}$) in the dataset.

The U-Net model is trained by minimizing $L\left(P,\hat{P}\right)$ with respect to the model parameters. To train the neural network, we use the following hyperparameters: Adaptive Moment Estimation (ADAM) algorithm for optimization Kingma and Ba (2014), batch size = 128, learning rate = 0.001, and batch normalization. We used NVIDIA A100 GPU hardware with the Keras API Chollet et al. (2015).

We performed analysis at low and high count rates ($\bar{k}$) and trained respective models. For the low- $\bar{k}$ data, 100,000 training data points were simulated based on the detector parameters in Table 1 and uniformly selecting $\bar{k}$ in range [0.025, 0.2] and $C\left(q,t\right)$ over the range [0.1, 1.0]. For the high- $\bar{k}$ analysis, 300,000 data points were used for training with an equal proportion of datapoints coming from $\bar{k}$ $\in$ [0.025, 0.2], [1.0, 2.0] and [0.025, 2.0], respectively, with the the $C\left(q,t\right)$ randomly chosen from the range [0.1, 1.0]. To select between competing trained models, the optimal neural network was selected based on maximizing the correlation between the estimated and the true contrast on held-out validation sets of size 5000 for contrast values in the range [0.1, 1.0] with increments of $0.05$. Here, it is worth emphasizing that the metric used to evaluate the photonizing task is important. For example, the overall accuracy is not necessarily a good metric since many photon maps have a small number of photons. Therefore, a model which uniformly predicts $0$ for each pixel will show a uninformatively high accuracy, which is clearly not the desired performance and will lead to poor statistics. Similar issues have been documented in problems with high class imbalances Luque et al. (2019), and correlation based similarity metrics for evaluation are recommended therein Herlocker et al. (2004). Since our final goal is to obtain a good estimate of the contrast, it is useful to use this information directly in the evaluation metric. For the low- $\bar{k}$ analysis, validation datasets were simulated in the range [0.025-0.2]. For the high- $\bar{k}$ analysis, two additional datasets corresponding to $\bar{k}$ in the ranges [1.0, 2.0] and [0.025, 2.0] were used. The evaluation metric for this analysis was the average correlation for the three different $\bar{k}$ ranges. An example of a sample validation plot is shown in Appendix A.

Finally, to obtain an estimate of the contrast from the predicted photon maps we use the maximum likelihood estimation procedure on the negative binomial distribution. In general, the negative binomial distribution is a function of both $C\left(q,t\right)$ and $\bar{k}$. However, we directly use a per-image estimate for $\bar{k}$ and therefore the MLE procedure reduces to a 1D optimization in $C\left(q,t\right)$ Roseker et al. (2018).

As X-ray sources and detectors move towards faster repetition rates nearing 1 MHz, it is important to preserve the possibility of live data analysis. Here, we compare the speed of an optimized GGG droplet algorithm Burdet et al. (2021b) against the trained CNN model (see Table 2). On 1 CPU, the CNN outperforms the GGG algorithm by roughly an order of magnitude. This advantage stretches to two orders of magnitude when comparing GGG parallelized across multiple CPU cores to the CNN running on one NVIDIA A100 GPU.

The observed speedup presented here is consistent with intuition. At inference, the trained neural network, which consists primarily of matrix multiplication operations, is efficiently parallelized over thousands of GPU processes Kirk and Wen-Mei (2016). In contrast, the GGG algorithm requires for-loop operations at the level of each droplet. For this reason, one additional beneficial property of the CNN model is that the prediction rate does not depend on the content of the XPFS frames and is consequently independent of $\bar{k}$. In contrast, for the GGG algorithm, the run time scales linearly with $\bar{k}$ (Figure 3). Here, it is worth mentioning that the GGG algorithm is already orders of magnitudes faster than the Droplet Least Squares algorithm Hruszkewycz et al. (2012), which is exponential in computational complexity.

Finally, since the neural network architecture follows a fully convolutional paradigm, it is possible to make predictions on large input / detector sizes than those used in the training set. This is enabled by the fact that fully-convolutional architectures learn local spatial filters which apply to the full image, making the learning process relatively efficient. For this analysis, the neural network can handle input sizes of ($N_{f}$, 90$a$, 90$b$, 1), where $a$, $b$ are positive integers and $N_{f}$ denotes a variable number of frames. Note, that this allows for any detector size to be used after zero-padding to nearest (90$a$, 90$b$) frame size. We calculate the average time to make predictions on datasets of dimensionality [100, 90, 90, 1], [100, 270, 270, 1] and [100, 900, 900, 1]. We observe rates of 3.4, 3.1 and 0.3 kHz, respectively. As the rate only decreases a factor of 10 between a frame size of (90,90) and (900,900), it appears that we do not observe quadratic scaling that would be observed using droplet-based algorithms. Furthermore, the ability to analyze such data in a single-shot manner is a significant advantage over the sliding approach for droplet analysis which has been developed Abarbanel (2019). A representative example of the CNN prediction using an input resolution of 270x270 pixels is shown in Figure 4.

In this section, we compare the prediction quality of the CNN model against the GGG algorithm on data which simulates LCLS experiments and for which the ground truth contrast is known Burdet et al. (2021a). We begin with an analysis of low count-rate ($\bar{k}$ $\in$ [0.025, 0.2]) data and subsequently present results for higher count-rate data ($\bar{k}$ $\in$ [0.2, 2.0]). To quantify performance, we show parity plots for the predicted and true contrast on datasets with varying contrasts.

At low $\bar{k}$ and high contrasts, the CNN and GGG algorithm give good predictions for the contrast. Although, it is worth pointing out that in this regime the CNN algorithm systematically underpredicts the contrast and has a slightly larger bias than the GGG algorithm. However, at low contrast levels, the GGG algorithm exhibits much greater bias than the CNN (Figure 5a). One possible reason for the overall superior performance of the CNN is that it takes into account variation in photon charge cloud sizes during training. Here, it is worth emphasizing that even after optimizing droplet parameters on simulated data with known detector parameters, the GGG algorithm still exhibits large bias for lower contrast values, indicating that the algorithm may not have the complexity required to fully treat such data.

As $\bar{k}$ increases, the CNN performance decreases on an average photon map accuracy basis (Figure 5c). To further examine the CNN errors, we clipped the output photon map to the range of [0,8] (i.e. no photon map has more than eight photons or less than zero photons) and analyzed the confusion matrix (CF_i,j) of the predictions (Figure 5d). The diagonal of the confusion matrix represents per-class accuracy. For instance, CF_2,2 represents the accuracy of prediction for pixels containing two photons. By examining the diagonal elements, it is clear that the CNN model makes a greater proportion of errors for higher photon counts; note the trend does not hold for the eight photon event due to the clipping operation. The off-diagonals of the confusion matrix indicate how the model makes errors. For example, the CF_3,6 term indicates the probability of the model assigning three photons to a pixel when the true number of photons was actually equal to six. From these elements, we see that the CNN tends to systematically under-predict high photon events. Taken together, these observations suggest that there is a dataset imbalance issue owing to the fact that low photon events are more probable in the training set.

Although at high $\bar{k}$, the CNN is less accurate on a per-photon map basis (relative to its low $\bar{k}$ performance), this does not necessarily imply inferior contrast predictions. In fact, the parity plots are similar for low and high $\bar{k}$ cases (Figure 5b). This observation stems from the trade-off between information content and accuracy at high $\bar{k}$ (Appendix B). For a fixed dataset size, it is harder to estimate photon counts correctly, but the counts have significantly more information about the unknown $C\left(q,t\right)$ parameter.

In Figure 6, we quantify the performance of the CNN model and the GGG algorithm at different $\bar{k}$ ranges using the correlation in the contrast-contrast parity plot as our metric. It is evident that the GGG algorithm is slightly biased across all $\bar{k}$ levels and performs poorly for $\bar{k}>2.0$. This is unsurprising, as droplet-based algorithms were designed to cope with small droplets with relatively few overlapping charge clouds. Furthermore, this implies further development of the CNN algorithm will be capable of handling large $\bar{k}$ data sets.

In this section, we consider a neural network ensemble approach to quantify the uncertainty in the predicted photon maps and contrasts. The motivation for such an analysis is that, while deep learning models have exhibited significant successes in their application to scientific problems, they have a tendency to engender overconfident predictions that may be inexact. As an example, neural networks are unable to recognize Out Of Distribution (OOD) instances and habitually make erroneous predictions for such cases with high confidence Amodei et al. (2016); Nguyen, Yosinski, and Clune (2015); Hendrycks and Gimpel (2016). In reliability-critical tasks, such errors and uncertainties in model predictions have led to undesirable outcomes (2016) (NTSB); (2020) (NTSB); (2017) (NTSB). In this context, quantifying the uncertainties in deep learning model predictions is highly desirable.

There are two sources of predictive uncertainty that need to be considered: Epistemic and Aleatoric. Epistemic uncertainty Smith (2013) (reducible or subjective uncertainty) arises due to lack of knowledge regarding the dynamics of the system under consideration, or an inability to express the underlying dynamics accurately using models. Epistemic uncertainties can lead to biases in the predictions. Aleatoric uncertainty Smith (2013) (irreducible uncertainty or stochastic uncertainty) arises due to noise in the training data, projection of data onto a lower space, absence of important features, etc. Aleatoric sources can lead to variances in the predictions.

For our analysis, we use an ensemble of neural networks to make a point prediction of the contrast $C\left(q,t\right)$ as well as to give an estimate of statistical uncertainty. This is in line with model ensembling based uncertainty quantification (UQ) methods validated in literature Heskes (1996); Efron (1992). Such ensembling accounts for aleatoric uncertainties due to the data and weight uncertainties. In our investigation, the neural network ensemble is formed via sequential sampling, wherein ten partially decorrelated models were sampled during the model training. Contiguous samples were spaced by ten optimization epochs each. The contrast is calculated for each model via a maximum likelihood procedure and the contrast point prediction is taken as the median predicted value. To estimate the model uncertainty, we provide a $95\%$ contrast prediction interval using the standard deviation of the predicted contrasts and making the assuming that the predictive distribution follows a $t$-distribution with nine degrees of freedom (Figure 7).

We see that the error bars are larger at lower contrasts and correctly captures the notion that the prediction task is harder at lower contrasts Burdet et al. (2021b). We also notice a systematic bias in the CNN models at high contrasts. This bias may arise due to the epistemic uncertainties due to the model form (structural uncertainty). Such structural uncertainties in deep learning models cannot be accounted for by any extant approach, including the procedure used in this investigation. However, this indicates a need for more refinement on the present approach (for instance, more fine grained optimization of the model architecture, etc) and will be explored more fully in future work.

Another interesting avenue is to look at the median predicted photon map and the predicted standard deviation map to examine where the neural networks lack consensus. An example of this pair of outputs are shown in Figure 8. Evidently, the ensemble predicts insignificant uncertainty for the majority of the image with the exception of a few pixels with relatively high uncertainty. An interesting future strategy could involve using the CNN model as a fast, initial approach and subsequently run more complex fitting algorithms on regions of the image with high predicted uncertainty.