Cross-project vulnerability detection method based on transfer learning and feature fusion

Abstract

The invention provides a cross-project vulnerability detection method based on transfer learning and feature fusion, and belongs to the technical field of computers. The problem that a feature extraction part in a cross-project vulnerability detection model lacks expert features is solved, and the marked source project data is better utilized to train the model. According to the technical scheme, the method comprises the following steps that S1, a data set is formed; s2, extracting expert features by using an Understand tool; s3, using an oversampling method SMOTE to process a training data class imbalance problem; s4, respectively constructing models based on expert features and semantic features; s5, fusing the two models through a weighted average method; and S6, detecting whether new codes in the target project contain vulnerabilities or not by using the fused model. The cross-project vulnerability detection method has the beneficial effects that the expert features and the semantic features are fused, and the accuracy of cross-project vulnerability detection can be improved only by marking a small amount of data in the target project.

Description

Technical Field

[0001] The present invention relates to the field of computer technology, and in particular to a cross-project vulnerability detection method based on transfer learning and feature fusion. Background Art

[0002] Software vulnerabilities refer to errors, defects or weaknesses in computer programs or operating systems that may cause unexpected behavior or security vulnerabilities in the program under certain conditions. During software development and maintenance, if developers can identify vulnerable modules in advance, they can allocate testing resources more reasonably, thereby improving the quality and reliability of the software. When performing vulnerability detection on a new project, a large amount of labeled data is required for training. However, labeling data requires a lot of time for domain experts. A feasible solution is to use labeled data from old projects to assist new projects in vulnerability detection, that is, cross-project software vulnerability detection (CSVD).

[0003] In previous work, semantic features were mainly relied on to build CSVD models. However, these studies overlooked another important type of indicators, which are expert indicators manually designed by domain experts. The analysis perspectives of these two types of indicators, expert features and semantic features, are not exactly the same. Researchers have found that expert features and semantic features may be complementary. Therefore, it is necessary to study a method that can fuse the two features to better solve the problem of cross-project vulnerability prediction. In previous work, researchers mainly used labeled data in the source project and unlabeled data in the target project to build CSVD models. However, there is still a large difference in data distribution between the source project and the target project. If a small amount of data in the target project is labeled, the performance of the CSVD model may be effectively improved.

[0004] How to solve the above technical problems becomes the subject faced by the present invention. Summary of the invention

[0005] The purpose of the present invention is to provide a cross-project vulnerability detection method based on transfer learning and feature fusion. The method can detect whether the function code in the target project contains vulnerabilities.

[0006] The idea of ​​the present invention is: the present invention proposes a cross-project vulnerability detection method based on transfer learning and feature fusion, that is, by fusing features of two different dimensions, semantic features and expert features, and building a cross-project vulnerability detection model based on the transfer learning algorithm TrAdaBoost, using the labeled data of the source project and a small amount of labeled data in the target project for training, so as to more accurately detect whether the target project function code contains vulnerabilities. The method of the present invention achieves better performance than the benchmark method on the same data set.

[0007] The present invention is achieved by the following measures: a cross-project vulnerability detection method based on transfer learning and feature fusion, which includes the following steps:

[0008] (1) Collect the C language source code in the source project and the target project. The labels of all functions in the source project (i.e., whether they contain vulnerabilities) have been marked by security experts. Then, a small number of functions in the target project are selected and assigned to experts for labeling.

[0009] (2) Use the CodeBERT pre-trained model to extract the semantic features of the function codes in the source project and the target project, and use BERT-Whitening to reduce the dimension of the semantic features. Use the Understand tool to extract the expert features of the function codes in the source project and the target project, and finally obtain the training dataset D. The format of the dataset is set to <code semantic feature value, code expert feature value, whether it contains vulnerabilities>. The specific steps are as follows:

[0010] (2-1) For a given function code segment, first split it according to the camel case naming rule to obtain the input sequence;

[0011] (2-2) Input the processed sequence into the CodeBERT pre-trained model, extract the hidden states of the first and last layers in the output, and take their average to obtain the semantic feature vector value of the function;

[0012] (2-3) Use BERT-Whitening to process the semantic vector, extract key features and reduce dimension through linear transformation, and obtain the final semantic feature value vector;

