# Particle Multi-Axis Transformer for Jet Tagging

Muhammad Usman4, M Husnain Shahid4, Maheen Ejaz4, Ummay Hani4, Nayab Fatima4,
Abdul Rehman Khan1, Asifullah Khan1,2,3*, Nasir Majid Mirza4
1 Pattern Recognition Lab, Department of Computer & Information Science,
Pakistan Institute of Engineering & Applied Sciences, Nilore, Islamabad 45650,
Pakistan
2 PIEAS Artificial Intelligence Center (PAIC),
Pakistan Institute of Engineering & Applied Sciences, Nilore, Islamabad 45650,
Pakistan
3Center for Mathematical Sciences,
Pakistan Institute of Engineering & Applied Sciences, Nilore, Islamabad 45650,
Pakistan
4Department of Physics & Applied Mathematics,
Pakistan Institute of Engineering & Applied Sciences, Nilore, Islamabad 45650,
Pakistan
Corresponding Author: *Asifullah Khan<EMAIL_ADDRESS>
###### Abstract

Jet tagging is an essential categorization problem in high energy physics. In
recent times, Deep Learning has not only risen to the challenge of jet tagging
but also significantly improved its performance. In this article, we propose
an idea of a new architecture, Particle Multi-Axis transformer (ParMAT) which
is a modified version of Particle transformer (ParT). ParMAT contains local
and global spatial interactions within a single unit which improves its
ability to handle various input lengths. We trained our model on JETCLASS, a
publicly available large dataset that contains 100M jets of 10 different
classes of particles. By integrating a parallel attention mechanism and
pairwise interactions of particles in the attention mechanism,ParMAT achieves
robustness and higher accuracy over the ParT and ParticleNet. The scalability
of the model to huge datasets and its ability to automatically extract
essential features demonstrate its potential for enhancing jet tagging.

_K_ eywords Jet Tagging $\cdot$ Deep Learning $\cdot$ Transformer $\cdot$
Multi-Axis-Attention

##  1 Introduction

The largest and most robust particle accelerator in the world is the Large
Hadron Collider (LHC), which is located at CERN. It is made up of a circular
tube that is 27 km long and situated between 50 and 175 meters below the
surface. New particles are created inside the collider when beams collide with
each other. Reaching an energy level of 14 TeV in the center of the mass
reference frame is the main goal of the Large Hadron Collider. The
breakthrough of Higgs Boson in 2012 was one of its most important
accomplishments. Nine experiments are carried out at the LHC which are
facilitated by collisions within the accelerator. A variety of particle
detectors are used in this research. Four of these are large-scale
detectors.Two of them are general-purpose detectors such as ATLAS (A toroidal
LHC Apparatus) and CMS (Compact Muon Solenoid), and two are research-focused
detectors such as LHC-b and ALICE (A Large Ion Collider Experiment). These
huge detectors are positioned beneath the LHC in vast caverns. The Large
Hadron Collider Forward (LHCf), the Forward Search Experiment (FASER), the
Scattering and Neutrino Detector at the LHC (SND@LHC), the MoEDAL (Monopole
and Exotics Detector at the LHC), and TOTEM (Total, elastic and diffractive
cross-section measurement) are the five small detectors. Among these, the
ATLAS stands as one of the largest detector, measuring 46m in length, 25m in
height, and 25m in breadth. All detectors transfer some part of or all of the
radiation energy to the detector mass. The form in which energy is converted
depends on the type of detector. A gaseous detector directly collects
ionization electrons to form an electric signal. In Scintillation detectors,
both excitation and ionization induce molecular transition which results in
the emission of light. These detectors are capable of producing a usable
signal for a given type of radiation and energy. The Solenoid controls the
path of particles inside the detector. Neutrino remains invisible in the
detector.

Machine learning has enhanced the possibility of discovering new fundamentals
of nature and has transformed the methodology for analyzing large-scale data
samples in High Energy Physics. At the LHC, proton beams are propelled to
almost the velocity of light and then smash at 40 million times per second (40
MHz). These energetic proton-proton smashes yield new, short-lived particles
that subsequently decay, producing a concentrated spray of outgoing particles
referred to as jets. Particle detectors such as CMS or ATLAS are tasked with
detecting these jets. Jet Tagging is to identify and classify jets that are
produced by particles such as quark and gluons[1], the W or Z boson, or the
Higgs boson[2]. It is a critical task in High energy physics because the
particles that creates jets radiates and these radiated particles produces
more outgoing particles. It will be difficult to identify the original
particle because these radiations distorts their original features.
Traditional approaches to jet tagging depend on handcrafted features inspired
by concept of Quantum Chromodynamics(QCD). QCD is a theoretical framework that
describes how particles moves within a jet. Different methods have been
adopted so far for the identification of jets. The dominant representation
describe a jet as a cloud of particles, characterized by an unorganized and
variable sized collection of particles that are released from detectors.

Deep learning has transformed Jet tagging.We harnessed the idea of a new
architecture, Particle Multi-axis Transformer (ParMAT) for the classification
tasks of jets, which fits with the JETCLASS data sets. It simplifies the
complex attention mechanism used in the Multi-Axis Vision
Transformer(MaxViT)[3]. We used JETCLASS dataset for training this model. It
is large and detailed dataset consisting of 100M jets for training. It has 10
classes and it also includes the jets that have not been explored yet for
tagging. The resulting Particle Multi-axis transformer(ParMAT) achieves higher
performance than particle transformer(ParT), PFN, P-CNN and ParticleNet.
Training our model led to notable increase in accuracy.

##  2 Related Work

Deep Learning for Jet tagging: Deep learning has become the go-to approach for
jet tagging as it has significantly improved particle identification in recent
years. The efficiency and effectiveness of DL techniques on jet physics has
improved with the advancement of two factors.

  * •

Jet Representation: Previously jets were represented as images[4], graphs[5],
sequences[6] and trees[5] which worked well with corresponding architectures
like graph neural networks, 2D Convolutional Neural networks and recurrent
networks. In the recent times, jet as particle cloud[7][8] has been
introduced, comparable to point clouds, which view a jet as an invariant set
of particles under permutations. Based on this representation, architectures
like Energy Flow Networks[7], Graph Attention based Point Neural Network[9],
Point Cloud Transformer[10] were introduced. Among these, the new Particle Net
remained the predominant DL architecture.

  * •

Jet Tagging Datasets: Simulated jet tagging datasets has proved to be more
convenient than the real ones, For this investigation, a number of datasets
have been released so far. These include Top quark[11], Quark gluon[7], and
Higgs Bosons[12] tagging datasets. These along with JetNet[13] and multiclass
dataset[14] contain information about kinematics and particle identification
only. Architectures like 1D CNNs, RNN used these datasets for studies.
However, uptil now, the recently introduced JetClass has the most information
and maximum number of different types of jets . This dataset contains immense
information; thus it needs relatively bigger models like the recent Particle
transformer.

Deep-learning-powered jet tagging techniques have become integral to data
analysis at the LHC, notably through innovations like the DeepAK8 algorithm
developed by the CMS Collaboration[15]. This algorithm, employing a 1D CNN[16]
based on the ResNet architecture[17], effectively identifies jets arising from
top quarks or bosons like the Higgs, W, or Z. Such advancements have markedly
heightened the potential for discovering new heavy particles, as evidenced by
CMS publications in 2021[18] and 2022[19]. Additionally, ParticleNet, another
deep-learning tool utilized by CMS, enabled groundbreaking achievements,
including the first observation of Z boson decay into charm quarks and setting
stringent constraints on Higgs decay to charm quarks[19]. Furthermore,
ParticleNet’s application in probing the quartic Higgs and vector bosons
interaction has led to the indirect confirmation of its existence. These
strides underscore the crucial role of improved jet tagging methods in
advancing our knowledge of elementary particles, the fundamental constituents
of the universe.

