Exa.TrkX: GNN-Powered HL-LHC Track Reconstruction

• Exa.TrkX is a graph-based machine learning pipeline for charged-particle tracking at HL-LHC, transforming detector spacepoints into candidate trajectories using learned embeddings.
• It employs a multi-stage process including learned graph construction, edge classification with GNNs, and clustering methods like DBSCAN or connected components for efficient track labeling.
• The framework achieves high physics and technical efficiencies via scalable GPU inference and seamless integration with C++ frameworks, addressing challenges of high-dimensional sparse detector data.

Exa.TrkX is an end-to-end machine-learning-based track finding pipeline for charged-particle reconstruction in high-energy physics detectors. Developed for HL-LHC-like conditions, it operates on detector spacepoints and formulates track finding as a graph problem in which detector measurements are grouped into candidate trajectories through learned embedding, edge classification, and graph partitioning or labeling. In the literature, Exa.TrkX is presented as a graph-neural-network-based framework for prompt tracking, displaced large-radius tracking, accelerated GPU inference, C++ integration with ACTS, inference-as-a-service deployment, and scalable GNN training (Ju et al., 2021, Wang et al., 2022, Zhao et al., 9 Jan 2025, Tripathy et al., 7 Apr 2025).

1. Historical lineage and problem setting

Exa.TrkX emerged from the earlier HEP.TrkX program, which was founded to develop machine learning solutions for High Energy Physics tracking. Early HEP.TrkX work examined image-based methods inspired by computer vision, but the program concluded that such approaches were strained by detector data that are high-dimensional, extremely sparse, and irregularly shaped in geometry. This motivated a transition toward models that operate directly on hits or space points, represented as sequences or graphs rather than rasterized images (Farrell et al., 2018).

That shift was driven by the scale of HL-LHC tracking. One early formulation described collision events with about 200 interactions, roughly $O(10k)$ particles, and roughly $O(100k)$ detector space-point hits, while later Exa.TrkX studies emphasized HL-LHC pileup conditions with roughly 5,000 charged particles and 100,000 spacepoints per event (Farrell et al., 2018, Ju et al., 2021). In this regime, the dominant pattern-recognition stages of conventional tracking are difficult to parallelize efficiently, and classical seeding and track building can scale quadratically or worse with occupancy (Farrell et al., 2018).

The project’s relation to conventional Kalman-filter-based reconstruction is explicit throughout the literature. Conventional algorithms are described as excellent for prompt tracks originating near the interaction point, but displaced or large-impact-parameter tracks require dedicated configurations and additional computing time. The large-radius tracking study cites ATLAS as an example in which including large-radius tracking increased total event reconstruction time by about a factor of 2.5 relative to the standard configuration (Wang et al., 2022). Exa.TrkX was therefore developed not only as a physics-reconstruction method but also as an alternative computational formulation of pattern recognition at HL-LHC scale.

2. Core algorithmic structure

The literature does not present a single immutable stage count for Exa.TrkX. In the large-radius study, the pipeline is described in three discrete stages: graph construction, edge classification, and track labeling (Wang et al., 2022). In the accelerated-inference and scaling studies, it is described as a five-step pipeline: data encoding or preprocessing, graph building, edge filtering, GNN inference, and track labeling (Lazar et al., 2022, Tripathy et al., 7 Apr 2025). In the full TrackML performance study, it appears as six stages because preprocessing and post-processing are separated more finely, and track seeding is treated as an additional downstream use of the classified graph (Ju et al., 2021). This variation reflects decomposition choices rather than a change in the basic computational motif.

At its core, Exa.TrkX takes detector spacepoints as graph nodes and candidate relations between pairs of spacepoints as graph edges. One formalization writes the GNN input as a graph

$G = (V, E),$

with node features

$X \in \mathbf{R}^{n\times f}$

and edge features

$Y \in \mathbf{R}^{m\times f}.$

The downstream task is binary edge classification: edges are scored according to whether they connect measurements from the same true charged-particle trajectory (Tripathy et al., 7 Apr 2025).