[0013] (2-4) 39 expert features were identified from five dimensions: code length, code complexity, code readability, code maintainability, and code executableness. The Understand tool was used to extract the values ​​of these 39 expert features from the function code and form an expert feature value vector.

[0014] (3) For the training dataset D, the oversampling method SMOTE is used to deal with the class imbalance problem in the dataset, that is, the number of functions containing vulnerabilities is much smaller than the number of functions without vulnerabilities;

[0015] (4) Use the transfer learning method TrAdaBoost to establish a cross-project vulnerability detection model M based on semantic features S and a cross-project vulnerability detection model M based on expert features E , the specific steps are as follows:

[0016] (4-1) To build a cross-project vulnerability detection model M based on semantic features S , using the transfer learning algorithm TrAdaBoost to reduce the weight of the data in the source project that has low relevance to the target project data and increase the weight of the data in the source project that has high relevance to the target project data, so that an effective cross-project vulnerability detection model M for semantic features can be trained S ;

[0017] (4-2) Cross-project vulnerability detection model M based on expert features E , using the transfer learning algorithm TrAdaBoost to reduce the weight of the data in the source project that has low relevance to the target project data and increase the weight of the data in the source project that has high relevance to the target project data, so that an effective cross-project vulnerability detection model M can be trained for expert features. E .

[0018] (5) Use the weighted summation method to fuse the features of the two dimensions at the model level to form the final fusion detection model M Final , the specific steps are as follows:

[0019] (5-1) For the function code that needs vulnerability detection, a cross-project vulnerability detection model M based on semantic features is used S Get the probability P that it contains a vulnerability Ms ;

[0020] (5-2) Then the cross-project vulnerability detection model M based on expert features is used E Get the probability P that it contains a vulnerability Me ;

[0021] (5-3) Through the weighted summation formula P final =P Ms ×W Ms +P Me ×W Me Calculate the fusion model M Final The probability P of containing a vulnerability final , where W Ms and W MeThey are cross-project vulnerability detection models M S and M E The weight ratio in the fusion model.

[0022] (6) When faced with new unlabeled function codes in the target project, semantic features and expert features are extracted from the codes and a vector representation corresponding to the codes is formed. Finally, the fusion prediction model M is used Final To predict whether the function code contains a vulnerability.

[0023] The optimal parameter values ​​of the cross-project vulnerability detection method based on transfer learning and feature fusion are set as follows:

[0024] Dimension K of BERT-Whitening operation: 256;

[0025] Cross-project vulnerability detection model M based on semantic features S The weight coefficient W Ms : 0.6;

[0026] Cross-project vulnerability detection model M based on expert features E The weight coefficient W Me :0.4.

[0027] Compared with the prior art, the beneficial effects of the present invention are as follows: the cross-project vulnerability detection method based on transfer learning and feature fusion proposed by the present invention not only considers the semantic features of the code, but also considers the expert features of the code, uses the pre-trained model CodeBERT that can better represent the semantics of the code, and also uses the Understand tool to extract expert features from the code. In addition, the semantic features and expert features are fused, and the cross-project vulnerability detection model is established using the transfer learning method TrAdaBoost, so that the final detection result is more accurate. BRIEF DESCRIPTION OF THE DRAWINGS

[0028] The accompanying drawings are used to provide further understanding of the present invention and constitute a part of the specification. They are used to explain the present invention together with the embodiments of the present invention and do not constitute a limitation of the present invention.

[0029] Figure 1 A system framework diagram of a cross-project vulnerability detection method based on transfer learning and feature fusion provided by the present invention. DETAILED DESCRIPTION

[0030] In order to make the purpose, technical solution and advantages of the present invention more clearly understood, the present invention is further described in detail below in conjunction with the accompanying drawings and embodiments. Of course, the specific embodiments described here are only used to explain the present invention and are not used to limit the present invention.

[0031] Example 1

[0032] See also Figure 1 As shown, this embodiment provides a cross-project vulnerability detection method based on transfer learning and feature fusion, which specifically includes the following contents:

[0033] (1) Collect the C language source code in the source project and the target project. The labels of all functions in the source project (i.e., whether they contain vulnerabilities) have been marked by security experts. Then select a small number of functions from the target project and arrange for experts to label them;

