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

TL;DR

This work introduces an end-to-end differentiable neural tracking pipeline for charged particle trajectories by embedding a differentiable linear assignment (LSAP) within a graph neural network framework that operates on a line-graph representation of detector hits. By comparing predict-and-track (PAT) with the traditional predict-then-track (PTT) paradigm and analyzing loss landscapes, the authors show that E2E training achieves comparable reconstruction quality while providing gradient information through the discrete assignment step, enabling robust optimization and potential downstream-task integration. They demonstrate strong global connectivity of local minima and substantial prediction stability concerns that depend on the interpolation parameter lambda, highlighting the need for careful hyperparameter tuning in safety-critical applications. The study also reveals that end-to-end differentiability can guide robust tracking within component-based pipelines, potentially mitigating error propagation to downstream tasks such as image reconstruction in proton computed tomography (pCT). Overall, these findings lay groundwork for integrating tracking with downstream objectives and for broader adoption of E2E differentiable reconstruction pipelines in high-energy physics and related fields.

Abstract

Measurement and analysis of high energetic particles for scientific, medical or industrial applications is a complex procedure, requiring the design of sophisticated detector and data processing systems. The development of adaptive and differentiable software pipelines using a combination of conventional and machine learning algorithms is therefore getting ever more important to optimize and operate the system efficiently while maintaining end-to-end (E2E) differentiability. We propose for the application of charged particle tracking an E2E differentiable decision-focused learning scheme using graph neural networks with combinatorial components solving a linear assignment problem for each detector layer. We demonstrate empirically that including differentiable variations of discrete assignment operations allows for efficient network optimization, working better or on par with approaches that lack E2E differentiability. In additional studies, we dive deeper into the optimization process and provide further insights from a loss landscape perspective. We demonstrate that while both methods converge into similar performing, globally well-connected regions, they suffer under substantial predictive instability across initialization and optimization methods, which can have unpredictable consequences on the performance of downstream tasks such as image reconstruction. We also point out a dependency between the interpolation factor of the gradient estimator and the prediction stability of the model, suggesting the choice of sufficiently small values. Given the strong global connectivity of learned solutions and the excellent training performance, we argue that E2E differentiability provides, besides the general availability of gradient information, an important tool for robust particle tracking to mitigate prediction instabilities by favoring solutions that perform well on downstream tasks.

Figures (13)

• Figure 1: Schematic representation of the combinatorial complexity of reconstructing particle tracks over multiple subsequent layers. Marked in bold lines is the correct particle track. Marked in dashed lines are only possible next track segments originating from the previous correct segment. During reconstruction, all possible combinations of segments over all hits are considered.
• Figure 2: Schematic representation of directed hit graph ($\mathcal{G}_H$) and undirected line graph ($\mathcal{G}_L:= L(\mathcal{G}_H$)) using detector readout data simulated for the Bergen pCT prototype detector over multiple detector layers.
• Figure 3: Combination of interaction network style architecture and layerwise combinatorial assignment for particle tracking, providing gradient information using linear interpolations of the optimization mapping.
• Figure 4: True positive and false positive rate of assignments together with purity and efficiency of track reconstruction, evaluated for predict-and-track and predict-then-track as a function of training steps.
• Figure 5: Two-dimensional loss surfaces of PAT framework in logarithmic scale with annotated contour lines. Marked with $\bigstar$ are the trained network parameters.