Generalizable Foundation Models for Calorimetry via Mixtures-of-Experts and Parameter Efficient Fine Tuning
Authors: Carlos Cardona-Giraldo, Cristiano Fanelli, James Giroux, Cole Granger, Benjamin Nachman, Gerald Sabin
Organizations: RNET Technologies Inc., Dayton, OH 45429 · William & Mary, Department of Data Science, Williamsburg, VA 23185, USA · Fundamental Physics Directorate, SLAC National Accelerator Laboratory, Menlo Park, CA 94025, USA · Department of Particle Physics and Astrophysics, Stanford University, Stanford, CA 94305, USA
Modern particle physics experiments face an increasing demand for high-fidelity detector simulation as luminosities rise and computational requirements approach the limits of available resources. Deep generative models have emerged as promising surrogates for traditional Monte Carlo simulation, with recent advances drawing inspiration from large language models (LLM) and next-token prediction paradigms. In this work, we introduce a generalizable foundation model for calorimetry built on next-token transformer backbones, designed to support modular adaptation across materials, particle species, and detector configurations. Our approach combines Mixture-of-Experts pre-training with parameter-efficient fine-tuning strategies to enable controlled, additive model expansion without catastrophic forgetting. A pre-trained backbone is trained to generate electromagnetic showers across multiple absorber materials, while new materials are incorporated through the addition and tuning of lightweight expert modules. Extensions to new particle types are achieved via parameter-efficient fine-tuning and modular vocabularies, preserving the integrity of the base model. This design enables efficient, incremental knowledge integration as new simulation datasets become available, a critical requirement in realistic detector-development workflows. In addition, we demonstrate that next-token calorimeter models are computationally competitive with standard generative approaches under established LLM optimization procedures. These results establish next-token architectures as a viable path toward extensible, physics-aware foundation models for calorimetry and future high-energy physics experiments.
Figures & tables
Figure 1: Overview of the modular architecture for material and particle-species adaptation : The framework utilizes a core transformer backbone consisting of cross-attention and self-attention decoder blocks that remain frozen during secondary adaptation phases. Material extensibility is achieved through a Mixture-of-Experts (MoE) layer where a router directs inputs to specialized modules, allowing for the addition of new materials by fine-tuning only a singular new expert. When transitioning to different particle species, the model employs a parameter-efficient strategy using LoRA modules and expanded particle-specific vocabulary heads for pixel and energy prediction while the base photon model parameters remain static. For subsequent material expansion of an adapted model, a new expert is integrated while the previously tuned LoRA and vocabulary components are frozen to preserve the learned particle-specific features. The system is conditioned throughout these stages by a combination of spatial, kinematic, and energy query embeddings alongside a unique particle identifier.
Figure 2: Validation of generative shower modeling for photons in tungsten: Distribution level comparison between ground truth Geant4 reference sample (gray shaded), against the proposed method (blue) and the Omnijet- αc baseline (green) for photons in tungsten. The top row depicts observables including visible cell energy (a.u.), total energy sum, and total number of hits. The bottom row depicts spatial shower profiles including the longitudinal center of gravity ( Z ), energy deposition per layer, and radial energy distribution. The lower panels show the ratio of each model to the Geant4 reference, where a ratio of 1.0 (dashed line) indicates perfect agreement. A 3σ uncertainty profile is provided in the ratio plots, accounting for statistical error in both the generated, and ground truth sample.
Figure 3: Validation of generative shower modeling for photons in tantalum: Distribution level comparison between ground truth Geant4 reference sample (gray shaded) and the proposed method (blue) for photons in tantalum.
Figure 4: Validation of energy kinematic conditioning through prepended context at 50 GeV : Distribution level comparison between ground truth Geant4 reference sample (gray shaded) and the proposed method (blue) at a fixed initial energy of 50GeV for photons in tungsten.
Figure 5: Fine-tuning efficiency for photon showers in lead across varying sample sizes: Comparison between the ground truth Geant4 reference (black) and the fine-tuned model at 1k (blue), 10k (green), and full (red dashed) sample sizes. Lower panels display the ratio to the Geant4 reference, where a ratio of 1.0 (dashed line) indicates perfect agreement. Shaded areas in the ratio plots depict the statistical uncertainty from data samples in quadrature with the uncertainty from bootstrapped fine-tunings, represented at the 3σ level. Results demonstrate that freezing the base tantalum/tungsten MoE backbone and updating only the lead-specific expert allows for high-fidelity adaptation even with limited training data.
Figure 6: Transfer learning efficiency for electron showers in tungsten across varying sample sizes: Comparison between the ground truth Geant4 reference (grey shaded) and the fine-tuned model at 10k (blue), 50k (orange), and full (green) sample sizes. Lower panels display the ratio to the Geant4 reference, where a ratio of 1.0 (dashed line) indicates perfect agreement. Shaded areas in the ratio plots depict the statistical uncertainty from data samples in quadrature with the uncertainty from bootstrapped fine-tunings, represented at the 3σ level.
Figure 7: Transfer learning efficiency for electron showers across varying sample sizes. Comparison between the ground truth Geant4 reference (grey shaded) and the fine-tuned model at 1k (blue), 10k (orange), and full (green) sample sizes for (a) tantalum and (b) lead.
Model
Inference Time (ms)
Hardware
Speedup Factor
Geant4 [ 4 , 5 , 3 ]
4100
CPU
N/A
OmniJet- αc [ 12 ]
3000
A100
1.39
CaloClouds II [ 17 ]
6.12
A100
670
L2LFlows [ 23 ]
3.24
A100
1260
Ours
10.46
A100
392
Table 1: Inference Speed Comparison : Comparison of inference speeds of various generative modeling approaches, and their relative speedup to Geant4 . All generative modeling approaches are tested with identical hardware (Nvidia-A100). Inference time values for Geant4 . Omnijet- αc , CaloCloudsII and L2LFlows are directly taken from [ 12 ] .
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Validation of energy kinematic conditioning through prepended context at 20 GeV : Distribution level comparison between ground truth Geant4 reference sample (gray shaded) and the proposed method (blue) at a fixed initial energy of 20GeV for photons in tungsten.
Figure 9: Validation of energy kinematic conditioning through prepended context at 80 GeV : Distribution level comparison between ground truth Geant4 reference sample (gray shaded) and the proposed method (blue) at a fixed initial energy of 80GeV for photons in tungsten.
Figure 10: Effect of calibration step on generated electrons in tantalum. Comparison between the ground truth Geant4 reference (grey shaded), the calibrated generations (blue), and the uncalibrated generations (green).
Figure 11: Effect of calibration step on generated electrons in lead. Comparison between the ground truth Geant4 reference (grey shaded), the calibrated generations (blue), and the uncalibrated generations (green).
William & Mary, Department of Data Science, Williamsburg, VA 23185, USA · University of Glasgow, School of Physics and Astronomy, Glasgow G12 8QQ,14 UK · Thomas Jefferson National Accelerator Facility, Newport News, VA 23606, USA +1
School of Physics and Astronomy, University of Edinburgh, Edinburgh EH9 3FD, United Kingdom · Department of Physics, University of Wisconsin-Madison, Madison, WI 53706, USA · Fermi National Accelerator Laboratory, Batavia, IL 60510, USA +2