CMS Replaces Hand-Coded Physics Rules with a Transformer That Reads Whole Collisions

Every proton-proton collision at the Large Hadron Collider detonates into a spray of billions of potential signals in a fraction of a nanosecond, and before physicists can say anything about a collision, software must turn raw detector electronics into a list of particles. For more than a decade, the Compact Muon Solenoid (CMS) experiment at CERN has relied on an algorithm called particle-flow, or PF, a meticulously hand-engineered chain of rules that reconciles tracks in the silicon tracker with energy deposits in the calorimeters and hits in the muon system. In a new paper published in The European Physical Journal C, the CMS Collaboration reports the first application of a machine-learning-based reconstruction pipeline to data collected by a hadron collider experiment, and the results mark a turning point in how the largest detectors in physics plan to interpret their data.

The new algorithm, dubbed MLPF, replaces the modular, physics-motivated heuristics of standard PF with a transformer neural network trained end-to-end on simulated collision events. Instead of linking tracks and calorimeter clusters through geometric proximity rules refined over years of tuning, the model learns those correlations directly from data-like simulation, processing the entire event in a single inference step rather than through sequential iterations. The architecture embeds each reconstructed track and calorimeter cluster into a 256-dimensional latent space and passes them through parallel transformer encoders for classification and regression, using scaled dot-product attention to learn pairwise correlations across the full event. With roughly 5.1 million trainable parameters, the model is compact by modern deep-learning standards, comparable in size to the unified particle transformer CMS already uses for jet flavor tagging.

Training such a model required solving a deceptively hard problem: defining what the network should learn. The collaboration constructed a particle-level training target from the Geant4 detector simulation, selecting only particles that genuinely leave detectable signals, excluding those that fall outside the detector acceptance or below energy thresholds, and treating charged hadrons that fail to produce a reconstructed track as neutral hadrons. A per-particle loss function assigns each target particle to exactly one input element, its primary track or cluster, combining a binary classification loss, a particle-identification loss computed with focal loss to handle class imbalance, and a regression loss on the four-momentum expressed in logarithmic ratios relative to the primary element. A transverse-momentum-dependent weighting factor emphasizes high-energy particles in the training tails, and all three loss components converge smoothly with no signs of overfitting.

The training itself ran on eight AMD MI250X GPUs at the LUMI supercomputer in Finland, taking about six hours per epoch for ten epochs using the PyTorch framework and a LAMB optimizer with cosine annealing. Simulated samples were generated under LHC Run-3 2023 conditions at a center-of-mass energy of 14 TeV, with testing events generated under 2024 conditions at 13.6 TeV, deliberately shifted from the training configuration to demonstrate that the model generalizes across detector conditions and collision energies rather than memorizing a single operating point. Events contained between 55 and 75 additional proton-proton interactions per bunch crossing, the pileup that contaminates real LHC collisions and makes reconstruction dramatically harder.

The physics performance numbers are striking. In simulated top quark-antiquark events, MLPF improves jet energy resolution by 10 to 20 percent compared to standard PF for jets with transverse momentum between 30 and 100 GeV, an improvement that holds consistently in both top and quantum chromodynamics multijet samples. Because the model is trained only on individual particles and never explicitly optimized on jets, the better jet performance emerges purely from improved particle-level reconstruction. MLPF also achieves higher efficiency for reconstructing neutral hadrons while maintaining the same misidentification rate as standard PF, a combination that is difficult to achieve with rule-based linking, and the reconstructed missing transverse momentum distributions remain generally consistent with the standard algorithm.

Crucially, the team did not stop at simulation. The model was commissioned on a small sample of real CMS data collected in 2024 with a dijet trigger, corresponding to an integrated luminosity of 375 inverse picobarns. Running both PF and MLPF on the same data events, the collaboration found comparable jet distributions and consistent dijet asymmetry, an observable particularly sensitive to relative jet resolutions. The MLPF-reconstructed missing transverse momentum spectrum was somewhat harder in the tails, matching the trend seen in simulation. This data validation makes MLPF the first machine-learning-based full-event reconstruction pipeline ever applied to data from a hadron collider experiment, a milestone that could reshape reconstruction strategies across the field.

The computational gains are equally significant. Standard PF runs on CPUs and took approximately 110 milliseconds per event in the benchmark, while MLPF performed inference on an Nvidia L4 GPU with a median of 20 milliseconds per event, a roughly fivefold speedup. Because of the iterative nature of the baseline PF block algorithms, no performant GPU implementation exists for standard PF; MLPF unlocks accelerator use for this critical step. The FlashAttention kernel keeps memory and runtime scaling manageable despite the nominally quadratic cost of full-event attention, and a single L4 GPU supported 64 parallel inference streams using less than 10 GB of memory. The model also produces fully reproducible outputs, with repeated evaluations on the same GPU yielding identical results and CPU and GPU inference agreeing within numerical precision.

Integration into the CMS software framework was designed to be minimally invasive. The model is exported as a frozen Open Neural Network Exchange file and loaded through ONNX Runtime, taking reconstructed tracks and calorimeter clusters as inputs and producing particle candidates that downstream modules consume without modification. Upstream tracking and cluster-building algorithms remain untouched, making MLPF fully complementary to ongoing GPU-based developments elsewhere in the reconstruction chain. This separation of concerns means analyses can switch between PF and MLPF with a configuration flag, enabling the side-by-side comparisons that underpinned the validation campaign.

The motivation extends well beyond Run 3. At the high-luminosity upgrade of the LHC, expected to begin operation later this decade, average pileup will reach up to 200 simultaneous interactions per bunch crossing, roughly three times the conditions tested here. Reconstruction complexity grows rapidly with pileup, and hand-engineered algorithms whose combinatorics scale poorly with detector occupancy face diminishing returns. Machine-learning approaches offer favorable scaling, faster retraining when detector geometries change, and even the possibility of baking analysis-specific optimizations directly into reconstruction, as CMS already demonstrated by improving missing transverse momentum resolution for its W boson mass measurement. For proposed future colliders, ML-based particle-flow can accelerate detector design studies through rapid iteration and differentiable programming pipelines.

For now, the collaboration is careful about what has and has not been proven. Electron reconstruction sequences have not yet been retuned for MLPF, the forward detector region needs re-optimization ahead of the high-granularity endcap calorimeter upgrade, and deploying the algorithm in production requires further validation against simulation across a wide range of samples. But the demonstration is complete in its essentials: a transformer trained on simulation runs faster, generalizes across operating conditions, matches the physics performance of decades of hand-tuned software, and has already interpreted real collision data. The era in which the world’s most complex detectors are read out by learned models rather than rulebooks has effectively begun.

Subject of Research: Machine-learning-based particle-flow reconstruction in the CMS detector

Article Title: Full event interpretation with machine-learning-based particle-flow reconstruction in the CMS detector

Article References: Full event interpretation with machine-learning-based particle-flow reconstruction in the CMS detector. (n.d.). https://doi.org/10.1140/epjc/s10052-026-15754-5

Image Credits: AI Generated

DOI: 10.1140/epjc/s10052-026-15754-5

Keywords: CMS, particle-flow reconstruction, machine learning, transformers, CERN, LHC, jet energy resolution, GPU computing, high-energy physics, pileup, HL-LHC, graph attention