UPDF AI

ReGVD: Revisiting Graph Neural Networks for Vulnerability Detection

Van-Anh Nguyen,D. Q. Nguyen,3 Auteurs,Dinh Q. Phung

2021 · DOI: 10.1145/3510454.3516865
158 citations

TLDR

This work considers vulnerability detection as an inductive text classification problem and proposes ReGVD, a simple yet effective graph neural network-based model for the problem, which outperforms the existing state-of-the-art models and obtains the highest accuracy on the real-world benchmark dataset from CodeXGLUE for vulnerability detection.

Résumé

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: https://github.com/daiquocnguyen/GNN-ReGVD.