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

Technical Summary: Prompting Particle Physics with Tokenized Multi-modal Foundation Models

Problem Statement
Historically, extracting physics from collider data (e.g., at the LHC) has relied on long chains of specialized, heuristic algorithms, each tuned to a single step (e.g., tracking, clustering, simulation). While machine learning has successfully optimized individual links in these chains, it remains an open question whether a single model can replace multiple links simultaneously while maintaining interpretability through intermediate objects (tracks, clusters, particles). Current approaches often treat reconstruction and simulation as distinct, single-stage problems, failing to leverage the potential of a unified framework that can handle diverse data modalities from low-level sensor features to high-level physics objects.

Methodology
The authors propose treating particle physics reconstruction and simulation as a single multi-modal input-output problem. The core methodology involves tokenization and the training of foundation models capable of mapping any subset of modalities to any other.

1. Tokenization Strategy:

   • Every object in a jet (tracks, topoclusters, particles, jets, calorimeter cells, etc.) is converted into a discrete sequence of "tokens" from a modality-specific codebook.
   • Content: Features (e.g., $p_T$, $d_0$, energy) are quantized using a Residual Vector Quantized Variational Autoencoder (RVQ-VAE), producing a sequence of residual tokens.
   • Position: Spatial attributes ($\eta, \cos\phi, \sin\phi$) are binned into fixed grids.
   • Modalities: The framework utilizes seven modalities: Tracks (TK), Topoclusters (CL), Truth Particles (TP), Reconstructed Particles (RP), Truth Jets (TJ), Calorimeter Cells (CE), and Charged-Energy Fraction (CF).
   • Task Definition: A task is defined purely by the choice of input modalities (the "prompt") and output modalities (the "completion"). This allows for 1,302 possible task directions (combinations of up to three inputs and three outputs).

2. Model Architectures:
   The authors train two distinct foundation models on the same tokenized data:

   • nanoHEP (Autoregressive): A decoder-only transformer (GPT-2 family) that generates output tokens sequentially. It uses a shared vocabulary for all modalities and structural masks to enforce sequence grammar (e.g., ensuring the correct number of objects and modality transitions). It learns the joint distribution factorized as $q_\theta(y|x) = \prod q_\theta(y_t | x, y_{<t})$.
   • HEP4M (Parallel): An encoder-decoder transformer (adapted from the 4M architecture) that predicts all output tokens in a single forward pass. It uses cardinality heads to determine the number of objects per modality and cross-attention mechanisms to generate tokens in parallel.

3. Training Regimen:

   • The models are trained on a dataset of single-jet events (88.9M training events) generated by COCOA (a Geant4-based simulation).
   • During training, the sampler uniformly selects up to three input modalities and up to three output modalities to form a task.
   • The particle flow task ($\{TK, CL\} \to \{TP\}$) is over-sampled (~10% of steps) to ensure stability, though it constitutes only ~2% of the total compute due to sequence length.

Key Contributions

1. Dual Foundation Models: Introduction of nanoHEP (autoregressive) and HEP4M (parallel), both capable of operating from sensor-level readouts to high-level physics features.
2. Unified Multi-Modal Framework: Demonstration that a single set of weights can perform reconstruction, simulation, and mixed-modality tasks (e.g., charged energy subtraction, flash simulation) simply by changing the input/output prompt.
3. Systematic Evaluation: A comprehensive evaluation across all 1,302 task directions, assessing performance on rarely seen tasks and generalization to unseen input combinations (zero-shot composition).
4. Performance Trade-off Analysis: A detailed comparison of the computational cost versus fidelity between autoregressive and parallel decoding strategies.

Results

• Particle Flow Reconstruction:
  • nanoHEP-pflow outperforms the state-of-the-art HGPflow algorithm in classifier AUC (0.691 vs. 0.902), indicating its output is harder to distinguish from truth when conditioned on event inputs. The paper notes that nanoHEP outperforms HGPflow in reconstructing jet momentum in both resolution level and classifier score, though HGPflow shows a slightly higher IQR (0.0754) compared to nanoHEP-pflow (0.0733) in Table 2.
  • nanoHEP-multi achieves even better performance (IQR 0.069) when provided with additional inputs (tracks, topoclusters, HGPflow particles, and cells), demonstrating the benefit of multi-modal fusion.
• Simulation:
  • nanoHEP-sim generates tracks and topoclusters that closely match Geant4 distributions (AUC ~0.54 for tracks, near the tokenization limit).
  • HEP4M models reproduce marginal distributions well but struggle with event-level structure, resulting in high AUCs (~0.98) when distinguishing generated events from truth, indicating a loss of correlation between objects.
• Generalization and Transfer:
  • Task Transfer: Models trained on the full task matrix outperform models trained from scratch on specific tasks, even when the from-scratch models are given 10x the compute budget for that specific task.
  • Zero-Shot Composition: The autoregressive model successfully generalizes to a task with four input modalities (tracks + topoclusters + HGPflow particles + cells), a combination never seen during training (which was capped at three inputs). This "zero-shot" task achieved an IQR of 0.0688, which matches or improves upon the performance of trained subsets (e.g., the trained "tracks + topoclusters + cells" task had an IQR of 0.0704).
• Limitations:
  • Generating calorimeter cells (CE) from truth particles remains difficult; the generated cells are easily separable from truth (AUC 0.998).
  • Parallel decoding (HEP4M) struggles with cardinality (multiplicity) and event-level correlations compared to the autoregressive approach.

Significance and Claims
The paper claims that suitably high-fidelity tokenization enables the training of stable, multi-modal foundation models for particle physics with minimal hyperparameter tuning. The authors posit that:

1. Unification: Reconstruction and simulation can be unified into a single "any-to-any" generative framework, reducing the complexity of maintaining separate specialized algorithms.
2. Efficiency: Multi-task training is computationally more efficient than training separate models for each task, offering better performance at lower compute budgets.
3. Iterative Necessity: Autoregressive decoding appears necessary to capture the full multimodal posterior and event-level structure, suggesting that future foundation models should be evaluated on Pareto frontiers of fidelity versus computational cost rather than raw performance alone.
4. Generalization: The models demonstrate a capacity for zero-shot generalization to unseen modality combinations, a capability well-documented in NLP but novel in particle physics.

The work concludes that while limitations exist (particularly in high-granularity generation), the broad capabilities of these models point toward a compelling future for foundation models in high-energy physics.