[论文解读] Equivariant Graph Neural Networks for Charged Particle Tracking

Daniel Murnane, S. J. Thais+1 more|arXiv (Cornell University)|Apr 11, 2023

Particle Detector Development and Performance被引用 5

一句话总结

本文提出 EuclidNet，一种用于高能物理中带电粒子追踪的新型 SO(2)-等变图神经网络，旨在利用束流轴周围的旋转对称性。通过使用标量化的输入特征和欧几里得等变模块强制实现等变性，EuclidNet 在模型参数量低于 1,000 的情况下实现了接近最先进水平的性能（AUC = 0.9913），优于非等变基线模型，同时保持了计算效率。

ABSTRACT

Graph neural networks (GNNs) have gained traction in high-energy physics (HEP) for their potential to improve accuracy and scalability. However, their resource-intensive nature and complex operations have motivated the development of symmetry-equivariant architectures. In this work, we introduce EuclidNet, a novel symmetry-equivariant GNN for charged particle tracking. EuclidNet leverages the graph representation of collision events and enforces rotational symmetry with respect to the detector's beamline axis, leading to a more efficient model. We benchmark EuclidNet against the state-of-the-art Interaction Network on the TrackML dataset, which simulates high-pileup conditions expected at the High-Luminosity Large Hadron Collider (HL-LHC). Our results show that EuclidNet achieves near-state-of-the-art performance at small model scales (<1000 parameters), outperforming the non-equivariant benchmarks. This study paves the way for future investigations into more resource-efficient GNN models for particle tracking in HEP experiments.

研究动机与目标

• 开发一种资源高效的图神经网络架构，以尊重粒子探测器几何结构中固有的旋转对称性。
• 通过强制实现 SO(2) 旋转等变性，提升高亮度强子对撞机（HL-LHC）典型高堆积极环境下的可扩展性和推理速度。
• 评估对称性等变模型是否能在小模型规模下匹配或超越非等变模型的追踪性能。
• 研究真实世界数据集中近似对称性对严格等变模型性能的影响。

提出的方法

• 模型采用碰撞事件的图表示，节点特征包括三维打点位置和标量化的分量（z 坐标），以实现 SO(2) 等变性。
• 采用欧几里得等变块（EEBs），通过旋转不变操作的消息传递层处理节点和边特征。
• 边和节点消息通过应用于拼接特征的可学习多层感知机（MLP）计算，更新过程在 SO(2) 变换下保持不变。
• 网络架构由一系列 EEB 组成，包含初始嵌入层和最终解码层，确保所有层均保持等变性。
• 通过数值测试验证等变性，结果显示在 SO(2) 旋转下相对偏差小于 10⁻⁸。
• 模型在单张 A100 GPU 上使用二元交叉熵损失和 Adam 优化器进行训练，共 200 个周期，并采用学习率衰减策略。

实验结果

研究问题

• RQ1 与非等变基线相比，SO(2)-等变 GNN 是否能在小模型规模下实现具有竞争力的追踪性能？
• RQ2 当数据集的旋转对称性仅为近似时，等变模型与非等变模型（如 Interaction Network）的性能表现如何比较？
• RQ3 在存在探测器颗粒度和磁场不均匀性等真实世界缺陷的情况下，强制实现严格等变性是否会限制模型容量？
• RQ4 在粒子追踪任务中，模型规模在等变与非等变模型之间的性能差距中扮演何种角色？
• RQ5 即使底层物理过程仅表现出近似对称性，非约束模型是否仍能学习到等变模型无法捕捉的非对称修正项？

主要发现

• EuclidNet 在 TrackML 数据集上模型参数量低于 1,000 时达到 AUC = 0.9913，优于非等变的 Interaction Network（AUC = 0.9849）。
• 模型对输入旋转表现出鲁棒性，在旋转后的测试集上保持一致性能，证实其 SO(2) 等变性达到数值精度。
• 尽管强制实现严格等变性，EuclidNet 仍表现出接近最先进水平的性能，表明对称性归纳偏置可显著提升样本效率。
• 在更大模型规模下，非等变的 Interaction Network 性能优于 EuclidNet，表明数据集中存在的近似对称性可能需要非对称修正。
• 等变性测试表明，EuclidNet 的输出相对于理想旋转的偏差小于 10⁻⁸，验证了其理论设计的正确性。
• 研究表明，当对称性仅为近似时，非约束模型可能更有效地学习对称与非对称特征，凸显了归纳偏置与表达能力之间的潜在权衡。