Transformers: In recent years, transformer models have achieved remarkable
success, initially in natural language processing and later expanding into
computer vision. The original Transformer architecture[20], along with its
derivatives like BERT[21], ViT[22], and Swin-Transformer[23], has consistently
set new performance benchmarks across diverse tasks, showcasing the
versatility of Transformers. Beyond traditional domains, Transformers, driven
by their attention mechanism, have proven adept at addressing fundamental
scientific challenges. Notably, AlphaFold2[24], a groundbreaking model in
protein structure prediction, leverages the attention mechanism, enhancing its
predictive accuracy[25]. Specifically, interpretability of the model is
improved by incorporating a pair bias into the self-attention mechanism, which
is obtained from pairwise features.. These advancements highlight the
transformative impact of Transformer architectures across both conventional
and scientific domains, highlights their universal applicability and
effectiveness.

Transformer in vision: In computer vision, the innovative perspective of ViT,
which takes image patches as visual words, has sparked an intense surge of
research interest in visual Transformers. Vit and its derivatives are commonly
utilized in image classification[26], segmentation[27][28], object detection
and visual recognition tasks. The variants like Deit[29], Beit[30], and
Swin[23] operate by iteratively focusing on pertinent segments of the input
sequences, updating representations via self-attention mechanisms to identify
complex relationships and patterns. Improved locality, pyramidal designs,
training techniques, sparse attention, and improved locality are some of the
more recent efforts aimed at improving model and data efficiency. In this
domain, Multi axis Vision transformer(MaxVit) has been introduced in 2022[3]
and RSIR Transformer[31] in 2023 which show superior performances as compared
to other common vision transformers.

##  3 Dataset

We used the JETCLASS dataset to train our proposed Particle Multi-axis
Transformer for the jet tagging task. JETCLASS is a large-scale, simulated
dataset for CMS detector, proposed and generated by Hulin Qu111CERN, Geneva,
Switzerland, Congqiao Li2 and Sitian Qian222School of Physics, Peking
University, Beijing, China[32], to facilitate deep learning research in jet
tagging. This dataset includes 10 types of jets, including signal jets from
the decay of $WZ$ Bosons, top quarks ($t$), Higgs ($H$) bosons and background
signal from quark-gluon pairs ($q/g$). Since, in particle collision
experiments, a particle has more than one mode of decay, which means each
decay mode results in a different type of jet. So, this dataset considers 5
decay modes of Higgs, two decay modes of top quarks, one decay mode for W and
Z bosons each, and background jets from quark-gluon plasma, resulting in 10
types of jets.

JETCLASS was generated by simulating jets in Monte Carlo event generators,
using MAD-GRAPH5_aMC@NLO[33], PYTHIYA[34] and DELPHES[35]. MAD-GRAPH_aMC@NLO
simulates the hard processes. Hard processes involve the production of
particles after collision and their decay into other particles. Simulation of
soft processes that include parton showering and hadronization, is carried out
using PYTHIA. In the end, the effects of the detector are simulated by using
DLPHES, and we get our final dataset. Important to note that only high-quality
jets with 500-1000 GeV transverse momentum and pseudorapidity less than 2 were
selected while compiling the dataset and the rest were ignored. Figure 1 shows
cloud represention of an example of jet from our dataset.

Figure 1: Cloud representation of an example of Higgs jet from JetClass. It
shows all particles included in an example from Higgs jet along with the
background noise. Higgs decays in to bottom quark anti quark pair which
further decays into hadrons that are detected by detectors and are shown here.

JETCLASS dataset is the TTree object of the ROOT library with 41 branches. 10
branches are for labels and have boolean values. The rest of the branches are
features and correspond to jets or individual particles. The following 3
categories have some most important features for our jet tagging task.

Kinematics of particles: In particle physics, we have a 4-vector
($E,p_{x},p_{y},p_{z}$) that includes the energy and momentum of particles and
is used to describe the kinematics of particles. 4-vector is directly measured
by our detector and is the most fundamental quantity which provides basis for
all other kinematics variables. That’s why it is a fundamental feature in jet
tagging studies.

Charge on particles: Particle charge is crucial in determining that either
particle is hadron or photon or electron etc. CMS uses it to identify the
particle, and in the dataset, one hot encoding is used to determine this. So
we have five variables for each particle type and each of them is hot encoded.
For example, to determine if a particle is a photon, we have variable
part_isPhoton and it is hot encoded such that in jet, if a particle is a
photon it would have a value of 1, otherwise it would be 0. So we have 5
variable as is_ChargedHadron, is_NeutralHadron, is_Photon, is_Electron, and
is_Muon for each particle. Recent studies showed that it is important for the
classification of decay processes including leptons[32].

Trajectory features: This feature is important for jets that involve bottom
and charm quarks[36]. This feature was not included in datasets before
JetClass. This feature is the impact parameter of particle trajectories which
is the smallest distance to the primary vertex value. We have two variables
for this, one for longitudinal and one for transverse impact parameters. We
also have two extra variables that record the uncertainties in both impact
parameters. We have an impact parameter only for charged particles while zero
is assigned to the neutral particles.

Training set: In the training dataset, we have 10 M examples of each type of
jet. So, we have 100 M jets in total to train our model.

Validation set: A separate dataset of 5 M jets would be used for validating
our model during training. This dataset also have equal contributions from
each class and has 500 K examples for each type of jet.

Test set: A test data set of 20 M jets (2 M from each jet type) would be used
after training the model, for model evaluation purposes.

##  4 Model Architecture

Figure 2: The model diagram of (a) Particle Multi-Axis Transformer, (b)
Particle Multi-Axis Attention block, (c) Multi-Head Attention block, (d)
Multi-Axis Attention block, and (e) Class Attention Block.

This paper introduces a modified transformer-based architecture called
Particle Multi-axis Transformer(ParMAT), a combination of Particle Transformer
and Multi-axis Transformer. We split attention into two types: local and
global.[3] Our approach uses layers of block and grid attention in sequence.
Block region operates within discrete, non-overlapping windows (small patches
on intermediate feature maps) to identify local patterns. In contrast, grid
region targets long-range (global) interactions by focusing on a uniformly
spaced, sparse grid. The sizes of the windows for both block and grid
attentions are adjustable parameters, allowing us to maintain linear
computational complexity relative to the size of the input. It serves as a
core component facilitating both local and global spatial interactions within
a single unit.ParMAT offers increased adaptability and effectiveness compared
to full self-attention, showcasing its ability to handle various input lengths
with linear complexity. Unlike traditional local attention mechanisms, ParMAT
expands model capacity by incorporating a global receptive field. ParMAT is
constructed by stacking a particle attention multi-axis block, particle multi-
head attention,multi-axis block attention,multi-axis grid attention and class
attention mechanism. ParMAT is a modified version of ParT inspired by
attention mechanism.[37]. The model of this architecture is shown in Figure
2(a). We consider $N$ number of particles in a jet.
Input Processing: ParMAT takes $Particles$ and $Interaction$ as an input.
$Particles$ contains a list of $R$ features for each input particle, which
forms a 2-dimensional array of a shape $(N,R)$. $Interactions$ comprises a
matrix with $R^{\prime}$ features for each particle’s pair, making a
3-dimensional array $(N,N,R^{\prime})$.

MLP Projection: Each set of inputs is passed through a Multi-Layer Perceptron
(MLP) to project the $R$ and $R^{\prime}$ features into $n$ and
$n^{\prime}$-dimensional embeddings, $x_{0}\in\textbf{R}^{N\times n}$ and
$W\in\textbf{R}^{N\times N\times n^{\prime}}$ for $particles$ and
$Interactions$ respectively, where $x0$ and $W$ are the inputs after an MLP
has processed them. As the order in which particles in a jet are processed
does not matter we do not need any positional encoding. So we feed $x_{0}$
into multiple Particle Multi-Axis Attention blocks. Each block output is fed
into the next block as an input. This produces a series of new embeddings
[38].

