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).
Accurate and efficient detector simulation is essential for modern collider experiments. To reduce the high computational cost, various fast machine learning surrogate models have been proposed. Traditional surrogate models for calorimeter shower modeling train separate networks for each particle species, limiting scalability and reuse. We introduce AllShowers, a unified generative model that simulates calorimeter showers across multiple particle types using a single generative model. AllShowers is a continuous normalizing flow model with a Transformer architecture, enabling it to generate complex spatial and energy correlations in variable-length point cloud representations of showers. Trained on a diverse dataset of simulated showers in the highly granular ILD detector, the model demonstrates the ability to generate realistic showers for electrons, photons, and charged and neutral hadrons across a wide range of incident energies and angles without retraining. In addition to unifying shower generation for multiple particle types, AllShowers surpasses the fidelity of previous single-particle-type models for hadronic showers. Key innovations include the use of a layer embedding, allowing the model to learn all relevant calorimeter layer properties; a custom attention masking scheme to reduce computational demands and introduce a helpful inductive bias; and a shower- and layer-wise optimal transport mapping to improve training convergence and sample quality. AllShowers marks a significant step towards a universal model for calorimeter shower simulations in collider experiments.
Thorsten Buss, Henry Day-Hall, Frank Gaede +2
Institut für Experimentalphysik, Universität Hamburg, Hamburg, Germany · Deutsches Elektronen-Synchrotron DESY, Hamburg, Germany
Modern particles physics experiments have demonstrated an increasing need for fast, high-fidelity detector simulation as detector components have improved and subsequent computational requirements approach the limits of available resources. Recently, deep generative models have emerged as a promising alternative to traditional Monte-Carlo methods, with recent works drawing inspiration from large language models (LLMs) and self-supervised next-token prediction methods. In this work, we present an application of a GPT-style autoregressive transformer as a fast surrogate model for the calorimeter inside the CLAS12 experiment at the Thomas Jefferson National Accelerator Facility. The model is conditioned on incident momentum and generates realistic detector hits autoregressively across all nine calorimeter layers as sequences of strip, ADC, and TDC tokens. We demonstrate that the model faithfully reproduces hit multiplicity, spatial distributions, energy deposits, and the energy-momentum response of the electromagnetic calorimeter. The generator achieves inference rates exceeding 700 events per second on a single GPU, providing a substantial speedup over traditional Geant4-based simulations while maintaining physics fidelity essential for high-luminosity experimental programs.
Cole Granger, James Giroux, Richard Tyson +2
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
High-precision calorimeter simulation at current and future colliders imposes rapidly growing computational demands, motivating the development of machine-learning surrogates for traditional Monte Carlo tools such as Geant4. Flow matching and diffusion-based generative models have become leading approaches for high-dimensional fast simulation because of their sample quality, but typically require O(100) function evaluations at inference and often rely on auxiliary networks to constrain global observables, compromising streamlined end-to-end generation. We introduce a unified framework that improves the balance between speed, shower quality, and physics fidelity. The method combines: (i) an average velocity field integrator that enables sampling in one or a few evaluations; (ii) a learned generative prior in shower space, constructed from data rather than random noise; and (iii) physics-guided loss terms that impose inductive biases on key observables during training. These elements are training time regularizers, preserving end-to-end inference with no additional cost. With only one or a few evaluation steps, the model achieves shower quality competitive with state-of-the-art flow and diffusion approaches, tested on several public high granularity calorimeter datasets. The results demonstrate inter-layer shower structure consistent with the underlying physics, providing a strong candidate for future fast simulation workflows.
Cheng Jiang, Sitian Qian, Kevin Pedro +3
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