Graph Neural Network Accelerates Particle Tracking for NICA Collider Experiments

Researchers Yauheni Talochka, Gennady Ososkov, and Nikolay Voytishin have introduced a Graph Neural Network incorporating an attention mechanism and two-stage aggregation to reconstruct particle trajectories within the Time-Projection Chamber of the Multi-Purpose Detector at the Nuclotron-based Ion Collider fAcility. The model achieved 96.2 percent accuracy in edge classification alongside matching 92.6 percent purity and efficiency metrics when evaluated at equal operating thresholds. Beyond accuracy, the architecture completed event evaluations in an average processing time of 50 milliseconds per collision on a high-performance graphics processing unit.

The study directly addresses the severe data processing demands characteristic of modern high-energy physics experiments operating at high-luminosity colliders. Facilities such as the future High-Luminosity Large Hadron Collider and NICA reach megahertz-level collision interaction frequencies, generating vast volumes of detector data and significant spatial overlap between particle trajectories from consecutive collision events. This phenomenon, recognized in the field as the tracking crisis, severely strains classical reconstruction methods such as the Kalman filter, which exhibit steep computational costs when handling dense particle environments despite retaining adequate tracking efficiency.

The primary tracking framework of the MPD experiment relies on TPC detectors arranged in cylindrical layers to capture the dense particle sprays produced by gold-gold ion collisions. To train and evaluate the proposed neural approach, the team simulated a dataset of 1,000 mean-bias Au-Au collision events using the official MPD software package MPDRoot. Within this cylindrical geometry, comprising 12 polar angle sections and 53 radial detector layers, individual collision events registered an average of 338 particles and 1.7 × 10^4 hit points, generating complex point clouds that required robust trajectory reconstruction.

To convert raw detector readings into graph representations, the authors defined each recorded hit as a graph node possessing cylindrical coordinates (r, ϕ, z) that were prenormalized prior to model ingestion. Edges represented potential trajectory segments connecting hits situated across adjacent detector layers. The edge feature vectors incorporated differences across hit coordinates, including pseudorapidity difference Δη, azimuthal angle difference Δϕ, radial-angular distance Δϱ, and longitudinal displacement Δz, providing explicit spatial relationships between neighboring sensor detections.

Before feeding the full graphs into the neural network, the authors applied a geometric filtering routine using non-prenormalized hit coordinates to discard physically implausible segment candidates. Connections were retained only if their parameters fell within bounded intervals: Δϕ between 0 and 0.06 radians, arctan(Δz / Δr) between 0 and 1.6, and edge lengths up to 12. This preliminary selection step established an initial edge purity of 26 percent. Edges belonging to particle tracks with transverse momenta below 150 megaelectronvolts were set to zero to remove spiral tracks from the learning objective, while the remaining true segments received positive binary labels.

The network architecture processes graph components through six core modular stages, starting with multi-layer perceptron encoders having dimensional architectures of [3, 64, 64, 32] for node coordinates and [10, 64, 64, 32] for edge features. An initial edge classifier MLP with dimensions [106, 64, 64, 1] assigns starting probability scores across all filtered connections. A specialized graph attentional convolution layer then aggregates surrounding node and edge representations using the predicted edge scores as attentional weights, running aggregations separately for adjacent layers differing by ±1 and ±2 layers to account for physical track curvature.

Following attentional convolution, the extended representation is processed through a [291, 64, 64, 32] MLP before updating edge embeddings via a [107, 64, 64, 32] MLP and evaluating final edge labels with a [107, 64, 64, 1] MLP. The convolution, edge updater, and classifier modules are chained and repeated across eight consecutive iterations, maintaining input node and edge attributes throughout the sequence to guarantee convergence without losing primary detector data. The complete model was implemented using the LibTorch C++ backend of PyTorch to enable native integration into the MPDRoot software environment.

Model training was conducted on the Govorun supercomputer using an NVIDIA A100-SXM4-80GB GPU, utilizing binary cross-entropy loss and the AdamW optimizer over 100 training epochs with a batch size of one across an 80-20 train-test split. In comparative ablation tests, the researchers evaluated a Simple GNN lacking attention and two-stage aggregation, an Attention-only GNN, and the Full GNN. While the Simple and Attention-only configurations contained 61,442 trainable parameters, the Full GNN scaled to 73,730 parameters and demonstrated the strongest edge classification performance, with each model completing full training within four hours.

Analysis of model score distributions showed pronounced separation between true and false connections, with true particle tracks peaking sharply between scores of 0.9 and 1.0, while false edges concentrated predominantly between 0.0 and 0.1. A secondary score peak around 0.6 emerged as an artifact of the training regime, corresponding to higher-momentum tracks above 150 megaelectronvolts among true edges and cross-track hit connections among false edges. When assessing reconstructed track integrity, the Full GNN maintained track reconstruction efficiency above 90 percent for integrity thresholds below 80 percent, though efficiency decreased to 48 percent at higher integrity requirements.

