Choose Your Diffusion: Efficient and flexible ways to accelerate the diffusion
model in fast high energy physics simulation

Cheng Jiang1$\star$, Sitian Qian2$\dagger$ and Huilin Qu3$\ddagger$

1 School of Physics and Astronomy, University of Edinburgh, EH9 3FD,
Edinburgh, United Kingdom

2 School of Physics, Peking University, 100871, Beijing, China

3 CERN, EP Department, CH-1121 Geneva 23, Switzerland

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

The diffusion model has demonstrated promising results in image generation,
recently becoming mainstream and representing a notable advancement for many
generative modeling tasks. Prior applications of the diffusion model for both
fast event and detector simulation in high energy physics have shown
exceptional performance, providing a viable solution to generate sufficient
statistics within a constrained computational budget in preparation for the
High Luminosity LHC. However, many of these applications suffer from slow
generation with large sampling steps and face challenges in finding the
optimal balance between sample quality and speed. The study focuses on the
latest benchmark developments in efficient ODE/SDE-based samplers, schedulers,
and fast convergence training techniques. We test on the public CaloChallenge
and JetNet datasets with the designs implemented on the existing architecture,
the performance of the generated classes surpass previous models, achieving
significant speedup via various evaluation metrics.

Copyright attribution to authors.
This work is a submission to SciPost Physics.
License information to appear upon publication.
Publication information to appear upon publication. Received Date
Accepted Date
Published Date

###### Contents

  1. 1 Introduction
  2. 2 Dynamics in the Diffusion Model
  3. 3 Dataset
    1. 3.1 CaloChallenge
    2. 3.2 JetNet
  4. 4 Training
    1. 4.1 CaloChallenge
    2. 4.2 JetNet
  5. 5 Performance
    1. 5.1 CaloChallenge
    2. 5.2 JetNet
  6. 6 Conclusion and outlooks
  7. A Hyperparameters of Samplers and Schedulers

##  1 Introduction

The Large Hadron Collider (LHC) [1] is one of the biggest particle accelerator
human have ever made. More than billions of events are generated and recorded
by experiments like ATLAS [2] and CMS [3] near the interaction points during
each run. Sufficient and precise simulation for comparing the data and
theoretical prediction is required to find the evidence of new physics at the
energy frontier.

The `Geant4` software toolkit [4] is applied for the detector simulation. The
simulated primary particles would propagate through specific material and
geometry, resulting vast amount of secondary particles produced in each
propagated steps to achieve precision. As a consequence, the detector
simulations for the modern high granularity detectors occupy the most
computation resources. Reconstructing the physics objects like jets with
complicated substructure for different physics process also plays a crucial
role in providing the enough statistics for the data-driven analysis. In the
upcoming High Luminosity LHC [5], the expected integrated luminosity will
scale up to ten times the current level. It is nearly impossible to perform
the full simulation for all the events with constrained computational budget,
manifesting the necessity of fast simulation.

To address with these challenges comprehensively, different techniques adopted
from deep generative modeling, where a lot of phenomenal improvements have
quickly emerged over the years, have become the popular approach as fast
surrogate models in high energy physics.

One latest category of generative models, the Diffusion Model (DM) as well as
its variations, has become increasingly dominant in various generative
applications [6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]. There has been
a few studies using DM in both fast detector and end-to-end event simulations.

For fast detector simulation, `CaloScore` [19] utilizes the score-based
generative model and the conditional UNet to produce almost indistinguishable
generated showers from truth ones. `CaloDiffusion` [20], on the other hand,
utilizes the denoising diffusion probabilistic model, incorporates novel
geometry adaptation, and enhances the UNet with middle self-attention and
cylindrical convolutional layers. This combination positions `CaloDiffusion`
as one of the top models in fast detector simulation. Moreover, `CaloClouds`
[21] offers an alternative way to generate geometry-independent point clouds
for great modelling of physically relevant distributions.

For end-to-end event simulation, there are a series of studies for generating
the point cloud with DM [22, 23]. Notably, those studies also use
architectures like transformer [22] or `EPiC` layer [23] that utilize the
permutation invariance of jet constituents.

Among those original studies, their first versions often suffer from a long
generation time, as desired results typically converge only after a large
number of sampling steps. Further studies involving the fresh distillation
methods like progressive distillation [24, 25], consistency model [20, 26, 27]
to get the result in only few steps with the exchange of longer training time,
and other ways [26, 28] to speedup the backward process.

Most of the time, different DM methods face a tradeoff between the quality and
speed of generated classes. There is a lack of a comprehensive study on the
training and sampling dynamics of diffusion models applied for various
purposes and datasets in high energy physics. In this study, we investigate
the effects of different training strategy by offering a soft way to expedite
the convergence of performance. We also adopt several popular training-free
ODE/SDE based samplers and noise schedulers in post-DDPM era [29, 30, 31, 32,
33]. By balancing the discretization errors and accumulated errors in the
sampling process, a flexible and efficient way can be offered to make the
diffusion model faster and more accurate.

We finally make our approaches portable by testing in different network and
different datasets using the public fast and high-fidelity calorimeter dataset
CaloChallenge [34] and the jet point cloud dataset JetNet [35]. The
performance surpasses the previous model by several evaluation metrics.

This paper is outlined as follows, Section 2 covers the fundamentals of
diffusion model, elucidating how data such as shower images or point clouds
are specifically generated through the backward process. Section 3 provides
the a brief introduction about the dataset used in this study. Section 4
discusses the training strategy applied. Various metrics will be presented to
evaluate the performance in Section 5. Lastly, Section 6 gives a empirical
discussion, conclusion and outlooks about the study.

##  2 Dynamics in the Diffusion Model

The diffusion model was initially introduced to sample complex data
distributions by drawing inspiration from concepts in non-equilibrium
thermodynamics [36]. At first, we gradually add noises to an initial data
sample from $x_{0}$ to $x_{T}$ until the time step $T$, usually denoted as the
forward process $q$ where the data points are degraded as $T$ increase. Then
it follows a backward process $p$ where noisy data could be denoised
iteratively to new observation $\sim x_{0}$ as the time goes reversely.

Two main classes are introduced later[37, 38]: Denoising Diffusion
Probabilistic Models (DDPMs) [39, 40]and Score-based Generative Models (SGMs)
[41].

For DDPM, the forward process transforms input distribution into a prior
distribution by

$\displaystyle
q(x_{t}|x_{t-1})=\mathcal{N}(x_{t}|\sqrt{1-\alpha_{t}}x_{t-1},\alpha_{t})$ (1)

where $\alpha_{t}$ is the noise variance in every time step in case of the
noise has unit variance $\epsilon\sim\mathcal{N}(0,\mathcal{I})$, the single
step formulation can be rewritten as:

$\displaystyle
q(x_{t}|x_{0})=\mathcal{N}(x_{t}|\sqrt{\tilde{\alpha_{t}}}x_{0},1-\tilde{\alpha_{t}})$
(2) $\displaystyle
x_{t}=\sqrt{\tilde{\alpha_{t}}}x_{0}+\sqrt{1-\tilde{\alpha_{t}}}\epsilon$ (3)

where $\tilde{\alpha_{t}}=\prod_{i=1}^{T}(1-\alpha_{i})$. The direct
calculation from $x_{0}$ can allow us to random sample time steps to add noise
during training. In the backward process, a probability function
$p(x_{t-1}|x_{t})$ is learned to denoise the noisy data at previous time step
until reach the initial input $x_{0}$.

The SGM is formularize similarly, in the score Stochastic Differential
Equations (SDEs), the forward process is defined by [41]:

$\displaystyle dx=f(x,t)dt+g(t)w_{t}$ (4)

Figure 1: Upper: The schematic diagram for a noisy shower image (energy
deposited in one layer of calorimeter) $x_{t}$ denoises along the ODE
curvature with varying time steps to clean shower close to $x_{0}$. Lower: ODE
and SDE ways to get the $x_{0}$, ODE flows along the certain direction toward
$x_{0}$, where SDE will have extra noise injection in each time step before
doing the ODE part. The Euler and Heun are the first and second order
calculation for updating $x_{i}$.

where $f$ and $g$ are drift and diffusion coefficients chosen differently for
variance preserving and variance exploding noise, $w$ is the standard Brownian
motion.

The reverse-time SDE is defined as:

$\displaystyle dx=[f(x,t)-g(t)^{2}\nabla_{x}\log{p_{t}(x)}]dt+g(t)d\tilde{w}$
(5)

unlike DDPM, the quantity we are solving is the score function (A vector field
indicating the direction toward regions of increased data density given a
certain level of noise.) $\nabla_{x}\log{p_{t}(x)}$. The learning process of
this gradient of log probability functions offers another possibility to use
not only SDE but also ODE to get the solution.

Then follows by the generalization from [42], we can write the reverse SDE
equation as a sum of probability flow ODE and Langevin diffusion SDE:

$\displaystyle
dx=-\dot{\sigma}(t)\sigma(t)\nabla_{x}\log{p_{t}(x)}dt-\beta(t)\sigma(t)^{2}\nabla_{x}\log{p_{t}(x)}dt+\sqrt{2\beta(t)}\sigma(t)dw_{t}$
(6)

