Calo-VQ: Vector-Quantized Two-Stage Generative Model in Calorimeter Simulation

Qibin Liu1 2$\star$, Chase Shimmin3, Xiulong Liu2, Eli Shlizerman2, Shu Li1
and Shih-Chieh Hsu2$\dagger$

1 Tsung-Dao Lee Institute, Shanghai Jiao Tong University

2 University of Washington, Seattle

3 Yale University

$\star$<EMAIL_ADDRESS>, $\dagger$<EMAIL_ADDRESS>
## Abstract

We introduce a novel machine learning method developed for the fast simulation
of calorimeter detector response, adapting vector-quantized variational
autoencoder (VQ-VAE). Our model adopts a two-stage generation strategy:
initially compressing geometry-aware calorimeter data into a discrete latent
space, followed by the application of a sequence model to learn and generate
the latent tokens. Extensive experimentation on the Calo-challenge dataset
underscores the efficiency of our approach, showcasing a remarkable
improvement in the generation speed compared with conventional method by a
factor of 2000. Remarkably, our model achieves the generation of calorimeter
showers within milliseconds. Furthermore, comprehensive quantitative
evaluations across various metrics are performed to validate physics
performance of generation.

###### Contents

  1. 1 Introduction
  2. 2 Calo-challenge Datasets
  3. 3 Two-stage Modeling Approach
  4. 4 First Stage: Compression with VQ-VAE
    1. 4.1 Data Pre-processing
    2. 4.2 Convolutional Network for Cylindrical Geometry
  5. 5 Second Stage: Latent Sampling with GPT
    1. 5.1 Training and Sampling of the Two-stage Model
  6. 6 Evaluation
    1. 6.1 Dataset 1: photon and pion
    2. 6.2 Dataset 2: electron
    3. 6.3 Dataset 3: electron
    4. 6.4 Summary Table
  7. 7 Conclusion
  8. A Hyper-parameters
  9. B Performance and Computing for Two Stages

##  1 Introduction

Simulation of detector response plays a crucial role for modern high energy
physics experiments [1, 2]. In the chain of physics event generation, from
theory model to the realistic detector output, the detector simulation stands
as the final and pivotal step for the interpretation of experimental
observation and hunting for the un-expected signature. Notably, calorimeter,
often the largest and most complex sub-system within the full detector, poses
unique challenges for the simulation due to their extensive spatial scale,
intricate structural composition, and the necessity for precise energy
deposition.

Standard way simulating the calorimeter response entails accurate tracing of
the propagation and evolution of secondary particles alongside their
interactions with detector materials and finally the recording of energy
deposited in each unit. Widely employed software like GEANT4 [3, 4, 5]
specializes in the simulation with Markov chain Monte Carlo process and is
highly optimized in the computing and precision. However, as modern
experiments continue to push the boundaries with increasingly large and high-
granularity calorimeter designs, the computational demands of accurate
simulation have emerged as a significant bottleneck [6].

The rapid development of machine learning enlightened new hope offering novel
avenues for fast simulation of calorimeter. Various methodologies have
emerged, leveraging generative neural networks based on architectures like GAN
[7, 8, 9], normalizing flow [10, 11, 12, 13], and diffusion [14, 15].

In this paper, we introduce a novel generative model based on Vector-Quantized
Variational Autoencoders (VQ-VAE) [16, 17, 18] to simulate calorimeter
response, achieving a remarkable speedup exceeding 2000 times compared to
GEANT4-based simulations on Calo-challenge datasets [19, 20, 21]. Our model
demonstrates realistic performance across physics variables[7] such as total
energy deposition and shower shape. Furthermore, we showcase the adaptability
of the VQ-VAE approach across a variety of encoder-decoder architectures,
ranging from fully-connected to convolutional networks. The former is optimal
for smaller and irregular geometries, whereas the latter excels in handling
datasets with regular structures and high granularity. Innovative designs
including FFT-resampling (Fast Fourier Transform based resampling),
cylindrical convolution, and physics-motivated normalization are developed to
enhance performance on physics metrics.

The paper is structured as following: Section 2 outlines the features of the
datasets, followed by Section 3 which summarizes the principle of two-stage
design for calorimeter simulation, based on VQ-VAE architecture. Section 4 and
5 detail the implementations of the two stages on different geometries.
Section 6 presents evaluation results encompassing both physics and
computational aspects, with concluding remarks and future outlook discussed in
Section 7.

##  2 Calo-challenge Datasets

The datasets of Fast Calorimeter Simulation Challenge 2022 project [22]are
used in this study. The datasets consists of 4 types of simulation from
different incident particle and detector setup. The different calorimeters
encompassa range of voxel counts, varying from 368 to 40,500 which reflects
various detector segmentation strategies. Specifically, one dataset captures
the response of the calorimeter from photons, another from pions, and two
others from electrons.

