# Particle Cloud Generation with Message Passing Generative Adversarial
Networks

Raghav Kansal, Javier Duarte, Hao Su
University of California San Diego
La Jolla, CA 92093, USA
&Breno Orzari, Thiago Tomei
Universidade Estadual Paulista
São Paulo/SP - CEP 01049-010, Brazil &Maurizio Pierini, Mary Touranakou
European Organization for Nuclear Research (CERN)
CH-1211 Geneva 23, Switzerland &Jean-Roch Vlimant
California Institute of Technology
Pasadena, CA 91125, USA &Dimitrios Gunopulos
National and Kapodistrian University of Athens
Athens 15772, Greece Also at National and Kapodistrian University of Athens,
Athens, Greece.

###### Abstract

In high energy physics (HEP), jets are collections of correlated particles
produced ubiquitously in particle collisions such as those at the CERN Large
Hadron Collider (LHC). Machine learning (ML)-based generative models, such as
generative adversarial networks (GANs), have the potential to significantly
accelerate LHC jet simulations. However, despite jets having a natural
representation as a set of particles in momentum-space, a.k.a. a particle
cloud, there exist no generative models applied to such a dataset. In this
work, we introduce a new particle cloud dataset (JetNet), and apply to it
existing point cloud GANs. Results are evaluated using (1) 1-Wasserstein
distances between high- and low-level feature distributions, (2) a newly
developed Fréchet ParticleNet Distance, and (3) the coverage and (4) minimum
matching distance metrics. Existing GANs are found to be inadequate for
physics applications, hence we develop a new message passing GAN (MPGAN),
which outperforms existing point cloud GANs on virtually every metric and
shows promise for use in HEP. We propose JetNet as a novel point-cloud-style
dataset for the ML community to experiment with, and set MPGAN as a benchmark
to improve upon for future generative models. Additionally, to facilitate
research and improve accessibility and reproducibility in this area, we
release the open-source JetNet Python package with interfaces for particle
cloud datasets, implementations for evaluation and loss metrics, and more
tools for ML in HEP development.

##  1 Introduction

Over the past decade, machine learning (ML) has become the de facto way to
analyze jets, collimated high-energy sprays of particles [1] produced at the
CERN Large Hadron Collider (LHC). To apply ML to jets, the most natural
representation is a _particle cloud_ , a variable-sized set of points in
momentum space, whose radiation pattern contains rich information about the
underlying physics known as quantum chromodynamics (QCD). A fundamental
question is whether ML algorithms can model this underlying physics and
successfully reproduce the rich high- and low-level structure in jets.

Answering this question affirmatively has important practical applications. At
the LHC, large simulated data samples of collision events111An event is a set
of observed particles with well-defined momenta, but not generally well-
defined positions. are generated using Monte Carlo (MC) methods in order to
translate theoretical predictions into observable distributions, and
ultimately perform physics analyses222See App. A for further details on the
downstream applications of simulations.. These samples, numbering in the
billions of events, require computationally expensive modeling of the
interaction of particles traversing the detector material. Recently developed
generative frameworks in ML such as generative adversarial networks (GANs), if
accurate enough, can be used to accelerate this simulation by potentially five
orders of magnitude [2].

In this work, we advocate for a benchmark jet dataset (JetNet) and propose
several physics- and computer-vision-inspired metrics with which the ML
community can improve and evaluate generative models in high energy physics
(HEP). To facilitate and encourage research in this area, as well as to make
such research more accessible and reproducible, we release interfaces for
public particle cloud datasets such as JetNet, implementations for our
proposed metrics, and various tools for ML in HEP development in the JetNet
library [3]. We next apply existing point cloud GANs on JetNet and find the
results to be inadequate for physics applications. Finally, we develop our own
message passing GAN (MPGAN), which dramatically improves results on virtually
every metric, and propose it as a benchmark on JetNet.

##  2 Jets

High-energy proton-proton collisions at the LHC produce elementary particles
like quarks and gluons, which cannot be isolated due to the QCD property of
color confinement [4]. These particles continuously radiate or “split” into a
set of particles, known as a parton shower. Eventually they cool to an energy
at which they undergo the process of hadronization, where the fundamental
particles combine to form more stable hadrons, such as pions and protons. The
final set of collimated hadrons produced after such a process is referred to
as a jet.

The task of simulating a single jet can be algorithmically defined as
inputting an initial particle, which produces the jet, and outputting the
final set of particles a.k.a. the jet constituents. Typically in HEP the
parton shower and hadronization are steps that are simulated sequentially
using MC event generators such as pythia [5] or herwig [6]. Simulating either
process exactly is not possible because of the complex underlying physics
(QCD), and instead these event generators fit simplified physics-inspired
stochastic models, such as the Lund string model for hadronization [7], to
existing data using MC methods. The present work can be seen as an extension
of this idea, using a simpler, ML-based model, also fitted to data, for
generating the jet in one shot, where we are effectively trading the
interpretability of the MC methods for the speed of GPU-accelerated ML
generators.

#### Representations.

As common for collider physics, we use a Cartesian coordinate system with the
$z$ axis oriented along the beam axis, the $x$ axis on the horizontal plane,
and the $y$ axis oriented upward. The $x$ and $y$ axes define the transverse
plane, while the $z$ axis identifies the longitudinal direction. The azimuthal
angle $\phi$ is computed with respect to the $x$ axis. The polar angle
$\theta$ is used to compute the pseudorapidity $\eta=-\log(\tan(\theta/2))$.
The transverse momentum ($p_{\mathrm{T}}$) is the projection of the particle
momentum on the ($x$, $y$) plane. As is customary, we transform the particle
momenta from Cartesian coordinates $(p_{x},p_{y},p_{z})$ to longitudinal-
boost-invariant pseudo-angular coordinates $(p_{\mathrm{T}},\eta,\phi)$, as
shown in Fig. 1.

Figure 1: The collider physics coordinate system defining
$(p_{\mathrm{T}},\eta,\phi)$ (left). The three jet classes in our dataset
(right). Gluon (g) and light quark (q) jets have simple topologies, with q
jets generally containing fewer particles. Top quark (t) jets have a complex
three-pronged structure. Shown also are the relative angular coordinates
$\eta^{\mathrm{rel}}$ and $\phi^{\mathrm{rel}}$, measured from the jet axis.

In the context of ML, jets can be represented in multiple ways. One popular
representation is as images [8, 9], created by projecting each jet’s particle
constituents onto a discretized angular $\eta$-$\phi$ plane, and taking the
intensity of each “pixel” in this grid to be a monotonically increasing
function of the corresponding particle $p_{\mathrm{T}}$. These tend to be
extremely sparse, with typically fewer than 10% of pixels nonempty [10], and
the discretization process can furthemore lower the resolution.

Two more spatially efficient representations are as ordered lists or unordered
sets [11, 12] of the jet constituents and their features. The difficulty with
the former is that there is no particular preferred ordering of the
particles—one would have to impose an arbitrary ordering such as by transverse
momentum [13]. The more natural representation is the unordered set of
particles in momentum space, which we refer to as a “particle cloud.” This is
in analogy to point cloud representations of 3D objects in position-space
prevalent in computer vision created, for example, by sampling from 3D
ShapeNet models [14].

Apart from how the samples are produced, significant differences between jets
and ShapeNet-based point clouds are that, firstly, jets have physically
meaningful low- and high-level features such as particle momentum, total mass
of the jet, the number of sub-jets, and $n$-particle energy correlations.
These physical observables are how we characterize jets, and hence are
important to reproduce correctly for physics analysis applications. Secondly,
unlike the conditional distributions of points given a particular ShapeNet
object, which are identical and independent, particle distributions within
jets are highly correlated, as the particles each originate from a single
source. The independence of their constituents also means that ShapeNet-
sampled point clouds can be chosen to be of a fixed cardinality, whereas this
is not possible for jets, which inherently contain varying numbers of
particles due to the stochastic nature of particle production.

#### JetNet.

We publish JetNet [15] under the CC-BY 4.0 license, to facilitate and advance
ML research in HEP, and to offer a new point-cloud-style dataset to experiment
with. Derived from Ref. [16]333This dataset was also released under the CC-BY
4.0 license., it consists of simulated particle jets with transverse momenta
$p_{\mathrm{T}}^{\mathrm{jet}}\approx 1\,\text{TeV}$, originating from gluons,
light quarks, and top quarks produced in $13\,\text{TeV}$ proton-proton
collisions in a simplified detector. Technical details of the generation
process are given in App. B. We limit the number of constituents to the 30
highest $p_{\mathrm{T}}$ particles per jet, allowing for jets with potentially
fewer than 30 by zero-padding. For each particle we provide the following four
features: the relative angular coordinates
$\eta^{\mathrm{rel}}=\eta^{\mathrm{particle}}-\eta^{\mathrm{jet}}$ and
$\phi^{\mathrm{rel}}=\phi^{\mathrm{particle}}-\phi^{\mathrm{jet}}\pmod{2\pi}$,
relative transverse momentum
$p_{\mathrm{T}}^{\mathrm{rel}}=p_{\mathrm{T}}^{\mathrm{particle}}/p_{\mathrm{T}}^{\mathrm{jet}}$,
and a binary mask feature classifying the particle as genuine or zero-padded.

