Source-linked AI summary
ReGVD: Revisiting Graph Neural Networks for Vulnerability Detection
Van-Anh Nguyen, Dai Quoc Nguyen, Van Nguyen, Trung Le, Quan Hung Tran, Dinh Phung
TL;DR
Source-code vulnerability detection is important but challenging, especially when methods depend on expert-crafted features or language-specific parsing. ReGVD constructs token-sequence graphs and applies residual GNN layers with graph-level pooling, reporting the highest accuracy, 63.69%, on the CodeXGLUE vulnerability-detection dataset.
Problem
Vulnerability detection is essential for protecting software systems, but existing approaches face burdensome feature engineering and parser limitations across source codes and programming languages.
Method
ReGVD treats vulnerability detection as inductive text classification, builds graphs from raw token sequences, initializes nodes with pre-trained programming-language token embeddings, and combines residual GNN layers with mixed sum-max pooling.
Results
63.69% accuracy is reported as ReGVD's highest result and a new state-of-the-art result on the CodeXGLUE vulnerability-detection dataset.
Takeaways & Limitations
ReGVD outperforms the up-to-date baseline models on the CodeXGLUE vulnerability-detection benchmark.
Takeaways & Limitations
Devign's parser-based graph construction is difficult to practice because no perfect parser reliably handles every programming language and varied source code without errors.
Abstract
from arXiv · showhide
Identifying vulnerabilities in the source code is essential to protect the software systems from cyber security attacks. It, however, is also a challenging step that requires specialized expertise in security and code representation. To this end, we aim to develop a general, practical, and programming language-independent model capable of running on various source codes and libraries without difficulty. Therefore, we consider vulnerability detection as an inductive text classification problem and propose ReGVD, a simple yet effective graph neural network-based model for the problem. In particular, ReGVD views each raw source code as a flat sequence of tokens to build a graph, wherein node features are initialized by only the token embedding layer of a pre-trained programming language (PL) model. ReGVD then leverages residual connection among GNN layers and examines a mixture of graph-level sum and max poolings to return a graph embedding for the source code. ReGVD outperforms the existing state-of-the-art models and obtains the highest accuracy on the real-world benchmark dataset from CodeXGLUE for vulnerability detection. Our code is available at: \url{https://github.com/daiquocnguyen/GNN-ReGVD}.
1 INTRODUCTION
Vulnerability detection is essential but difficult because vulnerabilities are increasing and effective detection often depends on specialized expertise and adaptable code representations. ReGVD addresses these challenges with a programming-language-independent GNN model that constructs graphs directly from raw token sequences.
- Vulnerability detection is essential because software vulnerabilities can cause substantial economic and social damage through cybersecurity attacks.
- Hand-engineered vulnerability-detection features require domain expertise, take time to create, and adapt poorly to evolving open-source code and libraries.
- NLP architectures and pre-trained programming-language models reduce reliance on manual feature engineering for source-code vulnerability detection.
- Graph neural networks provide practical low-dimensional representations for nodes and graphs and have achieved strong results in downstream text-classification tasks.
- Devign uses a programming-language parser to extract multi-edged graphs, but reliable parsers are unavailable for every language and diverse source code.
- ReGVD treats vulnerability detection as inductive text classification and builds graphs from flat token sequences, using token embeddings, residual GNN layers, and mixed sum-max pooling.
2 THE PROPOSED REGVD
ReGVD formulates function-level vulnerability detection as inductive graph classification by converting token sequences into graphs and applying GNNs with residual connections and mixed graph-level pooling.
- 2.1 Problem definition: ReGVD treats each source-code function as a binary vulnerability-classification instance and learns a graph-to-label mapping.Labels are 1 for vulnerable and 0 otherwise.
- 2.2 Graph construction: Raw source code is tokenized into a graph whose nodes represent either unique tokens or token positions, with sliding-window co-occurrences forming edges.The implementation first applies the corresponding pre-trained PL tokenizer before constructing the graph.
- 2.2 Graph construction: Node features are initialized using only the token embedding layer of a pre-trained programming-language model.This design retains pre-trained representations while supporting the paper’s comparison setting.
- 2.3 Graph neural networks with residual connection: GNN layers recursively aggregate neighbor representations, using GCNs or GGNNs as ReGVD’s base models.The initial node representation is H(0) = X.
- 2.3 Graph neural networks with residual connection: Residual connections carry lower-layer information into higher layers and provide direct gradient paths, while keeping hidden sizes fixed across layers.ReGVD adapts residual connections among its GNN layers to address vanishing- or exploding-gradient problems.
- 2.4 Graph-level readout pooling layer: ReGVD forms graph embeddings by mixing sum and max pooling, examining SUM, MUL, and CONCAT operators before softmax classification.The resulting embedding is passed through a fully connected layer and softmax to predict vulnerability.
3 EXPERIMENTAL SETUP AND RESULTS
Experiments on the CodeXGLUE function-level vulnerability dataset evaluate ReGVD against strong baselines and multiple graph settings. ReGVD achieves the strongest reported accuracy, while residual connections, graph-level readout, and limited-data training show consistent benefits.
- Experimental setup: 27,318 manually labeled functions from QEMU and FFmpeg form the real-world CodeXGLUE benchmark used for function-level vulnerability detection.The dataset contains vulnerable and non-vulnerable C functions extracted from security-related commits.
- Experimental setup: ReGVD is compared with BiLSTM, TextCNN, RoBERTa, Devign, CodeBERT, and GraphCodeBERT using the stated test-set accuracy protocol.Devign is reimplemented with the same training and evaluation protocols because its official implementation was not released.
- Main results: 63.69% is ReGVD’s highest accuracy, a new state-of-the-art result on the CodeXGLUE vulnerability detection dataset.The paper reports this as higher than the up-to-date baseline models.
- Main results: 2.20% higher accuracy is obtained by ReGVD than Devign in the matched GGNN+Idx+G-CB comparison, with accuracies of 63.51% and 61.31%, respectively.Across the remaining settings without residual connections, ReGVD’s graph-level readout also outperforms Devign’s pooling layer.
- Main results: 2.05% is the maximum accuracy gain attributed to residual connections, which improve ReGVD performance on seven settings.Across the mixture functions, MUL gives the highest accuracy on six settings and SUM on two remaining settings.
- Main results: 61.68% accuracy with 60% of the training set exceeds BiLSTM, TextCNN, RoBERTa, and Devign trained with all data.With 80% of the training set, ReGVD reaches 62.55%, exceeding CodeBERT and GraphCodeBERT.
4 CONCLUSION
ReGVD treats vulnerability identification as inductive text classification and represents raw source code as graphs initialized from pre-trained programming-language token embeddings. Its residual GNN layers and mixed sum/max pooling produce the highest reported accuracy on the CodeXGLUE benchmark, supporting a general, practical, programming-language-independent model.
- Conclusion: ReGVD frames vulnerability identification as inductive text classification and constructs a graph from each raw source-code sequence.Node features use only the token embedding layer of a pre-trained programming-language model.
- Conclusion: Residual connections among GNN layers and mixed sum/max pooling are used to learn graph representations for vulnerability detection.These design choices are part of ReGVD’s graph-based source-code representation.
- Conclusion: 63.69% accuracy is the highest reported result on the CodeXGLUE vulnerability detection benchmark.The experiments compare ReGVD with strong and up-to-date baseline models.
- Conclusion: ReGVD is presented as general, practical, and programming-language-independent for running on various source codes and libraries.The conclusion states this capability without requiring difficulty-free parsing of a particular programming language.