These datasets are generated with GEANT4, a Monte Carlo simulation program
commonly used to model the particle-material interaction in the research of
high energy physics. Dataset1-photon(referred to as "ds1_photon" hereafter)
and Dataset1-pion("ds1_pion") originate from open datasets of calorimeter from
ATLAS experiment. The showers induced by single photon or charged pion with 15
discrete incident energies ranging from 256 MeV to 4 TeV (in powers of 2) are
recorded in the dataset. The energy deposited in calorimeter is voxelized
based on the irregular detector geometry resulting in 368 voxels (5 layers)
for photon dataset and 533 (7 layers) for pion.

Dataset2("ds2") and Dataset3("ds3") capture the calorimeter response with
higher granularity. The showers stem from single electron with incident energy
logarithmically distributed from 1 GeV to 1 TeV. The detector’s cylindrical
geometry is described in three dimensions: depth ($Z$-direction where layers
arranged), angular ($A$-direction) and radical($R$-direction). Dataset2
corresponds to 45*16*9=6,480 voxels while Dataset3 comprises 45*50*18=40,500
voxels.

Dataset | Layers | Radial Seg. | Angular Seg. | Voxels | Condition Energy
---|---|---|---|---|---
ds1_photon | 5 | irregular | irregular | 368 | 15 bins
ds1_pion | 7 | irregular | irregular | 533 | 15 bins
ds2 | 45 | 16 | 9 | 6480 | continuous
ds3 | 45 | 50 | 18 | 40500 | continuous
Table 1: Details of four datasets from the Fast Calorimeter Simulation
Challenge 2022.

The configuration and statistical characteristics of each dataset are detailed
in Table 1. Each of the 4 types of datasets is given in the format of several
statistically independent and equally divided files which are naturally used
either for the training of generative model or the assessment of quality of
generated data. For all the study in the following, half of the statistics are
used for training and half for the reference of performance evaluation, of
each type of dataset.

##  3 Two-stage Modeling Approach

The calorimeter detector with high granularity typically entails large number
of channels (voxel/pixel) of the response up to tens of thousand. Direct
sampling in such high dimensionality and highly sparse space is frequently
inefficient and impractical. To address this challenge, we propose a two-stage
method, as depicted in Figure 1, drawing inspiration from advancements in the
field of computer vision [23].

Figure 1: Demonstration of the vector-quantization based two-stage generative
model. The upper and lower parts show the two stages of the model,
respectively.

The first stage is designed to reduce the dimensionality of data. The encoder
is utilized to transform the input into a different representation in latent
space followed by the decoder which reconstruct the input. In order to achieve
higher compression and more effective usage of latent space, vector
quantization technique is implemented in the latent space as in the model of
Vector Quantized Variational Auto Encoder (VQ-VAE). The same architecture had
been shown to be efficient in the jet physics for the tokenization of physics
object and benefits the classification and generation task [17, 18]. In our
study we focus on the generation of raw detector hit with higher dimensional
and low-level information. And we found that the encoder loss is not efficient
to describe the distance between input and reconstructed data in the sense of
high resolution calorimeter response. The discriminator loss is added to
improve the quality of reconstruction from VQ-VAE.

This stage enables the learning and sampling of the probabilistic distribution
in a highly reduced and regularized space. While the second stage focuses on
the modeling of latent space which is a sequence of discrete codes with fixed
length. GPT-based model is adapted to sample the latent codes conditioned on
the incident energy. As discussed in the previous section, the normalization
factor, which controls the energy deposition of each calorimeter layer, is
digitized into discrete codes and sampled at the same time of latent codes.

##  4 First Stage: Compression with VQ-VAE

The VQ-VAE model is tailored for utilization in the first stage. Following the
standard notation, the input data, energy deposited in each individual voxel
$i$, is named $E_{i}$ and after pre-processing $x_{i}$. The output of VQ-VAE
denoted as $\tilde{x_{i}}$. The index $i$ runs over all the numbers as shown
in the last column of Table 1. The global condition incident energy, is
denoted as $E_{c}$.

The first stage model includes the encoder network $q(z|x;E_{c})$ which
parameterize the posterior distribution of latent variable $z$ given the input
$x$, decoder network $p(x^{\prime}|z;E_{c})$ which runs inversely and the
prior distribution $P(z)$ which will be modelled by the second stage.

In contrast to the Variational Auto Encoder (VAE) [24], the latent
distribution $P(z)$ is not re-parameterized by independent Gaussian but
discrete index of quantized embedding vector. The embedding vector is defined
in the space $e\in R^{D}$ where $D$ is the size of latent dimension. The
output of encoder of each input is quantized to $L$ learnable vectors where
$L$ denotes the width of encoder output. In total $K$ unique vectors are
learned and they make up a "codebook", relate the index $k$ of vector to the
actual latent vector $e_{k}$. The index is also termed "code" hereafter. The
latent space is then represented as,

$z\equiv\\{e_{k(i)}\\}~{}\text{where}~{}i=1,2,...,L~{}\text{and}~{}k(i)\in[1,K]$
(1)

The back-propagation of the non-differentiable quantization operation is
implemented simply copies the gradients from the decoder to the encoder [16]
described as

$L_{VQ}=||x-p(q(x))||^{2}+||sg[q(z|x)]-e_{k}||^{2}+||q(z|x)-sg[e_{k}]||^{2}$
(2)

