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