Semi-supervised Graph Neural Network for Particle-level Noise Removal
Abstract
The high instantaneous luminosity of the CERN Large Hadron Collider leads to multiple proton-proton interactions in the same or nearby bunch crossings (pileup).
Advanced pileup mitigation algorithms are designed to remove this pileup particle noise and improve the performance of physics observables crucial to the science goals.
This study applies the semi-supervised graph neural network to particle-level pileup noise removal, by identifying the particles produced from pileup.
The graph neural network is trained on charged particles with well-known labels, which can be obtained from simulation truth information or measurements from data, and inferred on neutral particles of which such labeling is missing.
This semi-supervised approach does not depend on the ground truth information from simulation and thus allows us to perform training directly on real data.
The performance with this approach is found to be consistently better than widely-used domain algorithms and comparable to a fully supervised training approach.
The study serves as the ﬁrst attempt at applying semi-supervised learning on pileup mitigation, and opens up a new direction of fully data-driven pileup mitigation techniques.
Introduction
The high instantaneous luminosity of the CERN Large Hadron Collider (LHC) enables studies of the deep mysteries of our universe, such as the nature of the Higgs boson and dark matter as well as the origin of the matter-antimatter asymmetry.
The enormous amount of data coming from increasing noisy particle collisions, recorded by increasingly complex detectors, pose various challenges to the data collection and analysis.
Multiple collisions in the same or nearby proton beam crossings lead to overlapping particle interactions is referred to as pileup (PU).
To achieve the desired physics sensitivity with the LHC data, the PU noise needs to be identiﬁed and mitigated effectively in order to identify signals of interest, i.e., those from the primary interaction associated with the leading vertex (LV).
Various PU mitigation techniques have been developed over the past decade to tackle this important yet challenging task.
The current widely-used domain algorithms for PU mitigation include SoftKiller (1) and PUPPI (2), both of which depend on manually designed heuristics.
More recently, deep learning (DL) techniques (3; 4; 5; 6) have been applied to extract particle features for the noise removal task.
By
Original
LV Chg LV Neu PU Chg PU Neu - -   4 0.0 0.
PUPPI
LV Chg
LV Neu
PU Neu - -   4 0.0 0.
Fully Supervised GNN
LV Chg
LV Neu
PU Neu - -   4 0.0 0.
SSL GNN
LV Chg
LV Neu
PU Neu exploring a higher dimensional phase space of particle-level and event-level features with different network architectures, these DL algorithms are shown to perform better than the classical domain ones. However, these algorithms are trained with fast simulation datasets, fully supervised with the knowledge of the LV (signal) and PU (noise) labels as the ground truth information. This approach has two primary shortfalls: (1) in more realistic simulations and real data, the label information is harder or impossible to retrieve and deﬁne and (2) accurate modeling of the pileup noise interactions are difﬁcult to obtain, and any discrepancies between simulation and data can lead to additional biases or performance degradation in downstream physics observable (7; 8; 9).
In this study, we explore a novel approach of applying a semi-supervised machine learning technique (SSL) in pileup mitigation, taking advantage of the fact that, although LV/PU labels are not available for all particles, they can be precisely determined for the charged subset of the particles.
A graph neural network (GNN) is ﬁrstly trained on labeled charged particles, with the input features from both labeled charged particles and unlabeled neutral particles.
The trained GNN is then inferred on neutral particles to estimate the probability of each being produced from LV.
A randomized feature-masking technique is developed to handle the features that are not shared by charged and neutral particles.
In contrast to the expert-level ﬁxed form heuristics in SoftKiller and PUPPI, the GNN-based approach allows a more powerful, ﬂexible, and expressive formula that is purely data-driven with parameters directly learned from data.
The effectiveness of the SSL approach is studied and conﬁrmed by comparing its performance against methods with a fully-supervised GNN with the same architecture and the expert-level heuristics.
We conduct our studies in different experimental conditions with different PU/noise levels, and ﬁnd no signiﬁcant performance drop between the fully supervised training and semi-supervised training.
In all cases, the GNNs achieve better performance than the expert-level heuristics.
Figure 1 is an example of one event visualization that shows GNN methods remove PU neutral particles more effectively than PUPPI.
The SSL approach has no dependence on any ground truth information from simulations, and can be applied directly to real collision data, avoiding data-simulation differences.
Our initial results are promising and encouraging, and we hope to adopt and validate this approach in the real LHC collision data in the near future.
Formulating Pileup Mitigation as a Semi-supervised Learning Problem
Semi-supervised learning (SSL) is a machine learning technique that takes advantage of not only the training samples but also the features of testing samples to train a model (10).
The pileup mitigation problem can be formulated as an SSL problem by utilizing features of both charged particles (with labels, for training) and neutral ones (without labels, for testing) given their geometric relationship.
For each proton bunch crossing, multiple charged and neutral particles are produced and scattered in the pseudorapidity and azimuthal angle η − φ space.
Because of the interaction dynamics, particles from the LV (signal label) tend to be more localized.
Therefore, an effective learning procedure of a model should not only depend on the self features of a particle, but also the features of its neighboring charged and neutral particles.
To leverage neighboring particle features, one graph is constructed per proton-bunch-crossing event to establish the relations between particles and their neighbors: Particles are viewed as nodes and two particles are linked if their distance in the η − φ space, ∆ R = (cid:112) (∆ η ) 2 + (∆ φ ) 2 , is smaller than a threshold. ∆ R = 0 . 8 is chosen for the experiments in this study.
In practice, multiple events are collected and processed, and each of them forms a graph.
[Figure] There is a particular challenge in pileup mitigation as opposed to traditional SSL problems, where the training samples (charged particles) and testing samples (neutral particles) may have their speciﬁc features, such as the encoding of particle types.
These features can not be naively removed as they provide useful information to determine neighbors’ labels.
To ensure the model trained on charged particles can be well applied to neutral particles, a dedicated training strategy is needed.
Formally the problem setting is as follows.
There are M graphs { G 1 , · · · , G M } .
Each observed graph contains two types of nodes: charged particles V C and neutral particles V N .
For each charged particle v ∈ V C , both the features x v and its label y v are observed, while for a neutral particle u ∈ V N , only its features x u are observed.
Our goal is to infer the labels of neutral particles.
Each node feature can be categorized into two parts, where the ﬁrst part contains feature values shared by charged and neutral particles, while the rest contain the charged-or-neutral-speciﬁc feature values.
The Proposed Model
Training the model consists of four steps as illustrated in Figure 2.
Step (a) graph construction has been introduced.
Next, we focus on steps (b)-(d).
Step (b) Random Selection and Feature Masking.
The main goal of step (b) is to deal with the potential features not shared between charged and neutral particles.
These features in our simulation datasets are LV-or-PU label information where only charged particles hold.
In real experiments, they may include other features such as particle types as charged and neutral particles by nature are in different categories.
We term these features as charged-speciﬁc features and neutral-speciﬁc features respectively.
To prevent the model from overﬁtting such features, in each graph each training epoch, a certain portion of charged LV and PU particles are randomly selected as training samples.
The charged-speciﬁc features of these particles will be replaced by neutral-speciﬁc features of randomly selected neutral particles.
Note that such a portion should neither be too large nor too small, as otherwise, either the neighboring features of a particle get corrupted, or too many epochs are needed to fully utilize all labeled particles for training.
About 10 % of charged particles are selected in order to balance preserving neighboring structures and training convergence/speed.
Table 1 includes the numbers of selected charged LV and PU particles per graph per epoch and the total number of charged LV and PU particles per graph.
The effect of the selected charged LV and PU particles ratio has been studied by varying it in a certain region around known charged LV and PU particles ratio, and the effect is found to be small.
Step (c) GNN Encoding and Step (d) Prediction.
Any GNN architecture, like (11; 12; 13), can be applied to our SSL framework, though we focus on a variant of the gated GNN model (14).
Because, as shown in Fig. 1, there are some LV particles surrounded by PU particles, we use gates to control the messages from the neighbors to better ﬁt the problem.
Let h kv denote the node v representation at [Table] k -th layer.
The GNN follows
Edge message: m uv = (cid:2)
[Equation]
u , h k −
[Equation]
v , ∆ η uv , ∆ φ uv , ∆ R uv (cid:3) , Aggregation: m v = (cid:80) u ∈ N ( v ) g uv m uv , where g uv =
Sigmoid ( W 1 m uv + b 1 )
Node-level gate:
q v =
Sigmoid ( W 2 [ h k − 1 v , m v ] + b 2 ) Node update: h kv = ReLU ( q v ( W 3 h k − 1 v + b 3 ))
+ (1 − q v )( W 4 m v + b 4 )) , where ∆ η, ∆ φ, ∆ R are the geometric features that characterize the difference between the spatial coordinates η, φ of two particles and their distance (cid:112) ∆ η 2 + ∆ φ 2 respectively.
The node representations are initialized as particle features that in our experiments include the particle transverse momentum p T and one-hot label encoding, that is, (1 , 0 , 0) for LV charged particles, (0 , 1 , 0) for PU charged particles, (0 , 0 , 1) for neutral particles and masked charged particles.
The node representations of the selected particles in the ﬁnal layer of GNN are put through a multi-layer perceptron with two hidden layers to make the ﬁnal prediction.
Results
The datasets used in this study are generated using PYTHIA 8.223 (15) and DELPHES 3.3.2 (16).
Two pileup conditions are chosen to be studied: the number of pileup interactions n PU = 80 and 140, where the latter represents a more noisy experimental environment.
Table 1 shows the average numbers of different types of particles under different pileup conditions.
The experiments are designed to demonstrate the effectiveness of the model trained via SSL and its ability to be adapted to different n PU levels.
The model is trained and tested under three scenarios: (a) training the model on n PU = 80 with inference on n PU = 80 , (b) training the model on n PU =with inference on n PU = 140 , and (c) training the model on n PU = 80 with inference on n PU = 140 .
Experimental settings (a) and (b) are to demonstrate that the model trained on charged particles with
SSL can perform well on neutral particles in testing and work under different n PU levels.
Speciﬁcally, the SSL model is compared with the fully supervised model, which has the same architecture but is trained directly on neutral labels, and with the baseline algorithm PUPPI.
Experiment (c) is to check the adaptation ability of the model in different n PU conditions.
To evaluate the fully supervised model performance, the particles used for training and testing must come from different events, though this is not necessary in the SSL training.
Therefore, for experiments on n PU = 80 , there are 3000/1000/1000 events for training/validation/testing.
When n PU = 140 , 1000/400/800 events are used for training/validation/testing.
For the n PU = 140 scenario, there are more particles per event, so the total number of events is reduced to maintain reasonable memory usage.
During training, the model is trained until convergence, which normally takes about 5 epochs.
The testing results demonstrate the success of the model, and are shown in Fig.
3 ROC curves with
AUC scores.
In Fig.
3, (a) and (b) both indicate the SSL model outperforms the baseline algorithm
PUPPI (by 7 .
41% under n PU =80 and 6 .
22% under n PU =140).
Furthermore, the SSL performance is very similar to the fully-supervised model(decays by 1 .
41% under n PU =80 and 0 .
51% under n PU =140).
The comparison between (a) and (c) in Fig.
3 demonstrates the model can adapt between different n PU s since the model only degrades by 0 .
22% under SSL from (b) to (c) in Fig.
[Figure]