where $sg[X]$ denotes the stop-gradient operator which will not take the
gradient of $X$ into calculation. The first term is named "reconstruction
loss" which is designed to minimize the distance between input and output of
the whole VQ. The second term denotes the quantization loss which optimize the
codebook to better represent the latent space. The last term is called
commitment loss which limits the arbitrarily growth of embedding space and
makes the encoder output commit to a sets of vectors.

It is found in the study that the VQ-VAE model trained only with loss as
described in Formula 2 is not sufficient to obtain high fidelity due to the
limitation of L1/L2 loss. This is also observed in [23]. So an additional
discriminative regularization (loss) [25, 26] is implemented and updated
during the training as

$L_{GAN}=logD(x)+log(1-D(p(q(x))))$ (3)

where the $D$ denotes the discriminative network which is trained
simultaneously as the VQ network to capture the subtle difference between the
decoded output and input.

The total training objective of the first stage model is

$\underset{q,p,\mathcal{Z}}{arg\,min}\underset{D}{max}[L_{VQ}(q,p,\mathcal{Z})+\lambda
L_{GAN}(D)]$ (4)

where the $q,p$ denote the encoder and decoder, $\mathcal{Z}$ denotes the
trainable codebook and $\lambda$ denotes the adaptive factor [23] to balance
the updating of VQ and the discriminator.

###  4.1 Data Pre-processing

The data is pre-processed before feeding into the first stage model. The data
is firstly normalized per shower with the normalization factor denoted as $R$.
The factor is calculated per shower and needs to be sampled if new shower
being generated, which will be discussed in the next section of the second
stage 5. In the first stage the $R$ factor is assumed known information. Then
the data is transformed logarithmically with proper linear scaling and
shifting.

The pre-processed is described as formula,

$x_{i}=\frac{1}{c}\log\left(a+b\frac{E_{i}}{R_{l(i)}}\right)$ (5)

where $i$ is the index of voxel and $E_{i}$ is the original value of each
voxel. $a,b,c$ are hyper-parameter tuned according to inputs and the value
documented in Appendix A. The $R_{l(i)}$ denotes the corresponding
normalization factor for each voxel location with different normalization
scheme.

Two normalization methods are tested. Global method normalizes the whole
detector energy to 1 so that $R_{l(i)}=R=\sum_{j}{E_{j}}$ where $j$ runs over
all the voxels. Layer-wise method however performs normalization per detector
layer and $R_{l(i)}=\sum_{l(j)=l(i)}{E_{j}}$ where the summation collects all
the voxel indices in the same layer $l(i)$. The layer-wise normalization takes
the fact into consideration that the energy deposition is highly different
from layer to layer of the detector, sometime over order of magnitude in one
shower. Utilization of multiple normalization factors can achieve better
performance with the cost of more need to be learned and sampled in the the
second stage. This method basically trades speed for quality.

The input after pre-processing implies $x_{i}\simeq log(E_{i}/R_{l(i)})$ and
$\sum e^{x_{i}^{*}}=1$ where $x^{*}$ relates to $x$ with only scaling and
shifting using constants $a,b,c$. This inspires us using the softmax [27]
activation in the last layer of decoder. It is validated in experiments that
the softmax makes the training more stable and well performed than other
activation.

The post-processing defines the inverse operation and described as

$\tilde{E_{i}}=\tilde{x}_{i}\tilde{R}_{l(i)}$ (6)

where the $\tilde{x}_{i}$ denotes the decoder output and $\tilde{R}_{l(i)}$
the corresponding normalization factor which is sampled in second stage model.

###  4.2 Convolutional Network for Cylindrical Geometry

To enhance performance on particular calorimeter geometries, the encoder and
decoder model are implemented with various layers on different datasets. For
dataset1(photon) and dataset1(pion), 1D convolutional layer and fully
connected layer are combined to handle the irregular geometry. The
architecture of detaset1 model is shown in the Figure 2

Figure 2: Architecture of Stage-1 Model. Algorithm 1 Cylindrical 2D
Convolution

$x\in\mathbb{R}^{C_{in}\times Z\times A}$ $\triangleright$ Input data as 2D
array in $C_{in}$ channels

$W\in\mathbb{R}^{C_{out}\times C_{in}\times k_{z}\times k_{\alpha}}$
$\triangleright$ Network parameter for kernels

$y=\mathrm{\texttt{cyl\\_conv2d}}(x,W)$

$p_{\alpha}\leftarrow\lfloor\frac{(k_{\alpha}-1)}{2}\rfloor$ $\triangleright$
Compute padding for angular dimension

$a_{ij}\leftarrow\left(x_{ijk}\right)_{k=1}^{p_{\alpha}-1}$

$b_{ij}\leftarrow\left(x_{ijk}\right)_{k=A-p_{\alpha}+1}^{A}$

$y\leftarrow\mathrm{\texttt{concat}}([b,x,a],\mathrm{\texttt{axis}}=2)$
$\triangleright$ Pad x with cyclic boundary

