Deep Learning for Particle Track
Finding in High Energy Physics
Steve Farrell
AI-at-SLAC Seminar
2019-02-21

HEP.TrkX

Introduction• Machine learning is taking over the world• High Energy Physics is no stranger to ML• Modern techniques like Deep Learning gaining popularity in HEP• Numerous challenges remain in implementing DL solutions forvarious complex HEP problems

2019-02-21

S. Farrell, SLAC AI Seminar2

High Luminosity Large Hadron Collider• LHC upgrade for 2026
• Order of magnitude
increase in intensity• 200 averageinteractions perbunch crossing• 40 million bunch
crossings per second

2019-02-21

S. Farrell, SLAC AI Seminar

3

Particle tracking at theLarge Hadron Collider• Tracking detectors record“hits” from charged particles• O(100M) readout channels• Algorithms connect the dots
to reconstruct tracks• O(10k) tracks from O(100k) hits

2019-02-21

S. Farrell, SLAC AI Seminar

4

HL-LHC tracking detectors

ATLAS today• Cylindrical barrel and disk-shaped
endcap detector layers• Silicon pixel and strip detector
technologies• 100 million readoutchannels• Complex layouts

CMS proposed layout

S. Farrell, SLAC AI Seminar5

Today’s tracking algorithms• Hit clustering: cells → spacepoints (“hits”)• Seed finding: construct hit triplets• Track building: extend seeds and searchwith combinatorial Kalman Filter• Track fitting/selection: Resolve ambiguities,
fit track parameters

2019-02-21

S. Farrell, SLAC AI Seminar

6

Tracking challenges• Combinatorial explosion with increasing
occupancy• Track reconstruction will dominate CPU
consumption• Algorithms are• hard to parallelize• hard to run on SIMD
architectures

2019-02-21S. Farrell, SLAC AI Seminar

7

Thinking outside the box• The HEP.TrkX Project• DOE HEP-CCE pilot project to develop Deep Learningsolutions to particle track reconstruction• Collaboration between LBNL, Caltech, and FNAL

HEP.TrkX• The TrackML Challenge• Engaging with the broader DS/ML community todevelop solutions• Challenges hosted on Kaggle and Codalab

2019-02-21

S. Farrell, SLAC AI Seminar

8

Machine learning for tracking• What parts of the workflow could be replaced/augmented with ML?
• Hit clustering• Seed finding• Track extrapolation/evaluation (for track building)• Track fitting• What are the major challenges?
• Complex data

• Dimensionality, sparsity, geometry• Missing hits, double hits, shared hits

• Complex objectives• Constraints on compute and physics performance
2019-02-21

S. Farrell, SLAC AI Seminar

9

Deep Learning inspirations
Image segmentation

Video object tracking
https://arxiv.org/abs/1604.03635

https://arxiv.org/abs/1604.02135

2019-02-21

S. Farrell, SLAC AI Seminar

10

Data used for studies• 2D and 3D toy data (planes)• Simulated data with ACTS toolkit• Uses generic HL-LHC detector description

https://gitlab.cern.ch/acts/acts-core
https://www.kaggle.com/c/trackmlparticle-identification/data
2019-02-21

S. Farrell, SLAC AI Seminar

11

Image representations• Unroll cylindrical detector layers• Treat as multi-channel image• Apply convolutional and recurrent
neural networks

12Computer vision applications• Detector image segmentation with CNNs, LSTMs• Vertex finding
with CNNs

2019-02-21

particles hitting
3 layers

Vertex position
estimate13

Limitations of the image representation• Image-based methods face challenges scalingup to realistic conditions• High dimensionality and sparsity• Irregular detector geometry• Hence, hard to learn good feature representations• Instead, we can consider spacepoint-based
representations• Hit sequences• Hit graphs

2019-02-21

S. Farrell, SLAC AI Seminar

14

RNNs for hit sequences

Next-word
prediction• RNNs are good with sequence data
• NLP, temporal data, genomic• How could we utilize them for trackreconstruction?
• Next-step predictions (like Kalman Filter)• Track scoring• Track fitting

Next-hitprediction• Not a huge departure from today’s
tracking methods• But maybe there are benefits
2019-02-21

https://towardsdatascience.com/lstm-by-example-using-tensorflow-feb0c1968537

15

RNN hit predictor• Given a sequence of hits, predict
coordinates of the next hit• Trained as regression problem
with mean-squared-error loss onbarrel tracks

2019-02-21

S. Farrell, SLAC AI Seminar

16

RNN hit predictor examples

2019-02-21

