Exa.TrkX

GNN-based pipeline for particle tracking

As particle accelerators become more powerful, detectors become more complex with increasingly dense measurements: O(10M) particles/second, O(100M) measurements/second. Traditional algorithms generate all possible trajectories, which result in a combinatorial explosion.

With the DOE Exa.TrkX project, we pioneer the application of Geometric Deep Learning methods, specifically graph neural networks, to capture and regularize relationships between measurements. Unlike traditional algorithms, this pipeline scales linearly with data density which could increase the discovery capacity potential of future high energy physics experiments. The optimized pipeline runs end-to-end on the NVIDIA V100 GPU with a 20X speed-up wrto using a 48-core Xeon 8268s Cascade Lake CPU.

(Left) Reconstruction wall time per event as a function of the average number of interactions per bunch crossing ⟨μ⟩ with the default configurations of ATLAS Run 2 Inner Detector reconstruction. (Middle) Total inference time with Exa.TrkX as a function of number of space-points in each event for CPUs. (Right) Exa.TrkX inference time for GPUs.
Stages of the TrackML track formation inference pipeline. Light red boxes are trainable stages.