Code Vulnerability Detection via Nearest Neighbor Mechanism (2022)

Abstract: Code vulnerability detection is a fundamental and challenging task in the software security field. Existing research works aim to learn semantic information from the source code by utilizing NLP technologies. However, in vulnerability detection tasks, some vulnerable samples are very similar to non-vulnerable samples, which are difficult to identify. To address this issue and improve detection performance, we introduce the k -nearest neighbor mechanism which retrieves multiple neighbor samples and utilizes label information of retrieved neighbor samples to provide help for model predictions. Besides, we use supervised contrastive learning to make the model learn the discriminative representation and ensure that label information of retrieved neighbor samples is as consistent as possible with the label information of testing samples. Extensive experiments show that our method can achieve obvious performance improvements compared to baseline models.

## introduction code vulnerabilities generally denote security bugs or weaknesses in software. the presence of vulnerabilities in software is an inevitable problem because of design flaws of the program language and faults caused by programmers. traditional detection methods, such as static methods @xcite or dynamic methods @xcite @xcite , generally require a lot of human labor to summarize vulnerability rules or massive computational resources to trigger potential vulnerabilities by executing programs. to improve detection efficiency, there are multiple efforts @xcite @xcite have attempted to introduce deep learning and nlp techniques for automated vulnerability detection. deep learning-based detectors can extract semantic features from the source code and automatically learn potential vulnerability patterns. @xcite proposed a vulnerability detection framework named vuldeepecker, which used a binary classifier based on lstm to detect whether a piece of code is vulnerable or not. in order to learn comprehensive program semantics to characterize vulnerabilities of high diversity and complexity in the source code, @xcite utilized the graph neural network to extract code semantic representations. inspired by the big success of pre-training models on numerous natural language processing (nlp) tasks @xcite , code pre-training models for programming languages, such as codebert @xcite and graph-codebert (guo et al., 2020), were proposed and applied to a variety of code-related tasks such as code search, code completion, and code vulnerability detection, etc. although the above methods have achieved remarkable improvements, they still easily generate a large number of false predictions because of the relatively large classification uncertainty. in vulnerability detection tasks, vulnerabilities are sometimes subtle flaws. fig. 1 illustrates the comparison between the vulnerable and nonvulnerable code snippets. we can find that the difference between vulnerable and non-vulnerable samples is very small. this means that some vulnerable and non-vulnerable samples are close to the classifier boundary, which has the large classification uncertainty and easily results in false positive/negative predictions (as depicted in fig. 2 ). to alleviate this problem and improve classification performance, we argue that fully exploiting the label information of samples with the same class can assist the model to classify samples more accurately. that is, when predicting the label of a sample, the model can obtain a good guideline from its neighbor samples. motivated by this obervation, we present a novel code vulnerability detection method by introducing the k-nearest neighbor mechanism, which could fully utilize the label information of neighbor samples to perform more accurate vulnerability classification. specifically, it firstly retrieves multiple neighbor samples according to the representation of the sample and then integrate their labels into the final prediction of the model. furthermore, in order to retrieve more similar samples and make the label information of retrieved neighbor samples more effective, we introduce supervised contrastive learning (scl) to learn discriminative code representations. supervised contrastive learning (scl) could pull samples with the same class closer to each other and push away samples with the different classes. in short, our contributions are summarized as follows: ## methodology in this section, we will elaborate the details of the proposed method. the overview of the proposed method is illustrated in fig. 3 .