[1, 2]Department, of Mathematics, IIT Guwahati, Guwahati, 781039, Assam, India
[3]Corresponding Author : Arabin Kumar Dey; Email address :
<EMAIL_ADDRESS>
# Area-norm COBRA on Conditional Survival Prediction

Rahul Goswami<EMAIL_ADDRESS>Arabin Kr. Dey<EMAIL_ADDRESS>*  *

###### Abstract

The paper explores a different variation of combined regression strategy to
calculate the conditional survival function. We use regression based weak
learners to create the proposed ensemble technique. The proposed combined
regression strategy uses proximity measure as area between two survival
curves. The proposed model shows a construction which ensures that it performs
better than the Random Survival Forest. The paper discusses a novel technique
to select the most important variable in the combined regression setup. We
perform a simulation study to show that our proposition for finding relevance
of the variables works quite well. We also use three real-life datasets to
illustrate the model.

###### keywords:

Cox’s model, Weak learners, Survival Tree, Combined Regression Strategy

##  1 Introduction

Prediction of Conditional survival function is an important relevant
information for doctors, clinicians, and insurance companies. Function
prediction has a large number of applications in many areas of science and
technology. This is a challenging problem. Researchers attempted to provide a
reasonable generic solution over decades. After the Cox proportional hazard
model ([1], [2]), one of the breakthrough contributions to this problem is
Random Survival Forest by Ishawaran [3]. The solution encourages researchers
to explore the problem through different other neural network models [[4],
[5], [4, 5]] and ensemble techniques which include XGboost, Adaboost ([6])
etc. In technical language, this is a problem of constructing a functional
regression. Goswami et al ([7], [8]) provide different solutions to this
problem in a combined regression setup. However, the solution has issues and
demands a more detailed exploration to get a complete idea of the different
possible implementations and their benefits. The paper addresses a different
variation of the COBRA which works for certain aspects and certain types of
datasets. The paper proposes to use a different proximity measure for this
combined regression strategy and a specific set of weak learners which ensures
that the proposed algorithm works much better than the usual Random Survival
Forest. It ensures the algorithm will work at least as well as a random
survival forest if the dataset consists of the censored observation up to a
certain level.

All implementation in this paper assumes datasets are sufficiently large. This
ensures the proximity measure to include a reasonable number of points in the
prediction set. It uses Kaplan Meier estimation ([9]) or Nelson-Aalen [10, 11]
estimator to estimate the survival function or cumulative hazard function
respectively based on all similar observations. We use both concordance
measure and integrated Brier score in a censored setup to tune the
hyperparameters of the model and use them to understand the goodness of fit
for each dataset used in this paper. The approach does not use similar weak
learners as used by Goswami et al. The weak learner’s choices in this paper
follow the same as the original combined regression paper with its similar
extension in functional regression setup. Thus we take each base learner as a
survival predictor. We use different survival predictors instead of just one
type of survival predictor used by Goswami et al. : (1) Cox model (ridge) [12]
(2) Cox model (lasso) [13] (3) random survival forest [14] (4) KNN survival
([15], [16] ) (5) Survival Tree [14]. These are exact functional extensions of
the following models (1) Ridge (2) Lasso (3) Random Forest [14] (4) KNN [17]
(5) Decision tree used respectively in original cobra for usual regression. We
observe the choice can help us to ensure that the algorithm works better than
a Random survival forest for sufficiently large datasets with lesser censored
data. The experimental results open a large number of issues and questions. We
address the issues in other research articles.

The relevance of the covariates is also an important aspect of survival
analysis. Earlier statistical models like Cox Proportional Hazard Model [2]
and Accelerated Failure Time Model [18] have a clear idea of the covariates,
which are relevant for the happening of the event of interest at a specific
time. The importance of the covariates in predicting the survival function
measures the relevance of the covariates. Random Survival Forest [3] uses the
out-of-bag error ([14], [3], [19], [20], [21]) and minimal depth procedure.
Feature importance is crucial for providing the interpretability of the model.
Independent approaches like LIME [22] and Shapley Value [23] use variable
importance concepts in ordinary regression setup. The approaches like SurvLIME
([24, 25], [26]) and SurvSHAP ([27], [28]) serves the similar purpose of
bringing interpretability of survival model through measuring the importance
of features.

In this paper, we propose a novel technique to find out the most important
variable. It can also find a score to rank the importance. The technique
exploits the cobra structure to find the variable importance. The proposition
has support on a simulated framework.

The organization of the paper is as follows. Section 2 provides a brief
overview of the notations and basic framework of survival analysis and
original combined regressions strategy used for regression. We describe the
proposed framework for aggregating the survival-based functional models in
Section 3. Numerical results of the experiments conducted to evaluate the
performance of the proposed framework are available in Section 4. A novel
proposition for ranking the variables or selecting the most important variable
exploiting the Cobra Structure is available in Section 5. Finally, Section 6
concludes the paper and discusses future research directions.

##  2 Background

### Important Notations and Technical Specification of the Problem

Our goal is to develop individual survival function. Each patient $i$ have a
d-dimensional covariate $x_{i}\in\mathcal{X}$ which represents the unique
characteristics of the patient. Here, $\mathcal{X}$ is the covariate space. An
outcome of interest $T_{i}\in\mathbb{R}^{+}$, representing the time of
occurrence of the event of interest. The patient may drop out of the study
before the event of interest occurs. We denote this dropout time or censored
time as $C_{i}\in\mathbb{R}^{+}$. The observations carry another variable,
$\delta_{i}\in\\{0,1\\}$. $\delta_{i}=1$ when event of interest is available,
and $\delta_{i}=0$ if the event of interest is not available.

We can write $S(t\mid x_{i})$ as

$S(t\mid x_{i})=\mathbb{P}(T_{i}>t\mid x_{i})$ (1)

It represents the probability of survival of the patient $i$ beyond time $t$
given the covariate $x_{i}$. Note that, our aim is to estimate the survival
function $S$ specific to the patient $i$, not the unconditional survival
function.. The relationship between the patient covariates and survival
function is to be estimated by a dataset ${\cal D}_{n}$ comprising of $n$
individuals assumed to be drawn from the random tuple
$(X_{i},\delta_{i}T_{i}+(1-\delta_{i})C_{i},\delta_{i})$ . For simplicity we
denote the random tuple
$(X_{i},\delta_{i}T_{i}+(1-\delta_{i})C_{i},\delta_{i})$ by
$(X_{i},Y_{i},\delta_{i})$. Since, we do not observe both $T_{i}$ and $C_{i}$
concurrently in practice, we denote the observed time by
$Y_{i}=\delta_{i}T_{i}+(1-\delta_{i})C_{i}$

### Combined Regression Strategy for Usual Regression Framework

Let’s consider the regression setting, where ${\cal
D}_{n}=(X_{1},Y_{1}),\cdots,(X_{n},Y_{n})$ are independent and identically
distributed samples from $(\mathbf{X},Y)\in\mathbb{R}^{p}\times\mathbb{R}$. We
assume $\mathbb{E}Y^{2}<\infty$. In a typical regression setup we are
interested in estimating the conditional expectation of $Y$ given $X=x$ i.e
$E(Y\mid X=x)$.

In COBRA setup, we split the the dataset ${\cal D}_{n}$ into two parts, one
part ${\cal D}_{k}=(X_{1},Y_{1})\cdots(X_{k},Y_{k})$ and the other part ${\cal
D}_{l}=(X_{k+1},Y_{k+1}),\cdots,(X_{n},Y_{n})$ where $l=n-k\geq 1$. We denote
${\cal D}_{l}=(X_{1},Y_{1}),\cdots,(X_{l},Y_{l})$.

Suppose, we have M competing estimators, referred to as machines in COBRA
setup. Each machine $m$ is trained on ${\cal D}_{k}$ and referred as
$r_{k1},r_{k2},\cdots,r_{kM}$, here we note that $r_{ki}$ is a machine trained
only on ${\cal D}_{k}$ and capable of estimating $E(Y\mid X=x)$ for any
$x\in\mathbb{R}^{p}$. Then the estimate of $E(Y\mid X=x)$ by COBRA is given by

$\hat{E}_{COBRA}(Y\mid X=x)=\sum_{i=1}^{l}W_{n,i}(x)Y_{i}$ (2)

where $W_{n,i}(x)$ is the weight of the $i^{th}$ sample in ${\cal D}_{l}$ and
is given by

$W_{n,i}(x)=\frac{\mathbb{I}(\cap_{m=1}^{M}\mid
r_{k,m}(X_{i})-r_{k,m}(x)\mid\leq\epsilon)}{\sum_{j=1}^{l}\mathbb{I}(\cap_{m=1}^{M}\mid
r_{k,m}(X_{j})-r_{k,m}(x)\mid\leq\epsilon)}$ (3)

where, $\epsilon$ is a user specified parameter. The intuition behind the
weight is that if the $i^{th}$ sample is close to the query point $x$ in all
the machines, then the weight of the $i^{th}$ sample is 1 and if the $i^{th}$
sample is far from the query point $x$ in any machine, then the weight of the
$i^{th}$ sample is 0. This condition for consensus of a fraction of the
machines brings a new parameter $\alpha$, which tells us how many machines
should agree on the closeness of the $i^{th}$ sample to the query point
$x$.The weight of the $i^{th}$ sample is given by

$W_{n,i}(x)=\frac{\mathbb{I}(\sum_{m=1}^{M}\mathbb{I}(\mid
r_{k,m}(X_{i})-r_{k,m}(x)\mid\leq\epsilon)\geq
M\alpha)}{\sum_{j=1}^{l}\mathbb{I}(\sum_{m=1}^{M}\mathbb{I}(\mid
r_{k,m}(X_{j})-r_{k,m}(x)\mid\leq\epsilon)\geq M\alpha)}$ (4)

COBRA gives a good approximation of the conditional expectation of $Y$ given
$X=x$, there is both theoretical and experimental proof that COBRA perform
better that individual machine in terms of quadratic loss.

##  3 Proposed Algorithm for Functional Regression

The structure of the survival data takes the shape of ${\cal
D}_{n}:\\{X_{i},Y_{i},\delta_{i}\\}_{i=1}^{n}$, where $X_{i}$ signifies the
covariate under consideration. Correspondingly, $Y_{i}$ represents the lesser
of the event of interest occurrence time and the censoring time while
$\delta_{i}$ serves as the indicator that offers insight into whether we
observe the event of interest or not.

Let us consider of the set survival models $\mathcal{M}=\\{m_{1},m_{2}\dots
m_{\mid\mathcal{M}\mid}\\}$. Each model of the set $\mathcal{M}$ provides a
survival function. Since we use five fold cross validation, we divide the
dataset into five parts, where each part represents test set and rest as train
set. Suppose, we denote the train data as ${\cal D}_{n^{{}^{\prime}}}$ which
is of size $n^{{}^{\prime}}$ and split the dataset ${\cal
D}_{n^{{}^{\prime}}}$ into ${\cal D}_{k}$ and ${\cal D}_{l}$ such that
$k+l=n^{{}^{\prime}}$. Let us assume the train and test dataset as ${\cal
D}_{k}=\\{(X_{i},Y_{i},\delta_{i})\\}_{i=1}^{k}$ and ${\cal
D}_{l}=\\{(X_{i},Y_{i},\delta_{i})\\}_{i=k+1}^{n^{{}^{\prime}}}$. We use the
subdivision of the dataset ${\cal D}_{k}$ to train the base models, which
splits out individual survival function for a given query point $x$.
$S_{k,m}(t\mid x)$ denotes the estimated the survival function based on
$m^{th}$ base model, and $k$ is the dataset on which we train the base model.
We define a function as follows:

$\Gamma(x;X_{j},\mathcal{M},\epsilon,\alpha)=\mathbb{I}\biggl{(}\sum_{m=1}^{\mid\mathcal{M}\mid}\mathbb{I}(d(S_{k,m}(t\mid
x),S_{k,m}(t\mid X_{j}))\leq\epsilon)\geq\mid\mathcal{M}\mid\alpha\biggr{)}$
(5)

The k subscript notify that the prediction of the base models are trained on
the dataset ${\cal D}_{k}$. Here, $\mathbb{I}$ is the indicator function,
which is 1 if the condition is true and 0 if the condition is false.
$\epsilon$ is the threshold distance, $\alpha$ is the fraction of base models
in the in consensus of $\epsilon$ proximity of the query point $x$. basically
$\Gamma(.)$ is the indicator function, which is 1 if the $i^{th}$ individual
is in the $(\epsilon,\alpha)-$proximity of the query point $x$, otherwise 0.

In this paper we take the following area-norm to calculate the distance
function :

$d(S_{i}(t),S_{j}(t))=\frac{1}{t}\int_{0}^{t}\mid S_{i}(t)-S_{j}(t)\mid dt$
(6)

However, computing this area-norm is computationally expensive. Let’s assume
$t_{0}=0$. We use the following Euler approximation as distance function
between survival function :

$d(S_{i}(t),S_{j}(t))=\frac{1}{t_{n^{{}^{\prime}}}-t_{1}}\sum_{k=1}^{n^{{}^{\prime}}}\mid
S_{i}(t_{k})-S_{j}(t_{k-1})\mid(t_{k}-t_{k-1})$ (7)

The other part of the dataset ${\cal D}_{l}$ aggregate the survival function
of the base models. We calculate the event time, event count and survivor
count using the following equations,

$\mathcal{Y}_{\Gamma}(x;h)=\biggl{\\{}Y_{j}:j\geq
k\cap\delta_{j}\Gamma(x;X_{j},\mathcal{M},\epsilon,\alpha)=1\biggr{\\}}$ (8)

$\mathcal{D}_{\Gamma}(t\mid
x;h)=\sum_{j=k}^{n^{{}^{\prime}}}\Gamma(x;X_{j},\mathcal{M},\epsilon,\alpha)\delta_{j}\mathbb{I}(Y_{j}=t)$
(9)

$\mathcal{R}_{\Gamma}(t\mid
x;h)=\sum_{j=k}^{n^{{}^{\prime}}}\Gamma(x;X_{j},\mathcal{M},\epsilon,\alpha)\mathbb{I}(Y_{j}\geq
t)$ (10)

Then using the equations 8, 9 and 10, we can construct the aggregated survival
function $\hat{S}_{COBRA}(t\mid x;h)$ using the following equation

$\hat{S}_{COBRA}(t\mid
x;h)=\prod_{t^{{}^{\prime}}\in\mathcal{Y}_{\Gamma}(x;h)}\biggl{(}1-\frac{\mathcal{D}_{\Gamma}(t^{{}^{\prime}}\mid
x;h)}{\mathcal{R}_{\Gamma}(t^{{}^{\prime}}\mid
x;h)}\biggr{)}^{\mathbb{I}(t^{{}^{\prime}}\leq t)}$ (11)

The proposed algorithm is available in Algorithm 1.

Algorithm 1 Overall flow of the proposed Algorithm

1:

2:${\cal D}_{n}$ : Dataset of $n$ individuals

3:${\cal D}_{n^{{}^{\prime}}}$ : Dataset of $n^{{}^{\prime}}$ train
individuals

4:$\epsilon$ : Threshold distance

5:$\alpha$ : Fraction of base models in the consensus of $\epsilon$ proximity
of the query point $x$.

6:$M$ : Number of base models

7:

8:$\hat{S}(t\mid x)$ : Estimated survival function of the query point $x$ at
time $t$

9:

10:Partition the whole data into five parts. Take each part as test set.

11:Split train data ${\cal D}_{n^{{}^{\prime}}}$ into ${\cal D}_{k}$ and
${\cal D}_{l}$ such that $k+l=n^{{}^{\prime}}$

12:Train $M$ base models on ${\cal D}_{k}$ and predict the survival function
$S_{k,m}(t\mid x)$ for $m=1,2\dots M$

13:Calculate the distance $d(S_{k,m}(t\mid x),S_{k,m}(t\mid x_{i}))$ for
$m=1,2\dots M$ and $i=k+1,k+2,\cdots,n^{{}^{\prime}}$ using 7

14:Calculate the indicator function $\Gamma_{i}(x)$ for
$i=k+1,k+2,\cdots,n^{{}^{\prime}}$ using 5

15:Calculate the estimated survival function $\hat{S}(t\mid x)$ using 11

##  4 Dataset Descriptions and Numerical Experiments

###  4.1 Datasets

We evaluate the performance of the proposed method on the following datasets:

  * •

METABRIC 111Available at https://github.com/chl8856/DeepHit : The Molecular
Taxonomy of Breast Cancer International Con- sortium (METABRIC) dataset
comprises of gene expression profiles and clinical attributes aimed at
discerning breast cancer subtypes. Among the 1,981 patients included, 888
(44.8%) had complete follow-up data until mortality, while the remaining 1,093
(55.2%) were right-censored. Our analysis focuses on 21 publicly accessible
clinical features, including critical indicators such as tumor size and lymph
node involvement, with detailed descriptions available in Bilal et al. (2013)
[29]. Missing values underwent imputation to ensure data integrity, with real-
valued features substituted by their mean and categorical features by their
mode. We one-hot code the categorical variables, facilitating their
integration into subsequent analytical methodologies.

  * •

FLCHAIN 222Available at https://pypi.org/project/SurvSet/ [30, 31] : The
FLCHAIN3 [[30], [31]] dataset, introduced by Dispenzieri et al. [30],
constitutes a publicly accessible resource designed to investigate the
intricate association between serum free light chain levels and mortality.
This dataset encompasses a diverse array of covariates, including age, gender,
serum creatinine levels, and the presence of monoclonal gammopathy, all of
which are pivotal factors in comprehending the dynamics of this relationship.
This dataset is a critical resource for elucidating the multifaceted interplay
between serum free light chains and pertinent covariates in the context of
mortality studies.

  * •

RECID 333Available at https://github.com/georgehc/npsurvival : The Recidivism4
[32] dataset is a publicly available dataset, widely utilised in criminology
and social sciences, aimed at exploring the intricate dynamics of recidivism –
defined as the re-offending or re-incarceration of individuals previously
involved with the criminal justice system. This dataset offers a comprehensive
array of covariates essential for the analysis of recidivism, including
demographic information such as age, gender, socio-economic factors and
criminal history variables like prior convictions and offence type. This
dataset serves as a fundamental resource for gaining insights into the
multifaceted determinants of recidivism and contributes to informed policy and
intervention strategies in criminal justice studies.

Dataset | Total Observations | # Features∗ | % Censored | Five Number Summary (time)
---|---|---|---|---
metabric | 1981 | 79 | 55.2 | 3,1498,2632,4361,9218
flchain | 7874 | 10 | 72.5 | 0,2852,4302,4773,5215
recid | 1445 | 14 | 61.8 | 1,27,71,76,81

∗ The total number of features after preprocessing

###  4.2 Experimental Setup

We conduct the following experiments to evaluate the performance of the
proposed method. The experiment uses he five fold cross validation are
reported for each dataset, 3 and 2. We tune all parameters for each dataset
using hyperopt [33]. The parameters are given in table 1. Note that, we take
objective function in hyperopt as Integrate Brier Score and negative
Concordance Index.

Table 1: Parameters tuned with 1000 trials using hyperopt Parameter | Tuning Range
---|---
$\epsilon$ | 1e-300 - 0.9
$\alpha$ | 1/5,2/5,3/5,4/5,1
$l/n$ | 0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9

###  4.3 Performance Metrics

We evaluate the performance of the proposed method using the following two
metrics:

#### Concordance

Concordance [34] index is popular to evaluate the model performance in
survival models/functional regression setup. We define concordance index with
right censored data. The expression is as follows :

$C^{td}(t)=Pr(\hat{S}_{i}(t_{i})>\hat{S}_{j}(t_{i})\mid d_{i}=1,t\leq
t_{j},t_{i}<t_{j})$

where $S_{i}(t_{i})$ is the predicted survival probability of the $i^{th}$
individual at time $t_{i}$ and $S_{j}(t_{i})$ is the predicted survival
probability of the $j^{th}$ individual at time $t_{i}$. The time-dependent
concordance index given above measure the level of agreement between the
predicted survival probabilities of two individuals and actual event time of
the individuals. It is defined between 0.5 and 1, where 0.5 random prediction
and 1 is perfect prediction, represents that every prediction is in agreement
with the actual event time.

#### Integrated Brier Score

We use another benchmark as Integrated Brier Score. Since the data consists of
right censored observation we use the following modified version of the
Integrated Brier Score, we call it as $IBS^{c}$. $IBS^{c}$ computed between
time interval $[t_{1},t_{\max}]$ as

$IBS^{c}=\frac{1}{t_{\max}-t_{1}}\int_{t_{1}}^{t_{\max}}BS^{c}(t)dt$

. We estimate the integration numerically via trapezoidal rule. Expression of
Brier Score for the censored observation ($BS^{c}$) is :

$BS^{c}(t)=\frac{1}{N}\sum_{i=1}^{N}\left[I(y_{i}\leq
t,\delta_{i}=1)\frac{(0-\hat{S}(t\mid
x^{i}))^{2}}{\hat{G}(y_{i})}+I(y_{i}>t)\frac{(1-\hat{S}(t\mid
x^{i}))^{2}}{\hat{G}(t)}\right]$

where, $\hat{S}(t\mid x^{i})$ is the predicted conditional function and
$\frac{1}{\hat{G}(t)}$ is the inverse probability of the censoring weight. We
assume $C>t$, where $C$ is the censoring point.

### Base Models

The algorithm need base models to calculate the estimation on dataset ${\cal
D}_{l}$. We have used the following base models:

  1. 1.

Survival Tree : Survival Tree is a tree based method for survival analysis. We
split the data into different groups based on the log-rank test statistic
during the construction of survival tree. The log-rank statistic tells us how
different the survival curves of two groups are. Further, we grow the tree
recursively, till it meets a certain stopping criteria.

  2. 2.

Random Survival Forest : Random Survival Forest [3] is an extension Survival
Tree, The output of Random Survival Forest is the average of the output of
Survival Tree, obtained from different bootstrap samples of the data.

  3. 3.

Cox Model (Lasso and Ridge) : Cox Model is a semi-parametric model for
survival analysis, which is based on the proportional hazard assumption. The
Cox Model is given by

