Code vulnerability detection method based on efficient parameter fine tuning of code pre-training model

Abstract

The invention relates to the technical field of software reliability guarantee, in particular to a code vulnerability detection method based on efficient parameter fine tuning of a code pre-training model, which comprises the following steps: preprocessing source codes of a vulnerability data set; s2, converting the preprocessed source code, and extracting a token sequence and a code attribute graph path corresponding to the source code; s3, constructing a code vulnerability detection model, wherein the code vulnerability detection model comprises a semantic feature extractor, a structural feature extractor, a feature fusion layer and a classifier; s4, jointly training a code vulnerability detection model by using the token sequence corresponding to the extracted source code and the code attribute graph path; and S5, through the optimized code vulnerability detection model, outputting a result that whether the source code to be detected has vulnerabilities or not. According to the method, the model is finely adjusted through efficient parameter fine adjustment, training resources are saved, good performance is obtained, and the code vulnerability detection efficiency and accuracy are improved by fusing and learning semantic and structural features of source codes at the same time.

Description

Technical Field

[0001] The present invention relates to the technical field of software reliability assurance, and in particular to a code vulnerability detection method based on efficient parameter fine-tuning of a code pre-training model. Background Art

[0002] Software vulnerabilities are weak links in information systems, security procedures or internal control mechanisms that may be exploited or triggered by attackers, leading to serious security and privacy risks, and even posing a threat to personal safety. Software vulnerability detection aims to analyze possible security defects or weaknesses in software systems and patch them in time before they are exploited to avoid serious consequences.

[0003] To address the challenges in software vulnerability detection, the industry has introduced a variety of detection methods. Early detection methods mainly relied on security experts to manually audit source code based on domain knowledge to identify potential security vulnerabilities. However, this method requires extremely high professional capabilities of personnel, and the labor cost is huge in large-scale systems, which can easily lead to excessive burdens on analysts and a high false alarm rate. In order to reduce manual intervention, static analysis tools have gradually become an important means of vulnerability detection. However, static analysis tools rely on predefined vulnerability pattern libraries, resulting in a limited detection range. As the scale of software expands, methods that rely on manually defined vulnerability patterns appear inefficient and difficult to sustain when dealing with complex and changing security threats.

[0004] The rise of deep learning has greatly promoted the development of vulnerability detection. Through its powerful feature extraction capabilities, deep learning can identify vulnerability features from massive codes, reduce manual intervention, and improve the accuracy and coverage of detection. Currently, the most advanced method is to use the "pre-training and fine-tuning" strategy to detect vulnerabilities through code pre-training models. This model has rich semantic understanding capabilities and good generalization. It not only performs well in vulnerability detection, but is also widely used in tasks such as code completion.

[0005] However, code pre-training models still face some challenges in practical applications. First, such models consume a lot of resources, especially in the training and inference processes, which require high computing resources and storage. Second, although the model is good at capturing the semantic features of the code, it often ignores the structured information of the code, such as function call relationships, control flow, and data flow, which are crucial for vulnerability detection. In addition, there is a large task difference between the pre-training stage and the actual vulnerability detection task, which limits the migration effect of the model in downstream tasks, and to some extent hinders the widespread application and further development of code pre-training models in source code vulnerability detection. Summary of the invention

[0006] In order to solve the technical problems existing in the prior art, the present invention provides a code vulnerability detection method based on efficient parameter fine-tuning of a code pre-training model, which can make up for the missing code structure features in the pre-training model. At the same time, through efficient parameter fine-tuning, the consumption of training resources is reduced, and the gap between pre-training tasks and downstream tasks is reduced. The different features in the source code can be fully learned, the detection accuracy is high, and it has good generalization.

[0007] The purpose of the present invention can be achieved by adopting the following technical solutions:

[0008] The code vulnerability detection method based on efficient parameter fine-tuning of the code pre-training model includes the following steps

[0009] S1, build a vulnerability dataset and preprocess the source code of the vulnerability dataset;

[0010] S2. Transform the preprocessed source code to extract the token sequence and code attribute graph path corresponding to the source code;

[0011] S3. Build a code vulnerability detection model based on the UniXcoder model and the doc2vec model. The code vulnerability detection model includes a semantic feature extractor, a structural feature extractor, a feature fusion layer, and a classifier.

[0012] S4, jointly train the code vulnerability detection model with the token sequence and the code attribute graph path corresponding to the extracted source code, optimize the parameters of the model through the back propagation algorithm, and obtain the optimized code vulnerability detection model;

