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

技术摘要：基于 Transformer 方法增强 MEG II 中正电子追踪性能的研究

问题陈述
MEG II 实验旨在寻找带电轻子数破坏衰变 $\mu \to e\gamma$。实现实验目标灵敏度（$6 \times 10^{-14}$）的一个关键瓶颈是正电子谱仪在高堆积（pileup）条件下的性能。该谱仪使用圆柱形漂移室（CDCH），在目标缪子停止速率为 $5 \times 10^7 \mu/s$ 时，其打点占用率（hit occupancy）达到 35–50%。传统的重建算法依赖于局部模式识别来连接相邻的打点，由于受到堆积打点的干扰，其在轨迹查找阶段存在效率问题。这种低效迫使实验必须在降低的缪子停止速率（$4 \times 10^7 \mu/s$）下运行，以维持可接受的追踪效率，从而限制了可测量的总缪子衰变数。

方法论
为了解决局部模式识别的局限性，作者开发了一种基于 Transformer 的机器学习模型，将其设计为打点分类器。其核心概念是利用不同正电子轨迹转动段之间打点的全局相关性，而非仅仅依赖局部邻近性。

• 模型架构： 该模型采用具有两个不同注意力模块的 Transformer 架构：一个用于像素化计时计数器（pTC）打点，另一个用于 CDCH 打点。
  • 输入： 模型接收与特定正电子候选相关的 pTC 打点以及事件中的所有 CDCH 打点作为输入。
  • 特征工程： 为了实现有效的注意力机制，作者引入了特定的几何变换。这些变换包括通过共形映射（conformal mapping）对打点坐标进行对齐以匹配轨迹，以及基于第一个 pTC 打点的旋转对齐，以利用探测器的对称性。特征包括漂移时间、导线位置以及从 Michel 正电子轨迹导出的校准模式（$z_{turn;s}, \phi_{turn;s}$），用以区分属于特定转动段（例如，$s=1$ 为最后一个半转，$s=2$ 为前一个转动段）的打点。
  • 输出： 模型将每个 CDCH 打点分类为以下几类之一：堆积/噪声（$s=\emptyset$）或特定的轨迹段（$s \in \{1, 2_{in}, 2_{out}, 3_{in}, 3_{out}, 4_{in}, 4_{out}\}$）。
• 训练策略： 模型使用包含 90% 蒙特卡洛（MC）样本和 10% 低速率数据样本的数据集进行训练。为了防止对 MC 特有探测器对齐方式的过拟合，研究者使用了低速率数据（在此条件下传统重建高度可靠）来标注 CDCH 打点。这些数据驱动的标注随后与高速率背景数据混合，以模拟堆积环境。研究采用了两步训练过程：首先，学习低堆积水平下信号打点之间的相关性；其次，学习在高速率下区分信号与堆积。
• 集成： 训练好的模型充当预滤波器。在执行标准追踪查找和拟合算法（如文献 [5] 所述）之前，模型根据概率阈值选择 CDCH 打点。这使得输入到重建算法的打点集变为“纯度增强型”集合。

主要贡献

1. 新颖架构： 将 Transformer 模型应用于漂移室打点过滤，特别利用了 pTC 与 CDCH 打点之间的交叉注意力（cross-attention）以及 CDCH 打点之间的自注意力（self-attention），以识别全局轨迹模式。
2. 对堆积的鲁棒性： 该方法成功地将重建效率与由高占用率（35–50%）引起的性能退化解耦，使得算法在局部模式识别失效的情况下仍能有效运行。
3. 运行优化： 该实现使得实验能够在硬件极限的缪子停止速率（$5 \times 10^7 \mu/s$）下运行，而无需牺牲追踪性能。

结果
所提方法通过使用 MC 样本和 MEG II 收集的真实数据进行了评估：

• 打点过滤性能： 在缪子停止速率为 $5 \times 10^7 \mu/s$ 时，该模型实现了 98% 的信号打点效率，同时丢弃了 87% 的堆积打点（假阳性率为 13%）。
• 追踪效率： 与传统方法相比，使用 Transformer 滤波器使 $5 \times 10^7 \mu/s$ 下的正电子追踪效率提升了 15%。这一增益在 MC 和真实数据中均有观察到，且真实数据在测试速率范围内表现出一致的改进。
• 追踪分辨率： 正电子轨迹参数的分辨率提高了约 5%。具体而言，$z$ 分辨率（沿导线的空间位置）和动量分辨率（通过 52.8 MeV Michel 正电子端点谱推导）均有显著提升。动量分辨率的改进通过减少超过运动学极限的轨迹量化为约 10%。
• 计算成本： 尽管增加了机器学习推理步骤，总事件处理时间仍降至传统方法的 70–80%。这是因为减少的堆积打点显著加速了对占用率最敏感的种子构建（seed construction）和候选生成阶段。

意义
论文声称，这些性能提升直接改变了 MEG II 实验的物理研究潜力。通过使实验能够在更高的缪子停止速率（此前由于效率损失被限制在 $4 \times 10^7 \mu/s$）下运行，并同时提高追踪分辨率，实验预计 $\mu \to e\gamma$ 分支比测量的灵敏度将提高约 10%。因此，合作组已决定使用这种新方法重新处理截至 2022 年收集的数据集，并已为 2025 年的数据采集期增加了缪子束流速率。作者总结道，这种方法代表了显著的科学收益，也是将深度学习成功集成到高占用率追踪探测器重建中的范例。