Scaling Graph Neural Networks for Particle Track Reconstruction

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. Thethis http URLproject, 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 thethis http URLpipeline 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 augmentedthis http URLpipeline. These optimizations provide a $2\times$ speedup over our baseline implementation in PyTorch Geometric.

@article{tripathy2025_2504.04670,
  title={ Scaling Graph Neural Networks for Particle Track Reconstruction },
  author={ Alok Tripathy and Alina Lazar and Xiangyang Ju and Paolo Calafiura and Katherine Yelick and Aydin Buluc },
  journal={arXiv preprint arXiv:2504.04670},
  year={ 2025 }
}