We choose three jet classes, depicted in Fig. 1, to individually target the
unique and challenging properties of jets. Gluons provide a useful baseline
test, as they typically radiate into a large number of particles before
hadronization, largely avoiding the variable-sized cloud issue—at least with a
30 particle maximum, and have a relatively simple topology. Light quarks share
the simple topology, but produce fewer final-state particles, resulting in a
larger fraction of zero-padded particles in the dataset. They allow evaluation
of a model’s ability to handle variable-sized clouds. Finally, top quarks
decay into three lighter quarks through an intermediate particle, the W boson,
which each may produce their own sub-jets, leading to a complex two- or three-
pronged topology—depending on whether the jet clustering algorithm captures
all three or just two of these sub-jets. This results in bimodal jet feature
distributions (one peak corresponding to fully merged top quark jets and the
other to semi-merged, as seen in Fig. 3). Thus, top quark jets test models’
ability to learn the rich global structure and clustering history of a
particle cloud.

##  3 Related Work

#### Generative models in HEP.

Past work in this area has exclusively used image-based representations for
HEP data. One benefit of this is the ability to employ convolutional neural
network (CNN) based generative models, which have been highly successful on
computer vision tasks. Refs. [2, 17, 18, 19, 20], for example, build upon CNN-
based GANs, and Ref. [21] uses an auto-regressive model, to output jet- and
detector-data-images.

In addition to the issues with such representations outlined in Sec. 2, the
high sparsity of the images can lead to training difficulties in GANs, and the
irregular geometry of the data — a single LHC detector can typically have
multiple sections with differing pixel sizes and shapes — poses a challenge
for CNN GANs which output uniform matrices. While these can be mitigated to an
extent with techniques such as batch normalization [22] and using larger/more
regular pixels [18], our approach avoids both issues by generating particle-
cloud-representations of the data, as these are inherently sparse data
structures and are completely flexible to the underlying geometry.

#### GANs for point clouds.

There are several published generative models in this area, however the
majority exploit inductive biases specific to their respective datasets, such
as ShapeNet-based [23, 24, 25, 26] and molecular [27, 28, 29] point clouds,
which are not appropriate for jets. A more detailed discussion, including some
experimental results, can be found in App. C.

There do exist some more general-purpose GAN models, namely r-GAN [30],
GraphCNN-GAN [31], and TreeGAN [32], and we test these on JetNet. r-GAN uses a
fully-connected (FC) network, GraphCNN-GAN uses graph convolutions based on
dynamic $k$-nn graphs in intermediate feature spaces, and TreeGAN iteratively
up-samples the graphs with information passing from ancestor to descendant
nodes. In terms of discriminators, past work has used either a FC or a
PointNet [33]-style network. Ref. [34] is the first work to study point cloud
discriminator design in detail and finds amongst a number of PointNet and
graph convolutional models that PointNet-Mix, which uses both max- and
average-pooled features, is the most performant.

We apply the three aforementioned generators and FC and PointNet-Mix
discriminators as baselines to our dataset, but find jet structure is not
adequately reproduced. GraphCNN’s local convolutions make learning global
structure difficult, and while the TreeGAN and FC generator + PointNet
discriminator combinations are improvements, they are not able to learn multi-
particle correlations, particularly for the complex top quark jets, nor deal
with the variable-sized light quark jets to the extent necessary for physics
applications.

#### Message Passing Neural Networks.

We attempt to overcome limitations of existing GANs by designing a novel
generator and discriminator which can learn such correlations and handle
variable-sized particle clouds. Both networks build upon the generic message-
passing neural network (MPNN) [35] framework with physics-conscious design
choices, and collectively we refer to them as message-passing GAN (MPGAN). We
find MPGAN outperforms existing models on virtually all evaluation metrics.

###  3.1 Evaluating generative models.

Evaluating generative models is a difficult task, however there has been
extensive work in this area in both the physics and computer-vision
communities.

#### Physics-inspired metrics.

An accurate jet simulation algorithm should reproduce both low-level and high-
level features (such as those described in Sec. 2), hence a standard method of
validating generative models, which we employ, is to compare the distributions
of such features between the real and generated samples444More details on the
downstream validation procedure used in HEP analyses can be found in App. A.
[2, 17, 18, 19, 20, 36].

For application in HEP, a generative model needs to produce jets with physical
features indistinguishable from real. Therefore, we propose the validation
criteria that differences between real and generated sample features may not
exceed those between sets of randomly chosen real samples. To verify this, we
use bootstrapping to compare between random samples of only real jets as a
baseline.

A practically useful set of features to validate against are the so-called
“energy-flow polynomials” (EFPs) [37], which are a set of multi-particle
correlation functions. Importantly, the set of all EFPs forms a linear basis
for all useful jet-level features555In the context of HEP, this means all
infrared- and colinear-safe observables.. Therefore, we claim that if we
observe all EFP distributions to be reproduced with high fidelity and to match
the above criteria, we can conclude with strong confidence that our model is
outputting accurate particle clouds.

#### Computer-vision-inspired metrics

A popular metric for evaluating images which has shown to be sensitive to
output quality and mode-collapse, though it has its limitations [38], is the
Fréchet Inception Distance [39] (FID). FID is defined as the Fréchet distance
between Gaussian distributions fitted to the activations of a fully-connected
layer of the Inception-v3 image classifier in response to real and generated
samples. We develop a particle-cloud-analogue of this metric, which we call
Fréchet ParticleNet Distance (FPND), using the state-of-the-art (SOTA)
ParticleNet graph convolutional jet classifier [10] in lieu of the Inception
network. We note that the FPND and comparing distributions as above is
conceptually equivalent, except here instead of physically meaningful and
easily interpretable features, we are comparing those found to be
statistically optimum for distinguishing jets.

Two common metrics for evaluating point cloud generators are coverage (COV)
and minimum matching distance (MMD) [30]. Both involve finding the closest
point cloud in a sample $X$ to each cloud in another sample $Y$, based on a
metric such as the Chamfer distance or the earth mover’s distance. Coverage is
defined as the fraction of samples in $X$ which were matched to one in $Y$,
measuring thus the diversity of the samples in $Y$ relative to $X$, and MMD is
the average distance between matched samples, measuring the quality of
samples. We use both, and due to drawbacks of the Chamfer distance pointed out
in Ref. [30], for our distance metric choose only the analogue of the earth
mover’s distance for particle clouds a.k.a. the energy mover’s distance (EMD)
[40]. We discuss the effectiveness and complementarity of all four metrics in
evaluating clouds in Sec. 5.

##  4 MPGAN Architecture

We describe now the architecture of our MPGAN model (Fig. 2), noting particle
cloud-motivated aspects compared to its r-GAN and GraphCNN-GAN predecessors.

Figure 2: Top: The MP generator uses message passing to generate a particle
cloud. In blue is the initial latent vector and FC layer part of the MP-LFC
variant. Bottom: The MP discriminator uses message passing to classify an
input particle cloud as real or generated.

#### Message passing.

Jets originate from a single source particle decaying and hadronizing, hence
they end up with important high-level jet features and a rich global
structure, known as the jet substructure [1], stemming from the input
particle. Indeed any high-level feature useful for analyzing jets, such as jet
mass or multi-particle correlations, is necessarily global [37]. Because of
this, while past work in learning on point clouds [41, 42, 10], including
GraphCNN-GAN, has used a locally connected graph structure and convolutions
for message passing, we choose a fully connected graph, equally weighting
messages from all particles in the clouds. Rather than subtracting particle
features for messages between particles, useful in graph convolutions to
capture local differences within a neighborhood, the respective features are
concatenated to preserve the global structure (the difference between particle
features is also only physically meaningful if they are in the 4-vector
representation of the Lorentz group). During the update step in the message
passing we find it empirically beneficial to incorporate a residual connection
to previous particle features.

The operation can be described as follows. For an $N$-particle cloud
$J^{t}=\\{p_{1}^{t},\cdots,p_{N}^{t}\\}$ after $t$ iterations of message
passing, with $t=0$ corresponding to the original input cloud, each particle
$p^{t}_{i}$ is represented by features $\mathbf{h}^{t}_{i}$. One iteration of
message passing is then defined as

$\displaystyle\mathbf{m}^{t+1}_{ij}$
$\displaystyle=f_{e}^{t+1}(\mathbf{h}^{t}_{i}\oplus\mathbf{h}^{t}_{j}),$ (1)
$\displaystyle\mathbf{h}^{t+1}_{i}$
$\displaystyle=f_{n}^{t+1}(\mathbf{h}^{t}_{i}\oplus\sum_{j\in
J}\mathbf{m}^{t+1}_{ij})\,,$ (2)

where $\mathbf{m}^{t+1}_{ij}$ is the message vector sent from particle $j$ to
particle $i$, $\mathbf{h}^{t+1}_{i}$ are the updated features of particle $i$,
and $f_{e}^{t+1}$ and $f_{n}^{t+1}$ are arbitrary functions which, in our
case, are implemented as multilayer perceptrons (MLPs) with 3 FC layers.

#### Generator.

We test two initializations of a particle cloud for the MPGAN generator: (1)
directly initializing the cloud with $N$ particles with $L$ randomly sampled
features, which we refer to as the MP generator, and (2) inputting a single
$Z$-dimensional latent noise vector and transforming it via an FC layer into
an $N\times L$-dimensional matrix, which we refer to as the MP-Latent-FC (MP-
LFC) generator. The MP-LFC uses a latent space which can intuitively be
understood as representing the initial source particle’s features along with
parameters to capture the stochasticity of the jet production process. Due to
the complex nature of this process, however, we posit that this global,
flattened latent space cannot capture the full phase space of individual
particle features. Hence, we introduce the MP generator, which samples noise
directly per particle, and find that it outperforms MP-LFC (Table 5).