Interaction features: Pairwise interaction features are characteristics or
properties that emerge from the interaction between pairs of particles. ParMAT
is a powerful architecture to find these pairwise interactions, but for the
trivial case here we would consider the interaction features extracted from
the Energy-momentum 4-vector. For each particle in the Jet Tagging task, we
have $p=(E,p_{x},p_{y},p_{z})$. Using the 4-vector $p_{m},p_{n}$ for two
particles $m,n$, we calculate 4 features:

$\displaystyle\Delta$
$\displaystyle={\left((y_{m}-y_{n})^{2}+(\phi_{m}-\phi_{n})^{2}\right)}^{1/2},$
(1) $\displaystyle k_{T}$ $\displaystyle=\min(p_{T,m},p_{T,n})\Delta,$
$\displaystyle z$
$\displaystyle=\frac{\min(p_{T,m},p_{T,n})}{p_{T,m}+p_{T,n}},$ $\displaystyle
m^{2}$ $\displaystyle=(E_{m}+E_{n})^{2}-\|\mathbf{p}_{m}+\mathbf{p}_{n}\|^{2}$

where $y_{m,n}$ is the rapidity (the velocity of a particle in a way that is
invariant under the Lorentz transformations along the direction of motion),
$\phi_{m,n}$ is the azimuthal angle, $p_{T_{m},n}$ is the transverse momentum,
and $\mathbf{p}_{m,n}$ is the 3-momentum having norm $\|\mathbf{}\|$. we now
take the logarithm of these variables (ln$\Delta$, ln$k_{T}$, ln$z$,
ln$m^{2}$) as the interaction feature, as the variable has a long tail
distribution [39].

Particle Multi-Axis attention block: In Figure 2(b), the P-MAA block consists
of two parts. First part consists of a LayerNorm (LN) then a multi-head
attention (MHA) module and two Multi-axis attention, MAA (grid/global) and MAA
(Window/local) with a LayerNorm (LN) layer again. The the second part includes
the 2-layers of Multi-layer Perceptron having the LayerNorm (LN) before it and
GELU nonlinearity in between. Each layer is connected with the following one.
The overall block structure is taken from NormFormer. we modified MHA as P-MHA
(particle-Multi head attention), which deals with the pairwise particle
interaction. The P-MHA is computed as

$\text{P-MHA}(Q,K,V)=\text{SoftMax}(QK^{T}/({d})^{1/2}+W)V\hskip 14.22636pt$
(2)

where, Q, K, and V are the query, key, and value. These are the linear
projections of the particle embedding in our case $x^{L}$. We modify the pre-
softmax attention weights by adding the interaction matrix W. This increases
the expressiveness of the attention mechanism by enabling PMHA to alter the
dot-product attention weights and integrate particle interaction information.

To make our model more advanced and efficient, we added an extra layer of
Multi-Axis attention, which is also called relative attention. we can
calculate P-MAA using:

$\text{P-MAA}(Q,K,V)=\text{SoftMax}(QK^{T}/({d})^{1/2}+W)V\hskip 14.22636pt$
(3)

Class Attention Block: The class attention block as illustrated in figure 2(b)
has almost similar architecture as the particle multi-axis attention block. In
this block, we use simple MHA and calculate attention between particles and
Global class token $x_{class}$. MHA inputs are:

$\displaystyle Q=J_{qx_{class}}+b_{q},$ (4) $\displaystyle K=J_{kz}+b_{k},$
$\displaystyle V=J_{vz}+b_{v}.$

where $J_{q}$, $J_{k}$, and $J_{v}$ are the weight matrices that the model
learns for creating the query, key, and value vectors. $x_{class}$ and $z$ are
the input vectors. Here $z$ is the product of $x_{class}.x^{L}$.$b_{q}$,
$b_{k}$, and $b_{v}$ are the bias terms.

Model Implementation: We implement our ParMAT model in PyTorch[40].
Particularly, P-MHA and P-MAA is implemented using the PyTorch. The ParMAT
model includes 8 P-MAA blocks and 2 class attention blocks. we provide
interaction matrix $W$ as the attentionn mask input. It utilizes a
128-dimensional particle embedding, generated from the input particle features
through a 3-layer MLP with (128, 512, 128) nodes in each layer, utilizing GELU
nonlinearity, and LN for normalization. The interaction input features are
processed through a 4-layer pointwise 1D convolution with channel
configurations (64, 64, 64, 8), GELU nonlinearity, and batch normalization to
create an 8-dimensional interaction matrix. The P-MHA (MHA) in the particle
(class) attention blocks have 8 heads, each with a query dimension of 16 and
an expansion factor of 4 for the MLP. We use a dropout (a regularization
technique used in neural networks to prevent overfitting) of 0.1 for all P-MAA
blocks and no dropout for the class attention block.

##  5 Experiment

We perform experiments using the new JETCLASS dataset, a widely adopted
benchmark for jet tagging tasks in high-energy physics research. The dataset
comprises simulated particle collision events, with each event containing
detailed information about the particles produced in the collision, including
their momentum and energy. The results of our experiment are summarized in
Section VIII, where we also compare the performance to previous leading
methods.

Methodology: In the experiments performed using the JETCLASS dataset, we
employ the entire spectrum of particle characteristics as our inputs, covering
kinematics, particle identification, and trajectory displacement. These
comprise a comprehensive list of 17 features for each particle. Furthermore,
we incorporate the four interaction features introduced for the ParMAT. The
training process is executed on the entirety of the training set, which
consists of 100M jets. We utilized the Ranger optimizer, incorporating a
weight decay of 10-4. The initial learning rate of 0.001 schedule followed a
flat cosine curve with a warm-up period of 50,000 steps. 1024 was selected as
the batch size. We implemented it on an NVIDIA DGX system and utilized
multiple GPUs (0, 1, 2, and 3) to expedite training. Following training, we
evaluated the trained ParMAT model on a separate validation set from the
JetClass dataset.

Baselines: We assess the performance of PARMAT against four baseline models:
ParT[19], PFN[27] utilizing Deep Sets[41], P-CNN, as employed in the DeepAK8
algorithm of the CMS experiment[15], and ParticleNet[8], an advanced
architecture derived from DGCNN[14], as presented by Qu & Gouskos (2020). The
results documented in the ParT paper for all models trained comprehensively on
the JETCLASS dataset are employed for comparison.ParticleNet utilizes an
established PyTorch implementation, whereas, for PFN and P-CNN, the ParT paper
reconstructed them in PyTorch to guarantee the reproduction of published
outcomes. The optimizer and learning rate schedule are consistent with those
used in the training of ParMAT.The optimal (Learning rate, batch size)
combinations are (0.0010, 1024) for ParMAT, (0.0010, 512) for ParT, (0.010,
512) for ParticleNet, and (0.020, 4096) for PFN and P-CNN.

##  6 Evaluation Metrices

We recommend using key performance indicators to fully assess the performance
of ParMAT on the detailed JETCLASS dataset. Due to the fact that jet tagging
in this dataset is a classification that includes multiple classes, We advise
using two widely used metrics: area under the receiver operating
characteristic curve and accuracy. These measurements do a good job of
evaluating overall efficacy. Furthermore, We suggest that background rejection
which is the opposite of false positive rate should be taken into an account
at a specific signal optimization, denoted by the true positive rate also
written as TPR set at x%.

$Rej_{xy\%}=1/\text{FPR at TPR = x\%}$ (5)

We apply this equation to all kinds of signal jets. Quark-Gluon jets are, by
default, categorized as background jets, following standard procedure in the
majority of LHC data analysis. The remaining nine jet types are all regarded
as signal jets. However, for $H\rightarrow l\nu qq^{\prime}$($t\rightarrow
bl\nu$), this efficiency is increased to 99% (or 99.5%) as these types exhibit
more distinctive characteristics and are more readily distinguishable from q/g
jets. As the Rejection metric concentrates on two classes only – the
background class (b) and the signal class (s), the true positive rate and the
false positive rate are assessed by using dual class scoring method.

