Particle Cloud Generation with Message Passing Generative Adversarial Networks

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 . ⇤Also at National and Kapodistrian University of Athens , Athens , Greece . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Answering this question affirmatively has important practical applications . At the LHC , large simulated data samples of collision events2 are generated using Monte Carlo ( MC ) methods in order to translate theoretical predictions into observable distributions , and ultimately perform physics analyses3 . 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 can not 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 is computed with respect to the x axis . The polar angle ✓ is used to compute the pseudorapidity ⌘ = log ( tan ( ✓/2 ) ) . The transverse momentum ( pT ) is the projection of the particle momentum on the ( x , y ) plane . As is customary , we transform the particle momenta from Cartesian coordinates ( px , py , pz ) to longitudinal-boost-invariant pseudo-angular coordinates ( pT , ⌘ , ) , as shown in Fig . 1 . 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 ⌘- plane , and taking the intensity of each “ pixel ” in this grid to be a monotonically increasing function of the corresponding particle pT . 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 2An event is a set of observed particles with well-defined momenta , but not generally well-defined positions . 3See App . A for further details on the downstream applications of simulations . t→Wb→qqb h→bb t→Wb→qqb h→bb 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 ] 4 , it consists of simulated particle jets with transverse momenta pjetT ⇡ 1TeV , originating from gluons , light quarks , and top quarks produced in 13TeV 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 pT particles per jet , allowing for jets with potentially fewer than 30 by zeropadding . For each particle we provide the following four features : the relative angular coordinates ⌘rel = ⌘particle ⌘jet and rel = particle jet ( mod 2⇡ ) , relative transverse momentum prelT = pparticleT /p jet T , 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 . 4This dataset was also released under the CC-BY 4.0 license . Refs . [ 2 , 17–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–26 ] and molecular [ 27–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 variablesized 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 .

Summary: This paper proposes to explore an interesting application of existing point cloud generative models to the problem of generation of jets: high energy sprays of particles from elementary particle colliders. Jets are analogous in a way to point clouds with a difference in considered per-object features: per-particle momenta are considered for jets contrary to per-point 3D coordinates for point clouds. With sufficient level of quality such generative models could potentially be used instead of costly Monte Carlo simulations.  The authors explore existing point cloud gan-based generative models and propose an adaptation of existing message-passing networks to be used in generators and discriminators of resulting particle cloud gans. Several metrics are proposed for evaluation, including adaptation of existing point cloud generation task metrics, and the distance between momenta distributions for real and generated jets.