#### Discriminator.

We find the MP generator, in conjunction with a PointNet discriminator, to be
a significant improvement on every metric compared to FC and GraphCNN
generators. However, the jet-level features are not yet reproduced to a high
enough accuracy (Sec. 5). While PointNet is able to capture global structural
information, it can miss the complex interparticle correlations in real
particle clouds. We find we can overcome this limitation by incorporating
message passing in the discriminator as well as in the generator. Concretely,
our MP discriminator receives the real or generated cloud and applies MP
layers to produce intermediate features for each particle, which are then
aggregated via a feature-wise average-pooling operation and passed through an
FC layer to output the final scalar feature. We choose 2 MP layers for both
networks.

#### Variable-sized clouds.

In order to handle clouds with varying numbers of particles, as typical of
jets, we introduce an additional binary “masking” particle feature classifying
the particle as genuine or zero-padded. Particles in the zero-padded class are
ignored entirely in the message passing and pooling operations. The MP
generator adds mask features to the initial particle cloud, using an
additional input of the size of the jet $N$, sampled from the real
distribution, before the message passing layers based on sorting in particle
feature space. Ablation studies with alternative (as well as no) masking
strategies are discussed in App. E.

##  5 Experiments

#### Evaluation.

$W_{1}$ distances between real jet mass ($W_{1}^{\mathrm{M}}$), averaged
particle features ($W_{1}^{\mathrm{P}}$), and averaged jet EFPs
($W_{1}^{\mathrm{EFP}}$) distributions calculated as a baseline, for three
classes of jets.  Jet class $W_{1}^{\mathrm{M}}$ ($\times 10^{-3}$)
$W_{1}^{\mathrm{P}}$ ($\times 10^{-3}$) $W_{1}^{\mathrm{EFP}}$ ($\times
10^{-5}$) Gluon $0.7\pm 0.2$ $0.44\pm 0.09$ $0.62\pm 0.07$ Light quark $0.5\pm
0.1$ $0.5\pm 0.1$ $0.46\pm 0.04$ Top quark $0.51\pm 0.07$ $0.55\pm 0.07$
$1.1\pm 0.1$

We use four techniques discussed in Sec. 3.1 for evaluating and comparing
models. Distributions of physical particle and jet features are compared
visually and quantitatively using the Wasserstein-1 ($W_{1}$) distance between
them. For ease of evaluation, we report (1) the average scores of the three
particle features ($W_{1}^{\mathrm{P}}$) $\eta^{\mathrm{rel}}$,
$\phi^{\mathrm{rel}}$, and $p_{\mathrm{T}}^{\mathrm{rel}}$, (2) the jet mass
($W_{1}^{\mathrm{M}}$), and (3) the average of a subset of the EFPs666We
choose 5 EFPs corresponding to the set of loopless multigraphs with 4 vertices
and 4 edges. ($W_{1}^{\mathrm{EFP}}$), which together provide a holistic
picture of the low- and high-level aspects of a jet. The $W_{1}$ distances are
calculated for each feature between random samples of 10,000 real and
generated jets, and averaged over 5 batches. Baseline $W_{1}$ distances are
calculated between two sets of randomly sampled real jets with 10,000 samples
each, and are listed for each feature in Table 5. The real samples are split
70/30 for training/evaluation. We train ParticleNet for classification on our
dataset to develop the FPND metric. FPND is calculated between 50,000 random
real and generated samples, based on the activations of the first FC layer in
our trained model777ParticleNet training details are given in App. D.3. The
trained model is provided in the JetNet library [3].. Coverage and MMD are
calculated between 100 real and 100 generated samples, and averaged over 10
such batches. Implementations for all metrics are provided in the JetNet
package [3].

#### Results.

Figure 3: Comparison of real and generated distributions for a subset of jet
and particle features. We use the best performing model for each of the FC,
GraphCNN, TreeGAN, and MP generators, as per Table 5. Top: gluon jet features,
Middle: light quark jets, Bottom: top quark jets.

Six evaluation scores on different generator and discriminator combinations.
Lower is better for all metrics except COV.  Jet class Generator Discriminator
$W_{1}^{\mathrm{M}}$ ($\times 10^{-3}$) $W_{1}^{\mathrm{P}}$ ($\times
10^{-3}$) $W_{1}^{\mathrm{EFP}}$ ($\times 10^{-5}$) FPND COV $\uparrow$ MMD
Gluon FC FC $18.3\pm 0.2$ $9.6\pm 0.4$ $8.5\pm 0.5$ $176$ $0.24$ $0.045$
GraphCNN FC $2.6\pm 0.2$ $9.6\pm 0.3$ $12\pm 8$ $61$ $0.39$ $0.046$ TreeGAN FC
$41.9\pm 0.3$ $69.3\pm 0.3$ $14.2\pm 0.8$ $355$ $0.19$ $0.130$ FC PointNet
$1.3\pm 0.4$ $1.3\pm 0.2$ $1.5\pm 0.9$ $5.0$ $0.49$ $0.039$ GraphCNN PointNet
$1.9\pm 0.2$ $16\pm 6$ $200\pm 1000$ $7$k $0.46$ $0.040$ TreeGAN PointNet
$1.7\pm 0.1$ $4.0\pm 0.4$ $4\pm 1$ $84$ $0.37$ $0.042$ MP MP $0.7\pm 0.2$
$\mathbf{0.9\pm 0.3}$ $\mathbf{0.7\pm 0.2}$ $\mathbf{0.12}$ $\mathbf{0.56}$
$0.037$ MP-LFC MP $\mathbf{0.69\pm 0.07}$ $1.8\pm 0.2$ $0.9\pm 0.6$ $0.20$
$0.54$ $0.037$ FC MP $4.3\pm 0.3$ $21.1\pm 0.2$ $9\pm 1$ $368$ $0.11$ $0.085$
GraphCNN MP $2.5\pm 0.1$ $9.8\pm 0.2$ $13\pm 8$ $61$ $0.38$ $0.048$ TreeGAN MP
$2.4\pm 0.2$ $12\pm 7$ $18\pm 9$ $69$ $0.34$ $0.048$ MP FC $1.2\pm 0.2$
$3.7\pm 0.5$ $1.6\pm 0.8$ $39$ $0.44$ $0.040$ MP PointNet $1.3\pm 0.4$ $1.2\pm
0.4$ $4\pm 2$ $18$ $0.53$ $\mathbf{0.036}$ Light quark FC FC $6.0\pm 0.2$
$16.3\pm 0.9$ $3.9\pm 0.6$ $395$ $0.18$ $0.053$ GraphCNN FC $3.5\pm 0.2$
$15.1\pm 0.4$ $10\pm 50$ $100$ $0.25$ $0.038$ TreeGAN FC $31.5\pm 0.3$
$22.3\pm 0.4$ $9.3\pm 0.4$ $176$ $0.06$ $0.055$ FC PointNet $3.1\pm 0.2$
$4.5\pm 0.4$ $2.3\pm 0.6$ $17$ $0.37$ $0.028$ GraphCNN PointNet $4\pm 1$
$5.2\pm 0.5$ $50$k $\pm 100$k $316$ $0.37$ $0.031$ TreeGAN PointNet $10.1\pm
0.1$ $5.7\pm 0.5$ $4.1\pm 0.3$ $11$ $0.47$ $0.031$ MP MP $\mathbf{0.6\pm 0.2}$
$4.9\pm 0.5$ $\mathbf{0.7\pm 0.4}$ $0.35$ $0.50$ $0.026$ MP-LFC MP $0.7\pm
0.2$ $\mathbf{2.6\pm 0.4}$ $0.9\pm 0.9$ $\mathbf{0.08}$ $\mathbf{0.52}$
$\mathbf{0.024}$ FC MP $6.3\pm 0.2$ $16.5\pm 0.2$ $4.0\pm 0.8$ $212$ $0.11$
$0.070$ GraphCNN MP $3.5\pm 0.4$ $15.0\pm 0.3$ $10\pm 10$ $99$ $0.26$ $0.038$
TreeGAN MP $4.8\pm 0.2$ $33\pm 6$ $10\pm 2$ $148$ $0.22$ $0.041$ MP FC $1.3\pm
0.1$ $4.5\pm 0.4$ $2.2\pm 0.6$ $41$ $0.37$ $0.030$ MP PointNet $6.5\pm 0.3$
$23.2\pm 0.6$ $6\pm 1$ $850$ $0.18$ $0.034$ Top quark FC FC $4.8\pm 0.3$
$14.5\pm 0.6$ $23\pm 3$ $160$ $0.28$ $0.103$ GraphCNN FC $7.0\pm 0.3$ $8.0\pm
0.5$ $1$k $\pm 6$k $15$ $0.48$ $0.081$ TreeGAN FC $17.0\pm 0.2$ $19.6\pm 0.6$
$33\pm 2$ $77$ $0.39$ $0.083$ FC PointNet $2.7\pm 0.1$ $\mathbf{1.6\pm 0.4}$
$7.7\pm 0.5$ $3.9$ $0.56$ $0.075$ GraphCNN PointNet $11.3\pm 0.9$ $30\pm 10$
$37\pm 2$ $30$k $0.39$ $0.085$ TreeGAN PointNet $5.19\pm 0.08$ $9.1\pm 0.3$
$16\pm 2$ $17$ $0.53$ $0.079$ MP MP $\mathbf{0.6\pm 0.2}$ $2.3\pm 0.3$
$\mathbf{2\pm 1}$ $\mathbf{0.37}$ $0.57$ $\mathbf{0.071}$ MP-LFC MP $0.9\pm
0.3$ $2.2\pm 0.7$ $\mathbf{2\pm 1}$ $0.93$ $0.56$ $0.073$ FC MP $6.9\pm 0.1$
$39.1\pm 0.3$ $15\pm 1$ $81$ $0.26$ $0.120$ GraphCNN MP $6.7\pm 0.1$ $8.2\pm
0.5$ $40\pm 10$ $15$ $0.49$ $0.081$ TreeGAN MP $13.4\pm 0.4$ $45\pm 7$ $50\pm
30$ $66$ $0.29$ $0.101$ MP FC $12.9\pm 0.3$ $26.3\pm 0.4$ $46\pm 3$ $58$
$0.27$ $0.103$ MP PointNet $0.76\pm 0.08$ $\mathbf{1.6\pm 0.4}$ $4\pm 1$ $3.7$
$\mathbf{0.59}$ $0.072$

