ABCNet: an attention-based method for particle tagging

In high energy physics, graph-based implementations have the advantage of treating the input data sets in a similar way as they are collected by collider experiments. To expand on this concept, we propose a graph neural network enhanced by attention mechanisms called ABCNet. To exemplify the advantages and flexibility of treating collider data as a point cloud, two physically motivated problems are investigated: quark–gluon discrimination and pileup reduction. The former is an event-by-event classification, while the latter requires each reconstructed particle to receive a classification score. For both tasks, ABCNet shows an improved performance compared to other algorithms available.


Introduction
One of the main goals in modern machine learning is to be able to extract the maximum amount of information available from a data set. Successful implementations take advantage of the data structure for model building. In high energy physics (HEP), particle collisions in experiments are reconstructed by combining the energy deposits left by particles after crossing different parts of a detector. The information provided by sub-detectors can be further combined to give a full description of each particle produced. At the Large Hadron Collider (LHC) [1], jets are ubiquitous objects produced in proton-proton collisions. Jets are the byproducts of the hadronisation of quarks and gluons, resulting in an often collimated spray of particles. After each collision, O(1000) or more particles can be produced, making the task of identifying the original hard scattering objects challenging. The luminosity increase at the LHC will also increase the amount of multiple interactions per bunch crossing (pileup). For instance, event collisions recorded thus far by the ATLAS [2] and CMS [3] detectors at LHC measured an average of about 30 extraneous interactions. With the future upgrade, up to 200 pileup events per bunch crossing are expected, requiring new methods for particle identification and pileup suppression. In this paper, a new method for event classification in HEP is introduced. The attention-based cloud net (ABCNet) takes into account the data structure recorded by particle collision experiments, treating each interaction as an unordered set of points that defines a point cloud. This description is advantageous since the byproducts of each particle collision are treated in a similar fashion as they are collected by particle detectors. To enhance the extraction of local information, an attention mechanism is used, following closely the implementation developed in [4]. Attention mechanisms have proved to boost performance for different applications in machine learning by giving local and global context to the learning procedure. To show the performance and flexibility of the model, two critical problems are investigated: quark-gluon discrimination and pileup mitigation.

Related works
The main novelties introduced by ABCNet are the treatment of particle collision data as a set of permutation invariant objects, enhanced by attention mechanisms to filter out the particles that are not relevant for the tasks we want to accomplish. The usage of graph-based machine learning implementations is still a new concept in particle physics. Nevertheless, new implementations have already been proposed with promising results. ParticleNet [5] uses a similar approach, using point clouds for jet identification. The main difference between ABCNet and ParticleNet is that ABCNet takes advantage of attention mechanisms to enhance the local feature extraction, allowing for a more compact and efficient architecture. A theoryinspired approach was also developed in the framework of Deep Sets [6] using an infrared and collinear safe basis, developed in the context of Energy Flow Networks [7]. A messagepassing approach for jet tagging discussed in [8]. Interaction networks were also studied in the context of high-mass particle decays with JEDI-net [9]. Other graph-based implementations have also been presented in the context of signal and background classification [10,11], particle track reconstruction [12], and particle reconstruction on irregular calorimeters. [13]. In the context of pileup rejection, the GGNN implementation [14] shows promising results by combining graph nodes with GRU cells.

GAPLayer
ABCNet follows closely the implementation described for GAPNet [4], with key differences to adapt the implementation to our problems of interest. For clarity, the description of the essential aspects of the implementation are described. The key aspect of GAPNet is the development of a graph attention pooling layer (GAPLayer) using the edge convolution operation proposed in [15], which defines a convolution-like operation on point clouds together with attention mechanisms to operate on graph-structured data described in [16]. The point cloud is first represented as a graph with vertices represented by the points themselves. The edges are constructed by connecting the points to their k-nearest neighbours, while the edge features, y i j = (x i − x i j ), are taken as the difference between features of each point x i and its k-neighbours x i j . A GAPLayer is constructed by first encoding each point and edge to a higher-level feature space of dimension F using a single-layer neural network (NN), with learnable parameters θ , in the following form: where h() denotes the single-layer neural network operation. Self-and local coefficients are created by passing the transformed points and edges to a single-layer NN with output dimension of size one. Finally, the attention coefficients c i j are created by combining the newly created coefficients in the following way: where the nonlinear LeakyRelu operation is applied to the output of the sum. To align the attention coefficients between different points, a Softmax normalisation is applied to the coefficients c i j . At this moment, each point is associated with k attention coefficients. To compute a single attention feature for each point, a linear combination with a nonlinear activation function is defined asx To enhance the stability of the determination of the coefficientsx i , a multihead mechanism can be used. A M-head process repeats the same procedure described above, determininĝ Due to stackability properties, a GAPLayer output can be further used as an input to a subsequent GAPLayer or multilayer perceptron (MLP).

