# Learning to Quantize Vulnerability Patterns and Match to Locate Statement-
Level Vulnerabilities

Michael Fu1 Trung Le1 Van Nguyen1,2
Chakkrit Tantithamthavorn1 Dinh Phung1,3
1Monash University, Australia
2CSIRO’s Data61, Australia
3VinAI Research, Vietnam

###### Abstract

Deep learning (DL) models have become increasingly popular in identifying
software vulnerabilities. Prior studies found that vulnerabilities across
different vulnerable programs may exhibit similar vulnerable scopes,
implicitly forming discernible vulnerability patterns that can be learned by
DL models through supervised training. However, vulnerable scopes still
manifest in various spatial locations and formats within a program, posing
challenges for models to accurately identify vulnerable statements. Despite
this challenge, state-of-the-art vulnerability detection approaches fail to
exploit the vulnerability patterns that arise in vulnerable programs. To take
full advantage of vulnerability patterns and unleash the ability of DL models,
we propose a novel vulnerability-matching approach in this paper, drawing
inspiration from program analysis tools that locate vulnerabilities based on
pre-defined patterns. Specifically, a vulnerability codebook is learned, which
consists of quantized vectors representing various vulnerability patterns.
During inference, the codebook is iterated to match all learned patterns and
predict the presence of potential vulnerabilities within a given program. Our
approach was extensively evaluated on a real-world dataset comprising more
than 188,000 C/C++ functions. The evaluation results show that our approach
achieves an F1-score of 94% (6% higher than the previous best) and 82% (19%
higher than the previous best) for function and statement-level vulnerability
identification, respectively. These substantial enhancements highlight the
effectiveness of our approach to identifying vulnerabilities. The training
code and pre-trained models are available at
https://github.com/optimatch/optimatch.

##  1 Introduction

The number of software vulnerabilities has been escalating rapidly in recent
years. In particular, National Vulnerability Database (NVD) [6] reported
26,448 software vulnerabilities in 2022, soaring 40% from 18,938 in 2019. The
extensive use of open-source libraries, in particular, may contribute to this
rise in vulnerabilities. For instance, the Apache Struts vulnerabilities [31]
indicate that this poses a tangible threat to organizations. The root cause of
these vulnerabilities is often insecure coding practices, making the source
code exploitable by attackers who can use them to infiltrate software systems
and cause considerable financial and social harm.

To mitigate security threats, security experts leverage static analysis tools
that check the code against a set of known patterns of insecure or vulnerable
code, such as buffer overflow vulnerabilities and other common security flaws.
In contrast, deep learning-based vulnerability detection (VD) identifies
vulnerabilities at the file or function levels by implicitly learning
vulnerability patterns during training [33, 38, 28, 30]. DL-based
vulnerability detection (VD) methods have demonstrated higher accuracy
compared to static analysis tools that only target specific vulnerability
types [23, 17, 11]. Additionally, recent advancements have introduced fine-
grained VDs that offer statement-level vulnerability predictions, aiming to
minimize the manual analysis burden on security analysts. Previous studies
have employed graph structure of source code like the code property graph
[21], along with graph neural networks to detect vulnerabilities at the
statement level [22, 19]. Additionally, transformers have demonstrated their
capability to learn semantic features of code using self-attention, which is
particularly beneficial for handling long sequences compared to RNN models
[17, 13].

In this paper, we consider a vulnerable scope of a function as the collection
of all vulnerable statements in that function. As illustrated in Figure 1,
each function consists of two vulnerable statements that form similar
vulnerable scopes. This suggests that even if two functions contain the same
CWE-787 out-of-bound write vulnerability (the top-1 dangerous CWE-ID in 2022
[10]), the specific vulnerable statements can be written in different ways and
located in different parts of the code. Therefore, identifying vulnerabilities
at the statement level is challenging for both machine learning and deep
learning models. Despite this difficulty, our analysis reveals that state-of-
the-art VD approaches have not successfully leveraged the information
contained in vulnerable statements (that could be grouped to form vulnerable
scopes) to further improve the capability of machine learning and deep
learning vulnerability detection approaches at both the function and statement
levels.

To address this issue, we propose a novel DL-based framework that can
effectively utilize the information presented in vulnerable scopes. To achieve
this, we develop a method for quantizing similar vulnerable scopes that share
the same pattern into a vulnerability codebook consisting of common codewords
which represent common patterns. This codebook captures a diverse range of
vulnerabilities from the training data and facilitates the process of
vulnerability matching inspired by the pattern-matching concept utilized in
program analysis tools [1, 2, 3, 4]. Our approach is the first to successfully
exploit the benefits of vulnerability matching and codebook-based quantization
to enhance DL-based VD. This allows us to effectively identify vulnerabilities
in source code data, ultimately improving the overall capability of DL-based
VD.

Our approach involves collecting and quantizing a set of vulnerable scopes
from the training set before using the optimal transport (OT) [16] to cluster
this set into a vulnerability codebook consisting of a set of vulnerability
centroids (i.e., codewords). The vulnerable scopes (collected from the
training set) that share a similar pattern would stay closely in
representations, hence we cluster them into a centroid to summarize them. By
clustering the set of vulnerable scopes into a smaller set of centroids, we
reduce the dimensionality of the feature space and make it easier for the
model to perform matching during inference. Additionally, the use of centroids
ensures that similar vulnerable scopes are mapped to the same location in the
feature space. During training, we minimize the Wasserstein distance [16]
between the set of vulnerable scopes and the vulnerability codebook, which
allows us to effectively cluster vulnerable scopes and learn the
representative centroids in the codebook. During inference, our model matches
the input program against all centroids in the learned vulnerability codebook.
By examining all the vulnerability patterns in the codebook, the matching
process enables a thorough search for potential vulnerabilities. This explicit
matching method supports the identification of specific vulnerability patterns
and their associated statements, providing a comprehensive approach to
identifying vulnerabilities. We name this model OptiMatch, a function and
statement-level vulnerability identification approach via optimal transport
quantization and vulnerability matching.

In summary, our work presents several contributions: (i) an innovative
vulnerability-matching DL framework utilizing optimal transport and vector
quantization for function and statement-level vulnerability detection (VD);
(ii) a novel statement embedding approach using recurrent neural networks
(RNNs); and (iii) a thorough evaluation of our proposed method compared to
other DL-based vulnerability prediction techniques on a large benchmark
dataset of real-world vulnerabilities.

Figure 1: In the left function, writeToBuffer, if the sum of offset and i
exceeds or equals 20, it results in writing data beyond the buffer array’s
end. This overwrites memory beyond the array, posing a potential program
crash. Similarly, the copyToMemory function on the right uses the start index
to determine the starting point for copying data in memoryBlock. However, if
the sum of start and i surpasses or equals the size of memoryBlock, it leads
to overwriting memory beyond the array, causing an out-of-bounds write
vulnerability. Despite sharing the same vulnerability type and similar
vulnerable scopes, the vulnerable statements in each function are different in
their written form, variable names, and positions.