On each of JetNet’s three classes, we test r-GAN’s FC, GraphCNN, and TreeGAN
generators with rGAN’s FC and the PointNet-Mix discriminators, and compare
them to MPGAN’s MP generator and discriminator models, including both MP and
MP-LFC generator variations. Training and implementation details for each can
be found in App. D, and all code in Ref. [43].

We choose model parameters which, during training, yield the lowest
$W_{1}^{\mathrm{M}}$ score. This is because (1) $W_{1}$ scores between
physical features are more relevant for physics applications than the other
three metrics, and (2) qualitatively we find it be a better discriminator of
model quality than particle features or EFP scores. Table 5 lists the scores
for each model and class, and Fig. 3 shows plots of selected feature
distributions of real and generated jets, for the best performing FC,
GraphCNN, TreeGAN, and MP generators. We also provide in App. F discretized
images in the angular-coordinates-plane a.k.a “jet images”, however, we note
that it is in general not easy to visually evaluate the quality of individual
particle clouds, hence we focus on metrics and visualizations aggregated over
batches of clouds. Overall we find that MPGAN is a significant improvement
over the best FC, GraphCNN, and TreeGAN models, particularly for top and light
quark jets. This is evident both visually and quantitatively in every metric,
especially jet $W_{1}$s and FPND, with the exception of $W_{1}^{\mathrm{P}}$
where only the FC generator and PointNet discriminator (FC + PointNet)
combination is more performant.

We additionally perform a latency measurement and find, using an NVIDIA A100
GPU, that MPGAN generation requires 35.7 $\mu$s per jet. In comparison, the
traditional generation process for JetNet is measured on an 8-CPU machine as
requiring 46ms per jet, meaning MPGAN provides a three-orders-of-magnitude
speed-up. Furthermore, as noted in App. B, the generation of JetNet is
significantly simpler than full simulation and reconstruction used at the LHC,
which has been measured to require 12.3s [44] and 4s [45] respectively per top
quark jet. Hence in practical applications we anticipate MPGAN’s improvement
to potentially rise to five-orders-of-magnitude.

#### Real baseline comparison.

We find that MPGAN’s jet-level $W_{1}$ scores all fall within error of the
baselines in Table 5, while those of alternative generators are several
standard deviations away. This is particularly an issue with complex top quark
particle clouds, where we can see in Fig. 3 none of the existing generators
are able to learn the bimodal jet feature distributions, and smaller light
quark clouds, where we see distortion of jet features due to difficulty
reproducing the zero-padded particle features. No model is able to achieve
particle-level scores close to the baseline, and only those of the FC +
PointNet combination and MPGAN are of the same order of magnitude. We conclude
that MPGAN reproduces the physical observable distributions to the highest
degree of accuracy, but note, however, that it requires further improvement in
particle feature reconstruction before it is ready for practical application
in HEP.

#### Architecture discussion.

To disentangle the effectiveness of the MP generator and discriminator, we
train each individually with alternative counterparts (Table 5). With the same
PointNet discriminator, the GraphCNN and TreeGAN generators perform worse than
the simple FC generator for every metric on all three datasets. The physics-
motivated MP generator on the other hand outperforms all on the gluon and top
quark datasets, and significantly so on the jet-level $W_{1}$ scores and the
FPND. We note, however, that the MP generator is not a significant improvement
over the other generators with an FC discriminator. Holding the generator
fixed, the PointNet discriminator performs significantly better over the FC
for all metrics. With the FC, GraphCNN, and TreeGAN generators, PointNet is
also an improvement over the MP discriminator. With an MP generator, the MP
discrimimator is more performant on jet-level $W_{1}$ and FPND scores but, on
the top quark dataset, degrades $W_{1}^{\mathrm{P}}$ relative to PointNet.

We learn from these three things: (1) a generator or discriminator
architecture is only as effective as its counterpart—even though the MPGAN
combination is the best overall, when paired with a network which is not able
to learn complex substructure, or which breaks the permutation symmetry,
neither the generator or discriminator is performant, (2) for high-fidelity
jet feature reconstruction, both networks must be able to learn complex multi-
particle correlations—however, this can come at the cost of low-level feature
accuracy, and (3) MPGAN’s masking strategy is highly effective as both MP
networks are improvements all around on light quark jets.

#### Particle cloud evaluation metrics.

Figure 4: Correlation plots between pairs of evaluation metrics, evaluated on
400 separate batches of 50,000 MPGAN generated top quark jets.

We now discuss the merits of each evaluation metrics and provide suggestions
for their use in future work. Fig. 4 shows correlation plots between chosen
pairs of our evaluation metrics. As expected, we find W1-M and W1-EFP to be
highly correlated, as they both measure learning of global jet features. For
rigorous validation we suggest measuring both but for time-sensitive use-
cases, such as quick evaluations during model training, W1-M should be
sufficient. W1-M, FPND, and W1-P are all measuring different aspects of the
generation and are relatively uncorrelated. We expect FPND overall to be the
best and most discriminatory metric for evaluation, as it compares features
found by a SOTA classifier to be statistically optimum for characterizing
jets, while the W1 scores are valuable for their interpretability. Out of
these, W1-M/W1-EFP are the most important from a physics-standpoint, as we
generally characterize collisions by the high-level features of the output
jets, rather than the individual particle features.

MMD and coverage are both valuable for specifically evaluating the quality and
diversity of samples respectively, however we see from Fig. 4 that they
saturate after a certain point, after which FPND and $W_{1}$ scores are
necessary for stronger discrimination. We also note that in Table 5, models
with low $W_{1}$ scores relative to the baseline have the best coverage and
MMD scores as well. This indicates that the $W_{1}$ metrics are sensitive to
both mode collapse (measured by coverage), which is expected as in terms of
feature distributions mode collapse manifests as differing supports, to which
the $W_{1}$ distance is sensitive, as well as to individual sample quality
(measured by MMD), which supports our claim that recovering jet feature
distributions implies accurate learning of individual cloud structure.
Together this suggests that low $W_{1}$ scores are able validate sample
quality and against mode collapse, and justifies our criteria that a practical
ML simulation alternative have $W_{1}$ scores close to the baselines in Table
5. In conclusion, for thorough validation of generated particle clouds, we
recommend considering all three W-1 scores in conjunction with FPND, while MMD
and coverage, being focused tests of these aspects of generation, may be
useful for understanding failure modes during model development.

##  6 Summary

In this work, we publish JetNet: a novel particle cloud dataset to advance
machine learning (ML) research in high energy physics (HEP), and provide a
novel point-cloud-style dataset containing rich underlying physics for the ML
community to experiment with. We apply existing state-of-the-art point cloud
generative models to JetNet, and propose several physics- and computer-vision-
inspired metrics to rigorously evaluate generated clouds. We find that
existing models are not performant on a number of metrics, and fail to
reproduce high-level jet features—arguably the most significant aspect for
HEP. Our new message-passing generative adversarial network (MPGAN) model,
designed to capture complex global structure and handle variable-sized clouds
significantly improves performance in this area, as well as other metrics. We
propose MPGAN as a new baseline model on JetNet and invite others to improve
upon it.

#### Impact

With our JetNet dataset and library, we hope to lower the barrier to entry,
improve reproducibility, and encourage development in HEP and ML. Particularly
so in the area of simulation, where an accurate and fast ML particle cloud
generator will have significant impact in (1) lowering the computational and
energy cost of HEP research, as well as (2) increasing precision and
sensitivity to new physics at the Large Hadron Collider and future colliders
by providing more high-quality simulated data samples. One negative
consequence of this, however, may be a loss of interpretability, and hence
trustability, of the particle production generative model, which may
ultimately increase uncertainties—though the metrics we propose should
mitigate against this. More broadly, further advancements in the field of ML
point cloud generation may result in fake visual data generation for
proliferation of misinformation and impersonation/identity theft.

## Acknowledgments and Disclosure of Funding