$y\leftarrow\mathrm{\texttt{conv2d}}(y,W)$ $\triangleright$ Normal convolution
operation

Algorithm 2 Transposed Cylindrical 2D Convolution

$x\in\mathbb{R}^{C_{in}\times Z\times A}$ $\triangleright$ Input data as 2D
array in $C_{in}$ channels

$W\in\mathbb{R}^{C_{in}\times C_{out}\times k_{z}\times k_{\alpha}}$
$\triangleright$ Network parameter for kernels

$y=\mathrm{\texttt{cyl\\_conv\\_transpose2d}}(x,W)$

$y\leftarrow\mathrm{\texttt{conv\\_transpose2d}}(x,W)$ $\triangleright$ Normal
transpose-convolution operation

$e_{\alpha}\leftarrow\lfloor\frac{(k_{\alpha}-1)}{2}\rfloor$ $\triangleright$
Compute excess for angular dimension

$y_{ij}\leftarrow\left(y_{ijk}\right)_{k=e_{\alpha}+1}^{e_{\alpha}+A}$
$\triangleright$ Trim the excess

For dataset2 and 3 the 2D convolutional layer is adopted. The 3D calorimeter
data is interpreted as multi-channel 2D data considering that the
transnational symmetry only exists in the $Z$\- and $A$\- direction, as
described in Section 2. $R$\- direction is treated as channel dimension and
this image-like data is denoted "calo-image" in the following. It is notable
that the translation in angular direction is periodical bounded. A special
cylindrical convolution operation is defined on the $(Z,A)$ space, with
vanilla translation in one direction and circular wrapped in another. The
operation is described in Algorithm 1 and 2 where the handling of vanilla
padding in the $Z$-direction is trivial and omitted.

All the layers for the dataset2/3 model operating on $(Z,A)$ space are
implemented with cylindrical convolution. The equivariant property of the
calorimeter data is well kept when mapped in the cylindrical coordinate, as
shown in Figure 3.

Figure 3: Validation of cylindrical convolution. Equivariant property is well
kept for different transformation.

To achieve arbitrary up- or down-sampling along the angular direction with
circular or periodic boundaries, we employ FFT-resampling, as illustrated in
Figure 4. The data is firstly transformed into frequency space with discrete
FFT along the angular direction, then resampled with targeted resolution and
inversely transformed back.

Figure 4: Demonstration of FFT down-sampling.

##  5 Second Stage: Latent Sampling with GPT

In the second stage, we leverage auto-regressive model to learn and sample the
prior distribution of the discrete latent space.

GPT [28, 29], general pre-training transformer, is highly effective at
modeling the distribution of discrete token space. It is widely utilized in
natural language generation, as well as image and audio generation tasks. GPT
employs multi-layer transformer decoder blocks with a causal mask to capture
the sequence information of tokens. This makes it well-suited for the task of
the second stage in the fast calorimeter simulation context, as it can
efficiently sample a fixed length of discrete codes based on the given
conditions.

One notable advantage of the token model is the flexibility to append or co-
generate tokenized auxiliary information along with the token sequence. This
facilitates natural establishment of mutual information exchange. In our model
we explore to tokenize the continuous energy information and co-generate with
the latent codes from the first stage. One example is the shower-wise
normalization factor(s) $R$ which is mentioned in pre-processing section. The
co-generation of $R$ with the latent codes improves the physics performance as
it allows for better consideration of both the energy scale and spatial
distribution.

The details of tokenization and application for specific dataset are described
in Sec 5.1 and Sec 4.1.

The complete sequence for each sample to learn in the second stage is composed
like

$\underbrace{A_{1}A_{2}...}_{aux~{}tokens}\overbrace{B_{1}B_{2}B_{3}B_{4}....}^{shower~{}codes}$
(7)

and there is no special spacer, separator or terminator inside.

###  5.1 Training and Sampling of the Two-stage Model

The training of the two stages of calo-VQ model is performed sequentially,
taking the advantage to optimize each stage with well-defined and decomposed
tasks. The training of the first stage model as autoencoder-like architecture,
focuses on achieving high-quality reconstruction. Meanwhile the training of
the second stage model is executed in a sliding manner, aiming to maximize the
probability of predicting the next correct token based on all preceding ones.

As mentioned, auxiliary information can be sampled together with the quantized
codes in the second stage model. Here, we incorporate the $R$ factors as
auxiliary information tokenized with the digitization. Through
experimentation, we determined that utilizing 20bit (approximately 6 decimal
digits) ensures high-fidelity sampling of energy. With a codebook size of
1024, each scalar value corresponds to an additional 2 tokens in the sequence
to be learned which does not impose significant overhead. The detailed setup
such as number $R$ codes for different model is listed in Table 6

##  6 Evaluation

The fast calorimeter simulation model undergoes evaluation by comparing the
distribution of physics variables across a predefined set of events between
generated data and testing data, conditioned with same incident particle and
same energy range. The latter, produced with the same GEANT4 setup as the
training dataset but never seen during training. Physics variables, including
total energy, layer energy, shower shape variables and others are computed.
The complete list of variable definitions is outlined in Table 2. Given the
statistical nature of these physics variables, their distributions are
histogrammed to mitigate fluctuation effects. The disparity between two
histograms is quantified using the "Separation Metric" [30],

