arXiv:2504.04670cs.LGcs.CE2025-04被引 1

提升粒子轨迹重建的图神经网络效率,实现更快更准的计算。

Scaling Graph Neural Networks for Particle Track Reconstruction

  • 将轨迹重建转为图边分类,用改进的GNN模型处理复杂数据
  • 在更高精度和召回率下保持性能,支持更大规模输入
  • 引入训练优化技术,速度提升2倍,适合高能物理计算场景

粒子轨迹重建是高能物理中的关键问题,用于研究亚原子粒子特性。传统算法随粒子数量增加而计算效率下降。Exa.TrkX项目将轨迹重建转化为图上的边分类问题,利用图神经网络(GNN)生成粒子轨迹。然而,该方法内存消耗大,会跳过超出GPU内存的图。本文对Exa.TrkX流程进行改进,使其可在输入图样本上训练,并证明其在更高精度和召回率下具有泛化能力。此外,我们适配了针对GNN训练的性能优化技术,使新架构在PyTorch Geometric基线基础上实现2倍速度提升。

原文摘要 · Abstract (English)

Particle track reconstruction is an important problem in high-energy physics (HEP), necessary to study properties of subatomic particles. Traditional track reconstruction algorithms scale poorly with the number of particles within the accelerator. The Exa.TrkX project, to alleviate this computational burden, introduces a pipeline that reduces particle track reconstruction to edge classification on a graph, and uses graph neural networks (GNNs) to produce particle tracks. However, this GNN-based approach is memory-prohibitive and skips graphs that would exceed GPU memory. We introduce improvements to the Exa.TrkX pipeline to train on samples of input particle graphs, and show that these improvements generalize to higher precision and recall. In addition, we adapt performance optimizations, introduced for GNN training, to fit our augmented Exa.TrkX pipeline. These optimizations provide a $2\times$ speedup over our baseline implementation in PyTorch Geometric.

图神经网络轨迹重建高能物理性能优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。