This work was supported by the European Research Council (ERC) under the
European Union’s Horizon 2020 research and innovation program (Grant Agreement
No. 772369). R. K. was partially supported by an IRIS-HEP fellowship through
the U.S. National Science Foundation (NSF) under Cooperative Agreement
OAC-1836650, and by the LHC Physics Center at Fermi National Accelerator
Laboratory, managed and operated by Fermi Research Alliance, LLC under
Contract No. DE-AC02-07CH11359 with the U.S. Department of Energy (DOE). J. D.
is supported by the DOE, Office of Science, Office of High Energy Physics
Early Career Research program under Award No. DE-SC0021187 and by the DOE,
Office of Advanced Scientific Computing Research under Award No. DE-SC0021396
(FAIR4HEP). B. O and T. T are supported by grant 2018/25225-9, São Paulo
Research Foundation (FAPESP). B. O was also partially supported by grants
#2018/01398-1 and #2019/16401-0, São Paulo Research Foundation (FAPESP). J-R.
V. is partially supported by the ERC under the European Union’s Horizon 2020
research and innovation program (Grant Agreement No. 772369) and by the DOE,
Office of Science, Office of High Energy Physics under Award No. DE-SC0011925,
DE-SC0019227, and DE-AC02-07CH11359. D. G. is partially supported by the EU
ICT-48 2020 project TAILOR (No. 952215). This work was performed using the
Pacific Research Platform Nautilus HyperCluster supported by NSF awards
CNS-1730158, ACI-1540112, ACI-1541349, OAC-1826967, the University of
California Office of the President, and the University of California San
Diego’s California Institute for Telecommunications and Information
Technology/Qualcomm Institute. Thanks to CENIC for the 100 Gpbs networks.
Funding for cloud credits was supported by NSF Award #1904444 Internet2
supported E-CAS Exploring Clouds to Accelerate Science.

## References

  * [1] A. J. Larkoski, I. Moult, and B. Nachman, “Jet Substructure at the Large Hadron Collider: A Review of Recent Advances in Theory and Machine Learning”,  Phys. Rept. 841 (2020) 1, doi:10.1016/j.physrep.2019.11.001, arXiv:1709.04464.
  * [2] L. de Oliveira, M. Paganini, and B. Nachman, “Learning particle physics by example: Location-aware generative adversarial networks for physics synthesis”,  Comput. Softw. Big Sci. 1 (2017) 4, doi:10.1007/s41781-017-0004-6, arXiv:1701.05927.
  * [3] R. Kansal, “JetNet Library”. doi:10.5281/zenodo.5597893, https://github.com/jet-net/JetNet.
  * [4] R. K. Ellis, W. J. Stirling, and B. R. Webber, “QCD and collider physics”.  Cambridge monographs on particle physics, nuclear physics, and cosmology. Cambridge University Press, Cambridge, 2003.  Photography by S. Vascotto. doi:10.1017/CBO9780511628788.
  * [5] T. Sjöstrand et al., “An introduction to pythia 8.2”,  Comput. Phys. Commun. 191 (2015) 159, doi:10.1016/j.cpc.2015.01.024, arXiv:1410.3012.
  * [6] J. Bellm et al., “herwig 7.0/herwig ++3.0 release note”,  Eur. Phys. J. C 76 (2016), no. 4, 196, doi:10.1140/epjc/s10052-016-4018-8, arXiv:1512.01178.
  * [7] B. Andersson, G. Gustafson, G. Ingelman, and T. Sjostrand, “Parton Fragmentation and String Dynamics”,  Phys. Rept. 97 (1983) 31, doi:10.1016/0370-1573(83)90080-7.
  * [8] J. Cogan, M. Kagan, E. Strauss, and A. Schwarztman, “Jet-Images: Computer Vision Inspired Techniques for Jet Tagging”,  JHEP 02 (2015) 118, doi:10.1007/JHEP02(2015)118, arXiv:1407.5675.
  * [9] L. de Oliveira et al., “Jet-images — deep learning edition”,  JHEP 07 (2016) 069, doi:10.1007/JHEP07(2016)069, arXiv:1511.05190.
  * [10] H. Qu and L. Gouskos, “ParticleNet: Jet Tagging via Particle Clouds”,  Phys. Rev. D 101 (2020), no. 5, 056019, doi:10.1103/PhysRevD.101.056019, arXiv:1902.08570.
  * [11] J. Shlomi, P. Battaglia, and J.-R. Vlimant, “Graph neural networks in particle physics”,  Mach. Learn.: Sci. Technol. 2 (2021) 021001, doi:10.1088/2632-2153/abbf9a, arXiv:2007.13681.
  * [12] J. Duarte and J.-R. Vlimant, “Graph neural networks for particle tracking and reconstruction”, in  Artificial Intelligence for High Energy Physics, P. Calafiura, D. Rousseau, and K. Terao, eds.  World Scientific Publishing, 12, 2020.  arXiv:2012.01249.  Submitted to _Int. J. Mod. Phys. A_. doi:10.1142/12200.
  * [13] CMS Collaboration, “Identification of heavy, energetic, hadronically decaying particles using machine-learning techniques”,  JINST 15 (2020) P06005, doi:10.1088/1748-0221/15/06/P06005, arXiv:2004.08262.
  * [14] A. X. Chang et al., “ShapeNet: An information-rich 3D model repository”, Technical Report, 2015.  arXiv:1512.03012.
  * [15] R. Kansal et al., “JetNet”, 05, 2021.  doi:10.5281/zenodo.4834876.
  * [16] M. Pierini, J. M. Duarte, N. Tran, and M. Freytsis, “hls4ml LHC jet dataset (30 particles)”, 01, 2020.  doi:10.5281/zenodo.3601436.
  * [17] M. Paganini, L. de Oliveira, and B. Nachman, “CaloGAN : Simulating 3D high energy particle showers in multilayer electromagnetic calorimeters with generative adversarial networks”,  Phys. Rev. D 97 (2018), no. 1, 014021, doi:10.1103/PhysRevD.97.014021, arXiv:1712.10321.
  * [18] ATLAS Collaboration, “Fast simulation of the ATLAS calorimeter system with generative adversarial networks”, ATLAS Software Public Note ATL-SOFT-PUB-2020-006, 2020.
  * [19] M. Erdmann, J. Glombitza, and T. Quast, “Precise simulation of electromagnetic calorimeter showers using a Wasserstein generative adversarial network”,  Comput. Softw. Big Sci. 3 (2019) 4, doi:10.1007/s41781-018-0019-7, arXiv:1807.01954.
  * [20] F. Carminati et al., “Generative adversarial networks for fast simulation”,  J. Phys. Conf. Ser. 1525 (2020) 012064, doi:10.1088/1742-6596/1525/1/012064.
  * [21] Y. Lu, J. Collado, D. Whiteson, and P. Baldi, “Sparse autoregressive models for scalable generation of sparse images in particle physics”,  Phys. Rev. D 103 (2021), no. 3, 036012, doi:10.1103/PhysRevD.103.036012, arXiv:2009.14017.
  * [22] S. Ioffe and C. Szegedy, “Batch normalization: Accelerating deep network training by reducing internal covariate shift”, in  Proceedings of the 32nd International Conference on Machine Learning, F. Bach and D. Blei, eds., volume 37, p. 448.  PMLR, 2015.  arXiv:1502.03167.
  * [23] C.-L. Li et al., “Point cloud GAN”, in  2019 ICLR Workshop Deep Generative Models for Highly Structured Data.  2018\.  arXiv:1810.05795.
  * [24] G. Yang et al., “PointFlow: 3D point cloud generation with continuous normalizing flows”, in  2019 IEEE/CVF International Conference on Computer Vision (ICCV), p. 4540.  2019\.  arXiv:1906.12320.  doi:10.1109/ICCV.2019.00464.
  * [25] R. Klokov, E. Boyer, and J. Verbeek, “Discrete point flow networks for efficient point cloud generation”, in  2020 European Conference on Computer Vision (ECCV), A. Vedaldi, H. Bischof, T. Brox, and J.-M. Frahm, eds., p. 694.  Springer International Publishing, Cham, 2020.  arXiv:2007.10170.
  * [26] R. Cai et al., “Learning gradient fields for shape generation”, in  2020 European Conference on Computer Vision (ECCV), A. Vedaldi, H. Bischof, T. Brox, and J.-M. Frahm, eds., p. 364.  Cham, 2020.  arXiv:2008.06520.
  * [27] J. Köhler, L. Klein, and F. Noe, “Equivariant flows: Exact likelihood generative learning for symmetric densities”, in  Proceedings of the 37th International Conference on Machine Learning, H. D. III and A. Singh, eds., volume 119 of  Proceedings of Machine Learning Research, p. 5361.  PMLR, 2020.  arXiv:2006.02425.
  * [28] G. N. C. Simm, R. Pinsler, G. Csányi, and J. M. Hernández-Lobato, “Symmetry-aware actor-critic for 3d molecular design”, in  International Conference on Learning Representations.  2021\.
  * [29] N. Gebauer, M. Gastegger, and K. Schütt, “Symmetry-adapted generation of 3D point sets for the targeted discovery of molecules”, in  Advances in Neural Information Processing Systems 32, H. Wallach et al., eds., p. 7566.  Curran Associates, Inc., 2019.  arXiv:1906.00957.
  * [30] P. Achlioptas, O. Diamanti, I. Mitliagkas, and L. Guibas, “Learning representations and generative models for 3D point clouds”, in  Proceedings of the 35th International Conference on Machine Learning, J. Dy and A. Krause, eds., volume 80, p. 40.  PMLR, 2018.
  * [31] D. Valsesia, G. Fracastoro, and E. Magli, “Learning localized generative models for 3D point clouds via graph convolution”, in  International Conference on Learning Representations (ICLR) 2019.  2019\.
  * [32] D. W. Shu, S. W. Park, and J. Kwon, “3D point cloud generative adversarial network based on tree structured graph convolutions”, in  2019 IEEE/CVF International Conference on Computer Vision (ICCV), p. 3858.  2019\.  arXiv:1905.06292.  doi:10.1109/ICCV.2019.00396.
  * [33] C. R. Qi, H. Su, K. Mo, and L. J. Guibas, “PointNet: Deep learning on point sets for 3D classification and segmentation”, in  Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR).  July, 2017.
  * [34] H. Wang et al., “Rethinking sampling in 3D point cloud generative adversarial networks”, 2020. arXiv:2006.07029.
  * [35] J. Gilmer et al., “Neural message passing for quantum chemistry”, in  Proceedings of the 34th International Conference on Machine Learning, D. Precup and Y. W. Teh, eds., volume 70, p. 1263.  PMLR, 2017.  arXiv:1704.01212.
  * [36] CMS Collaboration, “Recent Developments in CMS Fast Simulation”,  PoS ICHEP2016 (2016) 181, doi:10.22323/1.282.0181, arXiv:1701.03850.
  * [37] P. T. Komiske, E. M. Metodiev, and J. Thaler, “Energy flow polynomials: A complete linear basis for jet substructure”,  JHEP 04 (2018) 013, doi:10.1007/JHEP04(2018)013, arXiv:1712.07124.
  * [38] A. Borji, “Pros and cons of GAN evaluation measures: New developments”, 2021. arXiv:2103.09396.
  * [39] M. Heusel et al., “GANs trained by a two time-scale update rule converge to a local nash equilibrium”, in  Advances in Neural Information Processing Systems, I. Guyon et al., eds., volume 30, p. 6626.  Curran Associates, Inc., 2017.  arXiv:1706.08500.
  * [40] P. T. Komiske, E. M. Metodiev, and J. Thaler, “Metric Space of Collider Events”,  Phys. Rev. Lett. 123 (2019), no. 4, 041801, doi:10.1103/PhysRevLett.123.041801, arXiv:1902.02346.
  * [41] Y. Wang et al., “Dynamic graph CNN for learning on point clouds”,  ACM Trans. Graph. 38 (2019) doi:10.1145/3326362, arXiv:1801.07829.
  * [42] F. Monti et al., “Geometric deep learning on graphs and manifolds using mixture model CNNs”, in  2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), p. 5425.  2017\.  arXiv:1611.08402.  doi:10.1109/CVPR.2017.576.
  * [43] R. Kansal, “MPGAN Code”. doi:10.5281/zenodo.5598407, https://github.com/rkansal47/MPGAN/tree/neurips21.
  * [44] Pedro, Kevin, “Current and future performance of the cms simulation”,  EPJ Web Conf. 214 (2019) 02036, doi:10.1051/epjconf/201921402036.
  * [45] C. Chen et al., “Data augmentation at the LHC through analysis-specific fast simulation with deep learning”, 2020. arXiv:2010.01835.
  * [46] K. Cranmer, J. Brehmer, and G. Louppe, “The frontier of simulation-based inference”,  PNAS 117 (2020) 30055, doi:10.1073/pnas.1912789117, arXiv:1911.01429.
  * [47] G. Cowan, “Statistics”, Ch. 40 in Particle Data Group, P. A. Zyla et al., “Review of Particle Physics”,  PTEP 2020 (2020) 083C01, doi:10.1093/ptep/ptaa104.
  * [48] CMS Collaboration, “Precise determination of the mass of the Higgs boson and tests of compatibility of its couplings with the standard model predictions using proton collisions at 7 and 8 TeV”,  Eur. Phys. J. C 75 (2015) 212, doi:10.1140/epjc/s10052-015-3351-7, arXiv:1412.8662.
  * [49] R. J. Barlow and C. Beeston, “Fitting using finite Monte Carlo samples”,  Comput. Phys. Commun. 77 (1993) 219, doi:10.1016/0010-4655(93)90005-W.
  * [50] J. Brehmer and K. Cranmer, “Simulation-based inference methods for particle physics”, in  Artificial Intelligence for High Energy Physics, P. Calafiura, D. Rousseau, and K. Terao, eds.  World Scientific Publishing, 10, 2020.  arXiv:2010.06439.  Submitted to _Int. J. Mod. Phys. A_. doi:10.1142/12200.
  * [51] J. Alwall et al., “The automated computation of tree-level and next-to-leading order differential cross sections, and their matching to parton shower simulations”,  JHEP 07 (2014) 079, doi:10.1007/JHEP07(2014)079, arXiv:1405.0301.
  * [52] R. D. Ball et al., “Parton distributions with LHC data”,  Nucl. Phys. B 867 (2013) 244–289, doi:10.1016/j.nuclphysb.2012.10.003, arXiv:1207.1303.
  * [53] P. Skands, S. Carrazza, and J. Rojo, “Tuning pythia 8.1: The Monash 2013 Tune”,  Eur. Phys. J. C 74 (2014), no. 8, 3024, doi:10.1140/epjc/s10052-014-3024-y, arXiv:1404.5630.
  * [54] M. Cacciari, G. P. Salam, and G. Soyez, “The anti-$k_{\mathrm{T}}$ jet clustering algorithm”,  JHEP 04 (2008) 063, doi:10.1088/1126-6708/2008/04/063, arXiv:0802.1189.
  * [55] M. Cacciari, G. P. Salam, and G. Soyez, “FastJet user manual”,  Eur. Phys. J. C 72 (2012) 1896, doi:10.1140/epjc/s10052-012-1896-2, arXiv:1111.6097.
  * [56] M. Cacciari and G. P. Salam, “Dispelling the $n^{3}$ myth for the $k_{\mathrm{T}}$ jet-finder”,  Phys. Lett. B 641 (2006) 57, doi:10.1016/j.physletb.2006.08.037, arXiv:hep-ph/0512210.
  * [57] X. Mao et al., “Multi-class generative adversarial networks with the L2 loss function”, 2016. arXiv:1611.04076.
  * [58] N. Srivastava et al., “Dropout: A simple way to prevent neural networks from overfitting”, J. Mach. Learn. Res. 15 (2014) 1929.
  * [59] I. Gulrajani et al., “Improved training of Wasserstein GANs”, in  Advances in Neural Information Processing Systems, I. Guyon et al., eds., volume 30, p. 5767.  Curran Associates, Inc., 2017.  arXiv:1704.00028.
  * [60] T. Miyato, T. Kataoka, M. Koyama, and Y. Yoshida, “Spectral normalization for generative adversarial networks”, in  6th International Conference on Learning Representations.  2018\.  arXiv:1802.05957.
  * [61] F. Schäfer, H. Zheng, and A. Anandkumar, “Implicit competitive regularization in GANs”, 2019. arXiv:1910.05852.
  * [62] T. Karras et al., “Training generative adversarial networks with limited data”, in  Advances in Neural Information Processing Systems, H. Larochelle et al., eds., volume 33, p. 12104.  Curran Associates, Inc., 2020.  arXiv:2006.06676.
  * [63] N.-T. Tran et al., “On data augmentation for GAN training”,  IEEE Trans. Image Process. 30 (2021) 1882, doi:10.1109/tip.2021.3049346, arXiv:2006.05338.
  * [64] Z. Zhao et al., “Image augmentations for GAN training”, 2020\. arXiv:2006.02595.
  * [65] M. Arjovsky, S. Chintala, and L. Bottou, “Wasserstein GAN”, 2017. arXiv:1701.07875.

