Scaling Graph Neural Networks for Particle Track Reconstruction

2025-04-07

AI Summary

Particle track reconstruction in high-energy physics faces significant computational overhead, with traditional algorithms and existing graph neural network (GNN) approaches—such as Exa.TrkX—suffering from GPU memory overflow and poor scalability on large-scale collision events. To address this, we propose a subgraph-sampling-based scalable training paradigm, enabling, for the first time, end-to-end generalizable training of Exa.TrkX on full-scale graphs while circumventing GPU memory constraints. Our method integrates GNN-based modeling, dynamic subgraph sampling, PyTorch Geometric framework optimizations, and computational graph acceleration techniques. Experimental results demonstrate a 2× speedup in training time over baseline methods, alongside substantial improvements in both precision and recall. The approach scales effectively to larger detector geometries and higher-resolution data, establishing a new scalable, high-fidelity paradigm for real-time track reconstruction in high-energy physics.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Data Mining & Knowledge Management: Graph Mining, Social Network Analysis & Community
• Computer Vision: Motion & Tracking

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

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

Problem

• Improving memory efficiency for GNN-based particle track reconstruction
• Enhancing precision and recall in particle track reconstruction
• Optimizing GNN training speed for high-energy physics applications

Innovation

• Uses graph neural networks for edge classification
• Trains on samples of input particle graphs
• Adapts optimizations for 2x speedup