PaperlayerView exact revision on arXiv

Source-linked AI summary

Prompting Particle Physics: Tokenized Multi-modal Foundation Models for Combinatorially Many Tasks

Nilotpal Kakati, Daniel Murnane, Baran Hashemi, Samuel Klein, Jeffrey Krupa, Eilam Gross, Lukas Heinrich, Michael Kagan

arXiv:2609.31862v1hep-phcs.LGhep-exphysics.data-an

TL;DR

Reconstruction and simulation at colliders commonly rely on long chains of specialized algorithms, and it remains open whether one model can span several stages while retaining inspectable intermediate objects. The paper tokenizes jet objects into a shared vocabulary and trains autoregressive and parallel decoders across modality-to-modality tasks. The models show broad task transfer, and autoregressive nanoHEP outperforms the state of the art on particle flow reconstruction.

  • Problem

    Collider reconstruction and simulation rely on long chains of specialized stages, and whether one model can replace several links while preserving inspectable intermediate objects remains open.

    Source: §1 · p. 2

  • Method

    The models tokenize jet objects into a shared vocabulary and learn mappings between modality subsets, using autoregressive or parallel decoding.

    Source: §1 · p. 2

  • Results

    nanoHEP outperforms the state of the art on particle flow reconstruction, while multi-task nanoHEP transfers to rarely seen directions and generalizes to an unseen four-modality combination.

    Source: §1 · p. 3, §5 · p. 21

  • Takeaways & Limitations

    The approach retains intermediate objects such as tracks and clusters for inspection and calibration rather than collapsing reconstruction into one opaque step.

    Source: §1 · p. 2

    Generated calorimeter cells from truth particles are fully separable from tokenized truth (AUC 0.998), and tokenization round-tripped cells are separable from raw cells (AUC 0.96).

    Source: §4.5 · p. 20

Abstract

from arXiv · show

Reconstruction and simulation at a collider experiment are long chains of specialised algorithms, each tuned to a single step. We explore how one model can serve many of those steps at once, while still producing the intermediate objects (tracks, calorimeter cells, clusters, particles and jets) that make the chain interpretable. To do so, we represent every object in a jet in a single shared token vocabulary and train one model to map any subset of these modalities to any other. A task is then only a choice of which modalities to provide and which to request: particle flow, detector simulation and charged energy subtraction are all directions through the same set of weights. We train over all modality combinations, with a decoder emitting tokens either autoregressively or in parallel. With tokenisation, both architectures train stably with little tuning. In evaluations, both models produce realistic reconstruction and simulation objects, with the autoregressive model particularly faithful to output from Geant4. The autoregressive model also outperforms a state-of-the-art particle flow algorithm on many typical jet reconstruction metrics.

1. Introduction

The paper asks whether one model can handle several reconstruction and simulation stages while retaining inspectable intermediate objects. It tokenizes jet modalities into a shared vocabulary and evaluates two decoding architectures across a broad set of tasks.

1. Introduction

  • Existing LHC workflows build event representations through long chains of tuned stages, while prior machine-learning models have generally addressed individual links.The paper seeks to combine several links without removing intermediate objects such as tracks, clusters and particles, which can still be inspected and calibrated.Source: §1 · pp. 1–2
  • A single model maps any subset of seven jet modalities to any other, with tasks specified by choosing input and output modalities.The modalities span truth-level particles, jets and charged-energy fractions, as well as detector-level cells, tracks, topoclusters and reconstructed particles.Source: §1 · p. 2, Figure 1 · p. 3
  • Tokenisation gives modalities a common input and output format, allowing one model to cover reconstruction and simulation tasks with stable classification-based training.The model maps continuous object features to discrete codebook entries and uses cross-entropy loss across modality combinations.Source: §1 · p. 2
  • Two models use different decoding strategies: nanoHEP emits output objects autoregressively, while HEP4M predicts them in one pass with parallel token sampling.Both are trained to reconstruct from and generate sensor-level detector readouts.Source: §1 · p. 2
  • Autoregressive nanoHEP outperforms the state of the art on particle flow reconstruction, and its simulated tracks and topoclusters closely match GEANT4 behavior.The paper compares this autoregressive model with a parallel-token-sampling architecture.Source: §1 · pp. 2–3
  • The work evaluates behavior across 1,302 task directions, including rarely seen tasks and generalization to an unseen input combination.Each direction uses at most three input and three output modalities.Source: §1 · p. 3

