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

Abstract

The application relates to the technical field of software reliability guarantee, and relates to a code vulnerability detection method based on efficient parameter fine-tuning of a code pre-training model. The method comprises the following steps: preprocessing source code of a vulnerability data set; S2, performing transformation processing on the preprocessed source code, and extracting token sequences and code attribute graph paths 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 the code vulnerability detection model by using the extracted token sequences and code attribute graph paths corresponding to the source code; and S5, outputting a result of whether the to-be-detected source code has a vulnerability by using the optimized code vulnerability detection model. The application fine-tunes the model by using efficient parameter fine-tuning, saves training resources and obtains good performance. By fusing and simultaneously learning semantic and structural features of the source code, the efficiency and accuracy of code vulnerability detection are improved.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of software reliability assurance, and particularly relates to a code vulnerability detection method based on code pre-training model parameter fine-tuning. BACKGROUND

[0002] Software vulnerabilities are weak links in information systems, security programs or internal control mechanisms that can be exploited or triggered by attackers, leading to serious security and privacy risks, and even threats to personal safety. Software vulnerability detection aims to analyze potential security defects or weaknesses in software systems and fix them 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 rely on security experts to manually audit source code based on domain knowledge to identify potential security vulnerabilities. However, this method requires a high level of professional expertise and is extremely costly in large-scale systems, which can lead to an excessive burden on analysts and a high rate of false positives. To reduce human intervention, static analysis tools have gradually become an important means of vulnerability detection. However, static analysis tools rely on predefined vulnerability pattern libraries, limiting their detection range. As software scales, methods that rely on manually defined vulnerability patterns are inefficient and difficult to sustain in the face of complex and changing security threats.

[0004] The rise of deep learning has greatly driven the development of vulnerability detection. Through its powerful feature extraction capabilities, deep learning can identify vulnerability features from vast amounts of code, reducing human intervention and improving detection accuracy and coverage. Currently, the most advanced method is to use a "pre-training and fine-tuning" strategy to detect vulnerabilities through code pre-training models. This model has rich semantic understanding capabilities and good generalizability, not only performing well in vulnerability detection, but also being widely used in code completion and other tasks.

[0005] However, code pre-training models still face some challenges in practical applications. First, these models consume a lot of resources, particularly in terms of computational resources and storage during training and inference. Second, while the model is good at capturing semantic features of code, it often overlooks structured information such as function call relationships, control flow and data flow, which are crucial for vulnerability detection. In addition, there is a significant task difference between the pre-training phase and the actual vulnerability detection task, which limits the model's transfer effect 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

[0006] To solve the technical problems existing in the prior art, the code vulnerability detection method based on code pre-training model parameter fine-tuning can make up for the missing code structure features in the pre-training model, reduce the consumption of training resources through parameter fine-tuning, reduce the gap between the pre-training task and the downstream task, fully learn different features in the source code, have high detection accuracy, and have good generalization.

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

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

[0009] S1, constructing a vulnerability data set, and preprocessing source code of the vulnerability data set;

[0010] S2, transforming the preprocessed source code, and extracting token sequences and code attribute graph paths corresponding to the source code;

[0011] S3, constructing a code vulnerability detection model based on a UniXcoder model and a doc2vec model, wherein the code vulnerability detection model comprises a semantic feature extractor, a structural feature extractor, a feature fusion layer and a classifier;

[0012] S4, jointly training the code vulnerability detection model by using the extracted token sequences and code attribute graph paths corresponding to the source code, optimizing the parameters of the model through a back propagation algorithm, and obtaining an optimized code vulnerability detection model;

[0013] S5, extracting token sequences and code attribute graph paths corresponding to the source code to be tested, and outputting a result of whether the source code to be tested has a vulnerability through the optimized code vulnerability detection model.

[0014] Specifically, the constructing of the vulnerability data set and the preprocessing of the source code of the vulnerability data set comprise:

[0015] Obtaining vulnerability-related data from open source projects Qemu and FFmpeg and constructing a vulnerability data set;

[0016] Preprocessing the source code of the vulnerability data set, replacing variable names and function names in the source code with corresponding symbolic representations, removing redundant space symbols and blank line symbols in the source code of the data set, and cleaning information irrelevant to vulnerability features in the data set.

