# End-to-End Jet Classification of Boosted Top Quarks
with the CMS Open Data

M. Andrews  Department of Physics, Carnegie Mellon University, Pittsburgh, USA
B. Burkle  Department of Physics, Brown University, Providence, USA  Y. Chen
Google Inc., Mountain View, CA, USA  D. DiCroce  Department of Physics and
Astronomy, University of Alabama, Tuscaloosa, USA  S. Gleyzer  Department of
Physics and Astronomy, University of Alabama, Tuscaloosa, USA  U. Heintz
Department of Physics, Brown University, Providence, USA  M. Narain
Department of Physics, Brown University, Providence, USA  M. Paulini
Department of Physics, Carnegie Mellon University, Pittsburgh, USA  N. Pervan
Department of Physics, Brown University, Providence, USA  Y. Shafi  Google
Inc., Mountain View, CA, USA  W. Sun  Google Inc., Mountain View, CA, USA  E.
Usai  Department of Physics, Brown University, Providence, USA  K. Yang
Google Inc., Mountain View, CA, USA

###### Abstract

We describe a novel application of the end-to-end deep learning technique to
the task of discriminating top quark-initiated jets from those originating
from the hadronization of a light quark or a gluon. The end-to-end deep
learning technique combines deep learning algorithms and low-level detector
representation of the high-energy collision event. In this study, we use low-
level detector information from the simulated CMS Open Data samples to
construct the top jet classifiers. To optimize classifier performance we
progressively add low-level information from the CMS tracking detector,
including pixel detector reconstructed hits and impact parameters, and
demonstrate the value of additional tracking information even when no new
spatial structures are added. Relying only on calorimeter energy deposits and
reconstructed pixel detector hits, the end-to-end classifier achieves an AUC
score of 0.975$\pm$0.002 for the task of classifying boosted top quark jets.
After adding derived track quantities, the classifier AUC score increases to
0.9824$\pm$0.0013, serving as the first performance benchmark for these CMS
Open Data samples. We additionally provide a timing performance comparison of
different processor unit architectures for training the network.

##  1 Introduction

The Large Hadron Collider (LHC) is a prolific top quark factory: since the
beginning of data-taking in 2010, over $10^{8}$ top quarks have been produced.
The measurement of the top quark’s properties and production rates at the LHC
remains one of the main research priorities at experiments like the Compact
Muon Solenoid (CMS) at the LHC. Moreover, investigating the resonant
production of top quarks offers potential hints of the presence of new physics
that may lie beyond the standard model (BSM).

Top quarks are unique in that they decay before they have time to hadronize,
always decaying to a bottom quark and a W-boson. During the top decay chain,
the W-boson will decay hadronically to quarks 66.5% or leptonically to a
lepton and neutrino pair 33.5% of the time[1]. At hadron colliders like the
LHC, the low production cross section of prompt electrons and muons can be
exploited to boost tagging efficiency when identifying top quarks with a
leptonically decaying W-boson in its decay chain. However, hadronic decays of
top quarks can be much harder to identify, since the primary features used to
identify them are the topology of its decay products and the track features of
the bottom quark decay products. In particular, at high transverse momenta,
the hadronic decay of highly a Lorentz-boosted top quark can lead to a single
merged cluster of particles in the detector, hereby referred to as jets,
offering a unique and challenging view into the study of the top quark’s
properties. Because of this, discriminating boosted top quark-jets from light
flavour- or gluon-jets has become an important challenge for the LHC
experiments, and a popular benchmark for data analysis techniques involving
machine learning (ML) algorithms in high-energy physics (HEP).

Most jet identification techniques rely on inputs provided by the Particle
Flow (PF) algorithm used to convert detector level information to physics
objects [2]. The Particle Flow algorithm has many advantages due to its
ability to greatly reduce the size and complexity of particle physics data
while providing a physically intuitive and easy to use representation in
physics analyses. Many of the modern machine learning approaches to jet
discrimination are based on PF-based inputs [3, 4, 5, 6, 7, 8, 9, 10].
However, there is some invariable loss of information from reducing the data
set complexity. Despite the very high reconstruction efficiency of PF
algorithms, some physics objects may fail to be reconstructed, are
reconstructed imperfectly, or exist as fakes [11]. For that reason it is
advantageous to consider end-to-end reconstruction that allows a direct
application of machine learning algorithms to low-level data representation in
the detector.

In this work, we extend the end-to-end deep learning approach for particle and
event classification [12]. Specifically, we extend the use of end-to-end jet
images introduced for quark- vs. gluon-jet discrimination [13] to the task of
boosted top quark- vs. light quark- or gluon-jet discrimination. In previous
work [13] we found that the track information was the leading contributor to
the classifier’s performance. Due to this insight and the importance of
identifying displaced tracks associated with bottom quark decays, this new
work introduces a number of key features from the CMS tracking detectors to
exploit the full topology of hadronically decaying top quarks.

This paper is arranged as follows: Section 2 describes the CMS Open Data
datasets used for this paper. In Section 3, we describe the construction of
the detector images with particular focus on the new tracking information
introduced in this work. In Section 4, we present the results of top jet
classification, and in Section 5 we provide an interpretation of the results.
In Section 6, we present a timing benchmark of the hardware platforms used to
train the models and in Section 7, we summarize our findings and conclusions.

##  2 Open Data Simulated Samples