$h(t\mid x)=h_{0}(t)\exp(\beta^{T}x)$ (12)

where $h(t\mid x)$ is the hazard function of an individual with covariate $x$
at time $t$, $h_{0}(t)$ is the baseline hazard function and $\beta$ is the
parameter vector. The Cox Model is usually estimated using partial likelihood
method. Regularisation brings two variations in the Cox Model, known as Lasso
and Ridge Cox Model. We use Breslow estimator [35] for calculation of the
conditional survival function from Cox model.

  4. 4.

K-Nearest Neighbour Survival : K-Nearest Neighbour Survival ( [15]) is a non-
parametric method to estimate conditional survival function. In K-Nearest
Neighbour Regression, the prediction of the query point is the average of the
K-Nearest Neighbours of the query point. Suppose, we have the following right
censored survival data,
$\left((X_{1},Y_{1},\delta_{1}),\cdots,(X_{n},Y_{n},\delta_{n})\right)\in\mathcal{X}\times\mathbb{R}^{+}\times\\{0,1\\}$
Then for specified Kernel function
$K:\mathbb{R}^{+}\rightarrow\mathbb{R}^{+}$, and bandwidth $h>0$, we can
measure the proximity between observation $X_{i}$ and $X_{j}$ by
$K\left(\frac{\rho(X_{i},X_{j})}{h}\right)$, where $\rho$ is the distance
function. We can create unique event time, event count and survivor count as
follows

$\mathcal{Y}_{K}(x;h)=\biggl{\\{}Y_{j}:j\in[n]\cap\delta_{j}K\left(\frac{\rho(x,X_{j})}{h}\right)>0\biggr{\\}}$

$\mathcal{D}_{K}(t\mid
x;h)=\sum_{j=1}^{n}K\left(\frac{\rho(x,X_{j})}{h}\right)\delta_{j}\mathbb{I}(Y_{j}=t)$

$\mathcal{R}_{K}(t\mid
x;h)=\sum_{j=1}^{n}K\left(\frac{\rho(x,X_{j})}{h}\right)\mathbb{I}(Y_{j}\geq
t)$

Then the survival function is given by

$\hat{S}_{K}(t\mid
x;h)=\prod_{t^{{}^{\prime}}\in\mathcal{Y}_{K}(x;h)}\biggl{(}1-\frac{\mathcal{D}_{K}(t^{{}^{\prime}}\mid
x;h)}{\mathcal{R}_{K}(t^{{}^{\prime}}\mid
x;h)}\biggr{)}^{\mathbb{I}(t^{{}^{\prime}}\leq t)}$ (13)

The K-Nearest Neighbour Survival can be derived from the above equation by
replacing the kernel function $K$ by the $\mathcal{I}$, where mathematical
expression of $\mathcal{I}(x;k,X_{i})$444Taking $\mathcal{I}(\cdot)=1\forall
i$ will lead to a population Kaplan-Meier Estimate is as follows :

$\mathcal{I}(x;k,X_{i})=\mathbb{I}\biggl{(}\sum_{j=1,j\neq
i}^{n}\mathbb{I}(\rho(x,X_{i})\leq\rho(x,X_{j}))\geq k\biggr{)}$ (14)

KNN survival implementation is available by George Chen at
https://github.com/georgehc/npsurvival.

### Numerical Results on the Real Datasets

Table 2 provides the information of the concordance of all the five models on
the three datasets, whereas Table 3 describes the integrated brier score of
all the five models. We use 5 fold cross validation to obtain the result We
observe that our proposed construction of Combined regression strategy
provides maximum concordance index and minimum integrated brier score for all
datasets.

Table 2: Concordance (Higher is Better) Dataset | METABRIC | FLCHAIN | RECID
---|---|---|---
Survival Tree | 0.553 | 0.780 | 0.544
Random Survival Forest | 0.669 | 0.899 | 0.629
Cox Model (Lasso) | 0.626 | 0.886 | 0.599
Cox Model (Ridge) | 0.589 | 0.901 | 0.551
KNN Survival Model | 0.561 | 0.829 | 0.540
proposed | 0.677 | 0.911 | 0.630
Table 3: Integrated Brier Scoe (Lower is better) Dataset | METABRIC | FLCHAIN | RECID
---|---|---|---
Survival Tree | 0.278 | 0.065 | 0.262
Random Survival Forest | 0.167 | 0.046 | 0.175
Cox Model (Lasso) | 0.160 | 0.079 | 0.183
Cox Model (Ridge) | 0.162 | 0.143 | .189
KNN Survival Model | 0.203 | 0.054 | 0.211
proposed | 0.148 | 0.043 | 0.175

We also perform a D-calibration test [36] at 0.05 level for each of the 5-fold
cross validation in all three datasets. The Table-4 reports the number of
times model passes D-calibration test. We observe that proposed model passes
the test maximum number of times than almost all weak learners. However, it
passes equal number of tests as the strongest weak learners Random Survival
Forest. The experiments validate not only the efficacy of the model over other
models, but also confirms that the model’s predicted probability calculations
are meaningful.

Table 4: Out of 5-fold number of times model passes D-calibration test at 0.05 level Models | METABRIC | FLCHAIN | RECID | total
---|---|---|---|---
Survival Tree | 0 | 2 | 0 | 2
Random Survival Forest | 5 | 4 | 5 | 14
Cox (lasso) | 4 | 3 | 5 | 12
Cox (ridge) | 5 | 0 | 5 | 10
KNN Survival | 0 | 2 | 0 | 2
proposed | 5 | 4 | 5 | 14

##  5 Covariate Relevance

The relevance of the covariates is also an important aspect of survival
analysis. COBRA Survival has an inbuilt mechanism to measure the relevance of
the covariates. We use a logistic regression to find the relevance of the
covariates for a query point $x$. A detailed overview of logistic regression
is available at [37], on the indicator function 15, where
$i=k+1,k+2,\cdots,n^{{}^{\prime}}$.

The logistic regression is given by

$\displaystyle\log\biggl{(}\frac{Pr(\Gamma_{i}(x;X_{j},\mathcal{M},\epsilon,\alpha)=1)}{1-Pr(\Gamma_{i}(x;X_{j},\mathcal{M},\epsilon,\alpha)=1)}\biggr{)}$
$\displaystyle=$
$\displaystyle\beta_{0}+\beta_{1}X_{1,i}+\beta_{2}X_{2,i}+\dots+\beta_{p}X_{p,i}=\beta^{T}\mathbf{X}_{i}$
$\displaystyle\forall i=k+1,k+2,\cdots,n^{{}^{\prime}}$

The coefficients $\beta_{1},\beta_{2}\dots\beta_{p}$ are the relevance of the
covariates $X_{1},X_{2},\cdots,X_{p}$ respectively. We recall the definition
of $\Gamma(\cdot)$ :

$\Gamma(x;X_{j},\mathcal{M},\epsilon,\alpha)=\mathbb{I}\biggl{(}\sum_{m=1}^{\mid\mathcal{M}\mid}\mathbb{I}(d(S_{k,m}(t\mid
x),S_{k,m}(t\mid X_{j}))\leq\epsilon)\geq\mid\mathcal{M}\mid\alpha\biggr{)}$
(15)

The k subscript notifies that the prediction of the base models takes training
on the dataset ${\cal D}_{k}$. Here, $\mathbb{I}$ is the indicator function,
which is 1 if the condition is true and 0 if the condition is false.
$\epsilon$ is the threshold distance, $\alpha$ is thetakes fraction of base
models in the consensus of $\epsilon$ proximity of the query point $x$.
basically $\mathcal{I}(.)$ is the indicator function, which is 1 if the
$i^{th}$ individual is in the $(\epsilon,\alpha)-$proximity of the query point
$x$, otherwise 0.

### Simulation Study

We use a synthetic population to demonstrate COBRA Survival’s ability to
discriminate between different factors and determine their relative
importance. Our strategy incorporates a non-linear data-generating procedure
supported by the following link rule555Similar Link rule is taken by Bellot
and van der Schaar

$\Lambda(\mathbf{X}_{i})=2+\log(13X_{0,i}+5X_{1,i}+7X_{2,i})+X_{3,i}$

We have generated 2000 sample, where 40 percent of the sample are censored.
The censoring time is generated by taking uniform distribution between 0 and
time of event. The time of event is generated by following weibuill
distribution

$T_{i}\sim\mathcal{W}(2,\Lambda(\mathbf{X}_{i}))$

There are 9 covariates generated, where $X_{0}$ , $X_{1}$ , $X_{2}$ , $X_{3}$
have non-linear effect on true event time, whereas $X_{4}$ , $X_{5}$ , $X_{6}$
, $X_{7}$ , $X_{8}$ have no effect on true event time. This is done to show
that COBRA Survival can identify the covariates which have effect on true
event time, and the covariates which have no effect on true event time. The
covariates are generated from uniform distribution. The estimates of the
covariate relevance is given in Figure 1.

Figure 1: Relevance of the covariates

##  6 Conclusion

A large number of variations of cobra structure can not able to beat the
performance of Random Survival Forest even with weak learners taken as Random
Survival Forest. The performance of the Cobra depends on the appropriate
choice of its parameters. We observe that the proposed method in the paper
works well for all three datasets. Also, a lesser number of observations in
proximity points affect the performance. Our proposed model in the variable
selection method works very well in the simulated framework. However, this
requires special attention in a separate research. We may pursue further
research on weighted version, problems on smaller observations, and more on
variable importance in this context. The work is in progress.