[0017] Specifically, the transforming of the preprocessed source code and the extraction of token sequences and code attribute graph paths corresponding to the source code comprise:

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

[0019] The source code is converted into a corresponding program dependence graph by using the Joern tool, all data paths and control paths in the program dependence graph are extracted, and a code property graph path of the source code is obtained.

[0020] Specifically, the semantic feature extractor adopts the UniXcoder model as a backbone network, and an adapter adapter layer is added in each ransformer encoder layer of the UniXcoder model, and the semantic feature vector of the source code is extracted through the semantic feature extractor;

[0021] The structure feature extractor adopts the doc2vec model, and the structure feature of the code property graph path is learned by using the doc2vec model, and a structure feature vector of the source code is obtained.

[0022] The structure feature vector and the semantic feature vector are spliced by the fusion layer to obtain a composite feature vector.

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

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

[0025] Specifically, the semantic feature vector of the source code is extracted by the semantic feature extractor, including:

[0026] The input source code is segmented by the tokenizer pre-trained by the UniXcoder model, and is converted into a sequence in token units;

[0027] The token sequence is converted into an initial vector representation, and the initial vector representation includes a word vector embedding and a position vector embedding.

[0028] The relationship between the tokens is learned by a 12-layer bidirectional transformer encoder to obtain a semantic feature vector.

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

[0030] Specifically, the joint training of the extracted source code corresponding token sequence and code attribute graph path and the code vulnerability detection model comprises: using multiple dropout layers to prevent overfitting, punishing different outputs of the same input of the model through KL divergence, determining a loss function L based on a binary cross-entropy loss function and a divergence loss function, and the loss function L is calculated according to the following formula:

[0031] L = L ce + beta * L kl ;

[0032] Wherein, L ce is a binary cross-entropy function, L kl is a divergence loss function, and beta is a hyperparameter used for adjustment.

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

[0034] The present application provides a code vulnerability detection method based on code pre-training model parameter fine-tuning, which adds an adapter adapter layer in each ransformer encoder layer of the UniXcoder model, alleviates the resources consumed by fine-tuning the pre-training model, improves the efficiency of model training, narrows the gap between the pre-training task and the downstream task, and obtains better performance while training a small part of parameters. The structural feature vector and the semantic feature vector are spliced into a composite feature vector through a fusion layer, which makes up for the defect that the pre-training model only uses code semantic features 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, and the semantic features and structural features are integrated, which greatly improves the performance of the model detection. The present application can fully learn different features in the source code, has high detection accuracy, can accurately predict whether there is a vulnerability in the source code, and has good generalization. BRIEF DESCRIPTION OF DRAWINGS

[0035] In order to more clearly illustrate the technical solutions in the embodiments of the present application or the prior art, the following will briefly introduce the drawings needed to be used in the embodiments or prior art description. Obviously, the drawings in the following description are only some embodiments of the present application, and for those skilled in the art, other drawings can also be obtained without creative labor based on the drawings shown.

[0036] Figure 1 is a flowchart of the code vulnerability detection method based on code pre-training model parameter fine-tuning in the embodiments of the present application;

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

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

[0039] Figure 4 This is a schematic diagram of the semantic extractor architecture in an embodiment of the present invention;

[0040] Figure 5 This is a schematic diagram of the adapter module structure inserted in the semantic extractor in an embodiment of the present invention. Detailed Implementation

[0041] The technical solution of the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. Obviously, the described embodiments are only some embodiments of the present invention, not all embodiments, and the implementation of the present invention is not limited thereto. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0042] Example 1:

[0043] This invention provides a code vulnerability detection method based on parameter fine-tuning of a pre-trained code model. It leverages the rich semantic knowledge of the pre-trained model, exhibiting good generalization ability. By fusing semantic and structural features, it compensates for the lack of structural features in the pre-trained model, improving its detection performance. Adapter-tuning is used to fine-tune the model, preserving its rich pre-trained knowledge while achieving good fit to the target task, saving training resources and obtaining good performance. Simultaneously learning the semantic and structural features of the source code through model fusion significantly improves the efficiency and accuracy of code vulnerability detection, providing valuable guidance for the field of vulnerability detection.