[0013] S5. Extract the token sequence and code attribute graph path corresponding to the source code to be tested, and output the result of whether the source code to be tested has vulnerabilities through the optimized code vulnerability detection model.

[0014] Specifically, the step of constructing a vulnerability dataset and preprocessing the source code of the vulnerability dataset includes:

[0015] Obtain vulnerability-related data from open source projects Qemu and FFmpeg and build a vulnerability dataset;

[0016] The source code of the vulnerability dataset is preprocessed, the variable names and function names in the source code are replaced with corresponding symbolic representations, the redundant space symbols and blank line symbols in the source code of the dataset are removed, and the information in the dataset that is not related to the vulnerability characteristics is cleaned.

[0017] Specifically, the preprocessed source code is transformed to extract the token sequence and code attribute graph path corresponding to the source code, including:

[0018] The tokenizer pre-trained by the UniXcoder model converts each word in the source code into smaller subwords to obtain the token sequence corresponding to the source code.

[0019] The Joern tool is used to convert the source code into the corresponding program dependency graph, extract all data paths and control paths in the program dependency graph, and obtain the code attribute graph path of the source code.

[0020] Specifically, the semantic feature extractor adopts the UniXcoder model as the backbone network, adds an adapter layer to each transmitter encoder layer of the UniXcoder model, and extracts the semantic feature vector of the source code through the semantic feature extractor;

[0021] The structural feature extractor adopts a doc2vec model, and uses the doc2vec model to learn the structural features in the code attribute graph path to obtain the structural feature vector of the source code;

[0022] The structural feature vector and the semantic feature vector are concatenated through the fusion layer to obtain a composite feature vector;

[0023] The classifier adopts a multi-layer perceptron, which learns the information of the composite feature vector and outputs the result of whether the source code has a vulnerability.

[0024] Specifically, the adapter layer consists of two linear layers and one nonlinear layer, and residual connections are used between each adapter layer.

[0025] Specifically, extracting the semantic feature vector of the source code by a semantic feature extractor includes:

[0026] The input source code is segmented by the tokenizer pre-trained by the UniXcoder model and converted into a sequence of tokens;

[0027] Convert the token sequence into an initial vector representation, which includes word vector embedding and position vector embedding;

[0028] A 12-layer bidirectional transformer encoder is used to learn the relationship between tokens and obtain the semantic feature vector.

[0029] Specifically, the transformer encoder includes a multi-head attention module and a residual normalization layer, there is an adapter layer between the multi-head attention module and the residual normalization layer, and there is an adapter layer between the residual normalization layer and the residual normalization layer.

[0030] Specifically, the token sequence and the code attribute graph path corresponding to the extracted source code are jointly trained to train the code vulnerability detection model, including: using multiple dropout layers to prevent overfitting, using KL divergence to penalize different outputs of the same input of the model, and determining the loss function L based on the binary cross entropy loss function and the divergence loss function. The loss function L is calculated as follows:

[0031] L=L ce +βL kl ;

[0032] Among them, L ce is the binary cross entropy function, L kl is the divergence loss function and β is the hyperparameter used for tuning.

[0033] Compared with the prior art, the present invention has the following advantages and beneficial effects:

[0034] The present invention provides a code vulnerability detection method based on efficient parameter fine-tuning of a code pre-training model. By adding an adapter layer to each encoder layer of the UniXcoder model, the resources consumed by fine-tuning the pre-training model are alleviated, the efficiency of model training is improved, the gap between pre-training tasks and downstream tasks is narrowed, and better performance is obtained while training a very small number of parameters. The structural feature vector and the semantic feature vector are spliced ​​together through a fusion layer to obtain a composite feature vector, which makes up for the defect that the pre-training model only uses the semantic features of the code and ignores the rich structural features in the code. The multi-head attention mechanism of UnixCoder is used to learn the features of different spaces in the code sequence, and the structural features in the code attribute graph are further learned through model fusion. The semantic features and structural features are integrated to greatly improve the performance of model detection. The present invention can fully learn different features in the source code, has a high detection accuracy, can accurately predict whether there are vulnerabilities in the source code, and has good generalization. BRIEF DESCRIPTION OF THE DRAWINGS

[0035] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on the structures shown in these drawings without paying creative work.

[0036] Figure 1 is a flow chart of a code vulnerability detection method based on efficient parameter fine-tuning of a code pre-training model in an embodiment of the present invention;

[0037] Figure 2is a schematic diagram of the flow of source code preprocessing in an embodiment of the present invention;