##  2 Related Work

Researchers have proposed various deep learning-based vulnerability detections
(VDs) such as convolutional neural networks (CNNs) [33], recurrent neural
networks (RNNs) [24, 29, 27], graph neural networks (GNNs) [38, 7, 22, 19, 30,
13], and pre-trained transformers [15, 18, 17, 13]. RNN-based methods [33, 24,
23] have been shown more accurate than program analysis tools such as
Checkmarx [1] and RATS [4] to predict function-level vulnerabilities. However,
RNNs face difficulty in capturing long-term dependencies in long sequences as
the model’s sequential nature may result in the loss of earlier sequence
information. Furthermore, function-level predictions lack the required
granularity to accurately identify the root causes of vulnerabilities. Thus,
researchers have proposed transformer-based methods that predict statement-
level vulnerabilities and capture long-term dependencies [13, 17] while ICVH
[28] leverages bidirectional RNNs with information theory to detect statement-
level vulnerabilities. On the other hand, Zhou et al. [38] embed the abstract
syntax tree (AST), control flow graph (CFG), and data flow graph (DFG) for a
code function and learn the graph representations for function-level
predictions. Nguyen et al. [30] proposed constructing a code graph as a flat
sequence for function-level predictions. Hin et al. [19] constructed program
dependency graphs (PDGs) for functions and predicted statement-level
vulnerabilities.

In contrast to the above methods, we propose a deep learning-based
vulnerability matching method inspired by the principles of program analysis
security tools. Specifically, we gather a group of vulnerability patterns from
the training set and develop a vulnerability codebook using optimal transport
[16] and vector quantization [35]. Our goal is to detect statements caused the
vulnerabilities by matching functions with the representative patterns we have
learned in our codebook.

##  3 Approach

Deep learning (DL) models have been proving their abilities in capturing
vulnerabilities more accurately than program analysis tools using implicit
vulnerability patterns learned from the training data set [11, 17]. However,
in real-world source code data sets, common vulnerable scopes would be written
in different styles (e.g., variable naming conventions) and appear at
different spatial locations in different vulnerable sections (i.e., functions
or programs) [28]. Existing DL-based VD approaches often fail to consider the
common vulnerable scopes (which could be clustered into patterns) that exist
in vulnerable functions or programs during both training and inference,
instead relying on implicit learning through supervised learning. To address
this limitation, we propose a novel DL framework that integrates vulnerable
scopes into centroids via a vulnerability codebook. The example in Figure 1
demonstrates that the two vulnerable functions have the same vulnerable scope,
consisting of two vulnerable statements, but are presented in different
variable names and spatial locations. To overcome this issue, we group these
vulnerable scopes with the same pattern and quantize them into a codebook
containing representative vulnerability centroids that can represent a set of
similar scopes. This codebook is then used to facilitate vulnerability
matching during the inference phase, effectively addressing the lack of
consideration for vulnerable scopes in existing DL-based VD approaches.

In general, our approach consists of two phases. The warm-up phase illustrated
in Figure 2 aims to gradually adjust the model parameters, with the goal of
improving the representation of embeddings for input programs and vulnerable
scopes. The main training phase is illustrated in Figure 3. The yellow section
on the left shows how we construct and learn our vulnerability codebook from
vulnerable scopes in our training data using optimal transport. The grey
section on the right shows how to utilize the codebook during training, which
matches functions with the learned vulnerability centroids in the codebook,
allowing us to identify and highlight the statements that caused the
vulnerabilities. Below, we first formulate our problem by defining common
notations followed by how we map textual source code to vector space and warm
up the embeddings. We then introduce the motivation and method on why and how
to learn a DL framework to achieve vulnerability matching.

###  3.1 Problem Statement

Let us consider a dataset of $N$ functions in the form of source code. The
data set includes both vulnerable and benign functions, where the function-
level and statement-level ground truths have been labeled by security experts.
We denote a function as a set of code statements,
$X_{i}=[\mathbf{x}_{1},...,\mathbf{x}_{n}]$, where $n$ is the max number of
statements we consider in a function. Let a sample of data be
$\bigl{\\{}(X_{i},y_{i},\mathbf{z}_{i}):X_{i}\in\mathcal{X},y_{i}\in\mathcal{Y},\mathbf{z}_{i}\in\mathcal{Z},i\in\\{1,2,...,N\\}\\}$,
where $\mathcal{X}$ denotes a set of code functions, $\mathcal{Y}=\\{0,1\\}$
with 1 represents vulnerable function and 0 otherwise, and
$\mathcal{Z}=\\{0,1\\}^{n}$ denotes a set of binary vectors with 1 represents
vulnerable code statement and 0 otherwise. Our objective is to identify the
vulnerability on both function and statement levels. We formulate the
identification of vulnerable functions as a binary classification problem and
the identification of vulnerable statements as a multi-label classification
problem. Given a function $X_{i}$, we first input to a statement embedding
layer ($SEMB$) to obtain statement embeddings, namely $S_{i}$ and $P_{i}$, as
specified in Equation 2 (refer to Section 3.2 for the embedding details).
$S_{i}\in\mathbb{R}^{n\times d}$ is the d-dimensional statement embedding
vectors for $X_{i}$. Prior studies have found that in a vulnerable function,
there are code statements associated with the vulnerabilities (i.e.,
vulnerable scopes) [28]. Let us denote $X_{i}^{vul}$ as a set of all
vulnerable statements in a vulnerable function. To explicitly capture
vulnerable scopes, we extract $X_{i}^{vul}$ from the vulnerable function and
encode those statements using d-dimensional statement embeddings as
$P_{i}\in\mathbb{R}^{q\times d}$. $q$ is the maximum number of vulnerable
statements we consider in a vulnerable function and we set $q=12$ by applying
truncation and padding because 95% of vulnerable functions in our data have
less than 12 vulnerable statements. Note that for each benign function without
any vulnerable statements, we leverage a special learnable embedding denoted
as $P_{benign}\in\mathbb{R}^{q\times d}$ to represent $P_{i}$. In addition, we
apply an RNN layer ($RNN_{vul}$) to summarize $P_{i}$ into a flat vector
denoted as $\mathbf{v_{i}}\in\mathbb{R}^{d}$, which can facilitate the
learning of our vulnerability codebook introduced in Section 3.4.2. Let us
denote a stack of transformer encoders as $\mathcal{F}$, we concatenate
$S_{i}$, $\mathbf{v_{i}}$, and feed them into transformer encoders as
$\mathcal{F}(S_{i},\mathbf{v_{i}})$. We then make function and statement-level
predictions based on the output of $\mathcal{F}$. The mapping from $X_{i}$ to
$y_{i}$ and $z_{i}$ is learned by minimizing the cross-entropy loss function,
denoted by $\mathcal{L}(\cdot)$, as follows:

$min\frac{1}{N}\sum_{i=1}^{N}\Bigl{[}\mathcal{L}_{function}\Bigl{(}\mathcal{F}(S_{i},\mathbf{v_{i}}),y_{i}|X_{i}\Bigl{)}+\mathcal{L}_{statement}\Bigl{(}\mathcal{F}(S_{i},\mathbf{v_{i}}),\mathbf{z}_{i}|X_{i}\Bigl{)}\Bigl{]}$
(1)

Figure 2: The overview of the warm-up phase in our approach. We tokenize each
statement in a vulnerable function (i.e., $X_{i}$) followed by an embedding
layer to map each token into a vector. We use $RNN_{statement}$ to summarize
the token embeddings and get the statement embedding ($S_{i}$, $P_{i}$). For
benign functions, $P_{i}$ is replaced by a special learnable embedding,
$P_{benign}$. Additionally, we use $RNN_{vul}$ to summarize vulnerable
statement embeddings $P_{i}$ to a vector $\mathbf{v_{i}}$ that represents the
vulnerable scope. We concatenate $S_{i}$ and $\mathbf{v_{i}}$ as the input to
transformer encoders to consider vulnerable scopes that arise in the function
and align with our vulnerability matching process introduced in Section 3.5.
We select the statement embeddings output from the last encoder, i.e.,
$H^{12}[1:n]$. Each statement embedding vector is mapped to a probability as
statement-level predictions, the function-level prediction is obtained by
summarising $H^{12}[1:n]$ to a vector using an $RNN_{function}$ and mapping it
to a probability.

###  3.2 Statement Embedding Using RNN

Figure 2 depicts the forward passes involved in our warm-up step to adjust the
embeddings for statements and vulnerable scopes. We now introduce our
motivations and method to embed statements and vulnerable scopes. Large
language models (LLMs) pre-trained for source code have been shown effective
to predict vulnerabilities [15, 18, 19, 17]. However, those LLMs leverage
token embeddings that only preserve 512 tokens (tokenized by the byte pair
encoding (BPE) algorithm [34]) for each input function while extra tokens need
to be truncated. This could lead to information loss for long functions with
more than 512 tokens. To address this limitation, we propose the statement
embedding layer $SEMB$ to encode a function (e.g., $X_{i}$) as a set of
statement embeddings:

$S_{i}=SEMB(X_{i}),\hskip 2.84526ptP_{i}=SEMB({X_{i}^{vul}}),\hskip
2.84526pt\textrm{where}\hskip 2.84526ptX_{i},X_{i}^{vul}\in\mathcal{X}$ (2)

Given $X_{i}=[\mathbf{x}_{1},...,\mathbf{x}_{n}]$, we use BPE to tokenize
$\mathbf{x}_{j}$ to a list of tokens, $[t_{1},...,t_{r}]$, where $r$ is the
number of tokens we consider in a code statement. We then obtain a token
embedding for each $t_{j}$ using an embedding layer $E\in\mathbb{R}^{v\times
d}$ where $v$ is the vocab size of the tokenizer. This results in a token
embedding matrix $\bar{S_{i}}\in\mathbb{R}^{n\times r\times d}$ for all
statements in $X_{i}$. Similarly, we obtain token embeddings of vulnerable
statements $X_{i}^{vul}$ as $\bar{P_{i}}\in\mathbb{R}^{q\times r\times d}$ to
represent a vulnerable scope. We apply truncation and padding to make $q$ a
constant for each vulnerable function.

With $n=155$ and $r=20$ (see Section 4.2), we can process 3,100 tokens per
function, which is six times more than the 512 tokens. Our statement embedding
method provides a more complete representation of code functions compared to
the token embedding method. Specifically, our method can fully represent 99%
of the functions in our dataset that have less than 2,700 tokens, while the
token embedding method can only fully represent around 85% of the functions
that have less than 500 tokens. Table 1 shows that our statement embedding
method results in a 33% and 32% enhancement in the performance of CodeBERT and
CodeGPT models for statement-level predictions.

Previous studies such as Sentence-BERT [32] leverage max or mean pooling to
aggregate token embeddings. The max pooling would lead to information loss
since it considers the maximum token embedding for each statement, discarding
all other token embeddings in the sequence. While the mean pooling considers
all token embeddings, it treats all the token embeddings equally regardless of
their importance or relevance to the statement they belong where the prominent
token features could be disregarded. In contrast, we propose to learn an RNN
[9] with $r$ (max number of tokens in each statement) time steps to aggregate
the token embeddings and obtain statement embeddings as below:

$S_{i}[j]=RNN_{statement}(\bar{S}[j,:,:]),\forall_{j}\in\\{1,\ldots,n\\}$ (3)
$P_{i}[j]=RNN_{statement}(\bar{P}[j,:,:]),\forall_{j}\in\\{1,\ldots,q\\}$ (4)

To acquire the $j^{th}$ statement embedding for $S_{i}$ and $P_{i}$, we
summarize the token embeddings of length $r$ using $RNN_{statement}$.
Following the convention of Python lists, we represent the $j^{th}$ statement
embeddings as $S_{i}[j]$. While mean or max pooling operations are not
learnable, the $RNN_{statement}$ layer allows us to learn to pool token
embeddings in each statement into a statement embedding vector while
preserving prominent token features and mitigating the potential information
loss. Finally, we use $RNN_{vul}$ to summarize our vulnerable scope $P_{i}$
into a flat vector $\mathbf{v_{i}}$ (see Section 3.4.1 for more details).

###  3.3 Training of Warm-Up Phase

To consider the statement embeddings and the vulnerable scope of $X_{i}$, we
concatenate $S_{i}$ and $\mathbf{v_{i}}$ to obtain the input to transformer
encoders as $H^{0}=S_{i}\oplus\mathbf{v_{i}}$. We select the statement
embeddings output from the trail encoder, i.e., $H^{12}[1:n]$ where the
$\mathbf{v_{i}}$ embedding is omitted. We provide details of the transformer
self-attention operation in Appendix A.1. We use $RNN_{function}$ with $n$
time steps to summarize statement embeddings into a vector and map it to the
function-level prediction $\hat{y_{i}}\in[0,1]$ as follows:

$\hat{y_{i}}=\sigma\Bigl{(}drop\bigl{(}tanh(drop(RNN(H_{1:n}^{12}))W^{G})\bigl{)}W^{U}\Bigl{)}$
(5)

where $W^{G}\in\mathbb{R}^{d\times d}$ and $W^{U}\in\mathbb{R}^{d\times 1}$
are model parameters, $drop$ is a dropout layer, and $\sigma$ is a sigmoid
function. We map statement embeddings to a statement-level prediction
$\hat{z_{i}}=[\hat{z_{i}}^{1},\ldots,\hat{{z_{i}}}^{n}]\in[0,1]^{n}$ via:

$\hat{z_{i}}=\sigma\Bigl{(}drop\bigl{(}tanh(drop(H_{1:n}^{12})W^{I})\bigl{)}W^{J}\Bigl{)}$
(6)

where $W^{I}\in\mathbb{R}^{d\times d}$ and $W^{J}\in\mathbb{R}^{d\times 1}$
are model parameters, and $\sigma$ is a sigmoid function.