## References

  *   * Cox [1975] Cox, D.R.: Partial likelihood. Biometrika 62(2), 269–276 (1975)
  * Cox [1972] Cox, D.R.: Regression models and life-tables. Journal of the Royal Statistical Society: Series B (Methodological) 34(2), 187–202 (1972)
  * Ishwaran et al. [2008] Ishwaran, H., Kogalur, U.B., Blackstone, E.H., Lauer, M.S.: Random survival forests (2008)
  * Faraggi and Simon [1995] Faraggi, D., Simon, R.: A neural network model for survival data. Statistics in medicine 14(1), 73–82 (1995)
  * Katzman et al. [2018] Katzman, J.L., Shaham, U., Cloninger, A., Bates, J., Jiang, T., Kluger, Y.: Deepsurv: personalized treatment recommender system using a cox proportional hazards deep neural network. BMC medical research methodology 18(1), 1–12 (2018)
  * Bellot and van der Schaar [2018] Bellot, A., Schaar, M.: Boosted trees for risk prognosis. In: Machine Learning for Healthcare Conference, pp. 2–16 (2018). PMLR
  * Goswami and Dey [2022a] Goswami, R., Dey, A.K.: Concordance based survival cobra with regression type weak learners. arXiv preprint arXiv:2209.11919 (2022)
  * Goswami and Dey [2022b] Goswami, R., Dey, A.K.: Integrated brier score based survival cobra–a regression based approach. arXiv preprint arXiv:2210.12006 (2022)
  * Kaplan and Meier [1958] Kaplan, E.L., Meier, P.: Nonparametric estimation from incomplete observations. Journal of the American statistical association 53(282), 457–481 (1958)
  * Aalen [1978] Aalen, O.: Nonparametric inference for a family of counting processes. The Annals of Statistics, 701–726 (1978)
  * Nelson [1972] Nelson, W.: Theory and applications of hazard plotting for censored failure data. Technometrics 14(4), 945–966 (1972)
  * Simon et al. [2011] Simon, N., Friedman, J., Hastie, T., Tibshirani, R.: Regularization paths for cox’s proportional hazards model via coordinate descent. Journal of statistical software 39(5), 1 (2011)
  * Tibshirani [1997] Tibshirani, R.: The lasso method for variable selection in the cox model. Statistics in medicine 16(4), 385–395 (1997)
  * Breiman [2001] Breiman, L.: Random forests. Machine learning 45, 5–32 (2001)
  * Chen [2019] Chen, G.: Nearest neighbor and kernel survival analysis: Nonasymptotic error bounds and strong consistency rates. In: International Conference on Machine Learning, pp. 1001–1010 (2019). PMLR
  * Beran [1981] Beran, R.: Nonparametric regression with randomly censored survival data (1981)
  * Fix and Hodges [1952] Fix, E., Hodges, J.L.: Discriminatory analysis: Nonparametric discrimination: Small sample performance (1952)
  * Kalbfleisch and Prentice [2011] Kalbfleisch, J.D., Prentice, R.L.: The Statistical Analysis of Failure Time Data. John Wiley & Sons, ??? (2011)
  * Ishwaran [2007] Ishwaran, H.: Variable importance in binary regression trees and forests (2007)
  * Grömping [2009] Grömping, U.: Variable importance assessment in regression: linear regression versus random forest. The American Statistician 63(4), 308–319 (2009)
  * Genuer et al. [2010] Genuer, R., Poggi, J.-M., Tuleau-Malot, C.: Variable selection using random forests. Pattern recognition letters 31(14), 2225–2236 (2010)
  * Ribeiro et al. [2016] Ribeiro, M.T., Singh, S., Guestrin, C.: ”why should i trust you?” explaining the predictions of any classifier. In: Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 1135–1144 (2016)
  * Strumbelj and Kononenko [2010] Strumbelj, E., Kononenko, I.: An efficient explanation of individual classifications using game theory. The Journal of Machine Learning Research 11, 1–18 (2010)
  * Kovalev et al. [2020] Kovalev, M.S., Utkin, L.V., Kasimov, E.M.: Survlime: A method for explaining machine learning survival models. Knowledge-Based Systems 203, 106164 (2020)
  * Kovalev and Utkin [2020] Kovalev, M.S., Utkin, L.V.: A robust algorithm for explaining unreliable machine learning survival models using the kolmogorov–smirnov bounds. Neural Networks 132, 1–18 (2020)
  * Utkin et al. [2020] Utkin, L.V., Kovalev, M.S., Kasimov, E.M.: Survlime-inf: A simplified modification of survlime for explanation of machine learning survival models. arXiv preprint arXiv:2005.02387 (2020)
  * Alabdallah et al. [2022] Alabdallah, A., Pashami, S., Rögnvaldsson, T., Ohlsson, M.: Survshap: A proxy-based algorithm for explaining survival models with shap. In: 2022 IEEE 9th International Conference on Data Science and Advanced Analytics (DSAA), pp. 1–10 (2022). IEEE
  * Krzyziński et al. [2023] Krzyziński, M., Spytek, M., Baniecki, H., Biecek, P.: Survshap (t): Time-dependent explanations of machine learning survival models. Knowledge-Based Systems 262, 110–234 (2023)
  * Bilal et al. [2013] Bilal, E., Dutkowski, J., Guinney, J., Jang, I.S., Logsdon, B.A., Pandey, G., Sauerwine, B.A., Shimoni, Y., Moen Vollan, H.K., Mecham, B.H., et al.: Improving breast cancer survival analysis through competition-based multidimensional modeling. PLoS computational biology 9(5), 1003047 (2013)
  * Dispenzieri et al. [2012] Dispenzieri, A., Katzmann, J.A., Kyle, R.A., Larson, D.R., Therneau, T.M., Colby, C.L., Clark, R.J., Mead, G.P., Kumar, S., Melton III, L.J., et al.: Use of nonclonal serum immunoglobulin free light chains to predict overall survival in the general population. Mayo Clinic Proceedings 87(6), 517–523 (2012)
  * Kyle et al. [2006] Kyle, R.A., Therneau, T.M., Rajkumar, S.V., Larson, D.R., Plevak, M.F., Offord, J.R., Dispenzieri, A., Katzmann, J.A., Melton III, L.J.: Prevalence of monoclonal gammopathy of undetermined significance. New England Journal of Medicine 354(13), 1362–1369 (2006)
  * Wooldridge et al. [2000] Wooldridge, J.M., et al.: Recid. Boston College Department of Economics, ??? (2000)
  * Bergstra et al. [2013] Bergstra, J., Yamins, D., Cox, D.D., et al.: Hyperopt: A python library for optimizing the hyperparameters of machine learning algorithms. In: Proceedings of the 12th Python in Science Conference, vol. 13, p. 20 (2013). Citeseer
  * Gerds et al. [2013] Gerds, T.A., Kattan, M.W., Schumacher, M., Yu, C.: Estimating a time-dependent concordance index for survival prediction models with covariate dependent censoring. Statistics in medicine 32(13), 2173–2184 (2013)
  * Breslow [1972] Breslow, N.E.: Contribution to discussion of paper by dr cox. Journal of the Royal Statistical Society, Series B 34, 216–217 (1972)
  * Haider et al. [2020] Haider, H., Hoehn, B., Davis, S., Greiner, R.: Effective ways to build and evaluate individual survival distributions. The Journal of Machine Learning Research 21(1), 3289–3351 (2020)
  * Cramer [2002] Cramer, J.S.: The origins of logistic regression. Tinbergen Institute Working Paper (2002)


# OMPGPT: A Generative Pre-trained Transformer Model for OpenMP

Le Chen<EMAIL_ADDRESS>,  Arijit Bhattacharjee<EMAIL_ADDRESS>Iowa
State UniversityAmesIowaUSA50010 ,  Nesreen Ahmed<EMAIL_ADDRESS>,
Niranjan Hasabnis<EMAIL_ADDRESS>Intel LabsUSA ,  Gal Oren
<EMAIL_ADDRESS>Technion - Israel Institute of TechnologyIsrael ,  Vy
Vo<EMAIL_ADDRESS>Intel LabsUSA and  Ali Jannesari<EMAIL_ADDRESS>Iowa State UniversityAmesIowaUSA50010

###### Abstract.

Large language models (LLMs), as epitomized by models like ChatGPT, have
revolutionized the field of natural language processing (NLP). Along with this
trend, code-based large language models such as StarCoder, WizardCoder, and
CodeLlama have emerged, trained extensively on vast repositories of code data.
Yet, inherent in their design, these models primarily focus on generative
tasks like code generation, code completion, and comment generation, and
general support for multiple programming languages. While the generic
abilities of code LLMs are useful for many programmers, the area of high-
performance computing (HPC) has a narrower set of requirements that make a
smaller and more domain-specific LM a smarter choice. This paper introduces
OMPGPT, a novel model meticulously designed to harness the inherent strengths
of language models for OpenMP pragma generation. Furthermore, we adopt and
adapt prompt engineering techniques from the NLP domain to create chain-of-
OMP, an innovative strategy designed to enhance OMPGPT’s effectiveness. Our
extensive evaluations demonstrate that OMPGPT outperforms existing large
language models specialized in OpenMP tasks and maintains a notably smaller
size, aligning it more closely with the typical hardware constraints of HPC
environments. We consider our contribution as a pivotal bridge, connecting the
advantage of language models with the specific demands of HPC tasks. The
success of OMPGPT lays a solid foundation, suggesting its potential
applicability and adaptability to a wider range of HPC tasks, thereby opening
new avenues in the field of computational efficiency and effectiveness.

HPC, LLM, Parallelization, OpenMP

††copyright: none††ccs: Do Not Use This Code Generate the Correct Terms for
Your Paper††ccs: Do Not Use This Code Generate the Correct Terms for Your
Paper††ccs: Do Not Use This Code Generate the Correct Terms for Your
Paper††ccs: Do Not Use This Code Generate the Correct Terms for Your Paper

##  1\. Introduction

In the realm of artificial intelligence and machine learning, large language
models (LLMs) have emerged as a groundbreaking technique. Originating from the
transformer model, which was introduced by Vaswani et al. (2017) in 2017, LLMs
represent a significant leap in the ability to process and generate human-like
text. These models are trained with large corpora of text, allowing them to
absorb a wide array of linguistic patterns, styles, and contexts. Their
ability to understand, interpret, and generate human language has opened up
new frontiers in NLP. The evolution of LLMs, from the first version to the
more sophisticated and powerful subsequent iterations, has continually pushed
the boundaries of what’s possible in machine learning and AI.

With the success of LLMs in the NLP field, the programming language (PL)
community has witnessed the advent of code-based Large Language Models. These
specialized LLMs, such as StarCoder (Li et al., 2023), WizardCoder (Luo et
al., 2023), and CodeLlama (Rozière et al., 2023), represent a tailored
evolution of the technology, specifically adapted for understanding and
generating programming code. These models can be categorized into two groups:
trained with programming language data only and trained with both natural
language and programming language data. Unlike their predecessors that
primarily focused on natural language, these models are trained extensively on
vast repositories of source code from a variety of programming languages.
Their design and training are centered around the unique structures, syntax,
and semantics inherent in programming languages, differentiating them from
traditional LLMs.

Figure 1. LLM prompt example for HPC questions. Typically, the prompt contains
both text in natural language to describe the question and the code in
programming language. Figure 2. Answers returned by different LLMs given the
prompt in Figure 1.

HPC tasks are critical in solving complex problems ranging from weather
forecasting to quantum simulations. Various studies have trained machine
learning models for specific HPC tasks such as auto parallelization (Chen et
al., 2023b), performance optimization (Dutta et al., 2023a), and power-
constrained autotuning (Dutta et al., 2023b). However, the inherent
characteristics of HPC tasks pose distinct challenges to current code-based
LLMs.

  1. (1)

Training data diversity: Natural Language (NL) vs. Programming (PL). Advanced
LLMs like GPT3 or CodeLlama have been trained with both NL and PL data. When
presented with prompts like the one in Figure 1, these models demonstrate the
ability to interpret questions framed in NL and respond with appropriate code
in PL. However, these models require more intricate training data preparation
and have a larger model size. In contrast, models like Starcoder are trained
with code only and take the comments as PL instead of NL. They are designed
for code generation tasks and cannot understand questions in PL.

  2. (2)

Performance consistency. The use of NL as input would naturally lead to
different users asking the same question (HPC-related or otherwise)
differently. Most likely, the output of LLMs will be different, leading to the
additional complexity for post-processing tools that can operate on the
output.

  3. (3)

Output processing. Pre-trained LLMs usually return answers in NL, adding the
effort to extract the information of interest from the answer.

