[论文解读] Track Seeding and Labelling with Embedded-space Graph Neural Networks

Nicholas Choma, Daniel Murnane+17 more|arXiv (Cornell University)|Jun 30, 2020

Advanced Graph Neural Networks参考文献 16被引用 11

一句话总结

本文提出了一种基于图神经网络（GNNs）和学习嵌入的端到端粒子轨迹重建框架，用于图结构的构建。通过利用嵌入空间中的非线性度量结构实现高效的邻域查询与聚类，该方法在HL-LHC数据上实现了高种子效率（>93%）和高纯度（>99%），在圆柱区域的TrackML得分为0.932，相较于传统方法在速度和可扩展性方面表现更优。

ABSTRACT

To address the unprecedented scale of HL-LHC data, the Exa.TrkX project is investigating a variety of machine learning approaches to particle track reconstruction. The most promising of these solutions, graph neural networks (GNN), process the event as a graph that connects track measurements (detector hits corresponding to nodes) with candidate line segments between the hits (corresponding to edges). Detector information can be associated with nodes and edges, enabling a GNN to propagate the embedded parameters around the graph and predict node-, edge- and graph-level observables. Previously, message-passing GNNs have shown success in predicting doublet likelihood, and we here report updates on the state-of-the-art architectures for this task. In addition, the Exa.TrkX project has investigated innovations in both graph construction, and embedded representations, in an effort to achieve fully learned end-to-end track finding. Hence, we present a suite of extensions to the original model, with encouraging results for hitgraph classification. In addition, we explore increased performance by constructing graphs from learned representations which contain non-linear metric structure, allowing for efficient clustering and neighborhood queries of data points. We demonstrate how this framework fits in with both traditional clustering pipelines, and GNN approaches. The embedded graphs feed into high-accuracy doublet and triplet classifiers, or can be used as an end-to-end track classifier by clustering in an embedded space. A set of post-processing methods improve performance with knowledge of the detector physics. Finally, we present numerical results on the TrackML particle tracking challenge dataset, where our framework shows favorable results in both seeding and track finding.

研究动机与目标

• 解决由于数据量庞大和高碰撞密度带来的高亮度LHC（HL-LHC）粒子追踪的计算可扩展性挑战。
• 利用机器学习，特别是图神经网络（GNNs），改进轨迹种子生成与标注性能，以替代传统且组合复杂度高的算法。
• 开发一个端到端的学习流程，整合学习到的点迹嵌入、嵌入空间中的图构建以及基于GNN的边分类，实现完整的轨迹重建。
• 通过在嵌入空间中进行高效聚类，并结合探测器物理知识的后处理步骤，进一步提升性能。
• 在TrackML粒子追踪挑战数据集上展示最先进结果，重点关注在具有真实数据约束的圆柱区域。

提出的方法

• 在学习到的嵌入空间中使用k近邻（KNN）构建图，将点迹投影到低维欧几里得空间，以实现高效的邻域查询。
• 使用度量学习训练神经网络，使来自同一粒子的点迹在嵌入空间中距离更近，从而形成非线性度量结构。
• 使用GNN对边进行分类，判断其为真实或虚假（即两个点迹是否属于同一条轨迹），利用来自探测器几何结构和嵌入的节点与边特征。
• 利用GNN的边预测结果生成高纯度轨迹种子，随后在嵌入空间中应用DBSCAN聚类，将点迹聚合成完整轨迹。
• 引入后处理步骤，通过结合探测器物理知识（如限制在圆柱区域|η| < 2.1并过滤碎片）来提升性能。
• 优化流水线以实现高速和并行化处理，实现每事件约5秒的处理速度，适用于HL-LHC的实时处理需求。

实验结果

研究问题

• RQ1 在非线性度量空间中学习到的嵌入是否能提升粒子轨迹重建中图构建的准确性和效率？
• RQ2 基于GNN的边分类器在双点对与三元组种子生成中，相较于传统手工设计的判据，能多大程度上实现性能超越？
• RQ3 在轨迹查找任务中，嵌入空间中的聚类与传统聚类流程相比，在纯度和效率方面表现如何？
• RQ4 限制在圆柱区域并去除端盖碎片对最终TrackML得分和重建性能有何影响？
• RQ5 该完整流水线是否能在极少依赖手工设计启发式规则的前提下实现高性能，从而为HL-LHC提供可扩展的端到端解决方案？

主要发现

• 当限制在圆柱区域（|η| < 2.1）并排除端盖碎片时，该方法在TrackML数据集上实现了0.932的TrackML得分，体现了在可重建性约束下的强劲性能。
• 种子效率超过93%，纯度高于99%，在准确性和计算可扩展性方面显著优于传统方法。
• 双点对GNN的TrackML得分为0.805，三元组GNN达到0.846，证明了GNN在轨迹种子边分类中的有效性。
• 在嵌入空间图上使用DBSCAN聚类表现出稳健性能，得分下降主要源于效率问题而非纯度，表明模型具有良好的泛化能力。
• 流水线每事件处理时间约为5秒，支持快速且可并行处理，适用于HL-LHC的数据率要求。
• 与传统聚类和基于卡尔曼滤波的方法相比，该框架表现出更优的性能，尤其在处理高密度事件和降低组合复杂度方面优势明显。