Exploring End-to-end Differentiable Neural Charged Particle Tracking – A Loss Landscape Perspective

The path of protons at energies relevant for pCT is mainly influenced by interactions of the particles with both electrons and atomic nuclei changing the path of the accelerated particle (Gottschalk, 2018; Groom & Klein, 2000), baring additional difficulties in recovering the traversal path of the particle throughout a multi-layer particle detector. In the following, we briefly describe all three relevant effects in increasing order regarding their influence on the particle trajectories:

1. 1.

   Ionizing energy loss: When interacting with negatively charged electrons in the atom’s outer shell, protons lose a fraction of their initial energy, due to attracting forces between the particles (Bloch, 1933). Due to the relatively small mass of the electron compared to the high momentum of the accelerated proton, the proton remains on its original path.

2. 2.

   Elastic nuclear interactions: However, when repeatedly interacting elastically with the heavier atomic nuclei, the proton is deflected from its original path due to the repelling forces between nuclei and particle. Integrated over multiple interactions, this effect, also referred to as Multiple-Coulomb scattering (MCS), follows an approximately Gaussian shape proportional to the particle’s energy as well as the amount and density of material traversed (Gottschalk, 2018).

3. 3.

   Inelastic nuclear interactions: Finally, on rare occasions, particles suffer from inelastic interactions with an atomic nucleus. In this case, the proton undergoes a destructive process, where the primary can knock out one or multiple secondary particles of various types from the target nucleus. Due to the stochastic nature of the event, the outgoing path of particles is highly complex in its nature and involves significantly larger angles compared to MCS (Gottschalk, 2018). Inelastic interactions are therefore extremely difficult to reconstruct while losing any meaningful information for image reconstruction due to its stochastic behavior.

Optimization of neural networks in high-dimensional parametric spaces is notoriously difficult to visualize and understand. While theoretical analysis is complex and usually requires multiple restrictive assumptions, recent work addressed this problem by providing empirical tools for loss landscape analysis providing compressed representations for characterizing and comparing different architectures or optimization schemes. We provide in the following section a brief introduction to a subset of methods, used for the analyses in the later sections. Here, we closely follow the methodology and taxonomy provided by Yang et al. (2021) helping us to characterize and assess differences and qualities of traditional and end-to-end differentiable particle tracking with graph neural networks by characterizing and comparing the general form of the loss landscape as well as the global connectivity of the landscape.

In this work, we parameterize the detector data as an undirected line graph $\mathcal{G}_{L}=(\mathcal{V}_{L},\mathcal{E}_{L})$ generated from an initially generated hit graph $\mathcal{G}_{H}=(\mathcal{V}_{H},\mathcal{E}_{H})$ containing all the detector readouts of a single readout frame. In an initial hit graph structure, we represent all particle hit centroids in a readout frame as a set of graph vertices and describe possible track connections (segments) between two consecutive layers in the detector by edge connections. To capture a richer representation of the scattering behavior, providing a strong inductive bias for the message passing mechanism of graph neural networks, we transform this initial representation into a line graph ($\mathcal{G}_{L}$), where each edge in $\mathcal{G}_{H}$ is transformed into a node in $G_{L}$ (ref. Figure 2). Further, edges are generated between nodes $v_{L,i},v_{L,j}$ that share a common vertex, constructing descriptions of track deflections of candidates over three consecutive layers. This transformed representation allows to efficiently aggregate information of scattering behavior (two nodes connected by an edge in $\mathcal{G}_{H}$) providing important information on possible track segments in a single GNN message, as supposed to the more complex aggregation in $\mathcal{G}_{H}$ over a two-hop neighborhood. To provide machine-readable and differentiable features (w.r.t. simulation output) we parameterize both vertices ($\bm{v}_{i}$) and edges ($\bm{e}_{ij}$) as

where $\Delta E_{H}$ is the deposited energy in the sensitive layers and $\langle x,y,z\rangle_{H}$ is the global position of the hit in the detector²²2We use the continuous energy deposition and position outputs of the Monte-Carlo simulations here to be able to calculate gradients w.r.t. position and energy deposition. The parameters can be replaced with the cluster size and cluster centroid position accordingly. However, this representation requires additional work to provide differentiable solutions for pixel clustering. for the adjacent nodes $i$ and $j$. Further, each edge is parameterized by a positional encoding, similar to Vaswani et al. (2017). However, instead of using discrete token positions, we leverage cosine similarities $s_{\cos}$ between two hit segments as proposed by Kortus et al. (2023). To reduce the complexity of the graph and to minimize the combinatorial explosion of edges in $\mathcal{G}_{L}$, we remove edges with high angles in $\mathcal{G}_{H}$, measured orthogonal to the sensitive layer. We find suitable thresholds as a trade-off between reducing the graph size and minimizing the number of removed true edges (see Appendix A).

