Transformer-Based Approach to Enhance Positron Tracking Performance in MEG II

Technical Summary: Transformer-Based Approach to Enhance Positron Tracking Performance in MEG II

Problem Statement
The MEG II experiment searches for the charged lepton flavor-violating decay $\mu \to e\gamma$. A critical bottleneck in achieving the experiment's target sensitivity ($6 \times 10^{-14}$) is the performance of the positron spectrometer under high pileup conditions. The spectrometer utilizes a Cylindrical Drift Chamber (CDCH) where the hit occupancy reaches 35–50% at the target muon stopping rate of $5 \times 10^7 \mu/s$. The conventional reconstruction algorithm, which relies on local pattern recognition to connect nearby hits, suffers from inefficiency in the track-finding stage due to contamination from pileup hits. This inefficiency forces the experiment to operate at a reduced muon stopping rate ($4 \times 10^7 \mu/s$) to maintain acceptable tracking efficiency, thereby limiting the total number of measurable muon decays.

Methodology
To address the limitations of local pattern recognition, the authors developed a Transformer-based machine learning model designed as a hit classifier. The core concept leverages global correlations between hits across different turn segments of the positron trajectory, rather than relying solely on local proximity.

• Model Architecture: The model utilizes a Transformer architecture with two distinct attention blocks: one for Pixelated Timing Counter (pTC) hits and one for CDCH hits.
  • Inputs: The model takes pTC hits (associated with a specific positron candidate) and all CDCH hits in an event as input.
  • Feature Engineering: To facilitate effective attention mechanisms, the authors introduced specific geometric transformations. These include conformal mapping of hit coordinates to align trajectories and rotational alignment based on the first pTC hit to exploit detector symmetry. Features include drift time, wire position, and calibrated patterns ($z_{turn;s}$, $\phi_{turn;s}$) derived from Michel positron tracks to distinguish hits belonging to specific turn segments (e.g., $s=1$ for the final half-turn, $s=2$ for the preceding turn).
  • Output: The model classifies each CDCH hit into one of several categories: pileup/noise ($s=\emptyset$) or specific track segments ($s \in \{1, 2_{in}, 2_{out}, 3_{in}, 3_{out}, 4_{in}, 4_{out}\}$).
• Training Strategy: The model was trained on a dataset comprising 90% Monte Carlo (MC) samples and 10% low-rate data samples. To prevent overfitting to MC-specific detector alignments, low-rate data (where conventional reconstruction is highly reliable) was used to annotate CDCH hits. These data-driven annotations were then mixed with high-rate background data to simulate the pileup environment. A two-step training process was employed: first, learning correlations between signal hits at low pileup, and second, learning to discriminate signal from pileup at high rates.
• Integration: The trained model acts as a pre-filter. Before the standard track-finding and fitting algorithms (described in Ref. [5]) are executed, the model selects CDCH hits based on a probability threshold. This reduces the input to the reconstruction algorithm to a "purity-enhanced" set of hits.

Key Contributions

1. Novel Architecture: The application of a Transformer model to drift chamber hit filtering, specifically utilizing cross-attention between pTC and CDCH hits and self-attention among CDCH hits to identify global track patterns.
2. Robustness to Pileup: The method successfully decouples the reconstruction efficiency from the degradation typically caused by high occupancy (35–50%), allowing the algorithm to function effectively where local pattern recognition fails.
3. Operational Optimization: The implementation enables the experiment to operate at the hardware limit of the muon stopping rate ($5 \times 10^7 \mu/s$) without sacrificing tracking performance.

Results
The proposed method was evaluated using both MC samples and real data collected by MEG II:

• Hit Filtering Performance: At a muon stopping rate of $5 \times 10^7 \mu/s$, the model achieved a signal hit efficiency of 98% while discarding 87% of pileup hits (false positive rate of 13%).
• Tracking Efficiency: The use of the Transformer filter resulted in a 15% increase in positron tracking efficiency at $5 \times 10^7 \mu/s$ compared to the conventional method. This gain was observed in both MC and data, with data showing a consistent improvement across the tested rate range.
• Tracking Resolution: The resolution of the positron track parameters improved by approximately 5%. Specifically, the $z$-resolution (position along the wire) and momentum resolution (inferred from the 52.8 MeV endpoint spectrum of Michel positrons) showed significant gains. The momentum resolution improvement was quantified as approximately 10% based on the reduction of tracks exceeding the kinematic limit.
• Computational Cost: Despite the addition of the ML inference step, the total event processing time decreased to 70–80% of the conventional approach. This is because the reduced number of pileup hits significantly accelerated the seed construction and candidate generation stages, which are the most sensitive to occupancy.

Significance
The paper claims that these performance enhancements are directly transformative for the MEG II experiment's physics reach. By enabling operation at the higher muon stopping rate of $5 \times 10^7 \mu/s$ (previously limited to $4 \times 10^7 \mu/s$ due to efficiency losses) and simultaneously improving the tracking resolution, the experiment expects an approximately 10% increase in the sensitivity of the $\mu \to e\gamma$ branching ratio measurement. Consequently, the collaboration has decided to reprocess the dataset collected up to 2022 using this new method and has already increased the muon beam rate for the 2025 data-taking period. The authors conclude that this approach represents a significant scientific gain and a successful integration of deep learning into high-occupancy tracking detector reconstruction.