Predicting Code Vulnerability Types via Heterogeneous GNN Learning

Graph Neural Network (GNN)-based techniques have demonstrated significant success in detecting code vulnerability. Due to the great variety of code vulnerabilities, the binary prediction outcomes offer limited support for accurate debugging and pinpointing vulnerabilities across different types. Additionally, the prevalent use of homogeneous GNN training has constrained the ability to effectively model the diverse relationships existing between code elements that contribute to various vulnerabilities. This paper introduces an innovative approach using a heterogeneous GNN framework for the prediction of code vulnerability types. Our method begins with the construction of Inter-Procedural Compressed Code Property Graphs (ICCPGs), designed to capture the structural and contextual attributes of code components. Then, we establish a heterogeneous GNN model to learn features by leveraging the diverse relationships within ICCPGs, thereby enabling the classification of vulnerabilities into specific types. The proposed approach achieved an accuracy of 92.8% in the C language, covering up to 108 different vulnerability types, and the top-2 classification further improved this score to 95.5%. This performance surpasses the previous state-of-the-art methods, marking a notable 4.6% improvement over the best prior results. Additionally, the method also demonstrated strong performance in Java with 90 types of vulnerabilities, achieving 94.1% accuracy, which is further enhanced to 96.7% with top-2 classification.