[0044] like Figure 1 The flowchart shown illustrates a code vulnerability detection method based on fine-tuning of pre-trained model parameters. The code vulnerability detection method based on fine-tuning of pre-trained model parameters described in this invention includes the following steps:

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

[0046] First, a vulnerability dataset was constructed by acquiring vulnerability-related data from the open-source projects Qemu and FFmpeg. The selection of the vulnerability dataset is crucial for validating the model proposed in this invention. Many studies have achieved good performance on synthetic datasets, such as the Juliet and SARD datasets. The code in these datasets is artificially synthesized code targeting specific vulnerabilities. However, because this synthetic code differs significantly from real code, the performance of these models on real datasets drops sharply. The dataset selected in this embodiment comes from two open-source projects, Qemu and FFmpeg. In these projects, all commits are labeled as vulnerability fix commits and non-vulnerability fix commits. All vulnerability fix commits contain security-related keywords such as DoS and injection. To ensure the quality of the dataset, four professional security researchers conducted two rounds of data annotation and cross-validation. For each commit, the source code before the commit was extracted and labeled. The dataset contains 27k C language functions with a defect rate of 45%. Each sample contains the corresponding function source code and label. The function source code, containing many redundant line breaks and spaces, will be cleaned up in subsequent processing. Training and validation on real datasets allows us to verify the model's real-world performance in practical applications and facilitates its application in real-world production environments.

[0047] Secondly, the source code of the vulnerability dataset is preprocessed by replacing variable and function names with their corresponding symbol representations, removing redundant spaces and blank lines, and cleaning up information irrelevant to vulnerability features. To avoid lexical explosion caused by an excessive number of variable names in the source code, which could hinder model learning, source code standardization is performed.

[0048] like Figure 2 The diagram illustrates the source code preprocessing workflow. Specifically, it includes: replacing variable and function names in the source code with their corresponding symbol representations VARi and FUNCi; removing redundant spaces, blank lines, and other symbols from the source code of the dataset; and cleaning parts of the dataset that are irrelevant to vulnerability characteristics.

[0049] S2. The preprocessed source code is transformed to extract the corresponding token sequence and code attribute graph path. In subsequent steps, the extracted tokens and code attribute graph paths can be used as semantic and structural inputs to the code vulnerability detection model, respectively.

[0050] S21, the tokenizer pre-trained by the UniXcoder model converts each word in the source code into smaller subwords to obtain a token sequence corresponding to the source code. After cleaning, the source code needs to be converted into the input form of the UniXcoder model, which is a sequence composed of source code tokens. Converting the source code into the corresponding token sequence specifically includes: using the tokenizer pre-trained by the UniXcoder model to convert each word in the source code into smaller subwords by the Byte Pair Encoding (BPE) algorithm to obtain a source code token sequence. The input of the conversion process is the source code function f, and the result of the conversion is a token sequence represented as T=(t1, t2, …, t n ). After obtaining the source code token sequence, the input of the UniXcoder model can be constructed, represented 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, and the SEP token represents the end position of the input. Because the input of the transformer model requires a fixed length, it is assumed that the input length is fixed as 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 L, the input sequence will be truncated to L, and the last token is SEP.

[0051] As shown in the source code in Figure 2 , the token sequence T of the segmented source code is:

[0052]

[0053] Among them, The symbol represents some tokens that appear very rarely. This processing reduces the vocabulary size of the vocabulary and reduces the difficulty of model training.

[0054] S22. Using the Joern tool, the source code is converted into a corresponding program dependency graph. All data paths and control paths in the program dependency graph are extracted to obtain the code attribute graph paths of the source code. To compensate for the lack of code structure features in the pre-trained UniXcoder model, path representations are extracted from the source code attribute graph as structural features. Specifically, the Joern tool is used to convert the source code functions into corresponding program dependency graphs (PDG), represented as G = (V, E), where V represents a series of nodes corresponding to statements, each node having two attributes: node type d and code token t, and E represents the data or control dependencies between statements. Therefore, the graph paths include two types: control flow paths and data flow paths. 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 the 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's 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 a corresponding code in the source code, and the number corresponds to the line number of the code. There are different dependencies between 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 denoted as S = [[1,2,4,5],[1,2,3,5],[2,5],[1,3,5]].

