Prompting Particle Physics: Tokenized Multi-modal Foundation Models for Combinatorially Many Tasks
Authors: Nilotpal Kakati, Daniel Murnane, Baran Hashemi, Samuel Klein, Jeffrey Krupa, Eilam Gross, Lukas Heinrich, Michael Kagan
Organizations: Weizmann Institute of Science, Rehovot, Israel · Niels Bohr Institute, University of Copenhagen, Copenhagen, Denmark · Lawrence Berkeley National Laboratory, Berkeley, CA, USA · Max Planck Institute for Mathematics in the Sciences, Leipzig, Germany · SLAC National Accelerator Laboratory, Menlo Park, CA, USA · Technical University of Munich, Munich, Germany
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.
Figures & tables
Figure 1: (a) The seven object modalities of a single jet, with the two-letter codes used in the figures of this paper. At truth level: the particles (TP), the jet (TJ) and the charged-energy fraction of each calorimeter cell (CF). At detector level: the calorimeter cells (CE), tracks (TK), topoclusters (CL) and the HGPflow particles reconstructed from them (RP). Each object is a short block of tokens. (b) A task is a choice of which modalities to provide and which to request. Four of the 1302 directions are shown.
modality
code
content features (jointly quantised)
content tokens
position tokens
per event
tracks
TK
pT , d0 ( ∣d0∣≤0.2 mm), z0 ( ∣z0∣≤0.4 mm)
3×256
3
3.6
topoclusters
CL
E , ρ , logit(EM fraction)
3×256
3
11.6
truth particles
TP
pT , E
3×128
3
6.7
HGPflow particles
RP
pT
3×128
3
6.5
truth jets
TJ
pT , E
8×128
3 per head
1
calorimeter cells
CE
E0.3 per cell, 8×8 patches ( 4×4 in the coarsest layer)
3×2048 per patch
3 (nanoHEP) / none (HEP4M)
156 patches
Table 1: The seven modalities, the two-letter codes used in the figures, and their tokenisers. Content tokens are written as nq×K (residual levels × codebook size); the jet has eight heads of one token each. Position tokens are the 1024 -bin encodings of (η,cosϕ,sinϕ) ; for a patch they are its binned centre, used by nanoHEP only, since in HEP4M a patch’s location is its index in the fixed 156 -patch grid. The last column is the mean number of objects per event that the models see, within the window around the jet axis, on the test sample (on average 135 of the at most 256 stored cells are non-empty).
Figure 2: The three tokeniser families. (a) Set-like objects (tracks shown; topoclusters and particles have the same form): each object becomes three residual content tokens and three binned position tokens (the residual cascade is drawn in Fig. 33 ). (b) The jet: a 128-dimensional latent split into eight heads, each matched to its own codebook; each head’s index is paired with the jet’s three position tokens. (c) The calorimeter-cell image (and the charged-energy fraction): six layers cut into 156 patches of three tokens each, with no position tokens in HEP4M since a patch’s location is its index (nanoHEP adds the binned patch centre).
Figure 3: A task is an input/output modality split of the same tokenised event, written as one token sequence: a task header (shaded) of the requested output tags closed by the task separator, the input blocks, the output blocks, and end-of-sequence. Top: particle flow (tracks and topoclusters in, truth particles out). Middle: detector simulation (truth particles in; tracks, cells and topoclusters out). Bottom: one sample from the generalist sampler, which draws up to three input and up to three output modalities from the seven available ( 1302 distinct tasks); here truth cell fractions, HGPflow particles and jets in, and topoclusters, tracks and truth particles out. Coloured tokens are objects; grey tokens are the structural header, role tags, and end-of-sequence.
Figure 4: Autoregressive (iterative) decoding model: a decoder-only GPT-2-class transformer over the shared token vocabulary. Each token attends causally to all previous tokens and the next token is sampled from the predicted distribution; a structural mask constrains which tokens are legal at each step, and a key-value cache avoids recomputation during decoding.
Figure 5: Parallel-decoding model. Input-modality token blocks are embedded and read by a bidirectional encoder; per-modality cardinality heads set the number of output slots; decoder query slots (seeded with modality and position embeddings) cross-attend to the encoder and predict all output tokens in one pass, which are detokenised back to objects.
Figure 6: Jet-level agreement on particle flow, for the single-task models (nanoHEP-pflow, HEP4M-pflow), the multi-task models (nanoHEP-multi, HEP4M-multi) and HGPflow: the reconstructed-to-target jet pT response (left; HEP4M models above, nanoHEP models below, each against HGPflow as the filled distribution; smoothed with a Gaussian kernel of width 0.0056 , the same for every curve) and its IQR/median in bins of target jet pT (right). The legend gives the median and IQR of the response. Two further nanoHEP-multi curves (purple) use the extra inputs of Table 4 : tracks, topoclusters and cells, and topoclusters, HGPflow particles and cells, with no further training.
jet pT response
classifier
model
median
IQR
AUC
HGPflow
0.986
0.0754
0.902
nanoHEP-pflow
0.978
0.0733
0.691
nanoHEP-pflow
0.982
0.0758
0.697
(matched, 29k steps)
nanoHEP-multi
0.990
0.0781
0.703
Table 2: Particle flow reconstruction performance, using argmax decoding. We report the jet pT response median and IQR, and the classifier AUC of Sec. 3.5 . The best value in each column is in bold, and all measures are in Table 5 .
Figure 7: Jet pT -response IQR versus approximate inference time per event for HEP4M-pflow, HEP4M (two-pass) of C , and nanoHEP-pflow. The dotted line is the HGPflow IQR, which is not timed on this axis.
model
tracks
topoclusters
nanoHEP-sim
0.541
0.517
nanoHEP-multi
0.554
0.525
HEP4M-sim
0.972
0.984
HEP4M-multi
0.987
0.985
Table 3: AUC of classifier two-sample tests for detector simulation (truth particles → tracks and topoclusters) separating generated tracks or topoclusters from Geant4’s, conditioned on the event’s truth particles. Decoded by sampling at temperature 1 . Mean of five retrainings shown, with a spread of at most 0.007 . Best value in each column in bold.
Figure 8: Generated track pT , η , ϕ , d0 and z0 for the single-task (solid) and multi-task (dashed) models, sampled at T=1 , against Geant 4. The KS distance to Geant 4 is given for each model.
Figure 9: Generated topocluster E , η , ϕ , EM fraction and ρ for the single-task (solid) and multi-task (dashed) models, sampled at T=1 , against Geant 4. The KS distance to Geant 4 is given for each model.
Figure 10: All 1302 task directions for HEP4M-multi (left) and nanoHEP-multi (right), on the grid of input subsets (rows) versus output subsets (columns): the fast-judge score (Sec. 4.3 ; compare values only within this figure). Grey cells show undefined tasks, since output should not contain an input modality. The multiplicity and marginal distances for the same directions are in Fig. 21 .
Figure 11: Typical tasks evaluated with nanoHEP-multi and with nanoHEP trained from scratch on each direction at nanoHEP-multi’s on-task budget ( 205 steps), at five times ( 1025 steps) and at ten times that budget ( 2000 steps). The radius is 1−AUC of the fast-judge score of Sec. 4.3 . Simulation-side directions are in italics and decoded by sampling at T=1 ; reconstruction-side directions are decoded by argmax. TK = tracks, CL = topoclusters, CE = calorimeter cells, TP = truth particles, TJ = truth jets.
jet pT response
classifier
inputs
trained
median
IQR [ 95% interval]
AUC
tracks + topoclusters (benchmark)
yes
0.990
0.0781 [ 0.0775 , 0.0787 ]
0.703
tracks + cells
yes
1.006
0.0875 [ 0.0867 , 0.0882 ]
-
tracks + topoclusters + cells
yes
0.999
0.0704 [ 0.0698 , 0.0710 ]
0.650
tracks + topoclusters + HGPflow particles
yes
0.994
0.0752 [ 0.0745 , 0.0758 ]
0.712
tracks + HGPflow particles + cells
yes
0.999
0.0732 [ 0.0726 , 0.0738 ]
0.672
Table 4: Particle flow by nanoHEP-multi from different input combinations, with no further training, argmax decoding: jet pT response median and IQR with the 95% bootstrap interval ( 1000 event-level resamples). The trained column marks whether the task was seen during training. The AUC is the set-transformer test of Sec. 3.5 conditioned on the benchmark inputs, five retrainings (spread ≤0.009 ).
Appendix figures & tables24 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 12: Jet pT response, response IQR/median versus target jet pT , and the ΔR , η and ϕ residuals between reconstructed and target jets for the reconstruction-refinement task.
Figure 13: Flash simulation (truth particles → HGPflow particles), argmax decoding: jet pT response, response IQR/median versus target jet pT , and the ΔR , η and ϕ residuals between reconstructed and target jets.
Figure 14: Jet pT response, response IQR/median versus target jet pT , and the ΔR , η and ϕ residuals between reconstructed and target jets for the jet-simulation task, decoded with argmax and with temperature-1 sampling.
Figure 15: Charged-energy subtraction by HEP4M-multi: the charged-fraction distribution for truth and prediction, the per-cell truth versus predicted fraction, and the per-event mean absolute error.
Figure 16: Per-event aggregates of the simulated detector response versus truth (sampled decoding): total cell energy, topocluster multiplicity, track multiplicity, and summed track pT , for nanoHEP-multi and HEP4M-multi.
Figure 17: Per-event truth versus generated multiplicity and summed magnitude for particles → tracks (left) and particles → topoclusters (right).
Figure 18: Example events in η - ϕ : truth, HEP4M-multi and nanoHEP-multi, sampled decoding. Cell energies are shown as heatmaps, topoclusters as circles (area proportional to energy) and tracks as crosses (size proportional to pT ).
Figure 19: Topoclustering (tracks and cells → topoclusters) by nanoHEP-multi and HEP4M-multi: single-object marginals, topocluster multiplicity, and the per-event summed-energy response.
Figure 20: Calorimeter cells before and after a tokeniser round trip (tokenised and decoded back, no model involved): cell energy spectrum (left), cells per event (middle), and the ratio of decoded to raw summed cell energy per event (right).
Figure 21: All 1302 task directions for HEP4M-multi (left) and nanoHEP-multi (right): normalised multiplicity Wasserstein-1 (top) and median marginal Kolmogorov-Smirnov distance (bottom), on the grid of Fig. 10 .
Figure 22: Per-particle pT , η , and ϕ distributions for HGPflow, HEP4M-pflow, and nanoHEP-pflow.
Figure 23: Per-event particle multiplicity for HGPflow, HEP4M-pflow, and nanoHEP-pflow.
Figure 24: Per-particle pT , η , and ϕ distributions for HGPflow, HEP4M-multi and nanoHEP-multi on particle flow.
Figure 25: ΔR , Δη and Δϕ between the reconstructed and target jets on particle flow, for HGPflow and the four particle flow models of Table 2 .
Figure 26: Per-event track (left) and topocluster (right) multiplicity, target versus generated (sampled decoding).
Figure 27: Track d0 and z0 and topocluster EM fraction and barycentre radius ρ in detector simulation, for nanoHEP-sim (top) and HEP4M-sim (bottom), decoded by argmax and by sampling at T=1 , against Geant4. The axes span the tokeniser’s range; the 1.3% of Geant4 tracks outside it are in the edge bins.
Figure 28: Track d0 and z0 widths (half the 16 to 84% interval) from nanoHEP-sim, and their Kolmogorov-Smirnov distance to Geant4, versus sampling temperature T , compared with argmax decoding and with Geant4.
Figure 29: Particle flow with argmax and temperature-1 sampled decoding, for nanoHEP-multi (top row) and nanoHEP-pflow (bottom row): jet pT response, particle multiplicity, and particle pT spectrum.
Figure 30: Jet pT response, response IQR/median versus target jet pT , and ΔR , Δη , Δϕ for the single-pass and two-pass 64 M HEP4M, and HGPflow.
Figure 31: Jet pT response IQR versus parameter count for nanoHEP trained on particle flow alone, with HGPflow as the dashed line.
Figure 32: Jet pT response, response IQR/median versus target jet pT , and the ΔR , η and ϕ residuals between reconstructed and target jets, for nanoHEP trained on particle flow alone with 107 and 8.9×107 events.
jet pT response
EMD (GeV)
classifier
compatibility
model
median
IQR
median
vs HGPflow
AUC
seed 1
seed 2
HGPflow
0.986
0.0754
2.09
0
0.902
+0.64
+0.98
nanoHEP-pflow
0.978
0.0733
2.09
+0.024±0.003
0.691
+0.53
−0.31
nanoHEP-multi
0.990
0.0781
2.16
+0.048±0.003
0.703
+0.52
−0.40
HEP4M-pflow
0.994
0.0836
2.62
+0.31±0.01
0.940
+0.30
−1.07
HEP4M-multi
0.986
0.0898
2.78
+0.46±0.01
0.939
+0.27
−1.48
Appendix
Table 5: All particle flow measures, argmax decoding. Jet pT response median and IQR as in Fig. 6 ; EMD median over events and median per-event difference to HGPflow with its bootstrap standard error; classifier AUC as in Table 2 ; compatibility as the judge logit difference, for two judge seeds. Best value in the jet, EMD-median and AUC columns in bold.
Figure 33: Residual VQ-VAE tokenisation of a single object. Top: the object’s content features are encoded to one latent z and quantised by a cascade of codebooks, each quantising the residual left by the previous, giving nq content tokens. Bottom left: the position coordinates (η,cosϕ,sinϕ) are binned independently. Bottom right: the resulting per-object token block.
nanoHEP-multi
HEP4M-multi
optimiser
AdamW [ 109 ]
schedule-free AdamW [ 110 ]
betas, weight decay
0.9/0.95 , 0.1
0.95/0.98 , 0.05
peak LR, warmup
1.2×10−3 , 2 k steps
1.0×10−3 , 2.5 k steps
schedule
cosine [ 111 ] to 1.2×10−4 , 285 k steps
schedule-free
effective batch
4096
16,384
checkpoint
296 k steps
lowest validation loss
Appendix
Table 6: Training settings of the multi-task models.
Full detector simulation and reconstruction of collider events are projected to become major bottlenecks at the High-Luminosity Large Hadron Collider, motivating the development of fast, ML-based surrogates. At the same time, LLMs have driven fast progress in generative discrete modeling: autoregressive transformers trained on tokenized data now represent the state of the art across a range of generative tasks. We extend the discrete modeling paradigm by introducing a particle-level generative model trained on residual-quantized full-event data. We demonstrate the ability of this model family to perform conditional generation from detector-stable particles; we study its scaling behavior across a range of dataset and model sizes, characterize the effects of repeated data exposure and demonstrate that token-level loss systematically predicts downstream physical fidelity. These results provide an empirical framework for scalable collider full-event generation based on residual-quantized representations.
Foundation models are increasingly being pursued in particle and nuclear physics, but existing approaches remain strongly tied to individual experiments through detector-specific architectures or pre-training objectives, limiting their reuse across sensing modalities. We show that a point cloud self-distillation framework yields a substantially more general sensor-level pre-training recipe. We show that the same refined architecture and objective can be independently pre-trained with minimal changes on three qualitatively different detector modalities: liquid argon time projection chamber (LArTPC), collider TPC, and water Cherenkov. Using 1,000 labeled images for downstream task adaptation, Panda V2 matches or exceeds specialized foundation-model baselines trained with orders of magnitude more supervision, matching state-of-the-art particle-clustering performance with 70x fewer labeled events on sPHENIX while substantially improving particle identification, and on LArTPC data matching Panda (arXiv:2512.01324) particle reconstruction with up to 1,000x fewer labels. Beyond reconstruction, simple linear probes reveal physically meaningful latent structure associated with particle causality and track curvature.
Samuel Young, César Jesús-Valls, Kazuhiro Terao
Stanford University · Stanford, CA 94305 · European Organization for Nuclear Research (CERN) +3
Recently observed empirical scaling laws describe the performance of foundation-type models as three independent key quantities -- dataset size, compute, and model parameters -- are modified. Extracting these scaling laws informs the training of large complex models for which the tuning of hyperparameters in traditional ways is not feasible. This work for the first time explores if scaling laws can also be observed for the task of particle jet generation -- both relevant as a pre-training objective for foundation models and as in-situ simulation by itself. We indeed replicate the key logarithmic scaling law behavior for model-size scaling. Beyond studying the next token prediction validation loss of the generative model, we also study the sliced Wasserstein distance of five physical quantities that are not immediately available to the model during training. Our study shows that this quantity is monotonically related to the next token prediction validation loss, meaning that this loss is indeed a good proxy for the physics performance. For the scaling with dataset size and compute, we observe substantially weaker scaling behavior of both the loss and the sliced Wasserstein distance. We analyze this behavior by introducing the concept of a learnable window, and argue that autoregressive next token prediction on jet constituents exhibits comparatively rapid saturation relative to language-model studies. We discuss possible origins of this behavior, including the stochastic nature of QCD radiation and differences between generative and supervised learning tasks in collider physics.
Oz Amram, Darius A. Faroughy, Tjarko Gerdes +5
Fermi National Accelerator Laboratory, Batavia, IL 60510, USA · NHETC, Dept. of Physics and Astronomy, Rutgers University, Piscataway, NJ 08854, USA · Institute for Experimental Physics, Universität Hamburg, Luruper Chaussee 149, 22761 Hamburg, Germany +1