A vulnerability detection algorithm based on residual graph attention networks for source code imbalance (RGAN)

Abstract

It is essential to detect potential vulnerabilities in software to ensure its safety. As software systems become more complex, traditional static vulnerability detection methods perform poorly. Currently, deep learning-based vulnerability detection models only extract source code vulnerability features using sequences or graphs. Sequential neural networks ignore structural information in the code, such as control flow diagrams and data flow diagrams. Additionally, graph neural networks cannot accurately extract features due to the lack of effective methods for extracting nodes’ features and aggregating global information. To address the above issue, we propose a vulnerability detection algorithm based on residual graph attention networks for source code imbalance (RGAN). Firstly, a local feature extraction module (PE-BL-A module) is designed. Using the sequence neural network, the module extracts various useful features, including node features in a control flow diagram based on local semantic features. Secondly, we present the Residual Graph Attention Network module (RGAT). To learn and update node features along the control flow direction, the module uses a graph attention network with residual connections. In this module, a mean biaffine attention pooling mechanism is proposed that can extract total graph vulnerability features more effectively. Thirdly, a dynamic cross-entropy loss function is designed. Using this function, it can handle sample imbalances during training. Finally, experiments conducted on several benchmark datasets demonstrate that the proposed model achieves state-of-the-art results.