Scaling Graph Neural Networks for Particle Track Reconstruction
This work addresses computational bottlenecks in particle track reconstruction for high-energy physics researchers, but it is incremental as it builds on the existing Exa.TrkX pipeline.
The paper tackles the memory and scalability issues of graph neural networks (GNNs) in particle track reconstruction for high-energy physics by introducing improvements to the Exa.TrkX pipeline, including training on graph samples and performance optimizations, resulting in a 2x speedup over the baseline.
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.