Sample | Description | Number of events
---|---|---
TTJets_HadronicMGDecays |  $\mathrm{t}\overline{\mathrm{t}}$, $p_{\mathrm{T}}$ (top)$>$400 GeV, all-jet top decays | 2969109
QCD_Pt-300to600_TuneZ2star_Flat | QCD, flat $300<\hat{p}_{T}<600$ GeV | 1498800
QCD_Pt-400to600_TuneZ2star_Flat | QCD, flat $400<\hat{p}_{T}<600$ GeV | 1989000
QCD_Pt-600to3000_TuneZ2star_Flat | QCD, flat $600<\hat{p}_{T}<3000$ GeV | 2975400
Table 1: List of the CMS Open Data simulated samples used and total number of
events considered for each of the datasets.

The end-to-end deep learning technique relies on high-fidelity simulated
detector data, which in this work comes from the simulated Monte Carlo in the
CMS Open Data Portal [14, 15]. We use a sample of SM top-antitop ($t\bar{t}$)
pair production as a source of boosted top quarks, where the $W$ boson from
the top quark decay is required to decay to quarks [16]. Additionally, the
reconstructed top quark transverse momentum ($p_{T}$) is required to be
greater than 400 GeV. At this momentum, we expect that a large fraction of the
W- and b-jets produced in the top quark decay chain will be suitably merged.
The Monte Carlo sample was generated with Madgraph 2.6.6 [17] and uses Pythia6
for parton showering [18] with the Z2Star tune. For the light-flavour and
gluon jets, we use three samples of QCD dijet production in different ranges
of the hard-scatter transverse momentum, $\hat{p}_{T}$:
$300<\hat{p}_{T}<600\,\text{GeV}$, $400<\hat{p}_{T}<600\,\text{GeV}$, and
$600<\hat{p}_{T}<3000\,\text{GeV}$ [19, 20, 21], each with a flat distribution
in $\hat{p}_{T}$. Like the $\mathrm{t}\overline{\mathrm{t}}$ sample, these
samples were generated and showered with Pythia6 [18] using the same Z2Star
tune. The full list of datasets and number of events is summarized in Table 1.
For all samples, the detector response is simulated using Geant4 with the full
CMS geometry and is processed through the CMS PF reconstruction algorithm
using CMSSW release 5_3_32 [22]. An average of ten additional background
collisions or pileup (PU) interactions are added to the simulated hard-scatter
event, which are sampled from a realistic distribution of simulated minimum
bias events. For this study, we use a custom CMS data format which includes
the low-level tracker detector information, specifically, the reconstructed
clusters from the pixel and silicon strip detectors [23]. From the tracker
clusters, we then do a parametric estimate of the position of the hit on the
sensor surface.

We take reconstructed jets clustered using the anti-kt algorithm [24] with a
radius parameter R of 0.8, or so-called AK8 jets, and require
$p_{\mathrm{T}}>400\,\text{GeV}$ and $|\eta|<1.37$ for our event selection.
Here, $\eta$ is the pseudorapidity and equates to the polar angle of the CMS
detector according to $\eta=-\ln(\tan{\frac{\theta}{2}})$. This $\eta$ cut is
to ensure that the jet image does not extend beyond the $|\eta|<2.4$
acceptance limit of the current CMS tracker. Additionally, for the top jets we
require the generator-level top quark, its bottom quark and W-boson daughters,
and W-boson daughters to be within an angular separation of $\Delta
R=\sqrt{\Delta\eta^{2}+\Delta\phi^{2}}<0.8$ from the reconstructed AK8 jet
axis, where $\phi$ is the azimuthal angle of the CMS detector. In order to
avoid biases caused by the different $p_{\mathrm{T}}$ distributions of the
jets, we pseudo-randomly drop jets from the three QCD samples such that the
total number of jets and $p_{\mathrm{T}}$ distribution of the
$\mathrm{t}\overline{\mathrm{t}}$ sample is reproduced. After this
$p_{\mathrm{T}}$-resampling, the QCD jets are mixed together with the
$t\bar{t}$ jets and split into training, validation, and test sets, as
detailed in Table 2.

Figure 1: Distribution of the transverse momentum of the selected jets for the $\mathrm{t}\overline{\mathrm{t}}$ and non-top multijet training samples. The non-top multijets are re-sampled to reproduce the $\mathrm{t}\overline{\mathrm{t}}$ transverse momentum distribution. Category | Top quark jets | QCD jets | Total Jets
---|---|---|---
Train | 1280830 | 1279170 | 2560000
Validation | 47859 | 48141 | 96000
Test | 319819 | 320181 | 640000
Table 2: Number of jets used for training, validation, and testing in the top
quark and non-top quark jet categories. Numbers are reported after the
$p_{\mathrm{T}}$-resampling procedure.

##  3 CMS Detector, Images, and Network Training

CMS is a multi-purpose detector composed of several cylindrical subdetector
layers, with both barrel and endcap sections, encasing a primary interaction
point. It features a large $B=3.8$ T solenoid magnet to bend the trajectories
of charged particles that aid in $p_{\mathrm{T}}$ measurement [25]. At the
innermost layers, close to the beamline, there is a silicon tracker used to
reconstruct the trajectory of charged particles and find their interaction
vertices. The tracker can be divided in two parts the silicon pixel detector
and silicon strip detector [26]. The first silicon pixel detector is the inner
most part and composed of three layers in the barrel region (BPIX) and three
disks in the endcap region (FPIX). Each layer is composed of pixel sensors
that provide a very precise position of the passage of a charged particle. The
pixel detector provides crucial information for vertexing and track seeding.
The outer part of the tracking system is composed of several layers of silicon
strip. These provide a precise position in the $\phi$ coordinate, but not in
the $\eta$ coordinate. This is followed by the electromagnetic calorimeter
(ECAL), made of lead-tungstate crystals, to measure the energy of
electromagnetically interacting particles, then the hardonic calorimeter
(HCAL), made of brass towers, to measure the energy of hadrons [27, 28]. These
are surrounded by the solenoid magnet which is finally encased by the muon
chambers to detect the passage of muons [29].