[0034] (2) Use the CodeBERT pre-trained model to extract the semantic features of the function codes in the source project and the target project, and use BERT-Whitening to reduce the dimension of the semantic features. Use the Understand tool to extract the expert features of the function codes in the source project and the target project, and finally obtain the training dataset D. The format of the dataset is set to <code semantic feature value, code expert feature value, whether it contains vulnerabilities>. The specific steps are as follows:

[0035] (2-1) For a given function code segment, first split it according to the camel case naming rule to obtain the input sequence;

[0036] (2-2) Input the processed sequence into the CodeBERT pre-trained model, extract the hidden states of the first and last layers in the output, and take their average to obtain the semantic feature vector value of the function;

[0037] (2-3) Use BERT-Whitening to process the semantic vector, extract key features and reduce dimension through linear transformation, and obtain the final semantic feature value vector;

[0038] (2-4) 39 expert features were determined from the five dimensions of code length, code complexity, code readability, code maintainability, and code executableness. See Table 1 for details. The Understand tool was used to extract the values ​​of these 39 expert features from the function code and form an expert feature value vector.

[0039] Table 1 39 expert features used in the method of this embodiment

[0040]

[0041]

[0042] (3) For the training dataset D, the oversampling method SMOTE is used to deal with the class imbalance problem in the dataset, that is, the number of functions containing vulnerabilities is much smaller than the number of functions without vulnerabilities;

[0043] (4) Use the transfer learning method TrAdaBoost to establish a cross-project vulnerability detection model M based on semantic features S and a cross-project vulnerability detection model M based on expert features E , the steps are as follows:

[0044] (4-1) To build a cross-project vulnerability detection model M based on semantic features S , using the transfer learning algorithm TrAdaBoost to reduce the weight of the data in the source project that has low relevance to the target project data and increase the weight of the data in the source project that has high relevance to the target project data, so that an effective cross-project vulnerability detection model M for semantic features can be trained S ;

[0045] (4-2) Cross-project vulnerability detection model M based on expert features E , using the transfer learning algorithm TrAdaBoost to reduce the weight of the data in the source project that has low relevance to the target project data and increase the weight of the data in the source project that has high relevance to the target project data, so that an effective cross-project vulnerability detection model M can be trained for expert features. E .

[0046] (5) Use the weighted summation method to fuse the features of the two dimensions at the model level to form the final fusion detection model M Final , the specific steps are as follows:

[0047] (5-1) For the function code that needs vulnerability detection, a cross-project vulnerability detection model M based on semantic features is used S Get the probability P that it contains a vulnerability Ms ;

[0048] (5-2) Then the cross-project vulnerability detection model M based on expert features is used E Get the probability P that it contains a vulnerability Me ;

[0049] (5-3) Through the weighted summation formula P final =P Ms ×W Ms +P Me ×W Me Calculate the fusion model M Final The probability P of containing a vulnerability final , where W Ms and W Me They are cross-project vulnerability detection models M S and M E The weight ratio in the fusion model.

[0050] (6) When faced with new unlabeled function codes in the target project, semantic features and expert features are extracted from the codes and a vector representation corresponding to the codes is formed. Finally, the fusion prediction model M is used Final To predict whether the function code contains vulnerabilities;

[0051] (7) The optimal parameter values ​​of this method are set as follows:

[0052] Dimension K of BERT-Whitening operation: 256;

[0053] Cross-project vulnerability detection model M based on semantic features S The weight coefficient W Ms : 0.6;

[0054] Cross-project vulnerability detection model M based on expert features E The weight coefficient W Me :0.4.

[0055] (8) The method of this embodiment and existing cross-project vulnerability detection methods (including Vuldeepecker, Dual-GD-DDAN, DAM2P and DTB) are evaluated on the same dataset, and three performance indicators (i.e., AUC, Popt, and R@20%E) from the cross-project vulnerability detection research field are used to automatically evaluate the quality of the model:

[0056] Table 2 Comparison of results between the method of this embodiment and other methods

[0057]

