中文

为粒子轨迹重建扩展图神经网络

机器学习 2025-04-08 v1 计算工程、金融与科学 分布式、并行与集群计算

摘要

粒子轨迹重建是高能物理学中的一个重要问题,必要用于研究亚原子粒子的性质。传统的轨迹重建算法随粒子数量而规模化差。Exa.TrkX 项目通过将粒子轨迹重建简化为图上的边分类,并使用图神经网络(GNN)来产生粒子轨迹来缓解这一计算负担。然而,这种基于 GNN 的方法在内存上具有限制性,并且会跳过那些超过 GPU 内存的图。我们对 Exa.TrkX 管道进行改进,以在输入粒子图的样本上进行训练,并展示这些改进可推广到更高的精度和召回率。此外,我们将为 GNN 训练引入的性能优化适应于我们增强的 Exa.TrkX 管道。这些优化相对于我们的基线实现在 PyTorch Geometric 中提供了 2×2\times 的加速。

关键词

引用

@article{arxiv.2504.04670,
  title  = {Scaling Graph Neural Networks for Particle Track Reconstruction},
  author = {Alok Tripathy and Alina Lazar and Xiangyang Ju and Paolo Calafiura and Katherine Yelick and Aydin Buluc},
  journal= {arXiv preprint arXiv:2504.04670},
  year   = {2025}
}