[0038] Figure 3 It is an overall network framework diagram of the code vulnerability detection model in an embodiment of the present invention;

[0039] Figure 4 is a flow chart of the semantic extractor architecture in an embodiment of the present invention;

[0040] Figure 5 It is a schematic diagram of the structure of an adapter module inserted in the semantic extractor in an embodiment of the present invention. DETAILED DESCRIPTION

[0041] The technical solution of the present invention will be further described in detail below in conjunction with the accompanying drawings and embodiments. It is obvious that the described embodiments are part of the embodiments of the present invention, not all of the embodiments, and the implementation of the present invention is not limited to this. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative work are within the scope of protection of the present invention.

[0042] Embodiment 1:

[0043] The present invention provides a code vulnerability detection method based on efficient parameter fine-tuning of a code pre-training model, which utilizes the rich code semantic knowledge of the code pre-training model and has good generalization. By fusing code semantics and structural features, the lack of structural features in the pre-training model is compensated, and the detection performance of the model is improved. The model is fine-tuned by efficient parameter fine-tuning adapter-tuning, while retaining the rich pre-training knowledge of the model, it fits the target task well, saves training resources and obtains good performance. By simultaneously learning the semantics and structural features of the source code through model fusion, the efficiency and accuracy of code vulnerability detection are greatly improved, which has certain guiding significance for the field of vulnerability detection.

[0044] like Figure 1 As shown, a flow chart of a code vulnerability detection method based on efficient parameter fine-tuning of a code pre-training model is provided. The code vulnerability detection method based on efficient parameter fine-tuning of a code pre-training model according to the present invention comprises the following steps:

[0045] S1. Build a vulnerability dataset and preprocess the source code of the vulnerability dataset;

[0046] First, a vulnerability data set is constructed, and vulnerability-related data is obtained from the open source projects Qemu and FFmpeg and a vulnerability data set is constructed. The selection of vulnerability data sets is very important for verifying the model proposed in the present invention. Many studies have obtained good performance on synthetic data sets, such as the Juliet data set and the SARD data set. The codes of these data sets are all artificially synthesized codes for certain vulnerabilities, but because these synthetic codes are quite different from the real codes, the performance of these models on real data sets drops sharply. The data set selected in this embodiment comes from two open source projects, Qemu and FFmpeg. In these two open source projects, all submissions are marked as vulnerability repair submissions and non-vulnerability repair submissions. All vulnerability repair submissions will contain security-related keywords such as Dos and injection, and in order to ensure the quality of the data set, two rounds of data annotation and cross-validation were performed by 4 professional security researchers. For each submission, the source code before submission was extracted and a label was assigned. The data set contains a total of 27k C language functions with a defect rate of 45%. Each sample contains the corresponding function source code and label. For the function source code, since it contains many redundant line breaks and spaces, it will be cleaned up in subsequent processing. By training and validating on real data sets, the actual performance of the model in practical applications can be verified, and it is beneficial to apply the model to the actual production environment.

[0047] Secondly, the source code of the vulnerability dataset is preprocessed, the variable names and function names in the source code are replaced with corresponding symbols, the redundant spaces and blank lines in the source code of the dataset are removed, and the information in the dataset that is not related to the vulnerability features is cleaned. In order to avoid vocabulary explosion caused by too many variable names in the source code, which brings difficulties to model learning, the source code is standardized.

[0048] like Figure 2 The figure shows the flow chart of source code preprocessing. Specifically, it includes: replacing the variable names and function names in the source code with the corresponding symbols VARi and FUNCi. Removing redundant spaces, blank lines and other symbols in the source code of the dataset, and cleaning the parts of the dataset that are not related to vulnerability features.

[0049] S2. Transform the preprocessed source code to extract the token sequence and code attribute graph path corresponding to the source code. In the subsequent steps, the tokens and code attribute graph paths of the source code are extracted, which can be used as the semantic input and structural input of the code vulnerability detection model respectively.