[0056] S3. Construct 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, the overall network framework diagram of the code vulnerability detection model is shown, the code vulnerability detection model includes a semantic feature extractor and a structural feature extractor. Specifically, the semantic feature extractor adopts a UniXcoder model as a backbone network, an adapter adapter layer is added in each Transformer Encoder layer of the UniXcoder model, and a semantic feature vector of source code is extracted 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 the specific task, thereby improving the efficiency and speed of fine-tuning, that is, the code vulnerability detection model is fine-tuned in the parameter fine-tuning adapter-tuning manner.

[0058] The UniXcoder model is mainly based on the Transformer structure and mainly includes a tokenizer, an input encoder and a decoder. The tokenizer is pre-trained on a large amount of code corpus and can map the tokenized code token to the corresponding vector representation. In the pre-training process, the input includes natural language and code abstract syntax tree, so it has a good understanding of the structural features of the code abstract syntax tree. After the tokenizer, the core components of the UniXcoder model are the encoder and the decoder. The encoder and the decoder are stacked by multiple transformer blocks and can work in three modes, namely the encoder-only architecture, the encoder-decoder architecture and the decoder-only architecture. In this embodiment, the mode adopted is the encoder-only (only the decoder part), which is used to obtain a rich vector representation of the source code. Under the encoder-only architecture, each transformer block mainly includes 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 to speed up the model convergence and improve the model effect. After learning through multiple transformer blocks, the final hidden layer output is a semantic feature vector containing rich feature representation.

[0059] The UniXcoder model is a pre-trained code model. It is pre-trained on a large amount of source code, including source code text and Abstract Syntax Trees (ASTs), possessing rich code semantic knowledge. The code token sequence extracted in step S1 is used as model input. It first undergoes word embedding and position embedding, mapping each token to a high-dimensional vector. Then, the initial vector is passed through a 12-layer bidirectional Transformer Encoder layer to fully learn the interdependencies between tokens and the semantic features of the source code. The bidirectional Transformer Encoder uses a bidirectional multi-head attention mechanism to learn the relationships between preceding and following sequences in the source code. The hidden layer vector of the Nth layer can be represented as... Typically, the vector representation of the last layer is used as the semantic representation of the source code function. UniXcoder has 125 million training parameters, a massive number that consumes huge resources and hinders the model's fit to downstream tasks.

[0060] To obtain a better source code vector representation, an adapter layer is added to each Transformer Encoder. This adapter layer consists of two linear layers and one non-linear layer, with residual connections between each adapter layer to prevent adverse effects on model performance. During model parameter fine-tuning, the pre-trained parameters of the UniXcoder model remain unchanged; only the parameters of the adapter layer are updated. This significantly reduces the number of parameters that need updating, accelerates training iterations, and allows for good fitting to downstream tasks. Most current fine-tuning methods are based on full-parameter fine-tuning, which involves fine-tuning all parameters of the pre-trained model on the downstream task dataset. Taking UniXcoder as an example, the model has 125 million parameters. Training all parameters would place high demands on computing power and GPU memory, and would be time-consuming. Furthermore, downstream tasks often have limited data; without sufficient data, it's difficult to achieve good results when training so many parameters. Parameter fine-tuning aims to achieve the same effect by training a very small number of parameters while keeping a large number of pre-trained parameters unchanged. In UniXcoder, using adapter-tuning only requires training 2.87 million parameters, which is 2% of the number of parameters in full parameter fine-tuning. However, when performing detection on downstream datasets, the performance is slightly better than that of full parameter fine-tuning. The effect of training a very small number of parameters is comparable to that of full parameter fine-tuning.

[0061] like Figure 4 The diagram illustrates the semantic extractor architecture. Specifically, the process of the semantic feature extractor extracting semantic feature vectors from the source code includes:

[0062] First, the source code in the vulnerability dataset is input, segmented by the UniXcoder model pre-trained segmenter, and converted into a sequence of tokens.

[0063] Then the token sequence is converted into an initial vector representation, which includes word vector embedding and position vector embedding. Word vector embedding is a vector representing the semantics of the token, while position embedding is a vector representing the position of each token in the entire sequence, as the transformer model does not have token position information. Through the two embedding methods, the initial vector representation contains the semantics and position information of each token.

[0064] After obtaining the initial vector representation, a 12-layer bidirectional transformer encoder is needed to learn the relationship between tokens and obtain semantic feature vectors rich in semantic features, and learn the vulnerability feature information therein. As shown in Figure 4 , the Transformer encoder includes a multi-head attention module and a residual normalization layer, and the multi-head attention module and the residual normalization layer have an adapter adapter layer between them, and the residual normalization layer and the residual normalization layer have an adapter adapter layer between them. The multi-head attention mechanism is the key to its strong learning ability, which first maps the input to Q, K and V through linear mapping, and the conventional attention mechanism is:

[0065]

[0066] Where k represents the dimension of the hidden layer. The multi-head attention mechanism is to obtain attention through multiple attentions of the hidden layer vector, and then learn various vulnerability feature information from different subspaces of the code representation, which is specifically represented as for Q, K and V vectors with dimension d, the vectors are divided into h heads, each head has a dimension of d / h, and after each head is operated by self-attention, the output of all heads is spliced back to the original dimension, represented as MultiHead=Concat(head1,head2,…,head h )W O , where head i is the self-attention output of each head, which is calculated by the attention mechanism Attention formula above, and the combined output is input into the inserted adapter layer. As shown in Figure 5As shown, the structural diagram of the adapter layer is composed of two linear layers and one nonlinear layer. The output of the multi-head attention layer is input into the adapter layer, which first passes through the lower mapping layer. The neuron nodes of this layer are usually determined by the reduction parameter. For example, if reduction = 4, the output node number of the lower mapping layer = input dimension / reduction. Then, a nonlinear layer is used for feature learning. Finally, the upper mapping layer is used to map it to the same dimension as the input layer. The adapter layer is provided with a residual connection to prevent parameter degradation. Finally, the residual normalization layer is used to prevent parameter degradation and accelerate model convergence. Similarly, the output of the residual normalization layer is also input into the adapter layer for mapping. The initial vector is subjected to the above forward propagation process for 12 layers, and finally a semantic feature vector rich in semantic features is obtained.

[0067] The structural feature extractor uses a doc2vec model to learn the structural features in the code property graph path and 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 a whole text and can process texts of different lengths, and is suitable for processing cases where the number and length of function code property graph paths are different. After extracting all the structural paths of the source code, all the paths S = (s1, s w ,…,s m ) of each function are taken as a whole to construct the training corpus of doc2vec. The doc2vec model takes all the paths S of a function as a whole to obtain the structural feature vector of the current function through the context of other path vectors. Doc2vec is unsupervised, and the training corpus is the structural features extracted from each function in the data set. After training, the model has good capturing ability for the structural features of the code.

[0068] The structural feature vector and the semantic feature vector are spliced by the fusion layer to obtain the final composite feature vector. After obtaining the semantic features and structural features of the source code, the two features are combined and learned by the fusion layer. First, a linear layer is used to map the structural features and semantic features to the same dimension vector to prevent learning bias caused by different dimensions. Then, the structural feature vector and the semantic feature vector are spliced to obtain the final composite feature vector.

[0069] The classifier adopts a multi-layer perceptron (MLP), 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 capability to enable it to complete various classification tasks. Since the semantic extractor and the structure 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 composed of an input layer, a hidden layer and an output layer, and 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 the weighted sum is added to the corresponding bias of each node. The weighted sum is nonlinearly transformed by an activation function as the output of the node to the next layer, wherein the nonlinear transformation is completed by the activation function of the hidden layer, and the present application adopts a Relu activation function. Therefore, the multi-layer perceptron can perform feature fusion layer by layer through the combination of multiple layers of neurons and nonlinear activation functions, and learn the complex patterns of high-dimensional data.