We construct the jet images using low-level detector information where each
subdetector is projected onto an image layer, or several layers in the case of
the tracker, in a grid of 125 x 125 pixels with the image centered around the
most energetic HCAL deposit of the jet. Each pixel corresponds to the span of
an ECAL barrel crystal which covers a 0.0174 $\times$ 0.0174 in the
$\eta-\phi$ plane, giving our images an effective $\Delta R$ of 2.175. For the
ECAL and HCAL images, each crystal or tower is directly mapped to one or more
image pixels containing the energy deposited in that crystal or tower, as
described in [13]. Reconstructed particle tracks are weighted by their
reconstructed $p_{\mathrm{T}}$ and their location is projected to an ECAL
crystal. In order to better overlap with the calorimeter images, the
$\eta-\phi$ position of the tracks are determined by assuming the track
originated from the primary vertex, the location of the collision with the
highest $\sum p_{\mathrm{T}}^{2}$, before being propagated to the ECAL
surface.

To improve the identification of tracks coming from the hadronization of b
quarks, we added additional layers motivated by the long flight distance of b
hadrons producing reconstructed tracks that do not converge to the primary
vertex. To make the network aware of this information, we tried two
approaches: a) additional two layers corresponding to the reconstructed tracks
weighted by their transverse ($d0$) and longitudinal ($dZ$) impact parameter
significance and b) additional layers from the BPix detector that contain low-
level representation of tracker RecHits. The impact parameter (IP) is defined
as the distance vectors of minimum approach between the track helix and the
primary vertex. To obtain the IP significance, the $d0$ and $dZ$ values are
divided by their respective uncertainties. These quantities are computed
without using approximations and are therefore accurate also for tracks
relatively far from the primary vertex. Any $d0$ ($dZ$) values larger than 10
cm (20 cm) are suppressed to zero to prevent training degradation caused by
the inclusion of tracks with superfluously large IP. Such tracks are expected
to originate from photon conversions in the tracker or from poor track
reconstruction, and these cuts are not expected to negatively impact network
performance. Finally, each of the layers is independently normalized such that
the value of the average cell, ignoring empty cells, is approximately unity in
order to facilitate training convergence.

In an effort to extract as much information as possible from the tracking
subdetector, we include the low-level detector information from this system:
the tracking hits traditionally used to reconstruct tracks in the detector.
There are multiple steps in the conversion from charge clusters produced via
charged particles passing through the tracker to fully reconstructed tracks.
In this study, we consider the reconstructed hit (RecHit) information from the
three layers of the BPIX, but not from the FPIX or the silicon strip detector,
network inputs. RecHits are obtained by first clustering nearby pixels of a
given sensor which pass an adjustable charge threshold. A straight line fits
the pixel cluster to center of the beam, and it’s angle with the sensor
surface is used to compute a hit location which is corrected for the Lorentz
drift the charges experience before being read off the sensor. Given the hit
location on the sensor and location of the sensor in the detector, the
location of the RecHit is obtained. For full track reconstruction, tracks are
first seeded by fitting to the RecHits in the silicon strip detector. Those
track seeds are then matched with RecHits in the pixel detector and used to
produce the full tracks. The RecHits serve as a good intermediary between raw
detector outputs and reconstructed track quantities, serving as a map between
a module based coordinate system and the detector coordinates.

Figure 2: Jet images of the 1st layer of the BPIX (top left), 1st & 2nd layer2
BPIX (top right), 1st, 2nd, & 3rd layers of the BPIX (bottom left), and all
layers of the BPIX with the reconstructed track $p_{\mathrm{T}}$ (bottom
right) of a simulated 703 GeV $p_{\mathrm{T}}$ merged top jet. Images are
created at the nominal ECAL resolution.

Figure 3: Composite images of a simulated boosted top quark event. Images are
produced for a single cropped jet image (left) and the full CMS detector
(right).

For this study, one additional step is performed on the RecHits prior to
producing the image layers. The $\eta$ and $\phi$ position of the RecHit is
re-calculated with respect to the primary vertex of the collision rather than
the geometric center of the detector. This is done so that the $\eta$ and
$\phi$ of the RecHits better match the $\eta$ and $\phi$ of their
corresponding tracks when reaching the ECAL, which would otherwise deviate due
to the pixel detectors closeness to the beamline. After these computations are
performed, image layers are produced where each pixel intensity is set to one
if the image pixel contains a RecHit and zero otherwise. We generate three
different image layers, one for each of the three concentric layers of the
BPIX. Figure 2 shows the three pixel layers separately, and together with the
track $p_{\mathrm{T}}$. The effect of the magnetic field, bending the
trajectories of charged particles in the $\phi$ coordinate can be seen by
comparing hits in the three pixel layers. Figure 3 shows an end-to-end image
featuring all the image layers considered in this work for a single jet and
the full detector. The only layers that cannot be seen are the track $d0$ and
$dZ$ values because the perfectly overlap with the track $p_{\mathrm{T}}$
layer. A full list of all image layers along with their description can be
found in Table 3.