S. Farrell, SLAC AI Seminar

17

RNN Gaussian hit predictor• Augment the model to give Gaussian distributions for its predictionsCovariance matrix• Train with Gaussian log-likelihood loss:• Now we have a model which can tell us its uncertainty
2019-02-21

S. Farrell, SLAC AI Seminar

18

RNN Gaussian hit predictor• Predictions show
mostly expectedbehavior• Pull distributions
have some nonGaussian effects

2019-02-21

RNN tracking outlook• RNNs show potential for modeling andanalyzing tracking hit sequences• A bit more work is needed to determine
whether it could actually replace theKalman Filter• Currently probably not competitive

2019-02-21

S. Farrell, SLAC AI Seminar

20

Graph representation• Detector geometry

2019-02-21

S. Farrell, SLAC AI Seminar

21

Graph representation• True tracks

2019-02-21

S. Farrell, SLAC AI Seminar

22

Graph representation• What our data looks like

2019-02-21

S. Farrell, SLAC AI Seminar

23

Graph representation• Connect compatible hits together to construct a graph• Learn on this representation with Graph Neural Networks
2019-02-21

S. Farrell, SLAC AI Seminar24

Geometric Deep Learning

http://geometricdeeplearning.com/• Deep Learning for graphs and manifolds• Message-passing architectures which can build localized,
hierarchical feature representations• Can have analogous properties to CNNs

• Some examples:• Semi-supervised Classification with Graph Convolutional Networks• Neural Message Passing for Jet Physics• Relational inductive biases, deep learning, and graph networks

GDL applications
Recommender systems /matrix completion

https://arxiv.org/abs/1704.06803

2019-02-21

Shape analysis

https://arxiv.org/pdf/1611.08097.pdf

S. Farrell, SLAC AI Seminar

Modeling traffic

https://medium.com/syncedreview/shanghai-tests-graph-recurrentneural-networks-for-traffic-prediction-fdd4c2182b53

26

Hopfield networks for tracking (~1990)
Input
segments

Finalsegments• Identify true segments in a graph of connected hits
• No learned parameters, but solved via annealing with an energy loss
https://www.sciencedirect.com/science/article/pii/0010465588900045
function
https://www.sciencedirect.com/science/article/pii/0168900289913004
https://www.sciencedirect.com/science/article/pii/001046559190048P

2019-02-21

S. Farrell, SLAC AI Seminar

27

Graph tracking tasks• Hit classification

2019-02-21• Segment classification

S. Farrell, SLAC AI Seminar

28

Graph network architecture
Two components operate on graph:

2019-02-21

S. Farrell, SLAC AI Seminar

29

Graph network architecture
Two components operate on graph:• Edge network uses the attachednode features to compute edge
scores2019-02-21

S. Farrell, SLAC AI Seminar

30

Graph network architecture
Two components operate on graph:• Edge network uses the attachednode features to compute edge
scores• Node network aggregates forward
and backward neighbor nodefeatures with edge scores andcomputes new node features
2019-02-21

S. Farrell, SLAC AI Seminar

31

Graph network architecture
Two components operate on graph:• Edge network uses the attachednode features to compute edge
scores• Node network aggregates forward
and backward neighbor nodefeatures with edge scores andcomputes new node features
2019-02-21

S. Farrell, SLAC AI Seminar

32

Putting it together• Edge and node networks applied in recurrent, alternating fashion• With each “layer”, the model propagates and accumulates information
through the graph, strengthening/pruning connections adaptively• A Recurrent Graph Neural Network with Attention• Application-specific output• Hit classifier: binary node classifier output layer• Segment classifier: final application of EdgeNetwork for edge scores2019-02-21

S. Farrell, SLAC AI Seminar

33

Architecture details• Inputs(N x D) node feature matrix(N x E) association matrix of nodes to input edges(N x E) association matrix of nodes to output edges• EdgeNetwork is a 2-layer MLP with tanh and sigmoid activations(E) edge weight array• NodeNetwork is a 2-layer MLP with tanh activations(N x D) node features
2019-02-21

S. Farrell, SLAC AI Seminar

34

hit classification

Constructing the graph• For hit classification, select hits in theregion around a track seed
• connect all hits between adjacent layers• For segment classification, use all hitsand geometric constraints to select edges

segment classification• Δphi, Δr, Δz, adjacent layers• Prototype results are using lowoccupancy data
2019-02-21

S. Farrell, SLAC AI Seminar

35

Hit classification• 7-layer model with 26k
parameters• Performs well, with good
purity and efficiency
Test set metrics
Accuracy:
0.9942
Purity:0.9918
Efficiency: 0.9793