Graph construction is learned rather than hand-engineered. In the large-radius study, metric learning is implemented as an MLP that maps the 3D spacepoint coordinates to an 8D metric space in which spacepoints from the same true track are close together and spacepoints from different tracks are far apart, without knowing tracking information (Wang et al., 2022). The TrackML performance study likewise uses a six-layer MLP with 512 hidden channels, hyperbolic tangent activations, layer normalization, and output dimension 8, trained with a pairwise hinge loss (Ju et al., 2021). Graph edges are then constructed through nearest-neighbor search in the learned latent space, using either fixed-radius nearest-neighbor search or learned-radius neighborhood queries (Lazar et al., 2022, Ju et al., 2021).

After graph construction, Exa.TrkX introduces learned sparsification before the main GNN. A filter MLP acts as a binary classifier to suppress obvious fake edges, reducing the graph to a size manageable by the more memory-intensive GNN stage (Lazar et al., 2022, Ju et al., 2021). The GNN itself has appeared in several variants: an attention-based Graph Network in the displaced-track study, an Interaction Network GNN in later service and scaling work, and both Interaction Network and Attention GNN variants in the TrackML performance study (Wang et al., 2022, Zhao et al., 9 Jan 2025, Ju et al., 2021). In one accelerated implementation, message passing is repeated eight times, with node hidden states concatenated to nearby edges and edge hidden states aggregated back to nodes (Lazar et al., 2022).

The final stage converts high-score edges into track candidates. In different papers this is done either by DBSCAN clustering on a sparse score graph or by weakly connected components on the thresholded edge graph (Wang et al., 2022, Lazar et al., 2022). The large-radius study used DBSCAN with minPoints = 2 and scanned epsilon from 0.1 to 1.0 in steps of 0.1, finding epsilon = 0.4 to give the highest track efficiency and purity (Wang et al., 2022). The accelerated-inference work replaced CPU DBSCAN with GPU weakly connected components through RAPIDS cuGraph (Lazar et al., 2022).

3. Prompt tracking, seeding, and TrackML performance

Exa.TrkX was first established as a competitive prompt-tracking pipeline on TrackML-like data. The 2021 performance study on the full TrackML dataset reported, averaged over 50 test events, a physics efficiency of $88.7 \pm 0.3\%$ and technical efficiency of $97.6 \pm 0.3\%$ for particles with $p_T > 500$ MeV. Without any fiducial $p_T$ cut, the same study reported physics efficiency $67.2 \pm 0.1\%$, technical efficiency $O(100k)$0, tracking purity $O(100k)$1, and a TrackML challenge score of $O(100k)$2. When trained with a dataset that includes noise hits, the TrackML score reached $O(100k)$3 (Ju et al., 2021).

That study also characterized the pipeline’s dependence on kinematics and detector conditions. For $O(100k)$4 in $O(100k)$5 MeV, physics efficiency was about 43%, while for $O(100k)$6 MeV it was above 88%. Technical efficiency was about 82% for $O(100k)$7 MeV and rose to above 97% for $O(100k)$8 MeV. Efficiency and purity degraded smoothly as occupancy increased, by roughly $O(100k)$9, when the number of spacepoints per event was used as a proxy for occupancy (Ju et al., 2021).

A preceding Exa.TrkX study focused more explicitly on embedded-space graph construction, doublet and triplet classification, and seeding. There, the learned embedding retained 96% of the true edges while keeping only 30.3% of all edges true. The best doublet GNN achieved ROC AUC $G = (V, E),$0; at threshold 0.5 it obtained accuracy $G = (V, E),$1, purity $G = (V, E),$2, and efficiency $G = (V, E),$3. The triplet GNN was stronger still, with accuracy $G = (V, E),$4, purity $G = (V, E),$5, and efficiency $G = (V, E),$6 at threshold 0.5. On 100 test events, the triplet-based seeding algorithm achieved efficiency $G = (V, E),$7 and purity $G = (V, E),$8 (Choma et al., 2020).