where $\sigma(t)$ is the noise schedule, $\beta(t)=\dot{\sigma}(t)/\sigma(t)$,
the first term is the marginally equivalent ODE, the second and third terms
are added as the deterministic flow and noise injection term for Langevin
diffusion, the schematic diagram as shown in Fig. 1.

This determines the characteristics of the general ODE/SDE solver. ODE solvers
often have smaller discretization errors than SDE solvers in smaller sampling
steps. However, as the sampling steps become larger, the accumulated error for
SDEs would be much smaller than that for ODEs, resulting in a smaller total
error (better performance) in large sampling steps. The noise scheduler
determines how fast those samplers should learn along different time steps. In
the paper, we will cover up to seven mainstream samplers nowadays (four ODEs
and three SDEs) over different schedulers.

The schedulers and samplers directly decide how the noisy data $x_{t}$ will
reach $x_{0}$ with proper noise levels. Our choices for schedulers and
samplers are listed belows:

  * •

Schedulers:

    * –

The default cosine beta scheduler for discrete time DDPM pretrained model,
this is the baseline to compare with performance for other schedulers and
samplers. $\bar{\alpha}=\cos{(\frac{\pi(t/T+s)}{2(1+s)})}$ where $s=0.008$
that determines the minimum noise for last steps. We can tell by the function,
the noises are added by a linearly decrease rate during intermediate steps to
avoid too much noise in the last few steps.

    * –

Variance Preserving schedulers used in [41], defined as a continous time step
$t$ and
$\sigma=\sqrt{(\exp{\frac{\beta_{\mathrm{max}}t^{2}}{2}+\beta_{\mathrm{min}}t}-1)}$.

    * –

The schedulers proposed by Karras [42].
$t_{i<N}=(\sigma_{\mathrm{max}}^{1/\rho}+\frac{i}{N-1}(\sigma_{\mathrm{min}}^{1/\rho}-\sigma_{\mathrm{max}}^{1/\rho}))^{\rho}$
and $\sigma_{N}=0$. It’s the monotonically decreasing function, so the noise
added in a decreasing rate, $\rho$ determines the curvature. Larger $\rho$
means more steps near $\sigma_{\mathrm{min}}$ are shortened. Reasonable value
for $\sigma_{\mathrm{min}},\sigma_{\mathrm{max}},\rho$ in our study should
around $[0.01,0.001],[20,80],[5,10]$ respectively. This scheduler is helpful
for generating good quality of data in low sampling step region.

    * –

Combine Karras schedulers proposed with Lu [43] by transforming the $\sigma$
to log scale.
$t_{i<N}=(\log{(\sigma_{\mathrm{max}})}^{1/\rho}+\frac{i}{N-1}(\log{(\sigma_{\mathrm{min}})}^{1/\rho}-\log{(\sigma_{\mathrm{max}})}^{1/\rho}))^{\rho}$.
The values for $\sigma_{\mathrm{min}},\sigma_{\mathrm{max}}$ remains the same
region, $\rho=1$ to keep a healthy noise level range. This scheduler is even
more decreasing than Karras schedulers, so the step size changes more rapidly
in low sampling step while more slowly in last steps.

  * •

Samplers:

    * –

The same DDPM samplers are used in [20, 39].

    * –

EDM Stochastic sampler with Heun’s method, as illustrated in Fig. 1, a noise
injection step is performed to current time step $\hat{t_{i}}=t_{i}+\gamma
t_{i}$ where $\gamma$ is determined by three parameters
$\gamma=\min{(S_{churn}/N,\sqrt{2}-1)}$ when $S_{\mathrm{min}}\leq t_{i}\leq
S_{\mathrm{max}}$. The $S_{\mathrm{max}},S_{\mathrm{min}}$ is the range for
injection, $S_{churn}$ controls how fast it converges to zero noise. The
optimal sets of this can make the sampling faster and better. Contrarily, a
bad sets of those stochastic term will make possible degradation of the
generated sample. In our study, such degradation only happens when extreme
values are selected. The difference between this approach and Euler Maruyama
is that the denoising steps in EDM samplers happens in some intermediate state
$\hat{t_{i}}$, the difference become diminishing when N goes large. This SDE
sampler, unlike other SDEs that only give satisfied result after a very large
sampling steps, can converge quickly becuase of this.

    * –

Restart sampler proposed in [44]. This sampler introduces a way to strongly
contract the accumulated error based on EDM ODE Heun sampler. A back-and-forth
ODE step will repeat $K$ times in the predefined single or multiple intervals
$[t_{\mathrm{max}},t_{\mathrm{min}}]$.

    * –

DPM-Solver samplers proposed in [43, 45]. A popular ODE solver that can
speedup the sampling process by applying change-of-variable to noise level. We
choose to use the latest multistep solver that calculate iteratively by the
knowledge of current and next timestep.

    * –

Combined DPM-Solver and EDM SDE samplers. Instead of using Euler and Heun
step, it uses the "mid-point" step by change-of-variable to intermediate state
$\hat{t_{i}}$.

    * –

LMS sampler. For each iteration, the sampler will denoise the input data using
the model and updates the $x_{i}$ using the linear multistep integration
method with the calculated coefficients of predefined orders.

    * –

Uni-PC sampler introduced in [46]. A unified predictor-corrector framework
that can be applied on many existing samplers. We apply it with DPM-Multistep-
sampler with both varying coefficient and $B(h)$ solver proposed in the paper.

##  3 Dataset

###  3.1 CaloChallenge

Fast Calorimeter Simulation Challenge 2022 [34] is the `GEANT4` simulation for
particle showers deposited in the calorimeter. The purpose for this challenge
is to develop and benchmark the fast and high-fidelity calorimeter simulation
using modern techniques from deep learning. Three datasets are provided
ranging from easy to hard. The complexity of dataset is determined by the
dimensionality of the calorimeter showers, including factors such as the
number of layers and voxels on each layer.

In each dataset, particles are simulated to propagate along the z-axis inside
the cylindrical and concentric detector. Similar to the sampling layers in the
common calorimeter, a discrete number of layers are placed with specific
radial and angular bins $r$ and $\alpha$. Two classes of quantities are
stored: the incident particle energies and voxel energies stored in each
shower. By mapping the voxel location with given $r$ and $\alpha$ bins, one
can reconstruct the layer deposited energy inside the shower. This offers
flexibility for designing the generative modelling tasks using either low or
high level features.

This study will primarily focus on testing dataset 2 which simulates a
physical detector with 45 layers of absorber material (Tungsten) and active
material (Silicon). A total of 100,000 electrons in each file enter the
calorimeter with a direction. The incoming particle energies span over a log-
uniform distribution from 1 GeV and 1 TeV. Each layer contains 16 angular and
9 radial bins, so 16$\times$9 $=$ 144 cells, and a total number of
144$\times$45 $=$ 6480 voxels in each shower. A minimal cutoff around 15 keV
is applied to voxel energy. 111The rationale behind selecting this dataset is
that previous studies indicate it exhibits a comparable level of difficulty to
the more finely granulated dataset 3, all while achieving training and
inference speeds that are ten times faster.

###  3.2 JetNet

JetNet [35] is the point cloud like dataset which consists of 200,000
simulated events. Each event contains the corresponding jet types with
transverse momenta $p_{\mathrm{T}}\sim 1$ TeV, originating from either gluons,
top or light quarks. Additionally, a narrow kinematic window is imposed,
ensuring that the relative transverse momenta of the generated parton and
gauge boson $\Delta p_{\mathrm{T}}/p_{\mathrm{T}}\leq 0.1$. The jet clustering
is performed using the anti$-k_{\mathrm{T}}$ algorithm [47] within a cone size
of $R$ = 0.4.

For up to 30 particle constituents stored in one jets by descending
$p_{\mathrm{T}}$ order, three features are provided: the relative angular
coordinate of particles w.r.t. the jets:
$\eta^{\mathrm{rel}}=\eta^{\mathrm{particle}}-\eta^{\mathrm{jet}}$,
$\mod{(\phi^{\mathrm{rel}}=\phi^{\mathrm{particle}}-\phi^{\mathrm{jet}},2\pi)}$
and relative transverse momentum
$p^{\mathrm{rel}}_{\mathrm{T}}=p^{\mathrm{particle}}_{\mathrm{T}}/p^{\mathrm{jet}}_{\mathrm{T}}$.
There is also a fourth variable that masks whether the particles is genuine or
zero-padded when there are fewer than 30 constituents stored in the jets.

We will show preliminary results for testing on gluon and top quark jets. The
gluon jet can serve as a baseline test since it usually produces the dominant
background for many physics processes and has a relatively simple topology.
The top jet has more subtle substructure and presents challenging tasks for
generative models.

##  4 Training

In this section, we provide an overview for the data processing, and training
details on two datasets mentioned in Section 3.

###  4.1 CaloChallenge

We follow a common preprocessing procedure for CaloChallenge dataset.

  * •

First, the voxel energies $v_{i}$ are normalized by the incident particle
energy.

  * •

Then a logit-norm transformation is applied to ensure input data has zero mean
and unit variance.

$\displaystyle u_{i}$ $\displaystyle=\log{(\frac{x}{1-x})},x=v_{i}-2\delta
v_{i}+\delta$ (7) $\displaystyle u^{\prime}_{i}$
$\displaystyle=\frac{u_{i}-\bar{u}}{\sigma_{u}}$

Here, $\delta=10^{-6}$ is introduced to avoid the discontinuity.