[0050] S21. Use the word segmenter pre-trained by the UniXcoder model to convert each word in the source code into smaller sub-words to obtain the token sequence corresponding to the source code. After cleaning, the source code needs to be converted into the input form of the UniXcoder model. The input form of the UniXcoder model is a sequence of source code tokens. Converting the source code into the corresponding token sequence specifically includes: using the word segmenter pre-trained by the UniXcoder model to convert each word in the source code into smaller sub-words to a source code token sequence using the Byte Pair Encoding (BPE) algorithm. The input of the conversion process is the source code function f, and the result of the conversion is a token sequence represented by T=(t1, t2,…, t n ). After obtaining the source code token sequence, the input of the UniXcoder model can be constructed, expressed as I = {[CLS], T, [SEP]}, where CLS and SEP are two special tokens. The model output corresponding to the CLS token will be used as the semantic representation of the entire sequence. The SEP token indicates the end position of the input. Because the input of the transformer model requires a fixed length, assuming that the input length is fixed to L, if the input length is less than L, the sequence will be supplemented to L after SEP. The supplemented token is PAD, which is equivalent to filling 0 in the sequence and will not affect the final output. If the input length is greater than the token, the input sequence will be truncated to L, and the last token is SEP.

[0051] by Figure 2 As shown in the source code, the token sequence T of the source code after word segmentation is:

[0052]

[0053] in, Symbols represent tokens that appear very rarely. This will reduce the size of the vocabulary and make model training easier.

[0054] S22. Use the Joern tool to convert the source code into the corresponding program dependency graph, extract all data paths and control paths in the program dependency graph, and obtain the code attribute graph path of the source code. In order to make up for the lack of code structure features in the pre-trained model UniXcoder model, the path representation is extracted from the source code attribute graph as a structural feature. Specifically, the Joern tool is used to convert the source code function into the corresponding program dependency graph (Program Dependency Graph, PDG), expressed as G = (V, E), where V represents a series of nodes corresponding to statements, each node has two attributes: node type d and code token t, and E represents the data or control dependency between statements. Therefore, the graph path includes two categories: control flow path and data flow path. The path is represented as p i =(t1,l1,t2,l2,…,t k ,l k ), where t i Indicates the type of node i, l i Indicates the type of edge between node i and node i+1, l i ∈(data,control). For each source code function, extract all data dependency paths and control dependency paths in the function program dependency graph, and represent the structural features of the source code function as S=(s1,s2,…,s m ), where s i =(s,p,e) represents a data flow path or control flow path between nodes s and e.

[0055] like Figure 2 As shown, the source code is first converted into a program dependency graph. Each node corresponds to the corresponding code in the source code, and the number corresponds to the line number where the code exists. There are different dependencies between the nodes. The extracted control dependency paths are [1,2,4,5] and [1,2,3,5], and the data dependency paths are [2,5] and [1,3,5]. Therefore, the final structural feature is recorded as S = [[1,2,4,5], [1,2,3,5], [2,5], [1,3,5]].

[0056] S3. Build a code vulnerability detection model based on the UniXcoder model and the doc2vec model. The code vulnerability detection model includes a semantic feature extractor, a structural feature extractor, a feature fusion layer and a classifier.

[0057] like Figure 3As shown in the figure, the overall network framework of the code vulnerability detection model includes a semantic feature extractor and a structural feature extractor. Specifically, the semantic feature extractor uses the UniXcoder model as the backbone network, adds an adapter layer to each Transformer Encoder layer of the UniXcoder model, and extracts the semantic feature vector of the source code through the semantic feature extractor. The adapter is inserted into the middle layer of the pre-trained model to allow only a small number of parameters to be modified when fine-tuning a specific task, thereby improving the efficiency and speed of fine-tuning, that is, fine-tuning the code vulnerability detection model in an efficient parameter fine-tuning adapter-tuning manner.

[0058] The UniXcoder model is mainly based on the Transformer structure, which mainly consists of a tokenizer, an input encoder, and a decoder. The tokenizer is pre-trained on a large number of code predictions and can map the tokenized code to the corresponding vector representation. During the pre-training process, its input includes the abstract syntax tree of natural language and code, so it also has a good understanding of the structural characteristics of the abstract syntax tree in the code. After the tokenizer, there are the encoder and decoder, the core components of the UniXcoder model. Both the encoder and decoder are stacked by multiple transformer blocks, and can work in three modes, namely encoder-only architecture, encoder-decoder (encoder-decoder) architecture and decoder-only architecture. In this embodiment, the encoder-only mode (only the decoder part) is adopted to obtain a rich vector representation of the source code. Under the encoder-only architecture, each transformer block is mainly composed of a multi-head attention layer and a residual normalization layer. The multi-head attention layer is used to learn different features in different subspaces of the input vector and can learn features of different dimensions. The residual normalization layer is used to normalize the vector, accelerate model convergence, and improve the model effect. After learning through multiple layers of transformer blocks, the final hidden layer output is a semantic feature vector containing rich feature representations.

