arXiv · 2504.04670
Scaling Graph Neural Networks for Particle Track Reconstruction
Abstract
Particle track reconstruction is an important problem in high-energy physics (HEP), necessary to study properties of subatomic particles. Traditional track reconstruction algorithms scale poorly with the number of particles within the accelerator. The Exa.TrkX project, to alleviate this computational burden, introduces a pipeline that reduces particle track reconstruction to edge classification on a graph, and uses graph neural networks (GNNs) to produce particle tracks. However, this GNN-based approach is memory-prohibitive and skips graphs that would exceed GPU memory. We introduce improvements to the Exa.TrkX pipeline to train on samples of input particle graphs, and show that these improvements generalize to higher precision and recall. In addition, we adapt performance optimizations, introduced for GNN training, to fit our augmented Exa.TrkX pipeline. These optimizations provide a $2\times$ speedup over our baseline implementation in PyTorch Geometric.
Explore related subjects
Keep this discovery
Alok Tripathy, Alina Lazar, Xiangyang Ju, Paolo Calafiura, Katherine Yelick, Aydin Buluc. 2025-04-07. Scaling Graph Neural Networks for Particle Track Reconstruction. https://arxiv.org/abs/2504.04670
Cite the original work for its findings. Save a collection to share your selection of sources.