Author: Jian Park
Mentors: Kilian Lieret, Gage DeZoort
Organization: CERN-HSF(High-Energy Physics Software Foundation)
Project Source Repository: GNN Tracking
Project Title: Graph Construction for Charged Particle Tracking using Point Cloud Networks
Graph neural networks (GNNs) have shown great potential in solving the complex problem of charged particle tracking, which is an essential task in high-energy physics. Charged particle tracking entails reconstructing the trajectory of charged particles produced after particle collisions, and the scalability of GNNs could enable this task to be performed with high efficiency.
Creating particle tracking models with improved accuracy brings practical benefits to particle physics researchers who study the behavior of particles generated after collision events, and my work during this program explores a promising avenue for performing this task.
The current work being performed within this project involves exploring different model architectures, training procedures, and loss functions in order to create a system that most accurately reconstructs particle tracks. My personal construbition to this project during the duration of GSoC involved exploring the capability of GravNet-based models in performing this same task. GravNet is a type of graph neural network that utilizes the k-Nearest Neighbor algorithm to generate graphs, upon which messages are propagated. The kNN is performed on a learned latent space, enabling the model to learn complex relations between different objects in the data [Qasim, et. al.]. One benefit to GravNet compared to current models is its simplicity, as previous solutions to particle tracking have involved an ensemble of separately trained models, while GravNet is trained in one step.
After developing and repeatedly refining the model/training hyperparameters throughout the course of GSoC, I have been able to gradually increase the performance of GravNet-based particle track construction models, and the most recent models achieve accuracy metrics on par with the other models being developed to solve this problem.
Throughout the duration of the GSoC program, I have also seperately contributed to improving the memory consumption and speed of clustering algorithms and loss functions used in the process of training and using charged particle tracking models. Details are provided below.
- Created multiple iterations of the GravNet model using PyTorch in order to test different hyperparameters combinations and alterations to the base GravNet architecture (explored the addition of multi-layer perceptrons in the message passing operation and the effects of adding multiple heads to the kNN operation).
- Improved upon the PyTorch implementation of the GravNet layer by replacing the PyTorch cluster implementation of the k-nearest neighbor algorithm with an implementation provided by the FAISS library.
- This change decreased the overall training time of GravNet models by 20%.
- The best performing GravNet model I have created and trained to date contained 1.7 million parameters, achieving 88.0% on the "double majority" accuracy metric and 61.5% on the "perfect" accuracy metric.
- The "perfect" accuracy metric measures the percentage of particle labels generated by the model where two hits share the same label if and only if they are from the same particle. "double majority" is a more lenient metric for measuring accuracy compared to the perfect metric [see here for more information].
- This result was achieved after training the model on a Nvidia A100 GPU for 24 hours.
- Designed and implemented a faster and more memory efficient way to compute object condensation loss, the main loss function used across multiple patricle tracking models currently being developed within the GNN Tracking project.
- Runtime for computing the loss function decreased by 50%, and memory consumption decreased approximately 80% compared to the prior impelementation when training GravNet models.
- Designed and implemented
DBSCANFastRescan, a faster method to search for optimal hyperparameters when performing DBSCAN clustering.- DBSCAN clustering is used frequently during the last step of most track construction models, and
DBSCANFastRescansped up the runtime performance during training by ~90%.
- DBSCAN clustering is used frequently during the last step of most track construction models, and
As the work on developing GravNet-based particle tracking models is ongoing research, there are several improvements that may increase the accuracy of these models. One avenue for further inquiry is to perform more hyperparameter optimization, specifically in regards to the number of GravNet layers contained inside these models. I was previously under the impression that ~4 layers was sufficient for completing the task of track construction, but recent models suggest that further increasing the size of these models may yield further accuracy improvements.
Another avenue for further inquiry is to address issues of instability that are experienced when training GravNet-based models. Often while training, large spikes are observed in the loss function for an unknown reason, and these spikes considerably worsen training performance. In addition to further exploring why these spikes occur, one solution to this problem may be to explore gradient clipping.
I will work on these tasks in the near future as I intend to continue contributing to this project after the completion of GSoC, and I hope persuing these improvements will further increase the accuracy of the GravNet models.
- This project was an excellent opportunity for me to learn more about PyTorch, PyTorch Geometric, and PyTorch Lightning. These new skills have enabled me to build machine learning models for personal projects and other jobs that I was working on during the summer.
- Working on the DBSCAN clustering and loss function optimization tasks have given me significant insight into how functions can be optimized for the GPU.
- It was quite challenging for me to switch from training models on my local device to training on the Princeton GPU clusters. However, instruction from my mentors enabled me to quickly transition and begin training my models on larger-scale datasets.
- Working on this project has taught me the value of using Weights and Biases in managing machine learning projects, and I hope to use this software when I work on future ml projects.
- I have gained a better understanding of Github and code formatting while making contributions to the main GNN Tracking repository.
- Repository for GravNet
- This repository contains all work I have done in relation to GravNet model training. This repository will be updated in the future as work on the this model continues.
- PR #327 - Added Remove_Isolated_Nodes
- PR #378 - Faster Way to Scan DBSCAN Eps
- PR #391 - Improve Memory Consumption of Condensation Losses
I would like to thank Kilian and Gage for their guidance and assistance throughout the duration of this program. I am especially thankful for their ability to find time in their already busy schedules to provide me with advice regarding the GravNet model training. I would also like to thank CERN-HSF for organizing this excellent program and providing me with a once-in-a-lifetime learning opportunity.