$Sep.\equiv\sum_{i=1..Nbin}\frac{(h^{\prime}_{i}-h_{i})^{2}}{2(h^{\prime}_{i}+h_{i})}$
(8)

$h^{\prime}_{i}$ and $h_{i}$ denote the i-th bin content in the histogram of
sampled and truth distribution, respectively.

Symbol | Name | Definition
---|---|---
$E_{i}$ | Cell energy for index $i$ | (raw data)
$E_{inc}$ | Incident energy | (condition)
$\eta_{i}(K)$/$\phi_{i}(K)$ | Location(mm) in $\eta$/$\phi$ at cell $i$ in layer $K$ | (geometry)
$\Delta\eta_{i}(K)$/$\Delta\phi_{i}(K)$ | Segment.(mm) in $\eta$/$\phi$ of cell $i$ in layer $K$ | (geometry)
$E_{tot}$ | Total energy | $\sum E_{i}$
$E_{tot}/E_{inc}$ | Total energy response | $\sum E_{i}/E_{inc}$
$E(K)$ | Layer energy (in layer $K$) | $\sum_{i\in K}E_{i}$
$\overline{\eta(K)}$ | Energy Center (EC) in $\eta$ (in layer $K$) | $\sum_{i\in K}\eta_{i}E_{i}/E(K)$
$\sigma_{\eta(K)}$ | Shower Width in $\eta$ (in layer $K$) | $\sum_{i\in K}\Delta\eta_{i}E_{i}/E(K)$
$\overline{\phi(K)}$ | Energy Center (EC) in $\phi$ (in layer $K$) | $\sum_{i\in K}\phi_{i}E_{i}/E(K)$
$\sigma_{\phi(K)}$ | Shower Width in $\phi$ (in layer $K$) | $\sum_{i\in K}\Delta\phi_{i}E_{i}/E(K)$
Table 2: Symbol and definition of physics variables.

###  6.1 Dataset 1: photon and pion

Two types of incident particles are studied in Dataset 1: photons and pions.
Models are trained and evaluated on each dataset separately. Figure 5
illustrates the averaged energy distribution in each calorimeter cell across
all sampled events for pion particles. Figure 6 and 7 display the distribution
of individual cell energies, total energy responses, and shower variables for
selected layers. The metrics of each layer are summarized in in Table 3, 4.
Generally, good agreement is observed across different incident energy
conditions, with metrics for energy response better than 0.01.

Figure 5: Average energy deposition of calorimeter cells. Each shower is
induced by single incident of pion. Generated on the top and reference on the
bottom.

Figure 6: Distribution of cell energy and total energy response of photon(top)
and pion(bottom) shower. Generated is shown in orange line and reference from
GEANT4 shown in solid blue.

Figure 7: Distribution of physics variables from calorimeter shower for pion incident. Layer 13 is selected for demonstration. Generated is shown in orange line and reference from GEANT4 shown in solid blue. From the top left to bottom right, the physics variables are layer energy, shower center and width in $\eta$ direction, shower center and width in $\phi$ direction.  Sep. for Layer $K$ | $E_{K}$ | $\overline{\eta(K)}$ | $\sigma_{\eta(K)}$ | $\overline{\phi(K)}$ | $\sigma_{\phi(K)}$
---|---|---|---|---|---
0 | 0.004 | - | - | - | -
1 | 0.001 | 0.018 | 0.017 | 0.015 | 0.023
2 | 0.001 | 0.007 | 0.012 | 0.005 | 0.012
3 | 0.001 | - | - | - | -
12 | 0.002 | - | - | - | -
Table 3: Separation metrics for different layer of photon dataset. For layer 0, 3 and 12 there is no segmentation along $\eta$ and $\phi$ direction so no definition of some shower variables.  Sep. for Layer $K$ | $E_{K}$ | $\overline{\eta(K)}$ | $\sigma_{\eta(K)}$ | $\overline{\phi(K)}$ | $\sigma_{\phi(K)}$
---|---|---|---|---|---
0 | 0.007 | - | - | - | -
1 | 0.006 | 0.015 | 0.037 | 0.027 | 0.019
2 | 0.003 | 0.015 | 0.018 | 0.015 | 0.016
3 | 0.005 | - | - | - | -
12 | 0.002 | 0.012 | 0.011 | 0.014 | 0.020
13 | 0.001 | 0.013 | 0.006 | 0.009 | 0.008
14 | 0.001 | - | - | - | -
Table 4: Separation metrics for different layer of pion dataset. For layer 0,
3 and 14 there is no segmentation along $\eta$ and $\phi$ direction so no
definition of some shower variables.

###  6.2 Dataset 2: electron

The averaged image of all electron showers is depicted in Figure 8 for a
selected calorimeter layer.

Distributions and metrics are summarized in Figures 9, 10 and 11. Despite the
significantly increased granularity compared to Dataset 1, the model
demonstrates strong performance in terms of modeling physics variables, with
most metrics no worse than 0.01.

