SIFT: enhance the performance of vulnerability detection by incorporating structural knowledge and multi-task learning

Abstract

Software vulnerabilities pose significant risks to software systems, leading to security breaches, data loss, operational disruptions, and substantial financial damage. Therefore, accurately detecting these vulnerabilities is of paramount importance. In recent years, pre-trained language models (PLMs) have demonstrated powerful capabilities in code representation and understanding, emerging as a promising method for vulnerability detection. However, integrating code structure knowledge while fine-tuning PLMs remains a significant challenge. To alleviate this limitation, we propose a novel vulnerability detection approach called SIFT. SIFT extracts the code property graph (CPG) to serve as the source of graph structural information. It constructs a code structure matrix from this information and measures the difference between the code structure matrix and the attention matrix using Sinkhorn Divergence to obtain the structural knowledge loss. This structural knowledge loss is then used alongside the cross-entropy loss for vulnerability detection in a multi-task learning framework to enhance overall detection performance. To evaluate the effectiveness of SIFT, we conducted experiments on three vulnerability detection datasets: FFmpeg+Qemu, Chrome+Debian, and Big-Vul. The results demonstrate that SIFT outperforms nine state-of-the-art vulnerability detection baselines, achieving performance improvements of 1.74%, 10.19%, and 2.87% in terms of F1 score, respectively. Our study shows the effectiveness of incorporating structural knowledge and multi-task learning in enhancing the performance of PLMs for vulnerability detection.