Software Vulnerability Detection Using a Lightweight Graph Neural Network
Authors: Miles Farmer, Ekincan Ufuktepe, Anne Watson, Hialo Muniz Carvalho, Vadim Okun, Zineb Maasaoui, Kannappan Palaniappan
Organizations: Department of Electrical Engineering and Computer Science at the University of Missouri, MO, USA. · Software and Systems Division at the National Institute of Standards and Technology, Gaithersburg, MD, USA
Large Language Models (LLMs) have emerged as a popular choice in vulnerability detection studies given their foundational capabilities, open source availability, and variety of models, but have limited scalability due to extensive compute requirements. Using the natural graph relational structure of code, we show that our proposed graph neural network (GNN) based deep learning model VulGNN for vulnerability detection can achieve performance almost on par with LLMs, but is 100 times smaller in size and fast to retrain and customize. We describe the VulGNN architecture, ablation studies on components, learning rates, and generalizability to different code datasets. As a lightweight model for vulnerability analysis, VulGNN is efficient and deployable at the edge as part of real-world software development pipelines.
Figures & tables
Fig. 1 : VulGNN graph convolutional deep architecture with preprocessing (including language tokenization) to generate the input Code Property Graph (Row 1), the internal details of the graph neural network layers (Row 2), and the implementation of the GeneralConv operator with graph attention (Row 3). The output from the GNN after Global Mean Pooling is passed to a single linear classification layer. VulGNN does not include a hidden MLP layer after pooling.
Fig. 2 : Hybrid testing with different amounts of real-world data mixed with synthetic NIST Juliet Test Suite vulnerability data. VulGNN accuracy and F1-score (left graph) when trained on different percentages of real-world data randomly sampled from the training set of DiverseVul, combined with a baseline of synthetic Juliet Test Suite [ 32 ] which is a special subset of the NIST SARD/Juliet vulnerable code collection. That is, 0% is trained on a 33,360 sample subset of the Juliet dataset, while 100% is trained on the Juliet subset plus 419,164 training samples from DiverseVul. The right graph shows speed of training convergence using the DiverseVul validation set, and learning rate curves showing improvement with increasing percentage of real-world data.
Fig. 3 : Accuracy and F1 Score test results on different Vul:Non-Vul training ratios
Family
Architecture
Train on Prev+Diverse, Test on Unseen Projects
Train (80%) & Test (10%) on Prev+Diverse (no overlap)
F1
Acc
Prec
Rec
FPR
F1
Acc
Prec
Rec
FPR
GNN
ReVeal ∗
8.67
85.88
5.67
18.46
11.58
29.76
82.81
23.75
39.83
12.87
VulGNN (Ours)
18.17
93.17
16.18
20.72
4.08
40.35
89.48
36.84
44.6
6.63
BERT
RoBERTa*
4.40
95.59
10.46
2.78
0.90
34.98
91.68
46.02
28.22
2.85
CodeBERT ∗
11.94
94.19
13.34
10.80
2.65
37.85
90.48
39.25
36.54
4.87
CodeBERT †
14.74
92.16
12.21
18.60
5.06
41.72
89.39
36.97
47.89
7.04
TABLE I : Vulnerability detection performance of deep learning models for two scenarios. The best performing result in each column is shown in bold. Results in the “Train on Prev+Diverse, Test on Unseen Projects” column represent results from the experiment described in Section IV-D . Results in the “Train & Test on Prev+Diverse” column represent results from the experiment described in Section IV-C . All results, except for VulGNN, are from DiverseVul [ 5 ] (see their Tables 4 and 5) (test on same and unseen data respectively; marked with ∗ ), and Table 6 (optimized hyperparameters; marked with † ). See the aforementioned sections for further details on dataset split and experiment setup.
Family
Architecture
Number of Parameters
GNN
ReVeal
1.28 Million
VulGNN
1.10 Million
BERT
RoBERTa
125 Million
CodeBERT
125 Million
GraphCodeBERT
125 Million
GPT-2
GPT-2 Base
117 Million
TABLE II : The number of learnable weight parameters for each network.