2. Related Work

Prior machine-learning work has often targeted single stages, while end-to-end efforts span varied objectives, supervision levels and uses of physics information. This work proposes one model trained across input-output modality pairs, connecting truth-level information to detector hits and back.

2. Related Work

  • Earlier machine-learning examples address individual stages such as particle flow, tracking, flavour tagging and detector simulation.The paper presents combining several stages as a next step, noting that independent stages must each be built, validated and kept in sync.Source: §2 · pp. 3–4
  • End-to-end work spans pretraining objectives, supervision levels and the amount of physics information supplied to models.Examples include supervised classification, masked modelling, next-token prediction, contrastive and latent-space objectives, as well as models incorporating symmetries or differentiable physics.Source: §2 · p. 4
  • Parnassus and FlowSim are identified as generative models for simulation directions, each producing reconstructed objects conditioned on generator-level particles for one fixed direction.The paper also notes complementary efforts to make large open collider datasets available for foundation-model-scale pretraining.Source: §2 · p. 4
  • The proposed approach trains one model over input-output modality pairs, up to three modalities on each side, spanning truth-level information to detector hits and back.The paper connects this approach to pretraining strategies using supervision when truth modalities are included and physics knowledge through the choice of modalities and tasks.Source: §2 · p. 4

3. Methods

The methods build a shared token interface for seven event modalities, then use modality choices to define many input-output tasks. Two decoding architectures generate requested objects, and evaluation uses complementary measures of response, distributions, and input-conditioned mismatch.

3.1. Dataset and simulation

  • The dataset contains seven object modalities from single-jet events, with 88.9M training events and 100k each for validation and test.Events are produced with COCOA using Pythia 8 and a Geant4 detector model; tracks are parameterised fast-simulations of charged truth particles.Source: §3.1 · pp. 4–5

3.2. Tokenisation

  • Objects are represented by modality-specific content-codebook indices and fixed-grid position bins, giving one categorical token interface across modalities.Content features use residual VQ-VAE quantisation; η, cosφ, and sinφ are each binned into 1024 bins, with azimuth represented continuously across the φ = ±π boundary.Source: §3.2 · p. 5

3.3. Tasks

  • A task is defined by the input and requested output modalities, and training samples up to three of each side across 1302 possible tasks.The particle-flow direction is oversampled to about 10% of task samples, while accounting for around 2% of compute because its sequences are short.Source: §3.3 · pp. 6–7

3.4. Model architectures

  • Both architectures use frozen tokenisation, per-token categorical predictions, cross-entropy loss, and shared sampling and detokenisation procedures.Objects are ordered by descending pT or energy, and residual tokens within each object are decoded conditionally during training and inference.Source: §3.4 · p. 8

3.4.1. Iterative (autoregressive) decoding

  • nanoHEP is a decoder-only autoregressive transformer that samples one next token at a time under a structural mask.Its shared vocabulary assigns modality-specific ID ranges and structural tokens; the sequence header specifies requested modalities, and a key-value cache avoids recomputation during decoding.Source: §3.4.1 · p. 9

3.4.2. Parallel decoding

  • HEP4M uses a bidirectional encoder and modality-specific cardinality heads to allocate query slots that predict output tokens in one pass.Each slot has a fixed modality; sampled objects do not condition on one another, although residual levels within an object are decoded in dependency order.Source: §3.4.2 · p. 10