Figure 8: Average energy deposition of calorimeter cells in layer 20 of
dataset2. Each shower is induced by single incident of electron. Generated on
the right and reference on the left.

Figure 9: Distribution of cell energy and total energy response of electron
shower in dataset2. Generated is shown in orange line and reference from
GEANT4 shown in solid blue.

Figure 10: Distribution of physics variables from calorimeter shower for
electron incident in dataset2. Layer 20 is selected for demonstration.
Generated is shown in orange line and reference from GEANT4 shown in solid
blue. From the top left to bottom right, the physics variables are layer
energy, shower center and width in $\eta$ direction, shower center and width
in $\phi$ direction. Figure 11: Separation metric of different physics
variables of each layer. The lower separation means the smaller difference
between fast simulation and GEANT4 simulation of calorimeter response. The
definitions of variable follow those listed in Table 2.

###  6.3 Dataset 3: electron

The images, distributions, and metrics are summarized in Figures 12, 13, 14,
15 and 16. The architecture of the two-stage model can be adapted to different
auto-encoder backbones optimized for different geometries. This flexibility
allows for handling and maintaining performance with up to 40,500 channels,
employing a similar architecture as designed for Dataset 2 just with simple
scaling. This underscores the robustness and scalability of the design.

Two models are presented here, one with and one without layer-wise
normalization as mentioned in Section 4.1. It is observed that layer-wise
normalization improves performance, particularly in the distribution of energy
sums for each layer. However, this improvement comes at the cost of increased
latent size, resulting in slower generation.

Figure 12: Average energy deposition of calorimeter cells in layer 20 of
dataset3 comparing the reference and generated. Each shower is induced by
single incident of electron. The result from vanilla model is shown on the
left and layer-wise normalized model on the right.

Figure 13: Distribution of cell energy and total energy response of electron
shower in dataset3. Generated is shown in orange line and reference from
GEANT4 shown in solid blue. The vanilla model is shown on the top and layer-
wise normalized model on the bottom.

Figure 14: Distribution of physics variables from calorimeter shower for
electron incident in dataset3. Layer 20 is selected for demonstration.
Generated is shown in orange line and reference from GEANT4 shown in solid
blue. From the top left to bottom right, the physics variables are layer
energy, shower center and width in $\eta$ direction, shower center and width
in $\phi$ direction. he vanilla model is shown on the top and layer-wise
normalized model on the bottom. Figure 15: Separation metric of different
physics variables of each layer. The lower separation means the smaller
difference between fast simulation and GEANT4 simulation of calorimeter
response. The definitions of variable follow those listed in Table 2. Figure
16: Separation metric of different physics variables of each layer. The lower
separation means the smaller difference between fast simulation and GEANT4
simulation of calorimeter response. The definitions of variable follow those
listed in Table 2.

###  6.4 Summary Table

The number of parameters and generation time per shower for each model, along
with the worst and best metric across all layers, are summarized in Table 5.
The evaluation was conducted using a single NVIDIA V100 GPU with 512 showers
per batch. For comparison, the generation time per shower for the standard
method, based on GEANT4, varies widely depending on the energy, taking
approximately $O(1)-O(100)$ seconds for Dataset1 [9] and around $O(100)$
seconds for Dataset 2 and 3 on average [31]. The proposed fast simulation
method achieves a speedup of over 2000 times. The performance is consistently
maintained across datasets with varying granularity, scaling well up to 40,500
dimensions.

Furthermore, the study reveals distinct scaling behaviors between the two
stages of the model. The first stage model exhibits nearly constant scaling
with different dataset complexities, demonstrating the scalability of deep
convolutional networks for compression tasks. Conversely, the second stage
model is only sensitive to the size of the latent space and demonstrates good
generalization across different datasets. The generation quality is primarily
limited by the first stage model and it impacts the computing costs solely
through the latent size. Details of the performance and computing cost for
both the first and second models are tabulated in Appendix B

Dataset(Model) | Channel (D) | Latent Size (L) | N#Pars/M | Total time/ms | Metric(worst/best)
---|---|---|---|---|---
ds1(photon) | 368 | 42 | 4.06 | 0.25 | 0.023/0.001
ds1(pion) | 533 | 46 | 4.31 | 0.28 | 0.037/0.001
ds2 | 6480 | 70 | 3.32 | 0.63 | 0.057/0.001
ds3(vanilla) | 40500 | 184 | 2.16 | 1.14 | 0.133/0.002
ds3(LN) | 40500 | 274 | 2.77 | 36.10 | 0.095/0.001
Table 5: Number of parameters and generation time per shower of each model for
different datasets. "LN" denotes the layer-wise normalized model.

##  7 Conclusion

The paper introduces a novel two-stage vector-quantization based generative
model for calorimeter simulation, presenting its technical framework, physics
performance, and computational considerations.

Significantly, our model achieves a speedup exceeding 2000 times compared to
conventional calorimeter simulation methods, while maintaining discrepancies,
measured by separation metric, in total energy and energy per layer at 0.01
level and shower shape at 0.1 level. Although our model’s performance does not
surpass existing generative models based on normalizing flow or diffusion, it
introduces a new dimension to the accuracy versus time trade-off, thus
expanding the landscape of possibilities in this domain.

