DeepTreeGAN: Fast Generation of High Dimensional Point Clouds

. In High Energy Physics, detailed and time-consuming simulations are used for particle interactions with detectors. To bypass these simulations with a generative model, the generation of large point clouds in a short time is required, while the complex dependencies between the particles must be correctly modelled. Particle showers are inherently tree-based processes, as each particle is produced by the decay or detector interaction of a particle of the previous generation. In this work, we present a novel Graph Neural Network model (DeepTreeGAN) that is able to generate such point clouds in a tree-based manner. We show that this model can reproduce complex distributions, and we evaluate its performance on the public JetNet dataset.


Introduction
In particle physics, detailed, near-perfect simulations of the underlying physical processes, the measurements, and the details of the experimental apparatus are state-of-the-art.However, accurate simulations of large and complex detectors, especially calorimeters, using current Monte Carlo-based tools such as those implemented in the Geant4 [1] toolkit are computationally intensive.Currently, more than half of the worldwide Large Hadron Collider (LHC) grid resources are used for the generation and processing of simulated data [2].For the future High-Luminosity upgrade of the LHC (HL-LHC [3]), the computational requirements will exceed the computational resources without significant speedups, e.g., projecting the current technologies to the needs of the CMS experiment [4] by 2028.
The majority of these approaches represent the data as voxels living on three-dimensional grids.But especially for high-granularity calorimeters like the HGCAL, the calorimeter cells are highly irregular and the data is often sparse.In such a situation, it is beneficial to represent the data as Point Clouds (PC), e.g., tuples of cell energies and three-dimensional coordinates, i.e., four-dimensional point clouds [15][16][17].For modelling calorimeter showers, which are cascades of particles, it seems natural to model the creation of such point clouds as a tree-like process.Such an approach can be implemented by Graph Neural Networks (GNN) [18] and Graph Convolutions [19].For modelling PC data representing three-dimensional shapes, a similar approach [20] called tree-GAN [20] exists.For our use case, especially for modelling the large dynamic range of the energy dimension, this approach is not sufficient.Here we report on the development of a largely extended Graph GAN approach, which we named DeepTreeGAN that can be applied to multiple areas of PC generation.To demonstrate the abilities of our model, we present a first application to the JetNet dataset introduced in the next section.

JetNet and MPGAN
In this study, the JetNet [21] datasets are used.Each dataset contains PYTHIA [22] jets with an energy of about 1 TeV, with each jet containing up to 30 or 150 constituents (here: 30).The datasets differ in the jet-initiating parton.Here, datasets for top quark, light quarks and gluon-initiated jets are studied [23,24].Each dataset contains about 170k individual jets split as 110k/10k/50k for training/test/validation, where the validation dataset is used for our results.The jet constituents, the particles, are clustered with a cone radius of R = 0.8.These particles are considered to be massless and can therefore be described by their 3-momenta or equivalently by transverse momentum p T , pseudorapidity η, and azimuthal angle ϕ.In the JetNet datasets, these variables are given relative to the jet momentum: η rel i η i − η jet , ϕ rel i ϕ i − (ϕ jet mod 2π), and p rel T,i p T,i /p T,jet , where i runs over the particles in a jet.Calculating the invariant mass from these relative quantities, for example, for the jet mass, implies m rel = m jet /p T,jet .The JetNet library [25] provides several metrics that are used in this study.Additionally, the authors provide a baseline model called MPGAN [26].This dataset has gained popularity in the particle physics community as a benchmark for PC-based generative models [15][16][17][27][28][29][30][31][32][33][34].

Architecture
The generator of this GAN constructs point clouds by starting with a random vector as the root of a tree and then attaching one level of leaves after the other.The output of the generator is the last level of the tree.Starting with the root node, a Branching Layer (Figure 1) takes the current leaves from the tree, maps each leaf to a given number nodes n i and attaches these nodes as new leaves to the tree.Then further Branching Layers are applied until the desired number i n i of nodes (here: 30 nodes) is reached.For these projections, multiple feed-forward neural networks (FFNs) are used.
In between the Branching Layers a Message Passing Layer (MPL) passes the information down the tree.In this step, each of the nodes is updated with information from its ancestors.This Ancestor MPL is described in Figure 2.
A Global Feature Layer is introduced that condenses the information of the current leaves in a single vector: The feature vectors of the leaves are passed through a first FFN, summed up and passed through a second FFN.The resulting vector is then used -together with a global condition vector -to condition the following Branching Layer and Ancestor MPL.The global condition vector contains only the number of constituents in the event.For each of the levels of the tree, first the Global Feature Layer is executed, then the Branching Layer and then the Ancestor MPL.
The FFNs of this model consist of 3 hidden layers of 100 nodes without bias.The first two hidden layers are followed by a Batch Normalization [37] layer and LeakyReLU activation with a negative slope of 0.1.The critic of MPGAN [26] is employed.The vector of the parent is split into one part for each of the branches.To each of these vectors, the Global Feature (2) and the global condition vector (here: number of constituents) are concatenated.The Branching FFN maps the vectors to the same size as their parents.
Then the feature vector of the parents is added to each of its children.With the new children added as leaves to the tree, all the levels of the tree are stacked up and passed through a Dimensionality Reduction FFN.With each Branching Layer, the number of nodes increases (1 → 2 → 2*3 → 2*3*5) and the number of features decreases (64 → 33 → 20 → 3).The edges in the graph are constructed so that each node receives messages from each of its ancestors as well as itself.In the given example node 4 is updated with the aggregate of the messages from nodes 1, 2 and 4. As a messagepassing algorithm, GINConv [35] -as implemented in PyTorch Geometric [36] -is chosen.For GINConv, the messages are the features of the source node (here: the ancestor).These messages are aggregated by summing over all messages addressed to the target node.These aggregates are added to the feature vector of the target node (scaled with a learnable weight) and passed through an FFN.The nodes are then updated with the output of the FFNs plus, as a residual connection, the original feature vector of the nodes themselves.

