Vision Transformers and Graph Neural Networks for Charged Particle Tracking in the ATLAS Muon Spectrometer

This paper presents two machine-learning approaches for the ATLAS Muon Spectrometer to address High-Luminosity LHC challenges: a Graph Neural Network that improves background-hit rejection and reconstruction speed by 15%, and a Vision Transformer proof-of-concept achieving 98% tracking efficiency in just 2.3 ms.

The Gist

Imagine the Large Hadron Collider (LHC) as the world's most powerful particle accelerator, smashing protons together at nearly the speed of light. The ATLAS experiment is one of the giant detectors watching these collisions, looking for clues about the universe's deepest secrets.

Right now, the LHC is busy, but in the future (after 2030), it's going to get crazy busy. Think of it like a quiet library suddenly turning into a packed, noisy concert hall. Instead of 60 people bumping into each other at once, there will be 200. This is called "pileup."

The problem? The detector is drowning in data. It's like trying to find a specific friend's face in a crowd of 200 people, all while someone is throwing confetti, flashing strobe lights, and shouting random numbers. The "confetti" and "shouts" are background noise (false signals), and the "friend" is the muon (a particle the scientists care about).

This paper presents two new, super-smart AI tools designed to help the ATLAS detector find its friends in this chaotic crowd faster and more accurately.

1. The "Smart Bouncer" (Graph Neural Networks)

The Problem: Before the computer can even try to trace a particle's path, it has to sort through millions of tiny signals (hits). Most of these are just noise. The current method is like a security guard checking every single person in line, one by one, which takes too long.

The Solution: The team built a Graph Neural Network (GNN).

• The Analogy: Imagine the detector hits as people in a room. Instead of checking everyone individually, the GNN groups them into "buckets" (clusters) based on who is standing next to whom. It then looks at the relationships between these groups.
• How it works: It acts like a super-smart bouncer at a club. It quickly scans the groups and says, "Hey, this group looks like a party of noise; get out!" and "This group looks like a real guest; let them in."
• The Result: By kicking out the fake noise before the main tracking starts, the computer has much less work to do. It sped up the whole process by 15%, saving precious milliseconds. In the world of particle physics, saving time means you can catch more rare events.

2. The "Super-Scanner" (Vision Transformers)

The Problem: Even after the bouncer does their job, there's still a massive puzzle to solve: connecting the remaining dots to draw the path of the particle. Traditional algorithms do this like a detective trying to connect dots one by one, which is slow and gets confused in a crowded room.

The Solution: The team used a Vision Transformer (ViT), a type of AI famous for recognizing images (like identifying a cat in a photo).

• The Analogy: Instead of looking at the dots one by one, this AI looks at the entire picture at once, like a bird's-eye view. It uses a technique called "Flash Attention," which is like having a superpower that lets it instantly focus on the most important parts of the image while ignoring the rest.
• How it works: It treats every single signal hit as a "token" (like a word in a sentence). It reads the whole "sentence" of hits simultaneously to figure out which ones belong together to form a particle track. It's like looking at a messy scribble and instantly seeing the hidden drawing within it.
• The Result: This is a "proof of concept," meaning it's a prototype, but it's incredibly fast. It can reconstruct a muon's path in 2.3 milliseconds on a standard gaming graphics card (the kind you'd buy for a PC). That is roughly 100 times faster than the current standard method.

Why Does This Matter?

Think of the ATLAS detector as a camera taking a photo of a speeding car.

• The Old Way: The camera takes the photo, then a team of humans sits down to manually clean up the dust on the lens and trace the car's path. It takes too long, and by the time they finish, the car is gone.
• The New Way: The camera has a built-in AI. The "Smart Bouncer" instantly wipes the dust off the lens. Then, the "Super-Scanner" instantly draws the car's path in the blink of an eye.

The Bottom Line:
As the LHC gets busier, the data will become too heavy for old computers to handle in real-time. These new AI tools act as a high-speed filter and a rapid pattern-recognition engine. They don't just make the process faster; they make it possible to keep doing high-quality physics when the machine is running at its absolute limit.

The paper shows that by borrowing ideas from computer vision (like recognizing faces) and graph theory (like social networks), physicists can build a "digital nervous system" for their detectors that is fast, efficient, and ready for the future.

Technical Summary

1. Problem Statement

The ATLAS experiment at the Large Hadron Collider (LHC) faces significant challenges in charged particle reconstruction as it transitions to the High-Luminosity LHC (HL-LHC) era (post-2030).

• Increased Pileup: The number of proton-proton collisions per bunch crossing ($\langle\mu\rangle$) will increase from ~60 (Run 3) to up to 200.
• Muon Spectrometer Challenges: Unlike the Inner Tracker, the Muon Spectrometer has a very low signal-to-noise ratio (as low as 0.6% in some regions) and integrates five different sub-detector technologies with varying readout speeds and spatial precisions.
• Latency Constraints: The Event Filter (EF), the second level of the trigger system, must process 1 MHz of input data with strict latency requirements. The current baseline reconstruction chain takes approximately 255 ms per event on a single CPU thread at $\langle\mu\rangle=200$, which is insufficient for future throughput needs.
• Goal: Develop efficient, robust machine learning (ML) strategies to reject background hits and reconstruct muon tracks in real-time without compromising physics performance.

