Progress towards an improved particle flow algorithm at CMS with machine learning

The particle-flow (PF) algorithm is of central importance to event reconstruction in the CMS experiment at the CERN LHC, and has been a focus of developments in light of planned Phase-2 running conditions with an increased pileup and detector granularity. Current rule-based implementations rely on extrapolating tracks to the calorimeters, correlating them with calorimeter clusters, subtracting charged particle energy and inferring neutral particles from significant energy deposits. Such rule-based algorithms can be difficult to maintain and extend, may be computationally inefficient under high detector occupancy, and are challenging to port to heterogeneous computing architectures in full detail. In recent years, end-to-end machine learning approaches for event reconstruction have been proposed, with the possible advantages of directly optimizing for the physical quantities of interest, being highly reconfigurable to new conditions, and being a natural fit for deployment to heterogeneous accelerators. One such approach, the machine-learned particle-flow (MLPF) algorithm, consists of training a graph neural network to infer the full particle content of an event from the tracks and calorimeter clusters. We discuss progress in CMS towards an improved implementation of the MLPF reconstruction, now optimized using generator/simulation-level particle information as the target for the first time. This paves the way to potentially improving the detector response of in terms of physical quantities of interest. We describe the simulation-based training target, progress and studies on event-based loss terms, details on the model hyperparameter tuning, as well as physics validation with respect to the current PF algorithm in terms of high-level physical quantities such as the jet and missing transverse momentum resolutions.