The inherent flexibility of the architecture, coupled with the incorporation
of multiple designs inspired by physics, promises to enrich future
investigations in fast calorimeter simulation. Further optimization options
are available, leveraging the decomposed and adaptable nature of the two-stage
model. Particularly promising is the exploration of latent diffusion
architecture in the second stage, offering potential for enhanced scalability
and expressive design of the latent space as the simulation demands for high
granularity calorimeter increase.

In conclusion, our work serves as a stepping stone towards achieving a
comprehensive solution for fast and accurate calorimeter simulation, offering
insights and directions for future research in this critical area of high-
energy physics.

##  Appendix A Hyper-parameters

The key hyperparameters of the models described in the paper are listed in
Table 6

| ds1-photon | ds1-pion | ds2 | ds3 | ds3(normed)
---|---|---|---|---|---
log_transform a | 1 | 1 | 1 | 1 | 1
log_transform b | 8000 | 8000 | 3000 | 40000 | 40000
log_transform c | 10 | 10 | 7 | 10 | 10
Hidden layers | 5 | 5 | 7 | 10 | 10
VQ dim | 256 | 256 | 256 | 192 | 256
cond. dim | 3 | 3 | 1 | 1 | 1
codebook size | 1024 | 1024 | 1024 | 1024 | 1024
R codes/token | 10 | 14 | 2 | 2 | 90
shower codes | 32 | 32 | 68 | 182 | 624
stage1 #pars/M | 3.8 | 4.1 | 3.1 | 2.1 | 2.2
xformer layer | 2 | 2 | 2 | 1 | 1
xformer head | 2 | 2 | 2 | 1 | 1
xformer embed | 64 | 64 | 64 | 16 | 128
stage2 #pars/K | 231 | 231 | 235 | 38 | 551
Table 6: Main hyper-parameters and number of trainable parameters of pre-
processing, first stage model and second stage model of caloVQ. For
"hidden_layers" the numbers are halved since the encoder and decoder are
symmetry.

##  Appendix B Performance and Computing for Two Stages

The performance and computing cost for two stages are listed in Table 7

Dataset | N#Pars(S1)/M | S1 time/ms | N#Pars(S2)/M | S2 time/ms
---|---|---|---|---
ds1-photon | 3.83 | 0.02 | 0.23 | 0.23
ds1-pion | 4.08 | 0.02 | 0.23 | 0.26
ds2 | 3.08 | 0.17 | 0.23 | 0.46
ds3 | 2.12 | 0.35 | 0.04 | 0.79
ds3(normed) | 2.22 | 1.70 | 0.55 | 34.40
Table 7: Number of parameters and generation time per shower of each stage for
differrent model.