Layer | Description
---|---
Track $p_{\mathrm{T}}$ | track momentum, position at ECAL surface
Track $d0$ & $dZ$ | track impact parameter projections, w.r.t PV, at ECAL surface
BPIX1–3 | Barrel pixel hits, position, binary value, PV-shifted
ECAL | EM energy deposit, per crystal
HCAL | Had energy deposit, per tower
Table 3: Summary descriptions of the image layers considered.

The network architecture and hyperparameters used in this work closely follow
what was previously used in [12, 13], making use of a ResNet-15 CNN [30]
trained with the ADAM optimizer [31]. The initial learning rate is $5\times
10^{-4}$ and is explicitly reduced by half every 10 epochs. We found training
for 20 epochs to be sufficient for convergence. However, for our final network
evaluations we trained for an additional 20 epochs. The network was developed
using the TensorFlow library [32]. The hardware platforms used to train the
models are described in Section 6.

##  4 Jet Identification Results

Table 4 shows the area under the receiver operator curve (AUC) for the
different combinations of track and calorimeter layers at ECAL granularity.
Networks were trained on a sample of 2.56M jets for 20 epochs and evaluated on
a separate sample of 200k jets with AUC statistical uncertainty of 0.002.

Layer Combinations | ROC-AUC
---|---
Track $p_{\mathrm{T}}$ (baseline) | 0.955$\pm$0.002
Track $p_{\mathrm{T}}$ \+ ECAL + HCAL (nominal) | 0.967$\pm$0.002
Track $p_{\mathrm{T}}$ \+ $d0$ \+ $dZ$ | 0.972$\pm$0.002
Track $p_{\mathrm{T}}$ \+ $d0$ \+ $dZ$ \+ ECAL + HCAL | 0.981$\pm$0.002
Table 4: Performance of the classifier trained up to 20 epochs and evaluated
on a set of 200000 jets. Results are shown for different combinations of
tracking and calorimeter layers; AUC scores have a statistical uncertainty of
$\pm$0.002.

Our previous end-to-end deep learning results showed that the Track
$p_{\mathrm{T}}$ layer gave the best single layer performance for jet
discrimination [13]. Therefore, we choose track $p_{\mathrm{T}}$ layer
performance as a baseline for our models’ performance. We observe that the
largest single-subdetector performance increase comes with the inclusion of
the $d0$ and $dZ$ track information, leading to an AUC score improvement of
0.014–0.017. Comparing rows 2 and 3 in Table 4 shows that the combination of
track $p_{\mathrm{T}}$, $d0$, and $dZ$ outperforms the nominal combination
layers despite the fact that the $p_{\mathrm{T}}$ \+ $d0$ \+ $dZ$ images are
agnostic to neutral particles, since they do not produce charge clusters in
the tracker. What we observe is in agreement with [13] where the tracks were
observed as the most important feature for jet discrimination, as well as more
traditional jet tagging approaches which require the presence of a b-tagged
subjet tagged using IP variables [33, 34].

Table 5 shows the network performance when including the BPIX RecHits in the
jet images. On their own, the BPIX RecHits give a worse performance than the
track $p_{\mathrm{T}}$. However, we observe multiple improvements in network
performance after combining the BPIX RecHit images with other layers. When
training the network on images composed of BPIX1–3, ECAL, and HCAL layers we
find that it outperforms the nominal combination of layers, shown in the
second row of Table 4, and improves the AUC score by 0.008. Comparing row 4 of
Table 5 with row 3 of Table 4 shows that adding the BPIX RecHits to the track
$p_{\mathrm{T}}$ \+ $d0$ \+ $dZ$ images improves the AUC by 0.005. To study
the effect of BPIX RecHit resolution on network performance, we additionally
trained the network on images produced at sub-ECAL granularity. However, we
found that the higher granularity produced no significant changes in network
performance.

Layer Combinations | ROC-AUC
---|---
BPIX1–3 | 0.947$\pm$0.002
BPIX1–3 + Track $p_{\mathrm{T}}$ | 0.965$\pm$0.002
BPIX1–3 + ECAL + HCAL (no reconstructed variables) | 0.975$\pm$0.002
BPIX1–3 + Track $p_{\mathrm{T}}$ \+ $d0$ \+ $dZ$ | 0.977$\pm$0.002
BPIX1–3 + Track $p_{\mathrm{T}}$ \+ $d0$ \+ $dZ$ \+ ECAL + HCAL (full image) | 0.9824$\pm$0.0013
Table 5: Performance of the classifier after including BPIX layers. When
training on a subset of images, the network was trained for 20 epochs and
evaluated on 200k jets. Training on the full images was performed for 40
epochs and evaluated on 640k jets, giving an uncertainty of $\pm$0.0013.
Training the network on the full combination of images layers give an
efficiency at 1% misidentification of 66.41%.

The bottom row of Table 5 reports the performance of our network when trained
on all 8 image channels. The network was trained for 40 epochs and used the
training, validation, and testing dataset sizes listed in Table 2. When
evaluating the network, we find an AUC score of 0.9824$\pm$0.0013 and a signal
efficiency of 66.41% at 1% misidentification.

##  5 Interpretation and Discussion

