Tracking with Graph �Neural Networks
DANIEL MURNANE�BERKELEY LAB, CERN
1
Part 1: Fundamentals
HIGHRR LECTURE WEEK, HEIDELBERG UNIVERSITY
SEPTEMBER 13, 2023
HighRR Lecture Week - Heidelberg University - September 13, 2023
OVERVIEW
2
HighRR Lecture Week - Heidelberg University - September 13, 2023
WARNING: BIAS AHEAD
3
HighRR Lecture Week - Heidelberg University - September 13, 2023
HIGH LUMINOSITY TRACK RECONSTRUCTION
4
HighRR Lecture Week - Heidelberg University - September 13, 2023
DATA SCIENCE IN THE DISCOVERY PIPELINE
Simulation
Reconstruction
Analysis
Matrix-element Calculation
Parton-shower / Hadronization
Detector�Simulation
Digitization
Topoclusters�& Spacepoints
Track Finding�& Fitting
Jet Tagging�& Vertexing
Particle ID�& Particle Flow
Calibration
Likelihood Fitting
Unfolding
Statistical Techniques, Bayesian Inference
Numerical �Integration
Markov Chain�Monte Carlo
Topological clustering
Kalman Filtering�& Fitting
Conformal Fits�& Hough Transform
HighRR Lecture Week - Heidelberg University - September 13, 2023
ML TODAY & TOMORROW IN THE DISCOVERY PIPELINE
Simulation
Reconstruction
Analysis
Matrix-element Calculation
Parton-shower / Hadronization
Detector�Simulation
Digitization
Topoclusters�& Spacepoints
Track Finding�& Fitting
Jet Tagging�& Vertexing
Particle ID�& Particle Flow
Calibration
Likelihood Fitting
Unfolding
Omnifold and Likelihood-free Inference
Generative Models:�GANs, VAEs, Normalizing Flows and Diffusion
Metric Learning, Object Condensation
Deep Full Event Reconstruction
CNNs, Graph Neural Networks & Transformers
Autoencoders�& Anomaly Detection
Symmetric ML�& Equivariance
HighRR Lecture Week - Heidelberg University - September 13, 2023
TASKS IN AN HL-LHC DETECTOR
7
Vertex Reconstruction
Jet Tagging
Pile-up Removal
Missing Energy Reconstruction
HighRR Lecture Week - Heidelberg University - September 13, 2023
TASKS IN AN HL-LHC DETECTOR
8
Vertex Reconstruction
Jet Tagging
Pile-up Removal
Missing Energy Reconstruction
These all require accurate �track reconstruction
HighRR Lecture Week - Heidelberg University - September 13, 2023
WHAT IS TRACK RECONSTRUCTION
9
Given set of hits from particles in a detector, assign label(s) to each hit.
Perfect classification: All hits from a particle (and only those hits) share the same label
HighRR Lecture Week - Heidelberg University - September 13, 2023
THE IMPORTANCE OF TRACKING
10
ATL-PHYS-SLIDE-2023-048
HighRR Lecture Week - Heidelberg University - September 13, 2023
THE COST OF TRACKING
11
https://cds.cern.ch/record/2729668/files/LHCC-G-178.pdf
HighRR Lecture Week - Heidelberg University - September 13, 2023
COMPUTE SCALING FOR HIGH LUMINOSITY
12
ATLAS Computing Requirements Over Time
In other words…
Computing power
Time, Energy, Number of Collisions
Predicted
capacity
Traditional
methods
(scale quadratically)
HL-LHC, 14 TeV
2027
3 billion collision/second
HighRR Lecture Week - Heidelberg University - September 13, 2023
WHY HIGH LUMINOSITY PHYSICS?
13
…
HighRR Lecture Week - Heidelberg University - September 13, 2023
TRACKING 101
14
HighRR Lecture Week - Heidelberg University - September 13, 2023
THE GOAL OF TRACKING
The following slides have material borrowed from Heather Gray’s excellent talk @ Zurich: https://indico.cern.ch/event/504284/
15
HighRR Lecture Week - Heidelberg University - September 13, 2023
WITHOUT A DETECTOR...
Particles are well-described by helices in a magnetic field
But once they are in a detector…
16
HighRR Lecture Week - Heidelberg University - September 13, 2023
TYPES OF INTERACTION
17
HighRR Lecture Week - Heidelberg University - September 13, 2023
REPRESENTING A HELIX
18
HighRR Lecture Week - Heidelberg University - September 13, 2023
TRACKING TERMINOLOGY
19
HighRR Lecture Week - Heidelberg University - September 13, 2023
TRADITIONAL TECHNIQUES FOR TRACKING
20
HighRR Lecture Week - Heidelberg University - September 13, 2023
KALMAN FILTER: TRACKING AS NAVIGATION
21
https://www.kalmanfilter.net/multiSummary.html
HighRR Lecture Week - Heidelberg University - September 13, 2023
KALMAN FILTER: TRACKING AS NAVIGATION
22
HighRR Lecture Week - Heidelberg University - September 13, 2023
KALMAN FILTER: TRACKING AS NAVIGATION
23
https://cds.cern.ch/record/1281363/files/ATLAS-CONF-2010-072.pdf
HighRR Lecture Week - Heidelberg University - September 13, 2023
COMBINATORIAL KALMAN FILTER
24
https://www.researchgate.net/publication/344039130_Pattern_Recognition_and_Reconstruction
HighRR Lecture Week - Heidelberg University - September 13, 2023
TRACKING AS GRAPH SEGMENTATION
25
HighRR Lecture Week - Heidelberg University - September 13, 2023
COMPUTE SCALING FOR HIGH LUMINOSITY
26
ATLAS Computing Requirements Over Time
ML Image Classification Efficiency Over Time
62 million�parameters
5 million�parameters
HighRR Lecture Week - Heidelberg University - September 13, 2023
TEASER: GRAPH-BASED PIPELINE FOR TRACK RECONSTRUCTION
27
HighRR Lecture Week - Heidelberg University - September 13, 2023
HOW SHOULD WE REPRESENT PARTICLE COLLISIONS?
Assuming we want to use deep learning, how can we represent a particle collision?
28
Image?
Sequence?
Set/Point Cloud?
For event collision as point cloud, with relationships between points, this is a graph.
HighRR Lecture Week - Heidelberg University - September 13, 2023
A COLLECTION OF NODES
29
NODE
WHAT IS A GRAPH?
HighRR Lecture Week - Heidelberg University - September 13, 2023
AND EDGES
30
EDGE
WHAT IS A GRAPH?
HighRR Lecture Week - Heidelberg University - September 13, 2023
NODES + EDGES = DOUBLETS
31
DOUBLET
WHAT IS A GRAPH?
HighRR Lecture Week - Heidelberg University - September 13, 2023
NODES CAN HAVE FEATURES
32
NODE FEATURE
e.g. “West Oakland”
WHAT IS A GRAPH?
HighRR Lecture Week - Heidelberg University - September 13, 2023
EDGES CAN HAVE FEATURES
33
EDGE FEATURE
e.g. “Under Maintenance
– Single Track”
WHAT IS A GRAPH?
HighRR Lecture Week - Heidelberg University - September 13, 2023
THE WHOLE GRAPH CAN HAVE FEATURES
34
GRAPH FEATURE
e.g. “Sunday Timetable”
HighRR Lecture Week - Heidelberg University - September 13, 2023
GRAPHS ARE A NATURAL WAY TO REPRESENT TRACKS
35
Given hits on �layers of a detector
x direction
y direction
HighRR Lecture Week - Heidelberg University - September 13, 2023
GRAPHS ARE A NATURAL WAY TO REPRESENT TRACKS
36
Connect the
hits in some way
x direction
y direction
HighRR Lecture Week - Heidelberg University - September 13, 2023
GRAPHS ARE A NATURAL WAY TO REPRESENT TRACKS
37
x direction
y direction
HighRR Lecture Week - Heidelberg University - September 13, 2023
INTRO TO GRAPH NEURAL NETWORKS
38
HighRR Lecture Week - Heidelberg University - September 13, 2023
GRAPH NEURAL NETWORK APPLICATIONS
39
Travelling Salesman Problem
Knowledge Graph Comprehension
Image �Comprehension
Molecular �Chemistry
Protein�Comprehension
HighRR Lecture Week - Heidelberg University - September 13, 2023
GRAPH NEURAL NETWORK PROCEDURE
40
Node features
Encoder
Node channels
Message passing
Node Aggregation
Messages
Node channels
Task output layer
Node channels
HighRR Lecture Week - Heidelberg University - September 13, 2023
STEP 1: MESSAGE PASSING MECHANISM
41
0
2
1
3
4
Figure inspired by Koshi et. al.
0
2
1
3
4
Input channels
Encoded channels
HighRR Lecture Week - Heidelberg University - September 13, 2023
STEP 2: AGGREGATION
At each node:
Sum all messages
Note: Called isotropic �message passing. �Introduced as “Graph Convolution Network”
42
Figure inspired by Koshi et. al.
0
2
1
3
4
0
2
1
3
4
+
+
=
Encoded channels
Input channels
+
HighRR Lecture Week - Heidelberg University - September 13, 2023
EDGE CHANNELS
43
0
2
1
3
4
0
2
1
3
4
Encoded channels
Pre-encoded channels
HighRR Lecture Week - Heidelberg University - September 13, 2023
EDGE CHANNELS
44
0
2
1
3
4
0
2
1
3
4
Encoded channels
Pre-encoded channels
Paris
Hilton
France
Texas
Socialite
Small town
Capital
HighRR Lecture Week - Heidelberg University - September 13, 2023
THE LANDSCAPE OF GNNS
45
GNNs
CNNs
RNNs
Cellular
Automata
Transformers
Grid �adjacency
MLP + 1-pixel�Equivalency
Fully-�Connected
Linearly�Connected
= reduces to
Graphical
Automata
MLP + 1-pixel�Equivalency
Grid adjacency
GNNS IN TRACKING
46
HighRR Lecture Week - Heidelberg University - September 13, 2023
GNNS IN TRACKING
47
Tobias Kortus , Ralf Keidel and Nicolas R. Gauger, 2022
Våge, Liv CTD Proceedings 2022
HighRR Lecture Week - Heidelberg University - September 13, 2023
GNNS ELSEWHERE IN PARTICLE PHYSICS
48
HighRR Lecture Week - Heidelberg University - September 13, 2023
THE TOPOLOGY PROBLEM
49
HighRR Lecture Week - Heidelberg University - September 13, 2023
TRACKING WITH GRAPHS VS. POINT CLOUD
50
https://arxiv.org/pdf/2012.09164.pdf
HighRR Lecture Week - Heidelberg University - September 13, 2023
TOPOLOGY FOR MESSAGE PASSING AND SEGMENTATION
51
Message 0-1
Message 0-2
0
1
2
Probability 0-1
Probability 0-2
0
1
2
Topology for message passing
Topology for edge classification
HighRR Lecture Week - Heidelberg University - September 13, 2023
THREE WAYS TO BUILD A GRAPH
52
https://arxiv.org/abs/2103.16701
MODULE MAP
r
Embed into learned �latent space
Connect all spacepoints�within radius r
All spacepoint pairs�joined into graph
METRIC LEARNING
GEOMETRIC HEURISTICS
HighRR Lecture Week - Heidelberg University - September 13, 2023
THE GNN4ITK PIPELINE
53
HighRR Lecture Week - Heidelberg University - September 13, 2023
WHO IS INVOLVED?
54
HighRR Lecture Week - Heidelberg University - September 13, 2023
GRAPH REPRESENTATION OF AN EVENT
Given set of hits in a detector from particles, assign label(s) to each hit.
Perfect classification: All hits from a particle (and only those hits) share the same label
55
HighRR Lecture Week - Heidelberg University - September 13, 2023
PIPELINE OVERVIEW
56
1
2
3
Metric �Learning
Module�Map
or
Graph Neural�Network
Connected�Components
Connected�Components�+ Walkthrough
or
Graph�Construction
Edge�Classification
Graph�Segmentation
Hits
Graph
Edge Scores
Track Candidates
HighRR Lecture Week - Heidelberg University - September 13, 2023
DATASETS
57
TrackML
ATLAS ITk
HighRR Lecture Week - Heidelberg University - September 13, 2023
ATLAS ITK GEOMETRY
58
Cluster
Spacepoint
Silicon
Track
Pixel is trivial: Each spacepoint maps to one cluster, which can map to many particles
Strip: Each spacepoint maps to two clusters – one on either side of strip, which can map to many particles
*Thanks Noemi Calace
0: Pixel barrel
1: Pixel endcap
2: Strip barrel
3: Strip endcap
ATLAS ITk
2
3
0
1
HighRR Lecture Week - Heidelberg University - September 13, 2023
ATLAS ITK GEOMETRY
59
Cluster
Spacepoint
Silicon
Track
Ghost spacepoint: Incorrectly constructed from �clusters left by different particles
Cluster A
Cluster B
Cluster C
Cluster D
ATLAS ITk
HighRR Lecture Week - Heidelberg University - September 13, 2023
60
Metric �Learning
Module�Map
or
Graph Construction
Hits
Graph
1
HighRR Lecture Week - Heidelberg University - September 13, 2023
EDGE TRUTH DEFINITIONS
61
All Edges
Target particle
Non-target particle
ATLAS ITk
HighRR Lecture Week - Heidelberg University - September 13, 2023
MODULE MAP - DOUBLETS
62
Metric �Learning
Module�Map
or
Graph Construction
Hits
Graph
1
Modules
HighRR Lecture Week - Heidelberg University - September 13, 2023
MODULE MAP – TRIPLETS
63
Step 1
Step 2
Metric �Learning
Module�Map
or
Graph Construction
Hits
Graph
1
HighRR Lecture Week - Heidelberg University - September 13, 2023
METRIC LEARNING �INTUITION
64
MLP
MLP
Repulsive training
Source
Target
Source
Target
Attractive training
cat
cat
cat
dog
“Contrastive” hinge loss
HighRR Lecture Week - Heidelberg University - September 13, 2023
METRIC LEARNING
65
r
Embed into learned �latent space
Connect all space points�within radius r
All spacepoint pairs�joined into graph
Metric �Learning
Module�Map
or
Graph Construction
Hits
Graph
1
HighRR Lecture Week - Heidelberg University - September 13, 2023
METRIC LEARNING - FILTERING
66
Metric �Learning
Module�Map
or
Graph Construction
Hits
Graph
1
512
Radius
Graph
512
512
512
512
8
cell
Hinge Loss
512
512
512
Cross Entropy Loss
Metric Learning
Filtering
GNN
norm
norm
ATLAS ITk
HighRR Lecture Week - Heidelberg University - September 13, 2023
GRAPH CONSTRUCTION RESULTS
67
ATLAS ITk
HighRR Lecture Week - Heidelberg University - September 13, 2023
68
Graph
2
Graph Neural�Network
Edge Labeling
Edge Scores
HighRR Lecture Week - Heidelberg University - September 13, 2023
EDGE CLASSIFICATION WITH�GRAPH NEURAL NETWORK
69
INTERACTION
NETWORK
Battaglia, Peter, et al. "Interaction networks for learning about objects, relations and physics.“ 2016.
HighRR Lecture Week - Heidelberg University - September 13, 2023
LOSS FUNCTION DESIGN
70
Graph
2
Graph Neural�Network
Edge Labeling
Edge Scores
HighRR Lecture Week - Heidelberg University - September 13, 2023
GNN EDGE CLASSIFICATION RESULTS
71
ATLAS ITk
HighRR Lecture Week - Heidelberg University - September 13, 2023
GNN EDGE CLASSIFICATION RESULTS
72
ATLAS ITk
HighRR Lecture Week - Heidelberg University - September 13, 2023
73
3
Connected�Components
Connected�Components�+ Walkthrough
or
Graph Segmentation
Edge Scores
Track Candidates
HighRR Lecture Week - Heidelberg University - September 13, 2023
GRAPH PARTITIONING 101
74
3
Connected�Components
Connected�Components�+ Walkthrough
or
Graph Segmentation
Edge Scores
Track Candidates
HighRR Lecture Week - Heidelberg University - September 13, 2023
75
Component-based
Classified edges
Ignore cut edges
Label connected
components
Track #1
Track #2
1
1
1
2
2
2
E.g. connected components algorithm:
Path-based
E.g. walkthrough algorithm:
Classified edges,
Starting node
Choose high�score junctions
Remove a high-�scoring path
Track #1
1
1
1
1
2
2
3
3
3
Connected�Components
Connected�Components�+ Walkthrough
or
Graph Segmentation
Edge Scores
Track Candidates
TRACK CANDIDATES CONSTRUCTION
HighRR Lecture Week - Heidelberg University - September 13, 2023
76
Component-based
Classified edges
Ignore cut edges
Label connected
components
Track #1
Track #2
1
1
1
2
2
2
E.g. connected components algorithm:
Path-based
E.g. walkthrough algorithm:
Classified edges,
Starting node
Choose high�score junctions
Remove a high-�scoring path
Track #1
1
1
1
1
2
2
3
3
3
Connected�Components
Connected�Components�+ Walkthrough
or
Graph Segmentation
Edge Scores
Track Candidates
Both methods by construction�associate each hit with only one track
TRACK CANDIDATES CONSTRUCTION
HighRR Lecture Week - Heidelberg University - September 13, 2023
77
1. Connected Components
2. Walkthrough, a.k.a “Wrangler”
Classified edges
Cut score < 0.2
Track #1
Track #2
1
2
2
3
4
Label simple
candidates
Track #3
Assign longest path�as candidate
3
Connected�Components
Connected�Components�+ Walkthrough
or
Graph Segmentation
Edge Scores
Track Candidates
3
TRACK CANDIDATES CONSTRUCTION
HighRR Lecture Week - Heidelberg University - September 13, 2023
FINDING & FITTING PERFORMANCE
78
HighRR Lecture Week - Heidelberg University - September 13, 2023
TRACK MATCHING DEFINITIONS
79
Particle 1
Particle 2
Candidate 1
HighRR Lecture Week - Heidelberg University - September 13, 2023
TRACK RECONSTRUCTION RESULTS
80
ATLAS ITk
HighRR Lecture Week - Heidelberg University - September 13, 2023
TRACK RECONSTRUCTION RESULTS
81
ATLAS ITk
Observe that the GNN track candidates have fewer hits than CKF. Will return to this!
HighRR Lecture Week - Heidelberg University - September 13, 2023
TRACK FITTING 101
82
HighRR Lecture Week - Heidelberg University - September 13, 2023
83
HighRR Lecture Week - Heidelberg University - September 13, 2023
TRACK FITTING 101: FITTING WITH KALMAN FILTER
84
http://physik.uibk.ac.at/hephy/theses/diss_as.pdf
HighRR Lecture Week - Heidelberg University - September 13, 2023
FITTING PERFORMANCE OF TRACK CANDIDATES
85
HighRR Lecture Week - Heidelberg University - September 13, 2023