A universal vision transformer for fast calorimeter simulations

TL;DR

This paper introduces a universal vision transformer model that efficiently simulates calorimeter responses with high accuracy across various detector geometries, significantly speeding up the process compared to traditional methods.

Contribution

The paper extends the CaloDREAM architecture to demonstrate a scalable, geometry-agnostic vision transformer for fast, high-fidelity calorimeter simulations.

Findings

ViTs produce shower simulations statistically indistinguishable from Geant4.

Generation time is approximately 10-100 ms on a single GPU.

Pretraining and fine-tuning improve data efficiency and fidelity.

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.