An in depth look at the networks’ performance when trained on different layer
combinations provides an insight into the features that the network is
learning. We note that the strongest single subdetector performance comes from
the reconstructed tracks weighted by their $p_{\mathrm{T}}$ and IP variables.
This is in agreement with the expectation based on the current understanding
of high momentum top jets. We expect a large number of high $p_{\mathrm{T}}$
tracks, due to the jet containing three merged subjets, and a small subset of
tracks having large IP values, attributed to a decaying B-meson. What is
particularly interesting is that the network is able to successfully extract
this IP information from the addition of the $d0$ and $dZ$ layers to the track
$p_{\mathrm{T}}$ image layer. By design, these track-only images are composed
of a set of sparse layers with the same distribution of activated pixels.
Intuitively, extracting information from such images using 2D convolutions
becomes much more difficult than the traditional computer vision tasks.
However, in this difficult to parse regime, our network achieves an AUC of
0.972$\pm$0.002, outperforming the denser jet images used for our nominal
layer combination.

The second insight comes from the performance of the BPIX RecHits. As
mentioned in Section 4, the BPIX RecHits do not show strong standalone single-
layer performance. However, this is to be expected for multiple reasons. The
pixel detector has an $\eta$ and $\phi$ resolution of 10 $\mu$m, giving the
inner most layers a 1D spatial resolution that is almost eight times finer
than the ECAL [26, 27]; the ECAL resolution is too coarse to derive vertex
information from pixel hits. Furthermore, we only considered the barrel region
of the pixel detector, and do not include any RecHits from the forward region
of the pixel detector. Any jets that border the $\eta$ acceptance of this
study will only have RecHit information for a portion of the jet image.
Finally, our network is agnostic to each layer’s distance from the beamline,
giving the network incomplete information about the RecHits global
positioning. For example, the RecHits will drift in $\phi$ as the charged
particle bends in the CMS detector’s magnetic field. But unless more layers
are added to the image, the network does not have enough information to know
the order of each hit nor the direction in $\phi$ the particle is moving. But
despite the shortcomings of our current RecHit implementation, we find
remarkable results. With the exception of the final layer combination, where
BPIX RecHits are added to images composed of track $p_{\mathrm{T}}$ \+ $d0$ \+
$dZ$ \+ ECAL + HCAL information, we note that adding the RecHits gives a
significant increase in network performance. The most notable are cases where
BPIX RecHits are added on top of the tracking variables (1), and the case
where BPIX RecHits are used in lieu of the derived tracking information (2).

In the case of (1) we see that the network is able to use the BPIX RecHits to
derive new jet features which were not present in the derived track quantities
alone. One possible feature is the track charge, where motion through $\phi$
can be combined with the final location of the track to determine its
direction of curvature, and thus the charge, of the track. However, more
abstract features can also exist in these images. In the case of (2), the
network does not use any reconstructed variables for its inputs. We see that
despite the lack of derived variables, the network outperforms the track
$p_{\mathrm{T}}$ \+ $d0$ \+ $dZ$ images, and only performs marginally worse
than the final performance on the full images. The overall success of our
network’s ability to learn from BPIX RecHits paves the foundation for future
studies of end-to-end top taggers where no derived variables are used. In
addition to the forward region of the pixel detector, future work can include
RecHits from the silicon strip detector, which is used for track seeding and
track momentum measurement. We also look to explore new types of
architectures, such as graph neural networks [35], to exploit the full spatial
resolution of the CMS tracker and the 3D correlation of its layers to
complement existing architecture in other layers.

##  6 Timing Performance Comparison

A number of factors affect the training and evaluation time of the classifiers
and their memory requirements: The number of events used during training, the
number of image layers in each event, and the resolution of each of the
images. Compared to our previous work [13], the number of training events has
increased by a factor of 3.2 and the number of image layers has increased by a
factor 2.7. Meanwhile, the image resolution has remained the same. From the
increase in image layers, we can see that the size of an uncompressed image
has increased from 183 kB in [13] to 488 kB. For storage the images are
compressed using gunzip and have an average compressed size of 6.56 kB giving
a compression factor of approximately 75. These factors combined lead to a
significant computational cost increase for training the network, both in
training time and memory requirements. For this task training was carried out
on multiple accelerated hardware architectures as a benchmark to compare their
performance. The comparison was performed on two different graphical
processing units (GPU) and a tensor processor unit (TPU), whose specifications
are summarized in Table 6. Additionally, a test was performed where we trained
a single network on multiple GPUs in parallel to observe how the training
speeds scaled with number of processing units. Each device was accessed via
different computing clusters with varying I/O architectures described below.

Processor | Manufacturer | Year Released | HBM Memory | Performance
---|---|---|---|---
Tesla P100 | NVIDIA | 2016 | 16 GB | 9.3 Single-Precision TeraFLOPS
Tesla V100 | NVIDIA | 2017 | 32 GB / 16 GB | 125 Mixed-Precision TeraFLOPS
TPUv3-8 | Google | 2018 | 128 GB | 420 Mixed-Precision TeraFLOPS
Table 6: Comparison of architecture specification for the NVIDIA Tesla P100
[36], NVIDIA Tesla V100 [37], and Google TPUv3-8 [38] architectures. The
floating point operations per second speed (FLOPS) are quoted based on the
data type and architecture setup used during training.

