GNNs for HL-LHC Tracking ExaTrkX @ Berkeley Lab Daniel Murnane - PowerPoint PPT Presentation

GNNs for HL-LHC Tracking ExaTrkX @ Berkeley Lab Daniel Murnane Office of BERKELEY LAB 1 Science

Goal Sub-second processing of HL-LHC hit data into: • Seeds (i.e. triplets) for further processing with traditional techniques, AND/OR • Tracks, where each hit is assigned to exactly one track Office of BERKELEY LAB 2 Science

The Current Pipeline Filter likely, Filter, convert adjacent to triplets Raw hit data doublets embedded Train/classify Train/classify doublets in tripets in Apply cut DBSCAN GNN GNN for seeds for track labels Office of BERKELEY LAB 3 Science

Dataset • “ TrackML Kaggle Competition” dataset • Generated by simulation • 8000 collisions to train on • Each collision has up to 100,000 hits of around 10,000 particles Office of BERKELEY LAB 4 Science

Dataset • Ideal final result is a “ TrackML score” 𝑇 ∈ 0,1 • All hits belonging to same track labelled with same unique label ⇒ 𝑇 = 1 • We use the barrel as a test case, and ignore noise Office of BERKELEY LAB 5 Science

Embedding + MLP Construction • Won’t give any detail (Nick’s talk next on embeddings) • Generally: For each hit in event, embed features (co-ordinates, cell direction data, etc.) into N- 1. dimensional space Associate hits from same tracks as close in N-dimensional distance 2. Score each hit within embedding neighbourhood against the “seed” hit at centre 3. Filter by score, to create a set of doublets for the neighbourhood 4. All doublets in event generate a graph, 5. converted to a directed graph (by ordering layers) Office of BERKELEY LAB 6 Science

Segmentation A full graph from the embedding does not fit on a single GPU. Therefore the event graphs are segmented, according to how large the GNN model is expected to be. Hard cut One-directional soft cut Bi-directional soft cut Office of BERKELEY LAB 7 Science

Previous ML Approaches • Tracks as images (CNN) • Tracks as sequences of points (LSTM) Office of BERKELEY LAB 8 Science 8

Graph Neural Network for Edge Classification Classify edges with score between [0,1] score > cut: true score < cut: fake Office of BERKELEY LAB 9 Science

Passing information around the graph gives it learning power • Can make a node “aware” of its neighbours by concatenating the neighbouring hidden features • Iterating this neighbourhood learning passes information around the graph • Can be considered a generalisation of a flat CNN convolution Office of BERKELEY LAB 10 Science

GNN Edge prediction architecture • Message Passing • Attention Message Passing + + • Attention Message Passing with Residuals • Attention Message Passing with Recursion Office of BERKELEY LAB 11 Science 11 11

GNN Edge prediction architecture • Message Passing • Attention Message Passing Have found + + best efficiency • Attention Message Passing & purity with Residuals performance. • Attention Message Passing with Recursion Office of BERKELEY LAB 12 Science 12 12

Edge attention architecture • Input node features • Hidden node features 𝑦 • Hidden edge features 𝑧 𝑨 • Edge score 0 • Attention aggregation • New hidden node features x n iterations • New hidden edge features (hyperparameter) • New edge score Office of BERKELEY LAB 13 Science 13 13

Edge attention architecture • Input node features • Hidden node features 𝑦 • Hidden edge features 𝑧 𝑨 • Edge score 0 • Attention aggregation • New hidden node features ℎ 1 x n iterations … • New hidden edge features ℎ 𝑜 0 • New edge score Office of BERKELEY LAB 14 Science 14 14

Edge attention architecture • Input node features ℎ 1 … • Hidden node features ℎ 𝑜 1 • Hidden edge features ℎ 1 • Edge score … ℎ 𝑜 • Attention aggregation 0,1 • New hidden node features ℎ 1 x n iterations … • New hidden edge features ℎ 𝑜 0 • New edge score Office of BERKELEY LAB 15 Science 15 15

Edge attention architecture • Input node features • Hidden node features 0.6 • Hidden edge features ℎ 1 • Edge score … ℎ 𝑜 • Attention aggregation 0,1 • New hidden node features x n iterations • New hidden edge features • New edge score Office of BERKELEY LAB 16 Science 16 16

Edge attention architecture • Input node features • Hidden node features 0.1 0.4 0.6 • Hidden edge features • Edge score • Attention aggregation • New hidden node features x n iterations • New hidden edge features 0.9 0.1 0.8 0.8 0.4 • New edge score Office of BERKELEY LAB 17 Science 17 17