[0059] The UniXcoder model is a code pre-training model. The UniXcoder model is pre-trained on a large number of source code predictions, including source code text and abstract syntax trees (AST), and has rich code semantic knowledge. The code token sequence extracted in step S1 is used as the model input. It will first undergo word embedding and position embedding. Each token is mapped to a high-dimensional vector. After that, the initial vector passes through a 12-layer bidirectional Transformer Encoder layer to fully learn the interdependencies between tokens and the source code semantic features. The bidirectional transformer encoder uses a bidirectional multi-head attention mechanism to learn the relationship between the source code and the previous and next sequences. The hidden layer vector of the Nth layer can be expressed as The vector representation of the last layer is usually used as the semantic representation of the source code function. UniXcoder has 125 million training parameters. The huge number of parameters consumes huge resources and hinders the model from fitting in downstream tasks.

[0060] In order to obtain a better representation of the source code vector, an adapter layer is added to each Transformer Encoder. The adapter layer consists of two linear layers and one nonlinear layer, and residual connections are used between each adapter layer to prevent adverse effects on model performance. During the fine-tuning of model parameters, the pre-trained parameters of the UniXcoder model remain unchanged, and only the parameters of the adapter layer are updated, which greatly reduces the parameters that need to be updated, speeds up the training iteration speed, and can fit the downstream tasks well. At present, most fine-tuning methods are based on full parameter fine-tuning, that is, fine-tuning all parameters of the pre-trained model on the downstream task data set. Taking UniXcoder as an example, the model has 125 million parameters. If all parameters are trained, first of all, it has high requirements on the computing power and video memory of the device, and it takes a long time. At the same time, in downstream tasks, there is a problem of small amount of data. If the amount of data is insufficient, it is difficult to obtain good results when training so many parameters. Efficient parameter fine-tuning aims to achieve the same effect by training a very small number of parameters while retaining a large number of original pre-trained parameters. In UniXcoder, using adapter-tuning only requires training 287W parameters, which is 2% of the parameters of full parameter fine-tuning. However, when tested on downstream datasets, the performance obtained is slightly better than that of full parameter fine-tuning. The effect of training a very small number of parameters is equivalent to that of full parameter fine-tuning.

[0061] like Figure 4 As shown, the flowchart of the semantic extractor architecture is shown in FIG. 1 . Specifically, the process of extracting the semantic feature vector of the source code by the semantic feature extractor includes:

[0062] First, input the source code in the vulnerability dataset, segment it using the tokenizer pre-trained by the UniXcoder model, and convert it into a sequence of tokens.

[0063] The token sequence is then converted into an initial vector representation, which includes word vector embedding and position vector embedding. Word vector embedding is a vector that represents the semantics of a token. At the same time, since the transformer model does not have the position information of a token, position embedding represents the position information of each token in the entire sequence. Through these two embedding methods, the initial vector representation contains the semantics and position information of each token.

[0064] After obtaining the initial vector representation, it is necessary to use a 12-layer bidirectional transformer encoder to learn the relationship between tokens, obtain a semantic feature vector rich in semantic features, and learn the vulnerability feature information. Figure 4 As shown in the figure, the Transformer encoder includes a multi-head attention module and a residual normalization layer. There is an adapter layer between the multi-head attention module and the residual normalization layer, and there is an adapter layer between the residual normalization layer and the residual normalization layer. The multi-head attention mechanism is the key to its powerful learning ability. First, the input is linearly mapped to Q, K and V. The conventional attention mechanism is:

[0065]

[0066] Where k represents the dimension of the hidden layer. The multi-head attention mechanism uses multiple attentions to obtain attention from the hidden layer vector, and then can learn various vulnerability feature information from different subspaces represented by the code. Specifically, for Q, K and V vectors of dimension d, these vectors are divided into h heads, and the dimension of each head is d / h. After each head performs self-attention operation, the output of all heads is spliced ​​back to the original dimension, which is expressed as MultiHead=Concat(head1,head2,…,head h )W O , where head i is the self-attention output of each head, calculated by the above attention mechanism Attention formula, and the combined output is input to the inserted adapter layer. Figure 5As shown in the figure, the schematic diagram of the adapter layer consists of two linear layers and one nonlinear layer. The output of the multi-head attention layer is input to the adapter layer (adapter layer), and first passes through the lower mapping layer. The neuron nodes of this layer are usually determined by the reduction parameter. Taking reduction = 4 as an example, the number of output nodes of the lower mapping layer = input dimension / recution, and then passes through a nonlinear layer for feature learning. Finally, it is mapped to the same dimension as the input layer through the upper mapping layer. The adapter layer sets a residual connection to prevent parameter degradation. Finally, it passes through the residual normalization layer to prevent parameter degradation and accelerate model convergence. Similarly, the output of the residual normalization layer will also be input to the adapter layer for mapping. The initial vector will go through the above-mentioned 12-layer forward propagation process, and finally obtain a semantic feature vector rich in semantic features.