$\text{Score}_{svsb}\equiv\frac{\text{Score}(s)}{\text{Score}(s)+\text{Score}(b)}$
(6)

In the context of jet tagging, separating signal jets (class s) from
background jets (class b) is crucial. This equation, inspired by the
methodologies of the CMS experiment at the LHC, utilizes softmax outputs
(Score(s) and Score(b)) to achieve optimal discrimination between the two
classes. Background rejection is challenging metric in jet tagging. This is
due to significant improvement in background rejection directly translates to
a substantial increase in the LHC’s discovery potential. For example, doubling
the background rejection can result in a 40% increase in discovery potential,
which is comparable to either doubling the dataset size or operating the LHC
for twice the duration.

333The ROC AUC function available in scikit-learn library may be used to
calculate the area under the curve by passing in the arguments, average equals
to ’macro’ and multi-class equals to ’ovo’.

##  7 Results

Table 1: The performance of our Particle Multi-axis transformer is evaluated and contrasted with the results obtained by Qu et al. in their 2022 study[19],where they employed the particle transformer for jet tagging. Our model’s performance is also compared with Particle flow network(Komsike et al.,2019b)[42],P-CNN, as reported by the CMS Collaboration (2020b)[15] and ParticleNet by Qu and Gouskos,2020[8]. ParMAT model has demonstrated higher performance using the JETCLASS dataset. Model |  Accuracy (%) |  AUC (%) |  Z→qq Rej50% |  t→bqq Rej50% |  t→blv Rej99.5% |  W→qq Rej50% |  H→bb Rej50% |  H→cc Rej50% |  H→gg Rej50% |  H→4q Rej50% |  H→lvqq Rej99%
---|---|---|---|---|---|---|---|---|---|---|---
Particle Flow Network(PFN) |  77.20 |  97.10 |  159.0 |  797.0 |  721.0 |  189.0 |  2924.0 |  841.0 |  75.0 |  198.0 |  265.0
PCNN |  80.90 |  97.80 |  204.0 |  2907.0 |  2304.0 |  241.0 |  4890.0 |  1276.0 |  88.0 |  474.0 |  947.0
ParticleNet |  84.40 |  98.40 |  283.0 |  10526.0 |  11173.0 |  347.0 |  7634.0 |  2475.0 |  104.0 |  954.0 |  3339.0
ParT(Plain) |  84.90 |  98.50 |  311.0 |  17699.0 |  12987.0 |  384.0 |  9569.0 |  2911.0 |  112.0 |  1185.0 |  3868.0
ParT |  86.10 |  98.70 |  402.0 |  32787.0 |  15873.0 |  543.0 |  10638.0 |  4149.0 |  123.0 |  1864.0 |  5479.0
ParMAT |  86.23 |  98.75 |  116.3 |  32906.7 |  24793.3 |  160.0 |  4484.3 |  1509.9 |  227.4 |  503.1 |  4211.5

Evaluation metrics described in Section VII are used to determine the
performance of our proposed model on JETCLASS dataset, with the outcomes
compiled in Table 1.This new proposed architecture sets a new benchmark by
achieving the highest performance of 86.32%, surpassing all others. ParMAT
performs with higher accuracy than Particle transformer.ParMAT significantly
outperform the earlier ParT,ParticleNet and P-CNN models on this comprehensive
dataset, showing background rejection differences that may exceed an order of
magnitude. This significant enhancement in the Particle Multi-Axis Transformer
is projected to significantly enhance the discovery capabilities for physics
research at the LHC.

We are plotting two graphs here: one with the x-axis representing time
(epochs) and the y-axis representing the model’s accuracy. These graphs
visualizes the performance of Particle Multi Axis Transformer models on
JETCLASS dataset. Figure 3 illustrates the plot of training accuracy versus
epochs, while Figure 4 displays the graph of validation accuracy versus time
(epochs). As training progresses, the graphs illustrate how the model’s
performance improves. The y-axis ranges from 0 to 100% but accuracy levels off
at around 85% ,indicating that its capacity may not be sufficient to learn
beyond that point. This could also suggest that the model might be over-
fitting.By analyzing the curves, it’s apparent that ParMAT consistently yields
higher accuracy suggesting its potential efficacy in particle collider studies
for distinguishing between different types of jets.

Figure 3: Training Accuracy vs Epochs. Figure 4: Validation Accuracy vs
Epochs.

Tagging performance shows notable variations across different signal types.
Exceptional discrimination against background quark/gluon jets is observed
with $t\rightarrow bl\nu$ and $H\rightarrow l\nu qq^{\prime}$ signals.
Employing latest ParMAT architecture allows for almost complete detection of
these signals, achieving over 99% efficiency with minimal contamination from
background jets. This breakthrough, utilizing jet types previously not being
discovered for tagging at the LHC, paves the way for new opportunities in this
domain.

Effect of training dataset Size on ParT Performance: To assess how the dataset
size influences jet tagging performance, we conducted further training
sessions and replicated the results of the ParT model using incomplete
data[43]. The original configuration of the particle transformer included 8
attention blocks with 8 attention heads each, utilizing the GeLU activation
function with a learning rate (LR) of 0.001. The Ranger optimizer was used to
enhance the model’s performance. Batch size was set to 512. In our initial
experiment, preserving this original setup, we achieved an accuracy of 85.43%.
However, when the training dataset size was reduced to 1%, there was a 0.67%
decrease in accuracy. In the second experiment, we increased the number of
attention blocks and heads to 16 each and switched to the Mish activation
function, while continuing to use only 1% of the dataset. This modification
resulted in an improved accuracy of 86.15%, with the same optimizer, learning
rate, and batch size.

Effect of training dataset Size on ParMAT Performance:To investigate the
impact of reduced class-specific data on overall performance, we retrained the
model over 20 epochs utilising only 50% of the original JETCLASS dataset in
our most recent training iteration. 8 attention heads and the same number of
layers made up the architecture’s configuration. We employed the Mish
activation function due to its proven effectiveness in handling non-
linearities. To reduce overfitting, a dropout rate of 0.20 was used in the
training, and the Ranger optimizer was used due to its effectiveness in
convergent to minima. The batch size was set to 1024.The model’s accuracy
decreased noticeably to 85.86% in spite of these configurations.This reduction
may be attributed to the decreased volume of training data for the JETCLASS,
suggesting that further adjustments in data handling or model parameters might
be necessary to restore or enhance performance.

##  8 Conclusion

We proposed Multi-axis particle transformer in this paper for jet tagging task
on particle collision data from CMS detector. Our model performed well on
JETCLASS dataset by achieving a reasonable accuracy for classifying our
signals into 10 classes and surpassed many previous ML models proposed for jet
tagging including the four baselines models. The AUC of 98.76% reflects the
strong discriminatory power of our model and its robustness. It also shows the
model’s ability to correctly identify our signals while minimizing the false
positives. It also suggest the ability of our model for jet tagging task where
precise classification is crucial and false positives can lead to significant
drawbacks. Ranger optimizer performed better than Adam for training this model
(Complete details about training parameters have been covered in section VI:
Experiment). Background rejection matrix shows that our model significantly
improved on background noise rejection for both decay modes of top quarks
$(t\xrightarrow{}bqq^{\prime}$ & $t\xrightarrow{}blv)$ and one decay mode of
Higgs$(H\xrightarrow{}gg)$. This improvement resulted from integration of
multi axis attention block which uses multi axis attention mechanism to
explore relations between different particles and jets by considering multiple
axis at a time. Results support the fact that transformers can indeed be a
good solution for classification tasks in particle physics where we have large
amounts of data generated by detectors. Furthermore, larger datasets would be
more effective for training such models in future[44].

## Acknowledgements

We thank Pattern Recognition Lab (PR-Lab) and the Pakistan Institute of
Engineering and Applied Sciences (PIEAS) for providing the necessary
computational resources and a healthy research environment.

