FastGraph: GPU-Accelerated Graph Construction for Particle Reconstruction

Researchers at ETP, CMU and UZH have recently introduced a new GPU-based algorithm package called FastGraph, developed to accelerate one of the most time-consuming steps in modern graph-based reconstruction methods. When particles traverse a detector, they leave behind a large number of individual signals (“hits”). To reconstruct the particle properties, these hits must be combined and interpreted together. A powerful way to do this is to let hits exchange information with one another: if we connect nearby hits, they form a graph on which graph neural networks (GNNs) can operate to rebuild the underlying particle structure.

The challenge is that creating this graph is far from trivial. A commonly used strategy is to connect each hit to its k nearest neighbours in a learned feature space. However, finding nearest neighbours becomes expensive very quickly: in the naïve approach every hit must be compared to every other, which scales with the square of the number of hits. And numbers are large — depending on the detector region or simulation setup, a single event can easily contain hundreds of thousands of hits. Repeating this process many times during neural-network training can dominate the entire runtime and use all available memory.

FastGraph addresses precisely this bottleneck. It performs the neighbour search directly on the GPU and is tailored for the low-dimensional spaces used inside GNNs. Instead of comparing each hit to all others, FastGraph partitions the space into small bins and restricts the search to nearby regions. Together with a carefully optimised GPU implementation, this allows the algorithm to compute exact nearest neighbours extremely efficiently, with full compatibility to modern machine-learning frameworks.
In the range most relevant for detector-related graph reconstruction, FastGraph achieves speed-up factors between 8 and 300 compared to state-of-the-art libraries, while keeping memory overhead essentially negligible. This means that the bottleneck of graph construction — often the slowest part of graph-based reconstruction pipelines — can be reduced by orders of magnitude.

For graph-based particle-reconstruction research, this opens several new possibilities: larger and more detailed graphs, faster training cycles, and the ability to explore more complex network architectures without prohibitive computational cost. FastGraph is available open-source under the MIT license and can be integrated directly into existing PyTorch-based workflows.

Contact: Prof. Dr. Jan Kieseler