# Accelerating Self-Supervised Learning via Efficient Training Strategies

Mustafa Taha Koçyiğit Timothy M. Hospedales Hakan Bilen

University of Edinburgh, UK
<EMAIL_ADDRESS>

###### Abstract

Recently the focus of the computer vision community has shifted from expensive
supervised learning towards self-supervised learning of visual
representations. While the performance gap between supervised and self-
supervised has been narrowing, the time for training self-supervised deep
networks remains an order of magnitude larger than its supervised
counterparts, which hinders progress, imposes carbon cost, and limits societal
benefits to institutions with substantial resources. Motivated by these
issues, this paper investigates reducing the training time of recent self-
supervised methods by various model-agnostic strategies that have not been
used for this problem. In particular, we study three strategies: an extendable
cyclic learning rate schedule, a matching progressive augmentation magnitude
and image resolutions schedule, and a hard positive mining strategy based on
augmentation difficulty. We show that all three methods combined lead up to
2.7 times speed-up in the training time of several self-supervised methods
while retaining comparable performance to the standard self-supervised
learning setting.

##  1 Introduction

Learning representations without manual human annotations that can be
successfully transferred to various downstream tasks has been a long standing
goal in machine learning [17, 48]. Self-supervised learning (SSL) aims at
learning such representations discriminatively through pretext tasks such as
identifying the relative position of image patches [14] and solving jigsaw
puzzles [36]. The recent success of SSL methods [23, 6, 8] builds on
contrastive learning where the representations are in a latent space invariant
to various image transformations such as cropping, blurring and colour
jittering. Contrastive learned representations have been shown to obtain on
par performance with their supervised counterparts when transferred to various
vision tasks including image classification, object detection, semantic
segmentation [7, 22], and extended to medical imaging [1] as well as multi-
view [45] and multi-modal learning [37].

Figure 1: Classification accuracy v.s. training cost of various SSL methods,
which are trained on a subset of ImageNet by using ResNet50 (see Sec. 4.2 for
experimental details). Our method (drawn in blue) successfully significantly
accelerates the SSL methods (drawn in red) without any significant drop in
their performance.

Despite the remarkable progress, an important downside of SSL methods is their
high training cost which hampers the development and adoption of these
promising techniques. Even the most efficient SSL methods require at least an
order of magnitude more computation to reach the performance of supervised
methods, e.g. training BYOL [22] requires 23 times more computation resulting
from 8.8 times more iterations and 2.6 times more computation at each
iteration than its supervised counterpart to reach similar accuracy. The large
training cost is largely due to the challenging task of learning invariant
representations over a large set of images and various augmentation
transforms. In this paper, we focus on developing algorithms that can speed up
the training of SSL while maintaining their performance.

Prior work in efficient supervised training reduced the cost of training by
gradually increasing the training resolution together with the augmentations
magnitude using Progressive Learning [44]. While augmentation is a
regularization mechanism in supervised learning it is the main source of
supervision in SSL and therefore plays a much more important role.
Additionally, Super Convergence [41] is introduced which uses a cyclic
learning rate and anti-phased momentum schedule that is used to accelerate
convergence in supervised learning tasks. The longer duration of training
makes the application of Super Convergence more difficult and using it
together with Progressive Learning causes instabilities in the early stage of
training.

In this work, we aim to increase the training efficiency of self-supervised
training methods by optimizing the learning rate, resolution, and augmentation
schedules to allow reaching the same level of performance with a smaller
computational budget. Inspired from [41] and [44] that were used in supervised
learning we propose a combined learning rate and resolution schedule for self-
supervised learning. Additionally, we modify the augmentation strategy used in
self-supervised learning for faster training.

Our contributions are three folds. Fixed 1-cycle Learning Rate schedules (see
Sec. 3.1) allows Super-Convergence [41] to take place for longer training
times by fixing the duration of its warm-up phase while extending the decay
phase and comparing it against alternative learning rate schedules used in
self-supervised learning. Super Progressive Learning (see Sec. 3.2) proposes a
suitable Progressive Learning schedule [44] for SSL by adding a warm-up stage
with full resolution at the beginning of the resolution schedule while
increasing the augmentation magnitude gradually during training so that it
works well with our proposed learning rate schedule. Hard Augment (see Sec.
3.3) selects a hard augmentation pair dynamically from multiple low resolution
augmentations in order to boost the learning signal and regularize the
training.

We show that our method accelerates the training of different self-supervised
learning methods and architectures while maintaining comparable performance to
standard training. We provide ablation studies for analysing the contribution
of different methods and determining the optimum hyperparameters of our
methods and add theoretical justifications.

##  2 Related Works

Generative models like DBN [25], VAE [32] and GANs [20] have been used for
unsupervised representation learning however discriminative SSL have proven to
be more effective. While early work on SSL focused on designing various
pretext tasks such as solving jigsaw puzzles [36], colouring [51, 33, 52] and
predicting rotation [18], more recent techniques focused on learning
augmentation invariant representation [49] [46] to replace the loss functions
that require manual supervision. Recently, Contrastive learning methods[23, 6,
8] have shown promising results for various target tasks and caused renewed
interest in this area. However, at least an order of magnitude, more
computation is required when training these methods in order to get comparable
performance to supervised learning. In fact, we demonstrate our method in
these recent models and show that SimSiam [9], BYOL [22], MoCov2 [8], SimCLR
[6] and DINO [4] training can significantly be sped up.

By adaptation gradient noise during training large scale data parallel
training has been used for increasing training speed in large deep learning
methods. Using many accelerators together and increasing the total batch size
training speed can be scaled up almost linearly with the number of
accelerators [35]. This way, supervised ImageNet training can be done in
minutes [42] or hours [21] instead of days even though the total training cost
remains roughly the same. This required adapting the hyperparameters so that
gradient noise was managed properly [43, 28, 5] by using the linear
relationship between the batch size and learning rate [21, 42]. It was shown
that training with larger learning rates had also a regularization effect [28,
5, 34] and prevented convergence to sharper minima [31]. We use the
relationship between gradient noise and learning rate schedules to optimize
for training efficiency instead of speed where we synchronize the learning
rate schedule with our resolution schedule to reduce the total cost of
training. We use a cyclic learning rate schedule [41] which combines cyclic
learning rate and cyclic momentum adaptation together as our learning rate
schedule due to its efficient training performance and adapt it to longer
training regimes common in SSL training.

Curriculum learning [2] has been used to accelerate convergence for deep
learning methods and when applying it to self-supervised learning,
augmentation strength and image resolution are the most relevant parameters.
Image resolution directly affects accuracy [26, 47] while reducing run-time
quadratically. Previous work [27] on efficient supervised learning gradually
increased the image resolution in the DAWN benchmark [11] to accelerate
training while having slightly lower performance. Recently [44] proposed
gradually increasing the training image resolution while increasing
augmentations strength to allow for both faster and accurate training. While
augmentations are used for regularization in supervised learning they are the
main source of supervision for recent self-supervised methods and determine
task difficulty. We build on this intuition and adapt Progressive Learning for
efficient self-supervised training since in order to reduce the training cost
especially in the initial phase while increasing the difficulty of the task
gradually. To the best of our knowledge having a good schedule for
augmentation magnitude and its relation to resolution hasn’t been examined for
self-supervised learning.

Hard positive mining is another technique that can be used to accelerate
training. Importance sampling based on sample loss has been shown to
accelerate supervised learning [30, 29]. In object detection [38] selects
regions with the highest loss in order to select useful positives and balance
positive to negative regions effectively. [19] applied multiple augmentations
and back-propagated only the augmented sample with the maximum loss in order
to improve the adversarial robustness of their method while increasing
supervised performance.[3] have introduced a multi-crop strategy for self-
supervised learning that combines low resolution and high resolution crops in
order to show increased performance while trying to keep the computational
cost limited. We mine hard augmentations during training by utilizing the loss
based importance sampling technique to dynamically select the most useful
augmentations. Since our focus is on efficient training, we used down sampled
versions of the augmented images in the selection pass to minimize the
overhead.

##  3 Method

In this section we introduce the techniques used to accelerate self-supervised
training. Each technique can be used to accelerate training on its own but
they can be combined together in a synergistic way to enable faster training.
We introduce Fixed 1-cycle Learning Rate Schedule in Sec. 3.1, Super
Progressive Learning in Sec. 3.2 and Hard Augment in Sec. 3.3.

We formulate a contrastive loss function that measures difference between
pairs of images and optimizing this loss allows us to learn parameters
$\theta$ for a deep neural network $f_{\theta}$ that produces similar
representations for augmentations of the same image and have representation
that is going to be useful for various target tasks. Let $D$ be an unlabeled
dataset consisting of $|D|$ images with resolution $r$. We randomly sample two
image transformations $\tau$ and $\bar{\tau}$ from the transformation space
$T$ for each training image $\bm{x}$, apply them to $\bm{x}$ to obtain two
views $\bm{v}$ and $\bm{\bar{v}}$, extract their features through the deep
neural network $\bm{z}=f_{\theta}(\bm{v})$ and
$\bm{\bar{z}}=f_{\theta}(\bm{\bar{v}})$ respectively. To learn the network
weights $\theta$, we minimize the loss function $L$ that represents the
mismatch between two representations $\bm{z}$ and $\bm{\bar{z}}$ over the
dataset and transformation space:

$\min_{\theta}\mathbb{E}_{x\sim D,(\tau,\bar{\tau})\sim
T}L(\bm{z},\bm{\bar{z}})$ (1)

We use minibatch Stochastic Gradient Descent (SGD) optimizer with momentum:

$L^{(B)}(\theta)=\frac{1}{|B|}\sum_{\bm{x}\sim B,(\tau,\bar{\tau})\sim
T}L(\bm{z},\bar{\bm{z}}),$ (2)

where the minibatch $B$ consists of $|B|$ randomly sampled images from $D$ and
the loss is averaged over the samples to obtain a noisy yet unbiased estimate
of the true gradient.

The update rule for $\theta$ is given as:

$\displaystyle\mu_{t}$
$\displaystyle=\beta_{t}\mu_{t-1}-\epsilon_{t}\nabla_{\theta_{t}}L^{(B)}(\theta_{t})$
(3) $\displaystyle\theta_{t+1}$ $\displaystyle=\theta_{t}-\mu_{t}$

where $\epsilon_{t}$ is the learning rate and $\beta_{t}$ is the momentum
weighs at step $t$. The values for the learning rate and momentum is given by
a learning rate scheduler $(\epsilon_{t},\beta_{t})=S(t)$.

Eq. 3 can be interpreted as stochastic differential equation (_e.g_. [43]):

$\displaystyle\frac{d\theta}{dt}$ $\displaystyle=\mu,$
$\displaystyle\frac{d\mu}{dt}$
$\displaystyle=\beta_{t}\mu-\frac{dL}{d\theta}+\eta_{t}$ (4)