We choose to use the network from [20] which is composed of a conditional UNet
with middle layer self-attention layer, cynlindircal convolution and Geometry
Latent Mapping (GLaM). The latter two techniques are found to be beneficial
for generating quantities in $r$ and $\alpha$ bins. In order to demostrate the
portability of our studies, the neural network architect remains unchanged
from the version in the literature [20].

###  4.2 JetNet

For JetNet dataset, we follow the preprocessing procedure from [22, 23]. The
Equivariant Point Cloud (EPiC) layer based on the deep sets framework is
chosen to be used as it shows great quality of generated classes while having
lighter architecture than transformer-based models [23].

Training the diffusion model is essential for generating higher quality
showers across different noise levels. We delve into two questions: 1. How can
the data with added noise be correctly weighted and learned by the model so it
can be denoised more easily in the backward process? 2. How can the training
converges faster to a better reliable result?

We follow the way from [42] to rewrite the score function to a denoiser
function which minimizing the $L2$ denoising error at each noise level:

$\displaystyle\nabla_{x}\log{p_{t}(x)}=(D(x)-x)/t^{2}$ (8)

A set of preconditioning parameters used to predict noise $\epsilon$ with this
denoiser function:

$\displaystyle\mathcal{L}=\mathbb{E}_{t,\epsilon}[w(t)||F_{\theta}(c_{in}*x_{t},t,z)-\frac{1}{c_{out}}(x_{0}-c_{\mathrm{skip}}(t)*(x_{t}))||^{2}_{2}]$
(9)

where $c_{in},c_{out}$ maps the input and output data magnitude, and
$c_{\mathrm{skip}}$ controls the skip connection of the network.
$c_{in}=1/\sqrt{t^{2}+\sigma_{c}^{2}}$ and
$c_{out}=(t*\sigma_{c})/\sqrt{t^{2}+\sigma_{c}^{2}}$, $\sigma_{c}$ is the
constant value, a reasonable value is around $[0.5,1.5]$ to avoid large
variation for gradient scale. The
$c_{\mathrm{skip}}=\sigma_{c}^{2}/(t^{2}+\sigma_{c}^{2})$ is important
quantity to partially mask the skip connection when the noise level goes
large. This can prevent the network from inaccurate prediction. The setting is
important in our case given noise varies widely in most of samplers we show
later. When the $t$ goes higher, the $c_{\mathrm{skip}}$ becomes smaller. In a
way, this compels the skip connection to be deactivated, encouraging the
network to prioritize learning the signal for overriding the input rather than
predicting noise to correct it. We choose the continuous timestep training
where $\log{t}=\mathcal{N}(P_{\mathrm{mean}},P^{2}_{\mathrm{std}})$ where
$P_{\mathrm{mean}},P_{\mathrm{std}}$ is -1.2 and 1.2 correspondingly.

From our study, the choice of the weighting function is very important as it
not only determines the speed of loss convergence but also the quality of the
outputs and how the specific energy distribution tends to evolve across
different time steps. To mitigate the instability of the optimization over the
training, and balance the gradient magnitude over diverse noise levels, one
loss weight category is often mentioned $-$ Signal to noise ratio (SNR)
weighting [29, 48, 49, 50]. We choose the min-SNR weighting as it is found to
be beneficial for training high-resolutions dataset. The original min-SNR
weighting proposed from [51] defined to be min{SNR(t),$\gamma$}, where
$\gamma$ is the upper limit for weighting value, usually set to be 2-5. This
could avoid the model putting too much weight on the low noise level.

In stead of directly clamping a constant value, we propose a softer version of
min-SNR weighting which is more flexible to find the optimal value for better
training. Specifically, the weight function is

$\displaystyle w(t)=(t*\sigma_{c})^{2}/(t^{2}+\sigma_{c}^{k})^{2}$ (10)

that $t$ and $\sigma_{c}$ is the same as the preconditioning parameters, the
reasonable value of k in the range of $[2,4]$. This function is even with one
minimum value of 0 around zero $t$ and two local maxima on positive and
negative t-axes. The constant k determines the local maximum weight value, and
the SNR value where the maxima start to smoothly fall down to 0. The
advantages of this weighting scheme will be shown later in Sec 5 with
comparison of default constant (unit) weighting strategy.

Figure 2: The distribution for total shower energy divided by the incident
particle energy, tale dotted line is the GEANT4 reference, red solid line is
the generated distribution by 79 steps EDM samplers. Upper: 30, 70, 100, 200
epochs default denoiser training with unit weight. Lower: 30, 70, 100, 200
epochs training with our proposed min-SNR weighting. The default training
still struggles to learn this variable after 200 epochs whereas the training
with proposed weighting strategy already gives good shape alignment after 100
epochs.

##  5 Performance

We have previously discussed various schemes for the backward diffusion
process. In order to comprehensively compare their performance and enhance
understanding of how to choose the appropriate samplers and schedulers for a
specific purpose, we offer several methods for evaluation.

###  5.1 CaloChallenge

There are multiple important distributions related to the layer and voxel
energies for CaloChallenge dataset. We compare the generated samples with
`GEANT4` as reference. The voxel energy distribution defined as the
distribution of raw generated low level voxel energies. The total energy
distribution defined as the distribution of total energy of the generated
shower. The center energy of the shower defined as $\bar{x}=\frac{\langle
x_{i}E_{i}\rangle}{\sum E_{i}}$ for cell location $x_{i}$ and energy $E_{i}$,
and shower width defined as $\sqrt{\bar{x^{2}}-\bar{x}^{2}}$, we will show
these two quantities at the polar coordinate $\alpha$ and $r$. The mean layer
energy distribution defined as the mean of generated layer energies. The
shower response $E_{\mathrm{ratio}}$ defined as the total shower energy
divided by the incident particle energies, this quantity is extremely
important as a correction factor for energy calibration of many particles at
LHC. We will also evaluate the separation power between generated
$E_{\mathrm{ratio}}$ and reference one.

In addition, a DNN classifier with 2 layers, each comprising 2048 nodes and a
0.2 dropout rate, is employed to assess the area under the receiver operating
characteristic curve (AUC) for all high-level features, including center
energies, shower width, and layer energies, of both generated samples and
references. An AUC value that close to 0.5 suggests that the generated and
target classes are nearly indistinguishable. Conversely, a higher AUC value
indicates a more evident distinction between the quantities of the two
classes.

Figure 3: Comparison plots for center energy of 5000 showers in radial bins.
left: EDM with 30 steps, middle: EDM with 79 steps, right: Restart with 30
steps.

First, a comparison plots for $E_{\mathrm{ratio}}$ with default EDM training
loss weight and our proposed training weight shown in Fig. 2. Several
observations can be made from the plots. The choice of the weighting function
influences how the kinematics distribution evolves over time. After 200
training epochs, the model trained with the default weighting still struggles
to learn the $E_{\mathrm{ratio}}$ distribution. In contrast, employing
proposed min-SNR weighting results in the generated $E_{\mathrm{ratio}}$
aligning with the reference shape after only 100 epochs. This observation
affirms that min-SNR weighting expedites the convergence of high-level
features to achieve better performance compared to constant weighting. 222A 1D
layer diffusion can generate better $E_{\mathrm{ratio}}$ distribution. But it
might be over simplified as the model only considers the layer deposited
energy. The paper only considers the low level voxel diffusion as a more
realistic and challenging case for fast calorimeter simulation

Few plots for center energy in $r$ direction are selected in Fig. 3. Notably,
the EDM samplers and Restart samplers with Karras schedulers can learn this
quantities really well. The ratio difference is almost indistinguishable with
naked eyes. We can tell from the small ($\sim$30) steps EDM, the differences
in last three radial bins are larger meaning it underestimates the lower
energy region. When the steps become larger, we see the EDM predicts well in
all radial bins. The Restart sampler can achieve this with only 30 steps.
Because the Restart samplers works with a predefined parameters for
configuring the EDM ODE. The generation time is directly proportional with the
restart iterations in the configuration, usually longer than default EDM
samplers. A lack of principled way for these hyperparameters is the inherent
limitation for this samplers. Later we will see the best performance is
achieved around 30-70 steps as we only uses specific configurations in this
range.

Then the mean layer energies, total shower energies and voxel energies for the
same three cases are shown in Fig. 4, 5, 6 (right).

Figure 4: Comparison plots for mean layer energies of 5000 showers. left: EDM
with 30 steps, middle: EDM with 79 steps, right: Restart with 30 steps. Figure
5: Comparison histograms for 5000 total shower energies. left: EDM with 30
steps, middle: EDM with 79 steps, right: Restart with 30 steps.

Figure 6: Comparison histograms for raw voxel energies of 5000 showers: Left:
30 steps EDM, 79 steps EDM, 30 steps Restart samplers; Right: 400 steps DDPM,
36 steps LMS, and 15 steps DPM-Solver $++$ samplers.

Most of the samplers need a relatively long sampling time to learn the mean
energy of the last few layer correctly since they have voxels with lowest
energy thresholds. The low steps EDM sampler has good performance on low
energy voxel while mediocre performance on high energy tails of voxel
distributions and low energy shower. DPM-solver samplers with Karras scheduler
does very well on predicting the voxel energy around 15 steps, achieving an
exceptional agreement on both low and high energy tails shown in the left plot
of Fig. 6. All proposed samplers have comparable and better performance with
smaller steps. Indeed some of them are struggling to match the low voxel
energy, the presence of the tail is probably a consequence of model itself and
too low energy threshold. It is expected this deficiency should not have big
impact on the downstream reconstruction of shower.