[0067] The structural feature extractor adopts the doc2vec model, and uses the doc2vec model to learn the structural features in the code attribute graph path to obtain the structural feature vector of the source code. Doc2Vec is a text representation model based on word vectors, which can provide semantic vector representation of the entire text and can process texts of different lengths. It is suitable for processing the situation where the number and length of each function code attribute graph path are different. After extracting all the structural paths of the source code, all the paths S = (s1, s2, ..., s m ) as a whole to build the training prediction of doc2vec. The Doc2vec model takes all the paths S of a function as a whole and obtains the structural feature vector of the current function through other path vectors in the context. Doc2vec is based on unsupervised learning. The training corpus is the structural features extracted from each function in the data set. After training, the model has a good ability to capture the structural features of the code.

[0068] The structural feature vector and the semantic feature vector are concatenated through the fusion layer as the final composite feature vector. After obtaining the semantic features and structural features of the source code, the two features are combined for learning through the fusion layer. First, the structural features and semantic features are mapped to vectors of the same dimension through a linear layer to prevent learning bias caused by different dimensions. Then the structural feature vector and the semantic feature vector are concatenated as the final composite feature vector.

[0069] The classifier adopts a multi-layer perceptron, which fully learns the information of the composite feature vector through the multi-layer perceptron and outputs the result of whether the source code has a vulnerability. The multi-layer perceptron is a classic feedforward neural network, which uses its multi-layer structure and nonlinear feature capabilities to enable it to complete a variety of classification tasks. Since the semantic extractor and the structural extractor have converted the semantic features and structural features in the code into vector representations, the multi-layer perceptron can further learn the relevant knowledge of the vulnerabilities in the features for classification. The multi-layer perceptron is divided into an input layer, a hidden layer, and an output layer. Each layer is fully connected, that is, each node is connected to each node of the previous layer, and each connection has a corresponding weight. Each node receives the output of all nodes from the previous layer, and after weighted summation according to the weight, the corresponding bias of each node is added. The weighted sum is nonlinearly transformed by the activation function and passed to the next layer as the output of the node, wherein the nonlinear transformation is completed by the activation function of the hidden layer, and the present invention adopts the Relu activation function. Therefore, the multi-layer perceptron can perform feature fusion layer by layer through the combination of multi-layer neurons and nonlinear activation functions, and learn the complex patterns of high-dimensional data.

[0070] S4. The token sequence and code attribute graph path corresponding to the extracted source code are jointly used to train the code vulnerability detection model, and the parameters of the model are optimized through the back propagation algorithm to obtain the optimized code vulnerability detection model.

[0071] Specifically, the token sequence obtained in step S2 is converted into a high-dimensional vector through an embedding layer, with a fixed batch size and the corresponding label as input, represented as (batch, seq, dimension, label), and the output is (batch,, dimension_1, label). The code attribute graph path obtained in S1 takes the same batch size as the input of the doc2vec model, converted into the corresponding structural feature vector, represented as (batch,, dimension_2, label), and then undergoes a linear mapping layer to be converted to (batch,, dimension_1, label), and the final concatenation result is (batch,, 2*dimension_1, label). First, as a classification task, the training loss function uses the binary cross entropy function to measure the difference between the predicted value and the true label:

[0072]

[0073] Among them, y i represents the true label, which is 0 or 1 in this task, 0 means that the code is predicted to have no vulnerability, and 1 means that the code contains a vulnerability; pi Represents the model prediction value, usually the probability that the sample belongs to category 1, and is a floating point number between 0 and 1. N is usually the number of samples in the training batch. Through the binary cross entropy loss function, the loss between different samples can be balanced, so that the model gradually converges and the performance of the model classification is optimized.

[0074] Specifically, the process of training the code vulnerability detection model includes using multiple dropout layers to prevent overfitting. KL divergence can be used to penalize different outputs of the same input of the model. The loss function is determined based on the binary cross entropy loss function and the divergence loss function que. Using multiple dropout layers to prevent overfitting will result in slight differences in the output of the same input. Therefore, KL divergence is applied to penalize the inconsistent output of the model under the same input. KL divergence is used to measure the difference between two distributions:

[0075]

[0076] Where P and Q are for input x i Specifically, P(x i ) and Q(x i ) are different outputs of the model for the same input.

[0077] The KL loss calculation formula is as follows:

[0078]

[0079] Among them, P1 and P2 are two different model output distributions for the same input. For example, if the input sample is 5, then P(x i ) and Q(x i ) is the model output obtained by inputting these 5 samples into the model twice. In the binary classification task, it is a 5*2-dimensional vector. The first column of the vector represents the probability that the sample is predicted to be 0, and the second column represents the probability that the sample is predicted to be 1.

[0080] Finally, the loss function calculation formula is:

[0081] L=L ce +βL kl

[0082] Among them, β is a hyperparameter used for adjustment, which is a parameter between 0 and 1. This loss function combines the binary cross entropy loss function and the divergence loss function, which can improve the model's fitting effect on the real data distribution, ensure the accuracy of the model output, and help improve the robustness of the model. The binary cross entropy loss function can penalize model classification errors, and since UniXcoder introduces a large number of dropout layers to prevent overfitting, the KL divergence can penalize different outputs of the same input of the model.

[0083] After determining the model loss function, the back propagation algorithm is used to optimize the model parameters to obtain an optimized code vulnerability detection model. Model training is mainly divided into two steps: forward propagation and back propagation. Forward propagation is to input our source code token into Adapter-UniXcoder, convert it into a vector through word embedding, and then obtain its semantic vector representation through a 12-layer bidirectional Transformer Encoder; in addition, the structural feature representation of the code is extracted through the doc2vec model, and the two models are combined as the feature representation of the code. The feature representation is converted into the final model prediction output after passing through multiple linear layers, and then the loss between the model output and the actual label value is calculated through the above loss function, and the forward propagation ends. After obtaining the loss, back propagation is performed, specifically, the loss is derived for the parameters of each layer, and the parameters of the layer are updated after the derivative is obtained. During the training process, we set the epoch to 10, that is, 10 rounds of iterations. Batch-size is set to 32, that is, parameters are updated sequentially for every 32 samples. After training every 100 batches, the validation loss on the validation set is calculated, and the model parameter value with the smallest loss on the validation set is finally saved as the optimal model. Subsequent predictions will load the saved optimal model parameters.

[0084] S5. Extract the token sequence and code attribute graph path corresponding to the source code to be tested, and output the result of whether the source code to be tested has vulnerabilities through the optimized code vulnerability detection model.

[0085] Specifically, the code vulnerability detection model after training can be used for vulnerability detection tasks. The vulnerability detection process includes: inputting the source code sequence to be tested, the detection framework will extract the code token and code attribute graph path corresponding to the source code to be tested, and then the code semantic feature extractor and structural feature extractor will obtain the corresponding semantic feature vector and structural feature vector respectively. After splicing to obtain the composite feature vector, it is input into the trained classifier, and finally a prediction value of whether there is a vulnerability will be obtained. The preset threshold is set to 0.5. If the prediction value is greater than the preset threshold 0.5, it is judged that the source code has a vulnerability. If the prediction value is less than or equal to the preset threshold 0.5, it is judged that the source code does not have a vulnerability.

[0086] In summary, the present invention provides a code vulnerability detection method based on efficient parameter fine-tuning of a code pre-training model. By adding an adapter layer to each encoder layer of the UniXcoder model, the resources consumed by fine-tuning the pre-training model are alleviated, the efficiency of model training is improved, the gap between pre-training tasks and downstream tasks is narrowed, and better performance is obtained while training a very small number of parameters. The structural feature vector and the semantic feature vector are spliced ​​together through the fusion layer to obtain a composite feature vector, which makes up for the defect that the pre-training model only uses the semantic features of the code and ignores the rich structural features in the code. The multi-head attention mechanism of UnixCoder is used to learn the features of different spaces in the code sequence, and the structural features in the code attribute graph are further learned through model fusion. The semantic features and structural features are integrated to greatly improve the performance of model detection. The present invention can fully learn different features in the source code, has a high detection accuracy, can accurately predict whether there are vulnerabilities in the source code, and has good generalization.

[0087] The above embodiments are preferred implementation modes of the present invention, but the implementation modes of the present invention are not limited to the above embodiments. Any other changes, modifications, substitutions, combinations, and simplifications that do not deviate from the spirit and principles of the present invention should be equivalent replacement methods and are included in the protection scope of the present invention.