## Declarations

### Funding/Competing interests

The authors declare that they have no known competing financial interests or
personal relationships that could have appeared to influence the work reported
in this paper.

## References

  * [1] Minxuan He and Daohan Wang.  Quark/gluon discrimination and top tagging with dual attention transformer.  The European Physical Journal C, 83(12):1116, 2023.
  * [2] Michael Spira, Abdelhak Djouadi, D Graudenz, and RM Zerwas.  Higgs boson production at the lhc.  Nuclear Physics B, 453(1-2):17–82, 1995.
  * [3] Zhengzhong Tu, Hossein Talebi, Han Zhang, Feng Yang, Peyman Milanfar, Alan Bovik, and Yinxiao Li.  Maxvit: Multi-axis vision transformer.  In European conference on computer vision, pages 459–479. Springer, 2022.
  * [4] Luke de Oliveira, Michael Kagan, Lester Mackey, Benjamin Nachman, and Ariel Schwartzman.  Jet-images—deep learning edition.  Journal of High Energy Physics, 2016(7):1–32, 2016.
  * [5] Isaac Henrion, Johann Brehmer, Joan Bruna, Kyunghyun Cho, Kyle Cranmer, Gilles Louppe, and Gaspar Rochette.  Neural message passing for jet physics.  Workshop on Deep Learning for Physical Sciences (DLPS 2017), NIPS 2017, 2017.
  * [6] Daniel Guest, Julian Collado, Pierre Baldi, Shih-Chieh Hsu, Gregor Urban, and Daniel Whiteson.  Jet flavor classification in high-energy physics with deep neural networks.  Physical Review D, 94(11):112002, 2016.
  * [7] P Komiske, E Metodiev, and J Thaler.  Pythia8 quark and gluon jets for energy flow.  Zenodo. doi, 10, 2019.
  * [8] Huilin Qu and Loukas Gouskos.  Jet tagging via particle clouds.  Physical Review D, 101(5):056019, 2020.
  * [9] Can Chen, Luca Zanotti Fragonara, and Antonios Tsourdos.  Gapointnet: Graph attention based point neural network for exploiting local feature of point cloud.  Neurocomputing, 438:122–132, 2021.
  * [10] Meng-Hao Guo, Jun-Xiong Cai, Zheng-Ning Liu, Tai-Jiang Mu, Ralph R Martin, and Shi-Min Hu.  Pct: Point cloud transformer.  Computational Visual Media, 7:187–199, 2021.
  * [11] Gregor Kasieczka, Tilman Plehn, Anja Butter, Kyle Cranmer, Dipsikha Debnath, Barry M Dillon, Malcolm Fairbairn, Darius A Faroughy, Wojtek Fedorko, Christophe Gay, et al.  The machine learning landscape of top taggers.  SciPost Physics, 7(1):014, 2019.
  * [12] CMS Collaboration, J Duarte, et al.  Sample with jet, track and secondary vertex properties for hbb tagging ml studies (higgstobbntuple_higgstobb_qcd_runii_13tev_mc).  CERN Open Data Portal, 10, 2019.
  * [13] Raghav Kansal, Javier Duarte, Hao Su, Breno Orzari, Thiago Tomei, Maurizio Pierini, Mary Touranakou, Dimitrios Gunopulos, et al.  Particle cloud generation with message passing generative adversarial networks.  Advances in Neural Information Processing Systems, 34:23858–23871, 2021.
  * [14] Eric A Moreno, Olmo Cerri, Javier M Duarte, Harvey B Newman, Thong Q Nguyen, Avikar Periwal, Maurizio Pierini, Aidana Serikova, Maria Spiropulu, and Jean-Roch Vlimant.  Jedi-net: a jet identification algorithm based on interaction networks.  The European Physical Journal C, 80:1–15, 2020.
  * [15] V Drugakov, CMS collaboration, et al.  Identification of heavy, energetic, hadronically decaying particles using machine-learning techniques.  Journal of Instrumentation, 2020.
  * [16] Muhammad Asam, Saddam Hussain Khan, Altaf Akbar, Sameena Bibi, Tauseef Jamal, Asifullah Khan, Usman Ghafoor, and Muhammad Raheel Bhutta.  Iot malware detection architecture using a novel channel boosted and squeezed cnn.  Scientific Reports, 12(1):15498, 2022.
  * [17] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
  * [18] S Belforte, V Candelise, M Casarsa, A DA ROLD, F Vazzoler, et al.  Search for top squark production in fully-hadronic final states in proton-proton collisions at sqrt (s)= 13 tev.  Physical review D, 104:1–46, 2021.
  * [19] Huilin Qu, Congqiao Li, and Sitian Qian.  Particle transformer for jet tagging.  In International Conference on Machine Learning, pages 18281–18292. PMLR, 2022.
  * [20] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [21] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  arXiv preprint arXiv:1810.04805, 2018.
  * [22] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  arXiv preprint arXiv:2010.11929, 2020.
  * [23] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 10012–10022, 2021.
  * [24] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al.  Highly accurate protein structure prediction with alphafold.  Nature, 596(7873):583–589, 2021.
  * [25] Maqsood Hayat and Asifullah Khan.  Mem-phybrid: Hybrid features-based prediction system for classifying membrane protein types.  Analytical biochemistry, 424(1):35–44, 2012.
  * [26] Asifullah Khan, Zunaira Rauf, Anabia Sohail, Abdul Rehman Khan, Hifsa Asif, Aqsa Asif, and Umair Farooq.  A survey of the vision transformers and their cnn-transformer based variants.  Artificial Intelligence Review, 56(Suppl 3):2917–2970, 2023.
  * [27] Asifullah Khan, Zunaira Rauf, Abdul Rehman Khan, Saima Rathore, Saddam Hussain Khan, Najmus Saher Shah, Umair Farooq, Hifsa Asif, Aqsa Asif, Umme Zahoora, et al.  A recent survey of vision transformers for medical image segmentation.  arXiv preprint arXiv:2312.00634, 2023.
  * [28] Abdul Rehman Khan and Asifullah Khan.  Maxvit-unet: Multi-axis attention for medical image segmentation.  arXiv preprint arXiv:2305.08396, 2023.
  * [29] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou.  Training data-efficient image transformers & distillation through attention.  In International conference on machine learning, pages 10347–10357. PMLR, 2021.
  * [30] Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei.  Beit: Bert pre-training of image transformers.  arXiv preprint arXiv:2106.08254, 2021.
  * [31] Zhemin Zhang and Xun Gong.  Rsir transformer: Hierarchical vision transformer using random sampling windows and important region windows.  arXiv preprint arXiv:2304.06250, 2023.
  * [32] Huilin Qu, Congqiao Li, and Sitian Qian.  Particle Transformer for Jet Tagging, January 2024.  arXiv:2202.03772 [hep-ex, physics:hep-ph, physics:physics].
  * [33] J. Alwall, R. Frederix, S. Frixione, V. Hirschi, F. Maltoni, O. Mattelaer, H.-S. Shao, T. Stelzer, P. Torrielli, and M. Zaro.  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(7):79, July 2014.
  * [34] An introduction to PYTHIA 8.2 - ScienceDirect.
  * [35] The DELPHES 3 collaboration, J. De Favereau, C. Delaere, P. Demin, A. Giammanco, V. Lemaître, A. Mertens, and M. Selvaggi.  DELPHES 3: a modular framework for fast simulation of a generic collider experiment.  Journal of High Energy Physics, 2014(2):57, February 2014.
  * [36] CMS Collaboration.  Identification of heavy, energetic, hadronically decaying particles using machine-learning techniques.  Journal of Instrumentation, 15(06):P06005–P06005, June 2020.  arXiv:2004.08262 [hep-ex, physics:physics].
  * [37] Huilin Qu, Congqiao Li, and Sitian Qian.  Particle transformer for jet tagging, 2024.
  * [38] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou.  Going deeper with image transformers, 2021.
  * [39] Frédéric A. Dreyer and Huilin Qu.  Jet tagging in the lund plane with graph networks, 2021.
  * [40] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Z. Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d’Alché Buc, Emily B. Fox, and Roman Garnett, editors, NeurIPS, pages 8024–8035, 2019.
  * [41] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola.  Deep sets.  Advances in neural information processing systems, 30, 2017.
  * [42] Patrick T Komiske, Eric M Metodiev, and Jesse Thaler.  Energy flow networks: deep sets for particle jets.  Journal of High Energy Physics, 2019(1):1–46, 2019.
  * [43] Maja Karwowska, Łukasz Graczykowski, Kamil Deja, Miłosz Kasak, and Małgorzata Janik.  Particle identification with machine learning from incomplete data in the alice experiment.  arXiv preprint arXiv:2403.17436, 2024.
  * [44] E. Fermi, P Pasta, S Ulam, and M Tsingou.  Studies of the nonlinear problems.  In STUDIES OF THE NONLINEAR PROBLEMS, pages LA–1940, 4376203, 1955.  Report Number: LA-1940, 4376203.