However, DPM-solver sampler seems to perform poorly on predicting the low and
high-level features simultaneously. The Uni-PC sampler has very similar
performance with DPM-solver sampler and only gives slightly better performance
over very low sampling (5-10) steps. The lowest AUC for high level feature
classifiers is around 60 at 20th step. We showed the AUC plots for evaluating
different samplers 20 times in Fig. 7. The EDM and Restart samplers surpasses
the previous 400 steps DDPM samplers at 50, 25 steps correspondingly, combined
with other plots we showed, achieving the same and better performance with a
3-8 times speedup. The LMS ODE has a better performance than EDM below 45
steps, but reaches a worse convergence than other samplers. The best AUC value
for Restart samplers achieves at 30-36 and 79 steps, this might because we use
the predefined configuration specifically for those steps. Most of the
samplers we use actually have more parameters to tune than the baseline DDPM.
For the common ODE samplers, a closer matching noise level during sampling
would often yield better results., LMS sampler involves an additional
parameter "order" of the coefficient which makes the generation time longer as
it increases. A possible future development is to fine-tune the configuration
in higher steps.

Figure 7: AUC values for different samplers as a function of sampling steps, the shaded area is possible AUC value range for 20 evaluations through independent classifier training. Figure 8: Separation power for $E_{\mathrm{ratio}}$ with different samplers and schedulers as a function of sampling steps, the shaded area is possible values of separation power for 20 independent samplings. Figure 9: The distribution for total shower energy divided by the incident particle energy by left: 400 steps DDPM, middle: 30 steps Restart samplers, left: high step EDM samplers Samplers(steps) | AUC | Separation power | FPD
---|---|---|---
DDPM(79) | 55.3 | 0.0810 | 0.074
DDPM(200) | 53.2 | 0.0344 | 0.046
DDPM(400) | 52.6 | 0.0155 | 0.043
EDM(36) | 54.1 | 0.0256 | 0.035
EDM(79) | 52.0 | 0.0076 | 0.027
EDM(200) | 51.5 | 0.0055 | 0.023
EDM(400) | 51.1 | 0.0039 | 0.021
EDM_DPM++(79) | 52.3 | 0.0103 | 0.026
EDM_Lu(79) | 52.2 | 0.0086 | 0.026
Restart(18) | 55.2 | 0.0169 | 0.059
Restart(36) | 52.0 | 0.0073 | 0.025
Restart(79) | 51.8 | 0.0057 | 0.022
LMS(36) | 53.8 | 0.0305 | 0.095
DPM++(20) | 59.8 | 0.0534 | 0.146
UniPC(20) | 60.3 | 0.1304 | 0.152
Table 1: Summary table for high-level feature classifier mean AUC,
$E_{\mathrm{ratio}}$ separation power, and Frechét Particle Distance [52] of
of different samplers. (400-step DDPM as baseline)

Now we take a look at one of the most difficult and important quantities for
the low level voxel diffusion – $E_{\mathrm{ratio}}$. In Fig.8, we choose 6
cases in our study. First, much faster convergences have been seen from all
new introduced samplers by 5-20 times. The result matches with the Table.3
mentioned in Song’s [53]. For EDM Heun SDE sampler, the separation power
quickly drops around 60 steps and achieves better performance than 400 steps
DDPM. If we zoom-in the figure, at around 100-200 steps, the separation power
becomes 3-5 times weaker than the baseline DDPM samplers. In Fig. 9, we have
seen really aligned shape between generated and reference classes for Restart
and EDM samplers. As the steps goes even higher, the discrepancy between two
histogram become smaller while the shape remains aligned. The histogram
geneated from Restart samplers does not fit perfectly on lower ratio region,
but still better than DDPM in other regions. This is crucial for us to perform
accurate energy calibration from low-level fast calorimeter simulation later.

The preliminary Frechét Particle Distance (FPD) [52] values for some samplers
are also provided in Table. 1. The FPD metric is relatively biased and
unstable among sub samples we choose. So we follow the way from [20] to
subtract the non-zero FPD value between two `GEANT4` sample. Surprisingly, we
find the EDM sampler converges even faster in this metric, surpassing the DDPM
sampler around 36 steps. The 18-step Restart sampler has slightly worse
performance, this is mainly because we choose the predefined configuration for
smallest $E_{\mathrm{ratio}}$ separation power, a throughout fine-tune can
make FPD value 20-30 percent smaller with almost same generation time.
However, it still obtains the best FPD results with larger steps.

We choose Karras and Lu schedulers to illustrate the impacts of different
noise schedulers on the same samplers. Changing the Karras scheduler to Lu
will decrease the separation power in smaller steps ($\leq 100$) but increase
it in larger steps. The choice of schedulers provides a training-free method
to balance and control the quality and speed of the backward diffusion
process. The EDM/DPM-Solver combined sampler with Karras scheduler performs
pretty similarly as the EDM sampler with Lu scheduler. The best performances
of Restart samplers again are achieved in 30-36 steps. It is crafted for
efficiently conducting fast ODE sampling with outstanding performance and can
be fine-tuned with post-training configuration. The DPM-solver exhibits a
rapid drop in separation power within 10-20 steps but ceases to improve beyond
those steps, achieving better performance than the 100-step DDPM around the
15th step. The plot further confirms that ODE sampler does better in low
sampling steps as the discretization error of them are smaller than SDE ones
in this region, the contraction of accumulated errors from SDE sampling
process become much stronger than ODE so that a better convergence observed in
higher sampling steps than ODE ones. Some ODE samplers, such as LMS and DPM-
solver, perform well on some metrics but not all. There is not a clearly
superior option that can achieve the best quality while having the lowest
generation time.

Figure 10: Comparison plots for relative angular coordinate of particle
constituents inside gluon initiated jets with 200 steps DDIM Euler Sampler, 18
steps multi-level Restart Sampler, 50 steps EDM samplers . left: $\Delta\eta$,
right: $\Delta\phi$.

Figure 11: Comparison plots for relative angular coordinate of particle
constituents inside top quark jets with 200 steps DDIM Euler Sampler, 18 steps
multi-level Restart Sampler, 50 steps EDM samplers . left: $\Delta\eta$,
right: $\Delta\phi$.

Figure 12: Comparison plots for relative jet mass with 200 steps DDIM Euler Sampler, 18 steps multi-level Restart Sampler, 50 steps EDM samplers . left: gluon initiated jet, right: top quark jet. Samplers(steps) |  AUC (Gluon Jet) |  FPD (Gluon Jet) |  AUC (Top Jet) |  FPD (Top Jet)
---|---|---|---|---
Euler(200) | 52.6 | 0.15 | 53.5 | 0.19
EDM(50) | 51.7 | 0.05 | 51.9 | 0.07
Restart(18) | 51.6 | 0.04 | 51.5 | 0.05
Restart_multi(18) | 51.2 | 0.02 | 51.2 | 0.02
Table 2: Summary table for low-level feature classifier mean AUC and Frechét
Particle Distance calculated from jet-level classifier of different samplers
on gluon-initiated jets top quark jets

###  5.2 JetNet

We employ the same network training techniques and representative samplers for
the JetNet dataset to demonstrate that these methods are entirely portable
across various datasets. As depicted in Fig. 10\- 11, the low level feature
generations with a 50-step EDM sampler and 18-step Restart sampler show better
performance than a 200-step DDIM [33] Euler sampler on gluon jets. For more
complicated topology, EDM and Restart have advantages in simulating the tails
of the distribution over large $\eta$ region. EDM sampler has small deficiency
for underestimating the particles from top quark jets after $\Delta\phi>0.2$.
This would have impact on reconstructing the jet level variables as shown in
Fig. 12. When simulating the complex distribution in a highly boosted
environment, the EDM and DDIM Euler samplers fail to achieve satisfactory
alignment around the two asymmetric peaks in the relative jet mass
distribution. The low step Restart sampler gives obviously better generated
distribution than others. We evaluate the samplers via the three low level
features with the same DNN classifier, and the values for the Fréchet Particle
distance (FPD) calculated from the Particle [54] multi-classifier with N = 6.
Performances are summarized on Table. 2.

The boosts brought by those changes isn’t as pronounced as they are in
CaloChallenge. It may be because methods are more applicable to UNet and
pixelated data than point clouds network.

##  6 Conclusion and outlooks

In this study, we conducted a comprehensive investigation into the utilization
of various state-of-the-art samplers and schedulers across diverse datasets in
collider physics. We introduced a soft version of loss weighting techniques to
expedite the convergence of physics-informed quantities towards a better
reliable result. Our study presents several approaches for evaluating
generated samples across different scenarios, providing a range of options for
accelerating diffusion model without additional training or achieving superior
results with fewer enhancements. Innovative sampling methods, such as Restart,
exhibit better performance with just 18-30 steps, leading to a
$\mathcal{O}(10)$ times speeding up without a significant performance loss. We
showcase a potentially new baseline for training and sampling diffusion models
in this rapidly evolving field.

Future development will aim to identify the optimal sets for these samplers on
each dataset to achieve the maximum possible boost. Additionally, exploration
into other areas, such as score distillation sampling and model fine-tuning,
will be undertaken to obtain similar performance levels with even faster
simulations.

