# Scaling Graph Neural Networks for Particle Track Reconstruction

arXiv: https://arxiv.org/abs/2504.04670 · HTML: https://mlsanity.com/paper/2504.04670/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 11 of its 18 references that are in the map.
- First posted on arXiv 7 Apr 2025.

## Builds on (most cited first, 11 of 11)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [GraphSAINT: Graph Sampling Based Inductive Learning Method](https://mlsanity.com/paper/1907.04931/) (2019)
- [Interaction Networks for Learning about Objects, Relations and Physics](https://mlsanity.com/paper/1612.00222/) (2016)
- [Layer-Dependent Importance Sampling for Training Deep and Large Graph Convolutional Networks](https://mlsanity.com/paper/1911.07323/) (2019)
- [Decoupling the Depth and Scope of Graph Neural Networks](https://mlsanity.com/paper/2201.07858/) (2022)
- [Performance of a Geometric Deep Learning Pipeline for HL-LHC Particle Tracking](https://mlsanity.com/paper/2103.06995/) (2021)
- [Accelerating Graph Sampling for Graph Machine Learning using GPUs](https://mlsanity.com/paper/2009.06693/) (2020)
- [Distributed Matrix-Based Sampling for Graph Neural Network Training](https://mlsanity.com/paper/2311.02909/) (2023)

## Built on by (most cited first, 1 of 1)

- [Locality-Sensitive Hashing-Based Efficient Point Transformer for Charged Particle Reconstruction](https://mlsanity.com/paper/2510.07594/) (2025)

## Abstract (arXiv)

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× speedup over our baseline implementation in PyTorch Geometric.