Learning to Reconstruct: A Differentiable Approach to Muon Tracking at the LHC

TL;DR

The paper tackles the challenge of high-precision muon tracking at the HL-LHC by proposing an end-to-end differentiable pipeline that embeds physics priors into a tracking model. It combines a Graph Attention Network for hit scoring with differentiable clustering and a differentiable $\chi^2$ fit to jointly recover track parameters and the transverse momentum $p_\mathrm{T}$, trained via a composite loss that back-propagates physics constraints. On a Geant4-based toy detector with 16 layers in a 1 T field, the end-to-end approach outperforms a factorized baseline in hit classification and $p_\mathrm{T}$ resolution, as evidenced by ROC improvements and narrower $q/p_\mathrm{T}$ residuals. This work demonstrates the potential of physics-informed differentiable programming to enhance reconstruction quality and trigger-level data selection in future collider experiments.

Abstract

Reconstructing the trajectories of charged particles in high-energy collisions requires high precision to ensure reliable event reconstruction and accurate downstream physics analyses. In particular, both precise hit selection and transverse momentum estimation are essential to improve the overall resolution of reconstructed physics observables. Enhanced momentum resolution also enables more efficient trigger threshold settings, leading to more effective data selection within the given data acquisition constraints. In this paper, we introduce a novel end-to-end tracking approach that employs the differentiable programming paradigm to incorporate physics priors directly into a machine learning model. This results in an optimized pipeline capable of simultaneously reconstructing tracks and accurately determining their transverse momenta. The model combines a graph attention network with differentiable clustering and fitting routines, and is trained using a composite loss that, due to its differentiable design, allows physical constraints to be back-propagated effectively through both the neural network and the fitting procedures. This proof of concept shows that introducing differentiable connections within the reconstruction process improves overall performance compared to an equivalent factorized and more standard-like approach, highlighting the potential of integrating physics information through differentiable programming.

Figures (5)

• Figure 1: Layout of the detector geometry simulated in the toy model
• Figure 3: Schematic representation of the end-to-end model developed for track reconstruction. Both the clustering and fitting routines depend on the GAT learnable parameters through the predicted node labels $w_{i,\phi}$. The regression of the particle's $p_\text{T}$ is back-propagated to the patter recognition module as visualized in the figure and described in more details within the text.
• Figure 4: ROC distributions for the hit classification task for both the baseline and the end-to-end models for tracking. The end-to-end model directly incorporates physics priors and perform simultaneously both the hit classification and the transverse-momentum regression. The performance clearly shows how the end-to-end approach boosts the accuracy of the hit classification task.
• Figure 5: Normalized $q/p_\text{T}$ residual distributions, with $q$ being the charge of the reconstructed muon. The three distributions correspond to using the truth labels as input weights in the clustering step, using the baseline model predictions and those from the proposed end-to-end model. The end-to-end model demonstrates improved accuracy in $p_\text{T}$ regression, approaching the limit of an ideal detector.
• Figure 6: $p_\text{T}$ residual distributions as a function of the truth $p_\text{T}$ values