## References

  * [1] H. S. Foundation, :, J. Apostolakis, M. Asai, S. Banerjee, R. Bianchi, P. Canal, R. Cenci, J. Chapman, G. Corti, G. Cosmo, S. Easo _et al._ ,  _Hep software foundation community white paper working group - detector simulation_ (2018), 1803.04165.
  * [2] G. A. Stewart, A. Valassi, J. McFayden, E. Yazgan, W. Pokorski, G. Corti, P. Canal, C. Doglioni, A. Dziurda, D. Lange, P. J. Laycock, T. J. Khoo _et al._ ,  _Hl-lhc computing review: Common tools and community software_ ,  Tech. rep.,  10.5281/ZENODO.4009114 (2020).
  * [3] S. Agostinelli, J. Allison, K. Amako, J. Apostolakis, H. Araujo, P. Arce, M. Asai, D. Axen, S. Banerjee, G. Barrand, F. Behner, L. Bellagamba _et al._ ,  _Geant4—a simulation toolkit_ ,  Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment 506(3), 250 (2003),  https://doi.org/10.1016/S0168-9002(03)01368-8.
  * [4] J. Allison, K. Amako, J. Apostolakis, P. Arce, M. Asai, T. Aso, E. Bagli, A. Bagulya, S. Banerjee, G. Barrand, B. Beck, A. Bogdanov _et al._ ,  _Recent developments in geant4_ ,  Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment 835, 186 (2016),  https://doi.org/10.1016/j.nima.2016.06.125.
  * [5] J. Allison, K. Amako, J. Apostolakis, H. Araujo, P. Arce Dubois, M. Asai, G. Barrand, R. Capra, S. Chauvie, R. Chytracek, G. Cirrone, G. Cooperman _et al._ ,  _Geant4 developments and applications_ ,  IEEE Transactions on Nuclear Science 53(1), 270 (2006),  10.1109/TNS.2006.869826.
  * [6] _ATLAS HL-LHC Computing Conceptual Design Report_ ,  Tech. rep., CERN, Geneva (2020).
  * [7] M. Paganini, L. de Oliveira and B. Nachman,  _Calogan: Simulating 3d high energy particle showers in multilayer electromagnetic calorimeters with generative adversarial networks_ ,  Physical Review D 97(1) (2018),  10.1103/physrevd.97.014021.
  * [8] E. Buhmann, S. Diefenbacher, E. Eren, F. Gaede, G. Kasieczka, A. Korol and K. Krüger,  _Getting high: High fidelity simulation of high granularity calorimeters with high speed_ ,  Computing and Software for Big Science 5(1) (2021),  10.1007/s41781-021-00056-0.
  * [9] _Fast simulation of the ATLAS calorimeter system with Generative Adversarial Networks_ ,  Tech. rep., CERN, Geneva,  All figures including auxiliary figures are available at https://atlas.web.cern.ch/Atlas/GROUPS/PHYSICS/PUBNOTES/ATL-SOFT-PUB-2020-006 (2020).
  * [10] C. Krause and D. Shih,  _Fast and accurate simulations of calorimeter showers with normalizing flows_ ,  Physical Review D 107(11) (2023),  10.1103/physrevd.107.113003.
  * [11] C. Krause and D. Shih,  _Caloflow ii: Even faster and still accurate generation of calorimeter showers with normalizing flows_ (2023), 2110.11377.
  * [12] M. R. Buckley, C. Krause, I. Pang and D. Shih,  _Inductive caloflow_ (2023), 2305.11934.
  * [13] F. Ernst, L. Favaro, C. Krause, T. Plehn and D. Shih,  _Normalizing flows for high-dimensional detector simulations_ (2023), 2312.09290.
  * [14] V. Mikuni and B. Nachman,  _Score-based generative models for calorimeter shower simulation_ ,  Physical Review D 106(9) (2022),  10.1103/physrevd.106.092009.
  * [15] O. Amram and K. Pedro,  _Denoising diffusion models with geometry adaptation for high fidelity calorimeter simulation_ (2023), 2308.03876.
  * [16] A. Van Den Oord, O. Vinyals _et al._ ,  _Neural discrete representation learning_ ,  Advances in neural information processing systems 30 (2017).
  * [17] J. Birk, A. Hallin and G. Kasieczka,  _Omnijet-alpha: The first cross-task foundation model for particle physics_ (2024), 2403.05618.
  * [18] L. Heinrich, T. Golling, M. Kagan, S. Klein, M. Leigh, M. Osadchy and J. A. Raine,  _Masked particle modeling on sets: Towards self-supervised high energy physics foundation models_ (2024), 2401.13537.
  * [19] M. F. Giannelli, G. Kasieczka, C. Krause, B. Nachman, D. Salamani, D. Shih and A. Zaborowska,  _Fast Calorimeter Simulation Challenge 2022 - Dataset 1_ ,  10.5281/zenodo.8099322 (2023).
  * [20] M. Faucci Giannelli, G. Kasieczka, C. Krause, B. Nachman, D. Salamani, D. Shih and A. Zaborowska,  _Fast Calorimeter Simulation Challenge 2022 - Dataset 2_ ,  10.5281/zenodo.6366271 (2022).
  * [21] M. Faucci Giannelli, G. Kasieczka, C. Krause, B. Nachman, D. Salamani, D. Shih and A. Zaborowska,  _Fast Calorimeter Simulation Challenge 2022 - Dataset 3_ ,  10.5281/zenodo.6366324 (2022).
  * [22] C. K. B. N. D. S. D. S. Michele Faucci Giannelli, Gregor Kasieczka and A. Zaborowska,  _Fast calorimeter simulation challenge 2022_ (2022).
  * [23] P. Esser, R. Rombach and B. Ommer,  _Taming transformers for high-resolution image synthesis_ ,  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 12873–12883 (2021).
  * [24] D. P. Kingma and M. Welling,  _Auto-encoding variational bayes_ (2022), 1312.6114.
  * [25] A. Lamb, V. Dumoulin and A. Courville,  _Discriminative regularization for generative models_ (2016), 1602.03220.
  * [26] F. Mentzer, G. Toderici, M. Tschannen and E. Agustsson,  _High-fidelity generative image compression_ (2020), 2006.09965.
  * [27] J. S. Bridle,  _Probabilistic interpretation of feedforward classification network outputs, with relationships to statistical pattern recognition_ ,  In _Neurocomputing: Algorithms, architectures and applications_ , pp. 227–236. Springer (1990).
  * [28] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei and I. Sutskever,  _Language models are unsupervised multitask learners_ (2019).
  * [29] A. Karpathy,  _mingpt: A pytorch re-implementation of gpt_ ,  https://github.com/karpathy/minGPT (2013).
  * [30] S. Diefenbacher, E. Eren, G. Kasieczka, A. Korol, B. Nachman and D. Shih,  _Dctrgan: improving the precision of generative models with reweighting_ ,  Journal of Instrumentation 15(11), P11004–P11004 (2020),  10.1088/1748-0221/15/11/p11004.
  * [31] M. R. Buckley, I. Pang, D. Shih and C. Krause,  _Inductive simulation of calorimeter showers with normalizing flows_ ,  Physical Review D 109(3) (2024),  10.1103/physrevd.109.033006.