# Thermal Analysis for NVIDIA GTX480 Fermi GPU Architecture

Savinay Nagendra
Department of Computer Science
The Pennsylvania State University
<EMAIL_ADDRESS>
###### Abstract

In this project, we design a four-layer (Silicon—TIM—Silicon—TIM), 3D floor
plan for NVIDIA GTX480 Fermi GPU architecture and compare heat dissipation and
power trends for matrix multiplication and Needleman-Wunsch kernels. First,
cuda kernels for the two algorithms are written. These kernels are compiled
and executed with the GPGPU Simulator to extract power logs for varying tensor
sizes. These power logs are converted to ptrace files with an automation
script written in Python. The 3D floor plan, along with the generated ptrace
files are given to HotSpot, which generates thermal heat maps to show heat
dissipation for various components of the Fermi architecture. These heat
dissipation trends for both the kernels are observed for multiple tensor sizes
to draw qualitative conclusions. The behavioral and execution patterns of both
the kernels are also observed with these varying heat dissipation trends. With
this project, we observe that increase in tensor size results in increase of
heat dissipation in components of the Fermi Architecture. However, the
temperature of the chip remains saturates after a particular tensor size, and
remains constant thereafter. Heat dissipation is non-uniform with smaller
tensor sizes, and becomes more uniform after a certain tensor size. This
means, after a particular tensor size, more cores of the architecture get
activated in the computations, thereby resulting in an almost constant
temperature. We also observe that Needleman Wunsch uses more data movement
between DRAM and caches, thereby showing higher heat dissipation patterns in
DRAMs when compared to Matrix multiplication for the same tensor size. Our
observations are in accordance with the theoretical concepts behind the
working of the two algorithms, thereby making our results consistent. The link
to the project files and scripts is
https://drive.google.com/file/d/1ukmSbWnk4ctFHF87x3lNsMU5qijbEAPu/view?usp=share_link

##  1 Design

In this section, we will discuss the system design for thermal analysis.

###  1.1 Tools or Software packages Used

In this section, we will discuss the tools used in this project.

####  1.1.1 GPGPU-Sim

GPGPU-Sim [1] is a micro-architecture simulator that provides a detailed
simulation model of a contemporary GPU. GPGPU-Sim simulates a kernel by first
transferring data to GPU memory. Then, the GPU kernels run on the GPGPU-Sim
which reports statistics for the kernels. Finally, the data is transferred
back to CPU memory. This simulator provides mechanisms for more efficient SIMD
branch execution on GPUs [11, 9, 2, 6, 14, 10, 7, 12, 8, 13, 17]. This
simulator is capable of dynamically regrouping threads into new warps on the
fly following the occurrence of diverging branch outcomes. Power efficiency
has become a more crucial evaluation metric than peak performance for
mainstream GPGPU computing.

Figure 1: GPU Watch Power model integrated with GPGPU-Sim

GPU Watch [5] as shown in Fig. 1 is a GPGPU power model that is configurable,
capable of cycle-level calculations, and carefully validated against real
hardware measurements. It has been integrated with GPGPU-Sim. This power model
estimates the power consumed by the GPU according to the timing behavior. This
makes it ideal for evaluating fine-grained power management mechanisms. GPU
watch uses a bottom-up method to build the initial model. The simulated power
from this model is compared against the measured hardware power to identify
modeling inaccuracies. These are resolved with a special suite of 80
microbenchmarks that are designed to create a system of linear equations that
correspond to the total power consumption. In this project, we are using the
integrated GPU Watch/GPGPU-Sim to generate the power report for our kernels.

####  1.1.2 HotSpot

Hotspot [16] is a pre-RTL thermal simulator that is used in the early design
process to analyse heating patterns of the designed architectures. It supports
simulation of traditional 2D Integrated Circuits, 3D Integrated Circuits and
Microfluiding Cooling. Hotspot takes as input a power trace file with power
statistics of each component in the architecture and a 2D/3D floor plan of the
architecture to generate heat dissipation maps of the architecture. The color
variations indicate differential heating of different components, thereby
allowing researchers to improve their designs.

###  1.2 Cuda Kernels

In this section, we will discuss the implementation of matrix multiplication
and Needleman-Wunsch kernels in written in cuda. In cuda, there is the concept
of grids and blocks. Blocks are 3D arrays of threads and Grids are 3D arrays
of blocks. Each thread knows the $x$ and $y$ coordinates of the block it is in
the grid, and the coordinates of where the thread itself is in the block.
These positions are used to calculate a unique thread ID for every thread. The
total number of threads in a single block cannot exceed 1024.

####  1.2.1 Linearized Matrix Multiplication Kernel

Assume we have two matrices $A$ and $B$, where $A$ is a $n\times m$ matrix and
$B$ is a $m\times w$ matrix. Then, the result $M=A\times B$ is a $n\times w$
matrix. The first element of $M$ is the sum of all the element-wise
multiplications of the numbers in the first row of $A$ with the numbers in
first column of $B$. We have to compute every element in $M$, where each of
them is independent from the others. This means, matrix multiplication can be
effectively parallelized. We realized with experiments that regular matrix
multiplication with three for loops is not an effective way of implementing
the kernel. The simulation time this was taking was significantly high, not
allowing us to use tensor sizes beyond $80\times 80$.

Figure 2: Linearized Matrix Multiplication Kernel [15]

So, we decided to make use of 1D arrays for our matrices. The insight is the
way Cuda is developed, where the number of columns must be known before
compiling the program. Linearizing the matrices allows us to know this, and
dynamically calculate the rows and columns for indexing. Linearization was
done by stacking each row length-wise from first to last. The result for an
index is calculated by $C[row*N+col]=A[row*n+i]*B[i*N+col]$, where $i$ is the
index in the linearized array (thread number) as shown in Fig. 2 This cuda
kernel was referred from [15]. The distribution of work between threads across
blocks is make such that there is one thread responsible for one $(i,j)$
index. So, the total number of threads is equal to the number of elements in
the matrix.

####  1.2.2 Linearized Matrix Multiplication Kernel optimized by using
Scratchpad memory (Cache Tiling)

Figure 3: Linearized Matrix Multiplication Kernel optimized with Scratchpad
memory.

Without optimization, every single thread is marched along main memory,
computing the product of two elements and then accumulating into the result.
With Shared memory (scratchpad memory) optimization, we are going to do the
same operation, but in chunks of thread block dimension ($32\times 32)$
elements. Each thread block will compute chunks of 32 elements of matrix
multiplication at a time. This way, we know what exactly the memory we are
going to be accessing next. We are exploiting this fact to load the next tile
into the private cache of the Streaming Multiprocessor (SM). So, we will load
in 32 chunks into the shared memory as shown in Fig. 3 When all the threads
are accessing those chunks, instead of going into the main memory and hoping
to get lucky in the cache hierarchy, we can guarantee that the memory access
will be in the private L1 level, making it very fast. This way, we get the
performance speedup, DRAM accesses can be reduced. With this optimization, it
can be observed that this kernel executes in significantly lesser amount of
time when compared to the previous one, along with significantly lesser DRAM
access.

