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 introduces a new particle cloud dataset (JetNet) for high-energy physics. The authors set up a benchmark for particle cloud generation with a few physics-inspired and vision-related metrics. As existing point cloud GANs are not suitable for physics applications, they have also developed a new message-passing GAN (MPGAN) and have achieved good empirical performance.