AllShowers: One model for all calorimeter showers

TL;DR

AllShowers tackles the high computational cost of detector simulation by introducing a single, unified CNF-transformer that models electromagnetic and hadronic calorimeter showers across $12$ particle types. The two-stage architecture combines PointCountFM (layer-wise point counts) with a CNF-transformer (per-point $(x,y,e)$) conditioned on particle information and layer index, trained via conditional flow matching and enhanced by layer embeddings, fast attention masking, and layer-wise optimal transport. Empirical results show strong agreement with Geant4 across individual showers and ensemble distributions, and demonstrate competitive performance against specialized baselines like CaloClouds3 (photons) and CaloHadronic (pions), with energy calibration improving total shower energy consistency. While diffusion-based or specialized models can be faster in isolation, AllShowers offers a scalable, universal solution that reduces memory footprint and can be extended to additional detector geometries, with future work focusing on improving energy resolution and distillation for faster sampling.

Abstract

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.

Figures (10)

• Figure 1: Schematic overview of the AllShowers model architecture. The PointCountFM predicts the number of points per layer, which are then used by the CNF-transformer to generate the full shower. The incident particle information is provided to both models. The point layer index, needed by the CNF-transformer, can be computed from the number of points per layer.
• Figure 2: Schematic overview of the CNF-transformer architecture. The input $\mathbf{x}_t$ for $t=0$ is a standard normal sample, for $t=1$ it is the preprocessed shower. Since the calorimeter layer is a condition, $\mathbf{x}_t$ is a three-dimensional point-cloud $(x_k, y_k, e_k)$. The condition $c$ includes the incident particle information and the layer index. $t$ is the time variable of the neural ODE. The output $v_c(\mathbf{x}_t, t)$ is the vector field used in the CNF, e.g. the right-hand side of the neural ODE.
• Figure 3: 2D t-SNE Maaten:2008 visualization of the learned 64 dimension particle type embeddings. Small circles indicate electromagnetic showers, large circles indicate hadronic showers. One can see four distinct clusters: one for electromagnetic showers, one for positively charged hadrons, one for negatively charged hadrons, and one for neutral hadrons.
• Figure 4: Examples of attention masks for three different showers. Shown is a part of the full $6016 \times 6016$ attention matrix where each entry indicates whether two points can attend to each other. Red entries indicate allowed attention, white entries indicate masked attention. The black lines indicate the $128 \times 128$ blocks flex-attention Dong:2024 will compute. White entries within these blocks are computed but then masked out.
• Figure 5: Comparison of individual showers simulated with Geant4 and with AllShowers for different incident particles, energies and angles. The point size indicates the energy of each hit. For these showers, the number of points per layer was taken from the Geant4 simulation rather than generated by the PointCountFM to allow for a more direct comparison of the spatial and energy distribution of hits.