Figure 3: The overview of the main training phase in our approach. We
introduce how to learn our vulnerability codebook on the left. We first
collect a set of vulnerable statement embeddings from our training data. We
then use $RNN_{vul}$ to pool a set of statement embeddings from each
vulnerable function, forming a vulnerable scope represented by a vector
$\mathbf{v_{i}}$. The set of these scopes forms our vulnerability collection
$V=\\{\mathbf{v}_{1},\ldots,\mathbf{v}_{a}\\}$. Next, we learn vulnerability
centroids $\mathbf{c}_{j}$ using the Wasserstein distance metric to create a
more compact vulnerability codebook
$C=\\{\mathbf{c}_{1},\ldots,\mathbf{c}_{k}\\}$, where each centroid represents
a group of vulnerable scopes. During training, we minimize the Wasserstein
distance between each $\mathbf{v_{i}}$ and its corresponding vulnerability
centroid $\mathbf{c}_{\mathbf{v_{i}}}^{*}$. We illustrate this main training
phase on the right side which is the same as our warm-up phase except that we
concatenate $S_{i}$ and $\mathbf{c}_{\mathbf{v_{i}}}^{*}$ to obtain $H^{0}$ as
detailed in Section 3.4.3. To overcome the non-differentiability of the
$argmax$ operation in the networks, we copy the gradients from $\mathbf{v}$ to
$\mathbf{c}_{\mathbf{v_{i}}}^{*}$ to learn the statement embedding and pattern
summarization RNNs for vulnerability patterns.

###  3.4 Vulnerability Codebook and Subsequent Main Training Phase

Our model parameters are now warmed up to embed statements and vulnerable
scopes. Our objective is to achieve vulnerability matching using trainable
vulnerability centroids. In the following, we outline our motivations and
approach for creating, training, and employing our vulnerability codebook
during the primary training phase.

####  3.4.1 Collect vulnerable scopes from Vulnerable Functions

To exploit and capture common vulnerable scopes in source code, we aim to
learn a vulnerability codebook containing representative centroids that group
vulnerable scopes sharing the same pattern. Unlike those patterns in program
analysis tools, our vulnerability centroids are represented in vectors to
conform with DL models, whose representation is adjustable during training,
enabling the model to recognize typical vulnerability patterns that may occur
at various spatial locations within a vulnerable function.

Given training data consisting of $a$ vulnerability functions, we extract
vulnerable statements to form vulnerable scopes for each function as presented
in the very first left part of Figure 3. To simplify the process of building
our vulnerability codebook introduced in Section 3.4.2, we take two steps.
First, we use $RNN_{vul}$ to summarize our vulnerable scopes into flat
vectors. Then, we reduce the dimensionality of these vectors. This enables us
to easily group them into vulnerability centroids and construct our
vulnerability codebook. We have denoted our vulnerable scope as
$\mathbf{v_{i}}$ in Section 3.1. $\mathbf{v_{i}}$ is obtained by applying
$RNN_{statement}$ and $RNN_{vul}$ to get the vulnerable statement embeddings
and condense them into a flat vector. To reduce the dimensionality of
$\mathbf{\mathbf{v_{i}}}$, we linearly project the d-dimensional vector to the
h-dimensional and normalize it as $\mathbf{v_{i}}=LN(\mathbf{v_{i}}\cdot
W^{F})$ where $W^{F}\in\mathbb{R}^{d\times h}$ is model parameters and $LN$ is
layer normalization. We then accumulate each $\mathbf{v_{i}}$ extracted from
vulnerable functions to form a vulnerability collection denoted as
$V\in\mathbb{R}^{a\times h}$ where $a$ is the total number of vulnerable
functions in our training data.

####  3.4.2 Learn to Transport Vulnerable Scopes to Vulnerability Centroids in
Codebook

However, $V$ may consist of repeated or similar vulnerable scopes.
Additionally, the huge collection size of $V$ will also require many computing
resources during inference since we need to match each function with $a$
number of scopes (in our training data, $a=6,361$). To address such issues, we
propose to learn a vulnerability codebook denoted as
$C=[\mathbf{c}_{1},...,\mathbf{c}_{k}]$ where
$\mathbf{c}_{i}\in\mathbb{R}^{h}$ is a vulnerability centroid. Intuitively,
this codebook integrates similar vulnerable scopes and forms common
vulnerability patterns. In particular, we reduce the 6,361 number of
$\mathbf{v}$ vectors in our vulnerability collection to 150 vulnerability
centroids in our codebook.

To ensure that vulnerability centroids can represent a group of similar
vulnerable scopes, we leverage the optimal transport theory to transfer
vulnerability patterns to their corresponding vulnerability centroid. We
minimize the Wasserstein distance [36] using the Sinkhorn approximation [12]
between our vulnerability collection and codebook. Consequently, the
vulnerable scopes and their respective vulnerability centroids will converge
towards each other. Ultimately, our codebook will comprise vulnerability
centroids acting as representative patterns that symbolize different sets of
vulnerability scopes. This allows us to aggregate similar vulnerability
patterns based on Euclidean distance. We summarize the process as follows:

$min_{C}\hskip 2.84526ptW_{d}(P_{V},P_{C}),\hskip
5.69054pt\textrm{where}\hskip
5.69054ptP_{V}=\frac{1}{a}\sum_{i=1}^{a}\delta_{\mathbf{v}_{i}}\hskip
5.69054pt\textrm{and}\hskip
5.69054ptP_{C}=\frac{1}{a}\sum_{j=1}^{a}\delta_{\mathbf{c}_{j}}$ (7)

where $W_{d}$ is the Wasserstein distance [36] and $\delta$ represents the
Dirac delta distribution. According to the clustering view of optimal
transport [26, 20], when minimizing $min_{C}\hskip
2.84526ptW_{d}(P_{V},P_{C})$, the set of codebooks $C$ will become the
centroids of the clusters formed by $V$. This clustering approach ensures that
similar vulnerable scopes potentially sharing the same vulnerability pattern
are grouped together, leading to a quantized vulnerability codebook that is
more concise and effective. We randomly initialize the embedding space of our
vulnerability codebook as $C=[\mathbf{c}_{1},...,\mathbf{c}_{k}]$ with $k$
number of clusters.

####  3.4.3 Main Training Phase

The right part of Figure 3 highlighted in grey summarizes our main training
phase. We load the model parameters warmed up in our previous phase. By
employing the same statement embedding methodology introduced in Section 3.2,
we obtain the statement embeddings $S_{i}$ and a summarized vulnerable scope
vector $\mathbf{v_{i}}$ for the input function $X_{i}$.