####  1.2.3 Needleman-Wunsch Kernel

Figure 4: Needleman-Wunsch kernel optimized with Scratchpad memory.

The Needleman-Wunsch algorithm is a 2D dynamic programming algorithm that is
used in bioinformatics to align protein or necleotide sequences. The algorithm
divides the large problem into a series of smaller problems, and uses the
solutions to the smaller problems to find the optimal solution to the larger
problem. The cuda kernel shown in figure Fig. 4 uses the following psuedo-code
for implementation as shown in Fig. 5. The code for this kernel implementation
is taken from [3].

Figure 5: Needleman-Wunsch Psuedo Code

###  1.3 NVIDIA GTX480 Fermi

In this section we will discuss the GTX480 fermi architecture [4] and floor
plan design.

####  1.3.1 Fermi Architecture

The Fermi architecture, as shown in Fig. 6 has upto 512 CUDA cores organized
in 16 Streaming Multiprocessors (SM) of 32 cores each, where each CUDA core
executes a floating point or integer instruction per clock for a thread.

Figure 6: Fermi Architecture

The GPU has six 64-bit memory partitions, for a 384-bit memory interface,
supporting upto a total of 6GB of GDDR5 DRAM memory. PCI-Express bus connects
the GPU to the CPU. The GigaThread global scheduler is responsible for
distributing thread blocks to SM thread schedulers.

Figure 7: Fermi Streaming Multiprocessor

The SMs are the execution units of the GPU. Each SM has 32 CUDA cores. Each
CUDA processor has a fully pipelined ALU and FPU. Each SM has 16 load/store
units, allowing source and destination addresses to be calculated for 16
threads per clock. Each SM has its own 64KB L1 cache, an uniform cache, an
Instruction cache, a Scheduler-dispatcher unit and a $32768\times 32-bit$
register file as shown in Fig. 7.

####  1.3.2 Fermi Floor Plan

The floor plan was designed for TIM Layers 0 and 2 based on the diagram in
Fig. 6. The chip dimensions were taken to be $23mm\times 23mm$. The dimensions
for each component was proportionally calculated and put in two $.flp$ files,
which were fed to HotSpot. Only the size DRAMs and the L2 cache were placed on
Layer0. The rest of the components were made into a void placeholder. The 16
SMs and their respective components were place on Layer2. Only the components
for which the power values are generated by GPGPU-Sim were considered while
designing the floor plan.

##  2 Methodology

In this section, we will discuss the flow of the project. Fig. 8 shows the
flow of this project. First CUDA kernels as mentioned in the previous section
are written.

Figure 8: Block diagram showing the flow of this project. The green boxes
indicate the work I have done.

A kernel along with the GTX480 config file are given as inputs to the GPSPU-
Sim + GPUWatch interface to generate power report. This is repeated for all
the kernels and varying tensor sizes. A python script parses the power report
to extract power values for each component of the Fermi architecture. The
power report has three values - minimum, average and maximum, Two other values
are added by using interpolation, such that the average of the five values
remains the same as the previous average. The output of the python script is a
$.ptrace$ file that will be given as an input to HotSpot tool. I have also
written a python script to generate Layer0 and Layer2 floor plans as mentioned
in the previous section. This script uses the proportionality of component
dimensions in Fig. 6 and generates two $.flp$ files. The $.ptrace$ file along
with the two $.flp$ files are given as inputs to the HotSpot code. The
$run.sh$ file is modified accordingly. Running this file will generate heat
maps for the four layers. We only use/need the heat maps for Layer2 in all our
results. We observe that Linearized Matrix Multiplication kernel does not
provide significantly different results in terms of DRAM heat dissipation for
increasing tensor sizes. Hence, we also decrease the L2 cache size in the
config file (along with the corresponding L2 cache area in the floor plan) to
observe the DRAM’s heating patterns. Since heating patterns are more expressed
with reduced L2 cache size, the cache tiling results are also taken for
reduced L2 mode, so that we can compare how DRAM accesses are reduced with
CUDA kernel optimization. While running Needleman-Wunsch kernels, we maintain
the original L2 cache size. Matrix multiplication and Needleman-Wunsch kernels
are varied for tensor sizes of 100, 250, 400 and 800.

##  3 Evaluation/Results

In this section, we will observe and infer from the heat map dissipation
patterns for the varying tensor sizes for the three matrix multiplication
kernels.

###  3.1 Linearized Matrix Multiplication Kernel

In this section, we will look at the linearized matrix multiplication kernel’s
heat dissipation patterns. It can be observed that from figures 9 to 12 that
with increase in tensor size from 100 to 800, the heat dissipated becomes more
and more uniform, which leads to maintaining a constant temperature or even
reduced temperature for the chip after $400\times 400$. With an increase from
$400\times 400$ to $800\times 800$, it can be observed that the DRAM closest
to L2 cache dissipates more heat. This is consistent with the theoretical
observation where $800\times 800\times 4=2500KB$ data cannot fit on a $768KB$
L2 cache. Only this DRAM is used more so that the data movement between the
DRAM and cache can be minimized. With increase in tensor size, more cores get
activated, thereby making the heat dissipation uniform.

Figure 9: Linearized Multiplication Kernel: 100x100 Figure 10: Linearized
Multiplication Kernel: 250x250 Figure 11: Linearized Multiplication Kernel:
400x400 Figure 12: Linearized Multiplication Kernel: 800x800

###  3.2 Linearized Matrix Multiplication Kernel with reduced L2 cache size

Since execution times of the linearized kernel is high, I was only able to run
matrix multiplication kernel for a maximum tensor size of $800\times 800$.
However, reducing L2 size provides the same effect as increasing tensor size .
With this, I reduced the L2 size of the GPU to $\frac{1}{4}^{th}$ of its
original value of $768KB$. The results are shown in figures 13 to 16. It can
be seen clearly that, with reduced L2 cache size, the GPU is now accessing
DRAM significantly, leading to higher heat dissipation. Also, from Fig. 16, it
can be observed that all the three DRAMs on the right side are being used,
leading to higher heat dissipation in each of them. However, the DRAM closest
to L2 cache is heated up the most, as that is the one being used the most in
order to minimize data movement.

Figure 13: Linearized Multiplication Kernel with reduced L2: 100x100 Figure
14: Linearized Multiplication Kernel with reduced L2: 250x250 Figure 15:
Linearized Multiplication Kernel with reduced L2: 400x400 Figure 16:
Linearized Multiplication Kernel with reduced L2: 800x800

###  3.3 CUDA optimized Linearized Matrix Multiplication Kernel with reduced
L2 cache size

Figures 17 to 20 show the heat dissipation patterns for cuda optimized
linearized multiplication kernel. The optimization involves using shared or
scratchpad memory by a method called cache tiling. This leads to deterministic
data accesses, where the tiles can now be conveniently stored in the private
L2s of every SM. The experiments have been conducted with a reduced L2 cache
size to observe pronounced patterns in DRAM accesses. Theoretically, since the
data accesses are deterministic and efficient caching system can now be
adapted, DRAM accesses have to minimized, thus leading to lesser DRAM heat
dissipation. This has been observed clearly in the figures 17 to 20. Comparing
figure 16 and 20, it is clearly seen that the DRAM access pattern with cuda
optimization is as expected. Infact, the optimization is so powerful that the
GPU is behaving as if the L2 cache size is the same as the original, i.e., the
even with reduced cache size, DRAM accesses are minimal.