## Acknowledgements

We thank Kevin Pedro and Oz Amram for fruitful discussions about details of
`CaloDiffusion` project, and insightful comments for different evaluation
metrics.

## References

  * [1] L. Evans and P. Bryant,  _LHC Machine_ ,  Journal of Instrumentation 3(08), S08001 (2008),  10.1088/1748-0221/3/08/S08001.
  * [2] G. Aad _et al._ ,  _The ATLAS Experiment at the CERN Large Hadron Collider_ ,  JINST 3, S08003 (2008),  10.1088/1748-0221/3/08/S08003.
  * [3] S. Chatrchyan _et al._ ,  _The CMS Experiment at the CERN LHC_ ,  JINST 3, S08004 (2008),  10.1088/1748-0221/3/08/S08004.
  * [4] S. Agostinelli _et al._ ,  _GEANT4–a simulation toolkit_ ,  Nucl. Instrum. Meth. A 506, 250 (2003),  10.1016/S0168-9002(03)01368-8.
  * [5] G. Apollinari, O. Brüning, T. Nakamoto and L. Rossi,  _High Luminosity Large Hadron Collider HL-LHC_ ,  10.5170/CERN-2015-005.1 (2015).
  * [6] 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.
  * [7] 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.
  * [8] R. Zhang,  _AtlFast3, the ATLAS fast simulation tool in Run3_ (2023).
  * [9] E. Buhmann, S. Diefenbacher, E. Eren, F. Gaede, D. Hundhausen, G. Kasieczka, W. Korcari, K. Krüger, P. McKeown and L. Rustige,  _Hadrons, better, faster, stronger_ (2021), 2112.09709.
  * [10] I. Pang, J. A. Raine and D. Shih,  _Supercalo: Calorimeter shower super-resolution_ (2024), 2308.11700.
  * [11] M. F. Giannelli and R. Zhang,  _Caloshowergan, a generative adversarial networks model for fast calorimeter shower simulation_ (2023), 2309.06515.
  * [12] S. Hoque, H. Jia, A. Abhishek, M. Fadaie, J. Q. Toledo-Marín, T. Vale, R. G. Melko, M. Swiatlowski and W. T. Fedorko,  _Caloqvae : Simulating high-energy particle-calorimeter interactions using hybrid quantum-classical generative models_ (2023), 2312.03179.
  * [13] J. Birk, E. Buhmann, C. Ewen, G. Kasieczka and D. Shih,  _Flow matching beyond kinematics: Generating jets with particle-id and trajectory displacement information_ (2023), 2312.00123.
  * [14] S. Bieringer, A. Butter, S. Diefenbacher, E. Eren, F. Gaede, D. Hundhausen, G. Kasieczka, B. Nachman, T. Plehn and M. Trabs,  _Calomplification — the power of generative calorimeter models_ ,  Journal of Instrumentation 17(09), P09028 (2022),  10.1088/1748-0221/17/09/p09028.
  * [15] M. A. W. Scham, D. Krücker, B. Käch and K. Borras,  _Deeptreegan: Fast generation of high dimensional point clouds_ (2023), 2311.12616.
  * [16] K. Dohi,  _Variational autoencoders for jet simulation_ (2020), 2009.04842.
  * [17] C. Krause and D. Shih,  _Caloflow ii: Even faster and still accurate generation of calorimeter showers with normalizing flows_ (2023), 2110.11377.
  * [18] C. Krause, B. Nachman, I. Pang, D. Shih and Y. Zhu,  _Anomaly detection with flow-based fast calorimeter simulators_ (2023), 2312.11618.
  * [19] V. Mikuni and B. Nachman,  _Score-based generative models for calorimeter shower simulation_ ,  Physical Review D 106(9) (2022),  10.1103/physrevd.106.092009.
  * [20] O. Amram and K. Pedro,  _Denoising diffusion models with geometry adaptation for high fidelity calorimeter simulation_ (2023), 2308.03876.
  * [21] E. Buhmann, S. Diefenbacher, E. Eren, F. Gaede, G. Kasieczka, A. Korol, W. Korcari, K. Krüger and P. McKeown,  _Caloclouds: Fast geometry-independent highly-granular calorimeter simulation_ (2023), 2305.04847.
  * [22] M. Leigh, D. Sengupta, G. Quétant, J. A. Raine, K. Zoch and T. Golling,  _Pc-jedi: Diffusion for particle cloud generation in high energy physics_ (2023), 2303.05376.
  * [23] E. Buhmann, C. Ewen, D. A. Faroughy, T. Golling, G. Kasieczka, M. Leigh, G. Quétant, J. A. Raine, D. Sengupta and D. Shih,  _Epic-ly fast particle cloud generation with flow-matching and diffusion_ (2023), 2310.00049.
  * [24] V. Mikuni and B. Nachman,  _Caloscore v2: Single-shot calorimeter shower simulation with diffusion models_ (2023), 2308.03847.
  * [25] T. Salimans and J. Ho,  _Progressive distillation for fast sampling of diffusion models_ (2022), 2202.00512.
  * [26] E. Buhmann, F. Gaede, G. Kasieczka, A. Korol, W. Korcari, K. Krüger and P. McKeown,  _Caloclouds ii: Ultra-fast geometry-independent highly-granular calorimeter simulation_ (2023), 2309.05704.
  * [27] Y. Song, P. Dhariwal, M. Chen and I. Sutskever,  _Consistency models_ (2023), 2303.01469.
  * [28] M. Leigh, D. Sengupta, J. A. Raine, G. Quétant and T. Golling,  _Pc-droid: Faster diffusion and improved quality for particle cloud generation_ (2023), 2307.06836.
  * [29] R. Rombach, A. Blattmann, D. Lorenz, P. Esser and B. Ommer,  _High-resolution image synthesis with latent diffusion models_ (2022), 2112.10752.
  * [30] N. Ruiz, Y. Li, V. Jampani, Y. Pritch, M. Rubinstein and K. Aberman,  _Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation_ (2023), 2208.12242.
  * [31] L. Zhang, A. Rao and M. Agrawala,  _Adding conditional control to text-to-image diffusion models_ (2023), 2302.05543.
  * [32] S. Zhai, Y. Dong, Q. Shen, S. Pu, Y. Fang and H. Su,  _Text-to-image diffusion models can be easily backdoored through multimodal data poisoning_ (2023), 2305.04175.
  * [33] J. Song, C. Meng and S. Ermon,  _Denoising diffusion implicit models_ (2022), 2010.02502.
  * [34] M. F. Giannelli, G. Kasieczka, B. N. Claudius Krause, D. Salamani, D. Shih and A. Zaborowska,  _Fast calorimeter simulation challenge 2022_ (2022).
  * [35] R. Kansal, J. Duarte, H. Su, B. Orzari, T. Tomei, M. Pierini, M. Touranakou, J.-R. Vlimant and D. Gunopulos,  _Particle cloud generation with message passing generative adversarial networks_ (2022), 2106.11535.
  * [36] J. Sohl-Dickstein, E. A. Weiss, N. Maheswaranathan and S. Ganguli,  _Deep unsupervised learning using nonequilibrium thermodynamics_ (2015), 1503.03585.
  * [37] L. Yang, Z. Zhang, Y. Song, S. Hong, R. Xu, Y. Zhao, W. Zhang, B. Cui and M.-H. Yang,  _Diffusion models: A comprehensive survey of methods and applications_ (2023), 2209.00796.
  * [38] B. B. Moser, A. S. Shanbhag, F. Raue, S. Frolov, S. Palacio and A. Dengel,  _Diffusion models, image super-resolution and everything: A survey_ (2024), 2401.00736.
  * [39] J. Ho, A. Jain and P. Abbeel,  _Denoising diffusion probabilistic models_ (2020), 2006.11239.
  * [40] A. Nichol and P. Dhariwal,  _Improved denoising diffusion probabilistic models_ (2021), 2102.09672.
  * [41] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon and B. Poole,  _Score-based generative modeling through stochastic differential equations_ (2021), 2011.13456.
  * [42] T. Karras, M. Aittala, T. Aila and S. Laine,  _Elucidating the design space of diffusion-based generative models_ (2022), 2206.00364.
  * [43] C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li and J. Zhu,  _Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps_ (2022), 2206.00927.
  * [44] Y. Xu, M. Deng, X. Cheng, Y. Tian, Z. Liu and T. Jaakkola,  _Restart sampling for improving generative processes_ (2023), 2306.14878.
  * [45] C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li and J. Zhu,  _Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models_ (2023), 2211.01095.
  * [46] W. Zhao, L. Bai, Y. Rao, J. Zhou and J. Lu,  _Unipc: A unified predictor-corrector framework for fast sampling of diffusion models_ (2023), 2302.04867.
  * [47] M. Cacciari, G. P. Salam and G. Soyez,  _The anti-ktjet clustering algorithm_ ,  Journal of High Energy Physics 2008(04), 063–063 (2008),  10.1088/1126-6708/2008/04/063.
  * [48] P. Dhariwal and A. Nichol,  _Diffusion models beat gans on image synthesis_ (2021), 2105.05233.
  * [49] C. Meng, R. Rombach, R. Gao, D. P. Kingma, S. Ermon, J. Ho and T. Salimans,  _On distillation of guided diffusion models_ (2023), 2210.03142.
  * [50] J. Ho, T. Salimans, A. Gritsenko, W. Chan, M. Norouzi and D. J. Fleet,  _Video diffusion models_ (2022), 2204.03458.
  * [51] T. Hang, S. Gu, C. Li, J. Bao, D. Chen, H. Hu, X. Geng and B. Guo,  _Efficient diffusion training via min-snr weighting strategy_ (2023), 2303.09556.
  * [52] R. Kansal, A. Li, J. Duarte, N. Chernyavskaya, M. Pierini, B. Orzari and T. Tomei,  _Evaluating generative models in high energy physics_ ,  Physical Review D 107(7) (2023),  10.1103/physrevd.107.076017.
  * [53] Y. Song and P. Dhariwal,  _Improved techniques for training consistency models_ (2023), 2310.14189.
  * [54] H. Qu and L. Gouskos,  _Jet tagging via particle clouds_ ,  Physical Review D 101(5) (2020),  10.1103/physrevd.101.056019.

