[Paper Review] Hierarchical Graph Neural Networks for Particle Track Reconstruction

Ryan Liu, P. Calafiura+4 more|arXiv (Cornell University)|Mar 3, 2023

Machine Learning and Data Classification4 citations

TL;DR

This paper proposes a Hierarchical Graph Neural Network (HGNN) with a learnable pooling mechanism (GMPool) to improve particle track reconstruction in high-luminosity LHC environments. By creating super-nodes that represent tracks and enabling message passing across disconnected components, HGNN achieves 97.86% tracking efficiency and 36.71% fake rate on the TrackML dataset—surpassing flat GNNs and demonstrating robustness to edge inefficiencies.

ABSTRACT

We introduce a novel variant of GNN for particle tracking called Hierarchical Graph Neural Network (HGNN). The architecture creates a set of higher-level representations which correspond to tracks and assigns spacepoints to these tracks, allowing disconnected spacepoints to be assigned to the same track, as well as multiple tracks to share the same spacepoint. We propose a novel learnable pooling algorithm called GMPool to generate these higher-level representations called "super-nodes", as well as a new loss function designed for tracking problems and HGNN specifically. On a standard tracking problem, we show that, compared with previous ML-based tracking algorithms, the HGNN has better tracking efficiency performance, better robustness against inefficient input graphs, and better convergence compared with traditional GNNs.

Motivation & Objective

• Address the limitations of flat GNNs in particle tracking, particularly broken and merged tracks due to insufficient message passing across disconnected components.
• Reduce sensitivity to inefficient input graphs—common in high pile-up HL-LHC conditions—by introducing a hierarchical structure that enhances long-range connectivity.
• Develop a novel learnable pooling mechanism (GMPool) to generate super-nodes representing tracks, enabling flexible, differentiable clustering of spacepoints.
• Introduce a bipartite matching loss function tailored for track reconstruction, improving track candidate assignment and reducing false positives.
• Demonstrate that hierarchical GNNs can outperform flat GNNs and edge-based baselines in tracking efficiency, robustness, and convergence speed on real-world datasets.

Proposed method

• Propose GMPool, a learnable graph pooling method that clusters spacepoints into super-nodes using a differentiable KNN-based super-graph, enabling message passing between disconnected components.
• Construct a hierarchical graph by applying GMPool to generate higher-level representations (super-nodes) that correspond to reconstructed tracks.
• Use a bipartite classifier to match spacepoints (nodes) to super-nodes (tracks), with a novel loss function designed to optimize track assignment and reduce false matches.
• Integrate the HGNN into a pipeline similar to GNN4ITk, using metric learning for initial graph construction and edge filtering, followed by hierarchical refinement.
• Train the model end-to-end using hinge embedding loss for graph construction and a custom bipartite matching loss for track assignment.
• Apply KNN to build a super-graph among super-nodes to restore connectivity and enable long-range information flow, overcoming the limitations of standard GNN message passing.

(a) Flat GNN

Experimental results

Research questions

• RQ1 Can a hierarchical GNN architecture improve tracking efficiency by enabling message passing across disconnected components in particle track reconstruction?
• RQ2 Does the proposed GMPool pooling mechanism outperform existing graph pooling methods in terms of flexibility, differentiability, and adaptability to variable event sizes?
• RQ3 To what extent can the HGNN recover broken tracks and reduce merged tracks compared to flat GNNs and edge-based classifiers?
• RQ4 How robust is the HGNN to edge inefficiencies in the input graph, particularly under high pile-up conditions common in HL-LHC?
• RQ5 Can the proposed bipartite matching loss function improve track candidate assignment and reduce fake rates compared to standard connected components or edge-based methods?

Key findings

• The HGNN achieved 97.86% tracking efficiency and 36.71% fake rate on the TrackML-full dataset, outperforming the flat E-GNN (94.61% efficiency, 47.31% fake rate) and EC-GNN (96.35% efficiency, 55.58% fake rate).
• The Bipartite Classifier HGNN (BC-HGNN) maintained high efficiency (98.55% at 0% edge removal) and low fake rate (1.23%) even when 40% of edges were removed, demonstrating strong robustness to graph inefficiency.
• In contrast, the Truth-CC baseline showed a sharp decline in efficiency (from 98.72% to 64.81%) and a rise in fake rate (from 5.87% to 53.12%) under 50% edge removal, highlighting the vulnerability of edge-based methods.
• The HGNN model converged faster and showed better robustness than traditional GNNs, particularly in handling disconnected components and recovering broken tracks.
• The GMPool pooling mechanism enabled variable-sized super-graphs and subquadratic scaling, supporting end-to-end training and adaptive clustering without fixed super-node counts.
• The method was integrated into the GNN4ITk pipeline and is publicly available on GitHub, enabling broader adoption in high-energy physics track reconstruction.

(b) Hierarchical GNN