The NVIDIA Tesla P100 is a GPU that utilizes the Pascal architecture [36]. The
Tesla P100 GPU was accessed on a shared cluster at Fermilab via a dedicated
GPU worker node through a 12 GB/s PCIe connection using the CUDA Toolkit v9.1
drivers. During training, data was stored and read from an HGST 1W10002 hard
drive [39] located on the GPU machine. Images were uncompressed, pre-
processed, and provided to the GPU using a single Intel(R) Xeon(R) Silver 4110
8-core CPU [40].

The NVIDIA Tesla V100 uses the Volta architecture, incorporating eight tensor
cores and an all-around higher performance than the Pascal architecture [37].
Unlike the P100, the V100 is able to make use of mixed precision operations,
which were utilized for this comparison, to drastically increase the number of
floating point operations per second (FLOPS). During training, images were
read from a solid state drive (SSD) with better random input/output operations
per second. The computing node ran Cuda v11.0.2 drivers and used four Intel(R)
Xeon(R) Gold 5118 12-core CPUs [41] to perform data loading and pre-processing
tasks, which were parallelized using the NVIDIA DGX-1 architecture [42]. Of
the 48 available CPU cores, 20 were used for data loading. When training the
network on multiple GPUs we chose to train on one, two, four, and eight V100s
using the the Horovod framework [43]. For the scaling tests a larger batch
size of 1024 tf examples was used.

The TPU is a type of AI-accelerated architecture designed by Google with the
purpose of training and running inference on machine learning models [44], and
can be accessed through the Google Cloud platform [45]. TPUs boast a high
number of FLOPS, made possible by dedicated matrix multiplication units which
make use of the bfloat16 data type [46]. Unlike GPU based architectures, the
CPU cores used to fetch and pre-process batches all live on TPU board creating
a low latency I/O pipeline which does not have to compete for CPU resources.
For this work, a TPUv3-8 running TPU software v1.14 was run using a type
n1-standard-1 Google Cloud virtual machine node. Data used to train the
networks was stored in Google Cloud buckets located in the same region as the
VM and TPU to decrease time associated with transferring data during the
training. This gives the cloud storage buckets comparable performance to
persistent HDs in the local VM storage area [47].

###  6.1 GPU Scaling Performance

The Horovod framework [43] was used to perform the GPU scaling, as it provides
flexibility of scaling the training of the network to multiple GPUs. Horovod
takes advantage of the inter-GPU and inter-node communication methods such as
NCCL (Nvidia Collective Communications Library) and MPI (Message Passing
Interface) to distribute the deep learning model parameters between various
workers and aggregate them accordingly.

Figure 4: Scaling end-to-end deep learning training on multiple GPUs.

Figure 4 shows the training time when the final optimised model was scaled to
multiple GPUs. Going from one to two GPUs, we observe a roughly 50% decrease
in training speeds. However, as we continue to double the GPUs, we start to
see diminishing returns in improvement. This is primarily due to the
underlying input bottlenecks which are not addressed by improving computation
parallelization. Additionally, adding more GPUs will increase the output
latency associated with averaging gradients.

###  6.2 Computing Architecture Comparison

Table 7 provides a timing comparison for training machine learning models
using the three different architectures. The training performed on the P100
were drastically longer, primarily stemming from low I/O speeds. These low
speeds come from random read inefficiencies associated with disk HDs [48] as
well as a weaker CPU utilizing fewer CPU nodes when fetching batches and
sending them to the GPU. The Tesla V100 and TPUv3-8 give much stronger
performance and were accessed utilizing optimized I/O pipelines. Furthermore,
when training on the Tesla V100 we were able to take advantage of SSD storage,
leading to improved random read speeds. Comparing the trainings on the Tesla
V100 and TPUv3-8 architectures, we find that the I/O speeds were almost
identical. By subtracting the single batch I/O time from the single batch
train time, we obtain an approximate computation time. From this, we see that
the TPUv3-8 spent approximately 2.6 ms to perform forwards and backwards
propagation calculations, which is a factor of four faster than the 11 ms
required by the Tesla V100.

Category | Tesla P100 | Tesla V100 | TPUv3-8
---|---|---|---
Training Batch Size | 32 tf examples | 64 tf examples | 64 tf examples
Batches Per Epoch | 80k batches | 40k batches | 40k batches
I/O Time (single batch) | 0.119 s | 0.0180 s | 0.0176 s
Train Time (single batch) | 0.481 s | 0.0290 s | 0.0202 s
Train Time (one epoch) | 321 min | 19 min | 14 min
Table 7: Comparison of I/O and training time for different computing
architectures. A larger batch size was used when training on the Tesla V100
and TPUv3-8. Training times were found to vary by approximately 10% between
epochs.

As mentioned, comparing all of the architectures studied shows that one of the
largest training speed increases comes from improving the hardware associated
with reading, decompressing, and pre-processing data. The training pipelines
used to access the Tesla V100 and the TPUv3-8 used optimized data storage and
highly parallelized CPU architectures to meet I/O speeds that were 10-20x
faster than what was used to train on the Tesla P100. For future studies I/O
speeds can be further reduced by improving the way our data is stored
optimizing the compression factor used when storing the tf examples. When
training on TPUs data can be stored as a bfloat16 to decrease storage space,
the memory usage, decreasing decompression speeds while producing no overhead
associated with type conversions during data loading.

##  7 Conclusions

Figure 5: Comparison of time taken to train over a single batch on the GPU and
TPU architectures.