##  Appendix A Hyperparameters of Samplers and Schedulers

We will list more details about the hyperparamter range we ended up choosing
for each sampler.

Dataset | $S_{\mathrm{Churn}}$ | $S_{\mathrm{max}}$ | $S_{\mathrm{min}}$ | $S_{\mathrm{noise}}$ | $\sigma_{\mathrm{max}}$ | $\sigma_{\mathrm{min}}$
---|---|---|---|---|---|---
CaloCha. | [15-40] | [20-40] | [0.01-0.1] | [1.001-1.004] | [40-80] | [0.002-0.01]
JetNet | [20-50] | [20-39] | [0.01-0.1] | [1.001-1.006] | [30-80] | [0.002-0.01]
Table 3: Summary table for Stochastic parameter and scheduler of EDM Heun
Karras sampler

For the Stochastic parameter we listed on Table. 3, the choices of
$S_{\mathrm{Churn}}$ and $S_{\mathrm{min}}$ are very important. In the
CaloChallenge dataset, a too small $S_{\mathrm{Churn}}$ would make the
prediction worse over lower energy region, resulting a higher $E_{ratio}$
separation power, and a high $S_{\mathrm{Churn}}$ would make noise injection
power too high over low sampling step region to converge quickly. With our
training setup, we usually set a larger $S_{\mathrm{Churn}}$ on the Jetnet
dataset. That is also one reason why a smaller improvement is observed here.
The $S_{\mathrm{noise}}$ does not affect too much on both dataset.

We also find several ways to constrain the multilevel configuration for
generating satisfied results with Restart sampler by empirical evaluations.

level | $N_{\mathrm{steps}}$ | $K$ | $\sigma_{\mathrm{max}}$ | $\sigma_{\mathrm{min}}$
---|---|---|---|---
0 | [2-5] | [2-15] | [1.09,0.30] | [0.14, 0.06]
1 | [2-5] | [1-2] | [40.79] | [19.35]
2 | [3-6] | [1-2] | [1.92] | [1.09]
3 | [4-6] | [2-4] | [1.09,0.59,0.30] | [0.59,0.06]
4 | [4-6] | [1-4] | [0.59,0.30] | [0.30,0.06]
5 | [4-7] | [2-6] | [0.30] | [0.06]
Table 4: Summary table for good parameter of Restart Karras sampler, 0 denoted
as the single level Restart.

Oftentimes different $N_{\mathrm{steps}}$ across different level (range of
time steps) yield better FPD value. Gradually increasing $N_{\mathrm{steps}}$
toward last few steps give slightly better result than constant
$N_{\mathrm{steps}}$ among all levels. With a larger $K$ iteration value in
$3^{\mathrm{rd}}$ and $5^{\mathrm{th}}$, we obtain the best results in current
sampling steps, but also increase the generation time by few percent.

Figure 13: FPD values in CaloChallenge for 36 steps multilevel Restart
samplers with 25 different configurations

For DPM-solver, we choose a slightly different noise level. The
$\beta_{\mathrm{max}},\beta_{\mathrm{min}}$ for VP scheduler is $19.9,0.002$.
For Karras and Lu scheduler, $\sigma_{\mathrm{max}},\sigma_{\mathrm{min}}$ is
$[15,40],[0.001,0.01]$. For LMS, we choose $[3,4,6]^{\mathrm{th}}$ order of
calculation. The UniPC sampler is using $2^{\mathrm{nd}}$ order $B(h)$-$2$
solver with the same scheduler as DPM-solver.


# A contrastive learning approach for individual re-identification in a wild
fish population

Ørjan Langøy Olsen  Centre for Artificial Intelligence Research, University of
Agder, Norway  Tonje Knutsen Sørdalen  Centre for Coastal Research, University
of Agder, Norway  Morten Goodwin  Centre for Artificial Intelligence Research,
University of Agder, Norway  Ketil Malde  Institute of Marine Research,
Ecosystem Acoustics Group, Bergen, Norway  Kristian Muri Knausgård
Corresponding<EMAIL_ADDRESS>Top Research Centre Mechatronics, University
of Agder, Norway  Kim Tallaksen Halvorsen  Institute of Marine Research,
Ecosystem Acoustics Group, Bergen, Norway

###### Abstract

In both terrestrial and marine ecology, physical tagging is a frequently used
method to study population dynamics and behavior. However, such tagging
techniques are increasingly being replaced by individual re-identification
using image analysis.

This paper introduces a contrastive learning-based model for identifying
individuals. The model uses the first parts of the Inception v3 network,
supported by a projection head, and we use contrastive learning to find
similar or dissimilar image pairs from a collection of uniform photographs. We
apply this technique for corkwing wrasse, Symphodus melops, an ecologically
and commercially important fish species. Photos are taken during repeated
catches of the same individuals from a wild population, where the intervals
between individual sightings might range from a few days to several years.

Our model achieves a one-shot accuracy of 0.35, a 5-shot accuracy of 0.56, and
a 100-shot accuracy of 0.88, on our dataset.

##  1 Introduction

Physical tagging, using external or internal markings for individual
identification, is a widely used method for monitoring terrestrial and aquatic
animal populations. Information from resightings or recapture of the same
individuals can be used to estimate population size, survival and movement
patterns. However, most tagging methods are costly, intrusive, and labor-
intensive. To our beneift, many animals have natural markings or morphological
features that are unique to individuals that could be used for photo-
identification and replace the need for physical tags [22, 19]. However, for
ecologists, working with fish may mean keeping track of hundreds or
potentially thousands of individuals in a population, which makes manual
photo-identification challenging, if not impossible. For this reason, fully-
or semi-automatic tools for re-identification of individuals would be
immensely useful for ecologists.

Re-identification (re-ID) is different from normal classification in that it
is a few-shot learning problem. Few-shot problems are characterised by having
few samples per class, but there may be a large or indefinite number of
classes. One way to solve such problems is a technique called metric learning,
where data is transformed into embeddings of a lower dimension, that clusters
points from the same class together. Classification can then be performed on
the embeddings. Metric learning approaches have been proved to work well for
re-identification of animal species [18]. A crucial advantage with metric
learning approaches is that the network does not need to be retrained to be
able to add new classes.

Contrastive learning is a technique that can be used to solve few-shot
problems. Constrastive learning compares data and identifies whether they are
similar or dissimilar. A siamese network [1] is the most basic form and takes
two inputs through the same network with the same shared weights and gets an
embedding for both. During training, it tries to predict whether they are of
the same class or not. A major advantage here is that it does not need to know
which class an input belongs to, nor how many classes there are. Triplet
networks [10] are an improvement to the siamese network with three inputs.

The goal of this work was to test the applicability of image based re-ID
analysis for a commercially and ecologically important fish species, the
corkwing wrasse (Symphodus melops). The image dataset consists of standardized
photos of captures and recaptures of individuals in a wild population, where
the time between individual sightings spans from days to several years. The
first step is to detect a fish in an image with an object detector, followed
by a re-identification method. With high enough precision, computer vision re-
ID has the potential to replace physical tagging for individual identification
and may be applied in monitoring of survival rates, growth, movement, and
population size, key knowledge for sustainable management and conservation
[18] [4].

##  2 Related works

Advancements in machine learning have produced powerful techniques for
extracting ecologically important information from image and video data. For
instance, machine learning have successfully been utilized to detect fish
wounds [5], count and categorize organisms in digital photos and real-time
video [14], [12], identify species, [6], and discover, and count creatures
from digital images, [4], and even quantify their behaviour [3].

Some work on the topic of re-identification of fish has been conducted, but
work on wild teleost fish are lacking. Bruslund Haurum et al. [2] achieved an
mAP of 99% on Zebrafish using metric learning with 15 samples per class of 6
classes. Meidell and Sjøblom [16] reports a true positive rate of 96% on 225
thousand images of salmon divided between 715 individuals. Li et al. [13]
achieved an accuracy of 92% using 3412 images of 10 individuals using their
novel FFRNet network. These studies have in common that they were carried out
in captivity and are not using temporally independent observations. In other
words, the individuals did not change morphology through growth, maturation,
senescence, or similar biological processes. Moskvyak et al. [17] used a
metric learning approach on a dataset of 1730 images of 120 manta ray
individuals and achieved an accuracy@1 of 62% and an accuracy@10 of 97%.

##  3 Method