The same study also showed how classified edge scores could be converted into full track labels with DBSCAN using a precomputed sparse distance defined by $G = (V, E),$9. Reported TrackML scores were 0.805 for the doublet GNN and 0.815 for the triplet GNN, increasing to 0.876 with $X \in \mathbf{R}^{n\times f}$0 and to 0.932 under further adjacency restriction and reconstructability conditions (Choma et al., 2020). Those results were obtained under explicitly simplified conditions, including barrel-only hits and removal of noise hits, but they established the central Exa.TrkX claim that learned sparse graph construction and message passing can support both seeding and full track labeling.

4. Large-radius and displaced-track reconstruction

A distinguishing feature of Exa.TrkX is that it is designed to be agnostic about global track positions. In the displaced-track study, this property is attributed to the metric-learning stage, which projects 3D spacepoint positions into an 8D learned metric space and groups measurements through relational information rather than explicit assumptions about proximity to the interaction point (Wang et al., 2022). This allowed the same pipeline to reconstruct prompt tracks from the primary vertex and large-radius tracks from displaced vertices.

The displaced-track benchmark used simulated heavy-neutral-lepton events

$X \in \mathbf{R}^{n\times f}$1

with HNL mass 15 GeV and proper decay length $X \in \mathbf{R}^{n\times f}$2. Event generation used Pythia 8, detector response used Fast Track Simulation (FATRAS), and the detector geometry was a generic geometry in ACTS. The setup included no pileup, the average event contained about 140 tracks, all tracks in the event were used for training, and only the prompt and displaced muons from the HNL process were used to evaluate tracking efficiency. The dataset comprised 10,000 events for training and 5,000 events for evaluation, and no improvement was observed with more training events (Wang et al., 2022).

Performance in that study was reported at multiple levels. For graph construction, the reconstructed graph retained 99.9% of true edges, with 5.7% edge purity; the low purity was expected because the initial graph was intentionally over-complete (Wang et al., 2022). For edge classification at a threshold of 0.5, prompt tracks achieved edge efficiency $X \in \mathbf{R}^{n\times f}$3 and edge purity $X \in \mathbf{R}^{n\times f}$4, while displaced tracks achieved edge efficiency $X \in \mathbf{R}^{n\times f}$5 and edge purity $X \in \mathbf{R}^{n\times f}$6 (Wang et al., 2022). These figures support the paper’s statement that the classifier performs comparably well for prompt and displaced tracks.

The paper defined a track as reconstructable if $X \in \mathbf{R}^{n\times f}$7, $X \in \mathbf{R}^{n\times f}$8, and at least five spacepoints are present in the detector. The reconstruction efficiency was defined as

$X \in \mathbf{R}^{n\times f}$9

A reconstructed track was counted as matched if the number of shared spacepoints was greater than 50% of the total spacepoints in both the reconstructed and the true track, simultaneously (Wang et al., 2022). Under these definitions, the reported efficiencies were $Y \in \mathbf{R}^{m\times f}.$0 for all tracks in the event, $Y \in \mathbf{R}^{m\times f}.$1 for prompt HNL tracks, and $Y \in \mathbf{R}^{m\times f}.$2 for displaced HNL tracks (Wang et al., 2022).

The same study reported efficiency as a function of true $Y \in \mathbf{R}^{m\times f}.$3, true $Y \in \mathbf{R}^{m\times f}.$4, true $Y \in \mathbf{R}^{m\times f}.$5, and true production vertex radius, with the qualitative conclusion that the pipeline maintained good efficiency across these variables, including displaced-vertex radii (Wang et al., 2022). The paper’s comparison with conventional tracking remained qualitative rather than a head-to-head benchmark: Exa.TrkX was presented as reconstructing prompt and large-radius tracks simultaneously, with high efficiency and without extra computing cost specifically for large-radius tracking, but the paper also stated that direct computing comparisons against conventional algorithms remained future work (Wang et al., 2022).

5. High-performance implementations and deployment models