[0070] S4, the extracted token sequence corresponding to the source code and the code property graph path are jointly trained to obtain an 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. A fixed batch size and corresponding labels are used as inputs, represented as (batch, seq, dimension, label), and the output obtained is (batch, dimension_1, label). The code property graph path obtained in S1 adopts the same batch size as the input of the doc2vec model, and is converted into a corresponding structural feature vector, represented as (batch, dimension_2, label). Then, it is converted into (batch, dimension_1, label) through a linear mapping layer. Finally, the result of the splicing is (batch, 2*dimension_1, label). First, as a classification task, the training loss function adopts a binary cross-entropy function to measure the difference between the predicted value and the true label:

[0072]

[0073] wherein y i represents the true label, which is 0 or 1 in this task, 0 indicating that the code does not contain a vulnerability, and 1 indicating that the code contains a vulnerability; pi represents the model prediction value, usually represents the probability that the sample belongs to class 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 optimizes the performance of the model classification.

[0074] Specifically, in the process of training the code vulnerability detection model, multiple dropout layers are used to prevent overfitting, KL divergence is used to punish different outputs of the model for the same input, and 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 results in slight differences in output for the same input. Therefore, KL divergence is applied to punish the model for inconsistent output for the same input, and KL divergence is used to measure the difference between two distributions:

[0075]

[0076] where P and Q are different distributions 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] where P1 and P2 are two different model output distributions for the same input, for example, if the input sample is 5, P(x i ) and Q(x i ) are the model outputs obtained by inputting the 5 samples into the model twice. In a binary classification task, it is a 5*2-dimensional vector, the first list of which represents the probability of the sample being predicted as 0, and the second list represents the probability of the sample being predicted as 1.

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

[0081] L=L ce +βL kl

[0082] where, β is a hyperparameter used for adjustment, is a parameter between 0-1, the loss function combines the binary cross-entropy loss function and the divergence loss function, which can improve the fitting effect of the model to the real data distribution, ensure the accuracy of the model output, and help improve the robustness of the model. Through the binary cross-entropy loss function, the model classification error can be punished, and because a large number of dropout layers are introduced in UniXcoder to prevent overfitting, the different outputs of the model for the same input can be punished through KL divergence.

[0083] After determining the model loss function, the parameters of the model are optimized through the back propagation algorithm to obtain the optimized code vulnerability detection model. Model training mainly includes 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 12 layers of 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 through multiple linear layers, and then the loss of the model output and the actual label value is calculated through the loss function described above, and the forward propagation ends. After obtaining the loss, back propagation is performed, that is, the loss is differentiated with respect to the parameters of each layer, and the derivative is obtained to update the parameters of the layer. During training, we set epoch to 10, that is, 10 iterations. Batch-size is set to 32, that is, each 32 samples are updated in turn, and the validation loss on the validation set is calculated after training every 100 batches. Finally, the model parameter value with the smallest loss on the validation set is saved as the optimal model, and the saved optimal model parameters are loaded for subsequent prediction.

[0084] S5, extracting the token sequence and code attribute graph path corresponding to the to-be-tested source code, and outputting whether the to-be-tested source code has a vulnerability 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 to-be-tested source code sequence, the detection framework extracting the code token and code attribute graph path corresponding to the to-be-tested source code, then obtaining the corresponding semantic feature vector and structure feature vector by the code semantic feature extractor and structure feature extractor respectively, splicing to obtain a composite feature vector, and finally inputting into the trained classifier to obtain the prediction value of whether there is a vulnerability. 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, and 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 application provides a code vulnerability detection method based on code pre-training model parameter fine-tuning, by adding an adapter adapter layer in each ransformer encoder layer of the UniXcoder model, the resources consumed by fine-tuning the pre-trained 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 small part of parameters. The structural feature vector and the semantic feature vector are spliced into a composite feature vector through a fusion layer, which makes up for the defect that the pre-trained model only uses code semantic features 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, and the semantic features and structural features are integrated, which greatly improves the performance of model detection. The present application can fully learn different features in the source code, has high detection accuracy, can accurately predict whether there is a vulnerability in the source code, and has good generalization.

