Improved particle-flow event reconstruction with scalable neural networks for current and future particle detectors

TL;DR

The paper tackles full-event particle-flow reconstruction for future high-granularity detectors by framing it as a set-to-set prediction problem and comparing scalable ML architectures. It introduces a locality-sensitive hashing graph neural network (GNN) and a kernel-based transformer, both designed to avoid quadratic memory/computation for events with up to about $|X|,|Y| \simeq 10^4$ elements, using an EDM4HEP dataset. Hyperparameter optimization yields significant improvements, with the GNN achieving up to a $50\%$ improvement in jet $p_{\mathrm{T}}$ resolution over the baseline and outperforming the transformer on key event-level metrics. The work offers open datasets and code, demonstrates cross-hardware portability and scalable inference on GPUs, and outlines future directions including applying to Run-3 data, extending to raw detector hits, and integration into standard reconstruction frameworks.

Abstract

Efficient and accurate algorithms are necessary to reconstruct particles in the highly granular detectors anticipated at the High-Luminosity Large Hadron Collider and the Future Circular Collider. We study scalable machine learning models for event reconstruction in electron-positron collisions based on a full detector simulation. Particle-flow reconstruction can be formulated as a supervised learning task using tracks and calorimeter clusters. We compare a graph neural network and kernel-based transformer and demonstrate that we can avoid quadratic operations while achieving realistic reconstruction. We show that hyperparameter tuning significantly improves the performance of the models. The best graph neural network model shows improvement in the jet transverse momentum resolution by up to 50% compared to the rule-based algorithm. The resulting model is portable across Nvidia, AMD and Habana hardware. Accurate and fast machine-learning based reconstruction can significantly improve future measurements at colliders.

Figures (10)

• Figure 1: A conceptual overview of the machine-learned particle flow approach based on tracks, hits and clusters on one simulated $\mathrm{t\overline{t}}$ event.
• Figure 2: One layer of the locality sensitive hashing based graph neural network.
• Figure 3: The model validation loss and physics performance throughout training for the graph neural network and kernel-based transformer, before and after hypertuning.
• Figure 4: Performance of the particle flow (PF) and machine-learned particle flow (MLPF) algorithms on single particle gun samples.
• Figure 5: The generated (truth) and reconstructed kinematic distributions for baseline particle flow (PF) and the proposed machine-learned particle flow (MLPF) algorithm.