Instead of concatenating $S_{i}$ with $\mathbf{v_{i}}$, we employ a cluster
selection process to map the vulnerable scope $\mathbf{v_{i}}$ to its most
similar vulnerability centroid (denoted as
$\mathbf{c}_{\mathbf{v_{i}}}^{*}\in\mathbb{R}^{1\times h}$) selected from our
codebook. By doing so, the model inherently develops an understanding of the
vulnerability centroids stored in our vulnerability codebook, which are
closely linked to vulnerable functions. We utilize the cross-attention (see
Appendix A.2) between the vulnerable scope and the codebook and determine the
vulnerability centroid for $\mathbf{v_{i}}$ as
$\mathbf{c}_{\mathbf{v_{i}}}^{*}=argmax_{C}CrossAtt(\mathbf{v_{i}},C)$. The
$argmax$ function selects the index of the vulnerability centroid with the
highest attention score, which corresponds to the closest vector to
$\mathbf{v_{i}}$ in terms of similarity. We linearly project
$\mathbf{c}_{\mathbf{v_{i}}}^{*}$ from the factorized h-dimension to the
d-dimension to align with the dimension of our statement embeddings. Different
from our warm-up phase where we concatenate $S_{i}$ with $\mathbf{v_{i}}$, we
now concatenate $S_{i}$ with $\mathbf{c}_{\mathbf{v_{i}}}^{*}$ (the most
similar centroid to the vulnerable scope $\mathbf{v_{i}}$). Thus, the input to
the encoders becomes $H^{0}=S_{i}\oplus\mathbf{c}_{\mathbf{v_{i}}}^{*}$. The
subsequent forward passes are the same as our warm-up phase described in
Section 3.3.

Note that no real gradient is defined for $\mathbf{v_{i}}$ once we map it to a
$\mathbf{c}_{\mathbf{v_{i}}}^{*}$ via an $argmax$ operation that causes the
networks non-continuous and non-differentiable. To let the networks which
embed and summarize vulnerable statements be trainable via backpropagation, we
follow the idea in VQ-VAE [35] which was shown effective for vector
quantization. We approximate the gradient similar to the straight-through
estimator [5] and copy gradients from summarized vulnerable scope
$\mathbf{v_{i}}$ to selected vulnerability centroid
$\mathbf{c}_{\mathbf{v_{i}}}^{*}$. Below, we introduce how to leverage our
learned codebook for vulnerability matching during inference.

###  3.5 Vulnerability Identification Through Explicit Vulnerability Patterns
Matching

Our approach utilizes vulnerable patterns that are often ignored by existing
methods. By matching vulnerability centroids during inference, our approach
enables us to fully harness the capabilities of DL models for vulnerability
identification. We first obtain d-dimensional statement embeddings $S_{i}$
from an input function $X_{i}$ as described in Section 3.2. For each
vulnerability centroid $\mathbf{c_{j}}$ in our codebook, we linearly project
$\mathbf{c_{j}}$ from h-dimensional to d-dimensional space and concatenate it
with $S_{i}$ as $H_{j}^{0}=S_{i}\oplus\mathbf{c_{j}}$. We then pass
$H_{j}^{0}$ through transformer encoders ($\mathcal{F}$) to obtain function-
level and statement-level vulnerability predictions, which is summarized as
$P_{i}^{func},P_{i}^{stmt}=\mathcal{F}(S_{i},\mathbf{c_{j}})\hskip
5.69054pt\forall_{j}\in\\{1,\ldots,k\\}$ where $P_{ij}^{func}\in[0,1]$ and
$P_{ij}^{stmt}\in[0,1]^{n}$. Thus, we get $k$ (number of centroids in our
codebook) function and statement-level predictions. We use max pooling to pick
the most prominent vulnerability-matching results as
$\bar{P_{i}}^{func}=max_{k}P_{i}^{func}$ and predict if $X$ is a vulnerable
function using a probability threshold of 0.5. If $X$ is predicted as a benign
function, we directly output a zero vector as the statement-level prediction.
Otherwise, we employ mean pooling to consider the prediction from each
vulnerability centroid in our codebook as
$\bar{P_{i}}^{stmt}=\frac{1}{k}\sum_{j=1}^{k}P_{ij}^{stmt}$ and predict if
each statement is vulnerable using a probability threshold of 0.5.

##  4 Experiments

###  4.1 Experimental Dataset and Baseline Methods

To identify vulnerabilities on function and statement levels, we select the
Big-Vul data set created by Fan et al. [14] as it is one of the largest
vulnerability data sets with statement-level vulnerability labels and has been
used to assess statement-level vulnerability detection methods [19, 17]. The
data set was collected from 348 Github projects and consists of 188k C/C++
functions with 3,754 code vulnerabilities spanning 91 vulnerability types. The
data distribution in our experiments resembles real-world scenarios, where the
proportion of vulnerable to benign functions is 1:20. Our training data set
comprises 6,361 vulnerability scopes before we group them into patterns in our
codebook.

We compare our approach with (i) LLMs for code (i.e., CodeBERT [15] and
GraphCodeBERT [18]), (ii) Transformer-based VD (i.e., LineVul [17] and VELVET
[13]), (iii) GNN-based VD (i.e., LineVD [19], ReGVD [30], and Devign [38]),
(iv) RNN-based ICVH [28], and (v) CNN-based TextCNN [8]. More details of the
baselines are provided in Appendix A.3.

###  4.2 Parameter Settings and Model Training

We split the data into 80% for training, 10% for validation, and 10% for
testing. For both our approach and baselines, we consider $n=155$ statements
in each function and $r=20$ tokens in each statement as the descriptive
statistics of the whole data set suggest that 95% of source code functions
have less than 155 statements and 95% of statements have less than 20 tokens.
To initialize our transformer encoders, we make use of the pre-trained model
provided by Wang et al. [37]. This model has undergone pre-training through
various denoising objectives associated with programming languages. Details of
the hyperparameter settings for our method in both phases are provided in
Appendix A.4. In both training phases, we train our model through specific
epochs and select the model that demonstrates the highest F1 score for
statement-level prediction in the validation set. The experiments were
conducted on a Linux machine with an AMD Ryzen 9 5950X processor, 64 GB of
RAM, and an NVIDIA RTX 3090 GPU. The potential limitations imposed by our
experimental setup are discussed in Appendix A.5.

###  4.3 Main Results

We conduct our experiments several times and report the average numbers. The
experimental data set and baseline methods are detailed in Section 4.1. We
report accuracy (Acc), precision (Pre), recall (Re), and F1-score (F1) for
function-level and statement-level vulnerability prediction tasks for a
comprehensive evaluation of each approach. This enables us to assess the
models’ performance on both positive and negative classes, regardless of the
class imbalance between vulnerable and benign functions. Note that the
statement-level metrics are computed on the statement level instead of the
function level to determine if each statement is correctly predicted. The
experimental results are shown in Table 1. Our approach yields an improvement
in function-level F1-score of 6% to 65% and an improvement in statement-level
F1-score of 19% to 71%.These results highlight the effectiveness of our
approach in accurately predicting vulnerabilities, both at the function and
statement levels, outperforming all other state-of-the-art methods.
Furthermore, our RNN statement embedding method significantly enhances the
performance of CodeBERT (30% $\rightarrow$ 63%) and CodeGPT (12% $\rightarrow$
44%) in statement-level vulnerability prediction. This finding validates our
intuition that the statement embeddings learned by our method can capture
contextual information and locate statements associated with vulnerabilities
more accurately than token embeddings.

