StagedVulBERT: Multi-Granular Vulnerability Detection with a Novel Pre-trained Code Model

TL;DR

This study introduces StagedVulBERT, a novel vulnerability detection framework that leverages a pre-trained code language model and employs a coarse-to-fine strategy, which enhances the performance of vulnerability detection at both coarse- and fine-grained levels.

Abstract

The emergence of pre-trained model-based vulnerability detection methods has significantly advanced the field of automated vulnerability detection. However, these methods still face several challenges, such as difficulty in learning effective feature representations of statements for fine-grained predictions and struggling to process overly long code sequences. To address these issues, this study introduces StagedVulBERT, a novel vulnerability detection framework that leverages a pre-trained code language model and employs a coarse-to-fine strategy. The key innovation and contribution of our research lies in the development of the CodeBERT-HLS component within our framework, specialized in hierarchical, layered, and semantic encoding. This component is designed to capture semantics at both the token and statement levels simultaneously, which is crucial for achieving more accurate multi-granular vulnerability detection. Additionally, CodeBERT-HLS efficiently processes longer code token sequences, making it more suited to real-world vulnerability detection. Comprehensive experiments demonstrate that our method enhances the performance of vulnerability detection at both coarse- and fine-grained levels. Specifically, in coarse-grained vulnerability detection, StagedVulBERT achieves an F1 score of 92.26%, marking a 6.58% improvement over the best-performing methods. At the fine-grained level, our method achieves a Top-5% accuracy of 65.69%, which outperforms the state-of-the-art methods by up to 75.17%.

Figures (14)

• Figure 1: A real-world example (i.e., CVE-2019-13304) of stack-based buffer overflow, which occurs from writing excess data to an array due to improper buffer size check.
• Figure 2: Workflow of the proposed StagedVulBERT, which includes two steps: pre-training and fine-tuning. The Workflow of CodeBERT-HLS is illustrated at the top of the figure.
• Figure 3: The overall framework of the CodeBERT-HLS method, which is composed of the following three components: TETransformer, Token2Statement embedding, and SETransformer.
• Figure 4: Initialization of statement vector representations based on token vectors and statement-token correspondence
• Figure 5: Initializing statement vector representations using token vectors and statement-token correspondence, while accounting for the varied importance of tokens within statements