2019-02-21

Great separation

Segment classification• 4-layer model with 7kparameters• Performs well, with good
purity and efficiency
Test set metrics
Accuracy:
0.9952
Purity:0.9945
Efficiency: 0.9870

2019-02-21

Great separation

Visualization on 10-track sample• Edge scores after 1st layer
2019-02-21

S. Farrell, SLAC AI Seminar

38

Visualization on 10-track sample• Edge scores after 2nd layer
2019-02-21

S. Farrell, SLAC AI Seminar

39

Visualization on 10-track sample• Edge scores after 3rd layer
2019-02-21

S. Farrell, SLAC AI Seminar

40

Visualization on 10-track sample• Edge scores after 4th layer
2019-02-21

S. Farrell, SLAC AI Seminar

41

Visualization on 10-track sample• Edge scores after final layer
2019-02-21

S. Farrell, SLAC AI Seminar

42

Scaling up• The model is recently showing promising results on high-occupancy
data as well
• Good efficiency• Decent purity• Watch for next
publication for
details

2019-02-21

S. Farrell, SLAC AI Seminar

43

Graph network outlook and next steps• We think this is a promising approach for particle tracking• Near-perfect results on low-occupancy data with “simple” models• There’s still quite a bit of room for innovation (e.g. adding domain knowledge)• Already getting decent performance on high occupancy target data

• Nonetheless, challenges remain• Increasing occupancy increases edge density, making it harder to resolve thereal edges• Need robust preprocessing step for building the input graph• Need robust post-processing step for final track assignment• Computational: need optimized sparse matrix support to be scalable
2019-02-21

S. Farrell, SLAC AI Seminar44

TrackML Challenge• Can the broader data science community help us develop solutions?
• Side benefit: standardize the tracking problem and dataset!
• Previous HEP challenges have been popular: Higgs ML, Flavour of Physics• Particle tracking is trickier in certain ways
• Complex data and physics metrics• Compute vs. physics performance tradeoffs
https://sites.google.com/site/trackmlparticle/

2019-02-21

S. Farrell, SLAC AI Seminar45

TrackML Accuracy Phase• Featured competition on Kaggle, May – Aug 2018
• https://www.kaggle.com/c/trackml-particle-identification• Scoring metric based on reconstruction efficiency only (no timing)• Winners used combinatorial algorithms and “simple” ML
Logistic regression

Deep Neural Network2019-02-21

S. Farrell, SLAC AI Seminar46

TrackML Throughput Phase• Hosted on CodaLab platform, started in September• https://competitions.codalab.org/competitions/20112• Scoring based on reconstruction efficiency and timing• Low participation, but current leaders have fast solutions• Competition was extended to March 12, 201947

Conclusion• HEP particle tracking could be transformed by Deep Learning
methodologies• Image-based and sequence-based methods may be promising• Graph-based methods seem very promising• The TrackML challenge has been a great way to promote ourproblem and engage with the broader community• Though unclear if we will get revolutionary ideas yet48

NERSC is hiring!
• NESAP postdocs, to help prepare for our next-gen supercomputer,
Perlmutter (has an ML track):• http://www.nersc.gov/users/application-performance/nesap/nesappostdoctoral-fellowships/• NERSC data analytics student internships:• https://www.nersc.gov/research-and-development/data-analytics/dasinternships

2019-02-21

S. Farrell, SLAC AI Seminar

49

Backups

2019-02-21

S. Farrell, SLAC AI Seminar50

ML for pixel clustering• In dense environments we getmerged pixel clusters• Neural networks used toresolve these cases• Count particles• Estimate positions anduncertainties
https://arxiv.org/abs/1406.7690

51

TheProject

https://heptrkx.github.io/• A pilot project to develop machine learning algorithms for HEP
tracking• Funded by DOE ASCR, COMP HEP; part of HEP-CCE
LBL:Steve Farrell, Mayur Mudigonda, Prabhat, Paolo Calafiura
Caltech: Dustin Anderson, Jean-Roch Vlimant, Josh Bendavid, Maria Spiropulu, Stephan Zheng
FNAL: Aristeidis Tsaris, Giuseppe Cerati, Jim Kowalkowski, Lindsey Gray, Panagiotis Spentzouris• We’ve been trying lots of ideas to see what sticks• Representations: detector images, hit sequences, hit graphs• Architectures: CNNs, RNNs, Graph NNs• Problems: single/multi track building, vertex finding• Datasets: toy data, ACTS generic detector, TrackML dataset
52

