ParticleFlow model checkpoints
MLPF reconstructs particles from detector inputs. This repository provides trained models, configurations, training logs and evaluation results.
Full documentation covers installation, datasets, training, evaluation and the physics studies.
Reusable event views: CLD, CLIC, IDEA, MAIA and ColliderML, with white/transparent PNGs and vector SVGs.
Latest models: v3.2.0
| Detector | Inputs | Selected step | Release | Models |
|---|---|---|---|---|
| CMS | PF-element inputs | 100,000 | cms/v3.2.0 | Attention / HEPTv2 |
| CLIC | Tracks + calorimeter clusters | 50,000 | clic/clusters/v3.2.0 | Attention / HEPTv2 |
| CLIC | Raw tracker + calorimeter hits | 50,000 | clic/hits/v3.2.0 | Attention / HEPTv2 |
| CLD | Tracks + calorimeter clusters | 50,000 | cld/clusters/v3.2.0 | Attention / HEPTv2 |
| CLD | Raw tracker + calorimeter hits | 50,000 | cld/hits/v3.2.0 | Attention / HEPTv2 |
Run directories include checkpoints with optimizer/resume state, configs, model
kwargs, histories, training logs, TensorBoard streams and final plots.
All runs use seed 12345 and training revision
6b80b983.
See training datasets for the TFDS selections.
Training and validation
Training curves use an exponential moving average with a 2,000-step half-life;
validation is unsmoothed. CLIC/CLD use TensorBoard step/loss_Total_unweighted,
with separate panels for track/cluster and hit objectives.
CMS
CMS uses group/physical_nopu/loss_Total_unweighted and
group/physical_pu/loss_Total_unweighted. Training points average each group over
the preceding 1,000 steps before smoothing. The panels separate the differing
mixtures: training is 91% no pileup; validation weights the groups equally.
Each panel has its own loss scale; the bottom row zooms into steps 60kโ100k.
CLIC
CLD
CLIC and CLD physics performance
HEPTv2 improves final particle F1, jet IQR and MET error in both detectors and input representations. Hit HEPTv2 reaches the highest F1; track/cluster HEPTv2 retains better momentum and jet precision.
Particle matches require ฮR < 0.1 and |ฮpT|/pT < 0.5, with assignment costs capped at 2. Matched-pT MAE uses accepted pairs. Jet IQR and response-qualified recall average ttbar, WW fullhad and qq test summaries equally; the recall also requires |pT ratio โ 1| < 0.5. MET error compares magnitudes, and event p3 error measures the difference between summed predicted and target three-momenta.
These single-seed comparisons use different PF/hit target populations. Validation contains 1,530 PF entries or 1,536 padded hit entries; detailed momentum metrics use differently selected 512-entry subsets. The missing CLD hit HEPTv2 30k history appears as a gap. Curve data and provenance provide the plotted values and source files.
CLIC/CLD training resources
LUMI MI250X, global batch 128. Track/cluster runs use one logical device; hit runs use eight, with 16 events per device. HEPTv2 uses activation checkpointing.
| Detector | Model | Training s/step | Logical devices | Loader wait [ms] |
|---|---|---|---|---|
| CLIC | Track/cluster attention | 0.112 | 1 | 2.05 |
| CLIC | Track/cluster HEPTv2 | 0.400 | 1 | 1.99 |
| CLIC | Hits attention | 2.356 | 8 | 17.11 |
| CLIC | Hits HEPTv2 | 2.888 | 8 | 16.89 |
| CLD | Track/cluster attention | 0.112 | 1 | 2.07 |
| CLD | Track/cluster HEPTv2 | 0.412 | 1 | 2.01 |
| CLD | Hits attention | 2.378 | 8 | 16.91 |
| CLD | Hits HEPTv2 | 2.842 | 8 | 16.71 |
Step times are medians from 100-step intervals excluding validation boundaries, warmup and restarts. Loader wait is the median exposed rank-zero batch wait.
| Detector | Model | Tensor peak [GiB/device] | Reserved peak [GiB/device] | VRAM median / max | Mean GPU busy |
|---|---|---|---|---|---|
| CLIC | Track/cluster attention | 2.25 | 2.77 | 5% / 5% | Unavailable |
| CLIC | Track/cluster HEPTv2 | 3.03 | 3.76 | 7% / 7% | Unavailable |
| CLIC | Hits attention | 11.86 | 28.25 | 37% / 47% | 96.17% |
| CLIC | Hits HEPTv2 | 14.50 | 35.66 | 40% / 59% | 96.59% |
| CLD | Track/cluster attention | 2.27 | 2.79 | 5% / 5% | Unavailable |
| CLD | Track/cluster HEPTv2 | 3.03 | 3.81 | 7% / 7% | Unavailable |
| CLD | Hits attention | 11.69 | 28.48 | 37% / 47% | 95.17% |
| CLD | Hits HEPTv2 | 14.57 | 32.77 | 40% / 54% | 96.38% |
Tensor and reserved peaks come from PyTorch's historical allocation counters; reserved memory includes cached blocks. ROCm VRAM occupancy covers all participating devices across 500 snapshots. Its maximum is the largest observed reading.
GPU busy averages 3,500 readings from devices 1โ7, including intermittent zeros. Device 0 is excluded because rank-zero logging biases its readings. PF runs use only device 0, so their utilization is unavailable. These are arithmetic snapshot means. Memory and ROCm data include the raw readings and measurement details.
Previous releases
Browse earlier models under CLIC track/cluster, CLIC hits, CLD track/cluster, CLD hits, and CMS.
Papers
- 2026: Full event interpretation with machine-learning-based particle-flow reconstruction in the CMS detector.
- 2025: Fine-tuning machine-learned particle-flow reconstruction for new detector geometries in future colliders.
- 2024: Improved particle-flow event reconstruction with scalable neural networks for current and future particle detectors.
- 2023 proceedings: Machine Learning for Particle Flow Reconstruction at CMS. Initial study: CMS-DP-2021-030.
- 2022: CMS-DP-2022-061.
- 2021: MLPF: Efficient machine-learned particle-flow reconstruction using graph neural networks.
See the publication guide for study summaries, archived code, datasets and citation guidance.
Datasets
The dataset catalog is the main reference for dataset definitions, schema compatibility and published availability. Follow the download guide for retrieval and verification.
Datasets used for v3.2.0 training
Both backbones use the selections below, as recorded in each model directory's
train-config.yaml. All runs combine configuration partitions 1โ10.
| Detector / inputs | TFDS training datasets | TFDS version |
|---|---|---|
| CMS tracks/clusters | cms_pf_{ttbar,qcd,ztt} and cms_pf_{ttbar,qcd,ztt}_nopu |
2.8.0 |
| CLIC tracks/clusters | clic_edm_{ttbar,ww_fullhad,qq}_pf |
3.3.0 |
| CLIC hits | clic_edm_{ttbar,ww_fullhad,qq}_hits |
3.2.2 |
| CLD tracks/clusters | cld_edm_{ttbar,ww_fullhad,qq}_pf |
3.3.0 |
| CLD hits | cld_edm_{ttbar,ww_fullhad,qq}_hits |
3.2.2 |
TFDS schema versions and model release versions are independent. The current CMS production guide describes schema 3.2.0; these checkpoints were trained on 2.8.0. The catalog's published-availability section tracks downloadable selections, including partition coverage.