A motivating example. Figure 2 presents the responses generated by GPT-3.5
(20B), CodeLlama (13B), Starcoder (15.5B), and WizardCoder (15B) to the OpenMP
query illustrated in Figure 1. Starcoder, being pre-trained exclusively with
code, struggles to comprehend the query. In contrast, GPT-3.5, CodeLlama, and
WizardCoder, all pre-trained with both Natural Language (NL) and Programming
Language (PL) data, successfully grasp the prompt’s objective. However, the
quality of their outputs varies. CodeLlama’s response, while somewhat correct,
includes incorrect pragma and unnecessary text, necessitating post-processing
to extract the relevant code snippet for HPC users. WizardCoder, reflective of
its training methodology, offers multiple candidate outputs, adding to the
complexity of identifying the most suitable answer. GPT-3.5 emerges as the
most accurate in generating the correct code snippet. Nonetheless, its use is
limited by associated costs and requires post-processing to refine the output.

Figure 3.  (black): Traditional large language models require training on both
natural language and programming datasets to generate OpenMP pragmas for given
code segments effectively. This process introduces added complexity in data
handling, model training, and overall model size. Additionally, users often
need to meticulously formulate prompts and interpret the outputs to extract
accurate OpenMP pragmas.  (blue): OMPGPT is specifically trained with a
variety of HPC code, including almost a half of OpenMP code. We elaborately
align the OpenMP pragmas to follow their scope in the training process to meet
the instinct of GPT models. Besides overcoming traditional LLMs’
disadvantages, OMPGPT can be boosted with the proposed Chain-of-OMP prompt
engineering techniques.

Motivated by the limitations of the current LLMs for HPC tasks, this work
introduces OMPGPT model, along with the innovative chain-of-OMP prompt
engineering method, specifically targeting the HPC task of automatic
parallelization via OpenMP pragma generation. As illustrated in Figure 3,
OMPGPT is trained on an extensive HPC dataset, exclusively comprising C and
C++ code, which are pivotal in HPC applications. By harnessing the instinct of
LLMs, we covert OpenMP pragma generation problems into code generation
problems, eliminating the need for training NL data to understand the task
objective. The choice of training data and the specific design of OMPGPT helps
us in reducing the size of OMPGPT considerably, thereby enabling efficient
training and fine-tuning. On top of the pre-trained OMPGPT, we introduce the
chain-of-OMP prompt engineering method and the fine-tuning approach to enhance
the performance of the model. Specifically, our chain-of-OMP prompt
engineering method leverages ideas from the chain-of-thought prompt
engineering method and tailors it to OpenMP pragma generation. In particular,
we observed that chain-of-OMP improves OMPGPT’s performance on the OpenMP
pragma generation task. Intuitively, the chain-of-OMP method mimics OpenMP
experts by refining the inquiry with incrementally added preconditions for
OpenMP pragma generation. This strategy fits perfectly with the structure of
OpenMP pragmas and can be applied to other LLMs. Finally, our design of OMPGPT
and chain-of-OMP bridges the instinct of LLMs and OpenMP syntax, representing
a significant stride in aligning language model capabilities with the demands
of HPC tasks.

Contributions. The contributions of this work are as follows:

  * •

We develop OMPGPT, a 0.76B language model (smallest among our baseline models)
tailored for OpenMP pragma generation with competitive performance to large
LLMs.

  * •

To enhance OMPGPT, we develop a novel OpenMP clause-based prompting technique
that we call chain-of-OMP. Conceptually, chain-of-OMP represents a new design
point for LLMs for HPC in that it provides hints to the underlying model to
direct it toward the right answer. To the best of our knowledge, this is the
first application of chain-of-thought prompting to LLMs for HPC.

  * •

We perform a comprehensive evaluation of OMPGPT against the latest baseline
model (MonoCoder) and the popular GPT-3.5 model, showcasing its superior
performance over in the OpenMP pragma generation task. We then also analyze
the effectiveness of the chain-of-OMP method in improving the performance of
OMPGPT on the OpenMP pragma generation task.

##  2\. Background

This section briefly describes Large Language Models and their application in
High-Performance Computing tasks.

###  2.1. Generative Pre-trained Transformers

The advent of Generative Pre-trained Transformers (GPT) has marked a
revolutionary shift in the field of Natural Language Processing (NLP) and
beyond. Each GPT version has brought significant improvements in its
capabilities. These improvements are mainly attributed to their increased
size, which requires increasing amounts of training data. Inspired by their
success in the NLP area, studies have been attempting to apply GPT-like models
to programming languages. The most recent versions of GPT (3.5 and above) can
now generate source code in addition to natural language.
The left-to-right characteristic(Nichols et al., 2023) of GPT models refers to
the sequential nature in which these models generate text. GPT models,
including GPT-3.5, are autoregressive, meaning they generate outputs from left
to right one token at a time based on the input context. This left-to-right
generation process contributes to the models’ ability to comprehend and
generate diverse and contextually appropriate responses in natural language
tasks.

###  2.2. Code LLMs

Large Language Models (LLMs) for code, or Code LLMs, represent a category of
language models specifically designed to understand and generate programming
code. These models, built upon architectures like GPT (Generative Pre-trained
Transformer), are trained on large-scale datasets containing code snippets in
various programming languages. These models have the potential to transform
the way software is developed by providing code completion suggestions,
improving code quality, and overall streamlining the development process.
However, many challenges remain, including the need for models to understand
context-specific programming nuances and generate code that is not only
syntactically correct but also functionally meaningful in a given application
domain. Ongoing research in the field of Code LLMs aims to address these
challenges and further advance the capabilities of these models for a wide
range of coding tasks. Below, we describe some of the most well-known open-
source Code LLMs.

##### StarCoder:

StarCoder (Li et al., 2023) is a 15B parameter model trained for code
generation or completion. The training dataset, the Stack (Kocetkov et al.,
2022), has 1 trillion tokens sourced from a large collection of permissively
licensed GitHub repositories.

##### CodeLlama:

CodeLlama (Rozière et al., 2023) is an LLM capable of generating code from
natural language prompts. It is a code-specialized version of Llama2 (Touvron
et al., 2023) developed by Meta. It was derived from Llama2 by training it
with code-specific datasets. Code Llama was released in sizes of 7B, 13B, and
34B parameters. Each was trained with 500B tokens of code data.

##### WizardCoder:

Most code LLMs like StarCoder have been trained on raw code data without
instruction fine-tuning. WizardCoder(Luo et al., 2023), with its complex
instruction fine-tuning by adapting the Evol-Instruct methods for coding
tasks, has shown to have improved performance for code generation.

###  2.3. LLMs on code-related tasks for HPC

LLM models have shown promise in code-related High Performance Computing (HPC)
tasks, where computational efficiency and parallel processing are crucial. In
HPC scenarios, LLM models can be utilized for tasks such as optimizing code
for parallelization, improving algorithm performance, or suggesting
enhancements for better resource utilization. By understanding the intricacies
of HPC code, including parallelization strategies and optimization techniques,
LLM models can generate insightful suggestions to enhance code efficiency.
Their ability to capture complex patterns and dependencies in code can aid HPC
developers in creating algorithms that leverage the full potential of high-
performance computing architectures. While LLM models are not specifically
designed for HPC, their adaptability and generalization capabilities (Valero-
Lara et al., 2023) make them valuable tools for addressing code-related
challenges in the realm of high-performance computing.
High-performance computing (HPC) languages are predominantly based on C, C++,
and Fortran due to their well-established features that cater to the demands
of computationally intensive tasks. C and C++ are renowned for their low-level
control over system resources, efficient memory management, and ability to
express complex algorithms in a concise manner. These languages allow
developers to optimize code for parallelism and take advantage of hardware
architectures, making them ideal for applications requiring high computational
efficiency. A unique feature of HPC tasks is that the models need not be
trained on a multitude of programming languages but be centered around the
primary three. This means that models can be smaller and more efficient. HPC
tasks are also very specific which results in the input/prompt being direct
and not needing too much flexibility.

The application of LLMs in HPC tasks represents an emerging yet rapidly
evolving area of research. LM4HPC (Chen et al., 2023a) stands as the first
attempt to adapt LLMs to the HPC domain. Building upon this, a subsequent
study by Ding et al. ([n. d.]) introduced a Llama-based Question and Answer
model specifically tailored for HPC tasks. Despite these advancements, these
early efforts predominantly rely on existing LLMs, indicating a nascent stage
in the development of HPC-focused language models.

####  2.3.1. OpenMP Pragma

OpenMP (Open Multi-Processing) (Dagum and Menon, 1998) is an API that supports
multi-platform shared-memory multiprocessing programming in C, C++, and
Fortran. It enables developers to write parallel code that can take advantage
of multiple processors or cores on a single machine. The pragma directive in C
and C++ is a compiler directive that provides a hint or instruction to the
compiler about how to process certain parts of the code. In the context of
OpenMP, pragma omp is used to include OpenMP directives. OpenMP directives
actually follow the left-to-right instinct of a generative model. The first
part being the pragma omp followed by a directive, which is then followed, as
needed by the loop, a control variable.

###  2.4. Chain-of-Thought Prompting

Chain-of-Thought$($CoT$)$ (Wei et al., 2023) prompting is a series of
intermediate reasoning steps that greatly help a LLM’s ability to perform
complex reasoning. Rather than using a singular prompt to generate a response
from the LLM, this prompting involves guiding the LLM through multiple phases
of small hints. The hints help the model learn the deeper meaning of what’s
being asked in the query. An avenue that has been greatly helped with kind of
prompting is Arithmetic and Common Sense reasoning.

##  3\. Approach

This section outlines our approach to developing OMPGPT and Chain-of-OMP,
converging several key stages: dataset processing, OMPGPT training, design and
implementation of Chain-of-OMP, and the fine-tuning process for OMPGPT.

###  3.1. OMPGPT Design

As highlighted in Section 1, existing Large Language Models (LLMs) encounter
significant limitations when applied to HPC tasks, particularly in the realm
of OpenMP pragma generation. In designing OMPGPT, we focused on several key
metrics to ensure its suitability and efficacy for such specialized tasks:

  1. (1)

Training Data Relevance. The quality of the training dataset is essential to a
language model. OMPGPT should be trained on HPC code in the most common
programming languages in the HPC field.

  2. (2)

Model Compatibility. The model size should align with the hardware
configurations prevalent in most HPC clusters.

  3. (3)

Flexibility and Adaptability. OMPGPT is designed to be flexible and adaptable,
capable of handling a variety of OpenMP pragma generation tasks without the
need to craft prompts

  4. (4)

Performance Efficiency. OMPGPT is expected to outperform previous small
language models and be competitive compared to advanced LLMs.

  5. (5)

User Accessibility: Recognizing the varied nature of HPC ecosystems, OMPGPT is
designed for a diverse array of users, regardless of their OpenMP knowledge
level.

With these design principles in mind, we pre-train a publicly available
transformer architecture, GPT-Neo 2.7B, from scratch only using C and C++ HPC
code.

Figure 4. Distribution of OpenMP Pragmas in the HPCorpus OpenMP Subset. For
clarity, only the top 15 most frequently occurring pragmas are displayed.

###  3.2. OMPGPT Training