Figure 17: Cuda optimized Linearized Multiplication Kernel with reduced L2:
100x100 Figure 18: Cuda optimized Linearized Multiplication Kernel with
reduced L2: 250x250 Figure 19: Cuda optimized Linearized Multiplication Kernel
with reduced L2: 400x400 Figure 20: Cuda optimized Linearized Multiplication
Kernel with reduced L2: 800x800

###  3.4 Needleman-Wunsch Kernel

Needleman-Wunsch is a quadratic dynamic programming algorithm. Unlike matrix
multiplication, the thread blocks are not independent. At any given time, for
computing the score for $(i+1,j+1)$ index, we need data from rows $(i,j)$.
This causal correlation between the thread blocks essentially reduces
parallelization. As the tensor size increases, the DRAM gets significantly
heated up, especially ones closest to the L2 cache. DRAM accesses is
significantly more than that of linearized matrix multiplication kernel, for
the tensor size of $800\times 800$. This is because, same parts of L2 cache
are being accessed each time by multiple thread blocks, leading to conflicts
in L2 cache. Also, as the computation by different thread blocks is not
independent, for memory consistency data needs to be constantly written to
DRAM.

Figure 21: Needleman-Wunsch: 100x100 Figure 22: Needleman-Wunsch: 250x250
Figure 23: Needleman-Wunsch: 400x400 Figure 24: Needleman-Wunsch: 800x800
Figure 25: Needleman-Wunsch: 1200x1200

###  3.5 Comparison between CUDA kernels

Based on the heat map patterns, following are the conclusions we can draw
regarding the CUDA kernels:

  1. 1.

Linearized matrix multiplication takes significantly more time to execute when
compared to optimized linearized matrix multiplication. The computations are
faster due to lesser data movement and deterministic data accesses in the L1
cache of each SM.

  2. 2.

For the same tensor size, optimized linearized matrix multiplication
significantly uses lesser DRAM accesses when compared to regular linearized
matrix multiplication. This is because of the use of cache tiling or
scratchpad memory.

  3. 3.

Needleman-Wunsch kernel is memory intensive when compared to linearized matrix
multiplication. In matrix multiplication, since the threads can work
independently, and each subproblem is unrelated to the other, the L2 accesses
will have lesser conflicts when compared to Needlman-Wunsch, where the data in
the left, top and diagonally top-left are necessary for an element’s
computation. This is why for a tensor-size of $800\times 800$, DRAM gets
heated up more in Needleman-Wunsch when compared to matrix multiplication.

##  4 Learning and Key Takeaways

  1. 1.

This is the first time I am writing CUDA code. Initially, I was not aware of
the benefits of linearization of matrices for GPU threads. Eventually, when I
came accross linearization benefits, it was truly an eye-opener. I could run
the kernel for higher tensor sizes with lesser time.

  2. 2.

CUDA optimization using the concept of scratchpad memory was very intriguing
to me. I had a very good learning experience in applying what was learnt in
class to a practical problem, and seeing it perform significantly better than
non-optimized kernel in terms of the execution time and DRAM accesses.

  3. 3.

Since I work in the field of machine learning, learning the intricacies and
architecture of a GPU and being able to develop a floor plan was a great
takeaway.

  4. 4.

I also learnt to use GPGPU-Sim, which was again a first time for me. I feel
more confident in speaking about the practicals of computer architecture after
the assignments and projects in this course.

##  5 Acknowledgements

I would like to thank Prof. Vijay for the wonderful semester and intuitive
classroom sessions. This is my final course of my PhD career, and it was
definitely an amazing learning experience.

## References

  * [1] Wilson WL Fung, Ivan Sham, George Yuan, and Tor M Aamodt.  Dynamic warp formation and scheduling for efficient gpu control flow.  In 40th Annual IEEE/ACM International Symposium on Microarchitecture (MICRO 2007), pages 407–420. IEEE, 2007.
  * [2] Christopher Funk, Savinay Nagendra, Jesse Scott, Bharadwaj Ravichandran, John H Challis, Robert T Collins, and Yanxi Liu.  Learning dynamics from kinematics: Estimating 2d foot pressure maps from video frames.  arXiv preprint arXiv:1811.12607, 2018.
  * [3] GitHub.  Needleman-wunsch implementation : https://github.com/anujc25/parallel-nw-implementation.
  * [4] Peter N Glaskowsky.  Nvidia’s fermi: the first complete gpu computing architecture.  White paper, 18, 2009.
  * [5] Jingwen Leng, Tayler Hetherington, Ahmed ElTantawy, Syed Gilani, Nam Sung Kim, Tor M Aamodt, and Vijay Janapa Reddi.  Gpuwattch: Enabling energy optimizations in gpgpus.  ACM SIGARCH Computer Architecture News, 41(3):487–498, 2013.
  * [6] Jiangtao Liu, Chaopeng Shen, Te Pei, Kathryn Lawson, Daniel Kifer, Savinay Nagendra, and Srikanth Banagere Manjunatha.  A new rainfall-induced deep learning strategy for landslide susceptibility prediction.  In AGU Fall Meeting Abstracts, volume 2021, pages NH35E–0504, 2021\.
  * [7] Savinay Nagendra, S Banagere Manjunatha, Chaopeng Shen, Daniel Kifer, and Te Pei.  An efficient deep learning mechanism for cross-region generalization of landslide events.  In AGU Fall Meeting Abstracts, volume 2020, pages NH030–0010, 2020\.
  * [8] Savinay Nagendra and Daniel Kifer.  Patchrefinenet: Improving binary segmentation by incorporating signals from optimal patch-wise binarization.  In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 1361–1372, 2024.
  * [9] Savinay Nagendra, Daniel Kifer, Benjamin Mirus, Te Pei, Kathryn Lawson, Srikanth Banagere Manjunatha, Weixin Li, Hien Nguyen, Tong Qiu, Sarah Tran, et al.  Constructing a large-scale landslide database across heterogeneous environments using task-specific model updates.  IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, 15:4349–4370, 2022.
  * [10] S Nagendra, T Pei, S Banagere Manjunatha, G He, T Qiu, D Kifer, and C Shen.  Cloud-based interactive database management suite integrated with deep learning-based annotation tool for landslide mapping.  In AGU Fall Meeting Abstracts, volume 2020, pages NH030–0011, 2020\.
  * [11] Savinay Nagendra, Nikhil Podila, Rashmi Ugarakhod, and Koshy George.  Comparison of reinforcement learning algorithms applied to the cart-pole problem.  In 2017 international conference on advances in computing, communications and informatics (ICACCI), pages 26–32. IEEE, 2017.
  * [12] Savinay Nagendra, Chaopeng Shen, and Daniel Kifer.  Threshnet: Segmentation refinement inspired by region-specific thresholding.  arXiv preprint arXiv:2211.06560, 2, 2022.
  * [13] Savinay Nagendra, Chaopeng Shen, and Daniel Kifer.  Estimating uncertainty in landslide segmentation models.  arXiv preprint arXiv:2311.11138, 2023.
  * [14] Te Pei, Savinay Nagendra, Srikanth Banagere Manjunatha, Guanlin He, Daniel Kifer, Tong Qiu, and Chaopeng Shen.  Utilizing an interactive ai-empowered web portal for landslide labeling for establishing a landslide database in washington state, usa.  In EGU General Assembly Conference Abstracts, pages EGU21–13974, 2021.
  * [15] Quantstart.  Matmul linearized implementation: https://www.quantstart.com/articles/matrix-matrix-multiplication-on-the-gpu-with-nvidia-cuda/.
  * [16] Virginia Tech.  Hotspot: https://github.com/uvahotspot/hotspot.  GitHub.
  * [17] L Zhu, P Tilke, S Nagendra, M Etchebes, and M LeFranc.  A rapid and realistic 3d stratigraphic model generator conditioned on reference well log data.  In Second EAGE Digitalization Conference and Exhibition, volume 2022, pages 1–5. European Association of Geoscientists & Engineers, 2022.