where $\eta(t)\sim\mathcal{N}(0,g\bm{F}(\theta)/|D|)$ is an additive Gaussian
noise originating from the stochasticity and $\bm{F}(\theta)$ is the
covariance matrix for gradients of the samples and $g$ is the “noise scale”
which is given by $g\approx\frac{\epsilon_{t}|D|}{b(1-\beta_{t})}$. For our
analysis we will utilize the noise in the gradients which is proportional with
the learning rate $\epsilon_{t}$ while being inversely proportional with batch
size $b$ and $1-\beta_{t}$. This relationship plays an important role when
adapting hyperparameters to different setups and for giving us a theoretically
grounded understanding for learning rate schedules and their relationship with
Progressive Learning.

###  3.1 Fixed 1-cycle Learning Rate Schedule

Figure 2: Learning rate (left) and momentum (middle) for standard cosine
annealing learning rate schedule in blue and F1-CLR schedule in red are
visualized. Learning rate for F1-CLR schedule with different length of
training instances is visualized (right)

Cosine annealing learning rate schedule which has been used widely in SSL [23,
8, 9] decays the learning rate starting from a maximum learning rate value
$\epsilon_{max}$ using a cosine function and can be described as,

$\displaystyle\epsilon_{t}=\frac{1}{2}\epsilon_{max}(\cos({\frac{t}{L}\pi)}+1),$
(5)

where $L$ is the total number of iterations.

Methods that perform better with larger batch sizes train with larger learning
rates in order to maintain the gradient “noise scale”
$g\propto\frac{\epsilon_{t}}{b}$ during training. In order to prevent
instabilities in the early stages of training due to rapidly changing
parameters values [21] a linear warm-up phase for the learning rate is added
[6, 22, 10, 50, 3, 4],

$\displaystyle\epsilon_{t}=\begin{cases}\frac{t}{t_{w}}\epsilon_{max},&t<t_{w}\\\
\frac{1}{2}\epsilon_{max}(\cos({\frac{t-t_{w}}{L-t_{w}}\pi)}+1),&\text{otherwise}\end{cases}$
(6)

where $t_{w}$ is the number of warm up steps.

On the other hand, Cyclic Learning Rate (CLR) schedule introduced by [39] and
later refined as the 1-cycle learning rate (1-CLR) schedule [41, 40] starts
with a small learning rate and than increases the learning rate to a very high
value and than decays it gradually to a very small value while changing the
momentum at the opposite direction as the learning rate. The learning rates
and momentum values at each step using a cosine window can be calculated by

$\displaystyle\epsilon_{t}=\begin{cases}\epsilon_{max}-\frac{1}{2}\epsilon_{max}(\cos({\frac{t}{\rho
L}\pi)}+1)&t<\rho L\\\ \frac{1}{2}\epsilon_{max}(\cos({\frac{t-\rho
L}{L(1-\rho)}\pi)}+1)&\text{otherwise}\end{cases}$ (7)
$\beta_{t}=\begin{cases}\beta_{l}+\frac{1}{2}(\beta_{h}-\beta_{l})(\cos({\frac{t}{\rho
L}\pi)}+1)&t<\rho L\\\
\beta_{h}-\frac{1}{2}(\beta_{h}-\beta_{l})(\cos({\frac{t-\rho
L}{L(1-\rho)}\pi)}+1)&\text{otherwise}\end{cases}$

where $\rho$ is the time percentage of time allocated for the first phase,
while $\beta_{l}$ and $\beta_{h}$ are lower and higher limits for momentum
respectively. We use the learning rate (LR) range test proposed by [40] to set
the maximum learning rate $\epsilon_{max}$ please see the details in
supplementary material C.

A phenomenon called Super Convergence has been demonstrated in supervised
classification where using larger learning rates and the 1-CLR schedule
training time to reach a specified performance has been decreased dramatically
[27]. However, unlike supervised learning where longer training times do not
generally result in better performance and can sometimes even cause worse
performance due to over-fitting, in SSL the quality of the representation
typically improves with longer training time [8, 6]. A problem with the
current 1-CLR is that the percentage of the first phase is being determined in
proportion to the full training duration which causes an extremely long warm-
up phase in longer training settings wasting compute time. To address this
problem, we propose to extend the annealing phase of 1CLR while keeping the
warm up length $t_{w}$ the same which we call _Fixed 1-cycle Learning Rate_
(F1-CLR) schedule,