Training
Generator and Critic are trained using the Hinge [38] loss, with a ratio of 1 to 2 gradient steps.The optimizer for the Generator is Adam [39] with β 1 = 0.9, β 2 = 0.999, a weight decay of 10 −4 and a learning rate of 10 −5 .The optimizer for the critic is Adam with β 1 = 0.0, β 2 = 0.9, a weight decay of 10 −4 and a learning rate of 3 * 10 −5 .

Comparison to tree-GAN
Compared to the model proposed in [20], DeepTreeGAN introduces some major changes: The branching (Figure 1) uses two FFNs instead of matrix multiplication: one for increasing the number of nodes, the other for reducing the number of features.As the new leaves have the same size as their ancestors, a single Dimensionality Reduction FFN can be applied to all nodes of the tree.After the branching, both models implement a graph convolution step, passing information from the ancestors to their children.In tree-GAN, the leaves are updated using the leaf itself and its ancestors, each multiplied with a separate trainable matrix.In DeepTreeGAN, all nodes in the tree have the same dimension at all times.Thus, all nodes can be updated at the same time, using a single FFN in the Ancestor MPL (Figure 2).

Preprocessing & Postprocessing
For the training, η rel and ϕ rel are scaled separately to a normal distribution.The p rel T distribution is transformed using the Box-Cox transformation with standardizing implemented in [40].To produce samples, the inverse scaling is applied to the output of the generator.The variables of the constituents of the top quark dataset are shown in Figure 3.The generated distributions of the constituents closely match the distribution of the data.The jet variables shown in Figure 4 are significantly harder to model.Our model successfully reproduces the double peak structure in the relative jet mass, showing a significant advantage over tree-GAN as shown in [21,Figure 3].
In Table 1 the achieved metrics are shown.While our model shows an advantage over the MPGAN baseline on the top quark dataset and an at least similar performance on the light quarks' dataset, the performance is worse for the gluon dataset.At the current stage, the generator always produces the full 30 constituents for each of the events.Thus, the performance for events with less than 30 constituents is limited, as is frequently the case for the gluon jets.The models and code are available on GitHub1 .The 'Limit' row gives the measured in-sample distance by sampling datasets from the training dataset and calculating the metrics.Note that the MPGAN model was selected for W M 1 , and the testing sample was used to compute the given metrics.The W 1 metrics are calculated with 10000 samples, FPND is calculated with 50000 samples.

Conclusion
This paper presents a new model (DeepTreeGAN) for generating large point clouds using a novel upscaling method for point clouds.It successfully replicates the complex distributions of the JetNet datasets with high fidelity (Section 3) and demonstrating its suitability for generating high-dimensional point clouds.The model's tree-based structure promises good scaling behavior for even larger point clouds, such as particle showers in high-granularity calorimeters.However, the preliminary version presented here uses the MPGAN critic that scales with O(points 2 ) and thus, requires a new critic with better scaling behavior.In addition, the generator will need to be able to generate a variable number of points for each event, especially for the JetNet datasets with up to 150 particles and the calorimeter use case.This will be the subject of future publications.

Figure 2 :
Figure2: Ancestor Message Passing Layer.The edges in the graph are constructed so that each node receives messages from each of its ancestors as well as itself.In the given example node 4 is updated with the aggregate of the messages from nodes 1, 2 and 4. As a messagepassing algorithm, GINConv[35] -as implemented in PyTorch Geometric[36] -is chosen.For GINConv, the messages are the features of the source node (here: the ancestor).These messages are aggregated by summing over all messages addressed to the target node.These aggregates are added to the feature vector of the target node (scaled with a learnable weight) and passed through an FFN.The nodes are then updated with the output of the FFNs plus, as a residual connection, the original feature vector of the nodes themselves.

1 Figure 3 :
Figure 3: The global particle distributions for the top-quark test dataset, as described in section 1.1.
Branching Layer.In this example, the nodes of the 2nd level of the tree are produced and attached as the new leaves.With the nodes from level 1 (dark green / orange) as the parents, the children in level 2 (blue+green / red+yellow) are generated independently: StackFigure1:

Table 1 :
Figure 4: The distributions of the jet variables in the top-quark test dataset.Comparison of the proposed DeepTreeGAN to the MPGAN baseline with metrics introduced in [21].Significant advantages for a model in a given metric are highlighted bold.