A universal vision transformer for fast calorimeter simulations

TL;DR

This work tackles the computational bottleneck of Geant4 calorimeter simulations by introducing a universal vision-transformer surrogate built on CaloDREAM, trained with Conditional Flow Matching to model calorimeter showers across both regular and irregular geometries. The ViT-based shape network, complemented by an energy network, enables fast, high-fidelity generation of electromagnetic and hadronic showers, achieving Geant4-like distributions with generation times in the $\mathcal{O}(10-100)$ ms range on a single GPU. The study demonstrates strong fidelity across datasets (regular CaloChallenge ds2/ds3 and irregular ds1, plus hadronic tests) and shows substantial data-efficiency gains through pretraining on the LEMURS dataset followed by targeted fine-tuning, as well as effective transfer-learning between detectors and voxelizations. The results support deploying transformer-based calorimeter emulators in diverse detector geometries, with public code and data releases to facilitate community adoption and further development.

Abstract

The high-dimensional complex nature of detectors makes fast calorimeter simulations a prime application for modern generative machine learning. Vision transformers (ViTs) can emulate the Geant4 response with unmatched accuracy and are not limited to regular geometries. Starting from the CaloDREAM architecture, we demonstrate the robustness and scalability of ViTs on regular and irregular geometries, and multiple detectors. Our results show that ViTs generate electromagnetic and hadronic showers statistically indistinguishable from Geant4 in multiple evaluation metrics, while maintaining the generation time in the $\mathcal{O}(10-100)$ ms on a single GPU. Furthermore, we show that pretraining on a large dataset and fine-tuning on the target geometry leads to reduced training costs and higher data efficiency, or altogether improves the fidelity of generated showers.

Figures (15)

• Figure 1: Schematic diagram of the vision transformer peebles2023scalablediffusionmodelstransformers, which highlights the detector-specific and the universal part of the architecture. The color coded detector-specific steps (see text for more details) indicate the components which may be reinitialized during fine-tuning. The universal ViT block only contains learnable transformations at patch-level objects. These weights, trained on a large corpus of data, can learn general features of calorimeter showers which are used as initializations for other detectors.
• Figure 2: Summary of the sliced evaluation for $E_\text{inc}\in [0.33\cdot E_\text{i-max}, 0.66\cdot E_\text{i-max})$ GeV, $\theta\in[1.52, 1.62)$), where $E_\text{i-max}$ is 100 GeV for the FCC detectors and 1TeV for the Par04 and ODD detectors. We show the visible energy, the energy profile in the z direction, and the energy profile in the radial direction: (top) Par04SiW, (middle) ODD, and (bottom) FCCeeALLEGRO detectors.
• Figure 3: Summary of the sliced evaluation for $E_\text{inc}\in [0.66\cdot E_\text{i-max}, E_\text{i-max})$ TeV, $\theta\in[2.1, 2.2)$), where $E_\text{i-max}$ is 100 GeV for the FCC detectors and 1TeV for the Par04 and ODD detectors. We show the visible energy, the energy profile in the z direction, and the energy profile in the radial direction: (top) Par04SiW, (middle) ODD, and (bottom) FCCeeALLEGRO detectors.
• Figure 4: Summary of the evaluation on the CaloChallenge-ds1-$\gamma$ dataset. We show the center of energy and the shower width in layer-1, the AUC scores of a low- and high-level neural classifier, and the generation time on CPU, with batch size 1, and GPU, with batch size 100.
• Figure 5: Summary of the evaluation on the CaloChallenge-ds1-$\pi^{+}$ dataset. We show the center of energy and the shower width in layer-1, the AUC scores of a low- and high-level neural classifier, and the generation time on CPU, with batch size 1, and GPU, with batch size 100.