3.5. Evaluation metrics

  • The evaluation combines high-level response and resolution, per-object marginal comparisons, and input-conditioned classifier tests to assess generated objects from different angles.The C2ST AUC is 0.5 when the classifier cannot distinguish predictions from references; it is a lower bound under the stated classifier and sample protocol.Source: §3.5 · pp. 11–12

4. Results

The results show that the multi-task models span reconstruction and simulation directions, with nanoHEP-multi scoring closer to truth across the task matrix and generalising to rare and unseen input combinations. On particle flow, nanoHEP compares favourably with HGPflow, while simulation quality and computational cost expose trade-offs and limits.

4.1. The Particle Flow Task

  • nanoHEP-pflow was the best particle-flow model, outperforming HGPflow in jet-momentum resolution and classifier score; its object-level AUC was 0.69 against HGPflow’s 0.90.NanoHEP-pflow and HGPflow were within 0.002 in jet IQR and 1% in EMD.Source: §4.1 · p. 14, Table 2 · p. 14

4.1.1. Computational performance

  • Partial autoregression in HEP4M improved fidelity at modest additional computational cost, yielding an emerging Pareto frontier between inference cost and fidelity.This approach generates the leading particle individually and the remaining particles in parallel.Source: §4.1.1 · p. 15, Figure 7 · p. 15

4.2. The Simulation Task

  • Both simulators’ single-object distributions aligned closely with full simulation, but only nanoHEP-sim reproduced topocluster multiplicities; HEP4M-sim underproduced high-multiplicity events.The autoregressive model performed better than HEP4M at increased inference time; cardinality remained difficult, with broad per-event residuals for both models.Source: §4.2 · p. 15, Figure 8 · p. 16, Figure 9 · p. 17

4.3. All tasks

  • nanoHEP-multi had median fast-judge score 0.54 and was within 0.05 of chance on 779 of 1302 directions; HEP4M-multi scored 0.95 median, reproducing marginals but missing event-level structure.The most difficult tasks generated calorimeter cells: their median was 0.58, compared with 0.52 for the rest.Source: §4.3 · p. 17, Figure 10 · p. 18

4.4. Cross-task and unseen-task generalisation

  • At its on-task training budget, nanoHEP-multi was less distinguishable from truth than scratch-trained models on all eight tested tasks; at ten times that budget, it remained ahead on five reconstruction directions.At five times the budget, it was still ahead on six of eight directions; for three generation directions at ten times the budget, scratch-trained models were slightly closer to truth. Cell-input comparisons matched steps but not tokens because cell inputs are four times longer.Source: §4.4 · p. 18, Figure 11 · p. 19
  • Without further training, nanoHEP-multi generalized to an unseen four-modality particle-flow-like input combination, reaching IQR 0.069 and improving on the trained tracks, topoclusters, and cells combination.The authors describe this as the first particle-physics demonstration of a pretrained model performing at least as well on an input combination absent from training as on every trained one.Source: §4.4 · pp. 19–20, Table 4 · p. 20

4.5. Limitations

  • Generating calorimeter cells from truth particles remains a limitation: nanoHEP-multi’s outputs were fully separable from truth (0.998), and tokenisation separately made raw and round-tripped cells separable (AUC 0.96).The authors leave both issues to future work.Source: §4.5 · p. 20

5. Conclusions

The paper concludes that tokenisation supports stable training across many tasks and that the models can perform particle-flow reconstruction competitively. It also argues that fidelity should be considered alongside computational cost and reports transfer to rare and unseen task combinations.

5. Conclusions

  • HEP4M-pflow was within 11% of state of the art in jet resolution, while nanoHEP-pflow had the best jet resolution and lowest classifier AUC on the particle-flow task.For particle reconstruction from four inputs, nanoHEP-multi reached IQR 0.069 and AUC 0.66.Source: §5 · p. 20
  • The authors argue that foundation models should be compared on Pareto frontiers of fidelity against computational cost, rather than on outright performance alone.They link this recommendation to the cost of iterative approaches and report that realistic spectra and jets can still be distinguished from truth by an object-level classifier.Source: §5 · p. 21
  • NanoHEP-multi transferred to rarely seen directions and generalized zero-shot to an unseen four-modality combination, matching or improving on every trained subset.At its on-task budget it beat scratch-trained models on rarely seen tasks; the unseen combination required no further training.Source: §5 · p. 21