In this work we have extended the end-to-end deep learning technique to top
quark jet classification. To enhance the performance of the classifier we
added additional layers containing information about track parameters and
pixel detector reconstructed hits, marking the first top-tagging algorithm
which uses tracking RecHits as input variables. The model was trained using
CMS Open Data datasets containing low-level tracking information [16, 19, 20,
21].

The end-to-end classifier trained on all input features achieves the
performance of AUC of 0.9824$\pm$0.0013. We find that the addition of $d0$ and
$dZ$ variables gives the largest boost to network performance when compared to
subdetector information used in previous end-to-end jet discrimination studies
[13]. At ECAL image granularity, BPIX RecHits do not provide the network with
information that is not present in the combination of track $pT$, $d0$, $dZ$,
ECAL, and HCAL layers. However, we find that it still improves subgroups of
these layers, and the network achieves an AUC score of 0.975$\pm$0.002 when
training on images void of derived variables. These findings lay the ground
work for future studies which look to incorporate RecHits from the full CMS
tracker, higher-resolution training, and to explore new deep learning
architectures that can fully exploit the tracker granularity.

Identical copies of the model were trained on a single NVIDIA Tesla P100 GPU,
a single NVIDIA Tesla V100 GPU, and a TPUv3-8 accessed through the Google
Cloud platform. We observe that the single largest training speed improvements
come from the optimized I/O infrastructures. We find that both the Tesla V100
GPU and the TPUv3-8 offer a significant training time improvement over a Tesla
P100 GPU when training on end-to-end jet images. We also trained an identical
copy on the network in parallel on multiple NVIDIA Tesla V100 GPUs to observe
the relationship between number of GPUs and training speeds. From this, we
were able to see substantial improvements when going to higher number of GPUs.

## Acknowledgments

We would like to thank the CMS Collaboration and the CERN Open Data group for
releasing their simulated data under an open access policy. We strongly
support initiatives to provide such high-quality simulated datasets that can
encourage the development of novel but also realistic algorithms, especially
in the area of machine learning. We believe their continued availability will
be of great benefit to the high energy physics community in the long run.

We would also like to thank Google for access to the Google Cloud which helped
us speed up the training of machine learning models used in this study.

