We study scalable machine learning models for full event reconstruction in high-energy electron-positron collisions based on a highly granular detector simulation. Particle-flow (PF) reconstruction can be formulated as a supervised learning task using tracks and calorimeter clusters or hits. We compare a graph neural network and kernel-based transformer and demonstrate that both avoid quadratic memory allocation and computational cost while achieving realistic PF reconstruction. We show that hyperparameter tuning on a supercomputer significantly improves the physics performance of the models. We also demonstrate that the resulting model is highly portable across hardware processors, supporting Nvidia, AMD, and Intel Habana cards. Finally, we demonstrate that the model can be trained on highly granular inputs consisting of tracks and calorimeter hits, resulting in a competitive physics performance with the baseline. Datasets and software to reproduce the studies are published following the findable, accessible, interoperable, and reusable (FAIR) principles.
翻译:我们研究了基于高粒度探测器模拟的高能正负电子对撞中全事件重建的可扩展机器学习模型。粒子流(PF)重建可被表述为利用径迹和量能器簇团或击中点的监督学习任务。我们对比了图神经网络与基于核的Transformer模型,并证明两者在实现实际PF重建的同时,均避免了二次方的内存分配与计算开销。研究表明,在超级计算机上进行超参数调优可显著提升模型的物理性能。我们还证明所得模型在高性能计算硬件处理器间具有高度可移植性,支持Nvidia、AMD及Intel Habana系列加速卡。最后,我们验证了模型能够基于由径迹和量能器击中点构成的高粒度输入进行训练,并获得与基准方案相当的竞争性物理性能。用于复现研究的公开数据集与软件均遵循可发现、可访问、可互操作与可复用(FAIR)原则发布。