Code and data availability

The authors provide code, trained models, tokenisers, evaluation tools, and datasets for reproducing and using the work.

Code and data availability

  • The Apache-2.0 code repository includes both models, tokenisers, training and inference for modality choices, and evaluation metrics including the classifier two-sample test.Installation tests check reference predictions, a training loop, and token reproduction against the release.Source: Code and data availability · pp. 21–22
  • The Zenodo data release provides tokenised training, validation, and test sets for all seven modalities, with 88.9M, 100k, and 100k events respectively.It also contains seven tokeniser checkpoints, raw COCOA files, HGPflow predictions, and trained checkpoints for models in Tables 2 and 3.Source: Code and data availability · p. 22
  • A loader in the code repository rebuilds the exact token layout used for training.Source: Code and data availability · p. 22

Appendix A. Additional task results

The appendix tests additional task directions for the multi-task models, spanning reconstruction refinement, flash and jet simulation, and charged-energy subtraction. Results vary by direction: some comparisons favor one model or decoding choice, while refinement does not improve HGPflow resolution and charged-fraction prediction remains weak.

Appendix A. Additional task results

  • The appendix evaluates additional task directions beyond the main-text examples, with most directions trained through the multimodal sampler.The reported tasks include refinement, simulation, and charged-energy subtraction.Source: Appendix A · p. 31

Appendix A.1. Reconstruction refinement: HGPflow particles →truth particles

  • Neither model improves on HGPflow input resolution when mapping its particles to same-event truth particles: jet-response IQRs are 0.085 for nanoHEP-multi and 0.093 for HEP4M-multi.Figure A1 reports jet response and resolution versus target jet pT, plus angular residuals.Source: Appendix A.1 · p. 31, Figure A1 · p. 31

Appendix A.2. Flash simulation: truth particles →HGPflow particles

  • 0.076 versus 0.082 jet-response IQRs: HEP4M-multi beats nanoHEP-multi on flash simulation from truth particles to HGPflow particles.This direction skips the detector and is the paper’s one direction where the parallel model is ahead.Source: Appendix A.2 · p. 31, Figure A2 · p. 32

Appendix A.3. Jet simulation: truth jet →truth particles

  • 0.028 versus 0.096 IQRs with argmax, but 0.042 versus 1.24 with temperature-1 sampling: nanoHEP-multi is tighter for jet simulation, while sampled HEP4M-multi lies above the plotted range.The task generates a jet’s constituent particles from the jet alone.Source: Appendix A.3 · p. 32, Figure A3 · p. 33

Appendix A.4. Charged-energy subtraction: cell →charged fraction

  • 0.35 mean absolute error: HEP4M-multi beats the 0.45 best-constant and 0.49 other-cell baselines, but charged-fraction prediction remains weak rather than solved.The truth fraction is bimodal, with 35% of cells below 0.01 and 42% above 0.99; per-cell correlation is 0.36.Source: Appendix A.4 · p. 32, Figure A4 · p. 33

Appendix A.5. Detector simulation: additional detail

Additional evaluations examine detector simulation, topoclustering, tokeniser fidelity, aggregate task metrics, and particle-flow and simulation views. They expose both model behavior and a measurable distortion introduced by the calorimeter-cell tokeniser round trip.

Appendix A.5. Detector simulation: additional detail

  • 2.0 median per-event ratio: nanoHEP-multi produces about twice the truth topocluster multiplicity in the three-output detector-simulation direction, while track multiplicity matches.The appendix also reports per-event aggregates, per-modality aggregates, and example events.Source: Appendix A.5 · p. 34, Figure A5 · p. 34, Figure A6 · p. 34

