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

hep-phSep 25, 2026

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

Abstract

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	$\sqrt{p_{T}}$ , $d_{0}$ ( $\|d_{0}\|\leq 0.2$ mm), $z_{0}$ ( $\|z_{0}\|\leq 0.4$ mm)	$3\times 256$	$3$	$3.6$
topoclusters	CL	$\sqrt{E}$ , $\rho$ , logit(EM fraction)	$3\times 256$	$3$	$11.6$
truth particles	TP	$\sqrt{p_{T}}$ , $\sqrt{E}$	$3\times 128$	$3$	$6.7$
HGPflow particles	RP	$\sqrt{p_{T}}$	$3\times 128$	$3$	$6.5$
truth jets	TJ	$\sqrt{p_{T}}$ , $\sqrt{E}$	$8\times 128$	$3$ per head	$1$
calorimeter cells	CE	$E^{0.3}$ per cell, $8\times 8$ patches ( $4\times 4$ in the coarsest layer)	$3\times 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

n_{q}\times K

(residual levels

\times

codebook size); the jet has eight heads of one token each. Position tokens are the

1024

-bin encodings of

(\eta,\cos\phi,\sin\phi)

; 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).

Panel of Figure 2 — 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

p_{T}

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

p_{T}

(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.

(matched, 29k steps)	$0.982$	$0.0758$	$0.697$
	jet $p_{T}$ 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$
nanoHEP-multi	$0.990$	$0.0781$	$0.703$

Table 2: Particle flow reconstruction performance, using argmax decoding. We report the jet

p_{T}

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 $p_{T}$ -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

\rightarrow

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

p_{T}

\eta

\phi

d_{0}

and

z_{0}

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

\eta

\phi

, EM fraction and

\rho

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-\mathrm{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 $p_{T}$ 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

p_{T}

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

\leq 0.009

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 12: Jet $p_{T}$ response, response IQR/median versus target jet $p_{T}$ , and the $\Delta R$ , $\eta$ and $\phi$ residuals between reconstructed and target jets for the reconstruction-refinement task.

Figure 13: Flash simulation (truth particles $\rightarrow$ HGPflow particles), argmax decoding: jet $p_{T}$ response, response IQR/median versus target jet $p_{T}$ , and the $\Delta R$ , $\eta$ and $\phi$ residuals between reconstructed and target jets.

Figure 14: Jet $p_{T}$ response, response IQR/median versus target jet $p_{T}$ , and the $\Delta R$ , $\eta$ and $\phi$ 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 $p_{T}$ , for nanoHEP-multi and HEP4M-multi.

Panel of Figure 17 — Figure 17: Per-event truth versus generated multiplicity and summed magnitude for particles $\rightarrow$ tracks (left) and particles $\rightarrow$ topoclusters (right).

Figure 18: Example events in $\eta$ - $\phi$ : 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 $p_{T}$ ).

Figure 19: Topoclustering (tracks and cells $\rightarrow$ 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 $p_{T}$ , $\eta$ , and $\phi$ 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 $p_{T}$ , $\eta$ , and $\phi$ distributions for HGPflow, HEP4M-multi and nanoHEP-multi on particle flow.

Figure 25: $\Delta R$ , $\Delta\eta$ and $\Delta\phi$ between the reconstructed and target jets on particle flow, for HGPflow and the four particle flow models of Table 2 .

Panel of Figure 26 — Figure 26: Per-event track (left) and topocluster (right) multiplicity, target versus generated (sampled decoding).

Figure 27: Track $d_{0}$ and $z_{0}$ and topocluster EM fraction and barycentre radius $\rho$ 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 $d_{0}$ and $z_{0}$ 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 $p_{T}$ response, particle multiplicity, and particle $p_{T}$ spectrum.

Figure 30: Jet $p_{T}$ response, response IQR/median versus target jet $p_{T}$ , and $\Delta R$ , $\Delta\eta$ , $\Delta\phi$ for the single-pass and two-pass $64$ M HEP4M, and HGPflow.

Figure 31: Jet $p_{T}$ response IQR versus parameter count for nanoHEP trained on particle flow alone, with HGPflow as the dashed line.

Figure 32: Jet $p_{T}$ response, response IQR/median versus target jet $p_{T}$ , and the $\Delta R$ , $\eta$ and $\phi$ residuals between reconstructed and target jets, for nanoHEP trained on particle flow alone with $10^{7}$ and $8.9\times 10^{7}$ events.

	jet $p_{T}$ 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\pm 0.003$	0.691	$+0.53$	$-0.31$
nanoHEP-multi	$0.990$	$0.0781$	$2.16$	$+0.048\pm 0.003$	$0.703$	$+0.52$	$-0.40$
HEP4M-pflow	0.994	$0.0836$	$2.62$	$+0.31\pm 0.01$	$0.940$	$+0.30$	$-1.07$
HEP4M-multi	$0.986$	$0.0898$	$2.78$	$+0.46\pm 0.01$	$0.939$	$+0.27$	$-1.48$

Appendix

Table 5: All particle flow measures, argmax decoding. Jet

p_{T}

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

n_{q}

content tokens. Bottom left: the position coordinates

(\eta,\cos\phi,\sin\phi)

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\times 10^{-3}$ , $2$ k steps	$1.0\times 10^{-3}$ , $2.5$ k steps
schedule	cosine [ 111 ] to $1.2\times 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.

Explore similar work

Scaling Collider Event Generation with Residual-Quantized Tokens
Sep 30, 2026Dan Godi, Dmitrii Kobylianskii, Eilam GrossAutoregressive Generation Conditional Generative Modeling
Weizmann Institute of Science, Rehovot, Israel
Panda Diplomacy: Foundation Model Pre-training across Particle Imaging Detectors for High Energy and Nuclear Physics
Sep 1, 2026Samuel Young, César Jesús-Valls, Kazuhiro TeraoSelf-Supervised Pre-Training Point Cloud Learning
Stanford University · Stanford, CA 94305 · European Organization for Nuclear Research (CERN) +3
Neural Scaling Laws for Jet Generation
May 27, 2026Oz Amram, Darius A. Faroughy, Tjarko Gerdes +5High-Energy Physics Autoregressive Generation
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

Sep 30, 2026J/K move · Enter open · S save

Scaling Collider Event Generation with Residual-Quantized Tokens

Dan Godi, Dmitrii Kobylianskii, Eilam Gross

Weizmann Institute of Science, Rehovot, Israel

Autoregressive Generation Conditional Generative Modeling

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Scaling Collider Event Generation with Residual-Quantized Tokens

Panda Diplomacy: Foundation Model Pre-training across Particle Imaging Detectors for High Energy and Nuclear Physics

Neural Scaling Laws for Jet Generation