Claims

1. A code vulnerability detection method based on efficient parameter fine-tuning of a code pre-training model, characterized in that: The following steps are included S1, build a vulnerability dataset and preprocess the source code of the vulnerability dataset; S2. Transform the preprocessed source code to extract the token sequence and code attribute graph path corresponding to the source code; S3. Build a code vulnerability detection model based on the UniXcoder model and the doc2vec model. The code vulnerability detection model includes a semantic feature extractor, a structural feature extractor, a feature fusion layer, and a classifier. S4, jointly train the code vulnerability detection model with the token sequence and the code attribute graph path corresponding to the extracted source code, optimize the parameters of the model through the back propagation algorithm, and obtain the optimized code vulnerability detection model; S5. Extract the token sequence and code attribute graph path corresponding to the source code to be tested, and output the result of whether the source code to be tested has vulnerabilities through the optimized code vulnerability detection model.

2. The code vulnerability detection method based on efficient parameter fine-tuning of code pre-training model according to claim 1 is characterized in that: The step of constructing a vulnerability dataset and preprocessing the source code of the vulnerability dataset includes: Obtain vulnerability-related data from open source projects Qemu and FFmpeg and build a vulnerability dataset; The source code of the vulnerability dataset is preprocessed, the variable names and function names in the source code are replaced with corresponding symbolic representations, the redundant space symbols and blank line symbols in the source code are removed, and the information irrelevant to the vulnerability features in the vulnerability dataset is cleaned.

3. The code vulnerability detection method based on efficient parameter fine-tuning of code pre-training model according to claim 2 is characterized in that: The preprocessed source code is converted to extract the token sequence and code attribute graph path corresponding to the source code, including: The tokenizer pre-trained by the UniXcoder model converts each word in the source code into smaller subwords to obtain the token sequence corresponding to the source code. The Joern tool is used to convert the source code into the corresponding program dependency graph, extract all data paths and control paths in the program dependency graph, and obtain the code attribute graph path of the source code.

4. The code vulnerability detection method based on efficient parameter fine-tuning of code pre-training model according to claim 3 is characterized in that: The semantic feature extractor uses the UniXcoder model as the backbone network, adds an adapter layer to each transmitter encoder layer of the UniXcoder model, and extracts the semantic feature vector of the source code through the semantic feature extractor; The structural feature extractor adopts a doc2vec model, and uses the doc2vec model to learn the structural features in the code attribute graph path to obtain the structural feature vector of the source code; The structural feature vector and the semantic feature vector are concatenated through the fusion layer to obtain a composite feature vector; The classifier adopts a multi-layer perceptron, which learns the information of the composite feature vector and outputs the result of whether the source code has a vulnerability.

5. The code vulnerability detection method based on efficient parameter fine-tuning of code pre-training model according to claim 4 is characterized in that: The adapter layer consists of two linear layers and one nonlinear layer, and residual connections are used between each adapter layer.

6. The code vulnerability detection method based on efficient parameter fine-tuning of code pre-training model according to claim 5 is characterized in that: The step of extracting the semantic feature vector of the source code by using a semantic feature extractor includes: The input source code is segmented by the tokenizer pre-trained by the UniXcoder model and converted into a token sequence; Convert the token sequence into an initial vector representation, which includes word vector embedding and position vector embedding; A 12-layer bidirectional transformer encoder is used to learn the relationship between tokens and obtain the semantic feature vector of the source code.

7. The code vulnerability detection method based on efficient parameter fine-tuning of code pre-training model according to claim 5 is characterized in that: The transformer encoder includes a multi-head attention module and a residual normalization layer, an adapter layer is provided between the multi-head attention module and the residual normalization layer, and an adapter layer is provided between the residual normalization layer and the residual normalization layer.

8. The code vulnerability detection method based on efficient parameter fine-tuning of code pre-training model according to claim 5 is characterized in that: The token sequence and the code attribute graph path corresponding to the extracted source code are jointly trained into a code vulnerability detection model, including: using multiple dropout layers to prevent overfitting, using KL divergence to penalize different outputs of the same input of the model, and determining the loss function L based on the binary cross entropy loss function and the divergence loss function. The loss function L is calculated as follows: L=L ce +βL kl ; Among them, L ce is the binary cross entropy function, L kl is the divergence loss function and β is the hyperparameter used for tuning.

Images