Appendix A.6. Topoclustering

  • 1.03 median summed-energy response and 0.10 IQR: sampled nanoHEP-multi topoclustering has cluster count within 1% of the target.The task maps tracks and cells to topoclusters, with response measured against COCOA topoclusters.Source: Appendix A.6 · p. 34, Figure A8 · p. 36

Appendix A.7. Calorimeter-cell tokeniser round trip

  • The tokeniser round trip yields about 4% fewer cells, a 2.5% low median summed event energy, and more than 8% energy loss in a sixth of events; raw and decoded cells are separable at AUC 0.96.Below about 5 MeV, the decoded cell spectrum is too low; no model is involved in this round trip.Source: Appendix A.7 · pp. 34–35, Figure A9 · p. 36

Appendix A.8. All directions: multiplicity and marginals

  • Figure A10 reports normalised multiplicity Wasserstein-1 and median marginal Kolmogorov-Smirnov distance across all 1302 task directions for both multi-task models.These are the two measures not shown in Figure 10.Source: Appendix A.8 · p. 37, Figure A10 · p. 37

Appendix A.9. Additional particle flow and simulation views

  • The appendix supplies particle-flow marginal, multiplicity, and angular-distance views for single-task and multi-task models, plus sampled track and topocluster multiplicity comparisons.Figures A11–A14 cover particle-flow distributions and jet angular distances; Figure A15 covers simulation multiplicities.Source: Appendix A.9 · pp. 37–38, Figure A11 · p. 38, Figure A12 · p. 38

Appendix B. Decoding studies: argmax versus sampling

The decoding study contrasts argmax and sampling for reconstruction and generation. Argmax collapses sharply peaked generated marginals, while temperature-1 sampling restores track-width agreement with Geant4; the preferred decoding also depends on the task and fidelity measure.

Appendix B. Decoding studies: argmax versus sampling

  • At T = 1, sampled track impact-parameter widths match Geant4 to within about 1%, whereas argmax makes d0 and z0 far too narrow for both models.The widths vary smoothly with sampling temperature.Source: Appendix B · p. 40, Figure B1 · p. 41, Figure B2 · p. 41
  • Sampling preserves the track and topocluster angular identity in nanoHEP-sim, with cos²φ + sin²φ = 1.000 ± 0.011 and 1.000 ± 0.006, respectively; independent HEP4M sampling does not enforce it.Before renormalisation, HEP4M sums are 0.97 ± 0.21 for tracks and 0.97 ± 0.30 for topoclusters; HEP4M topocluster marginals are also too flat between peaks.Source: Appendix B · p. 40
  • On particle flow, argmax gives tighter nanoHEP-pflow jet response (IQR 0.074 versus 0.093 with sampling) but biased median response (0.978 versus 0.998) and collapsed multiplicity; sampling improves distributional fidelity.For nanoHEP-multi, AUC is 0.52 ± 0.01 with sampling versus 0.70 with argmax; on the Table 4 input combinations, argmax gives AUC 0.65 to 0.71 versus about 0.55 with sampling, at about 0.025 IQR cost.Source: Appendix B · p. 40, Figure B3 · p. 42

Appendix C. Two-pass HEP4M decoding

The appendix tests whether a second HEP4M decoder pass can condition particle generation on a predicted leading particle. The two-pass setup improves jet-response resolution over the second model’s single-pass unconditional mode, recovering about 70% of the difference from HGPflow at the cost of another pass.

Appendix C. Two-pass HEP4M decoding

  • The second pass conditions on the first pass’s predicted leading particle; no truth is used at inference, and both passes use argmax decoding.The 64M-parameter particle-flow-only HEP4M versions differ in whether the first output slot is trained to hold the true leading-pT particle on 80% of events or left empty.Source: Appendix C · p. 42
  • 0.0785 jet-response IQR with two passes versus 0.0865 for one unconditional pass, recovering about 70% of the difference from HGPflow’s 0.075 at the cost of a second decoder pass.Figure C1 compares single-pass and two-pass 64M HEP4M with HGPflow.Source: Appendix C · pp. 42–43, Figure C1 · p. 43