HPCorpus. Kadosh et al. (2023a) introduced an extensive HPC database called
HPCorpus, derived from GitHub repositories containing C, C++, and Fortran
code. Notably, 45% of HPCorpus repositories that use parallel programming APIs
use OpenMP as their primary programming language, making it an ideal
foundation for training models specifically targeted at OpenMP tasks. The
success of MonoCoder (Kadosh et al., 2023b), a language model trained using
HPCorpus, particularly in specific tasks of parallel code generation, further
underscores the potential of this database. In light of these findings, we
have utilized HPCorpus to create our training and fine-tuning datasets, aiming
to leverage its diverse and rich content for our model development. We utilize
the C and C++ code from HPCorpus for our training data. HPCorpus contains
$144,522$ repositories with C code and $150,481$ repositories with C++ code,
amounting to a total of $8,781,759$ C code functions and $68,233,984$ C++ code
functions. The total disk storage cost is $72.39$ GB. Figure 4 illustrates the
distribution of OpenMP pragmas in the HPCorpus OpenMP Subset. This subset is
later used for OMPGPT fine-tuning. For both the training and fine-tuning
phases, we randomly allocate 10% of the data as a test set. This ensures that
the model does not encounter this subset during the training or fine-tuning
stages, maintaining the integrity of our evaluation.

Data processing. To align with the design of OMPGPT, our data pre+processing
phase includes:

  1. (1)

NL trimming: Designed as a code-based model trained with code only, we remove
all the natural language text, such as comments.

  2. (2)

Moving OpenMP pragmas from the start of the loop to the end: Differing from
previous studies, our approach capitalizes on the instinct of GPT models to
overcome the constraints present in the current work. In our pre-processing of
the OpenMP code, we position the OpenMP pragmas after their respective scopes.
This pre-processing strategy aligns with the findings of previous work of
Nichols et al. (2023), which tested similar methodologies using simple fine-
tuning techniques. Our approach diverges from prior efforts. By training
OMPGPT from the ground up, we specifically enable it to generate OpenMP
pragmas, an ability not inherently present in conventional pre-trained models.
This fundamental difference in training methodology not only tailors OMPGPT
more closely to the unique demands of OpenMP pragma generation but also sets
it apart from existing models in its ability to intuitively understand and
process HPC-specific code structures.

  3. (3)

Filtering code by size: In line with practices established in previous
research like PolyCoder (Xu et al., 2022), we apply size-based filtering to
the code from HPCorpus. Specifically, we exclude code segments exceeding 100
tokens or larger than 1 MB in size.

Figure 5. Basic prompt for OMPGPT. Prompt (a): using code as the prompt for
code generation. Prompt (b): using code and the #pragma omp hint as the prompt
for OpenMP pragma generation. The highlighted part in Yellow is the content of
interest in the output answer.

####  3.2.1. Compute hardware

For all the experiments in this work, we used a singular node on the cluster,
which consists of a dual-socket AMD EPYC Rome 7402 with 24 cores per socket.
It is complimented with 512 GB of DDR4-3200 RAM. Each node also has 4 x NVIDIA
A100 40GB HBM2e Tensor Core GPUs connected via NVLink3 to each other. For
network, it has 4 × Mellanox HDR200 InfiniBand ConnectX 6 $($200 Gbit/s
each$)$.

Training OMPGPT: OMPGPT’s training leverages the GPT-Neo 2.7B (Black et al.,
2021) model architecture as a foundation. However, we have tailored the model
to better suit our needs by downsizing the number of layers to 8 while
maintaining 32 attention heads per layer. The model features a hidden
dimension of 2560. We employ the StarCoder tokenizer for tokenization, which
utilizes a vocabulary comprising 49,152 tokens. This configuration results in
a total parameter count of 0.76 billion for OMPGPT, making it more compact
than our baseline models.