Classification: quark-gluon tagging
Quark-gluon tagging refers to the task of identifying the origin of a jet as produced from the hadronisation of a gluon or a quark. The data set used for the studies are available from [7]. It consists of stable particles clustered into jets, excluding neutrinos, using the antik T algorithm [17] with R = 0.4. The quark-initiated sample (signal) is generated using a Z(νν) + (u, d, s) while the gluon-initiated data (background) are generated using Z(νν) +g processes. Both samples are generated using Pythia8 [18] without detector effects. Jets are required to have transverse momentum p T ∈ [500, 550] GeV and rapidity |y| < 1.7 for the reconstruction. For the training, testing and evaluation of the method, the recommended splitting is used with 1.6M/200k/200k events, respectively. For every reconstructed jet, up to 200 constituents are saved. Each constituent contains the four momentum and the expected particles type (electron, muon, photon, or charged/neutral hadrons). A typical jet has O (10) to O(100) particles. To simplify the implementation, ABCNet uses the first 100 constituents sorted by p T from highest to lowest. If the jet has less than 100 constituents, the event is padded with zeros; if there are more than 100 constituents, the event is truncated.
To enhance the non-local information extraction, global features can also be added to ABCNet. The approach is similar to the one described in [19], where global information is used to parameterise the network, improving the generalisation and performance as a function of the global parameters.
The features used to describe each constituent are listed in Table 1.

Network architecture
The network layout used is shown in Fig. 1. The first step is to calculate the distances between the constituents in the pseudorapidity-azimuth (η − φ) space of the form R = η 2 + φ 2 . From the distances, we create the first GAPLayer by associating each particle to its nearest 10 neighbours. While different choices for k were tested, the overall performance did not improve with the addition of more neighbours. The encoding channel size of Logarithm of the ratio between the constituent's E and the jet E R Distance in the η − φ space between the constituent and the jet PID Particle type identifier as described in [20] m(jet) Jet mass p T (jet) Jet transverse momentum the GAPLayer F is selected to be 32 with a 1-head. The attention features created by the GAPLayer are then passed through two MLPs with node sizes (128,128). The distances used for the second GAPLayer are calculated using the full-feature space produced in the output of the last MLP, allowing the network to learn distances in the transformed feature space. To achieve a robust estimation, the encoding channel size is selected to be 64 with the number of heads determined to be two. The newly created attention features are passed through two MLPs of node sizes each of 128. In parallel, ABCNet also takes additional global inputs in the form of the jet mass and transverse momenta. The global inputs are first transformed by means of a single-layer MLP with small node size of 16. The two graph features and the output of each MLP are concatenated with the transformed global features and fed to a MLP of node size 128. An average pulling is applied, and the result is further passed to 2 additional MLPs of node sizes (128,256) interleaved by two dropout layers. A Softmax operation is applied to the output result.

Results
The performance of ABCNet is compared to the methods implemented in [5] and [7], using the same data set.  Table 2. Even though the accuracy obtained by ABCNet is numerically the same as the one reported by ParticleNet, ABCNet excels on the other figures of merit, improving the background rejection, at 30% signal efficiency, by 15-20%. The use of attention coefficients allow the model complexity of ABCNet to be reduced, having 40% less parameters compared to ParticleNet.  The uncertainty quoted corresponds to the standard deviation of nine trainings with different random weight initialisation. If the uncertainty is not quoted, then the variation is negligible compared to the expected value