Appendix D. Data and model scaling

In a particle-flow-only scaling study, nanoHEP’s jet pT response IQR improves modestly as parameter count grows and changes little with more training events.

Appendix D. Data and model scaling

  • 0.084 at 6M parameters falls to 0.075 at 22M and 0.074 at 94M, with little further change at the largest model.Figure D1 plots jet pT response IQR against parameter count, with HGPflow as a dashed reference.Source: Appendix D · p. 43, Figure D1 · p. 44
  • 0.077 with 10^7 events falls to 0.075 with 8.9×10^7 events, while the two runs use different learning-rate schedules.Figure D2 also shows response and residual distributions for the two training-set sizes.Source: Appendix D · p. 43, Figure D2 · p. 44

Appendix E. Particle-cloud measures of particle flow

This block reports particle-cloud measures and implementation details spanning Appendix E and Appendices F.1–F.4, including a particle-flow comparison, tokenisation, dataset scope, cardinality prediction, and token factorisation.

Appendix E. Particle-cloud measures of particle flow

  • EMD measures the minimum angular-distance-weighted cost of moving reconstructed-particle pT onto truth particles, using the unbalanced form with R = 1.The nanoHEP tokeniser round trip alone contributes a median EMD of 0.17 GeV, an order of magnitude below model values.Source: Appendix E · p. 45
  • The compatibility judge separates true input–truth pairings from mismatched pairings with AUC 0.98, and reports each reconstruction’s mean logit difference from its own truth.A positive difference means the reconstruction scores as more strongly bound to the inputs; the value is an ordering, not an accuracy, and its sign depends on judge seed.Source: Appendix E · p. 45
  • The nanoHEP-pflow jet response is narrower than HGPflow by a paired IQR difference of −0.0021 (95% interval [−0.0028,−0.0013]), while their EMDs are within 1%.HEP4M models are 0.3 to 0.5 GeV behind on EMD, in the same order as their jet resolution; HGPflow has the highest AUC among HGPflow and nanoHEP rows, while its point estimates are as good as theirs.Source: Appendix E · p. 45, Table E1 · p. 45

Appendix F.1. Tokenisers

  • Content tokenisers use residual VQ-VAEs with three dimension-8 codebooks, while jet tokenisation uses eight parallel dimension-16 codebooks; position coordinates are uniformly binned.Tokenisers are trained per modality, then frozen; the figure shows residual codebooks for content and independent position binning.Source: Appendix F.1 · p. 46, Figure F1 · p. 46

Appendix F.2. Dataset

  • Each event contains one jet whose pT is the vector sum of truth-particle pT, with at most 256 cells stored.Source: Appendix F.2 · p. 46

Appendix F.3. Cardinality head

  • HEP4M predicts each modality’s object count with a softmax over 0,…,Nmax; inference masks the zero bin in rare directions where argmax would produce empty output.A regression head collapsed to a biased point estimate without spread, while a per-slot indicator head did not control total count.Source: Appendix F.3 · p. 46

Appendix F.4. Token factorisation

  • Both models sample an object’s residual token levels in order, but only nanoHEP conditions each object on preceding objects.Position tokens are sampled from separate heads.Source: Appendix F.4 · p. 47

Appendix F.5. Training and inference

Training uses different losses for the autoregressive and parallel models, and inference uses argmax for reconstruction and sampling at temperature 1 for generation.

Appendix F.5. Training and inference

  • nanoHEP trains with next-token cross-entropy on the output region, while HEP4M sums masked-target per-codebook cross-entropy and cardinality loss.Training uses bf16 and evaluation fp32; inference uses a key-value cache and a well-formed-sequence mask for nanoHEP.Source: Appendix F.5 · p. 47
  • Reconstruction is decoded with argmax, while generation is sampled at T = 1.Table F1 collects the multi-task models’ training settings.Source: Appendix F.5 · p. 47, Table F1 · p. 47

Related papers

Loading 2609.31862v1…