[0087] The above embodiments are the preferred embodiments of the present application, but the embodiments of the present application are not limited by the above embodiments, and any changes, modifications, substitutions, combinations, simplifications made without departing from the spirit and principles of the present application shall be equivalent replacement methods, and shall be included in the protection scope of the present application.

Claims

1. A code vulnerability detection method based on code pre-training model parameter fine-tuning, characterized in that, The method comprises the following steps S1, constructing a vulnerability dataset, preprocessing the source code of the vulnerability dataset; S2, transforming the preprocessed source code to extract the token sequence and code attribute graph path corresponding to the source code; The transformation of the preprocessed source code to extract the token sequence and code attribute graph path corresponding to the source code comprises: Each word in the source code is converted into a smaller subword by a tokenizer pre-trained by a UniXcoder model to obtain the token sequence corresponding to the source code; The source code is converted into a corresponding program dependency graph by using a Joern tool, and all data paths and control paths in the program dependency graph are extracted to obtain the code attribute graph path of the source code; S3, constructing a code vulnerability detection model based on the UniXcoder model and the doc2vec model, the code vulnerability detection model comprising a semantic feature extractor, a structural feature extractor, a feature fusion layer, and a classifier; The semantic feature extractor adopts the UniXcoder model as a backbone network, and an adapter adapter layer is added in each bidirectional transformer encoder layer of the UniXcoder model, and the semantic feature vector of the source code is extracted by the semantic feature extractor; The structural feature extractor adopts the doc2vec model, and the structural features in the code attribute graph path are learned by using the doc2vec model to obtain the structural feature vector of the source code; The structural feature vector and the semantic feature vector are spliced by the fusion layer to obtain a composite feature vector; The classifier adopts a multilayer perceptron, learns the information of the composite feature vector by the multilayer perceptron, and outputs the result of whether the source code has a vulnerability; S4, jointly training the extracted token sequence and code attribute graph path corresponding to the source code to optimize the parameters of the model by a back propagation algorithm to obtain an optimized code vulnerability detection model; S5, extracting the token sequence and code attribute graph path corresponding to the to-be-tested source code, and outputting the result of whether the to-be-tested source code has a vulnerability by the optimized code vulnerability detection model.

2. The code vulnerability detection method based on code pre-training model parameter fine-tuning according to claim 1, characterized in that, The preprocessing of the source code of the vulnerability dataset comprises: Obtaining vulnerability-related data from open source projects Qemu and FFmpeg and constructing a vulnerability dataset; The source code of the vulnerability dataset is preprocessed, the variable names and function names in the source code are replaced by corresponding symbolic representations, the redundant space symbols and 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 code pre-training model parameter fine-tuning according to claim 1, characterized in that, The adapter adapter layer is composed of two linear layers and one nonlinear layer, and residual connections are used between the adapter adapter layers.

4. The code vulnerability detection method based on code pre-training model parameter fine-tuning according to claim 3, characterized in that, The semantic feature vector of the source code extracted by the semantic feature extractor comprises: The input source code is tokenized by the tokenizer pre-trained by the UniXcoder model to convert the token sequence in token units; The token sequence is converted into an initial vector representation, and the initial vector representation comprises a word vector embedding and a position vector embedding; The relationship between the tokens is learned by a 12-layer bidirectional transformer encoder to obtain a semantic feature vector of the source code.

5. The code vulnerability detection method based on code pre-training model parameter fine-tuning according to claim 4, characterized in that, The bidirectional transformer encoder comprises a multi-head attention module and a residual normalization layer, and 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 layers.

6. The code vulnerability detection method based on code pre-training model parameter fine-tuning according to claim 4, characterized in that, The extracted token sequence corresponding to the source code and the code property graph path are jointly trained to obtain a code vulnerability detection model, which comprises a plurality of dropout layers to prevent overfitting, a KL divergence penalty model to punish different outputs of the same input, and a loss function L determined based on a binary cross-entropy loss function and a divergence loss function, and the loss function L is calculated according to the following formula: L = L ce + βL kl ； where, L ce is a binary cross-entropy function, L kl is a divergence loss function, β is a hyperparameter used to adjust

Images