Inference with OMPGPT. By replacing the OpenMP pragmas in the training set, we
convert the problem of OpenMP pragma generation into a code generation task.
Figure 5 illustrates the inference with OMPGPT for code generation and OpenMP
pragma generation. Originating from a GPT-based model, OMPGPT is inherently
proficient in the generation tasks, given an input code. Because OMPGPT is
trained on pre-processed OpenMP code containing OpenMP pragmas that are moved
towards the end of the loop. When provided with a prompt that concludes with
the prefix of an OpenMP directive (i.e., #pragma omp), OMPGPT demonstrates its
specialized ability to generate relevant OpenMP pragmas. This functionality
underscores the model’s aptitude for intuitively continuing code sequences,
specifically tailored to the syntax and structure of OpenMP directives,
thereby showcasing its practical utility in real-world programming scenarios.

Figure 6. Chain-of-OMP. Start phase: using the standard OpenMP pragma
generation prompts for OMPGPT. Directive/clause generation phase: the first
generated Directive/clause is augmented to the previous input as the input for
the next chain. This phase is updated till it number of chains reaches
parameter $N_{\text{chain}}$. The final Control variable generation phase
generates the last component to complete the chain-of-OMP and generate the
OpenMP pragma.

###  3.3. Chain-of-OMP

The concept of chain-of-thought (CoT), as introduced by Wei et al. (2023),
facilitates complex reasoning in models by guiding them through intermediate
steps. Inspired by this ”think step by step” methodology, we have developed
Chain-of-OMP, a novel prompt engineering technique specifically designed to
augment OMPGPT’s capabilities in generating OpenMP pragmas.

We consider the OpenMP pragma as three major components.

  * •

Directive Prefix (#pragma omp): Serving as the universal starting point for
all OpenMP directives, this prefix is critical in signaling to the compiler
the presence of an OpenMP directive. Within our model, this prefix is used in
prompts to orient OMPGPT towards generating OpenMP pragmas, steering it away
from generic code generation. We represent it as <prefix> below.

  * •

Directives and Clauses. This segment outlines the type of parallel or work-
sharing construct being utilized and details its specific behavior. We
represent them as <dc> below.

  * •

Control Structure. A control structure typically immediately follows the
directives. It is essential for the behavior of directives like #omp parallel
reduction. We represent it as <cs> below.

In OpenMP syntax, these three components are listed in a sequential manner
(i.e., <prefix><dc><cs>), perfectly fitting the instinct of LLMs. Leveraging
this, Chain-of-OMP operates as an automated sequence of prompts, as depicted
in Figure 6. The process begins with standard OpenMP pragma generation prompts
(i.e., <prefix>). The output from the first stage is then refined, retaining
only the generated directive/clause, which is subsequently passed to the next
OMPGPT client to generate the following component. In other words, we only
retain <prefix><dc> part of the generated output and ignore the rest. This
retained part is then fed as the prompt for the next stage. In general, the
components of an OpenMP pragma expand incrementally, with one element added at
each stage, until it reaches a user-defined limit, $N_{\text{chain}}$. The
input $I_{n}$ for a chain component chainn is defined by Equation 1.

(1) $\text{I}_{n}=concat(\text{I}_{n-1},\text{first generated directive/clause
in O}_{n-1})$

where $I_{n}$, $O_{n}$ stands for the input and output of chainn. $concat$
stands for the concatenate operation. OpenMP pragmas have a different number
of components. Consequently, we need a different number of chains to generate
the complete OpenMP pragma. We let the user set the maximum number of chains
limit by setting $N_{\text{chain}}$. The default value of $N_{\text{chain}}$
is 256 to avoid an infinity loop. The value aligns with the commonly used
maximum output length of LLMs.

The design of Chain-of-OMP has the following advantages:

  * •

Mimicking Expert User Inquiry: This approach closely replicates the querying
process employed by experienced OpenMP users. Instead of requesting a complete
pragma with a basic prompt (using just #pragma omp), skilled users often
provide more specific information, such as directives (e.g., #pragma omp for),
to refine their inquiry about the remaining parts of the pragma. Chain-of-OMP
embodies this nuanced approach, leading to more targeted and accurate pragma
generation.

  * •

Enhancing Performance Across Various LLMs: By selectively retaining the
initial directive/clauses and discarding the rest, Chain-of-OMP effectively
narrows down the search space for subsequent chains. This approach is akin to
the strategy in the classic Monty Hall problem, where a ”pick again” method
theoretically increases the chances of accuracy. Such a strategy is
anticipated to improve OpenMP pragma generation accuracy not just for OMPGPT,
but for other LLMs as well.

  * •

Automation of the Process: A key strength of Chain-of-OMP is its fully
automated chain generation capability. This stands in contrast to traditional
chain-of-thought techniques, where users typically need to craft prompts for
each step manually. The automation in Chain-of-OMP streamlines the process,
making it more efficient and user-friendly.

###  3.4. Fine-tuning

Due to the complexity of OpenMP pragmas, previous works either train an ML
model or fine-tune a language model for specific tasks covering a limited
selection of pragmas. In our work, we have employed a strategic fine-tuning
approach to demonstrate the enhanced performance of OMPGPT after fine-tuning
and to facilitate a comprehensive comparison with baseline models.

Following the setup of Monocoder, We fine-tune the pre-trained OMPGPT model
using the AdamW optimizer, a variant of the Adam optimizer known for its
effectiveness in large models and datasets. This process incorporates a linear
warmup phase over the initial 100 steps. The warmup phase gradually increases
the learning rate from zero to the initial learning rate set for training,
helping to stabilize the model’s learning process in its early stages.
Following the warmup, we implement a linear decay in the learning rate for the
remaining steps. This decay approach gradually reduces the learning rate,
allowing for finer adjustments to the model’s weights as it converges toward
optimal performance.

##  4\. Evaluation

The section delves into the evaluation of OMPGPT, encompassing task
definition, experiment setup, and evaluation results. We first evaluate our
base model, focusing on its language perplexity and the ability to generate
any OpenMP pragma. This is followed by the evaluation of the proposed Chain-
of-OMP methodology with the OMPGPT base model. Subsequently, we explore the
practical performance of OMPGPT after fine-tuning it for specific targeted
tasks. Each of these steps is critical in demonstrating the effectiveness and
versatility of our model in real-world HPC applications. We use the same
hardware as used in the training for evaluation.

###  4.1. Model Perplexity

Task Definition. Perplexity is a widely used metric in natural language
processing, particularly in the evaluation of language models. It’s a measure
of how well a probability model predicts a sample. Perplexity is defined as
the exponential of the average negative log-likelihood of a sequence of words.
If a model has low perplexity, it means that it assigns higher probabilities,
on average, to the test samples.

Evaluation Setup. We measure OMPGPT on the test set of HPCorpus and calculate
its perplexity, $PP$, using equation 2.

(2) $PP=\exp\left(-\frac{1}{N}\sum_{i=1}^{N}\log
P(w_{i}|w_{1},w_{2},\ldots,w_{i-1})\right)$

Where $P(w_{i},w_{1},w_{2},\ldots,w_{i-1})$ is the probability of the first
word $w_{i}$ occurring, given the sequence of the subsequent words
$w_{1},w_{2},\ldots,w_{i-1}$. We use prompt (a) in Figure 5 for general code
generation and calculate the perplexity. We compare the perplexity result of
OMPGPT with MonoCoder and other open-source language models.

Results. The perplexity scores for each model, as presented in Table 1, offer
a revealing comparison between OMPGPT and other prominent open-source language
models. Notably, OMPGPT, with a model size of 0.76B, achieves perplexity
scores of 3.54 and 3.66 for C and C++, respectively. This performance is
remarkably competitive, especially when considering its relatively smaller
size compared to other models. For instance, Monocoder (Kadosh et al., 2023b),
with a slightly larger size (0.89B), shows similar perplexity scores of 3.51
(C) and 3.69 (C++). However, larger models such as PolyCoder (2.7B) (Xu et
al., 2022), GPT-Neo (2.7B) (Gao et al., 2020), GPT-J (6B) (Wang and
Komatsuzaki, 2021), CodeX (12B) (Chen et al., 2021), StarCoder (15.5B) (Li et
al., 2023), and GPT-NeoX (20B) (Black et al., 2022) generally demonstrate
lower perplexity values, indicating a better performance in understanding and
generating code. This trend suggests a correlation between model size and
perplexity, but the efficiency of OMPGPT is notable, given its significantly
smaller size. It competently bridges the gap with larger models, indicating a
promising direction for efficient model design in the code generation tasks.

Table 1. Perplexity Comparison across Open-source Language Models Model | Size (B) | C | C++
---|---|---|---
OMPGPT | 0.76 | 3.54 | 3.66
MonoCoder | 0.89 | 3.51 | 3.69
PolyCoder | 2.7 | 2.33 | 2.99
GPT-Neo | 2.7 | 3.69 | 2.87
GPT-J | 6 | 2.82 | 2.47
CodeX | 12 | 2.55 | 1.95
StarCoder | 15.5 | 1.71 | 2.01
GPT-NeoX | 20 | 2.37 | 2.32

###  4.2. OpenMP Pragma Generation with OMPGPT Base Model

Task Definition. We trained OMPGPT with the preprocessed OpenMP code from
HPCorpus. The capability of generating any OpenMP pragma distinguishes OMPGPT
from the majority of the existing code-oriented Large Language Models (LLMs).
In this task, we rigorously evaluate the pragmas generated by OMPGPT.

Evaluation Setup. For a comprehensive assessment, we focus on the 15 most
prevalent OpenMP pragmas found within HPCorpus. We use these 15 OpenMP pragmas
in the HPCorpus test set as our benchmark. We follow the Prompt (b) in Figure
5, where the input code is followed by the hint #pragma omp to let OMPGPT know
that we are performing the pragma prediction task instead of generic code
generation task. Evaluation metrics are centered around accuracy, specifically
gauging the congruence between OMPGPT’s outputs and the original OpenMP
pragmas.

Results. Figure 7 presents the outcomes of the pragma prediction for the top
15 most prevalent OpenMP pragmas. Given that the evaluation criteria focus on
exact matches, the results demonstrate OMPGPT’s capability to generate a
diverse range of OpenMP pragmas.

Figure 7. Pragma Generation Accuracy with Base Model.

###  4.3. OpenMP Pragma Generation using Chain-of-OMP

Task Definition. As demonstrated in Section 4.2, OMPGPT is capable of
generating OpenMP pragmas for a given code snippet. However, the observed
accuracy is modest, primarily attributable to the intricate syntax associated
with OpenMP. This task is dedicated to a thorough evaluation of our proposed
chain-of-OMP methodology. We aim to assess its performance and draw a
comparative analysis with the basic prompts. This comparison will provide
critical insights into the effectiveness of Chain-of-OMP in addressing the
syntactical complexities inherent in OpenMP pragma generation.

Evaluation Setup. We evaluate chain-of-OMP for two specific OpenMP pragmas:
#pragma omp for schedule and #pragma omp target. These pragmas were chosen due
to their frequent usage in HPC tasks and their pragma alignment with the
three-component structure we previously outlined. We assess the effectiveness
of chain-of-OMP by comparing its accuracy with that of the base prompt,
thereby quantifying its performance enhancement.

Results. Table 2 illustrates the performance boosts provided by chain-of-OMP.
For #pragma omp for schedule, the base model achieved an accuracy of 1.3%,
while Chain-of-OMP significantly improved this to 7.8%, resulting in a net
boost of 600%. Similarly, with the #pragma omp target pragma, the base model’s
accuracy of 2.9% was enhanced to 16.9% by Chain-of-OMP, yielding a boost of
approximately 580%.

Table 2. chain-of-OMP results | Basic Prompt | chain-of-OMP
---|---|---
for schedule: | 1.3% | 7.8%
target | 2.9% | 16.9%
Table 3. Fine-tuning results for private, reduction, and SIMD clauses on the first test (P = Precision, R = Recall, Acc = Accuracy) Model | private | reduction | SIMD
---|---|---|---
| P | R | F1 | Acc | P | R | F1 | Acc | P | R | F1 | Acc
OMPGPT | 0.94 | 0.91 | 0.92 | 0.96 | 0.92 | 0.91 | 0.915 | 0.98 | 0.88 | 0.82 | 0.85 | 0.93
MonoCoder | 0.89 | 0.83 | 0.86 | 0.94 | 0.99 | 0.81 | 0.891 | 0.98 | N/A | N/A | N/A | N/A
GPT-3.5 | 0.4 | 0.16 | 0.23 | 0.77 | 0.53 | 0.57 | 0.549 | 0.92 | 0.51 | 0.55 | 0.53 | 0.84
Table 4. Fine-tuning results for private, reduction, and SIMD clauses on the second test with Accuracy % Model | private | reduction | SIMD
---|---|---|---
OMPGPT | 51.41 | 56.84 | 41.85
Chain-of-OMP | 63.77 | 60.1 | 58.19
MonoCoder | 48.55 | 51.78 | N/A
GPT-3.5 | 14.39 | 51.26 | 39.81

###  4.4. Fine-tuning OMPGPT Base Model for Specific Pragma Generation

Task Definition. As we state in Section 3.4, a fine-tuning phase is required
due to the complexity of OpenMP. This task focuses on fine-tuning OMPGPT for
three pragma generation sub-tasks: private, reduction, and SIMD. The
performance of OMPGPT in these sub-tasks will be evaluated against two
baseline models: MonoCoder, which has undergone fine-tuning for specific
pragma generation, and GPT-3.5, which, due to time and cost constraints, has
not been fine-tuned for these tasks. Additionally, we applied chain-of-OMP
with fine-tuned OMPGPT for each sub-task with $N_{\text{chain}}=2$.

Evaluation Setup. Our evaluation for each sub-task utilizes the test set from
our fine-tuning dataset. We employ four commonly used metrics to assess the
results: precision, recall, F1-score, and accuracy. Because these clauses also
have variables and other parameters, we break down our evaluation into two
subtests: the clause and the control structure. To be precise, in the first
subtest, we only check if the predicted clause matches with the expected
clause. Then, in the second subtest, we check that both the clause and its
control structure match with the expected clause and its control structure. It
is easy to see that the second test is stricter than the first test. These
metrics are defined as follows:

(3) Precision $\displaystyle=\frac{\text{True Positives}}{\text{True
Positives}+\text{False Positives}}$ (4) Recall $\displaystyle=\frac{\text{True
Positives}}{\text{True Positives}+\text{False Negatives}}$ (5) F1-score
$\displaystyle=2\times\frac{\text{Precision}\times\text{Recall}}{\text{Precision}+\text{Recall}}$
(6) Accuracy $\displaystyle=\frac{\text{True Positives}+\text{True
Negatives}}{\text{Total Number of Cases}}$

Results. Table 3 shows the performance of different models on the first
subtest for private, reduction, and SIMD clauses. MonoCoder does not support
the generation of SIMD clause. Hence, the table contains no entries for it for
that clause. Overall, we find that OMPGPT outperforms both MonoCoder (although
marginally) as well as GPT-3.5 models.

Table 4, on the other hand, shows the performance of those models on the
second subtest for private, reduction, and SIMD clauses. As expected, the
performance of all the models drop considerably in comparison to the first
subtest. Nonetheless, we find that OMPGPT still outperforms MonoCoder as well
as GPT-3.5. We would also like to point out that our chain-of-OMP technique
helps OMPGPT in this subtest as it improves the accuracy of predicting clauses
as well as the control structure correctly.

###  4.5. Observations

####  4.5.1. Performance of OMPGPT base model for OpenMP pragma generation.

The experiments detailed in Section 4.2 are designed to assess the capability
of OMPGPT in accurately generating various OpenMP pragmas. For these tests, we
adopted a stringent evaluation criterion where a generated pragma is deemed
correct only if it matches exactly with the corresponding pragma in the test
dataset. For example, the two OpenMP pragmas below are effectively the same in
terms of functionality but not considered correct. This high bar of precision
provides a robust assessment of OMPGPT’s capabilities. However, as indicated
by the results in Figure 7, this strict criterion results in relatively low
scores, reflecting the challenge of meeting such exacting standards in pragma
generation.

⬇

#pragma omp parallel for reduction(+:sum) private(var)

#pragma omp parallel for private(var) reduction(+:sum)

####  4.5.2. Effectiveness of chain-of-OMP

. The effectiveness of Chain-of-OMP is showcased through the analysis of two
selected OpenMP pragmas, as detailed in Section 4.3. Furthermore, we extended
our evaluation to include a smaller set comprising ten sample data of the top
15 most common pragmas. This testing led to two significant observations:

  * •

Consistent Performance Enhancement: In all the sets tested, Chain-of-OMP
consistently maintained or improved performance levels. There was no instance
where its application resulted in a performance decline.

  * •

Notable Improvement in Majority of Tasks: Impressively, Chain-of-OMP enhanced
the performance in 11 out of the 15 pragma generation tasks. This indicates a
substantial efficacy of Chain-of-OMP in improving OpenMP pragma generation
across a diverse range of tasks.

##  5\. Related Work

Machine learning approaches to develop tools for HPC programmers have been
investigated in recent years. The most recent work has focused on the
potential of language models for a variety of downstream HPC tasks, as
described in detail in (Kadosh et al., 2023b). For example, a few works have
fine-tuned existing code LLMs for tasks such as race detection (Ding et al.,
[n. d.]) and generating parallel code such as MPI routines (Nichols et al.,
2023). The latter work also explored the task of generating accurate OpenMP
pragmas. In contrast, (Kadosh et al., 2023c) formulated the task of predicting
OpenMP pragmas as a multi-label classification task. This work relied on a
multimodal transformer that also used graph inputs rather than a code LLM that
only represents code as tokens in a vocabulary.
Chain-of-thought prompting is a technique to improve the ability of LLMs to
reason over their inputs. The first work published on this displayed good
performance on math word problems and commonsense reasoning problems (Wei et
al., 2023). A few works have investigated the application of chained prompting
to complex programming problems. For example, (Rukmono et al., 2023) use this
prompting strategy to help generate summaries of software system components.
It could also be useful to generate incremental metadirective clauses
(Bhattacharjee et al., 2023) in OpenMP. Other works focus on devising natural
language prompts for models such as ChatGPT to improve code generation (Li et
al., [n. d.]; Liu et al., [n. d.]) and the reproduction of bugs from bug
reports (Feng and Chen, [n. d.]). The most related work applied this style of
prompting to correcting compilable code (Chen et al., [n. d.]).

##  6\. Conclusion and Future Work

Large language models, such as GPT-3.5, have revolutionized the field of
natural language processing, and their applications to programming problems
such as code generation, code completion, etc., look equally promising. These
models however, perform poorly when applied to problems from specific domains
such as HPC. Moreover, their large size (partly because of natural language
support) demands expensive computing resources for training as well as
deployment. Realizing these limitations, we developed OMPGPT, a 0.76B
parameter model that is trained on an HPC-specific code dataset only and for
the HPC-specific task of OpenMP pragma generation. Although, OMPGPT delivers
reasonable performance in comparison to the large LLMs in terms of perplexity,
we found that it performs poorly in accurately generating OpenMP pragmas. Our
analysis led to our novel prompting technique, called chain-of-OMP, that
assists OMPGPT in improving its accuracy by using prompts in a step-by-step
fashion. Overall, we believe that these findings could help in improving the
performance of smaller LLMs on domain-specific tasks, thereby paving the way
to democratize LLMs.

Although our results are promising, we envision several avenues for future
work. Specifically, we plan to evaluate the performance of the chain-of-OMP
technique on other OpenMP clauses as well. Additionally, we plan to look into
extending the context window beyond the existing snippet of the for loop. In
OpenMP, sometimes the expressions before the for loop also play a role in the
generation of the correct clauses. Another interesting avenue to explore is to
train a model like GPT-3.5 from scratch on our training data and chain-of-OMP
prompting technique and to check for performance gains, albeit at a higher
cost. Finally, we plan to explore this research direction to approach other
HPC tasks and assist HPC developers.

## References

  * (1)
  * Bhattacharjee et al. (2023) Arijit Bhattacharjee, Christopher S. Daley, and Ali Jannesari. 2023.  OpenMP Offload Features and Strategies for High Performance across Architectures and Compilers. In _2023 IEEE International Parallel and Distributed Processing Symposium Workshops (IPDPSW)_. 564–573.
  * Black et al. (2022) Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. 2022.  GPT-NeoX-20B: An Open-Source Autoregressive Language Model.  arXiv:2204.06745 [cs.CL]
  * Black et al. (2021) Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. 2021.  _GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow_.  https://doi.org/10.5281/zenodo.5297715 If you use this software, please cite it using these metadata..
  * Chen et al. ([n. d.]) Le Chen, Arijit Bhattacharjee, Nesreen K. Ahmed, Niranjan Hasabnis, Gal Oren, Bin Lei, and Ali Jannesari. [n. d.].  _CompCodeVet: A Compiler-guided Validation and Enhancement Approach for Code Dataset_.  https://doi.org/10.48550/arXiv.2311.06505 arXiv:2311.06505 [cs]
  * Chen et al. (2023a) Le Chen, Pei-Hung Lin, Tristan Vanderbruggen, Chunhua Liao, Murali Emani, and Bronis de Supinski. 2023a.  Lm4hpc: Towards effective language model application in high-performance computing. In _International Workshop on OpenMP_. Springer, 18–33.
  * Chen et al. (2023b) Le Chen, Quazi Ishtiaque Mahmud, Hung Phan, Nesreen Ahmed, and Ali Jannesari. 2023b.  Learning to Parallelize with OpenMP by Augmented Heterogeneous AST Representation.  _Proceedings of Machine Learning and Systems_ 5 (2023).
  * Chen et al. (2021) Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert-Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. 2021.  Evaluating Large Language Models Trained on Code.  arXiv:2107.03374 [cs.LG]
  * Dagum and Menon (1998) L. Dagum and R. Menon. 1998.  OpenMP: an industry standard API for shared-memory programming.  _IEEE Computational Science and Engineering_ 5, 1 (1998), 46–55.  https://doi.org/10.1109/99.660313
  * Ding et al. ([n. d.]) Xianzhong Ding, Le Chen, Murali Emani, Chunhua Liao, Pei-Hung Lin, Tristan Vanderbruggen, Zhen Xie, Alberto Cerpa, and Wan Du. [n. d.].  HPC-GPT: Integrating Large Language Model for High-Performance Computing. In _Proceedings of the SC ’23 Workshops of The International Conference on High Performance Computing, Network, Storage, and Analysis_ (Denver CO USA, 2023-11-12). ACM, 951–960.  https://doi.org/10.1145/3624062.3624172
  * Dutta et al. (2023a) Akash Dutta, Jordi Alcaraz, Ali TehraniJamsaz, Eduardo Cesar, Anna Sikora, and Ali Jannesari. 2023a.  Performance Optimization using Multimodal Modeling and Heterogeneous GNN. In _Proceedings of the 32nd International Symposium on High-Performance Parallel and Distributed Computing_. 45–57.
  * Dutta et al. (2023b) Akash Dutta, Jee Choi, and Ali Jannesari. 2023b.  Power Constrained Autotuning using Graph Neural Networks.  _arXiv preprint arXiv:2302.11467_ (2023).
  * Feng and Chen ([n. d.]) Sidong Feng and Chunyang Chen. [n. d.].  Prompting Is All You Need: Automated Android Bug Replay with Large Language Models. arXiv.  https://doi.org/10.48550/arXiv.2306.01987 arXiv:2306.01987 [cs]
  * Gao et al. (2020) Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. 2020\.  The Pile: An 800GB Dataset of Diverse Text for Language Modeling.  _arXiv preprint arXiv:2101.00027_ (2020).
  * Kadosh et al. (2023a) Tal Kadosh, Niranjan Hasabnis, Timothy Mattson, Yuval Pinter, and Gal Oren. 2023a.  Quantifying openmp: Statistical insights into usage and adoption. In _2023 IEEE High Performance Extreme Computing Conference (HPEC)_. IEEE, 1–7.
  * Kadosh et al. (2023b) Tal Kadosh, Niranjan Hasabnis, Vy A Vo, Nadav Schneider, Neva Krien, Mihai Capota, Abdul Wasay, Nesreen Ahmed, Ted Willke, Guy Tamir, et al. 2023b.  Domain-Specific Code Language Models: Unraveling the Potential for HPC Codes and Tasks.  _arXiv preprint arXiv:2312.13322_ (2023).
  * Kadosh et al. (2023c) Tal Kadosh, Nadav Schneider, Niranjan Hasabnis, Timothy Mattson, Yuval Pinter, and Gal Oren. 2023c.  Advising OpenMP Parallelization via a Graph-Based Approach with Transformers.  _arXiv preprint arXiv:2305.11999_ (2023).
  * Kocetkov et al. (2022) Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Ferrandis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, Dzmitry Bahdanau, Leandro von Werra, and Harm de Vries. 2022.  The Stack: 3 TB of permissively licensed source code.  _Preprint_ (2022).
  * Li et al. ([n. d.]) Jia Li, Ge Li, Yongmin Li, and Zhi Jin. [n. d.].  _Structured Chain-of-Thought Prompting for Code Generation_.  arXiv:2305.06599 [cs] http://arxiv.org/abs/2305.06599
  * Li et al. (2023) Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Marc Marone, Christopher Akiki, Jia Li, Jenny Chim, Qian Liu, Evgenii Zheltonozhskii, Terry Yue Zhuo, Thomas Wang, Olivier Dehaene, Mishig Davaadorj, Joel Lamy-Poirier, João Monteiro, Oleh Shliazhko, Nicolas Gontier, Nicholas Meade, Armel Zebaze, Ming-Ho Yee, Logesh Kumar Umapathi, Jian Zhu, Benjamin Lipkin, Muhtasham Oblokulov, Zhiruo Wang, Rudra Murthy, Jason Stillerman, Siva Sankalp Patel, Dmitry Abulkhanov, Marco Zocca, Manan Dey, Zhihan Zhang, Nour Fahmy, Urvashi Bhattacharyya, Wenhao Yu, Swayam Singh, Sasha Luccioni, Paulo Villegas, Maxim Kunakov, Fedor Zhdanov, Manuel Romero, Tony Lee, Nadav Timor, Jennifer Ding, Claire Schlesinger, Hailey Schoelkopf, Jan Ebert, Tri Dao, Mayank Mishra, Alex Gu, Jennifer Robinson, Carolyn Jane Anderson, Brendan Dolan-Gavitt, Danish Contractor, Siva Reddy, Daniel Fried, Dzmitry Bahdanau, Yacine Jernite, Carlos Muñoz Ferrandis, Sean Hughes, Thomas Wolf, Arjun Guha, Leandro von Werra, and Harm de Vries. 2023.  StarCoder: may the source be with you!  arXiv:2305.06161 [cs.CL]
  * Liu et al. ([n. d.]) Chao Liu, Xuanlin Bao, Hongyu Zhang, Neng Zhang, Haibo Hu, Xiaohong Zhang, and Meng Yan. [n. d.].  _Improving ChatGPT Prompt for Code Generation_.  https://doi.org/10.48550/arXiv.2305.08360 arXiv:2305.08360 [cs]
  * Luo et al. (2023) Ziyang Luo, Can Xu, Pu Zhao, Qingfeng Sun, Xiubo Geng, Wenxiang Hu, Chongyang Tao, Jing Ma, Qingwei Lin, and Daxin Jiang. 2023.  WizardCoder: Empowering Code Large Language Models with Evol-Instruct.  arXiv:2306.08568 [cs.CL]
  * Nichols et al. (2023) Daniel Nichols, Aniruddha Marathe, Harshitha Menon, Todd Gamblin, and Abhinav Bhatele. 2023.  Modeling Parallel Programs using Large Language Models.  arXiv:2306.17281 [cs.DC]
  * Rozière et al. (2023) Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. 2023.  Code Llama: Open Foundation Models for Code.  arXiv:2308.12950 [cs.CL]
  * Rukmono et al. (2023) Satrio Adi Rukmono, Lina Ochoa, and Michel RV Chaudron. 2023.  Achieving High-Level Software Component Summarization via Hierarchical Chain-of-Thought Prompting and Static Code Analysis. In _2023 IEEE International Conference on Data and Software Engineering (ICoDSE)_. IEEE, 7–12.
  * Touvron et al. (2023) Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. 2023.  Llama 2: Open Foundation and Fine-Tuned Chat Models.  arXiv:2307.09288 [cs.CL]
  * Valero-Lara et al. (2023) Pedro Valero-Lara, Alexis Huante, Mustafa Al Lail, William F. Godoy, Keita Teranishi, Prasanna Balaprakash, and Jeffrey S. Vetter. 2023.  Comparing Llama-2 and GPT-3 LLMs for HPC kernels generation.  arXiv:2309.07103 [cs.SE]
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  _Advances in neural information processing systems_ 30 (2017).
  * Wang and Komatsuzaki (2021) Ben Wang and Aran Komatsuzaki. 2021.  GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model.  https://github.com/kingoflolz/mesh-transformer-jax.
  * Wei et al. (2023) Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed Chi, Quoc Le, and Denny Zhou. 2023.  Chain-of-Thought Prompting Elicits Reasoning in Large Language Models.  arXiv:2201.11903 [cs.CL]
  * Xu et al. (2022) Frank F. Xu, Uri Alon, Graham Neubig, and Vincent Josua Hellendoorn. 2022.  A systematic evaluation of large language models of code. In _Proceedings of the 6th ACM SIGPLAN International Symposium on Machine Programming_ (San Diego, CA, USA) _(MAPS 2022)_. Association for Computing Machinery, New York, NY, USA, 1–10.  https://doi.org/10.1145/3520312.3534862