Following the basic idea of track reconstruction using an edge classification scheme with edge scores (Heintz et al., 2020; DeZoort et al., 2021), we compose a graph neural network based on the interaction network (IN) architecture (Battaglia et al., 2016; 2018) as proposed in (DeZoort et al., 2021). However, we use a slightly modified architecture to predict node scores on the line graph representation $\mathcal{G}_{L}$, which directly correspond to the edge scores in $\mathcal{G}_{H}$ (see Figure 3). Following is the formulation of edge and node updates in a generic message passing formulation, as well as the final scoring function based on node representation on $\mathcal{G}_{L}$. In a first step, updated edge representations are calculated using a concatenated vector containing the edge attributes, describing the scattering behavior of the triplet-candidate, together with the node attributes $\bm{v}_{i}^{(0)}$ and $\bm{v}_{j}^{(0)}$ of adjacent nodes according to

Here $\phi_{R,1}$ is a multilayer feed-forward network mapping the concatenated input to a fixed size representation with $\phi_{R,1}:\mathbb{R}^{2d_{\text{node}}+d_{\text{edge}}}\rightarrow\mathbb{R}^{d_{\text{hidden}}}$.

In a subsequent step, all node features of $\mathcal{G}_{L}$ are updated in an aggregation step using the current node representation and edge information aggregated from the direct neighborhood of the respective node $v_{i}$:

Given the dynamic number of incoming edges for each graph vertex, either due to the changing numbers of hit readouts in subsequent layers of $\mathcal{G}_{H}$ or the required generalization ability to different particle densities of the proton beam, we use a node-agnostic max-aggregation scheme, providing empirically better generalization abilities compared to non-node-agnostic aggregation schemes (Joshi et al., 2021).

To predict for each edge in $\mathcal{G}_{H}$ an edge score, we perform a final reasoning step, transforming the node embeddings of each node in $\mathcal{G}_{L}$ into a single scalar output using a feedforward network $\phi_{R,2}$:

We generate the final edge scores $s_{ij}\in[0,1]$ using a sigmoid activation function as proposed in Heintz et al. (2020); DeZoort et al. (2021). We augment the multi-layer feed forward networks using batch normalization (Ioffe & Szegedy, 2015), to mitigate the risk of vanishing gradients and smoothen the optimization landscape (Santurkar et al., 2018), providing better training and inference performance.

Generating discrete track candidates from the continuous edge-score predictions, requires the application of a discrete algorithm that matches the tracks with the lowest total cost. Optimally, this would require a multidimensional assignment problem (MAP) to minimize the total cost of all assignments over the entire detector geometry. This exact definition is, however, impractical in real-world usage since solving a MAP is NP-hard (Pierskalla, 1968), requiring an approximate solution to solve. In the related literature, different optimization procedures, such as DBSCAN or union-find or partitioning algorithms, are used frequently (DeZoort et al., 2021; Ju et al., 2021; Lieret & DeZoort, 2023). In this work, however, we rely on a relaxation of the original MAP formulation to a layer-wise linear sum assignment problem (LSAP), finding a solution $y(\bm{\mathcal{C}})$ that minimizes

for every detector layer. Here, $c_{ij}\in\bm{\mathcal{C}}$ is the assignment cost for the edge combination $i,j$, where $i\in\mathcal{V}_{S}$ is a node from the set of source nodes and $j\in\mathcal{V}_{T}$ is a node from the set of target nodes respectively. Further, $\pi_{ij}$ is the respective assignment policy defining whether the edge should be assigned under the reconstruction policy. While the results in Section 5.1 demonstrate good performance for the choice of the LSA assignment, we want to note here that an increasing amount of detector noise and particle scattering can significantly influence the performance due to the required assignment without any notion of noise.

While all previous building blocks are shared between the proposed predict-and-track and predict-then-track framework, the main difference lies in the optimization procedure of the network. We aim to optimize the prediction quality of both networks by reducing a specific loss function minimizing the disagreement between prediction and ground truth, which is defined as a binary label for every edge in $\mathcal{G}_{H}$. We treat splitting particle tracks containing secondary particles as noisy labels due to the relatively low production rate and generally short lifetime of most secondary particles (e.g., electrons). This avoids the need for computationally costly tracing algorithms that follow the particle generation of the Monte-Carlo (MC) simulation for the creation of the training data and avoids generating assignment rules for complex particle interactions. Given the true edge labels, we minimize for the predict-and-optimize scheme the hamming loss of the assignments over a batch of $N$ randomly sampled hit graphs according to