##  Appendix A Simulation-Based Inference in HEP

The primary motivation for producing simulations with Monte Carlo (MC)
generators, and potentially by machine learning (ML) generators as well, in
experimental high energy physics is to develop likelihood models for new
fundamental physics theories, with unknown values of parameters of interest
[46, 47]. These models are compared with experimental data, such as that
collected at the LHC, to perform hypothesis tests of the theories as well as
estimate and develop limits and confidence intervals for physical parameters
[47, 48]. Increasing the number of simulations and their accuracy can reduce
the statistical uncertainties in our models, thus allowing for higher
precision measurements and potential discovery of new physics [49, 50].

In practice, in each analysis we perform rigorous checks of our simulations,
including comparisons in “control regions” (selected portions of the data with
a known composition) to check for MC discrepancies with real data. Mismatches
are corrected via reweighting the events, i.e. unrealistic simulated jets will
be given less weight in the overall analysis, and are factored into the final
uncertainties in the analysis results. The same procedure should be followed
for data created by generative ML methods, however, validation before this
step using the metrics proposed in Sec. 3, particularly the $W_{1}$ scores,
should mitigate the possibilities of such outliers.

##  Appendix B JetNet Generation

The so-called parton-level events are first produced at leading-order using
MADGRAPH5_aMCATNLO 2.3.1 [51] with the NNPDF 2.3LO1 parton distribution
functions [52]. To focus on a relatively narrow kinematic range, the
transverse momenta of the partons and undecayed gauge bosons are generated in
a window with energy spread given by $\Delta
p_{\mathrm{T}}/p_{\mathrm{T}}=0.01$, centered at $1\,\text{TeV}$. These
parton-level events are then decayed and showered in pythia 8.212 [5] with the
Monash 2013 tune [53], including the contribution from the underlying event.
For each original particle type, 200,000 events are generated. Jets are
clustered using the anti-$k_{\mathrm{T}}$ algorithm [54], with a distance
parameter of $R=0.8$ using the FastJet 3.1.3 and FastJet contrib 1.027
packages [55, 56]. Even though the parton-level $p_{\mathrm{T}}$ distribution
is narrow, the jet $p_{\mathrm{T}}$ spectrum is significantly broadened by
kinematic recoil from the parton shower and energy migration in and out of the
jet cone. We apply a restriction on the measured jet $p_{\mathrm{T}}$ to
remove extreme events outside of a window of
$0.8\,\text{TeV}<p_{\mathrm{T}}<1.6\,\text{TeV}$ for the
$p_{\mathrm{T}}=1\,\text{TeV}$ bin. This generation is a significantly
simplified version of the official simulation and reconstruction steps used
for real detectors at the LHC, so as to remain experiment-independent and
allow public access to the dataset.