Table 1: (Main Results) We compare our OptiMatch approach against other
baseline methods and present results in percentage.

|  | Function Level | Statement Level
---|---|---|---
Method | Embedding | Acc | Pre | Re | F1 | Acc | Pre | Re | F1
OptiMatch(ours) | Statement | 99.45 | 97.66 | 89.83 | 93.58 | 99.65 | 86.8 | 77.96 | 82.14
CodeBERT + our embedding | Statement | 98.91 | 92.15 | 82.89 | 87.28 | 99.19 | 59.39 | 67.84 | 63.33
CodeBERT | Token | 98.75 | 93.9 | 77.27 | 84.78 | 96.89 | 19.29 | 63.54 | 29.6
CodeGPT + our embedding | Statement | 98.95 | 91.25 | 84.81 | 87.91 | 98.23 | 32.54 | 67.34 | 43.88
CodeGPT | Token | 95.69 | 56.18 | 19.02 | 28.42 | 98.48 | 14.4 | 9.7 | 11.6
GraphCodeBERT | Token | 95.51 | 50.11 | 27.03 | 35.12 | 96.94 | 10.56 | 26.34 | 15.08
LineVul | Token | 98.61 | 89.25 | 78.47 | 83.51 | - | - | - | -
VELVET | Statement | 98.88 | 93.37 | 80.86 | 86.67 | 98.5 | 38.19 | 73.5 | 50.26
LineVD | Statement | - | - | - | - | 95.19 | 27.1 | 53.3 | 36
ReGVD | Token | 97.12 | 77.92 | 50.24 | 61.09 | - | - | - | -
Devign | Token | 96.9 | 72.29 | 50.24 | 59.28 | - | - | - | -
ICVH | Statement | 96.56 | 77.44 | 33.25 | 46.53 | 97.77 | 21.31 | 43.17 | 28.53
TextCNN | Statement | 95.95 | 62.31 | 25.12 | 35.81 | 98.15 | 21.03 | 28.91 | 24.34

###  4.4 Ablation Study

To assess the effectiveness of the proposed components in our OptiMatch
approach, we conduct an ablation study. Specifically, we compare our RNN
statement embedding method with mean or max pooling methods. Furthermore, we
examine the impact of our vulnerability codebook and matching by comparing our
approach with a variant that employs the same model architecture and pre-
trained weights, but without using the vulnerability codebook and matching.
Finally, we demonstrate the impact of the number of vulnerability centroids
(i.e., $k$) on the performance of our approach.

The experimental results are shown in Table 2. The utilization of mean or max
pooling to summarize token embeddings into statement embeddings results in a
slight decrease of 1.75% and 0.45% in function-level F1-score and 4.6% and
4.12% in statement-level F1-score, respectively, as compared to using an RNN.
The results confirm the effectiveness of our RNN statement embedding method,
indicating that it is more effective in summarizing token embeddings by
retaining token features at each time step. The performance significantly
deteriorates by 33.58% and 45.2% for function and statement-level predictions
when the vulnerability codebook and matching components are removed. This
underscores the importance of these components in achieving high-performance
levels. The results suggest that the vulnerability codebook plays a crucial
role in our approach, which is responsible for retaining and leveraging the
vulnerability patterns information present in vulnerable functions. This
information is then utilized to identify vulnerable statements effectively
during the vulnerability-matching inference. The lower section of Table 2
illustrates the impact of the number of vulnerability centroids on our
approach. The results demonstrate that our approach attains favorable
statement-level F1-scores for $k\in[100,150,200]$, and we set $k=150$ as it
produces the optimal statement-level F1-score. Notably, $k$ represents a
crucial factor, where a small value of $k$ (e.g., 50) may result in
unsatisfactory performance due to the grouping of too many vulnerability
patterns together, resulting in an inadequate representation of each pattern.
Conversely, a large value of $k$ (e.g., 400) leads to a substantial embedding
space of our codebook, making it challenging to update during the backward
process.

Table 2: (Ablation Results) We compare our proposed method to other variants
to investigate the impact of the individual components. The metrics are
reported as percentages.

| Function Level | Statement Level
---|---|---
Method | Acc | Pre | Re | F1 | Acc | Pre | Re | F1
OptiMatch (ours) | 99.45 | 97.66 | 89.83 | 93.58 | 99.65 | 86.8 | 77.96 | 82.14
w/o RNN embedding (mean pooling applied) | 99.31 | 98.49 | 86 | 91.83 | 99.59 | 90.4 | 67.89 | 77.54
w/o RNN embedding (max pooling applied) | 99.4 | 96.53 | 89.95 | 93.13 | 99.56 | 79.7 | 76.4 | 78.02
w/o vulnerability codebook & matching | 94.81 | 45.91 | 86.6 | 60 | 98.19 | 28.77 | 51.57 | 36.94
OptiMatch wt 50 vulnerability centroids | 85.9 | 23.95 | 98.21 | 38.51 | 95.5 | 16.92 | 86.13 | 28.28
OptiMatch wt 100 vulnerability centroids | 99.38 | 98.13 | 87.92 | 92.74 | 99.64 | 88.14 | 74.98 | 81.03
OptiMatch wt 150 vulnerability centroids (ours) | 99.45 | 97.66 | 89.83 | 93.58 | 99.65 | 86.8 | 77.96 | 82.14
OptiMatch wt 200 vulnerability centroids | 99.45 | 96.69 | 90.91 | 93.71 | 99.63 | 83.44 | 80.02 | 81.69
OptiMatch wt 400 vulnerability centroids | 98.28 | 99.05 | 62.32 | 76.51 | 99.54 | 81.91 | 70.47 | 75.76

##  5 Conclusion

This paper presents a novel vulnerability-matching method for function and
statement-level vulnerability detection (VD). Our approach capitalizes on the
vulnerability patterns present in vulnerable programs, which are typically
overlooked in deep learning-based VD. To be specific, we collect vulnerability
patterns from the training data and learn a more compact vulnerability
codebook from the pattern collection using optimal transport (OT) and vector
quantization. During inference, the codebook is utilized to match all learned
patterns and detect potential vulnerabilities within a given program. Our
comprehensive evaluation, conducted on over 188,000 real-world C/C++
functions, demonstrates that our method surpasses other competitive baseline
techniques, while our ablation study confirms the soundness of our approach.