where $\hat{y}_{ij}$ is the assignment of the $k$-th edge in hit graph $\mathcal{G}_{H}$ and $y_{ij}$ is the ground truth assignment of the track. We perform additional weighting of the reconstruction loss for tracking and calorimeter layer to account for the different scattering behavior. Similarly, we use the binary-cross entropy (BCE) loss of the predicted edge scores for the predict-then-optimize version for comparison. To reduce the overall memory footprint of the optimization, we implement the gradient calculation of a single batch using gradient accumulation as N individual predictions over a single graph. We use for all following studies an RMSProp optimizer (Hinton et al., 2012) with a learning rate of $1\mathrm{e}{-3}$, which we found to be significantly more stable than Adam (Kingma & Ba, 2015) and more robust to the selection of the learning rate compared to standard SGD. We train all network variants PAT and PTT for 10,000 training iterations, where in each iteration a random minibatch of N graphs is sampled from the training set. In initial experiments, we observed that the E2E architectures showed the tendency to get unstable after reaching a certain reconstruction performance. We thus perform selective “early stopping”, where we used the training checkpoint (evaluated every 100 iterations) with the best validation performance, determined based on the purity of the reconstructed tracks.

Figure 4 visualizes the intermediate model performances for every 100 training iterations, determined by true-positive rate, false-positive rate, purity and efficiency, both on the training and validation dataset. Additionally, marked with horizontal dotted lines is the iteration of each model run with the highest validation purity. Noticeably, all model instances (PAT and PTT) display similar training performances with a significant improvement of all metrics over the first 1000-2000 training iterations. Continuing from there, the end-to-end trained networks, especially with a lambda of 50 and 75, shows slightly improved performance compared to the model minimizing the prediction loss. However, this performance advantage does not translate to the validation dataset, where once again all models show almost identical results. This lack of generalization ability can be likely traced back to the absence of cost margins, as described in Section 4.3, which we had to remove due to the amount of instability introduced by this mechanism. The replacement of this mechanism with an equivalent, more stable mechanism is thus likely beneficial.

Further, Figure 4 demonstrates that the end-to-end differentiable versions with an interpolation factor $\lambda$ of 50 and 75 can converge faster than the PTT variant. PAT ($\lambda=50$) and ($\lambda=75$) converges on average in $7000\pm 1512$ and $4780\pm 1263$ training steps, respectively, while PTT requires on average $7580\pm 1536$ steps to converge. In contrast, PAT ($\lambda=25$) requires with $9080\pm 1151$ more steps than PPT.

Additional performance results are provided in Table 1. Here, the performance of the trained models is compared for different phantom and particle density configurations. In addition, the performance results of a traditional track follower procedure, which was developed for the particular use-case of particle tracking in the DTC prototype, are provided as a baseline. Similar to the results provided in Figure 4, all model configurations, trained on the task- and the prediction-loss, demonstrate near identical reconstruction purity and efficiency over all tested phantom and particle configurations. Further, the graph neural networks outperform the baseline algorithm in almost all performance metrics over all configurations, with only the configuration of 150 $p^{+}/F$ and 100 mm water phantom as a slight outlier in this tendency. Here, the track follower and RL-based tracker demonstrate a higher reconstruction efficiency, while the reconstruction purity is still higher for the GNN architectures.

To analyze and characterize the training and inference behavior of the end-to-end and two-step tracking approaches and gain an understanding of the effects of end-to-end optimization and its hyperparameters, we visualize (Li et al., 2018) and characterize the loss landscape structure, closely following the methodology and taxonomy by Yang et al. (2021) for characterizing the optimization loss landscapes. We specifically focus on the global connectivity of the loss landscape as it allows us to infer conclusions and comparisons between different model types, given the same parameterization along all configurations. We thus further augment the evaluation by additional analysis of functional similarities (Klabunde et al., 2023) of the optimized models. All used algorithms are detailed in Section 3.2. Additional implementation details are listed in Appendix D.

The similar training and inference performance (Section 5.1) as well as the strong correlation of prediction and task loss surfaces (Section 5.2) indicate strong connections between the E2E and two-step training paradigms. However, previous results lack information about parameter and prediction level similarities and differences, thus providing only limited information of the connectivity in the global loss landscape. We thus continue analyzing the inherent similarities of the networks to get a more profound understanding of existing structural differences in both training paradigms.