$\displaystyle\epsilon_{t}=\begin{cases}\epsilon_{max}-\frac{1}{2}\epsilon_{max}(\cos({\frac{t}{t_{w}}\pi)}+1&t<t_{w}\\\
\frac{1}{2}\epsilon_{max}(\cos({\frac{t-t_{w}}{L-t_{w}}\pi)}+1)&\text{otherwise}\end{cases}$
(8)
$\beta_{t}=\begin{cases}\beta_{l}+\frac{1}{2}(\beta_{h}-\beta_{l})(\cos({\frac{t}{t_{w}}\pi)}+1)&t<t_{w}\\\
\beta_{h}-\frac{1}{2}(\beta_{h}-\beta_{l})(\cos({\frac{t-t_{w}}{L-t_{w}}\pi)}+1).&\text{otherwise}\end{cases}$

Fig. 2 shows a comparison between cosine schedule and F1-CLR schedule in terms
of learning rate (left) and momentum (middle) as well as illustrating the
F1-CLR learning rate schedule for different lengths of training (right). The
anti-phased movement of momentum in F1-CLR allows larger learning rates to be
achieved while keeping gradient noise in check
($g\propto\epsilon_{t}/(1-\beta_{t})$ see Eq. 4). Note that we do not extend
the warm up phase, as the warm up is used as a stabilizer, while the learning
rate and gradient noise scale are increasing ($g\propto\epsilon_{t}$ see Eq.
4).

###  3.2 Super Progressive Learning

Figure 3: Augmentation magnitude and input resolution schedule where
augmentation magnitude is gradually increased while the input size starts
large and than gradually increases from its minimum value mimicking the
inverse of the F1-CLR schedule. (left) Augmentation Magnitude, (middle) Super
Progressive Learning resolution schedule, (right) Progressive Learning
resolution schedule

Input resolution is very important factor that effects training time.
Typically image resolution is a trade-off hyperparameter between performance
and computational load, i.e. higher resolution, higher performance but also
more computations. In SSL, the standard practice [23, 8, 9, 6, 22, 10, 50, 3,
4] is to train with fixed resolution and fixed augmentation settings.

We hypothesize that higher resolution inputs are only necessary when there is
small amount of noise in the gradients, and hence, propose to employ a
learning rate aware progressive learning strategy inspired by [44]. However,
as starting training with small input resolution as in [44] results in
additional noise in the gradients which prevent it from being trivially
incorporated to F1-CLR due to the instabilities in the warm-up phase. Hence,
we propose train with full resolution images during the warm-up phase before
going into the linearly increasing resolution schedule while we gradually
increase the augmentation magnitude as shown in Fig. 3 (left, middle). We call
this strategy _Super Progressive Learning_. When the learning rate is high
there is a large amount of noise in the gradients $g\propto\epsilon_{t}$ and
this allows us to use low resolution images at that point without resulting in
bad performance however since the training starts with small learning rates we
need to adapt our schedule and use large resolution images.

Note that the speedup of this strategy depends on the input resolution
schedule ($r_{t}$). Assuming a quadratic relationship between training time
and resolution, the speedup $M$ can be calculated as:

$M=\frac{1}{L}\sum_{t}(\frac{r_{max}}{r_{t}})^{2},$ (9)

where $r_{t}$ is the input image resolution at time $t$, $r_{\text{max}}$ is
the final resolution. We discretize the resolution steps by 32 due to the fact
that most network architectures have 5 pooling/dilation layers. The minimum
value for the resolution schedule is determined empirically.

###  3.3 Hard Augment

So far we focused on varying the effective step size and input resolution.
Informativeness of training samples is another important factor that can speed
up the training by providing more efficient gradients. We propose Hard Augment
to boost the learning signal in self-supervised learning. We generate $m$
augmentations $\tau^{m}=\\{\tau_{1..m}:\tau_{1..m}\sim T\\}$ for each image
and evaluate the loss for all pairs of augmentations $p=\binom{\tau^{m}}{2}$
using a forward pass to select the augmentation pair that results in the
largest amount of loss for back-propagation. By focusing on the samples that
produce the most lost and ignoring a large fraction of the augmentations that
produce a small loss and do not make a significant contribution to the
training we can accelerate training dramatically. The overall objective that
we optimize is

$\displaystyle L^{(B)}_{ha}(\theta)=\frac{1}{|B|}\sum_{\bm{x}\sim
D,\tau_{1..m}\sim T}\left[\max_{\tau_{p}\in
p}L(\bm{z_{i}},\bm{z_{j}})\right].$ (10)

[19] have shown that using the maximum of multiple augmentations applies a
regularisation on the gradient-norm with respect to the input images
$\left\|\nabla_{\bm{x}}L(\bm{z_{i}},\bm{z_{j}})\right\|_{2}$ that is in the
order of $\sigma\sqrt{\log m}$ where $\sigma$ is the strength of the
augmentation under the assumption that
$\tau(x)\sim\mathcal{N}(x,\sigma^{2}I)$. The regularization effect in our
method can be seen as a corollary of their theorem that uses a pair of
augmentations instead with a number of equivalent augmentations as the number
of pairs $|p|$.

Crucially in order to decrease the overhead of selection we propose to down
sample the images during the selection pass to $r_{sel}$. The image pairs that
have the highest loss will than be used in training in their original
resolution $r$. We do not need to have full resolution for the selection pass
as we only need to have sufficient resolution to rank the augmentation pairs
with respect to their loss and to find the highest loss pair which requires
far less precision than is required for obtaining good quality gradient
updates during training.

$M_{ha}=\frac{r^{2}C}{r^{2}C+mr_{sel}^{2}}$ (11)

We find the minimum resolution $r_{sel}$ that we can use for the selection
pass empirically. On a typical setting in ImageNet where the original
resolution is $r=224$ and the reduced selection resolution is $r_{sel}=64$
with $m=4$ augmentations and 6 possible pairs and the fraction for full
training iteration cost to a forward pass cost is $C=6$ we only add
$1-M_{ha}=\%5$ overhead on the training while significantly boosting the
training speed by focusing on the most valuable augmentations.

Algorithm 1 Efficient SSL Training, PyTorch-like

⬇

# f: backbone + projection layers

# criterion: loss function for SSL method

# aug: augmentation function

# F1CLR: Super Convergence Schedule (Section 3.1)

# SP: Super Progressive Schedule (Section 3.2)

t=0 # iteration

for e in range(epochs):

for x in loader: # load a minibatch x with |B| samples

eta, beta = F1CLR(t) # Update learning rate and momentum weight

r, m = SP(t) # Update resolution(r) and augmentation magnitude(m)

vs = [aug(x, r, m) for i in range(n)] # n random augmentation

vi, vj = hard_augment(vs) # select hardes augmentations

zi, zj = f(vi), f(vj) # forward-propagation

L = criterion(zi,zj) # Loss calculation

L.backward() # back-propagation

update(f, eta, beta) # SGD update

t+=1

def hard_augment(vs): # negative cosine similarity

vs_small = interpolate(vs, r_min) # down sample views to r_min

zs = f(vs) # forward prop for all views

indexes = arg_max_loss_pair(zs) # determine pairs with maximum loss

vi, vj = hard_select(vs, indexes) # select pairs by tensor indexing

return vi, vj

We have provided the pseudo code for our method in Algorithm 1. Our method can
easily be adopted to different SSL methods without any major difficulty as it
operates at the level of augmentations, optimizer parameters and sampling
strategies.

| Baseline | Efficient
---|---|---
Architecture | Acc.(%) | Cost (pF) | Acc.(%) | Speedup
ResNet18 | 51.1 | 3756 | 51.3 | 2.3x
ResNet50 | 69.2 | 8494 | 69.3 | 2.3x
Table 1: Our efficient strategies applied to different backbone architectures using SimSiam model on Imagenet dataset. Speed up is reported based on the reduction in training cost which is measured in terms of petaFLOPs | Baseline | Efficient
---|---|---
Model | Acc.(%) | Cost (pF) | Acc.(%) | Speedup
SimSiam | 90.9 | 248 | 91.4 | 2.3x
BYOL | 92.0 | 371 | 92.6 | 2.7x
SimCLR | 91.3 | 247 | 90.4 | 2.3x
MoCov2+ | 92.4 | 371 | 92.8 | 2.7x
DINO | 89.8 | 372 | 91.2 | 2.4x
Table 2: Our efficient strategies applied to different SSL models training on
ResNet50 backbone architecture on Imagenette dataset. Speed up is reported
based on the reduction in training cost which is measured in terms of
petaFLOPs

##  4 Experiments

In this section, we first study evaluate our method on different backbone
architectures in Sec. 4.1 and on various existing SSL methods in Sec. 4.2,
finally we analyze the effect of each proposed component in our method in Sec.
4.3 and run ablation studies that analyze the hyperparameters and different
trade-offs of our method in Sec. 4.4.

We evaluate the experiments on the ImageNet dataset [13] in terms of linear
probing accuracy. This is done by freezing the backbone after pre-training and
training a linear classification layer in a supervised way following [9]. The
evaluation on the Imagenette dataset111ImageNet subset that consist of 10
classes (tench, English springer, cassette player, chain saw, church, French
horn, garbage truck, gas pump, golf ball, parachute) avalilable at
https://github.com/fastai/imagenette is done using online kNN accuracy of the
feature maps. When calculating speed up we consider the number of steps used
in training and multiply that by the number of floating point calculations
made in each step and compare it against the baseline.

#### Implementation Details

Here we provide details on the default parameters that we have used during our
experiments where we use the SimSiam [9] model with ResNet50 backbone using
the SGD optimizer with .9 momentum. On ImageNet experiments we use a batch
size of 512 and weight decay of 0.0001. For the baseline experiments we use
the Cosine Annealing learning rate schedule with learning rate of 0.1
($\eta_{max}$) and trained them for 200 epochs ($t_{max}$) and for our
accelerated F1-CLR schedule we use a learning rate of 0.16 ($\eta_{max}$) with
10 epoch warm-up ($t_{w}$) and train them for 120 epochs ($t_{max}$). For
Imagenette experiments we use a batch size of 128, weight decay of 0.0005,
learning rate of 0.1 ($\eta_{max}$) with Cosine Annealing scheduler and train
them for 800 epochs ($t_{max}$). For our accelerated setting we use the F1-CLR
schedule with learning rate of 0.2 ($\eta_{max}$) with 80 epochs for warm up
($t_{w}$) and train for 480 epochs ($t_{max}$). For Hard Augment we use 6
augmentations generated using SimCLR[6] transforms and scale the colour
jittering linearly where the standard scale is defined as 5. For Super
Progressive learning we use 96 as the minimum resolution and use 6 stages each
a multiple of 32. Additional details can be seen in Supplementary Sec. B.

We have observed that GPU pre-processing plays an important role in obtaining
speedups in real time. Since loading images and making multiple augmentations
can become a bottleneck when training with faster GPU cards like the Nvidia
V100. We use the Nvidia DALI222available at https://github.com/NVIDIA/DALI
library for fast loading and augmentation of the images in the accelerators.

###  4.1 Backbone Architecture Analysis

In order to show the generality of our method to different different backbone
architectures we applied our training strategies to SimSiam [9] model on the
ImageNet dataset and report linear evaluation results and speed-up in Tab. 1.

Our results show that our method can accelerate both ResNet-18 and ResNet-50
architectures by 2.3 times while maintaining the baseline accuracy.

###  4.2 Model Experiments

We have evaluated our acceleration strategies on different SSL models in Tab.
2 by pre-taining on a ResNet 50 backbone on the Imagenette dataset. We used
the default hyperparameters for SimSiam[9], BYOL[22], MoCov2+ [8], DINO [4]
and SimCLR [6] models while we used the same optimizer and augmentation hyper-
parameters for all of them which can be seen in the supplementary Sec. B.

LR schedule | Sup. Prog. | Hard Aug. | kNN-Acc.(%) | Speedup
---|---|---|---|---
CA | - | - | 90.9 | 1.0x
F1-CLR | - | - | 90.2 (-.7) | 1.7x
CA | - | ✓ | 89.8 (-1.1) | 1.8x
F1-CLR | ✓ | - | 90.1 (-.8) | 2.7x
F1-CLR | - | ✓ | 90.3 (-.6) | 1.5x
F1-CLR | ✓ | ✓ | 91.4 (+.5) | 2.4x
Table 3: Comparing the contribution of each component in our experimental
setup. Speed up is reported based on the reduction in training cost. Using
SimSiam model with ResNet50 backbone architecture on Imagenette dataset. CA
stands for the baseline Cosine Annealing learning rate schedule.

We have observed that that our method can accelerate the training of different
models. We note that methods with a momentum encoder typically have a higher
accuracy and training cost even when trained with the same number of steps due
to the additional forward passes made with the momentum encoder which helps to
stabilize the training. Our method is especially effective for models with a
momentum encoder where we can accelerate the training up to 2.7 times.

###  4.3 Experimental Setup

Here we compare the contribution of each component in our experimental setup
in Tab. 3 training with ResNet50 backbone on the Imagenette dataset using
default hyperparameters whenever possible. We have observed that our
improvements are generally compatible with each other and allow us to train
close to the same accuracy while reducing the training cost by 2.4 times.

While Super Progressive Learning is the main driver in reducing computational
cost, Hard Augment allows us to improve performance dramatically by
prioritizing useful augmentations. We note that Super Progressive Learning and
Hard Augment also work well energetically since Hard Augment can select from a
harder set of alternatives towards the end of the training which further boost
performance.

###  4.4 Ablations

In the ablation studies use the Imagenette dataset with the parameters
described in the previous section. We aim to identify various parameters of
our model and examine its behaviour in different circumstances.

#### Learning Rate Schedule Comparison

We compared different learning rate schedulers in isolation in order to see
what percentage of improvement is attributable to the different learning rate
schedulers on the ImageNet dataset in Tab. 4. We have seen that our Extended
Super Convergence method still gives a reliable performance increase in this
setup compared to Cosine Annealing and Cosine Annealing with linear warm up.

Schedule | Warm-Up | Epochs | Acc.(%)
---|---|---|---
Cosine Decay | - | 200 | 67.7
Cosine Decay | - | 100 | 66.1
Cosine Decay | Linear-10 | 100 | 66.0
F1-CLR (3.1) | F1-CLR (3.1) | 100 | 66.4
Table 4: Comparing different learning rate schedules. Cosine decay with linear
warm-up and Super convergence both have 10 epochs of warm-up in order to make
the comparison more direct.

#### F1-CLR Warm Up

Here we make an ablation study where we train the SimSiam model for 320 epochs
with different number of warm up epoch lengths in Tab. 5. We observe that 80
epochs is the optimal warm up length and we use this setting in our future
experiments and ablations that utilize F1-CLR schedule.

Warm Up | 32 | 48 | 64 | 80 | 96 | 112
---|---|---|---|---|---|---
Acc(%) | 86.8 | 86.5 | 86.9 | 87.1 | 86.8 | 86.6
Table 5: Super Convergence training where the number of warm up epochs is
changed.

#### Minimum Resolution

An important hyper-parameter that determines how much speed-up can be achieved
with Super Progressive Learning is the minimum resolution. We made a ablation
study in Tab. 6 to see the maximum speedup that can be achieved with our
training strategies by increasing the minimum resolution by 32 increments
starting from 64. The smallest resolution that does not result in a drop in
performance is 96 which we used in our other experiments while 64 was giving
better speedups it wasn’t maintaining the level of performance we are seeking.

Min. Res. | Speed up | Acc.(%)
---|---|---
128 | 1.40x | 86.6
96 | 1.60x | 86.6
64 | 1.82x | 86.1
0 | 3x | N/A
Table 6: Ablation study for the Super Progressive learning speedup and
accuracy trade-off with different values for minimum resolution using 224 as
the maximum resolution.

#### F1-CLR Training Length

Baseline | Super Convergence
---|---
Epochs | Acc.(%) | Speedup | Acc.(%)
160 | 68.8 | 1.61x | 84.0
200 | 77.9 | 1.53x | 86.1
240 | 82.7 | 1.37x | 86.6
280 | 85.6 | 1.37x | 87.6
320 | 86.4 | 1.42x | 88.1
360 | 87.3 | 1.41x | 88.4
400 | 87.8 | 1.59x | 89.1

Table 7: (left):Online validation accuracy of training instances of different
lengths that are trained with F1-CLR (right): Total length of training given
the fixed duration for warm-up of 80 epochs. Speedups are calculated based on
linearly interpolated training length in the baseline given that the baselines
has at 800 epochs has %90.0 accuracy

To build a better understanding of 1CLR schedule we change the length of the
training from 160 epochs to 800 epochs using the cosine annealing learning
rate schedule and compare against how the proposed extended super convergence
schedule performs. Tab. 7 shows that the difference in performance is stark
when the training duration is short and gradually becomes smaller.
Additionally, our F1-CLR setting leads to between 1.3x-1.6x speedup of the
training to reach equivalent performance on the baseline.

#### Number of Positives

In order to understand the effect of number of augmentations we have made an
ablation study that examines the trade-off between increased overhead and the
effect on accuracy in Tab. 8. We have used the SimSiam [9] model with F1-CLR
schedule and our default hyperparameters. We have observed that 6 positives
gives good balance between maximizing accuracy while minimizing the overhead
so we use this setting in other experiments.

# of Positives | 4 | 5 | 6 | 8
---|---|---|---|---
Accuracy(%) | 89.1 | 89.0 | 89.6 | 89.8
Overhead(%) | 5.4 | 6.8 | 8.2 | 10.9
Table 8: Ablation study for the numbers of positives used in Hard Augment. 6
augmentations result in good performance increase while maintaining reasonable
overhead.

#### Multi-Crop Comparison

Multi-crop augmentation proposed by SwAV [3] uses multiple augmentations with
a combination of high resolution and low resolution crops where only high
resolution crops are used for back-propagation and low resolution ones are
used only as extra comparison targets. A slight modification to our Hard
Augmentation method will produce an intermediate algorithm useful for a
comparison where the additional low resolution augmentations that do not have
the highest loss can be used as extra targets. We made an ablation study that
compares the multi-crop augmentation strategy at the last line of Tab. 9
against our method and our method with extra targets. We train our method with
Extended Super Convergence for 400 epochs and used 4 augmentations for each
image.

Min Res. | Selection | Extra Target | Acc(%)
---|---|---|---
64 | Hard |  | 89.1
64 | Hard | ✓ | 88.8
96 | Hard | ✓ | 89.7
96 | Random | ✓ | 88.4
Table 9: Progressive learning speedup with different values for minimum to
maximum resolution fraction

We have seen that extra targets are useful when 96 resolution is used for the
selection pass and hard augment with extra targets and hard augment
outperforms multi-crop augmentation. We have also seen that the standard
resolution of 64 that is enough for ranking loss pairs in Hard Augment is not
enough for generating accurate targets.

Since increasing selection resolution has a direct effect on the overhead for
Hard Augment and the additional improvement that can be obtained from extra
targets is small we have decided not to include extra targets in our method.

##  5 Conclusion

We have demonstrated that self-supervised training can be accelerated by
adapting the learning rate, augmentation and resolution schedules for self
supervised training and boosting the training signal by hard positive mining
on the augmentations. Our method have shows an training speed-up between 2.3
and 2.7 times. which allows a much wider community to reproduce and contribute
to the self-supervised learning literature, reduce the financial and carbon
cost of training these models. However there is a risk that using efficient
training methods the community will adopt larger and computationally more
expensive benchmarks which will eliminate some of the intended benefits of our
method. As future work we aim to apply our method to recent methods [4, 10]
that used Vision Transformer [15] backbones.

## References

  * [1] Shekoofeh Azizi, Basil Mustafa, Fiona Ryan, Zachary Beaver, Jan Freyberg, Jonathan Deaton, Aaron Loh, Alan Karthikesalingam, Simon Kornblith, Ting Chen, et al.  Big self-supervised models advance medical image classification.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 3478–3488, 2021.
  * [2] Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston.  Curriculum learning.  In Proceedings of the 26th annual international conference on machine learning, pages 41–48, 2009.
  * [3] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin.  Unsupervised learning of visual features by contrasting cluster assignments.  arXiv preprint arXiv:2006.09882, 2020.
  * [4] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin.  Emerging properties in self-supervised vision transformers.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9650–9660, 2021.
  * [5] Pratik Chaudhari and Stefano Soatto.  Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks.  In 2018 Information Theory and Applications Workshop (ITA), pages 1–10. IEEE, 2018.
  * [6] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  In International conference on machine learning, pages 1597–1607. PMLR, 2020.
  * [7] Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey E Hinton.  Big self-supervised models are strong semi-supervised learners.  Advances in neural information processing systems, 33:22243–22255, 2020.
  * [8] Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He.  Improved baselines with momentum contrastive learning.  arXiv preprint arXiv:2003.04297, 2020.
  * [9] Xinlei Chen and Kaiming He.  Exploring simple siamese representation learning. corr abs/2011.10566 (2020).  arXiv preprint arXiv:2011.10566, 2020.
  * [10] Xinlei Chen, Saining Xie, and Kaiming He.  An empirical study of training self-supervised vision transformers.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9640–9649, 2021.
  * [11] Cody Coleman, Deepak Narayanan, Daniel Kang, Tian Zhao, Jian Zhang, Luigi Nardi, Peter Bailis, Kunle Olukotun, Chris Ré, and Matei Zaharia.  Dawnbench: An end-to-end deep learning benchmark and competition.  Training, 100(101):102, 2017.
  * [12] Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le.  Randaugment: Practical automated data augmentation with a reduced search space. arxiv e-prints, page.  arXiv preprint arXiv:1909.13719, 2019.
  * [13] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee, 2009.
  * [14] Carl Doersch, Abhinav Gupta, and Alexei A Efros.  Unsupervised visual representation learning by context prediction.  In Proceedings of the IEEE international conference on computer vision, pages 1422–1430, 2015.
  * [15] 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.
  * [16] William Falcon et al.  Pytorch lightning.  GitHub. Note: https://github. com/PyTorchLightning/pytorch-lightning, 3:6, 2019.
  * [17] Kunihiko Fukushima and Sei Miyake.  Neocognitron: A self-organizing neural network model for a mechanism of visual pattern recognition.  In Competition and cooperation in neural nets, pages 267–285. Springer, 1982.
  * [18] Spyros Gidaris, Praveer Singh, and Nikos Komodakis.  Unsupervised representation learning by predicting image rotations.  arXiv preprint arXiv:1803.07728, 2018.
  * [19] Chengyue Gong, Tongzheng Ren, Mao Ye, and Qiang Liu.  Maxup: A simple way to improve generalization of neural network training.  arXiv preprint arXiv:2002.09024, 2020.
  * [20] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  Advances in neural information processing systems, 27, 2014.
  * [21] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He.  Accurate, large minibatch sgd: Training imagenet in 1 hour.  arXiv preprint arXiv:1706.02677, 2017.
  * [22] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al.  Bootstrap your own latent: A new approach to self-supervised learning.  arXiv preprint arXiv:2006.07733, 2020.
  * [23] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.  Momentum contrast for unsupervised visual representation learning.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9729–9738, 2020.
  * [24] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition. arxiv 2015.  arXiv preprint arXiv:1512.03385, 2015.
  * [25] Geoffrey E Hinton, Simon Osindero, and Yee-Whye Teh.  A fast learning algorithm for deep belief nets.  Neural computation, 18(7):1527–1554, 2006.
  * [26] Elad Hoffer, Berry Weinstein, Itay Hubara, Tal Ben-Nun, Torsten Hoefler, and Daniel Soudry.  Mix & match: training convnets with mixed image sizes for improved accuracy, speed and scale resiliency.  arXiv preprint arXiv:1908.08986, 2019.
  * [27] Jeremy Howard.  Training imagenet in 3 hours for usd 25; and cifar10 for usd 0.26, 2018\.
  * [28] Stanisław Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey.  Three factors influencing minima in sgd.  arXiv preprint arXiv:1711.04623, 2017.
  * [29] Angela H Jiang, Daniel L-K Wong, Giulio Zhou, David G Andersen, Jeffrey Dean, Gregory R Ganger, Gauri Joshi, Michael Kaminksy, Michael Kozuch, Zachary C Lipton, et al.  Accelerating deep learning by focusing on the biggest losers.  arXiv preprint arXiv:1910.00762, 2019.
  * [30] Angelos Katharopoulos and François Fleuret.  Not all samples are created equal: Deep learning with importance sampling.  In International conference on machine learning, pages 2525–2534. PMLR, 2018.
  * [31] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.  On large-batch training for deep learning: Generalization gap and sharp minima.  arXiv preprint arXiv:1609.04836, 2016.
  * [32] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  arXiv preprint arXiv:1312.6114, 2013.
  * [33] Gustav Larsson, Michael Maire, and Gregory Shakhnarovich.  Colorization as a proxy task for visual understanding.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 6874–6883, 2017.
  * [34] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari.  The large learning rate phase of deep learning: the catapult mechanism.  arXiv preprint arXiv:2003.02218, 2020.
  * [35] Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al.  Pytorch distributed: Experiences on accelerating data parallel training.  arXiv preprint arXiv:2006.15704, 2020.
  * [36] Mehdi Noroozi and Paolo Favaro.  Unsupervised learning of visual representations by solving jigsaw puzzles.  In European conference on computer vision, pages 69–84. Springer, 2016.
  * [37] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In International Conference on Machine Learning, pages 8748–8763. PMLR, 2021.
  * [38] Abhinav Shrivastava, Abhinav Gupta, and Ross Girshick.  Training region-based object detectors with online hard example mining.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 761–769, 2016.
  * [39] LN Smith.  Cyclical learning rates for training neural networks. arxiv.  Preprint at https://arxiv. org/abs/1506.01186, 2015.
  * [40] Leslie N Smith.  A disciplined approach to neural network hyper-parameters: Part 1–learning rate, batch size, momentum, and weight decay.  arXiv preprint arXiv:1803.09820, 2018.
  * [41] Leslie N Smith and Nicholay Topin.  Super-convergence: Very fast training of neural networks using large learning rates. arxiv e-prints, page.  arXiv preprint arXiv:1708.07120, 2017.
  * [42] Samuel L Smith, Pieter-Jan Kindermans, Chris Ying, and Quoc V Le.  Don’t decay the learning rate, increase the batch size.  arXiv preprint arXiv:1711.00489, 2017.
  * [43] Samuel L Smith and Quoc V Le.  A bayesian perspective on generalization and stochastic gradient descent.  arXiv preprint arXiv:1710.06451, 2017.
  * [44] Mingxing Tan and Quoc V Le.  Efficientnetv2: Smaller models and faster training.  arXiv preprint arXiv:2104.00298, 2021.
  * [45] Yonglong Tian, Dilip Krishnan, and Phillip Isola.  Contrastive multiview coding.  In European conference on computer vision, pages 776–794. Springer, 2020.
  * [46] Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola.  What makes for good views for contrastive learning?  Advances in Neural Information Processing Systems, 33:6827–6839, 2020.
  * [47] Hugo Touvron, Andrea Vedaldi, Matthijs Douze, and Hervé Jégou.  Fixing the train-test resolution discrepancy: Fixefficientnet.  arXiv preprint arXiv:2003.08237, 2020.
  * [48] Laurenz Wiskott and Terrence J Sejnowski.  Slow feature analysis: Unsupervised learning of invariances.  Neural computation, 14(4):715–770, 2002.
  * [49] Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin.  Unsupervised feature learning via non-parametric instance discrimination.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3733–3742, 2018.
  * [50] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny.  Barlow twins: Self-supervised learning via redundancy reduction.  In International Conference on Machine Learning, pages 12310–12320. PMLR, 2021.
  * [51] Richard Zhang, Phillip Isola, and Alexei A Efros.  Colorful image colorization.  In European conference on computer vision, pages 649–666. Springer, 2016.
  * [52] Richard Zhang, Phillip Isola, and Alexei A Efros.  Split-brain autoencoders: Unsupervised learning by cross-channel prediction.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 1058–1067, 2017.

##  Appendix A Speed Up Calculation

While Super Progressive Learning reduces the number of floating point
calculations made by using a smaller resolution, Hard Augment adds an
additional overhead for training which can be calculated by equations 9 and 11
respectively. Generally we find that we can reduce the number of steps by 1.7
times while reduce the number of floating point operations at each step by 1.4
times.

##  Appendix B Additional Implementation Details

| Imagenette | ImageNet
---|---|---
Hyper-parameter | Baseline | Efficient | Baseline | Efficient
Number of classes | 10 | 10 | 1000 | 1000
Batch size | 128 | 128 | 512 | 512
Weight decay | $5\text{\times}{10}^{-4}$ | $5\text{\times}{10}^{-4}$ | $1\text{\times}{10}^{-4}$ | $1\text{\times}{10}^{-4}$
Learning rate | 0.05 | 0.1 | 0.1 | 0.16
LR schedule | CA | ESC | CA | ESC
Warmup epochs | 0 | 80 | 10 | 10
Momentum weight | 0.9 | 0.85-0.95 | 0.9 | 0.85-0.95
Learning rate | 0.05 | 0.05 | 0.05 | 0.16
Min aug. magnitude | 5 | 4 | 5 | 4
Max aug. magnitude | 5 | 6 | 5 | 6
Min view resolution | 224 | 96 | 224 | 96
Number of positives | 2 | 6 | 2 | 6
Selection resolution | N/A | 64 | N/A | 64
Table 10: Hyper parameters used in various setups. CA: Cosan Annealing, ESC:
Extended Super Convergence

We have trained all our experiments in the single node setting with maximum 8
GPUs which restricted us to use 512 as the maximum batch size. Some of the SSL
methods like SimCLR [6] and BYOL [22] have been shown to perform better with
larger batch sizes which we could not replicate and restricted ourselves to
the same training setting in all our experiments.

In some of the ablations and analysis experiments where we study the effect of
a certain parameter different values can be used. For Linear evaluation we use
freeze the encoder and re-initialize the last layer and train with batch size
2048 and learning rate 0.8 for 90 epochs using the LARS optimizer.

Our implementation is based on PyTorch Lightning [16] where we define Hard
Augmentation and Super Progressive Learning Schedule as callbacks. We adapt
the OneCycle Learning rate schedule
implementation333https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.OneCycleLR.html
for Extended Super Convergence. We will provide the code for our method upon
publication. See the table for a full list of parameters Tab. 10.

##  Appendix C Learning rate range finder

We have used the Learning rate range finder test proposed by [41] to find good
values for minimum and maximum learning rates. This test increase learning
rate exponentially and keeps track of on the training and validation loss. The
minimum value and maximum value are determined by the point at which the
validation loss starts to decrease and the point at which is starts to
diverge. We calculate the validation loss at each step on a batch of 4096
validation samples in order to keep the test manageable.

In the test shown in Figure 4 we increase the learning rate from
$1\text{\times}{10}^{-3}$ to $1$ in 200 steps. Both the training and
validation losses plateau close to step 135 and learning rate .1 which we use
in our experiments on Imagenette dataset.

Figure 4: Learning rate finder plot for training (left) and validation (right)
losses

##  Appendix D Augmentation Resolution Relationship

In order to examine the relationship between resolution and augmentation
magnitude for self-supervised learning we have conducted a series of
experiments. We trained a BYOL [22] model on the Imagenette dataset using the
ResNet18 [24] architecture with RandAugment [12] augmentation with the cosine
annealing learning rate schedule where we change the augmentation strength $m$
while we change the input image resolution. Each entry in the table is a
separate experiment with fixed input resolution and augmentation magnitude.
The results in Table 11 confirms that we can apply higher magnitude
augmentations only on higher resolution images and the optimum augmentation
magnitude increases with resolution. This confirms our intuition behind the
relationship between resolution and augmentation magnitude in SSL and provides
motivation for applying Super Progressive Learning and allows us to confirm
the findings of [44] for SSL.

| m=3 | m=5 | m=7 | m=10 | m=15
---|---|---|---|---|---
128 | 85.0 | 84.0 | 81.4 | 77.7 | 75.3
192 | 88.8 | 87.7 | 88.9 | 86.8 | 87.3
300 | 89.5 | 89.7 | 89.9 | 88.3 | 87.7
Table 11: Resolution and Augmentation Magnitude relationship shown by training
a self-supervised learning method on the combination of resolution and
magnitudes and measuring the online classification accuracy of a linear layer

##  Appendix E Progressive Augmentation Curriculum

However what values should be used as the minimum and maximum augmentation
magnitude is an important question. Since we use the linearly scaled SimCLR
[6] augmentations in our method we made a hyper-parameter search on these
values in order to determine augmentation magnitudes empirically. We train our
method for 320 epochs with Super Progressive learning schedule using 128 as
the minimum resolution for various augmentation magnitudes and compare against
the default augmentation magnitude of 5. Our experiment shows that a minimum
value of 4 and maximum value of 6 perform better in our setting and outperform
the fixed augmentation setting.

$m_{min}$ | $m_{max}$ | Acc(%)
---|---|---
5 | 5 | 88.6
2.5 | 4 | 88.4
3 | 4 | 88.7
4 | 5 | 88.7
5 | 6 | 88.7
4 | 6 | 88.9
Table 12: Maximum and minimum augmentation magnitude when trained with Super
Progressive learning

##  Appendix F Future Work

This idea can be extended to other modalities where resolution is defined in
different manners for example on sound with the sampling rate analogue which
in a similar sense accelerates the training however than the amount of
acceleration and cost benefit calculation will change. A similar case can also
be made for tabular data where the most important features are processed first
and later additional features are added however this is much more difficult to
implement would probably require additional structure. Similarly training with
a small vocabulary and than enlarging that can have a similar effect as well.


# Your Instructions Are Not Always Helpful: Assessing the Efficacy of
Instruction Fine-tuning for Software Vulnerability Detection

Imam Nur Bani Yusuf
Singapore Management University
<EMAIL_ADDRESS>
&Lingxiao Jiang
Singapore Management University
<EMAIL_ADDRESS>

###### Abstract

Software, while beneficial, poses potential cybersecurity risks due to
inherent vulnerabilities. Detecting these vulnerabilities is crucial, and deep
learning has shown promise as an effective tool for this task due to its
ability to perform well without extensive feature engineering. However, a
challenge in deploying deep learning for vulnerability detection is the
limited availability of training data. Recent research highlights the deep
learning efficacy in diverse tasks. This success is attributed to instruction
fine-tuning, a technique that remains under-explored in the context of
vulnerability detection. This paper investigates the capability of models,
specifically a recent language model, to generalize beyond the programming
languages used in their training data. It also examines the role of natural
language instructions in enhancing this generalization. Our study evaluates
the model performance on a real-world dataset to predict vulnerable code. We
present key insights and lessons learned, contributing to understanding the
deep learning application in software vulnerability detection.

##  1 Introduction

Software is integral to various technologies, such as smartphones and
e-learning systems. Although software brings many benefits, it presents
potential attack surfaces for cyber-attacks due to vulnerabilities. These
vulnerabilities, in extreme cases, can lead to significant financial losses.
For instance, Revolut experienced a $20 million loss when attackers exploited
weaknesses in its payment system.111https://thehackernews.com/2023/07/hackers-
steal-20-million-by-exploiting.html Consequently, the detection of software
vulnerabilities is a critical task.

Program analysis techniques, such as static analysis and dynamic analysis,
have been used to detect vulnerabilities in code Zhang et al. (2018, 2019);
Zhan et al. (2021). However, these techniques suffer from false positives
(i.e., detect benign code as vulnerable) or false negatives (i.e., miss many
vulnerabilities) Chakraborty et al. (2022). Machine learning has been proposed
as an alternative Narayanan et al. (2014); Liu et al. (2015); Ma et al.
(2016), but it requires extensive feature engineering. Recently, deep learning
has emerged as a promising technique due to its good performance without the
need for extensive feature engineering Li et al. (2018); Fu et al. (2022);
Chakraborty et al. (2022).

The training of deep learning-based vulnerability detection tools faces a
significant challenge due to limited data availability. Typically, the
instances of vulnerable code are much less frequent than non-vulnerable code.
This disparity can lead to a bias towards the majority class, potentially
causing the model to overlook the minority class, which in this context is the
vulnerable code Johnson and Khoshgoftaar (2019).

Despite this challenge, recent researches have demonstrated the effectiveness
of deep learning in various tasks, including those not encountered during
training. This success is largely attributed to instruction fine-tuning, a
technique that trains models using task descriptions and corresponding input-
label pairs Wei et al. (2022a); Mishra et al. (2022); Honovich et al. (2023);
Wang et al. (2023a). This method leverages task descriptions as a guiding
framework for task execution. During fine-tuning, the model learns to match a
given task description with specific inputs to generate appropriate outputs.
Consequently, the model can infer the correct output based on the provided
instruction and input.

Despite its potential, the application of instruction fine-tuning in
vulnerability detection remains largely unexplored. This paper aims to
investigate how well models can generalize beyond the programming languages
present in their training data and examine the impact of natural language
instructions on improving generalization in the vulnerability detection task.

We have conducted the study by evaluating a recent language model to predict
vulnerable code from the real-world dataset. Our study yields three insights.
First, models perform more effectively in scenarios where the language is the
same as in the training data (intra-lingual) compared to different languages
(cross-lingual), although the difference in performance is not substantial.
Second, models that do not use natural language instructions outperform those
that do in both intra-lingual and cross-lingual settings. Third, when multiple
programming languages are added to the training set without natural language
instructions, there’s a decline in model performance. Conversely, when natural
language instructions are combined with multiple programming languages in the
fine-tuning process, the models show better performance.

The contributions of this paper is as follows.

  1. 1.

We are the first to study how well models can generalize beyond the
programming languages present in their training data in the vulnerability
detection task.

  2. 2.

We are the first to study the impact of natural language instructions on
improving generalization in the vulnerability detection task.

The structure of this paper is organized as follows: Section 2 outlines the
Research Questions (RQs). Section 3 describes the process of dataset
selection. Section 4 elaborates the methodology used for the model selection.
Section 5 outlines the experimental settings. Section 6 presents the results.
Section 7 discusses the manual case study, lessons learned, and potential
threats to validity. Finally, Section 8 concludes the paper.

##  2 Research Questions

We frame the vulnerability prediction task as a binary classification problem.
The classification decision can be formalized as Equation 1.

$v=\arg\max_{v}P(v\mid c)$ (1)

The model receives a code snippet $c$ as the input. The objective of the model
is to make prediction $v$, where $v$ $\in$ (vulnerable, non-vulnerable). The
prediction is based on the conditional probability $P(v|c)$, which represents
the likelihood of the code being vulnerable given the code snippet $c$.

Based on the problem formulation above, this study aims to examine model
generalization beyond the languages in the fine-tuning data and the role of
natural language instructions in improving generalization performance. We
define the following Research Questions (RQs).

RQ1: How does the model perform in the intra-lingual and cross-lingual
settings? The intra-lingual setting tests the model on the same programming
language used during fine-tuning, providing a baseline for performance. The
result in the intra-lingual setting will serve a baseline for comparing with
the cross-lingual setting. In contrast to the intra-lingual setting, the
cross-lingual evaluation tests the model with programming languages different
from the fine-tuning language. This question aims to assess the model’s
ability to generalize across various programming languages.

RQ2: How do natural language instructions influence model performance in the
intra-lingual and cross-lingual settings? This question assesses whether
incorporating natural language instructions enhances the model’s vulnerability
prediction capabilities within the intra-lingual and cross-lingual settings.

RQ3: How does the diversity of programming languages in the fine-tuning
dataset affect performance in the intra-lingual and cross-lingual settings?
This question investigates the influence of the programming language diversity
in the fine-tuning dataset on the model performance to predict vulnerabilities
in the intra-lingual and cross-lingual settings.

##  3 Dataset

This section describes the dataset selection, statistics, and preprocessing
steps.

###  3.1 Dataset Selection

A number of datasets are available in the domain of vulnerability detection,
including SARD222https://samate.nist.gov/SARD/, VulDeePecker Li et al. (2018),
Draper Russell et al. (2018), D2A Zheng et al. (2021), Devign Zhou et al.
(2019), BigVul Fan et al. (2020), Reveal Chakraborty et al. (2022), DiverseVul
Chen et al. (2023), CrossVul Nikitopoulos et al. (2021), and CVEFixes Bhandari
et al. (2021).

VulDeePecker, SARD, BigVul, Draper, D2A, Devign, BigVul, Reveal, and
DiverseVul are limited to two programming languages, which are C and C++.
Moreover, Draper and D2A are annotated by static analyzers. Hence, Draper and
D2A may contain false positive and false negative cases. SARD covers five
programming languages. However, SARD is a synthetic dataset and may not
reflect real-world vulnerabilities. Similarly, CrossVul and CVEFixes also
contain at least five programming languages. However, these datasets reflect
real-world vulnerabilities better than SARD because both are curated from
real-world security issues. Nevertheless, we choose CVEFixes because it
contains a broader range of vulnerability types (i.e., CWE Id) than CrossVul,
i.e., 193 vs 167 types.

###  3.2 Dataset Preprocessing

The initial dataset contains 116,057 instances. We exclude instances with more
than 990 code tokens to adhere to the maximum context length of 1024 in our
experiment, due to the memory constraint. The first exclusion step results in
67,692 instances. Then, we only include dominant languages: C, PHP, and
Python. We choose these three languages because they represent different
programming use cases and have distinct syntax. Then, we sample randomly 2,750
instances from each language for fine-tuning set and 250 for testing set. We
choose this number such that the proportion between the vulnerable and non-
vulnerable instances are the same across different programming languages
between the fine-tuning and testing sets. In the end, we have 8,250 instances
for fine-tuning and 750 instances for testing. Table 1 show the statistics for
fine-tuning and testing sets. We will use this dataset for the subsequent
experiments.

Statistics | Fine-tuning | Testing
---|---|---
#Instances | 8250 | 750
Average #tokens | 271.80 | 268.12
#Vulnerable | 4110 | 390
#Non-vulnerable | 4140 | 360
#C | 2750 | 250
#PHP | 2750 | 250
#Python | 2750 | 250
Table 1: Dataset statistics used in our study. "#" represents "number of."

##  4 Model Selection

Before conducting the experiments to answer the RQs, we perform a preliminary
experiment to select the model with a strong base performance for the
vulnerability detection task. We explain the steps as follows.

###  4.1 Selecting Model Candidates

Models are evaluated against 3 criteria: they must be open-source for custom
data fine-tuning, capable of being fine-tuned on a single 24GB memory GPU, and
not previously instruction-tuned to ensure performance improvements are due
solely to our fine-tuning efforts. From the first filtering, five models are
considered: CodeLlama (7B) Rozière et al. (2023), CodeT5+ (770M) Wang et al.
(2023b), CodeT5 (770M) Wang et al. (2021), Llama (7B) Touvron et al. (2023),
Mistral (7B) Jiang et al. (2023), and Yi (6B)333https://github.com/01-ai/Yi,
CodeLlama, CodeT5+, and CodeT5 are chosen because they are specifically
pretrained for code-related tasks. Conversely, Llama 2, Mistral and Yi are
chosen as representatives of models pre-trained on general corpora. The
selection excludes the larger variants of CodeT5+ due to the lack of
optimization with Flash Attention 2 Dao (2023), which restricts fine-tuning on
our GPU.

###  4.2 Preliminary Experimental Setting

We fine-tune and evaluate the model using all languages in our dataset, which
includes C, PHP, and Python, to identify a model with a strong base
performance. We leverage accuracy as the evaluation metric. We compute
accuracy by calculating the ratio of correctly predicted instances to the
total number of predictions made. This metric is chosen because it reflects
the model’s ability to accurately predict both vulnerable and non-vulnerable
instances. Ideally, the base model should demonstrates high performance in
both classes. Note that the testing set has been balanced to contain an equal
proportion of vulnerable and non-vulnerable instances, ensuring that accuracy
reflects an unbiased assessment of the model’s predictive capabilities.

###  4.3 Preliminary Experiment Result

Model | Accuracy
---|---
CodeLlama | 0.535
CodeT5+ | 0.520
CodeT5 | 0.480
Llama 2 | 0.511
Mistral | 0.531
Yi | 0.480
Table 2: The preliminary experiment result. CodeLlama performs the best among
the others.

Table 2 reveals the results of the preliminary experiments. CodeLlama performs
the best among the other models with accuracy 0.535. Therefore, we select
CodeLlama as the base model to answer our Research Questions (RQs) in the
subsequent experiments.

##  5 Study Experimental Setting

This section discusses the pipeline in the fine-tuning and inference stages.
Then, this section elaborates the implementation details, including the model
implementation, training setting, and how to generate the natural language
instructions.

###  5.1 Fine-tuning

Let $c$ represent the input and $v$ its corresponding label in the fine-tuning
process. A transformation is applied using a template function $Q$, which
takes $c$ and a label $l$ as inputs and produces an output $w$. This output is
denoted as $s=Q(c,l)$. The function $Q$ is responsible for formatting $c$ and
$l$ into a specific template. The choice of template depends on whether the
experiment incorporates natural language instructions. The details of this
template are depicted in Figure 1.

with instruction
<s> [INST] {insert instruction} [/INST]
### Code: {insert code}
### Answer: {insert label} </s>
without instruction
### Code: {insert code}
### Answer: {insert label} </s> Figure 1: The distinct formats for the
template function $Q(c,l)$, showing variations with and without the inclusion
of natural language instructions.

Given a sequence $s$, a function $F$ generates a set of input-label pairs,
$F(s)=\\{((s_{1},\dots,s_{j}),s_{j+1})\mid 1\leq j<n\\}$, where each pair
consists of a subsequence of $s$ and the following word as the label. These
pairs are then processed in a batch where a tokenizer function $\tau$ converts
each item in the batch into a sequence of token ids $t=\tau(w)$. An embedding
function $\epsilon$ maps each token id $t_{i}\in t$ to a sequence of vectors
$t^{\prime}=\epsilon(t)$. The transformer architecture (either encoder or
decoder, depending on the model), represented as a function $\phi$, processes
$t^{\prime}$ to produce a probability distribution over the vocabulary for
each token, denoted as $P=\phi(t^{\prime})$. The prediction of the next token
$\hat{t}_{i+1}$ is determined by $\hat{t}_{i+1}=\operatorname{argmax}(P)$. A
loss function $L$, i.e., Cross Entropy Cox (1958), is used to compute the
loss. It takes as input the predicted token $\hat{t}_{i+1}$ and the ground
truth next token $t_{i+1}$, yielding the loss value $L(\hat{t}_{i+1},t_{i+1})$
to optimize the model.

###  5.2 Inference

Given a code snippet $c$, the transformation function $Q$ formats the code
snippet $c$ into a specific format, resulting in $w=Q(c)$. This process
differs from the training stage, as the label $v$ is not included in the input
to $Q$. Next, the tokenizer function processes $w$ to produce a sequence of
token ids $t=\tau(w)$. The embedding layer converts these token ids into
vectors $t^{\prime}=\epsilon(t)$. The transformer block takes these vectors
and computes a probability distribution $P=\phi(t^{\prime})$ over the entire
vocabulary. Notably, the token with the highest probability,
$\operatorname{argmax}(P)$, may not match the target labels ’True’ or ’False’.
To address this, we only extract the probabilities of the ’True’ and ’False’
tokens, i.e., $P_{\text{True}}$ and $P_{\text{False}}$. The final output is
the token with the greater probability, determined by
$\text{max}(P_{\text{True}},P_{\text{False}})$.

###  5.3 Evaluation Metrics

We choose the F1 score as the evaluation metrics because it has been widely
used to evaluate vulnerability detection tools Fu and Tantithamthavorn (2022);
Steenhoek et al. (2023, 2022). The F1 score considers both the precision and
the recall of the test to compute the score. Precision is the number of
correct positive results divided by the number of all positive results, and
recall is the number of correct positive results divided by the number of
positives that should have been returned. Here, positive results refers to
vulnerable code. The F1 score is the harmonic mean of precision and recall,
giving both measures equal weight.

###  5.4 Implementation Details

####  5.4.1 Models and Training Setting

For all models, we use the official implementation from HuggingFace. The link
to the model can be found in the Appendix.

We fine-tune the model using the LoRA technique Hu et al. (2022) with $r=16$
and $\alpha=32$, where $r$ is chosen as the largest value that does not
trigger an out-of-memory error and $\alpha$ is double the value of $r$, as
recommended. To accommodate the model within our hardware limitations, we
implement several optimizations: loading the model with 4-bit precision,
limiting the maximum context length to 1024 tokens, utilizing Flash Attention
2, and applying gradient checkpointing. The learning rate is set to $3\times
10^{-4}$ and the batch size is chosen as the largest multiple of 2 that fits
our machine, i.e., $8$. We employ a cosine learning rate scheduler with warmup
steps of 100 to mitigate unstable learning that may result from arbitrary
initialization of optimizer states. We conduct our fine-tuning process on a
machine with Linux 22.04.2 and an RTX 3090 GPU with 24GB of memory. We perform
all experiments three times with three different random seeds to avoid a lucky
trial.

####  5.4.2 Generating and Choosing Instructions

Initially, we create a set of instructions manually. These initial
instructions are then refined using GPT-4, yielding 15 refined versions. From
these, we select 5 instructions for our final instruction pool. This selection
is based on evaluating the quality and diversity through human judgment. The
rationale for limiting to 5 instructions is to prevent an excessive number
that might impede the model’s learning ability, potentially treating the
instructions as noise. During the experiment using natural language
instructions (RQ3), the template function $Q$ chooses one instruction from the
instruction pool using random sampling.

##  6 Experimental Results

Figure 2 illustrates the model’s performance in the intra-lingual and cross-
lingual settings. In the intra-lingual context, the C$\rightarrow$C model
(i.e., the model fine-tuned and tested on C) exhibits the highest performance,
closely followed by the (PHP$\rightarrow$PHP) and (Python$\rightarrow$Python)
models. When comparing intra-lingual to cross-lingual performance for each
programming language, models generally show better results in the intra-
lingual setting. For example, the C$\rightarrow$C model outperforms the
C$\rightarrow$PHP and C$\rightarrow$Python models. This pattern is consistent
when the model is tested using PHP and Python languages. A higher performance
in the intra-lingual setting is expected, given the syntactical similarities
between the fine-tuning and testing datasets.

However, an interesting observation is the relatively small performance gap
between intra-lingual and cross-lingual settings. For example, the
PHP$\rightarrow$PHP model exhibits similar performance to the
PHP$\rightarrow$C model and the PHP$\rightarrow$Python model. A similar trend
is also observed in the C$\rightarrow$C model against the C$\rightarrow$PHP
and C$\rightarrow$Python models and Python$\rightarrow$Python against the
Python$\rightarrow$C and Python$\rightarrow$PHP models.

Answer to RQ1: The models perform better in the intra-lingual than cross-
lingual settings, as shown by Figure 2. However, while cross-lingual
performance is generally lower, the gap between intra-lingual and cross-
lingual results is notably small.

Figure 2: The F1 scores for models fine-tuned on C, PHP, and Python, across
different test languages. Two sets of bars are displayed for each test
language, representing model performance with (grey) and without (black)
natural language instructions.

Figure 2 also presents the evaluation results for intra-lingual and cross-
lingual settings, comparing scenarios with and without natural language
instructions. A significant decline in the model’s performance is evident when
inputs include natural language instructions. Such a case is indicated by the
lower F1 scores of grey bars compared to black bars in Figure 2. For instance,
the model fine-tuned and tested on C without instructions C$\rightarrow$C
yields better performance than the model fine-tuned and tested on C with
instructions (Cinstruct$\rightarrow$Cinstruct). The same result is also
observed in the Python$\rightarrow$Python model against the
Pythoninstruct$\rightarrow$Pythoninstruct model and the PHP$\rightarrow$PHP
model against the PHPinstruct$\rightarrow$PHPinstruct model. The cross-lingual
setting also exhibits the same trend. This decrease in performance might be
due to the natural language instructions inadvertently diverting the model’s
attention from important syntactical and semantic features of the code, which
are essential to predict the existence of vulnerabilities in code.

Answer to RQ2: The models without natural language instructions outperform
those with instructions across intra-lingual and cross-lingual settings. This
is exemplified by higher F1 scores in Figure 2 in scenarios such as
C$\rightarrow$C compared to Cinstruct$\rightarrow$Cinstruct, and similarly for
Python and PHP models.

Figure 3: Comparative Performance of Models in Different Testing Sets. These
grouped bar charts illustrate the F1 scores for models fine-tuned on various
language combinations (C, C-PHP, C-Python, C-Python-PHP) and tested on C, PHP,
and Python. Each chart contrasts the model’s performance with (grey bars) and
without (black bars) natural language instructions.

Figure 3 demonstrates the impact of language diversity in the fine-tuning set
on the model’s performance. For the model tested on the C language (the
leftmost chart of Figure 3), incorporating additional languages when fine-
tuning the model without natural language instructions leads to a decrease in
the model’s performance. This finding aligns with our expectations. However,
the results show that adding relevant languages when fine-tuning the model
without natural language instructions does not enhance performance but yields
a significant decline. Such a case is indicated by the higher black bars in
Figure 3 when the model is fine-tuned using a single language rather than
multiple languages. For example, in the middle chart of Figure 3, the
C-PHP$\rightarrow$PHP model exhibits poorer performance than the
C$\rightarrow$PHP model. The same trend is also observed in the
C-Python$\rightarrow$Python model against the C$\rightarrow$Python model.
Similarly, including all languages (C, PHP, and Python) when fine-tuning the
model without natural language instructions also does not help the model
perform better than only fine-tuning the model using a single language.

Conversely, a different pattern emerges when fine-tuning using natural
language instructions with more programming languages. Adding more programming
languages can improve the model’s performance across all test sets. This
improvement is illustrated by the higher grey bars in Figure 3 when the model
is fine-tuned using multiple languages than when only using a single language.
For instance, the C-Python-PHP$\rightarrow$C, C-Python$\rightarrow$C, and
C-PHP$\rightarrow$C models demonstrate better performance than the
C$\rightarrow$C model as indicated by the leftmost chart of Figure 3. The PHP
and Python languages also exhibit the same trend.

Answer to RQ3: Adding multiple programming languages to the fine-tuning set
without natural language instructions reduces model performance. This is shown
by higher black bars in the charts in Figure 3 for models fine-tuned on
multiple languages compared to those on a single language. Conversely,
incorporating natural language instructions and multiple languages in the
fine-tuning process improves the model’s performance, as indicated by higher
grey bars in the charts.

##  7 Discussion

This section discusses the manual case study, lessons learned, and threats to
validity of this study.

###  7.1 Can Language-specific Instructions Improve the Model’s Performance?

One hypothesis for the model’s poor performance when the model is fine-tuned
using natural language instructions is that the instruction pool lacks
language-specific cues, i.e., the same instructions are used across all
languages. To verify this, we conducted a check by modifying the experimental
setup from RQ3. We redo the experiments by explicitly stating the programming
language of the code snippet in the instruction, as illustrated in the updated
template shown in Figure 4. We conduct the new experiments three times using
three different random seeds.

The following code is written in {insert language}
<s> [INST] {insert instruction} [/INST]
### Code: {insert code}
### Answer: {insert label} </s>
Figure 4: The updated template that incorporates language information in the instruction. Fine-tuning set | F1_new | F1_ori
---|---|---
C | 0.054 | 0.156
C-PHP | 0.081 | 0.194
C-Python | 0.057 | 0.283
C-Python-PHP | 0.045 | 0.411
Table 3: The comparison results between the updated and old templates when
tested on C languages. F1_ori is the F1 score using the instruction in Figure
1, while F1_new is the F1 score using the instruction in Figure 4.

Table 3 presents a comparison between the new results obtained using the
updated template and the previous results from Figure 3. These results
demonstrate a performance decline when employing the updated template.
Specifically, the F1 scores for models tested on the C language with
instructions from Figure 4 are lower than those achieved with the original
instruction format in Figure 1. A similar trend is observed in tests involving
PHP and Python languages. This pattern suggests that adding more instructions
might detrimentally affect model performance. Furthermore, the significant
performance drop highlighted in Table 3 indicates a notable sensitivity of the
model to the instructions.

###  7.2 Is the Model Fine-tuned on Language X Able to Predict a Similar
Vulnerability in Language Y?

We investigate whether models fine-tuned in one programming language can
predict vulnerabilities in another by comparing intra-lingual and cross-
lingual settings within the same testing language. Specifically, we chose the
C$\rightarrow$C model without instructions for intra-lingual comparison due to
its superior performance as shown in Figure 2. For cross-lingual comparison,
we analyze models trained on a single language but tested on C, namely the
PHP$\rightarrow$C and Python$\rightarrow$C models. Both models are also fine-
tuned without instructions. We compare the top-3 correct predictions across
these models, which are consistent across all: CWE-125 (out of bounds read),
CWE-119 (buffer overflow), and CWE-20 (improper input validation), as detailed
in Table 4.

CWE-ID | C$\rightarrow$C | PHP$\rightarrow$C | Python$\rightarrow$C
---|---|---|---
CWE-125 | 36 | 36 | 19
CWE-119 | 23 | 23 | 7
CWE-20 | 20 | 20 | 6
Table 4: Number of correct instances for the top-3 predicted vulnerabilities
for the C$\rightarrow$C, PHP$\rightarrow$C, and Python$\rightarrow$C models.

Further analysis is conducted by examining the fine-tuning data for
occurrences of the CWEs listed in Table 4, as summarized in Table 5.

CWE-ID | C | PHP | Python
---|---|---|---
CWE-125 | 353 | 1 | 68
CWE-119 | 274 | 0 | 26
CWE-20 | 230 | 136 | 153
Table 5: Number of instances for each CWE-ID in the fine-tuning dataset.

The data reveals that models can recognize similar vulnerabilities even when
fine-tuned in different languages. For example, the Python$\rightarrow$C
model, initially fine-tuned with 68 instances of CWE-125 in Python,
successfully detects 19 instances in C, as shown in Table 4. Notably, the
PHP$\rightarrow$C model, despite being fine-tuned with only one instance of
CWE-125, matches the C$\rightarrow$C model’s performance in correctly
detecting 36 vulnerability instances in C language. Moreover, even with no
instances of CWE-119 in its fine-tuning data, the PHP$\rightarrow$C model
correctly predicts 36 instances when evaluated in C language. We conduct a
manual check and find that CWE-119 and CWE-125 are closely related, i.e.,
buffer overflow can potentially cause out-of-bound access. This could suggest
a correlation between CWE-119 and CWE-125, where learning to detect one may
improve the detection of the other.

###  7.3 Lessons Learned

This study yields several insightful lessons, which are outlined as follows:

Instruction-based fine-tuning may not always enhance performance. Our findings
suggest that fine-tuning language models using natural language instructions
does not invariably lead to improved performance. This observation contrasts
with prior studies Wei et al. (2022a); Honovich et al. (2023), which
highlighted the benefits of instruction-based fine-tuning. Consequently, we
encourage the research community to further investigate this area to gain a
deeper understanding of instruction-based fine-tuning.

Single-language fine-tuning outperforms multi-language fine-tuning. Our
findings reveal that training language models with data from multiple
languages can significantly diminish their performance. Therefore, it appears
more effective to fine-tune using data from a single language. This can be
attributed to several factors, such as syntax and grammatical complexity,
limited model capacity, and cross-lingual interference.

Predicting similar vulnerabilities across languages. Our results indicate the
feasibility of accurately predicting similar vulnerabilities in different
programming languages. This approach could be particularly advantageous when
dealing with low-resource languages. An alternative strategy might involve
identifying languages with a diverse range of vulnerabilities and using these
as the basis for fine-tuning, rather than relying solely on data from low-
resource languages.

###  7.4 Threats to Validity

Internal Validity: The choice of the model can significantly influence our
study. To minimize such threats, we systematically select the model through
initial filtering and preliminary experiments. Additionally, we employ the
official implementations available on HuggingFace to ensure consistency. We
also conduct the experiments three times, which helps to avoid bias that might
result in a particular trial favoring specific settings.

External Validity: The selection of our dataset can impact the study. To
minimize this threat, we choose the dataset systematically. Furthermore, the
class imbalance between vulnerable and non-vulnerable code is a potential
issue. We address this by carefully preprocessing the dataset to balance the
number of vulnerable and non-vulnerable code snippets across different
programming languages. We also ensure equal representation of each programming
language in the dataset. However, our study is limited to three languages: C,
PHP, and Python, considering the time and resource limitations. However, we
consider their syntax diversity and the programming use cases to make sure
that the chosen languages are representative. Moreover, the threat of
randomness is mitigated by setting a fixed random seed for each experiment
trial.

Construct Validity: The choice of evaluation metrics can potentially threaten
the validity of our study, as inappropriate metrics may not accurately capture
the intended outcomes. To minimize this threat, we conduct a manual case study
alongside the automated metrics in Section 7.

##  8 Conclusion and Future Work

This study scrutinizes the model generalization beyond the languages in the
fine-tuning data and the role of natural language instructions in improving
generalization performance in the vulnerability detection task. We have
conducted experiments using recent models to detect real-world
vulnerabilities. Our study yields three insights. First, models perform more
effectively in scenarios where the language is the same as in the training
data (intra-lingual) compared to different languages (cross-lingual), although
the difference in performance is not substantial. Second, models that do not
use natural language instructions outperform those that do in both intra-
lingual and cross-lingual settings. Third, when multiple programming languages
are added to the training set without natural language instructions, there’s a
decline in model performance. Conversely, when natural language instructions
are combined with multiple programming languages in the fine-tuning process,
the models show better performance.

Future Work. Future research can further investigate the generalization
capabilities of models in identifying different types of vulnerabilities. Our
initial case study indicates that the model can generalize from one
vulnerability type, such as a buffer overflow, to a closely related one, like
out of bounds read. However, comprehensive studies are required to confirm its
applicability to a broader range of vulnerabilities. Additionally, this work
has primarily focused on straightforward fine-tuning instructions. Exploring
advanced techniques like Chain-of-Thought Wei et al. (2022b), instruction with
reasoning Mitra et al. (2023), or instruction with demonstration Shao et al.
(2023), which have shown promise in other tasks, could offer valuable insights
in the context of vulnerability detection. Finally, expanding experiments to
include a diverse array of models and programming languages could provide a
more comprehensive understanding of the model’s effectiveness across different
scenarios.

## References

  * Bhandari et al. (2021) Guru Prasad Bhandari, Amara Naseer, and Leon Moonen. 2021.  Cvefixes: automated collection of vulnerabilities and their fixes from open-source software.  In _PROMISE_ , pages 30–39. ACM.
  * Chakraborty et al. (2022) Saikat Chakraborty, Rahul Krishna, Yangruibo Ding, and Baishakhi Ray. 2022.  Deep learning based vulnerability detection: Are we there yet?  _IEEE Trans. Software Eng._ , 48(9):3280–3296.
  * Chen et al. (2023) Yizheng Chen, Zhoujie Ding, Lamya Alowain, Xinyun Chen, and David A. Wagner. 2023.  Diversevul: A new vulnerable source code dataset for deep learning based vulnerability detection.  In _RAID_ , pages 654–668. ACM.
  * Cox (1958) David R Cox. 1958.  The regression analysis of binary sequences.  _Journal of the Royal Statistical Society Series B: Statistical Methodology_ , 20(2):215–232.
  * Dao (2023) Tri Dao. 2023.  Flashattention-2: Faster attention with better parallelism and work partitioning.  _CoRR_ , abs/2307.08691.
  * Fan et al. (2020) Jiahao Fan, Yi Li, Shaohua Wang, and Tien N. Nguyen. 2020.  A C/C++ code vulnerability dataset with code changes and CVE summaries.  In _MSR_ , pages 508–512. ACM.
  * Fu and Tantithamthavorn (2022) Michael Fu and Chakkrit Tantithamthavorn. 2022.  Linevul: A transformer-based line-level vulnerability prediction.  In _MSR_ , pages 608–620. ACM.
  * Fu et al. (2022) Michael Fu, Chakkrit Tantithamthavorn, Trung Le, Van Nguyen, and Dinh Q. Phung. 2022.  Vulrepair: a t5-based automated software vulnerability repair.  In _ESEC/SIGSOFT FSE_ , pages 935–947. ACM.
  * Honovich et al. (2023) Or Honovich, Thomas Scialom, Omer Levy, and Timo Schick. 2023.  Unnatural instructions: Tuning language models with (almost) no human labor.  In _ACL (1)_ , pages 14409–14428. Association for Computational Linguistics.
  * Hu et al. (2022) Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2022.  Lora: Low-rank adaptation of large language models.  In _ICLR_. OpenReview.net.
  * Jiang et al. (2023) Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. 2023.  Mistral 7b.  _CoRR_ , abs/2310.06825.
  * Johnson and Khoshgoftaar (2019) Justin M. Johnson and Taghi M. Khoshgoftaar. 2019.  Survey on deep learning with class imbalance.  _J. Big Data_ , 6:27.
  * Li et al. (2018) Zhen Li, Deqing Zou, Shouhuai Xu, Xinyu Ou, Hai Jin, Sujuan Wang, Zhijun Deng, and Yuyi Zhong. 2018.  Vuldeepecker: A deep learning-based system for vulnerability detection.  In _NDSS_. The Internet Society.
  * Liu et al. (2015) Bin Liu, Bin Liu, Hongxia Jin, and Ramesh Govindan. 2015.  Efficient privilege de-escalation for ad libraries in mobile apps.  In _MobiSys_ , pages 89–103. ACM.
  * Ma et al. (2016) Ziang Ma, Haoyu Wang, Yao Guo, and Xiangqun Chen. 2016.  Libradar: fast and accurate detection of third-party libraries in android apps.  In _ICSE (Companion Volume)_ , pages 653–656. ACM.
  * Mishra et al. (2022) Swaroop Mishra, Daniel Khashabi, Chitta Baral, and Hannaneh Hajishirzi. 2022.  Cross-task generalization via natural language crowdsourcing instructions.  In _ACL (1)_ , pages 3470–3487. Association for Computational Linguistics.
  * Mitra et al. (2023) Arindam Mitra, Luciano Del Corro, Shweti Mahajan, Andrés Codas, Clarisse Simões, Sahaj Agrawal, Xuxi Chen, Anastasia Razdaibiedina, Erik Jones, Kriti Aggarwal, Hamid Palangi, Guoqing Zheng, Corby Rosset, Hamed Khanpour, and Ahmed Awadallah. 2023.  Orca 2: Teaching small language models how to reason.  _CoRR_ , abs/2311.11045.
  * Narayanan et al. (2014) Annamalai Narayanan, Lihui Chen, and Chee Keong Chan. 2014.  Addetect: Automated detection of android ad libraries using semantic analysis.  In _ISSNIP_ , pages 1–6. IEEE.
  * Nikitopoulos et al. (2021) Georgios Nikitopoulos, Konstantina Dritsa, Panos Louridas, and Dimitris Mitropoulos. 2021.  Crossvul: a cross-language vulnerability dataset with commit data.  In _ESEC/SIGSOFT FSE_ , pages 1565–1569. ACM.
  * 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.  _CoRR_ , abs/2308.12950.
  * Russell et al. (2018) Rebecca L. Russell, Louis Y. Kim, Lei H. Hamilton, Tomo Lazovich, Jacob Harer, Onur Ozdemir, Paul M. Ellingwood, and Marc W. McConley. 2018.  Automated vulnerability detection in source code using deep representation learning.  In _ICMLA_ , pages 757–762. IEEE.
  * Shao et al. (2023) Zhihong Shao, Yeyun Gong, Yelong Shen, Minlie Huang, Nan Duan, and Weizhu Chen. 2023.  Synthetic prompting: Generating chain-of-thought demonstrations for large language models.  In _ICML_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 30706–30775. PMLR.
  * Steenhoek et al. (2022) Benjamin Steenhoek, Wei Le, and Hongyang Gao. 2022.  Deepdfa: Dataflow analysis-guided efficient graph learning for vulnerability detection.  _CoRR_ , abs/2212.08108.
  * Steenhoek et al. (2023) Benjamin Steenhoek, Md Mahbubur Rahman, Richard Jiles, and Wei Le. 2023.  An empirical study of deep learning models for vulnerability detection.  In _ICSE_ , pages 2237–2248. IEEE.
  * 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, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. 2023.  Llama 2: Open foundation and fine-tuned chat models.  _CoRR_ , abs/2307.09288.
  * Wang et al. (2023a) Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A. Smith, Daniel Khashabi, and Hannaneh Hajishirzi. 2023a.  Self-instruct: Aligning language models with self-generated instructions.  In _ACL (1)_ , pages 13484–13508. Association for Computational Linguistics.
  * Wang et al. (2023b) Yue Wang, Hung Le, Akhilesh Deepak Gotmare, Nghi D. Q. Bui, Junnan Li, and Steven C. H. Hoi. 2023b.  Codet5+: Open code large language models for code understanding and generation.  _CoRR_ , abs/2305.07922.
  * Wang et al. (2021) Yue Wang, Weishi Wang, Shafiq R. Joty, and Steven C. H. Hoi. 2021.  Codet5: Identifier-aware unified pre-trained encoder-decoder models for code understanding and generation.  In _EMNLP (1)_ , pages 8696–8708. Association for Computational Linguistics.
  * Wei et al. (2022a) Jason Wei, Maarten Bosma, Vincent Y. Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V. Le. 2022a.  Finetuned language models are zero-shot learners.  In _ICLR_. OpenReview.net.
  * Wei et al. (2022b) Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoc V. Le, and Denny Zhou. 2022b.  Chain-of-thought prompting elicits reasoning in large language models.  In _NeurIPS_.
  * Zhan et al. (2021) Xian Zhan, Lingling Fan, Sen Chen, Feng Wu, Tianming Liu, Xiapu Luo, and Yang Liu. 2021.  ATVHUNTER: reliable version detection of third-party libraries for vulnerability identification in android applications.  In _ICSE_ , pages 1695–1707. IEEE.
  * Zhang et al. (2019) Jiexin Zhang, Alastair R. Beresford, and Stephan A. Kollmann. 2019.  Libid: reliable identification of obfuscated third-party android libraries.  In _ISSTA_ , pages 55–65. ACM.
  * Zhang et al. (2018) Yuan Zhang, Jiarun Dai, Xiaohan Zhang, Sirong Huang, Zhemin Yang, Min Yang, and Hao Chen. 2018.  Detecting third-party libraries in android applications with high precision and recall.  In _SANER_ , pages 141–152. IEEE Computer Society.
  * Zheng et al. (2021) Yunhui Zheng, Saurabh Pujar, Burn L. Lewis, Luca Buratti, Edward A. Epstein, Bo Yang, Jim Laredo, Alessandro Morari, and Zhong Su. 2021.  D2A: A dataset built for ai-based vulnerability detection methods using differential analysis.  In _ICSE (SEIP)_ , pages 111–120. IEEE.
  * Zhou et al. (2019) Yaqin Zhou, Shangqing Liu, Jing Kai Siow, Xiaoning Du, and Yang Liu. 2019.  Devign: Effective vulnerability identification by learning comprehensive program semantics via graph neural networks.  In _NeurIPS_ , pages 10197–10207.