Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Software Vulnerability Detection Using a Lightweight Graph Neural Network

Miles Farmer, Ekincan Ufuktepe, Anne Watson, Hialo Muniz Carvalho, Vadim Okun, Zineb Maasaoui, Kannappan Palaniappan

Abstract

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.

Software Vulnerability Detection Using a Lightweight Graph Neural Network

Abstract

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.

Paper Structure

This paper contains 35 sections, 16 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Background and Fundamentals
  3. Code Graph Representations
  4. Abstract Syntax Trees
  5. Control-Flow Graphs
  6. Program Dependency Graphs
  7. Code Property Graphs
  8. Sentence Embedding Models
  9. Deep Learning Model Architecture and Parameters
  10. Overview
  11. Graph Representation
  12. Input Encoding
  13. Convolutional Blocks
  14. Graph Readout and Classification Head
  15. Training Objective
  16. ...and 20 more sections

Figures (3)

  • Figure 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.
  • Figure 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 vulcnn2022 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.
  • Figure 3: Accuracy and F1 Score test results on different Vul:Non-Vul training ratios