A major strand of Exa.TrkX research has concerned the transformation of the original Python prototype into a low-latency GPU pipeline integrated with existing tracking software. The accelerated-inference study described the original implementation as five discrete steps—data encoding, graph building, edge filtering, GNN inference, and track labeling—written in Python and run on both GPUs and CPUs (Lazar et al., 2022). The baseline used PyTorch 1.9.0 for embedding and filtering, TensorFlow 2.5.0 for the GNN, PyTorch Geometric RadiusGraph for graph construction, and scikit-learn DBSCAN for track labeling (Lazar et al., 2022).

The optimized implementation introduced Faiss for graph construction, cuGraph weakly connected components for track labeling, CuPy for GPU-accelerated preprocessing, PyTorch automatic mixed precision for embedding and filtering, and a custom FRNN library for a faster nearest-neighbor search (Lazar et al., 2022). These changes reduced average GPU inference from roughly 15 s/event to about 0.7 s/event on the TrackML benchmark dataset. The graph-building stage was the dominant initial bottleneck: build-edge time fell from about 12 s with PyTorch Geometric RadiusGraph to about 0.54 s with Faiss and about 0.04 s with the custom FRNN. Labeling fell from more than 2 s to about 0.11 s, and to about 0.09 s in the final optimized configuration (Lazar et al., 2022).

The same work reported that automatic mixed precision gave a 1.8–3× speedup for the embedding and filtering MLPs, reduced memory footprint, and incurred no loss in accuracy relative to full precision (Lazar et al., 2022). It also converted trained models to ONNX and executed them through the ONNX Runtime C++ API, ported FRNN to CUDA and libtorch, used the C++ API for cuGraph (libCuGraph), and integrated the complete C++ pipeline into ACTS as a standalone plugin component with a generic tracking interface whose input is a vector of spacepoints and whose output is a list of track candidates (Lazar et al., 2022). A stated limitation was that ONNX support for tensor reduction operators such as scatter_add remained problematic for GNN-specific operations, so the authors suggested PyTorch JIT/libtorch as a temporary alternative for some models (Lazar et al., 2022).

The TrackML performance study had already shown why such engineering mattered. Averaged over 500 events, end-to-end inference took $Y \in \mathbf{R}^{m\times f}.$6 seconds/event on an Nvidia V100 GPU with 16 GB memory, compared with $Y \in \mathbf{R}^{m\times f}.$7 seconds/event on a single CPU core. In synchronous step-by-step timings, graph construction and edge filtering were the principal bottlenecks, and one event typically had to be split into 7 mini-batches of 800k edges for filtering because the filtering input for a single event did not fit in GPU memory. Peak GPU memory usage was about 15.7 GB (Ju et al., 2021). Those studies together established Exa.TrkX as an algorithmic pipeline whose practical viability depends strongly on systems-level implementation choices.

A later development extended this engineering logic to remote service deployment. In the inference-as-a-service study, Exa.TrkX was deployed through NVIDIA Triton Inference Server as a custom backend, with ACTS on the client side sending gRPC requests to a remote Triton server that executed inference on the GPU (Zhao et al., 9 Jan 2025). The backend lifecycle was defined in three phases—initialization, execution, and termination—and the service model emphasized fine-grained control over memory copies and type conversion, support for complex sequential workflows, integration of existing C++ code, and lower latency by avoiding unnecessary framework overhead (Zhao et al., 9 Jan 2025).

The reported throughput on a single NVIDIA A100-SXM4 GPU reached about $Y \in \mathbf{R}^{m\times f}.$8 when there were two or more model instances on the GPU and sufficiently many concurrent requests. Throughput and GPU utilization reached their maxima when there were four or more model instances, and there was no significant throughput difference between A100 40 GB and A100 80 GB versions in that study. With multiple GPUs, throughput increased from about 1.6 events/s on one GPU to 4.6 events/s on four GPUs (Zhao et al., 9 Jan 2025). A direct timing comparison reported Exa.TrkX model inference times of 9.65 s for direct CPU, 2.42 s for direct GPU, and 2.24 s for GPU as a service, while client send/receive time, network transfer time, and server I/O overhead were described as negligible in the tested setup (Zhao et al., 9 Jan 2025).