##  Appendix C Point Cloud Generative Models

Figure 5: Comparison of real and PCGAN-generated distributions for a subset of
jet and particle features. Top: gluon jet features, Middle: light quark jets,
Bottom: top quark jets.

Apart from the GAN models discussed in Sec. 3, there are several published
generative models for point clouds which we argue are not applicable to jets.

###  C.1 ShapeNet Point Clouds

A number of successful generative models exploit a key inductive bias of
ShapeNet-based clouds: that the individual distributions of sampled points
conditioned on a particular object are identical and independent (the i.i.d
assumption). This assumption allows for hierarchical generative frameworks,
such as Point-Cloud-GAN (PCGAN) [23], which uses two networks: one to generate
a latent object-level representation, and a second to sample independent
points given such a representation. The PointFlow [24] and Discrete PointFlow
[25] models use a similar idea of sampling independently points conditioned on
a learnt latent representation of the shape, but with a variational
autoencoder (VAE) framework and using normalizing flows for transforming the
sampled points.

This hierarchical-sampling approach is appealing for ShapeNet clouds, however,
as discussed in Sec. 2 the key i.i.d. assumption is not applicable to jets
with their highly correlated particle constituents. In fact, in contrast to
ShapeNet objects which have a structure independent of the particular sampled
cloud, jets are entirely defined by the distribution of their constituents.

Another model, ShapeGF [26], uses an approach of again sampling points
independently from a prior distribution, but transforming them to areas of
high density via gradient ascent, maximizing a learnt log-density concentrated
on an object’s surface. This approach suffers as well from the i.i.d.
assumption in the context of jets, and additionally, unlike for ShapeNet point
clouds, there is no such high-density region in momentum-space where particles
tend to be concentrated, so learning and maximizing a log-density is not
straightforward.

To support our overall claim of the inviability of the i.i.d. assumption for
particle clouds, we train a PCGAN model on JetNet and show the produced
feature distributions in Fig. 5. We can see, as expected, while this network
is partially reproducing the particle feature distributions, it is entirely
unable to learn the jet-level structure in particle clouds.

###  C.2 Molecular Point Clouds

3D molecules are another common point-cloud-style data structure, and there
have been developments in generative models in this area as well. Kohler et
al. [27] introduce physics-motivated normalizing flows equivariant to
rotations around the center of mass, i.e. the SO(N) symmetries, for generating
point clouds. This is appealing as normalizing flows give access to the
explicit likelihood of generated samples, and having an architecture
equivariant to physical symmetries such as 3D rotations can improve the
generalizability and interpretability of the model. Since jets are
relativistic, however, we require an architecture equivariant to the non-
compact SO(3, 1) Lorentz group, to which this model has not been generalized
yet. Simm et al. [28] present a reinforcement-learning-based approach for
generating 3D molecules, using an agent to iteratively add atoms to a molecule
and defining the reward function as the energy difference between the new
molecule and the old with the new atom at the origin. This reward function is
not directly applicable to jets. where particle distributions are based on the
QCD dynamics rather than on minimizing the total energy. Finally, Gebauer et
al. [29] introduce G-SchNet, an autoregressive model for producing molecules
represented as point clouds, iteratively adding one atom at a time based on
the existing molecule. Their iterative procedure however was proposed for
point clouds of at most nine atoms, and does not scale well in terms of time
to larger clouds.

Overall, all the models discussed heavily incorporate inductive biases which
are specific to their respective datasets and don’t apply to JetNet. However,
they are extremely interesting approaches nonetheless, and adapting them with
jet-motivated biases should certainly be explored in future work. Indeed, a
significant contribution of our work is publishing a dataset which can
facilitate and hopes to motivate such development.

##  Appendix D Training and Implementation Details

PyTorch code and trained parameters for models in Table 5 are provided in the
MPGAN repository [43]. Models were trained and hyperparameters optimized on
clusters of NVIDIA GeForce RTX 2080 Ti, Tesla V100, and A100 GPUs.

###  D.1 MPGAN

We use the least squares loss function [57] and the RMSProp optimizer with a
two time-scale update rule [39] with a learning rate (LR) for the
discriminator three times greater than that of the generator. The absolute
rate differed per jet type. We use LeakyReLU activations (with negative slope
coefficient 0.2) after all MLP layers except for the final generator and
discriminator outputs where tanh and sigmoid activations respectively are
applied. We attempted discriminator regularization to alleviate mode collapse
via dropout [58], batch normalization [22], a gradient penalty [59], spectral
normalization [60], adaptive competitive gradient descent [61] and data
augmentation of real and generated graphs before the discriminator [62, 63,
64]. Apart from dropout (with fraction $0.5$), none of these demonstrated
significant improvement with respect to mode dropping or cloud quality.

We use a generator LR of $10^{-3}$ and train for 2000 epochs for gluon jets,
$2\times 10^{-3}$ and 2000 epochs for top quark jets, and $0.5\times 10^{-3}$
and 2500 epochs for light quark jets. We use a batch size of 256 for all jets.

###  D.2 rGAN, GraphCNNGAN, TreeGAN, and PointNet-Mix

For rGAN and GraphCNNGAN we train two variants: (1) using the original
architecture hyperparameters in Refs. [30, 31] for the 2048-node point clouds,
and (2) using hyperparameters scaled down to 30-node clouds—specifically: a 32
dimensional latent space, followed by layers of 64, 128, and 90 nodes for
r-GAN, or followed by two graph convolutional layers with node features sizes
of 32 and 24 respectively for GraphCNN-GAN. The scaled-down variant performed
better for both models, and its scores are the ones reported in Table 5. For
TreeGAN, starting from single vertex—in analogy with a jet originating from a
single particle—we use five layers of up-sampling and ancestor-descendant
message passing, with a scale-factor of two in each and node features per
layer of 96, 64, 64, 64, and 64 respectively. LRs, batch sizes, loss
functions, gradient penalties, optimizers, ratios of critic to generator
updates, activations, and number of epochs are the same as in the original
paper and code. We use the architecture defined in [34] for the PointNet-Mix
discriminator.

###  D.3 FPND

Apart from the number of input particle features (three in our case, excluding
the mask feature), we use the original ParticleNet architecture in Ref. [10].
We find training with the Adam optimizer, LR $10^{-4}$, for 30 epochs
outperformed the original recommendations on our dataset. Activations after
the first fully connected layer, pre-ReLU, are used for the FPND measurement.

###  D.4 PCGAN

We use the original PCGAN implementation for the sampling networks and
training, with a 256-dimensional latent object representation. For the latent
code GAN we use a 3 layer fully connected network for both the generator, with
an input size of 128 and intermediate layer sizes of 256 and 512, and
discriminator, with intermediate layer sizes of 512 and 256, trained using the
Wasserstein-GAN [65] loss with a gradient penalty.

##  Appendix E Masking Strategies

Figure 6: Particle $p_{\mathrm{T}}^{\mathrm{rel}}$ and relative jet mass
distributions of real jets and those generated by MPGAN without our masking
strategy. Left: gluon, right: light quark jets. We see that while for gluon
jets the generator learns distributions correctly, it struggles to learn the
discontinuous spike, due to the zero-padded particles, in the light quark
$p_{\mathrm{T}}^{\mathrm{rel}}$ distribution. This also leads to a distorted
mass distribution.

In JetNet, jets with fewer than 30 particles are zero-padded to fill the
30-particle point cloud. Such zero-padded particles pose a problem for the
generator, which is not able to learn this sharp discontinuity in the jet
constituents (Fig. 6).

Figure 7: The four alternative masking strategies which we test.

Figure 8: Loss curve of a training on light quark jets with masking strategy
3, typical of loss curves with all four strategies.

To counter this issue, we experiment with five masking strategies, out of
which the one described in Sec. 4 was most successful. The four alternatives,
which all involve the generator learning the mask without any external input,
are shown in Fig. 7.

Strategy 1 treats the mask homogeneously as an extra feature to learn. A
variation of this weights the nodes in the discriminator the mask. In strategy
2, a mask is calculated for each generated particle as a function of its
$p_{\mathrm{T}}^{\mathrm{rel}}$, based on an empirical minimum cutoff in the
dataset. In particular, both a Heaviside-step-function and a continuous mask
function as in the figure are tested. The standard MP discriminator, as
described in Sec. 4, is used. Strategy 3 sees the generator applying an FC
layer per particle in the initial cloud to learn their respective masks, with
both the MP discriminator, as well as a variant with the number of unmasked
nodes in the clouds added as an extra feature to the FC layer. In strategies 1
and 3 we test both binary and continuous masks. Finally, in strategy 4, we
train an auxiliary network to choose a number of particles to mask (as opposed
to sampling from the real distribution), which is then passed into the
standard MP generator.

