[Paper Review] Do graph neural networks learn traditional jet substructure?

Farouk Mokhtar, R. Kansal+1 more|arXiv (Cornell University)|Nov 17, 2022

Particle physics theoretical and experimental studies5 citations

TL;DR

This paper investigates whether Graph Neural Networks (GNNs), specifically ParticleNet, learn traditional jet substructure observables like the number of prongs in top quark jet tagging. Using layerwise relevance propagation (LRP), the authors show that trained ParticleNet models learn to prioritize edges connecting distinct subjets—particularly in three-pronged top quark jets—demonstrating that the model implicitly learns and utilizes known physics-based substructure features.

ABSTRACT

At the CERN LHC, the task of jet tagging, whose goal is to infer the origin of a jet given a set of final-state particles, is dominated by machine learning methods. Graph neural networks have been used to address this task by treating jets as point clouds with underlying, learnable, edge connections between the particles inside. We explore the decision-making process for one such state-of-the-art network, ParticleNet, by looking for relevant edge connections identified using the layerwise-relevance propagation technique. As the model is trained, we observe changes in the distribution of relevant edges connecting different intermediate clusters of particles, known as subjets. The resulting distribution of subjet connections is different for signal jets originating from top quarks, whose subjets typically correspond to its three decay products, and background jets originating from lighter quarks and gluons. This behavior indicates that the model is using traditional jet substructure observables, such as the number of prongs -- energetic particle clusters -- within a jet, when identifying jets.

Motivation & Objective

• To investigate whether state-of-the-art graph neural networks like ParticleNet learn traditional jet substructure observables used in high-energy physics.
• To understand the decision-making process of ParticleNet in top quark jet tagging by analyzing the relevance of learned graph edges.
• To determine if the model's attention is aligned with known physics features such as the number of prongs in a jet.
• To assess whether the model's learned connectivity patterns reflect known substructure differences between top quark and QCD jets.
• To evaluate the interpretability of GNNs in particle physics using explainable AI (XAI) techniques like layerwise relevance propagation (LRP).

Proposed method

• Trained a ParticleNet model on a benchmark dataset of 14 TeV jets from top quark decays and QCD dijets, with 200 leading constituents per jet.
• Applied layerwise relevance propagation (LRP) to assign relevance scores to edges in the graph, identifying which particle connections most influence the model's prediction.
• Computed the fraction of the N most relevant edges that connect different subjets (i.e., inter-subjet edges) across varying N values to assess substructure dependence.
• Used the anti-kT algorithm with R=0.8 to define subjets from jet constituents, enabling consistent substructure identification.
• Compared edge relevance distributions between untrained and trained models to observe changes in connectivity patterns during training.
• Averaged results over 10 randomly initialized untrained models and one fully trained model to ensure statistical robustness.

Experimental results

Research questions

• RQ1 Does the ParticleNet GNN learn to prioritize edges connecting different subjets when classifying top quark jets compared to QCD jets?
• RQ2 How does the distribution of relevant edges change during training, and does it reflect known substructure differences?
• RQ3 To what extent do the most relevant edges in the model correspond to physical observables like the number of prongs in a jet?
• RQ4 Is the model's attention on inter-subjet edges correlated with the presence of three-pronged substructure in top quark jets?
• RQ5 Can LRP-based interpretability reveal that the model implicitly learns physics-informed features such as N-subjettiness or energy correlation functions?

Key findings

• The trained ParticleNet model learns to assign higher relevance to edges connecting different subjets when classifying top quark jets compared to QCD jets.
• As the model trains, the fraction of the N most relevant edges that connect different subjets increases significantly for top quark jets, indicating learning of substructure-aware connectivity.
• The distribution of relevant edges for top quark jets shows a higher average ΔR separation between connected particles than for QCD jets, aligning with known three-pronged substructure.
• Untrained models treat both top quark and QCD jets similarly in terms of edge relevance, but the distinction emerges only after training.
• The model's learned edge relevance patterns are consistent with traditional substructure observables such as the number of prongs and N-subjettiness.
• The results demonstrate that ParticleNet implicitly learns and utilizes physics-informed features that are central to traditional jet substructure analysis.