Computer vision models

Outputs

Inputs

Detector layers

We can process detector images with• Convolutional neural networks (CNNs)• Recurrent neural networks (RNNs)
Can solve a variety of tasks
53Computer vision applications• Mapping images onto track parameters• De-noising54

Image based tracking• Image segmentation task on very
simple 2D and 3D toy data
• RNNs and CNNs can find the pixels
belonging to a desired track• Performance degrades with increased
occupancy

Track in noise

Multi-track
background

3D toy data
pixel classification accuracy
vs. occupancy
In 3D

CTD-WIT 2017: https://doi.org/10.1051/epjconf/201715000003
Code: https://github.com/HEPTrkX/heptrkx-ctd

55

ACTS detector images
Binned into 3 detector
imagesZ and phi projections56

Track parameter estimation[Work of Dustin Anderson]
CTD 2017• Use a basic CNN with downsampling andregression head to estimate a track’sparameters• could be an auxiliary target to guidetraining, or potentially useful as the
final output of tracking!
• Identifying straight line params in heavy
noise:57

Extending to variable number of tracks[Work of Dustin Anderson]
CTD 2017• Attach an LSTM to a CNN to emit parameters for a variable numberof tracks!
• The LSTM generates the sequence of parameters• Requires an ordering the model can learn• Should provide some kind of stopping criteria

58

Estimating uncertainties on parameters[Work of Dustin Anderson]
CTD 2017• Train the model to also estimate the uncertainties by adding additional targets:• Train using a log gaussian likelihood loss:• and voila!

59Computer vision applications• Estimating primary vertex position from detector image

particles hitting
3 layers

Vertex position
estimate

Use to constrain seed finding

Average 5 pileupinteractions

Julien Esseiva’s bachelor thesis
Shown at ACAT 201760

Hit sequence to track assignment
FCbi-GRU
bi-GRUbi-GRU• Sort all hits in an event according to position

Work of Daniel R. Zurawski, Keshav
Kapoor, interns at FNAL.
Shown at ACAT 2017

• Feed hits into a few layers of bi-directional recurrent net (GRU)• Output is a set of assignment probabilities to track groups• Ordering of output track categories is similarly sorted as hits
• Requires assumed maximum number of tracks
61

From Jean-Roch Vlimant (CHEP 2018)
6

CNNs for seed finding in CMS[not HEP.TrkX]

20-channel input image• Hit doublet classification• Uses pixel images of hitclusters• Classifies with a CNN
https://indico.cern.ch/event/587955/contributions/2937548/
63

RNN (TrackNet) for BM@N experiment[not HEP.TrkX]

https://indico.cern.ch/event/5879
55/contributions/2937545/
64

RNN Gaussian hit predictor examples

Coarse outer layers

Can these be used to build tracks?
• We tested on simple low-occupancy datawith naive approach• no combinatorial branching• But promising results• Full algorithm in development
65

RNN hit predictor models
HitPredictor((lstm): LSTM(3, 32, batch_first=True)(fc): Linear(in_features=32, out_features=2))
Parameters: 4802

HitGausPredictor((lstm): LSTM(3, 32, batch_first=True)(fc): Linear(in_features=32, out_features=5))
Parameters: 4901

66

GNN hit classifier
NodeClassifier((input_network): Sequential((0): Linear(in_features=4, out_features=64)(1): Tanh())(edge_network): EdgeNetwork((network): Sequential((0): Linear(in_features=136, out_features=64)(1): Tanh()(2): Linear(in_features=64, out_features=1)(3): Sigmoid()))(node_network): NodeNetwork((network): Sequential((0): Linear(in_features=204, out_features=64)(1): Tanh()(2): Linear(in_features=64, out_features=64)(3): Tanh()))(output_network): Sequential((0): Linear(in_features=68, out_features=1)(1): Sigmoid()))
Parameters: 2650267

GNN segment classifier
SegmentClassifier ((input_network): Sequential ((0): Linear (3 -> 32)(1): Tanh ())(edge_network): EdgeNetwork ((network): Sequential ((0): Linear (70 -> 32)(1): Tanh ()(2): Linear (32 -> 1)(3): Sigmoid ()))(node_network): NodeNetwork ((network): Sequential ((0): Linear (105 -> 32)(1): Tanh ()(2): Linear (32 -> 32)(3): Tanh ())))
Parameters: 6881

68

Slide by Andrey Ustyuzhanin

6

Slide by Andrey UstyuzhaninSlide by Andrey Ustyuzhanin7