We find that all such strategies are unable to produce accurate light quark
jets, and in fact trainings for each diverge in the fashion depicted in Fig.
8, even using each discriminator regularization method mentioned in App. D).
We conclude that learning the number of particles to produce is a significant
challenge for a generator network, but is a relatively simple feature with
which to discriminate between real and fake jets. To equalize this we use the
strategy in Sec. 4 where the number of particles to produce is sampled
directly from the real distribution, removing the burden of learning this
distribution from the generator network.

##  Appendix F Jet Images

Figs. 9–11 show samples of real and generated “jet images”: discretized
representations of jets in the angular-coordinate-plane.

Figure 9: Random samples of discretized images in the $\eta^{\mathrm{rel}}$
$-$$\phi^{\mathrm{rel}}$ plane, with pixel intensities equal to particle
$p_{\mathrm{T}}^{\mathrm{rel}}$, of real and generated gluon jets (left), and
an average over 10,000 such sample images (right).

Figure 10: Random samples of discretized images in the $\eta^{\mathrm{rel}}$
$-$$\phi^{\mathrm{rel}}$ plane, with pixel intensities equal to particle
$p_{\mathrm{T}}^{\mathrm{rel}}$, of real and generated light quark jets
(left), and an average over 10,000 such sample images (right).

Figure 11: Random samples of discretized images in the $\eta^{\mathrm{rel}}$
$-$$\phi^{\mathrm{rel}}$ plane, with pixel intensities equal to particle
$p_{\mathrm{T}}^{\mathrm{rel}}$, of real and generated top quark jets (left),
and an average over 10,000 such sample images (right).


# Astronomical Image Processing at Scale With Pegasus and Montage

G. Bruce Berriman1, John C. Good 2, Ewa Deelman3, Ryan Tanaka3, and Karan
Vahi3

###### Abstract

Image processing at scale is a powerful tool for creating new data sets and
integrating them with existing data sets and performing analysis and quality
assurance investigations. Workflow managers offer advantages in this type of
processing, which involves multiple data access and processing steps.
Generally, they enable automation of the workflow by locating data and
resources, recovery from failures, and monitoring of performance. In this
focus demo we demonstrate how the Pegasus Workflow Manager Python API manages
image processing to create mosaics with the Montage Image Mosaic engine. Since
2001, Pegasus has been developed and maintained at USC/ISI. Montage was in
fact one of the first applications used to design Pegasus and optimize its
performance. Pegasus has since found application in many areas of science.
LIGO exploited it in making discoveries of black holes. The Vera C. Rubin
Observatory used it to compare the cost and performance of processing images
on cloud platforms. While these are examples of projects at large scale, small
team investigations on local clusters of machines can benefit from Pegasus as
well.

1Caltech/IPAC-NExScI, Pasadena, CA, 91125<EMAIL_ADDRESS>
2Caltech/IPAC-NExScI, Pasadena, CA, 91125

3USC Information Sciences Institute, Marina del Rey, CA, 90292

##  1 Introduction

We aim to show how astronomers can take advantage of open source tools to
build workflows that can be executed on-premises machines/clusters and then
scale as far as possible without modification on distributed platforms. The
demo was a proof-of-concept image mosaic workflow run on a local laptop; the
same workflow was run without any modifications or optimizations on the Open
Science Grid (OSG) 111https://opensciencegrid.org/ high-throughput platform.

The tools we have chosen for this demo are the Pegasus Workflow Manager
222https://pegasus.isi.edu and the Montage Image Mosaic Engine
333http://montage.ipac.caltech.edu. Both are open source, operate by design on
multiple platforms, support a Python API, are (reasonably) easy to use, and
operate well with one another. The demo itself will create a 1∘ x 1 ∘ mosaic
of M17 in a Jupyter Notebook.

##  2 Pegasus and Montage

Pegasus is a fully featured workflow management system (Deelman et al. 2021)
that enables scientists to develop portable scientific workflows, structured
as directed acyclic graphs through Python, Java, or R APIs. These workflows
can then be run on heterogeneous resources, including local machines, clouds,
HPC systems, and campus clusters. Some important features of Pegasus are:
monitoring and debugging capabilities via command-line interface (CLI) tools
or a web dashboard; automatic handling of data movement between jobs (e.g.,
pulling input data from an AWS S3 bucket or sending one job’s output to
another job); and running workflows with data integrity and fault tolerance in
mind.

When working with Pegasus, scientists run a set of computations whose workflow
can, as a starting point, be sketched on paper. The abstract workflow can be
described programmatically using one of the APIs, in our case case Python.
Once that is done, the workflow will be serialized to a YAML representation,
which will be consumed by Pegasus and compiled into an executable workflow
intended for a specific execution environment, such as a cloud platform or a
campus cluster. HTCondor sits underneath Pegasus and is responsible for
running the jobs.

Montage is open source toolkit for assembling FITS images into custom mosaics
(Berriman & Good 2017). It is written in C for performance and portability.
Python binary extensions provide compiled performance on parallel platforms.
The modular design of Montage, with independent modules for each step in the
mosaic process, allows Montage to play well with workflow managers and enables
embarrassingly parallel processing.

##  3 Try It Out

We have created a Docker image with everything needed to run the Montage
workflow just described in a Jupyter Notebook 444 https://github.com/pegasus-
isi/ADASS21-montage-docker-image. This Docker image contains an installation
of Pegasus and HTCondor, the Montage binaries, and the workflow notebook. You
can interact with this container through Jupyter Notebooks via the browser.
This container environment supports testing different workflow configurations,
and contains Pegasus tutorial notebooks.

##  4 The Demo Notebook

The notebook demonstrates proof-of-concept parallelization through creation of
a small mosaic. The demo was run on a desktop machine for convenience. Because
of the overhead imposed on the process by HTCondor, the real power of the
approach will lie in its use on high-performance platforms.

The process involves populating three structures, maintained by Pegasus as
SQLite databases:

  * •

The "replica catalog:" The input files for the processing and how to access
them.

  * •

The "transformations:" The Montage modules and where to run them.

  * •

The "workflow" enumerates the specific steps in the processing, including
transforms, argument lists, which files are input and which are output, etc.
Pegasus deduces when jobs can be run in large part by when the precursor files
are available.

Then the Notebook asks Pegasus to devise a processing plan, and, if
successful, to submit that plan to HTCondor. Figure 1 shows a screenshot of
one step in the process.

Figure 1.: Screenshot of a section of the Notebook for building the Pegasus
Montage workflow.

##  5 Running the Workflow on a Parallel Platform: The Open Science Grid

We ran the Montage workflow on the OSG (Sfiligoi et al. 2009) with minimal
modifications to the Python workflow script to utilize OSG-specific features:

  * •

Data movements were managed by OSG’s StashCache data infrastructure, which
caches files in an opportunistic manner.

  * •

The workflow used RedHat Enterprise 7 Singularity containers, which are cached
within OSG’s CernVM-File system (CVMFS). Setup required changing only a few
lines in our workflow generation script.

Table 1 summarized metrics Pegasus reported on the 332 jobs (including jobs
added by Pegasus to handle data staging and cleanup).

Table 1.: Pegasus Montage Workflow Metrics on OSG Metric | Value
---|---
Number of Jobs | 332
Workflow Wall Time | 54m 43s
Cumulative Job Wall Time | 36m 48s
Cumulative Job Wall Time (from submit side) | 2h 38m
Integrity Checking Time Spent | 3.57s (for 357 files)

The wall times reveal considerable queuing overhead associated with running
the workflows on OSG. Because the overheads are incurred at the time of job
submission, it is beneficial to cluster jobs together to pay a single overhead
for a "cluster" of jobs. Pegasus can do this clustering automatically based on
user-specified criteria. Pegasus also adds overheads such as integrity
checking, but these, by contrast, are negligible and in this case took only
3.57 seconds.

##  6 Conclusions

We have successfully run Pegasus and Montage on a laptop and on the OSG and we
have made our setup available to others for replication and reuse. Our next
steps include optimizing performance on parallel platforms through clustering
and running the workflow on commercial cloud platforms.

### Acknowledgments

Pegasus is funded by the National Science Foundation under OAC SI2-SSI program
grant 1664162. Previously, NSF funded Pegasus under OCI SDCI program grant
0722019 and OCI SI2-SSI program grant 1148515. Montage is funded by the
National Science Foundation grant awards 1835379, 1642453 and 1440620. OSG is
supported by the National Science Foundation grant award 2030508.

## References

  * Berriman & Good (2017) Berriman, G. B., & Good, J. C. 2017, PASP, 129, 058006. 1702.02593
  * Deelman et al. (2021) Deelman, E., Ferreira da Silva, R., Vahi, K., Rynge, M., Mayani, R., Tanaka, R., Whitcup, W., & Livny, M. 2021, Journal of Computational Science, 52, 101200
  * Sfiligoi et al. (2009) Sfiligoi, I., Bradley, D. C., Holzman, B., Mhashilkar, P., Padhi, S., & Wurthwein, F. 2009, in 2009 WRI World Congress on Computer Science and Information Engineering, vol. 2 of 2, 428