## References

  * [1] Particle Data Group Collaboration, “Review of Particle Physics”,  Phys. Rev. D 98 (2018), no. 3, 030001, doi:10.1103/PhysRevD.98.030001.
  * [2] CMS Collaboration, “Particle-flow reconstruction and global event description with the CMS detector”,  JINST 12 (Jun, 2017) P10003. 82 p, doi:10.1088/1748-0221/12/10/P10003.
  * [3] CMS Collaboration Collaboration, “Boosted jet identification using particle candidates and deep neural networks”,.
  * [4] A. Butter et al., “The Machine Learning Landscape of Top Taggers”,  SciPost Phys. 7 (2019) 014, doi:10.21468/SciPostPhys.7.1.014, arXiv:1902.09914.
  * [5] G. Kasieczka, T. Plehn, M. Russell, and T. Schell, “Deep-learning Top Taggers or The End of QCD?”,  JHEP 05 (2017) 006, doi:10.1007/JHEP05(2017)006, arXiv:1701.08784.
  * [6] P. T. Komiske, E. M. Metodiev, and J. Thaler, “Energy flow networks: deep sets for particle jets”,  Journal of High Energy Physics 2019 (Jan, 2019) doi:10.1007/jhep01(2019)121.
  * [7] A. Butter, G. Kasieczka, T. Plehn, and M. Russell, “Deep-learned Top Tagging with a Lorentz Layer”,  SciPost Phys. 5 (2018), no. 3, 028, doi:10.21468/SciPostPhys.5.3.028, arXiv:1707.08966.
  * [8] CMS Collaboration Collaboration, “New Developments for Jet Substructure Reconstruction in CMS”,.
  * [9] ATLAS Collaboration, “Performance of top-quark and $\mathrm{W}$-boson tagging with ATLAS in Run 2 of the LHC”,  The European Physical Journal C 79 (Apr, 2019) doi:10.1140/epjc/s10052-019-6847-8.
  * [10] CMS Collaboration, “Identification of heavy, energetic, hadronically decaying particles using machine-learning techniques”,  JINST 15 (2020), no. 06, P06005, doi:10.1088/1748-0221/15/06/P06005, arXiv:2004.08262.
  * [11] CMS Collaboration, “CMS Physics: Technical Design Report Volume 1: Detector Performance and Software”.  Technical Design Report CMS. CERN, Geneva, 2006.
  * [12] M. Andrews, M. Paulini, S. Gleyzer, and B. Poczos, “End-to-End Physics Event Classification with CMS Open Data: Applying Image-Based Deep Learning to Detector Data for the Direct Classification of Collision Events at the LHC”, 2018.
  * [13] M. Andrews et al., “End-to-end jet classification of quarks and gluons with the CMS Open Data”,  Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment 977 (Oct, 2020) 164304, doi:10.1016/j.nima.2020.164304.
  * [14] CERN, “CERN OpenData portal”, (2019).
  * [15] CMS Collaboration, “CMS data preservation, re-use and open access policy”, 2014.  doi:10.7483/opendata.cms.udbf.jkr9.
  * [16] CMS Collaboration, “Tracker-hit-enriched TTJets_HadronicMGDecays_8TeV-madgraph”, 2019.  doi:10.7483/OPENDATA.CMS.OPKY.OJMJ.
  * [17] J. Alwall et al., “The automated computation of tree-level and next-to-leading order differential cross sections, and their matching to parton shower simulations”,  Journal of High Energy Physics 2014 (Jul, 2014) doi:10.1007/jhep07(2014)079.
  * [18] T. Sjöstrand, S. Mrenna, and P. Skands, “PYTHIA 6.4 physics and manual”,  Journal of High Energy Physics 2006 (May, 2006) 026–026, doi:10.1088/1126-6708/2006/05/026.
  * [19] CMS Collaboration, “Tracker-hit-enriched 300 to 600 bin of QCD_Pt-15to3000_TuneZ2star_Flat_8TeV_pythia6”, 2019.  doi:10.7483/OPENDATA.CMS.HUED.7R3E.
  * [20] CMS Collaboration, “Tracker-hit-enriched 400 to 600 bin of QCD_Pt-15to3000_TuneZ2star_Flat_8TeV_pythia6”, 2019.  doi:10.7483/OPENDATA.CMS.YWDZ.KSLK.
  * [21] CMS Collaboration, “Tracker-hit-enriched 600 to 3000 bin of QCD_Pt-15to3000_TuneZ2star_Flat_8TeV_pythia6”, 2019.  doi:10.7483/OPENDATA.CMS.CWTT.8Q3E.
  * [22] CMS Collaboration, “CMS Software Version 5_3_32 (CMSSW_5_3_32)”, 2016.  doi:10.7483/OPENDATA.CMS.WYJG.FYK9.
  * [23] E. Usai et al., “Samples with full event information including tracker hits for tracking, ML, and top quark tagging studies”, 2019\.  doi:10.7483/OPENDATA.CMS.CHC3.5KPG.
  * [24] M. Cacciari, G. P. Salam, and G. Soyez, “The anti-ktjet clustering algorithm”,  Journal of High Energy Physics 2008 (Apr, 2008) 063–063, doi:10.1088/1126-6708/2008/04/063.
  * [25] CMS Collaboration, “The CMS magnet project: Technical Design Report”.  Technical Design Report CMS. CERN, Geneva, 1997.
  * [26] CMS CollaborationV. Karimäki, et al., “The CMS tracker system project: Technical Design Report”.  Technical Design Report CMS. CERN, Geneva, 1997.
  * [27] CMS Collaboration, “The CMS electromagnetic calorimeter project: Technical Design Report”.  Technical Design Report CMS. CERN, Geneva, 1997.
  * [28] CMS Collaboration, “The CMS hadron calorimeter project: Technical Design Report”.  Technical Design Report CMS. CERN, Geneva, 1997.
  * [29] CMS Collaboration, J. G. Layter, “The CMS muon project: Technical Design Report”.  Technical Design Report CMS. CERN, Geneva, 1997.
  * [30] K. He, X. Zhang, S. Ren, and J. Sun, “Deep Residual Learning for Image Recognition”, 2015.
  * [31] D. P. Kingma and J. Ba, “Adam: A Method for Stochastic Optimization”, 2014.
  * [32] M. Abadi et al., “TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems”, 2015\.  Software available from tensorflow.org.
  * [33] C. Weiser, “A Combined Secondary Vertex Based B-Tagging Algorithm in CMS”, Technical Report CMS-NOTE-2006-014, CERN, Geneva, Jan, 2006.
  * [34] CMS Collaboration, “Boosted Top Jet Tagging at CMS”,.
  * [35] J. Zhou et al., “Graph Neural Networks: A Review of Methods and Applications”, 2018.
  * [36] “NVIDIA Tesla P100: The Most Advanced Data Center Accelerator”. Accessed: 28 August 2020.
  * [37] “NVIDIA V100 $|$ NVIDIA”. Accessed: 16 September 2020.
  * [38] “Cloud Tensor Processing Units (TPUs)”, https://cloud.google.com/tpu/docs/tpus.  Accessed: 30 August 2020.
  * [39] “Western Digital DC HA210 Datasheet, 3.5 Inch Data Center Hard Drives”. Accessed: 30 August 2020.
  * [40] “Intel® Xeon® Silver 4110 Processor (11M Cache, 2.10 GHz) Product Specifications”. Accessed: 16 September 2020.
  * [41] “Intel® Xeon® Gold 5118 Processor (16.5M Cache, 2.30 GHz) Product Specifications”. Accessed: 16 September 2020.
  * [42] “NVIDIA DGX-1: Deep Learning Server for AI Research”. Accessed: 16 September 2020.
  * [43] A. Sergeev and M. D. Balso, “Horovod: fast and easy distributed deep learning in TensorFlow”, 2018.
  * [44] N. P. Jouppi et al., “In-Datacenter Performance Analysis of a Tensor Processing Unit”, arXiv:1704.04760.
  * [45] “Google Cloud Computing Services”, https://cloud.google.com/.  Accessed: 30 August 2020.
  * [46] “BFloat16: The secret to high performance on Cloud TPUs $|$ Google Cloud Blog”. Accessed: 18 September 2020.
  * [47] “Storage classes — Google Cloud”, https://cloud.google.com/compute/docs/disks.  Accessed: 29 September 2020.
  * [48] S.-W. Lee, B. Moon, and C. Park, “Advances in Flash Memory SSD Technology for Enterprise Database Applications”, in  Proceedings of the 2009 ACM SIGMOD International Conference on Management of Data, SIGMOD ’09, p. 863–870.  Association for Computing Machinery, New York, NY, USA, 2009.  doi:10.1145/1559845.1559937.