## References

  * [1] Checkmarx.  https://checkmarx.com/.
  * [2] Cppcheck.  https://cppcheck.sourceforge.io/.
  * [3] Flawfinder.  https://dwheeler.com/flawfinder/.
  * [4] Rats.  https://code.google.com/archive/p/rough-auditing-tool-for-security/.
  * [5] Yoshua Bengio, Nicholas Léonard, and Aaron Courville.  Estimating or propagating gradients through stochastic neurons for conditional computation.  arXiv preprint arXiv:1308.3432, 2013.
  * [6] Harold Booth, Doug Rike, and Gregory Witte.  The national vulnerability database (nvd): Overview, 2013-12-18 2013.
  * [7] Saikat Chakraborty, Rahul Krishna, Yangruibo Ding, and Baishakhi Ray.  Deep learning based vulnerability detection: Are we there yet.  IEEE Transactions on Software Engineering, 2021.
  * [8] Yahui Chen.  Convolutional neural network for sentence classification.  Master’s thesis, University of Waterloo, 2015.
  * [9] Kyunghyun Cho, Bart van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio.  On the properties of neural machine translation: Encoder–decoder approaches.  In Proceedings of SSST-8, Eighth Workshop on Syntax, Semantics and Structure in Statistical Translation, pages 103–111, 2014.
  * [10] CWE Community.  2022 cwe top 25 most dangerous software weaknesses.  https://cwe.mitre.org/top25/archive/2022/2022_cwe_top25.html, 2022\.
  * [11] Roland Croft, Dominic Newlands, Ziyu Chen, and M Ali Babar.  An empirical study of rule-based and learning-based approaches for static application security testing.  In Proceedings of the 15th ACM/IEEE international symposium on empirical software engineering and measurement (ESEM), pages 1–12, 2021.
  * [12] Marco Cuturi.  Sinkhorn distances: Lightspeed computation of optimal transport.  In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013.
  * [13] Yangruibo Ding, Sahil Suneja, Yunhui Zheng, Jim Laredo, Alessandro Morari, Gail Kaiser, and Baishakhi Ray.  Velvet: a novel ensemble learning approach to automatically locate vulnerable statements.  In 2022 IEEE International Conference on Software Analysis, Evolution and Reengineering (SANER), pages 959–970. IEEE, 2022.
  * [14] Jiahao Fan, Yi Li, Shaohua Wang, and Tien N Nguyen.  Ac/c++ code vulnerability dataset with code changes and cve summaries.  In Proceedings of the 17th International Conference on Mining Software Repositories (MSR), pages 508–512, 2020.
  * [15] Zhangyin Feng, Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng, Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, et al.  Codebert: A pre-trained model for programming and natural languages.  In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 1536–1547, 2020.
  * [16] Jean Feydy, Thibault Séjourné, François-Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré.  Interpolating between optimal transport and mmd using sinkhorn divergences.  In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2681–2690. PMLR, 2019.
  * [17] Michael Fu and Chakkrit Tantithamthavorn.  Linevul: a transformer-based line-level vulnerability prediction.  In Proceedings of the 19th International Conference on Mining Software Repositories (MSR), pages 608–620, 2022.
  * [18] Daya Guo, Shuo Ren, Shuai Lu, Zhangyin Feng, Duyu Tang, LIU Shujie, Long Zhou, Nan Duan, Alexey Svyatkovskiy, Shengyu Fu, et al.  Graphcodebert: Pre-training code representations with data flow.  In International Conference on Learning Representations, 2021.
  * [19] David Hin, Andrey Kan, Huaming Chen, and M Ali Babar.  Linevd: statement-level vulnerability detection using graph neural networks.  In Proceedings of the 19th International Conference on Mining Software Repositories (MSR), pages 596–607, 2022.
  * [20] Nhat Ho, XuanLong Nguyen, Mikhail Yurochkin, Hung Hai Bui, Viet Huynh, and Dinh Phung.  Multilevel clustering via wasserstein means.  In International conference on machine learning, pages 1501–1509. PMLR, 2017.
  * [21] Joern.  Code property graph.  https://docs.joern.io/code-property-graph/, 2023.
  * [22] Yi Li, Shaohua Wang, and Tien N Nguyen.  Vulnerability detection with fine-grained interpretations.  In Proceedings of the 29th ACM Joint Meeting on European Software Engineering Conference and Symposium on the Foundations of Software Engineering, pages 292–303, 2021.
  * [23] Zhen Li, Deqing Zou, Shouhuai Xu, Hai Jin, Yawei Zhu, and Zhaoxuan Chen.  Sysevr: A framework for using deep learning to detect software vulnerabilities.  IEEE Transactions on Dependable and Secure Computing, 19(4):2244–2258, 2021.
  * [24] Zhen Li, Deqing Zou, Shouhuai Xu, Xinyu Ou, Hai Jin, Sujuan Wang, Zhijun Deng, and Yuyi Zhong.  Vuldeepecker: A deep learning-based system for vulnerability detection.  arXiv preprint arXiv:1801.01681, 2018.
  * [25] Shuai Lu, Daya Guo, Shuo Ren, Junjie Huang, Alexey Svyatkovskiy, Ambrosio Blanco, Colin Clement, Dawn Drain, Daxin Jiang, Duyu Tang, et al.  Codexglue: A machine learning benchmark dataset for code understanding and generation.  In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1).
  * [26] Tuan Nguyen, Trung Le, Nhan Dam, Quan Hung Tran, Truyen Nguyen, and Dinh Phung.  Tidot: A teacher imitation learning approach for domain adaptation with optimal transport.  In Zhi-Hua Zhou, editor, Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, IJCAI-21, pages 2862–2868. International Joint Conferences on Artificial Intelligence Organization, 8 2021.  Main Track.
  * [27] Van Nguyen, Trung Le, Olivier De Vel, Paul Montague, John Grundy, and Dinh Phung.  Dual-component deep domain adaptation: A new approach for cross project software vulnerability detection.  Pacific-Asia Conference on Knowledge Discovery and Data Mining, 2020\.
  * [28] Van Nguyen, Trung Le, Olivier De Vel, Paul Montague, John Grundy, and Dinh Phung.  Information-theoretic source code vulnerability highlighting.  In 2021 International Joint Conference on Neural Networks (IJCNN), pages 1–8. IEEE, 2021.
  * [29] Van Nguyen, Trung Le, Tue Le, Khanh Nguyen, Olivier DeVel, Paul Montague, Lizhen Qu, and Dinh Phung.  Deep domain adaptation for vulnerable code function identification.  In The International Joint Conference on Neural Networks (IJCNN), 2019.
  * [30] Van-Anh Nguyen, Dai Quoc Nguyen, Van Nguyen, Trung Le, Quan Hung Tran, and Dinh Phung.  Regvd: Revisiting graph neural networks for vulnerability detection.  In Proceedings of the ACM/IEEE 44th International Conference on Software Engineering: Companion Proceedings, pages 178–182, 2022.
  * [31] National Institute of Standards and Technology.  Apache struts vulnerability (cve-2021-31805).  https://nvd.nist.gov/vuln/detail/CVE-2021-31805, 2022.
  * [32] Nils Reimers and Iryna Gurevych.  Sentence-bert: Sentence embeddings using siamese bert-networks.  In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 3982–3992, 2019.
  * [33] Rebecca Russell, Louis Kim, Lei Hamilton, Tomo Lazovich, Jacob Harer, Onur Ozdemir, Paul Ellingwood, and Marc McConley.  Automated vulnerability detection in source code using deep representation learning.  In 2018 17th IEEE international conference on machine learning and applications (ICMLA), pages 757–762. IEEE, 2018.
  * [34] Rico Sennrich, Barry Haddow, and Alexandra Birch.  Neural machine translation of rare words with subword units.  In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1715–1725, 2016.
  * [35] Aaron Van Den Oord, Oriol Vinyals, et al.  Neural discrete representation learning.  Advances in neural information processing systems (NeurIPS), 30, 2017.
  * [36] Cédric Villani.  Optimal transport: Old and new.  2008\.
  * [37] Yue Wang, Weishi Wang, Shafiq Joty, and Steven CH Hoi.  Codet5: Identifier-aware unified pre-trained encoder-decoder models for code understanding and generation.  In the Proceedings of the Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 8696–8708, 2021.
  * [38] Yaqin Zhou, Shangqing Liu, Jingkai Siow, Xiaoning Du, and Yang Liu.  Devign: Effective vulnerability identification by learning comprehensive program semantics via graph neural networks.  Advances in neural information processing systems (NeurIPS), 32, 2019.