Figure 7 visualizes the representational similarities (measured as average CKA similarities between layer representations) and functional similarities (measured as min-max normalized disagreement in predictions). Here, the first row shows all elementwise similarities with an empty diagonal, and the second row shows group wise similarities across random initializations.

Noticeably, PAT-based models reveal consistently moderate to high functional similarities in all network modules. Especially the first two network layers, representing the message passing and aggregation steps of the GNN architecture, show significantly higher similarities compared to their two-step counterpart. This effect is consistent both across random initialization and different combinations of interpolation factors $\lambda$. The representational similarities in the final network module are comparable for both decision-focused and two-step learning, despite the two-step models exhibiting lower similarities in the preceding layers.

Overall, combinations of PTT and PAT demonstrate the lowest representational similarities, indicating that both E2E and two-step training can find similar performing network parameters, that, however, show significantly larger differences in their internal representation compared to models trained using the same optimization paradigm.

While the average reconstruction performance across multiple readout frames is unaltered by changes in the learned network representations (ref. Table 1), subtle prediction instability caused by different learned prediction mechanisms can have a noticeable impact on downstream tasks. Especially for medical applications such as pCT, robustness and stability of the classifier is essential to assure reproducibility and avoid unpredictable changes in the following downstream tasks. To avoid comparing the outcome of downstream tasks, which is computationally expensive and depends on multiple intermediate steps, we analyze functional similarities or prediction instability of the tracking as a helpful proxy, complementing the previous results.

As illustrated in Figure 7, both PTT and PAT-based models show substantial functional disagreement in a similar range, with median disagreement rates around 15%. We want to specifically point out the increase in disagreement for models with higher interpolation values (see Appendix F for the statistical significance analysis), suggesting choosing smaller $\lambda$ values for improved stability of the model. Further, the experiments demonstrate, similar to the representational similarity results, a significant increase in disagreement for combinations containing both PAT and PTT variants. Here, we can observe median disagreement rates of approximately 20%, corresponding to an increase in prediction instability of approximately 33%.

Given the findings of Section 5.3, demonstrating substantial disagreements in model predictions, we now aim to quantify the global shape or connectivity of the loss landscape. Table 2 and Table 3 summarize the mode connectivity, comparing connecting Bézier curves between model parameters as specified in Section 5.2. To get a more profound understanding of the global landscape for prediction and task loss, we generate curves using both the BCE-loss on the predicted edge-scores and the Hamming loss on the assigned edges. Here, we take advantage of our findings in Section 5.2.1, suggesting a strong agreement between Hamming-loss and BCE-loss for low-loss values. In addition, we provide the linear connectivity between two models as a baseline. Further, Neyshabur et al. (2020); Juneja et al. (2022) and Lubana et al. (2023) demonstrate in their work that trained models deficient in linear connectivity, are mechanistic dissimilar⁷⁷7Trained models using different mechanisms or representations for predictions.. This allows us to gain further insides on the stability of the trained models from a loss landscape perspective, supporting the results in Section 5.3.

Supplemented by the results in Table 2, we find that connecting curves with consistently low-loss values and high true positive rates are existent for models with the same training configurations. However, despite the existence of the excellent connectivity with Bézier curves, the linear connectivity is associated with persistently high losses indicating low mechanistic similarity. Analogous to the functional similarity results in Figure 7, we report increasing losses for higher interpolation values, strengthening the findings in Section 5.3 pointing towards the importance of the interpolation factor despite similar reconstruction performances.

Despite the lack of linear connectivity, given the marginal difference in non-linear connecting curves between all PAT-trained models, all parameter configurations seem to be strongly interlinked with globally well-connected minima. This property also holds for the PTT configuration; however, we were only able to find low-loss connecting curves by using the BCE-loss as the optimization criterion. Optimizing the Hamming-loss, in contrast, yielded substantially worse results, which we cannot explain.

We find similar results demonstrating strong nonlinear connectivity between trained models of PAT and PTT mixtures, indicating a good global connectivity of the modes without noticeable loss barriers that could indicate more pronounced differences in the optimization process (see Table 3). However, the high loss values of the linear connectivity ($\text{mean}_{PTT\rightarrow PAT}$: 117.42$\pm$76.26, vs. e.g., $\text{mean}_{PTT}$: 107.24$\pm$42.2) demonstrate a noticeable impact of mixing the training paradigms on the mechanistic similarity. In addition, we find similar to the connecting curves of PTT minimizing the Hamming-loss noticeable higher loss values compared to the BCE-loss.