Scalable Multi-Task Learning for Particle Collision Event Reconstruction with Heterogeneous Graph Neural Networks

TL;DR

This work tackles the challenge of reconstructing particle collision events in the high-multiplicity, high-background environment of the LHCb experiment. It introduces a heterogeneous graph neural network (HGNN) with integrated graph pruning and a multi-task loss to perform beauty-hadron reconstruction and primary-vertex association within a single framework. Across extensive ablation studies and performance metrics, the HGNN variants outperform the prior DFEI multi-stage GNN, achieving higher reconstruction accuracy and near-deterministic PV association while delivering substantial inference-time speedups at high track multiplicities. The proposed approach offers scalable, end-to-end reconstruction suitable for real-time triggers and offline analysis, with potential applicability to other collider experiments and future upgrades.

Abstract

The growing luminosity frontier at the Large Hadron Collider is challenging the reconstruction and analysis of particle collision events. Increased particle multiplicities are straining latency and storage requirements at the data acquisition stage, while new complications are emerging, including higher background levels and more frequent particle vertex misassociations. This in turn necessitates the development of more holistic and scalable reconstruction methods that take advantage of recent advances in machine learning. We propose a novel Heterogeneous Graph Neural Network (HGNN) architecture featuring unique representations for diverse particle collision relationships and integrated graph pruning layers for scalability. Trained with a multi-task paradigm in an environment mimicking the LHCb experiment, this HGNN significantly improves beauty hadron reconstruction performance. Notably, it concurrently performs particle vertex association and graph pruning within a single framework. We quantify reconstruction and pruning performance, demonstrate enhanced inference time scaling with event complexity, and mitigate potential performance loss using a weighted message passing scheme.

Figures (9)

• Figure 1: (a) Heterogeneous graph representation (above) and HGNN architecture (below) for a simultaneous beauty hadron reconstruction and PV association. In the graph, solid lines (true edges) are shown with their class labels and indicate a physical relationship between nodes: for PV-track edges, the track is associated with the corresponding PV, while for track-track edges, the tracks originate from the same beauty-hadron decay. Dashed lines (false edges) denote examples of connections where no such relationship exists and have a label 0. (b) Heterogeneous modification in blue to the GNN layer updates from Battaglia et al. battaglia2018relationalinductivebiasesdeep in red. Here, the inputs to the HGNN layer are the sets of node features for tracks and PVs, $V_{\rm tr}$ and $V_{\rm pv}$, sets of edge features for track-track and PV-track edges, $E_{\rm tr}$ and $E_{\rm pv-tr}$, and finally, global features, $u$. The sequence of computations corresponding to Equation \ref{['eq:GNeq2']} is illustrated, while the added pruning tasks are indicated in light brown.
• Figure 2: (a) and (b) Receiver Operating Characteristic curves for edge and node scores $\hat{y}^{e_{\rm tr}}$ and $\hat{y}^{v_{\rm tr}}$ at layers 1, 2, 3 and 8. The pruning selections used in the subsequent results are highlighted. (c) Confusion matrix (in percent) between the number of particles originating from beauty hadrons and the number of particles selected with a tight edge and node pruning, $\hat{y}^{e_{\rm tr}}_{8}>0.2$ and $\hat{y}^{v_{\rm tr}}_{8}>0.2$. (d) Track signal efficiency and background retention as functions of track multiplicity for tight and loose ($\hat{y}^{e_{\rm tr}}_{8}>0.01$ and $\hat{y}^{v_{\rm tr}}_{8}>0.01$) pruning selections.
• Figure 3: GNN and HGNN GPU (left) and CPU (right) mean inference time as a function of track multiplicity per event with and without early pruning ($\hat{y}^{e_{\rm tr} / v_{\rm tr} }_{1-3}>10^{-4}$).The error bands indicate $\pm 1$ standard deviation of the inference time distribution for events in each bin.
• Figure 4: Relationship between the average GPU timing (ms) for events with track multiplicity greater than 300 and percentage of each reconstruction category for HGNN (top) and WHGNN (bottom) models. The left panels show the fractions of Perfect and Complete reconstructions, while the right panels show the fractions of Partial and Not isolated reconstructions. The results are shown for configurations without early pruning (corresponding to the triangular points with the highest average GPU timing) and with edge pruning applied in layers 1-3 for selection thresholds ranging from $\hat{y}^{e_{\rm tr}}_{1-3} > 10^{-7}$ to $\hat{y}^{e_{\rm tr}}_{1-3} > 10^{-2}$ in factors of 10.
• Figure 5: Demonstration of the HGNN inference in an event with seven PVs. The HGNN is able to simultaneously associate all tracks to a PV (in blue), isolate the tracks from beauty hadrons via edge and node pruning (in green) and determine the decay hierarchy of both beauty hadrons via its LCAG prediction shown on the right in matrix form.