##  Appendix A Appendix

###  A.1 Self-Attention of Transformer Encoders

Given the input $H^{0}$, we leverage 12 layers of transformer encoders to
learn the representation of $x$ as follows:

$A^{t}=LN(MultiAttn(H^{t-1}))+H^{t-1},t\in\\{1,...,12\\}$ (8)
$H^{t}=LN(FFN(A^{t})+A^{t})$ (9)

where $A^{t}$ the self-attention output, $MultiAttn$ is a multi-head
attention, $FFN$ is feed-forward neural networks and $LN$ is layer
normalization.

###  A.2 Cross Attention for Selecting Vulnerability Centroids

$Q=\mathbf{v_{i}}\cdot W^{Q},K=C\cdot W^{K},V=C\cdot W^{V}$ (10)
$AttScore=Drop(\psi(Q\cdot K^{T})),AttScore\in\mathbb{R}^{1\times k}$ (11)
$\mathbf{c}_{\mathbf{v_{i}}}^{*}=argmax(AttScore)$ (12)

where query states $Q$ is obtained by linearly projected $\mathbf{v_{i}}$
using model parameters $W^{Q}\in\mathbb{R}^{h\times h}$, key states $K$ and
value states $V$ are obtained by linearly projected $C$ using model parameters
$W^{K},W^{V}\in\mathbb{R}^{h\times h}$. $\psi$ is a softmax function and
$Drop$ is a dropout layer. We use $argmax$ operation to obtain the codebook
embedding index and map $\mathbf{v_{i}}$ to the corresponding
$\mathbf{c}_{\mathbf{v_{i}}}^{*}$ having the maximum $AttScore$.

###  A.3 Details of Baseline Methods

We compare our OptiMatch approach with LLMs pre-trained on source code data,
state-of-the-art transformer-based, GNN-based, RNN-based, and CNN-based
vulnerability detection (VD) approaches. We reproduce each baseline based on
the code provided by the original authors.

LLMs for code: We include CodeBERT [15], CodeGPT [25], and GraphCodeBERT [18].
These models were pre-trained with token embeddings. We include an additional
trial for CodeBERT and CodeGPT using our RNN statement embedding method. Note
that the statement embedding is not compatible with GraphCodeBERT’s data flow
construction.

Transformer-based VD: LineVul [17] is designed to perform function-level
prediction by leveraging a pre-trained transformer model. Although it can also
provide statement-level predictions by interpreting and ranking the attention
scores of the transformer, this approach is not suitable for the statement-
level classification setting. To ensure a fair comparison, we only evaluate
our approach against LineVul on the function level. VELVET [13] is an ensemble
method that leverages a vanilla transformer with GNNs.

GNN-based VD: LineVD [19], ReGVD [30], and Devign [38] are GNN-based methods
that learn the graph property of source code. Note that ReGVD and Devign only
predict function-level vulnerabilities.

RNN-based and CNN-based VD: ICVH [28] leverages Bi-RNN with information theory
to detect statement-level vulnerabilities. ICVH was initially trained in the
unsupervised setting for statement-level vulnerability prediction, but we
found that it was not effective in our context. Therefore, we adopted the
original ICVH architecture and added a cross-entropy loss to train ICVH in the
supervised setting to achieve a fair comparison. On the other hand, TextCNN
[8] uses convolutional layers for sentence classification tasks.

###  A.4 Hyper-Parameter Settings of Our OptiMatch Approach

Table 3 lists the hyper-parameter settings required to reproduce our approach.
We have made our replication package available at
https://github.com/optimatch/optimatch, which includes all the experimental
scripts. We have included a comprehensive README file that contains all of the
details needed to reproduce the experimental results demonstrated in this
paper.

Table 3: The hyper-parameter settings of our OptiMatch approach.

Phase | Optimizer | Scheduler | LR | Grad Clip | Epochs | Stmt Len ($r$) | Max Num Stmt ($n$) | Num Centroids ($k$) | Batch
---|---|---|---|---|---|---|---|---|---
Warm-up | AdamW | Linear (4,650 warm-up steps) | 1e-4 | 1.0 | 20 | 20 tokens | 155 | - | 64
Main | AdamW | Linear (4,650 warm-up steps) | 1e-4 | 1.0 | 20 | 20 tokens | 155 | 150 | 64

###  A.5 Discussion

Inspired by program analysis tools for locating vulnerabilities based on
predefined vulnerability patterns, we proposed our innovative vulnerability-
matching deep-learning framework not only successfully utilizing optimal
transport and vector quantization for function and statement-level
vulnerability detection but also effectively leverage the information
presented in vulnerable statements and patterns to enhance deep learning-based
vulnerability detection. We found that the performance of our approach can be
affected by the chosen number of vulnerability centroids ($k$) used in the
codebook. While our approach has shown promising results in detecting
vulnerabilities in source code, the number of centroids is currently a
hyperparameter that requires manual tuning. This could be a limitation of our
approach when scaling up to larger datasets or more complex codebases, as it
may not be feasible to manually optimize the number of centroids for each
dataset. Thus, our future work should focus on developing automated methods
for selecting the optimal number of centroids or incorporating more advanced
techniques such as adaptive quantization to dynamically adjust the codebook
dictionary during training. Nevertheless, we conduct an ablation study to
reason the optimal solution in this paper. We assessed the effectiveness of
our method using the extensive Big-Vul dataset, which includes 188,000 C/C++
functions and 3,754 code vulnerabilities across 91 distinct CWE-IDs. This
dataset is expected to be comprehensive and inclusive, containing a range of
code patterns and vulnerabilities that are reflective of real-world scenarios.
Nevertheless, the performance of our method and the comparison baselines may
vary when tested on other datasets with different characteristics. We
acknowledge this limitation and the potential bias in our findings.