[0058] Experiments have shown that the cross-project vulnerability detection method based on transfer learning and feature fusion proposed in this embodiment can perform more effective cross-project vulnerability detection compared to the baseline method. Specifically, the method of this embodiment combines the expert features and semantic features of the code, and uses the transfer learning algorithm TrAdaBoost to surpass these baseline methods. Among them, for AUC, the method of this embodiment improves the performance by 7.48%-24.6%; for Popt, the method of this embodiment improves the performance by 2.3%-16.9%; for R@20%E, the method of this embodiment improves the performance by 3.43%-23%. These results show the effectiveness and competitiveness of the method proposed in this embodiment.

[0059] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principle of the present invention should be included in the protection scope of the present invention.

Claims

1. A cross-project vulnerability detection method based on transfer learning and feature fusion, characterized in that: The following steps are involved: (1) Collect the C language source code in the source project and the target project. All functions in the source project have been labeled by security experts. Then select a small number of functions from the target project and arrange for experts to label them. (2) Use the CodeBERT pre-trained model to extract the semantic features of the function codes in the source project and the target project, use BERT-Whitening to reduce the dimension of the semantic features, and use the Understand tool to extract the expert features of the function codes in the source project and the target project to obtain the training dataset D. The format of the dataset is set to <code semantic feature value, code expert feature value, whether it contains vulnerabilities>; (3) For the training dataset D, the oversampling method SMOTE is used to deal with the class imbalance problem in the dataset. The number of functions containing vulnerabilities is smaller than the number of functions without vulnerabilities. (4) Use the transfer learning method TrAdaBoost to establish a cross-project vulnerability detection model M based on semantic features S and a cross-project vulnerability detection model M based on expert features E ; (5) Use the weighted summation method to fuse the features of the two dimensions at the model level to form the final fusion detection model M Final ; (6) When faced with new unlabeled function codes in the target project, semantic features and expert features are extracted from the codes and a vector representation corresponding to the codes is formed. Finally, the fusion prediction model M is used Final To predict whether the function code contains a vulnerability.

2. According to claim 1, the cross-project vulnerability detection method based on transfer learning and feature fusion is characterized in that: The extraction of semantic feature values ​​and expert feature values ​​of the code function in step (2) comprises the following steps: (2-1) For a given function code segment, first split it according to the camel case naming rule to obtain the input sequence; (2-2) Input the processed sequence into the CodeBERT pre-trained model, extract the hidden states of the first and last layers in the output, and take their average to obtain the semantic feature vector value of the function; (2-3) Use BERT-Whitening to process the semantic vector, extract key features and reduce dimension through linear transformation, and obtain the final semantic feature value vector; (2-4) Multiple expert features are determined from five dimensions: code length, code complexity, code readability, code maintainability, and code executableness. The Understand tool is used to extract the values ​​of these expert features from the function code and form an expert feature value vector.

3. According to claim 1, the cross-project vulnerability detection method based on transfer learning and feature fusion is characterized in that: In step (4), the cross-project vulnerability detection model is constructed using the transfer learning method TrAdaBoost, which includes the following steps: (3-1) To build a cross-project vulnerability detection model M based on semantic features S , we use the transfer learning algorithm TrAdaBoost to reduce the weight of data in the source project that has low relevance to the target project data and increase the weight of data in the source project that has high relevance to the target project data, thereby training an effective cross-project vulnerability detection model M for semantic features. S ; (3-2) Cross-project vulnerability detection model M based on expert features E , we use the transfer learning algorithm TrAdaBoost to reduce the weight of data in the source project that has low relevance to the target project data and increase the weight of data in the source project that has high relevance to the target project data, thereby training an effective cross-project vulnerability detection model M for expert features. E .

4. The cross-project vulnerability detection method based on transfer learning and feature fusion according to claim 1 is characterized in that: In step (5), the features of the two dimensions are fused at the model level, including the following steps: (5-1) For the function code that needs vulnerability detection, a cross-project vulnerability detection model M based on semantic features is used S Get the probability P that it contains a vulnerability Ms ; (5-2) Then the cross-project vulnerability detection model M based on expert features is used E Get the probability P that it contains a vulnerability Me ; (5-3) Through the weighted summation formula P final =P Ms ×W Ms +P Me ×W Me Calculate the fusion model M Final The probability P of containing a vulnerability final , where W Ms and W Me They are cross-project vulnerability detection models M S and M E The weight ratio in the fusion model.

Images