The Full GNN demonstrated substantial computational speedups during inference, evaluating full collision events in 50 milliseconds on an NVIDIA A100 GPU. By comparison, traditional Kalman filter tracking algorithms under implementation in MPDRoot require an average of 18 seconds to process simulated bismuth-bismuth collision events while delivering 96 percent efficiency at a 50 percent integrity threshold. With source code made accessible on GitHub, the authors demonstrated that attention-driven graph neural networks offer a viable, accelerated tracking path for high-luminosity experiments, while noting that maintaining reconstruction efficiency at strict track integrity thresholds remains a key operational challenge.

In high-energy physics tracking, researchers have actively sought computational alternatives capable of approximating arbitrary nonlinear functions, scaling efficiently, and parallelizing across modern accelerator hardware. Prior work across the field has explored architectures based on Hopfield neural networks, recurrent neural networks, graph neural networks, and Transformer-type models utilizing attention mechanisms. By formulating track reconstruction as an edge-classification task over a graph, the authors build directly on these exploratory efforts, redesigning earlier graph architectures to address high particle multiplicity through attentional convolution and two-stage aggregation.

The spatial geometry of the Time-Projection Chamber requires explicit angular formulations to characterize particle movements across detection layers. When constructing edge features from hit pairs, the pseudorapidity difference incorporates the zenith angle theta through the expression where eta equals the negative natural logarithm of the tangent of half of theta. The angular-pseudorapidity distance Delta varrho combines these values with azimuthal differences via the square root of the sum of squared differences. Furthermore, the geometric bounds applied during coordinate filtration, specifically zero to 0.06 radians for azimuthal differences, zero to 1.6 for the arctangent of longitudinal displacement over radial displacement, and zero to 12 for edge length, were intentionally chosen to encompass the entire parameter domain of true physical particle tracks, preventing genuine trajectory segments from being discarded before neural evaluation begins.

To quantitatively evaluate the classification of detector connections, the study establishes precise mathematical definitions for efficiency, purity, and accuracy. Edge classification efficiency is computed as the number of true edges with scores exceeding a cut threshold divided by the total number of true edges. Purity denotes the ratio of true edges above that threshold to the total count of all edges surpassing the cut. Accuracy encompasses the overall proportion of correctly classified connections, combining true edges above the cut and false edges below it relative to the entire edge sample. Evaluating the Full GNN at the operating threshold where efficiency and purity are exactly balanced demonstrates that incorporating attention and two-stage aggregation provides a marginal enhancement over simpler baseline variants. In addition, the iterative architecture retains the input node and edge features throughout the eight message-passing repetitions, ensuring that original measurement information is preserved rather than washed out during successive graph convolutions.

The distribution of edge probability scores provides critical insight into the classifier’s discernment across different kinematic conditions. Beyond the clean separation where true edges populate high scores near 1.0 and false edges cluster near zero, the analysis examines the origin of the intermediate peak near 0.6. By segmenting true edges across three transverse momentum intervals, specifically below 150 megaelectronvolts, between 150 and 300 megaelectronvolts, and above 300 megaelectronvolts, the researchers demonstrated that all tracks above the 150 megaelectronvolts threshold contribute to this secondary feature as a by-product of the training procedure. Conversely, in the false-edge distribution, this intermediate peak arises from spurious combinatorial segments connecting hits produced by entirely different particles, reflecting the geometric ambiguity that persists in dense hit environments.

Trajectory reconstruction was further analyzed through track integrity metrics and two-dimensional spatial projections. Track integrity measures the percentage of edges assigned by the model to a given track relative to the total number of real physical edges composing that track, with overall efficiency evaluated as the proportion of primary particle tracks reconstructed above a given integrity threshold averaged across events. The visual verification projected onto the xy-plane maps reconstruction outcomes into four discrete categories: green lines representing true positive segments, purple lines denoting false positives, gray lines showing true negatives, and red lines marking false negatives. These projections confirm the model’s ability to disentangle overlapping trajectories, even as efficiency monotonically decreases at higher integrity thresholds. Finally, supporting reproducibility within the high-energy physics community, the underlying dataset is maintained as available from the corresponding author upon reasonable request.

Subject of Research: Particle track reconstruction using graph neural networks in high-energy physics detectors

Article Title: Graph Neural Network with Attention and Two-Stage Aggregation for Particle Track Reconstruction in the TPC MPD of the NICA Accelerator Complex

Article References: Talochka, Y., Ososkov, G., & Voytishin, N. (2025). Graph Neural Network with Attention and Two-Stage Aggregation for Particle Track Reconstruction in the TPC MPD of the NICA Accelerator Complex. Computing and Software for Big Science, 9(1), Article 22. https://doi.org/10.1007/s41781-025-00153-4

Image Credits: AI Generated

DOI: 10.1007/s41781-025-00153-4

Keywords: particle physics, graph neural networks, track reconstruction, NICA collider, time-projection chamber, machine learning, high-energy physics, data processing, Graph, Neural, Network, Attention