2. Methodology

The paper proposes two distinct ML-based approaches to address these challenges:

Approach A: Graph Neural Networks (GNN) for Background Rejection

• Concept: Integrate GNNs into the existing non-ML baseline reconstruction chain to filter out background hits before the pattern recognition stage.
• Graph Construction:
  • Nodes are constructed from "Muon Buckets" (clusters of hits within 30 cm longitudinally), rather than individual hits, to ensure computational efficiency.
  • Edges connect buckets in identical or adjacent sectors satisfying specific spatial proximity criteria ($|\Delta z| < 15000$ mm, $\sqrt{\Delta x^2 + \Delta y^2} < 6800$ mm).
• Architecture: Based on the EdgeConv architecture. It propagates information via message passing to classify buckets as containing signal segments, multiple segments, or background noise.
• Implementation: Integrated into the C++ ATLAS Athena framework using PyTorch Geometric and ONNX (with custom CUDA kernel implementations for EdgeConv).

Approach B: Vision Transformers (ViT) for End-to-End Tracking

• Concept: A proof-of-concept for a purely ML-based end-to-end tracking system using Mask2Former (a state-of-the-art panoptic segmentation architecture adapted for HEP).
• Architecture Adaptation:
  • Tokenization: Individual detector hits are treated as tokens rather than image patches.
  • Physics-Informed Prior: Hits are sorted by azimuthal angle ($\phi$). A Windowed Flash Attention mechanism is used, assuming hits within a certain $\phi$ proximity belong to the same track. This reduces complexity from $O(N^2)$ to $O(W \times N)$.
  • Two-Stage Pipeline:
    1. Hit Filtering: A global attention stage performs binary classification to remove background noise.
    2. Tracking: A compact Mask2Former decoder uses cross-attention between learnable track queries and filtered hit tokens to assign hits to tracks and regress track parameters ($p_T, \eta, \phi, q$).
• Training: Trained on full ATLAS simulations ($Z \to \mu\mu$, $J/\psi \to \mu\mu$, $t\bar{t}$) at $\langle\mu\rangle=200$. Uses a weighted multi-task loss (Binary Cross-Entropy for assignment, MSE for regression).

3. Key Contributions

1. GNN Integration: Successfully demonstrated the integration of GNNs into the standard ATLAS reconstruction chain (Athena) for real-time background rejection, overcoming ONNX compatibility hurdles.
2. End-to-End ViT Tracking: Pioneered the application of Mask2Former and Windowed Flash Attention to the Muon Spectrometer, treating sparse, heterogeneous detector hits as tokens for global combinatorial solving.
3. Hardware Efficiency: Validated both approaches on consumer-grade and high-end GPUs, demonstrating feasibility for real-time deployment in the Event Filter.

4. Results

GNN Results (Background Rejection)

• Speed: Reduced per-event reconstruction time from 255 ms to 217 ms (a 15% improvement) at $\langle\mu\rangle=200$.
• Performance: Achieved a 97% background bucket rejection rate at $\langle\mu\rangle=60$.
• Physics Impact: Preserved signal reconstruction efficiency and precision across the entire kinematic range with no negative bias.

ViT Results (End-to-End Tracking)

• Hit Filtering:
  • Achieved an AUC of 0.9997.
  • At 99% signal efficiency, background rejection reached 99.7%, reducing average occupancy from ~6,900 hits to 55 hits per event.
  • 99.7% of muon tracks remained reconstructable after filtering.
• Tracking Performance:
  • Signal Detection Efficiency: 98.0% (at a fake rate of 5.1%).
  • Double Matching Efficiency: 94.59% (tracks with >50% hit assignment efficiency and purity).
  • Hit Assignment: 92.9% efficiency and 88.90% purity.
  • Charge Classification: 96.35% accuracy.
• Inference Speed:
  • Consumer GPU (RTX 3090): 2.3 ms per event (batch size 200).
  • High-End GPU (H100): 0.9 ms per event (batch size 200).
  • Note: The tracking stage itself is extremely fast (~0.07 ms), but latency is currently dominated by GPU kernel launch overheads.

5. Significance and Future Outlook

• Scalability: The ViT approach offers a massive reduction in latency (from 255 ms to ~2.3 ms) compared to the CPU baseline, making it highly suitable for the high-throughput demands of the HL-LHC Event Filter.
• Algorithmic Robustness: The attention-based approach is less reliant on specific sub-detector timing information compared to traditional algorithms, showing robustness even when specific detector technologies (like RPCs) are excluded in ablation studies.
• Future Work:
  • Optimizing inference speed via model compilation (TorchCompile) and operator fusion to eliminate kernel launch overheads.
  • Integrating hit-to-track assignment directly into the Event Filter to assist global track fits.
  • Exploring applications for Long-Lived Particle (LLP) decays.
  • Improving track parameter regression precision to match baseline algorithms.

Conclusion: The paper establishes that ML-based methods, specifically GNNs for preprocessing and Vision Transformers for end-to-end tracking, are viable, high-performance solutions for the next generation of particle physics triggers, offering significant speedups while maintaining or improving reconstruction quality.