###  3.1 Data collection

The study species, S. melops, is a commercially and ecologically important
species in coastal ecosystems in the Northeastern Atlantic [7]. This species
have two distinct male morphs, colourful large males that build nest and care
for the eggs, and smaller sneaker males, with a more brown coloration
resembling the female morphology (brown and gray) [21]. The dataset was
collected in Austevoll, western Norway, 2018-2021, by catching corkwing wrasse
by fyke nets left in the sea overnight and marking all captured individuals
with uniquely coded passive integrated transponder (PIT) tags (11 mm tags,
RFID Solutions). The tags were implanted in the abdominal cavity of the fish,
see full sampling description in [8] and [9].

This method enabled us to collect independent observations of each individual
across time and for the dataset to encompass changes in the fish’s morphology.
At each capture, a few images were taken of the fish on both sides and the
images were tagged with an id based on the RFID. The images are captured with
the dorsal side of the fish facing up. After some filtration, a dataset that
could be used for the task was compiled. The final dataset consists of 2113
images from 513 unique individuals. As an added statistic, the mean between
the first and last capture-date of all the individuals is 230 days. Samples
from the dataset can be seen in Figure 1.

Figure 1: Samples from the unprocessed dataset.

###  3.2 Individual re-identification

The re-identification system consists of a pipeline of different components,
as illustrated in Figure 2. The components fall into two categories, a
preprocessing part and a re-identification part. As part of a preprocessing
step in the pipeline, the system takes an image as input and feeds it to a
object detection network to get an image crop, only containing the fish in the
frame. Then a different network, the direction component, classifies whether
the fish is facing right or left and passes this as metadata. For the re-
identification part, the preprocessed data is fed to a contrastive learning
network that learns to group embeddings for the same individual together and
different apart. Classification can then be performed on the embeddings. By
storing the embeddings of all previously observed individuals, re-
identification can be achieved by nearest neighbor methods.

Object detector Direction classifier Embedding ClassificationPreprocessing Re-
identification Figure 2: The network pipeline takes an image of a fish as
input and outputs the id of the individual.

The object detector uses YOLOv5 [11] with an image size of 416x416, a batch
size of 32 and is trained for 50 epochs. During training, the network was
provided with manually annoted bounding boxes enclosing the fish.

The direction network is an Inception v3 [20] model with all its weights
frozen. A global average pooling layer, a ReLU activated layer with 32
neurons, and a sigmoid activated output have been appended to the network. The
dataset used for the training is the images cropped to only contain the head.
The dataset is manually annotated with the direction.

The embedding network consists of a CNN model with a projection head. Its
constituent parts were found experimentally. The CNN model is an Inception v3
model pre-trained on ImageNet, with the layers after the fourth concatenation
layer (layer 46, or 132 if counting activation layers) removed. Appended at
the end is a 2D global average pooling layer and a 128-dimensional linear
projection that is normalized to the unit hypersphere. The network diagram is
shown in Figure 3. The input size of the network is 224, and the images are
resized accordingly before being fed into it. The network utilizes letter-
boxing to maintain aspect-ratio. For the training of the embedding network,
the dataset is split into a training and test set with a test set fraction of
0.3.

3x ConvolutionMaxPoolAvgPoolConcatenationGlobalAvg2DDenseL2 Normalization
Figure 3: The embedding network utilizes the first part of Inception v3 [20]
with a custom projection head. The dashed line marks where the Inception part
ends and the custom part starts. The grey part is repeated three times.

The training of the embedding network uses gradual unfreezing. The first 100
epochs have the layers before layer 29 frozen and a learning rate of 0.001,
and the next 100 epochs have the layers before layer 18 frozen and a learning
rate of 0.0001 for a total of 200 epochs. Layer 29 and layer 18 were selected
because they are concatenation bottlenecks in the network architecture (green
nodes in Figure 3). The loss function is online hard triplet mining with a
margin of 1.0. Hard triplet mining is a technique where loss is only
backpropagated for triplets where the negative is closer to the anchor than
the positive. Thus, the use of online mining circumvents the need for three
identical networks with shared weights. The training samples are randomly
applied with image augmentations. A number between -20 and 20 is added to the
hue and saturation. The image is rotated by a fraction between 0 and 0.1 in
either direction, and a scale transformation between 0 and 0.1 is applied. The
batch size used is 32.

Classification, and by extension re-identification, is done using a nearest
neighbor approach. And in this case it is useful to define the training set as
the support set and the test set as the query set. Nearest neighbor
classification is non-parametric and does not need to be trained through
optimization. The training step is simply to feed the support images through
the embedding network and store the associated embeddings for the inference
step. To classify an image, a query image is fed through the embedding network
and then simply select the class of the nearest point of the query embedding
to the support set embeddings. Source code for our implementation is available
at GitHub111https://github.com/orilan93/SiameseFish.

###  3.3 Method for experiments

The Symphodus melops have a distinct high-contrast pattern in the head region
(particularly on the operculum). For this reason, it would be useful to
explore whether the network performs better on head crops than on crops of the
whole body. The experiment is performed by training and evaluating the
embedding network on images that are cropped to either part.

The system can also treat each side of the fish as different classes, and thus
valuable information can be gained by doing inference on both, and then
combining the results in an ensemble classifier. For this experiment, the
dataset is split up into a left-sided section and a right-sided section such
that there is a pairwise correspondance between the images. Two models are
trained, where one is only for left-sided images and the other is only for
right-sided images. The embeddings in the support set is sorted by the
distance to the query image for each side. The predicted class is then the
class which appears first when both sorted collections are taken into account.

An experiment to evaluate how well the system is able to distinguish between a
re-sighted individual and an individual that has never been seen before was
also conducted. A query embedding is considered a new individual if its
distance is greater than a certain distance away from any support embedding.
The query set was split into a test set and a validation set. A grid search
was used to find a good distance threshold by maximizing the F1 score when
evaluating the test set. The validation dataset for this experiment contains
317 samples.

##  4 Results

The metrics we use are accuracy@1, accuracy@5, and mAP@5. Accuracy@1 shows the
correctness of the highest ranked category, i.e., the percentages of the
highest predicted class are equal to the true class. Accuracy@5 shows the
correctness of the five highest ranked categories, i.e., how many of the five
highest classes contain the true class. mAP@5 similarly shows the precision of
the five highest ranked categories, i.e., how many of the true categories are
among the five highest ranked categories.

###  4.1 Re-identification

The re-identifcation system was evaluated against both the head and body crop
datasets. Table 1 presents results from accuracy@1 and accuracy@5 and shows
that the model performs best on the head crops. Figure 4 shows how the model
performs as the number of accumulated attempts increase. This approach is
essential in practice because, instead of having an unsorted catalog of images
to go through, a professional biologist can go through a sorted catalog and
expect to find the correct individual after inspecting the $k$ most promising
images sorted based on the distance measure. The larger $k$ the higher
accuracy, and as the number of attempts are approaching the number of images
in the support set, the accuracy is approaching 100%.

Table 1: Results for re-identification on head and body crops. Type | Accuracy@1 | Accuracy@5 | mAP@5
---|---|---|---
Head | 0.3534 | 0.5647 | 0.4227
Body | 0.2043 | 0.3892 | 0.2690
Figure 4: The number of accumulated attempts (k) needed to attain a certain
accuracy (accuracy@k).

Table 2 shows four random image samples from the dataset, together with the
image the trained model predicts is the same individual and the ground truth.
The classification rank and the distance in the embedding space are also
shown.

Table 2: The retrieval rank and euclidean distance between the embedding of a query image and a correct image. Query |  Predicted |  Ground truth |  Rank |  Distance
---|---|---|---|---
|  |  |  1 |  0.65
|  |  |  217 |  1.27
|  |  |  1 |  0.61
|  |  |  1012 |  1.46

To gain insight into what the model focuses on when making its inferences, we
present some test set samples and the accompanying SHAP plot [15] in Figure 5.
The colored area shows that the model is indeed picking up on the pattern of
the fish.

Figure 5: SHAP plot showing which areas in the images that are most
influential for the decisions of the model.

###  4.2 Ensemble classifier

This experiment shows the results of training a new model for each side of the
fish and then combining their respective classifications. Table 3 shows that
this strategy can significantly increase performance. Note that the direction
component, that is required for the ensemble classifier, yielded an accuracy@1
of 99.38% on the validation set using the head cropped dataset.

Table 3: Ensemble classifier results. Type | Accuracy@1 | Accuracy@5 | mAP@5
---|---|---|---
Left | 0.3568 | 0.5463 | 0.4243
Right | 0.4097 | 0.5595 | 0.4623
As pair | 0.5286 | 0.7533 | 0.6140

###  4.3 New observations

As our previous experiments have shown, re-identification works relatively
well. We aim at using this model for distinguishing new individuals from
earlier observed individuals. To identify new individuals with the model, an
embedding distance threshold needs to be decided. Note that this relates to
the distance metric in Table 2. Using grid search, we found a threshold of
0.820 to yield the best performance score on the validation set. The system
predicted 95 individuals as new sightings and got a 62.78% accuracy@1 at this
task.

##  5 Discussion and conclusion

Table 4: Summary of results. Experiment | Result | Metric
---|---|---
Re-identification | 0.3534 | Accuracy@1
Object detector | 0.9951 |<EMAIL_ADDRESS>
Direction classifier | 0.9937 | Accuracy@1
Ensemble classifier | 0.5286 | Accuracy@1
New observations | 0.6278 | Accuracy@1