6. Scalability, caveats, and ongoing directions

The principal scalability issue identified in later Exa.TrkX work concerns the GNN stage itself. Full-graph training is memory-prohibitive because large graphs produce large activation tensors, especially per edge, and the original pipeline skipped graphs that would exceed GPU memory (Tripathy et al., 7 Apr 2025). This was treated as both a systems problem and a learning problem: it reduces the amount of training data and, according to the paper, behaves like large-batch training, which is often inferior to smaller minibatches (Tripathy et al., 7 Apr 2025).

The proposed remedy was minibatch training on sampled subgraphs rather than entire event graphs. Using ShaDow sampling, the pipeline sampled a subgraph for each minibatch by random walks with fanout $Y \in \mathbf{R}^{m\times f}.$9 and depth $88.7 \pm 0.3\%$0, then fed the induced subgraph to the GNN (Tripathy et al., 7 Apr 2025). The experimental setup used Perlmutter GPU nodes with 4 NVIDIA A100 GPUs per node, PyTorch 2.3.1, PyTorch Geometric 2.5.3, NCCL 2.21, batch size 256, hidden dimension 64, 30 epochs, ShaDow depth $88.7 \pm 0.3\%$1, fanout $88.7 \pm 0.3\%$2, and an 8-layer GNN. On the CTD and Ex3 datasets, the optimized implementation provided about $88.7 \pm 0.3\%$3 to $88.7 \pm 0.3\%$4 speedup over PyG’s ShaDow implementation, and the abstract highlighted a $88.7 \pm 0.3\%$5 speedup over the baseline PyTorch Geometric implementation (Tripathy et al., 7 Apr 2025).

The same paper reported that minibatch training with sampled subgraphs generalizes better than full-graph training, converging to higher precision and recall than the original full-graph Exa.TrkX pipeline, while not degrading performance relative to PyG’s minibatch version (Tripathy et al., 7 Apr 2025). It also optimized distributed data parallelism by stacking parameter matrices and performing a single all-reduce call, thereby reducing communication latency relative to separate all-reduces for each matrix (Tripathy et al., 7 Apr 2025). This suggests that Exa.TrkX’s future scalability depends as much on training-graph management and communication patterns as on the expressivity of the GNN architecture.

Across the literature, however, several caveats remain explicit. The displaced-track study included no pileup, used a generic detector geometry in ACTS rather than a full experiment-specific detector reconstruction chain, and evaluated only a specific HNL benchmark with $88.7 \pm 0.3\%$6 and $88.7 \pm 0.3\%$7 (Wang et al., 2022). The 2021 TrackML study described itself as a first step toward validating the approach on ATLAS and CMS data and stated that full validation on real experiment data, dense-environment studies, displaced-track performance, realistic pileup and noise studies, and systematic studies of detector misalignment, magnetic-field uncertainty, and geometry uncertainty were still needed (Ju et al., 2021). Even in that earlier study, the reported misalignment test was only a proxy test—shifting the $88.7 \pm 0.3\%$8-axis of inner detector layers by up to 1 mm—and the effect on efficiency was less than 0.1% (Ju et al., 2021).

Exa.TrkX is therefore best understood not as a single fixed model, but as a family of learned graph-based tracking pipelines with a stable conceptual core and an evolving implementation stack. The stable core is the reduction of track reconstruction to learned sparse graph construction, edge classification, and graph partitioning. The evolving stack includes attention-based and interaction-network GNNs, DBSCAN and weakly connected components labeling, ONNX- and libtorch-based deployment, Triton-backed inference services, and sampled-subgraph training regimes (Wang et al., 2022, Lazar et al., 2022, Zhao et al., 9 Jan 2025, Tripathy et al., 7 Apr 2025). The literature consistently treats these developments as steps toward HL-LHC-scale tracking that is simultaneously physics-capable, accelerator-friendly, and compatible with existing reconstruction frameworks.