Edge attention architecture • Input node features • Hidden node features ℎ 1 ℎ 2 0.4 0.6 • Hidden edge features • Edge score • Attention aggregation • New hidden node features x n iterations • New hidden edge features 0.4 0.1 0.8 • New edge score ℎ 3 ℎ 4 ℎ 5 Office of BERKELEY LAB 18 Science 18 18

Edge prediction architecture • Input node features • Hidden node features 0.4 0.6 • Hidden edge features • Edge score ℎ 1 ℎ 2 ℎ 1 • Attention aggregation + … ℎ 𝑜 ℎ 5 0 • New hidden node features ℎ 3 ℎ 4 x n iterations • New hidden edge features 0.4 0.1 0.8 • New edge score Office of BERKELEY LAB 19 Science 19 19

Edge attention architecture • Input node features ℎ 1 … • Hidden node features ℎ 𝑜 1 0.6 • Hidden edge features ℎ 1 • Edge score … ℎ 𝑜 • Attention aggregation 0,1 • New hidden node features ℎ 1 x n iterations … • New hidden edge features ℎ 𝑜 0 • New edge score Office of BERKELEY LAB 20 Science 20 20

Edge attention architecture • Input node features • Hidden node features 0.1 0.2 0.9 • Hidden edge features ℎ 1 • Edge score … ℎ 𝑜 • Attention aggregation 0,1 • New hidden node features x n iterations x n iterations • New hidden edge features (hyperparameter) 0.9 0.3 0.9 0.2 0.2 • New edge score Office of BERKELEY LAB 21 Science 21 21

Doublet GNN Performance Threshold 0.5 0.8 Accuracy 0.9761 0.9784 Purity 0.9133 0.9694 Efficiency 0.9542 0.9052 Two points to keep in mind • In the past, graphs have been constructed with a heuristic procedure that had much lower efficiency than the learned embedding. This GNN is classifying a ∼ 96% efficient doublet dataset • These metrics are not the end product: we use the scores of the doublets to create triplets without losing efficiency Office of BERKELEY LAB 22 Science

Why not simply join together our doublet predictions? x 4 x 3 detector centre Distance from 0.01 0.99 Pretty easy x 2 decision 0.99 x 1 Office of BERKELEY LAB 23 Science 23 23

Doublet choice can be ambiguous x 4 x 3 detector centre Distance from 0.87 0.84 Not so easy… x 2 so teach the network how to combine 0.99 x 1 Office of BERKELEY LAB 24 Science 24 24

But a GNN doesn’t know about “triplets” x 4 x 3 detector centre Distance from A GNN only knows x 2 about nodes ? ? and edge x 1 Office of BERKELEY LAB 25 Science 25 25

Moving to a “doublet graph” gives us back GNN power x 4 x 3 0.87 0.84 x 2 x 2 x 2 Now… 0.99 nodes represent doublets , edges represent triplets x 1 Office of BERKELEY LAB 26 Science 26 26

Moving to a “doublet graph” gives us back GNN power ( ) ( ) x 2 x 2 x 4 x 3 0.87 0.84 ( ) x 1 x 2 0.99 Now… nodes represent doublets , edges represent triplets Office of BERKELEY LAB 27 Science 27 27

Triplet Propaganda Threshold 0.5 0.8 Accuracy 0.9761 0.9784 Purity 0.9133 0.9694 Efficiency 0.9542 0.9052 * relative Doublet GNN Triplet GNN Threshold 0.5 0.8 Accuracy 0.9960 0.9957 Purity 0.9854 0.9923 Efficiency 0.9939 0.9850 * relative Office of BERKELEY LAB 28 Science

Triplet propaganda Gold : Unambiguously correct triplet or quadruplet Other colours : False positive/negative Key : Silver: Ambiguously correct triplet or quadruplet (i.e. edge shared by correct triplet and false positive triplet) Bronze dashed: Correct triplet, but missed quadruplet (i.e. edge shared by correct triplet and false negative triplet) Red: Completely false positive triplet Blue dashed: Completely false negative triplet Office of BERKELEY LAB 29 Science

Triplet propaganda Gold : Unambiguously correct triplet or quadruplet Other colours : False positive/negative Key : Silver: Ambiguously correct triplet or quadruplet (i.e. edge shared by correct triplet and false positive triplet) Bronze dashed: Correct triplet, but missed quadruplet (i.e. edge shared by correct triplet and false negative triplet) Red: Completely false positive triplet Blue dashed: Completely false negative triplet Office of BERKELEY LAB 30 Science

Triplet GNN improves doublet GNN results Black : Triplet classifier correctly labelled, doublet classifier mislabelled Red : Doublet classifier correctly labelled, triplet classifier mislabelled In this graph, triplet classifier Fixes 389 edges Worsens 10 edges Office of BERKELEY LAB 31 Science