Our experiments, summarized in Table 4, indicate that the system performs
better on the head crops of the fish than on the whole body. This is likely
because the pattern on the head is most distinct and thus an important
feature, and this will appear at a higher resolution for the algorithm when
resizing for the network input size. However, the drawback here is that the
network is exposed to less information available in the data.

By utilizing the existing system in a new way by training separate models for
each side of the fish, one can make an ensemble classifier. This method was
tested and gained a considerable improvement from 35% to 53% accuracy. This
shows how important it is to use all the information available to make good
predictions.

The accuracy of this system is not high enough for a fully automated system
with humans out-of-the-loop, which is required to replace the need for
physical tags in ecological studies. However, we believe that continued
collection of data can produce a dataset that is more temporally balanced to
enable the model to account for the growth and ageing of the individuals.

Automatization can produce great benefits and is increasingly being adopted by
many industries, and the field of ecology should be no different. A successful
Re-ID algorithm with high precision can provide a new method with improved
fish welfare, while also being cheaper (only a camera needed) and potentially
more accurate (no tag loss). In the future, we envision that re-ID can be
applied directly on live streams from under-water video cameras, removing the
need for capture and handling fish altogether. This would be a revolutionary
method that can drastically change how we can collect key information for
sustainable conservation and management of fish and other animals.

## Acknowledgements

We thank Torkel Larsen, Anne Berit Skiftesvik, Ovin Holm, Ylva Vik, Nicolai
Aasen, Ben Ellis, Vegard Omestad Berntsen, and Steve Shema for assistance in
collecting the photos and capture-recapture data in the field. This study
received funding from Centre for Artificial Intelligence Research (CAIR),
Centre for Coastal Research (CCR), and Top Research Centre Mechatronics (TRCM)
at University of Agder, the Institute of Marine Research (project 15638-01),
and the Research Council of Norway (CoastVision, project number 325862, and
CreateView, project number 309784).

## References

  * Bromley et al. [1993] J. Bromley, J. W. Bentz, L. Bottou, I. Guyon, Y. LeCun, C. Moore, E. Säckinger, and R. Shah.  Signature verification using a “siamese” time delay neural network.  _International Journal of Pattern Recognition and Artificial Intelligence_ , 7(04):669–688, 1993.
  * Bruslund Haurum et al. [2020] J. Bruslund Haurum, A. Karpova, M. Pedersen, S. Hein Bengtson, and T. B. Moeslund.  Re-identification of zebrafish using metric learning.  In _Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision Workshops_ , pages 1–11, 2020.  10.1109/WACVW50321.2020.9096922.
  * Ditria et al. [2021] E. M. Ditria, E. L. Jinks, and R. M. Connolly.  Automating the analysis of fish grazing behaviour from videos using image classification and optical flow.  _Animal Behaviour_ , 177:31–37, jul 2021.  ISSN 00033472.  10.1016/j.anbehav.2021.04.018.
  * Goodwin et al. [2022] M. Goodwin, K. T. Halvorsen, L. Jiao, K. M. Knausgård, A. H. Martin, M. Moyano, R. A. Oomen, J. H. Rasmussen, T. K. Sørdalen, and S. H. Thorbjørnsen.  Unlocking the potential of deep learning for marine ecology: overview, applications, and outlook.  _ICES Journal of Marine Science_ , 01 2022.  ISSN 1054-3139.  10.1093/icesjms/fsab255.
  * Gupta et al. [2022a] A. Gupta, E. Bringsdal, N. Salbuvik, K. M. Knausgård, and M. Goodwin.  An accurate convolutional neural networks approach to wound detection for farmed salmon.  In _International Conference on Engineering Applications of Neural Networks_ , pages 139–149. Springer, 2022a.  10.1007/978-3-031-08223-8_12.
  * Gupta et al. [2022b] A. Gupta, E. S. Kalhagen, Ø. L. Olsen, and M. Goodwin.  Hierarchical object detection applied to fish species: Hierarchical object detection of fish species.  _Nordic Machine Intelligence_ , 2(1), 2022b.  10.5617/nmi.9452.
  * Halvorsen et al. [2016] K. T. Halvorsen, T. K. Sørdalen, C. Durif, H. Knutsen, E. M. Olsen, A. B. Skiftesvik, T. E. Rustand, R. M. Bjelland, and L. A. Vøllestad.  Male-biased sexual size dimorphism in the nest building corkwing wrasse (Symphodus melops): Implications for a size regulated fishery.  _ICES Journal of Marine Science_ , 73(10):2586–2594, nov 2016.  ISSN 10959289.  10.1093/icesjms/fsw135.
  * Halvorsen et al. [2017] K. T. Halvorsen, T. K. Sørdalen, L. A. Vøllestad, A. B. Skiftesvik, S. H. Espeland, and E. M. Olsen.  Sex- and size-selective harvesting of corkwing wrasse (Symphodus melops)—a cleaner fish used in salmonid aquaculture.  _ICES Journal of Marine Science_ , 74(3):660–669, mar 2017.  ISSN 1054-3139.  10.1093/icesjms/fsw221.
  * Halvorsen et al. [2021] K. T. Halvorsen, T. Larsen, H. I. Browman, C. Durif, N. Aasen, L. A. Vøllestad, A. Cresci, T. K. Sørdalen, R. M. Bjelland, and A. B. Skiftesvik.  Movement patterns of temperate wrasses ( Labridae ) within a small marine protected area.  _Journal of Fish Biology_ , 99(4):1513–1518, jul 2021.  ISSN 0022-1112.  10.1111/jfb.14825.
  * Hoffer and Ailon [2015] E. Hoffer and N. Ailon.  Deep metric learning using triplet network.  In _International workshop on similarity-based pattern recognition_ , pages 84–92. Springer, 2015.  10.1007/978-3-319-24261-3_7.
  * Jocher et al. [2021] G. Jocher, A. Stoken, A. Chaurasia, J. Borovec, et al.  ultralytics/yolov5: v6.0 - YOLOv5n ’Nano’ models, Roboflow integration, TensorFlow export, OpenCV DNN support, Oct. 2021.  URL https://doi.org/10.5281/zenodo.5563715.
  * Knausgård et al. [2022] K. M. Knausgård, A. Wiklund, T. K. Sørdalen, K. T. Halvorsen, A. R. Kleiven, L. Jiao, and M. Goodwin.  Temperate fish detection and classification: a deep learning based approach.  _Applied Intelligence_ , 52(6):6988–7001, 2022\.  10.1007/s10489-020-02154-9.
  * Li et al. [2022] D. Li, H. Su, K. Jiang, D. Liu, and X. Duan.  Fish face identification based on rotated object detection: Dataset and exploration.  _Fishes_ , 7(5):219, 2022.  10.3390/fishes7050219.
  * Lopez-Vazquez et al. [2020] V. Lopez-Vazquez, J. M. Lopez-Guede, S. Marini, E. Fanelli, E. Johnsen, and J. Aguzzi.  Video image enhancement and machine learning pipeline for underwater animal detection and classification at cabled observatories.  _Sensors_ , 20(3):726, 2020.  10.3390/s20030726.
  * Lundberg and Lee [2017] S. M. Lundberg and S.-I. Lee.  A unified approach to interpreting model predictions.  _Advances in neural information processing systems_ , 30, 2017.
  * Meidell and Sjøblom [2019] E. Meidell and E. S. Sjøblom.  Fishnet: A unified embedding for salmon recognition.  Master’s thesis, NTNU, 2019.
  * Moskvyak et al. [2019] O. Moskvyak, F. Maire, F. Dayoub, A. O. Armstrong, and M. Baktashmotlagh.  Robust re-identification of manta rays from natural markings by learning pose invariant embeddings.  In _2021 Digital Image Computing: Techniques and Applications (DICTA)_ , pages 1–8. IEEE, 2019.  10.1109/DICTA52665.2021.9647359.
  * Schneider et al. [2019] S. Schneider, G. W. Taylor, S. Linquist, and S. C. Kremer.  Past, present and future approaches using computer vision for animal re-identification from camera trap data.  _Methods in Ecology and Evolution_ , 10(4):461–470, 2019.  10.1111/2041-210X.13133.
  * Schneider et al. [2020] S. Schneider, G. W. Taylor, and S. C. Kremer.  Similarity learning networks for animal individual re-identification-beyond the capabilities of a human observer.  In _Proceedings of the IEEE/CVF winter conference on applications of computer vision workshops_ , pages 44–52, 2020.
  * Szegedy et al. [2016] C. Szegedy, V. Vanhoucke, S. Ioffe, J. Shlens, and Z. Wojna.  Rethinking the inception architecture for computer vision.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , June 2016.
  * Uglem et al. [2000] I. Uglem, G. Rosenqvist, and H. S. Wasslavik.  Phenotypic variation between dimorphic males in corkwing wrasse.  _Journal of Fish Biology_ , 57(1):1–14, jul 2000\.  ISSN 00221112.  10.1006/jfbi.2000.1283.
  * Weinstein [2018] B. G. Weinstein.  A computer vision for animal ecology.  _Journal of Animal Ecology_ , 87(3):533–545, 2018\.  10.1111/1365-2656.12780.