Visualisation
A simple way to check what ABCNet is learning is to look at the self-coefficients of each point of the point cloud. First, we pre-processes the images in a similar fashion as [21], using the following steps: -Centre: All jet images are translated in the η − φ space to a common centre at (0,0). The centre of the jet is taken as its p T -weighted centroid. -Particle scale: Each particle constituent has its transverse momentum scaled such that jet p T,i = 1, where i is the i-th constituent of the jet. Other steps were adopted in [21]; however, since the goal is to have a simple visual cue, they were not used. The resulting jet images are shown in Fig. 2 for quark-and gluon-initiated jets on the upper and lower rows, respectively. The leftmost images correspond to the jets after the pre-processing. The subsequent columns show the same distribution, but only considering  Fig. 2 Distribution of the p T -scaled distribution of the jet constituents averaged over all images in the test sample. The leftmost images are the quark (top) and gluon (bottom) jet averages after the pre-processing. The first 5% of the jet constituents with the highest self-attention coefficients for the first and second GAPLayers are shown on the images in the centre and right, respectively particles whose self-attention coefficients, resulting from the first (middle column) and second (right column) GAPLayers, are higher than a certain value. This value is chosen such that only the first 5% of all particles with the largest self-attention coefficients are selected. The self-coefficients from the first GAPLayer have the effect of giving higher attention to high-p T particles, while soft-QCD with large angular variation has less importance. The second GAPLayer, where nearest neighbours are calculated in the feature space, has different distributions for quark-initiated and gluon-initiated jets. Quark-initiated jets have the highest coefficients in a confined radius with ∼ R = 0.1 around the centre, while gluon-initiated coefficients spam a bigger area around the centre with ∼ R = 0.3. That behaviour is expected since gluon jets have a larger colour factor compared to quark jets, typically resulting in a broader angular distribution compared to quark jets.

Pileup reduction using part segmentation
Another crucial problem in particle physics is how to identify the particles originated from high-p T collisions and separate them from unwanted additional interactions. Two traditional methods to accomplish this task are the SoftKiller [22] and the Pileup Per Particle Identification (PUPPI) [23] algorithms. These two algorithms are chosen since they represent the most common algorithms for pileup mitigation at the LHC. To test the performance of ABCNet in this context, we change the scope of a single jet classifier to a particle-by-particle classification (part segmentation). In this case, a probability is estimated for object, determining how likely it is for each particle to originate from the leading vertex (LV). The sample used for this study is available from [24], containing a set of qq light-quark-initiated jets coming Table 3 Variable description for each feature used to define a point in the point cloud implementation for the pileup mitigation problem. The latter two features are the global information added to the network Variable Description η Particle's pseudorapidity φ Particle's azimuthal angle log p T Logarithm of the particle p T Q Boolean flag identifying if the particle is charged Logarithm of the ratio between the particle p T and the associated jet p T log E

E(jet)
Logarithm of the ratio between the particle E and the associated jet E w PUPPI PUPPI weight for the particle w SoftKiller Boolean flag identifying if the particle passes the SoftKiller p T requirement NPU number of pileup interactions NPART Number of reconstructed particles associated with jets from the decay of a scalar particle with mass m φ = 500 GeV. The samples were generated using Pythia8 at √ s = 13 TeV. The pileup events were generated by overlaying soft QCD processes onto each event. Stable particles are clustered into jets, excluding neutrinos, using the anti-k T algorithm with R=0.4. At parton level, a p T requirement of at least 95 GeV was applied. Only jets satisfying p T >100 GeV and η ∈ [−2.5, 2.5] are considered. For each event, up to two leading jets as ordered in p T are stored. Two thousand events are generated, each with a different number of pileup interactions (NPU) ranging from 0 to 180. For the training and testing samples, events are randomly selected from the generated samples according to a Poisson distribution with average pileup <NPU> = 140, motivated by the expected pileup levels for future collisions at the LHC. The training and evaluation are done with 80% and 10% of the events with <NPU> = 140, respectively. For testing, two samples are created: one corresponding to the remaining 10% of the events and <NPU> = 140 and the other a sample with independent events generated at different NPU levels. For each event, up to 500 particles are stored as long as they are matched to one of the two leading jets. The features used to define each particle are described in Table 3. The feature choice is similar to the ones used for the classification task. The main difference is that for this sample, the PID information is not available, but replaced by a flag that identifies if a particle is charged or not. Since more than one jet can be reconstructed, a global zero is used for all events, instead of using the jet axis as a reference point. While no selection is applied to the particles used in ABCNet, the PUPPI weights and the SoftKiller decision flag are also used as input features. The global information added to the parameterisation is NPU and the number of reconstructed particles associated with jets.

