EggNet: An Evolving Graph-based Graph Attention Network for Particle Track Reconstruction

TL;DR

EggNet tackles the heavy computational load of particle track reconstruction by eliminating the need for a pre-constructed graph and instead evolving graphs on-the-fly within a graph-attention framework. It adopts an object-condensation style objective, operating directly on hit point clouds and iteratively updating a KNN graph in a latent space, with final embeddings clustered by DBSCAN to form track candidates. The loss combines attractive and repulsive terms over edges in the final embedding space $p_{-1}$, specifically $l_ ext{ε} = y_ ext{ε} d_ ext{ε}^2 + (1 - y_ ext{ε}) \\max^2(0, m - d_ ext{ε})$ with $d_ ext{ε}$ the Euclidean distance and $m=1$, promoting correct edge connectivity. On TrackML, EggNet with multiple iterations achieves high track efficiency ($\sim$0.996), very low duplication and fake rates, and a total inference time of around $0.26$ seconds per event on an NVIDIA A100, outperforming methods that rely on pre-constructed graphs when enough iterations are used; this demonstrates scalable, end-to-end track reconstruction suitable for HL-LHC-scale data.

Abstract

Track reconstruction is a crucial task in particle experiments and is traditionally very computationally expensive due to its combinatorial nature. Recently, graph neural networks (GNNs) have emerged as a promising approach that can improve scalability. Most of these GNN-based methods, including the edge classification (EC) and the object condensation (OC) approach, require an input graph that needs to be constructed beforehand. In this work, we consider a one-shot OC approach that reconstructs particle tracks directly from a set of hits (point cloud) by recursively applying graph attention networks with an evolving graph structure. This approach iteratively updates the graphs and can better facilitate the message passing across each graph. Preliminary studies on the TrackML dataset show better track performance compared to the methods that require a fixed input graph.

Figures (5)

• Figure 1: The proposed end-to-end tracking pipeline. Starting with a point cloud, where each point corresponds to a hit, we run an EggNet model that iteratively constructs a graph and performs message passing. EggNet outputs the embedding for each node, and the track candidates are extracted from the DBSCAN clusters that are obtained in the node embedding space.
• Figure 2: The EggNet architecture. $i$ corresponds to each EggNet iteration, and $j$ corresponds to each message passing step. An EggNet iteration generally consists of a graph attention block and a KNN. The first iteration does not perform graph attention and the last iteration does not perform KNN.
• Figure 3: Edge-wise efficiency and purity in the KNN graphs as a function of $k$. The EggNet here consists of 5 iterations ($i \leq 4$). The upper bound for each metric is also shown in the plot.
• Figure 4: Track efficiency, duplication rate and fake rate with track candidates obtained from the DBSCAN clusters as a function of $\epsilon$. The EggNet here consists of 5 iterations ($i \leq 4$).
• Figure 5: Inference time for each event versus the number of hits. The inference time is also broken down different components, including Graph Attention, KNN and DBSCAN. The EggNet here consists of 5 iterations ($i \leq 4$).