Network architecture
The network architecture for the part segmentation problem is similar to the setup used previously. The main differences are: -Number of considered neighbours increased from 10 to 50. -Additional MLPs after the attention features and after the pooling layer.
-Usage of only 1-head GAPLayers. The increase in expansion of neighbours and MLPs is chosen to increase the model's capacity to cover the larger amount of points per event. The architecture is shown in Fig. 3.

Results
The performance of ABCNet is compared to the performance achieved using PUPPI and SoftKiller. The default parameters for those methods are the same as the ones used in [24]: R 0 =0.3, R min = 0.02, w cut = 0.1, p T cut (NPU) = 0.1 + 0.007 × NPU (PUPPI), grid size = 0.4 (SoftKiller). First, the jet mass is reconstructed with the <NPU> = 140 evaluation sample, applying the different mitigation algorithms. Inspired by PUPPI, the output probabilities from ABCNet are used to reweight the four-momentum of each particle. The reconstructed dijet mass and the dijet mass resolution are shown in Fig. 4. The resolution is defined as: In Table 4, the width of the jet mass resolution, extracted by fitting the distributions in Fig. 4 (right) with a Gaussian function, is also listed.
ABCNet improves jet mass resolution compared to both PUPPI and SoftKiller by 75% and 83%, respectively. The robustness of each algorithm is also tested by comparing. The Pearson linear correlation coefficient (PCC) between the true jet mass and corrected jet masses for different NPU is generated. Figure 5 shows the result of the comparison using the test sample with NPU from 0 to 180. To investigate the power of ABCNet to generalise, a training sample with <NPU> = 20 is created and trained using the same architecture described previously. For both trainings, ABCNet shows a superior performance compared to PUPPI and SoftKiller for the entire NPU range. Furthermore, ABCNet is also remarkably robust for  Fig. 4 Distribution of the dijet mass using the different pileup mitigation algorithms (left) and the jet mass resolution (right). A narrower resolution peak means better performance. All distributions are normalised to unit pileup variations outside the training region due to the addition of the global parameters to the method.

Training details
ABCNet is implemented using Tensorflow v1.4 [25]. A Nvidia GTX 1080 Ti graphics card is used for the training and evaluation steps. For all tasks described in this paper, the Adam optimiser [26] is used. The learning rate starts from 0.001 and decreases linearly by a factor 10 every seven epochs, until reaching a minimum of 1e-7. The training is performed with a mini batch size of 64 to a maximum number of 50 epochs. The epoch with the highest accuracy on the evaluation is saved in the case of the quark-gluon classification task. For the pileup identification, the epoch with the lowest loss is stored.

Conclusion
In this document, a new machine learning implementation for data classification in HEP is introduced. The attention-based cloud net (ABCNet) takes advantage of the data structure commonly found in particle colliders to create a point cloud interpretation. An attention mechanism is implemented to enhance the local information extraction and provide a simple way to investigate what the method is learning. To capture the global information, direct connections for global input features can be directly added. ABCNet can be used for event-by-event classification problems or generalised to particle-by-particle classification. To exemplify the architecture flexibility, two example problems are investigated: quark-gluon classification and pileup mitigation. For both problems, ABCNet achieved an improved performance compared to other available methods. By using a graph architecture and interpreting each point in a point cloud as a particle, ABCNet can be readily adapted to other applications in HEP like jet-flavour tagging, boosted jet identification, or particle track reconstruction.