A source code vulnerability prediction method, device, electronic device and medium
Through the source code vulnerability prediction model of lightweight Transformer blocks and custom location encoding layer, traditional methods are solved in identifying a few categories of vulnerabilities and high computing resources, and efficient and accurate vulnerability prediction is achieved, reducing false positive rates and improving security.
Patent Information
- Application Number
- CN202411083393.7
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-08-08
- Publication Date
- 2025-07-25
- Estimated Expiration
- 2044-08-08
AI Technical Summary
When the prior art deals with highly unbalanced real-world vulnerability datasets, traditional machine learning methods are difficult to identify a few categories of vulnerabilities, and the traditional Transformer model consumes huge computing resources, making it impossible to effectively learn the syntax and structural dependencies of the source code.
The source code vulnerability prediction model is adopted with a lightweight Transformer block, a custom position coding layer, a convolutional layer and a global average pooling layer. Through text tokenization, one-hot encoding and multi-task learning, combined with the loss function of dynamic weights, the model training process is optimized to capture the long-distance dependencies and complex features of the source code.
It realizes high-accuracy and efficient vulnerability category prediction, reduces false positive rates, helps developers and security teams accurately detect system vulnerabilities and reduces security risks.
Smart Images
Figure CN118981773B_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the technical fields of program analysis and artificial intelligence, and particularly to a source code vulnerability prediction method, a source code vulnerability prediction device, an electronic device, and a computer-readable medium. Background Art
[0002] Source code vulnerability prediction is one of the key security measures in the software development process. During the software development process, the source code written by programmers may have various vulnerabilities, which may be exploited by hackers, leading to security vulnerabilities and system crashes. Traditional source code vulnerability prediction methods usually rely on using various machine learning algorithms to try to predict potential vulnerabilities from a large number of software attributes, such as using traditional machine learning techniques like decision trees, support vector machines, and neural networks. However, these existing methods often face difficulties in identifying minority-class vulnerabilities when dealing with highly imbalanced real-world vulnerability datasets, and traditional deep learning and machine learning methods are also unable to fully learn vulnerability features. Although the number of these minority-class vulnerabilities is small, they are often the most dangerous.
[0003] Multi-task learning is a strategy in machine learning. It can improve the generalization ability of a model by simultaneously learning multiple related tasks through a given model, and by sharing parameters between different tasks, forcing the model to learn cross-task general features. Especially when the data of some tasks is scarce, it can enable them to gain a deeper understanding from tasks with a larger amount of data, which helps the model better learn features from minority-class samples. However, there may be negative interactions between different tasks. Especially when the optimal representations of these tasks vary greatly, it may lead to a significant reduction in the performance of the model. In addition, the selection of different task weights will also have a great impact on the MTL model. Improper weight selection may cause the model to be biased towards certain tasks and thus ignore other tasks, resulting in some tasks not being able to learn feature data well, reducing the accuracy and generalization ability.
[0004] Program source code is highly structured, and the characteristics of syntax and structure determine the execution flow of the code. For example, a simple C language code segment can consist of the declaration and initialization of variable x, a conditional judgment, and a function call based on the conditional judgment to form a logical block, and the security of this block may be affected by the function return value. Without considering positional encoding, the self-attention mechanism equally considers all tokens within the entire sequence range, and the model will have difficulty understanding the direct dependency between the declaration and initialization of variables and the conditional judgment.
[0005] The Transformer model has become the mainstream model for processing sequence data, especially in the field of natural language processing (NLP), since it was proposed by Vaswani et al. in 2017. It is characterized by high computational efficiency, the ability to capture the relationship between any two elements, the ability to learn complex representations, and easy model expansion. However, traditional Transformer models often have a large number of parameters and consume huge computational resources, which will become a bottleneck when dealing with large-scale or complex source code datasets such as real-world vulnerability datasets. Summary of the Invention
[0006] In view of the above problems, the present invention is proposed to provide a source code vulnerability prediction method and a corresponding source code vulnerability prediction device, an electronic device, and a computer-readable medium that overcome the above problems or at least partially solve the above problems.
[0007] The present invention discloses a source code vulnerability prediction method, the method comprising:
[0008] Obtaining a source code dataset to be predicted;
[0009] Performing text tokenization processing on the source code dataset to be tested to obtain a plurality of source code token units;
[0010] Establishing a word frequency sorted vocabulary based on the source code text in the source code token units;
[0011] Converting the source code token units into source code feature vectors in one-hot encoding format according to the word frequency sorted vocabulary;
[0012] Inputting the source code feature vectors in one-hot encoding format into a source code vulnerability prediction model, and outputting a vulnerability category prediction result.
[0013] Optionally, the source code vulnerability prediction model includes a positional encoding layer, a lightweight Transformer block, a convolutional layer, a global average pooling layer, and a fully connected layer integrated into a Keras model; inputting the source code feature vectors in one-hot encoding format into the source code vulnerability prediction model and outputting a vulnerability category prediction result includes:
[0014] Mapping and converting the source code feature vectors in one-hot encoding format into source code embedding vectors.
[0015] Generating a corresponding position index for each source code embedding vector;
[0016] Embed the source code into a vector and input it into the positional encoding layer integrated into the Keras model. The positional encoding layer integrated into the Keras model calculates the positional encoding vector based on the total dimension of the source code embedding vector and the position index, and adds the positional encoding vector to the source code embedding vector to obtain a source code embedding vector containing positional information;
[0017] Pass the source code embedding vector containing positional information to the lightweight Transformer block to output a source code feature depth-enhanced fusion vector;
[0018] Process the source code feature depth-enhanced fusion vector through a convolutional layer and a global average pooling layer, and then pass it to multiple fully connected layers. The outputs of the multiple fully connected layers are respectively converted into the prediction probabilities of each vulnerability category through the softmax function.
[0019] Optionally, embed the source code into a vector and input it into the positional encoding layer integrated into the Keras model. The positional encoding layer integrated into the Keras model calculates the positional encoding vector based on the total dimension of the source code embedding vector and the position index, and adds the positional encoding vector to the source code embedding vector to obtain a source code embedding vector containing positional information, including:
[0020] Input the source code embedding vector into the positional encoding layer integrated into the Keras model. The positional encoding layer integrated into the Keras model calculates the sine and cosine values of each position in each dimension according to the total dimension of the source code embedding vector and the position index;
[0021] For each position, interleave and combine the calculated sine and cosine values into a positional encoding vector with the same dimension as the total dimension of the source code embedding vector;
[0022] Add the positional encoding vector to the source code embedding vector to obtain a source code embedding vector containing positional information.
[0023] Optionally, pass the source code embedding vector containing positional information to the lightweight Transformer block to output a source code feature depth-enhanced fusion vector, including:
[0024] Input the source code embedding vector containing positional information into the multi-head attention layer to output a multi-head attention fusion vector;
[0025] Perform a residual connection on the multi-head attention fusion vector and the source code embedding vector containing positional information and perform layer normalization to output a self-attention position-enhanced vector;
[0026] Pass the self-attention position-enhanced vector into a feed-forward network composed of a fully connected layer, a dropout layer, and a relu activation function to output a feed-forward network enhanced vector;
[0027] Perform residual connection and layer normalization on the feedforward network enhanced vector, and output the source code feature depth enhanced fusion vector.
[0028] Optionally, the generation of the source code vulnerability prediction model includes:
[0029] Obtain a source code public dataset; the source code public dataset includes data with CWE vulnerability labels and data without CWE vulnerability labels;
[0030] Identify the data without CWE vulnerability labels in the source code public dataset, randomly extract a preset number of data from the data without CWE vulnerability labels, and recombine the extracted preset number of data with the data with CWE vulnerability labels to generate an undersampled source code subset;
[0031] Divide the undersampled source code subset into a training dataset, a validation dataset, and a test dataset according to a preset ratio;
[0032] Perform text tokenization on the training dataset, the validation dataset, and the test dataset to obtain corresponding multiple source code token units;
[0033] Establish a corresponding word frequency sorted vocabulary based on the source code text in the source code token units;
[0034] Convert the training dataset, the validation dataset, and the test dataset into equal-length input matrices according to the source code token units;
[0035] Convert the input matrices of the training dataset, the validation dataset, and the test dataset into one-hot encoded source code feature vectors according to the word frequency sorted vocabulary;
[0036] Construct a source code vulnerability prediction model network;
[0037] Train the source code vulnerability prediction model network with the training dataset represented as one-hot encoded source code feature vectors, and perform validation and testing with the validation dataset and the test dataset represented as one-hot encoded source code feature vectors to obtain the source code vulnerability prediction model.
[0038] Optionally, training the source code vulnerability prediction model network with the training dataset represented as one-hot encoded source code feature vectors, and performing validation and testing with the validation dataset and the test dataset represented as one-hot encoded source code feature vectors to obtain the source code vulnerability prediction model includes:
[0039] The training dataset of source code feature vectors represented in one-hot encoding format is used to train the source code vulnerability prediction model network, and a loss function based on dynamic weights is used to calculate the total loss, and the parameters of the source code vulnerability prediction model network are updated based on the total loss;
[0040] During the training iteration process, when the change trend of the loss function on the validation dataset changes from decreasing in the previous iteration compared to the next iteration to increasing, record a state change. If the number of recorded times exceeds the preset early stopping threshold, stop training and save the model parameters when the loss function showed a decreasing trend for the last time during the current training process to obtain the source code vulnerability prediction model.
[0041] The present invention also discloses a source code vulnerability prediction device, and the device includes:
[0042] A to-be-predicted data acquisition module, configured to acquire a to-be-predicted source code dataset;
[0043] A first tokenization processing module, configured to perform text tokenization processing on the to-be-tested source code dataset to obtain a plurality of source code token units;
[0044] A first word frequency sorted vocabulary building module, configured to build a word frequency sorted vocabulary according to the source code text in the source code token units;
[0045] A first one-hot encoding module, configured to convert the source code token units into source code feature vectors in one-hot encoding format according to the word frequency sorted vocabulary;
[0046] A vulnerability prediction module, configured to input the source code feature vectors in one-hot encoding format into the source code vulnerability prediction model and output a vulnerability category prediction result.
[0047] Optionally, the source code vulnerability prediction model includes a positional encoding layer integrated into the Keras model, a lightweight Transformer block, a convolutional layer, a global average pooling layer, and a fully connected layer; the vulnerability prediction module includes:
[0048] An embedding vector generation sub-module, configured to map and convert the source code feature vectors in one-hot encoding format into source code embedding vectors.
[0049] A position index generation sub-module, configured to generate a corresponding position index for each source code embedding vector;
[0050] A positional encoding sub-module, configured to input the source code embedding vectors into the positional encoding layer integrated into the Keras model. The positional encoding layer integrated into the Keras model calculates a positional encoding vector according to the total dimension of the source code embedding vectors and the position index, and adds the positional encoding vector to the source code embedding vectors to obtain source code embedding vectors containing position information;
[0051] The self-attention fusion sub-module is used to pass the source code embedding vector containing position information to the lightweight Transformer block and output the source code feature depth-enhanced fusion vector;
[0052] The vulnerability category probability prediction sub-module is used to process the source code feature depth-enhanced fusion vector through a convolutional layer and a global average pooling layer and then pass it to multiple fully connected layers, and convert the outputs of the multiple fully connected layers into the prediction probabilities of each vulnerability category through the softmax function respectively.
[0053] Optionally, the position encoding sub-module includes:
[0054] The position sine and cosine value calculation unit is used to input the source code embedding vector into the position encoding layer integrated into the Keras model, and the position encoding layer integrated into the Keras model calculates the sine value and cosine value of each position on each dimension according to the total dimension of the source code embedding vector and the position index;
[0055] The position encoding vector generation unit is used to interleave and combine the calculated sine value and cosine value into a position encoding vector with the same dimension as the total dimension of the source code embedding vector for each position;
[0056] The position information fusion unit is used to add the position encoding vector and the source code embedding vector to obtain the source code embedding vector containing position information.
[0057] Optionally, the self-attention fusion sub-module includes:
[0058] The multi-head attention fusion vector generation unit is used to input the source code embedding vector containing position information into the multi-head attention layer and output the multi-head attention fusion vector;
[0059] The self-attention position enhancement vector generation unit is used to perform residual connection and layer normalization on the multi-head attention fusion vector and the source code embedding vector containing position information and output the self-attention position enhancement vector;
[0060] The feed-forward network enhancement vector generation unit is used to pass the self-attention position enhancement vector into the feed-forward network composed of a fully connected layer, a dropout layer, and a relu activation function and output the feed-forward network enhancement vector;
[0061] The source code feature depth-enhanced fusion vector generation unit is used to perform residual connection and layer normalization on the feed-forward network enhancement vector and output the source code feature depth-enhanced fusion vector.
[0062] Optionally, the device further includes:
[0063] The source code public dataset acquisition module is used to acquire the source code public dataset; the source code public dataset includes data with CWE vulnerability labels and data without CWE vulnerability labels;
[0064] The undersampling module is used to identify the data without CWE vulnerability labels in the source code public dataset, randomly extract a preset number of data from the data without CWE vulnerability labels, and recombine the extracted preset number of data with the data with CWE vulnerability labels to generate an undersampled source code subset;
[0065] The dataset splitting module is used to split the undersampled source code subset into a training dataset, a validation dataset, and a test dataset according to a preset ratio;
[0066] The second tokenization processing module is used to perform text tokenization processing on the training dataset, the validation dataset, and the test dataset to obtain corresponding multiple source code token units respectively;
[0067] The second word frequency sorted vocabulary building module is used to build a corresponding word frequency sorted vocabulary according to the source code text in the source code token units;
[0068] The equal-length matrix construction module is used to convert the training dataset, the validation dataset, and the test dataset into equal-length input matrices according to the source code token units;
[0069] The second one-hot encoding module is used to convert the input matrices of the training dataset, the validation dataset, and the test dataset into one-hot encoded source code feature vectors according to the word frequency sorted vocabulary;
[0070] The model network construction module is used to construct a source code vulnerability prediction model network;
[0071] The model training module is used to train the source code vulnerability prediction model network with the training dataset represented as one-hot encoded source code feature vectors, and perform validation and testing with the validation dataset and the test dataset represented as one-hot encoded source code feature vectors to obtain the source code vulnerability prediction model.
[0072] Optionally, the model training module includes:
[0073] The model network parameter optimization sub-module is used to train the source code vulnerability prediction model network with the training dataset represented as one-hot encoded source code feature vectors, calculate the total loss using a loss function based on dynamic weights, and update the source code vulnerability prediction model network parameters based on the total loss;
[0074] The source code vulnerability prediction model generation sub-module is used to record a state change during the training iteration when the change trend of the loss function on the validation data set changes from decreasing from the previous iteration to the next iteration to increasing. If the number of recorded times exceeds the preset early stopping threshold, the training is stopped, and the model parameters when the loss function shows a decreasing trend for the last time during the current training process are saved to obtain the source code vulnerability prediction model.
[0075] The present invention also discloses an electronic device, including a processor, a communication interface, a memory, and a communication bus. Among them, the processor, the communication interface, and the memory complete communication with each other through the communication bus;
[0076] The memory is used to store a computer program;
[0077] When the processor is used to execute the program stored on the memory, it implements the source code vulnerability prediction method as described in the present invention.
[0078] The present invention also discloses one or more computer-readable media, on which instructions are stored. When executed by one or more processors, the instructions cause the processors to execute the source code vulnerability prediction method as described in the present invention.
[0079] The present invention includes the following advantages:
[0080] In the source code vulnerability prediction method of the present invention, by obtaining a source code data set to be predicted, performing text tokenization processing on the source code data set to be measured to obtain multiple source code token units, establishing a word frequency sorted vocabulary table according to the source code text in the source code token units, converting the source code token units into source code feature vectors in one-hot encoding format according to the word frequency sorted vocabulary table, and inputting the source code feature vectors in one-hot encoding format into the source code vulnerability prediction model to output a vulnerability category prediction result. The present invention realizes accurate prediction of vulnerability categories by preprocessing the source code data set to be predicted and processing the preprocessed source code data set to be predicted through a trained source code vulnerability prediction model, helps developers and security teams accurately and efficiently discover vulnerabilities in the system, and reduces potential security risks. Description of the Drawings
[0081] Figure 1 is a flowchart of the steps of a source code vulnerability prediction method provided by an embodiment of the present invention;
[0082] Figure 2 is a flowchart of the training of a source code vulnerability prediction model provided by an embodiment of the present invention;
[0083] Figure 3 is a structural block diagram of a source code vulnerability prediction device provided by an embodiment of the present invention;
[0084] Figure 4 is a block diagram of an electronic device provided by an embodiment of the present invention;
[0085] Figure 5 is a schematic diagram of a computer-readable medium provided by an embodiment of the present invention. Detailed implementation manners
[0086] To make the above objects, features, and advantages of the present invention more obvious and understandable, the present invention will be further described in detail below with reference to the accompanying drawings and specific implementation manners.
[0087] Referring to Figure 1 , a step flowchart of a source code vulnerability prediction method provided by an embodiment of the present invention is shown, which may specifically include the following steps:
[0088] Step 101, obtain a source code dataset to be predicted;
[0089] Step 102, perform text tokenization processing on the source code dataset to be tested to obtain a plurality of source code token units;
[0090] Step 103, establish a word frequency sorted vocabulary based on the source code text in the source code token unit;
[0091] Step 104, convert the source code token unit into a one-hot encoded source code feature vector according to the word frequency sorted vocabulary;
[0092] Step 105, input the one-hot encoded source code feature vector into a source code vulnerability prediction model, and output a vulnerability category prediction result.
[0093] In the embodiment of the present invention, after obtaining the source code dataset to be predicted, some preprocessing can be performed on the source code dataset to be predicted to accurately predict the vulnerability category. Specifically, the source code dataset to be predicted can be first cleaned and standardized to ensure data consistency and accuracy, eliminate noise, and enable the model to better identify the vulnerability characteristics of the source code. Then, text tokenization processing can be performed on the pre-cleaned and pre-standardized source code dataset to be predicted, that is, convert the text data into a series of source code token units token for constructing a numerical input that the model can understand, and create a vocabulary sorted according to word frequency based on all source code texts in the training set. Further, according to the word frequency sorted vocabulary, the feature columns in the text-tokenized source code dataset to be predicted can be converted into a one-hot encoded format using one-hot encoding, so that each category has a clear target vector representation, further enhancing the data representation ability and laying a foundation for more accurate prediction of the vulnerability category.
[0094] Through these preprocessings of the source code dataset to be predicted, model input data that can improve the prediction accuracy and is suitable for the source code vulnerability prediction model of the present invention is obtained. Finally, the source code feature vectors in one-hot encoding format obtained by the preprocessing are input into the source code vulnerability prediction model. After the source code vulnerability prediction model processes the source code feature vectors in one-hot encoding format, it outputs the vulnerability category prediction results, thereby obtaining the probabilities of various vulnerability categories existing in the source code to be predicted, and realizing the accurate prediction of vulnerability categories through the source code vulnerability prediction model of the present invention, helping developers and security teams to accurately and efficiently discover vulnerabilities in the system and reduce potential security risks.
[0095] In an embodiment of the present invention, the source code vulnerability prediction model includes a positional encoding layer integrated into the Keras model, a lightweight Transformer block, a convolutional layer, a global average pooling layer, and a fully connected layer; inputting the source code feature vectors in one-hot encoding format into the source code vulnerability prediction model and outputting the vulnerability category prediction results includes:
[0096] Mapping and converting the source code feature vectors in one-hot encoding format into source code embedding vectors.
[0097] Generating a corresponding position index for each source code embedding vector;
[0098] Inputting the source code embedding vectors into the positional encoding layer integrated into the Keras model. The positional encoding layer integrated into the Keras model calculates the positional encoding vectors according to the total dimension of the source code embedding vectors and the position indexes, and adds the positional encoding vectors to the source code embedding vectors to obtain source code embedding vectors containing positional information;
[0099] Transmitting the source code embedding vectors containing positional information to the lightweight Transformer block and outputting source code feature depth-enhanced fusion vectors;
[0100] Processing the source code feature depth-enhanced fusion vectors through the convolutional layer and the global average pooling layer and then transmitting them to multiple fully connected layers, and converting the outputs of the multiple fully connected layers into the prediction probabilities of each vulnerability category through the softmax function respectively.
[0101] The source code vulnerability prediction model of the embodiment of the present invention may include a position encoding layer defined by being integrated into a Keras model, as well as a lightweight Transformer block, a convolutional layer, a global average pooling layer, and a fully connected layer. The source code vulnerability prediction model first receives a source code feature vector in one-hot encoding format, generates a source code embedding vector through mapping transformation, and uses the position encoding layer to generate and add position information to each embedding vector to capture the sequential characteristics and long-distance dependencies in the source code. Subsequently, these source code embedding vectors containing position information are fed into the lightweight Transformer block, which further enhances and fuses the deep features of the source code through the self-attention mechanism, thereby extracting richer context information.
[0102] After feature extraction and fusion, the model uses a convolutional layer and a global average pooling layer to further process the source code feature depth-enhanced fusion vector to reduce the number of parameters, prevent overfitting, and extract the most representative features. Finally, these features are passed to multiple fully connected layers, and the output of each fully connected layer is converted into the prediction probability of the corresponding vulnerability category through the softmax function. The vulnerability categories predicted by the source code vulnerability prediction model of the embodiment of the present invention have high accuracy, high efficiency, low false positives, and high true negatives, reducing the false alarm rate, helping developers and security teams to accurately and efficiently discover vulnerabilities in the system, and reducing potential security risks.
[0103] In an embodiment of the present invention, inputting the source code embedding vector into the position encoding layer integrated into the Keras model, the position encoding layer integrated into the Keras model calculates the position encoding vector according to the total dimension and position index of the source code embedding vector, and adds the position encoding vector to the source code embedding vector to obtain a source code embedding vector containing position information, including:
[0104] Inputting the source code embedding vector into the position encoding layer integrated into the Keras model, the position encoding layer integrated into the Keras model calculates the sine value and cosine value of each position in each dimension according to the total dimension and position index of the source code embedding vector;
[0105] For each position, the calculated sine value and cosine value are interleaved and combined into a position encoding vector with the same dimension as the total dimension of the source code embedding vector;
[0106] Adding the position encoding vector to the source code embedding vector to obtain a source code embedding vector containing position information.
[0107] The position encoding layer of the embodiment of the present invention is integrated into the Keras model, and the floating-point conversion of constants is used to improve the calculation accuracy. Specifically, the position encoding layer provides a unique position encoding PE for each position θ in the source code token sequence, where each position encoding PE is composed of the sine value H of each position θ z and the cosine value H y Combined generation, the sine value H z and the cosine value H y The formula is as follows:
[0108]
[0109] where d i is the dimension index of the θ position in the embedding vector. The sine value is calculated using formula (1) for even indices, and the cosine value is calculated using formula (2) for odd indices. h m is the total dimension of the embedding vector.
[0110] After obtaining the sine value H z and the cosine value H corresponding to each position θ y , calculate a vector of dimension 2×h composed of alternating sine and cosine values m , that is, the complete position encoding vector PE θ , the formula of PE θ is as follows:
[0111]
[0112] where, is the sine value calculated for each position θ, is the cosine value calculated for each position θ.
[0113] For each position index θ in the source code vulnerability prediction model, add the position encoding vector PE of this position θ to the embedding vector of this position to obtain the embedding vector X with position information added.
[0114] By adding a custom position encoding layer, the model can identify the relationship before and after the code line position, thereby understanding the role of variables in conditional judgments and capturing syntactic and structural dependencies. In addition, in the C language code segment, the semantics of each code line is closely related to its position in the code. Through the position encoding layer, the model can better understand the logical relationship before and after the code lines, reduce the confusion of the model's code logic understanding, and reduce the false alarm rate.
[0115] In an embodiment of the present invention, the source code embedding vector containing position information is passed to the lightweight Transformer block, and the source code feature depth enhanced fusion vector is output, including:
[0116] Embed the source code containing location information into the vector input of the multi-head attention layer, and output the multi-head attention fusion vector;
[0117] Perform residual connection on the multi-head attention fusion vector and the source code embedding vector containing location information and perform layer normalization, and output the self-attention position enhancement vector;
[0118] Input the self-attention position enhancement vector into the feed-forward network composed of a fully connected layer, a dropout layer, and a relu activation function, and output the feed-forward network enhancement vector;
[0119] Perform residual connection and layer normalization on the feed-forward network enhancement vector, and output the source code feature depth enhancement fusion vector.
[0120] The lightweight Transformer block linearly transforms the embedded vector X with location information added to obtain Q, K, and V, and then passes Q, K, and V through four attention heads respectively to obtain the output Att of each attention head i (i = 1, 2, 3, 4), then splice the outputs of all attention heads and obtain the multi-head attention output MH(Q, K, V) through a linear transformation, and then perform residual connection on the multi-head attention output MH(Q, K, V) and the input vector Q and perform layer normalization to obtain the output O attn , and finally input the output O after layer normalization attn into the feed-forward neural network (FFN), and then perform residual connection and layer normalization to obtain the output O t , O t The formula is as follows:
[0121]
[0122] O attn = LayerNorm(Q + MH(Q, K, V)) (5)
[0123] O t = LayerNorm(O attn + FFN·O attn ) (6)
[0124] Among them, W is the linear mapping matrix that maps the input to different representation spaces, W O is the matrix used for the mapping after splicing, LayerNorm is the layer normalization function, and FFN is the calculation function of the feed-forward network.
[0125] In an embodiment of the present invention, the generation of the source code vulnerability prediction model includes:
[0126] Obtain an open-source code public dataset; the open-source code public dataset contains data with CWE vulnerability labels and data without CWE vulnerability labels;
[0127] Identify the data in the open-source code public dataset that does not carry CWE vulnerability labels, randomly extract a preset number of data from the data without CWE vulnerability labels, and recombine the extracted preset number of data with the data with CWE vulnerability labels to generate an undersampled source code subset;
[0128] Divide the undersampled source code subset into a training dataset, a validation dataset, and a test dataset according to a preset ratio;
[0129] Perform text tokenization on the training dataset, the validation dataset, and the test dataset to obtain corresponding multiple source code token units;
[0130] Establish a corresponding word frequency sorted vocabulary based on the source code text in the source code token units;
[0131] Convert the training dataset, the validation dataset, and the test dataset into equal-length input matrices according to the source code token units;
[0132] Convert the input matrices of the training dataset, the validation dataset, and the test dataset into one-hot encoded source code feature vectors according to the word frequency sorted vocabulary;
[0133] Construct a source code vulnerability prediction model network;
[0134] Train the source code vulnerability prediction model network using the training dataset represented as one-hot encoded source code feature vectors, and perform validation and testing using the validation dataset and the test dataset represented as one-hot encoded source code feature vectors to obtain the source code vulnerability prediction model.
[0135] In an embodiment of the present invention, training the source code vulnerability prediction model network using the training dataset represented as one-hot encoded source code feature vectors, and performing validation and testing using the validation dataset and the test dataset represented as one-hot encoded source code feature vectors to obtain the source code vulnerability prediction model includes:
[0136] Train the source code vulnerability prediction model network using the training dataset represented as one-hot encoded source code feature vectors, calculate the total loss using a loss function based on dynamic weights, and update the parameters of the source code vulnerability prediction model network based on the total loss;
[0137] During the training iteration process, when the change trend of the loss function on the validation dataset changes from decreasing from the previous iteration to the next iteration to increasing, record a state change. If the recorded number exceeds the preset early stopping threshold, stop the training and save the model parameters when the loss function showed a decreasing trend for the last time during the current training process to obtain the source code vulnerability prediction model.
[0138] Referring to Figure 2 , the embodiment of the present invention uses the publicly disclosed real-world source code vulnerability dataset to train the constructed source code vulnerability prediction model network to obtain the source code vulnerability prediction model.
[0139] Specifically, the source code vulnerability prediction model network designed by the present invention is a multi-task learning model framework MTLPT based on a custom position encoding layer and lightweight Transformer blocks. The custom position encoding layer is integrated into the Keras model, and the floating-point conversion of constants is used to improve the calculation accuracy. The lightweight Transformer block includes a multi-head self-attention layer with 4 heads to reduce the computational complexity, a Dropout layer used to randomly discard a part of the features during training to prevent the MTLPT model from over-relying on the sample features of certain specific vulnerabilities in the training data and reduce the overfitting phenomenon, a layer normalization layer to maintain the stability of internal covariates and accelerate the convergence of the MTLPT model, a feed-forward network composed of two fully connected layers, a Dropout layer, and a relu activation function to ensure that the same fully connected neural network is applied at each position θ, thereby further enhancing the representation ability of the MTLPT model. It also includes a residual connection to help alleviate the problem of gradient vanishing or explosion in deep networks. And a convolutional layer and a global average pooling layer are added to improve the ability to capture dependencies and reduce the data dimension, reducing the computational amount. Finally, multiple fully connected Dense layers and softmax are used to make decisions for each task respectively to obtain the vulnerability prediction results. And a dynamic weight-based loss function is defined to dynamically adjust their weights in the total loss function according to the loss change rate of each task.
[0140] Before training the source code vulnerability prediction model network, some preprocessing is first performed on the source code public dataset. Specifically, undersampling is performed on the source code public dataset to alleviate the problem of data imbalance, that is, undersampling is performed on the data in the original dataset that does not contain CWE vulnerabilities. 10% of the data without CWE vulnerabilities in the original dataset is randomly selected and recombined with the data with CWE labels in the original dataset to generate an undersampled subset. And the subset of the undersampled source code dataset is cleaned and standardized to ensure the consistency and accuracy of the data, eliminate noise, and enable the model to better learn the vulnerability characteristics of the source code. And the training dataset, validation dataset, and test dataset are randomly divided in the ratio of 8:1:1.
[0141] Furthermore, perform text tokenization on the training dataset, validation dataset, and test dataset: convert the text data into a series of tokens for constructing numerical inputs that the model can understand, and create a vocabulary sorted by word frequency based on all the source code texts in the training set. Then, convert each dataset (training set, validation set, and test set) into an input matrix of equal length I max (D num , I max ), where D num is the number of data, ensuring the consistency of the model input. Use one-hot encoding to convert the feature columns in the dataset into one-hot encoding format, so that each category has a clear target vector representation, further enhancing the expressive ability of the data.
[0142] Use the one-hot encoded training set as the input of the model for training. At the end of each training iteration, call the loss function based on dynamic weights to calculate the total loss, learn the correlation between learning tasks, reduce the interference of differences between different tasks, improve the model performance, and control the saving of the best performance model through the change of the loss function of the model on the validation dataset after each iteration. Set the early stopping function threshold to prevent overfitting. The early stopping function threshold can be set according to the actual situation, such as 5, 8, 10, etc. The loss function L t based on dynamic weights is defined as follows:
[0143]
[0144] where L i is the loss function for each task, and the cross-entropy loss at the end of each iteration for each task is M i , N is the number of tasks, and S is a parameter used to control the smoothness of the weights (S ∈
[0145] (0, 5]), which is determined to be 0.7 after experimental hyperparameter tuning in the present invention.
[0146] The source code vulnerability prediction method of the present invention combines a multi-task learning framework (MTL, multi-task learning), a custom position encoding layer (PE, position encoding layer), and a multi-class vulnerability prediction model of Transformer blocks, and uses a dynamic weight allocation strategy to allocate task weights, so as to predict vulnerabilities more accurately and efficiently. This method mainly includes three parts: First, by utilizing the correlation between different vulnerabilities, a model simultaneously learns multiple related tasks to improve the model generalization ability, and by sharing parameters between different tasks, forcing the model to learn cross-task general features and improving the feature extraction ability of minority-class vulnerabilities. Second, through a custom position encoding layer and lightweight Transformer blocks, the number of parameters is reduced and the computing efficiency is improved, while retaining the core advantages and performance of Transformer, capturing syntax and structural dependencies, and effectively learning the complex features of source code vulnerabilities. Third, according to the loss change rate of each task, dynamically adjust their weights in the total loss function and adopt a multi-head attention mechanism to enable the model to automatically learn the correlation between tasks, reducing the interference of differences between different tasks. Through the multi-task learning framework, the present invention combines a custom lightweight Transformer block and a position encoding layer, thereby learning the potential relationship between different vulnerability types, effectively capturing the long-distance dependencies and complex context information of the source code, improving the accuracy and efficiency of vulnerability prediction, and enhancing the sensitivity of the model to rare but serious vulnerabilities. In addition, the present invention introduces a loss function based on dynamic weights, dynamically adjusts the loss weights according to the prediction difficulty of different tasks, effectively alleviates the challenges brought by unbalanced data, realizes high-accuracy, high-efficiency, low false positive, and high true negative vulnerability prediction, reduces the false alarm rate, helps developers and security teams accurately and efficiently discover vulnerabilities in the system, and reduces potential security risks.
[0147] It should be noted that for the method embodiments, for the sake of simple description, they are all expressed as a series of action combinations. However, those skilled in the art should know that the embodiments of the present invention are not limited by the described action sequence, because according to the embodiments of the present invention, certain steps can be performed in other sequences or simultaneously. Secondly, those skilled in the art should also know that the embodiments described in the specification are all preferred embodiments, and the actions involved are not necessarily essential for the embodiments of the present invention.
[0148] Referring to Figure 3 , a structural block diagram of a source code vulnerability prediction device provided in an embodiment of the present invention is shown, which may specifically include the following modules:
[0149] A to-be-predicted data acquisition module 301, configured to acquire a to-be-predicted source code data set;
[0150] The first tokenization processing module 302 is used to perform text tokenization processing on the to-be-tested source code dataset to obtain a plurality of source code token units;
[0151] The first word frequency sorted vocabulary building module 303 is used to build a word frequency sorted vocabulary according to the source code text in the source code token units;
[0152] The first one-hot encoding module 304 is used to convert the source code token units into source code feature vectors in one-hot encoding format according to the word frequency sorted vocabulary;
[0153] The vulnerability prediction module 305 is used to input the source code feature vectors in one-hot encoding format into the source code vulnerability prediction model and output the vulnerability category prediction result.
[0154] Optionally, the source code vulnerability prediction model includes a position encoding layer, a lightweight Transformer block, a convolutional layer, a global average pooling layer, and a fully connected layer integrated into the Keras model; the vulnerability prediction module includes:
[0155] The embedding vector generation sub-module is used to map and convert the source code feature vectors in one-hot encoding format into source code embedding vectors.
[0156] The position index generation sub-module is used to generate a corresponding position index for each source code embedding vector;
[0157] The position encoding sub-module is used to input the source code embedding vectors into the position encoding layer integrated into the Keras model. The position encoding layer integrated into the Keras model calculates the position encoding vectors according to the total dimension of the source code embedding vectors and the position indexes, and adds the position encoding vectors to the source code embedding vectors to obtain source code embedding vectors containing position information;
[0158] The self-attention fusion sub-module is used to pass the source code embedding vectors containing position information to the lightweight Transformer block and output the source code feature depth enhanced fusion vectors;
[0159] The vulnerability category probability prediction sub-module is used to pass the source code feature depth enhanced fusion vectors through the convolutional layer and the global average pooling layer and then to a plurality of fully connected layers, and convert the outputs of the plurality of fully connected layers into the prediction probabilities of each vulnerability category through the softmax function respectively.
[0160] Optionally, the position encoding sub-module includes:
[0161] A position sine and cosine value calculation unit is configured to input the source code embedding vector into the position encoding layer integrated into the Keras model. The position encoding layer integrated into the Keras model calculates the sine value and cosine value of each position in each dimension according to the total dimension of the source code embedding vector and the position index.
[0162] A position encoding vector generation unit is configured to, for each position, interleave and combine the calculated sine value and cosine value into a position encoding vector with the same dimension as the total dimension of the source code embedding vector.
[0163] A position information fusion unit is configured to add the position encoding vector and the source code embedding vector to obtain a source code embedding vector containing position information.
[0164] Optionally, the self-attention fusion sub-module includes:
[0165] A multi-head attention fusion vector generation unit is configured to input the source code embedding vector containing position information into the multi-head attention layer and output a multi-head attention fusion vector.
[0166] A self-attention position enhancement vector generation unit is configured to perform a residual connection between the multi-head attention fusion vector and the source code embedding vector containing position information and perform layer normalization, and output a self-attention position enhancement vector.
[0167] A feed-forward network enhancement vector generation unit is configured to input the self-attention position enhancement vector into a feed-forward network composed of a fully-connected layer, a dropout layer, and a relu activation function, and output a feed-forward network enhancement vector.
[0168] A source code feature depth enhancement fusion vector generation unit is configured to perform a residual connection and layer normalization on the feed-forward network enhancement vector, and output a source code feature depth enhancement fusion vector.
[0169] Optionally, the device further includes:
[0170] A source code public dataset acquisition module is configured to acquire a source code public dataset; the source code public dataset includes data with CWE vulnerability labels and data without CWE vulnerability labels.
[0171] An undersampling module is configured to identify the data without CWE vulnerability labels in the source code public dataset, randomly extract a preset number of data from the data without CWE vulnerability labels, and recombine the extracted preset number of data with the data with CWE vulnerability labels to generate an undersampled source code subset.
[0172] A dataset splitting module is configured to split the undersampled source code subset into a training dataset, a validation dataset, and a test dataset according to a preset ratio.
[0173] A second tokenization processing module, configured to perform text tokenization processing on the training data set, the validation data set, and the test data set, respectively obtaining a plurality of corresponding source code token units;
[0174] A second word frequency sorted vocabulary building module, configured to build a corresponding word frequency sorted vocabulary according to the source code text in the source code token units;
[0175] An equal-length matrix construction module, configured to convert the training data set, the validation data set, and the test data set into equal-length input matrices according to the source code token units;
[0176] A second one-hot encoding module, configured to convert the input matrices of the training data set, the validation data set, and the test data set into source code feature vectors in one-hot encoding format according to the word frequency sorted vocabulary;
[0177] A model network construction module, configured to construct a source code vulnerability prediction model network;
[0178] A model training module, configured to train the source code vulnerability prediction model network by using the training data set represented as source code feature vectors in one-hot encoding format, and perform validation and testing by using the validation data set and the test data set represented as source code feature vectors in one-hot encoding format, obtaining the source code vulnerability prediction model.
[0179] Optionally, the model training module includes:
[0180] A model network parameter optimization sub-module, configured to train the source code vulnerability prediction model network by using the training data set represented as source code feature vectors in one-hot encoding format, calculate the total loss by using a loss function based on dynamic weights, and update the source code vulnerability prediction model network parameters based on the total loss;
[0181] A source code vulnerability prediction model generation sub-module, configured to record a state change during the training iteration process when the change trend of the loss function on the validation data set changes from decreasing in the previous iteration compared to the next iteration to increasing. If the recorded number exceeds a preset early stopping threshold, stop the training, and save the model parameters when the loss function shows a decreasing trend for the last time during the current training process, obtaining the source code vulnerability prediction model.
[0182] For the apparatus embodiment, since it is basically similar to the method embodiment, the description is relatively simple. For related parts, refer to the partial description of the method embodiment.
[0183] In addition, an embodiment of the present invention further provides an electronic device, such as Figure 4As shown, it includes a processor 401, a communication interface 402, a memory 403, and a communication bus 404. Among them, the processor 401, the communication interface 402, and the memory 403 communicate with each other through the communication bus 404.
[0184] The memory 403 is used to store computer programs.
[0185] When the processor 401 is used to execute the program stored in the memory 403, it implements the source code vulnerability prediction method described in the above embodiments.
[0186] The communication bus mentioned in the above terminal can be a Peripheral Component Interconnect (PCI) bus or an Extended Industry Standard Architecture (EISA) bus, etc. This communication bus can be divided into an address bus, a data bus, a control bus, etc. For the sake of simplicity, only a thick line is used to represent it in the figure, but it does not mean that there is only one bus or one type of bus.
[0187] The communication interface is used for communication between the above terminal and other devices.
[0188] The memory may include a Random Access Memory (RAM), or may also include a non-volatile memory, such as at least one disk memory. Optionally, the memory may also be at least one storage device located far from the aforementioned processor.
[0189] The above-mentioned processor may be a general-purpose processor, including a Central Processing Unit (CPU), a Network Processor (NP), etc.; it may also be a Digital Signal Processor (DSP), an Application Specific Integrated Circuit (ASIC), a Field-Programmable Gate Array (FPGA), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components.
[0190] Such as Figure 5As shown, in another embodiment provided by the present invention, a computer-readable storage medium 501 is further provided. Instructions are stored in the computer-readable storage medium, and when it runs on a computer, it causes the computer to execute the source code vulnerability prediction method described in the above embodiment.
[0191] In another embodiment provided by the present invention, a computer program product containing instructions is further provided. When it runs on a computer, it causes the computer to execute the source code vulnerability prediction method described in the above embodiment.
[0192] In the above embodiment, it can be implemented in whole or in part by software, hardware, firmware, or any combination thereof. When implemented using software, it can be implemented in whole or in part in the form of a computer program product. The computer program product includes one or more computer instructions. When the computer program instructions are loaded and executed on a computer, the processes or functions described in the embodiments of the present invention are generated in whole or in part. The computer can be a general-purpose computer, a special-purpose computer, a computer network, or other programmable devices. The computer instructions can be stored in a computer-readable storage medium, or transmitted from one computer-readable storage medium to another. For example, the computer instructions can be transmitted from one website, computer, server, or data center to another website, computer, server, or data center in a wired manner (such as coaxial cable, optical fiber, Digital Subscriber Line (DSL)) or wirelessly (such as infrared, wireless, microwave, etc.). The computer-readable storage medium can be any available medium that the computer can access, or a data storage device such as a server or data center that includes one or more integrated available media. The available medium can be a magnetic medium (such as a floppy disk, hard disk, magnetic tape), an optical medium (such as a DVD), or a semiconductor medium (such as a Solid State Disk (SSD)).
[0193] It should be noted that in this article, relational terms such as first and second are only used to distinguish one entity or operation from another entity or operation, and do not necessarily require or imply any actual relationship or order between these entities or operations. Moreover, the term "comprising", "including", or any other variant thereof is intended to cover non-exclusive inclusion, so that a process, method, article, or device including a series of elements not only includes those elements, but also includes other elements not expressly listed, or elements inherent to such process, method, article, or device. Without further limitation, an element defined by the statement "including a..." does not exclude the existence of additional identical elements in the process, method, article, or device including the element.
[0194] Each embodiment in this specification is described in a related manner. For the same or similar parts among the embodiments, reference can be made to each other. Each embodiment focuses on the differences from other embodiments. In particular, for the system embodiment, since it is basically similar to the method embodiment, the description is relatively simple. For the related parts, reference can be made to the corresponding description in the method embodiment.
[0195] The above description is only for the preferred embodiments of the present invention and is not intended to limit the protection scope of the present invention. Any modifications, equivalent replacements, improvements, etc. made within the spirit and principle of the present invention are all included in the protection scope of the present invention.
Claims
1. A source code vulnerability prediction method, characterized in that, The method includes: Obtaining a source code dataset to be predicted; Performing text tokenization on the source code dataset to be predicted to obtain multiple source code token units; Establishing a word frequency sorted vocabulary based on the source code text in the source code token units; Converting the source code token units into source code feature vectors in one-hot encoding format according to the word frequency sorted vocabulary; Inputting the source code feature vectors in one-hot encoding format into a source code vulnerability prediction model and outputting a vulnerability category prediction result; The generation of the source code vulnerability prediction model includes: Obtaining a source code public dataset; the source code public dataset includes data with CWE vulnerability labels and data without CWE vulnerability labels; Identifying the data without CWE vulnerability labels in the source code public dataset, randomly extracting a preset number of data from the data without CWE vulnerability labels, and recombining the extracted preset number of data with the data with CWE vulnerability labels to generate an undersampled source code subset; Dividing the undersampled source code subset into a training dataset, a validation dataset, and a test dataset according to a preset ratio; Performing text tokenization on the training dataset, the validation dataset, and the test dataset to respectively obtain corresponding multiple source code token units; Establishing corresponding word frequency sorted vocabularies based on the source code text in the source code token units; Converting the training dataset, the validation dataset, and the test dataset into input matrices of equal length according to the source code token units; Converting the input matrices of the training dataset, the validation dataset, and the test dataset into source code feature vectors in one-hot encoding format according to the word frequency sorted vocabulary; Constructing a source code vulnerability prediction model network; Training the source code vulnerability prediction model network with the training dataset represented as source code feature vectors in one-hot encoding format, calculating the total loss using a loss function based on dynamic weights, and updating the parameters of the source code vulnerability prediction model network based on the total loss; During the training iteration process, when the change trend of the loss function on the validation dataset changes from decreasing in the previous iteration compared to the next iteration to increasing, record a state change. If the recorded number exceeds a preset early stopping threshold, stop training and save the model parameters when the loss function showed a decreasing trend in the last time during the current training process to obtain the source code vulnerability prediction model; Loss function L based on dynamic weights t is defined as follows: Among them, L i is the loss function for each task, and the cross-entropy loss of each task at the end of each iteration is M i , N is the number of tasks, and S is a parameter used to control the smoothness of the weights, where S ∈ (0, 5].
2. The method according to claim 1, characterized in that, The source code vulnerability prediction model includes a positional encoding layer, a lightweight Transformer block, a convolutional layer, a global average pooling layer, and a fully connected layer integrated into a Keras model; Inputting the source code feature vectors in one-hot encoding format into the source code vulnerability prediction model and outputting a vulnerability category prediction result includes: Mapping and converting the source code feature vectors in one-hot encoding format into source code embedding vectors; Generating a corresponding position index for each source code embedding vector; The source code embedding vector is input into the positional encoding layer integrated into the Keras model. The positional encoding layer integrated into the Keras model calculates the positional encoding vector according to the total dimension and position index of the source code embedding vector, and adds the positional encoding vector to the source code embedding vector to obtain a source code embedding vector containing positional information; The source code embedding vector containing positional information is passed to the lightweight Transformer block to output a source code feature depth-enhanced fusion vector; The source code feature depth-enhanced fusion vector is processed by a convolutional layer and a global average pooling layer and then passed to multiple fully connected layers. The outputs of the multiple fully connected layers are respectively converted into prediction probabilities for each vulnerability category through the softmax function.
3. The method according to claim 2, wherein The source code embedding vector is input into the positional encoding layer integrated into the Keras model. The positional encoding layer integrated into the Keras model calculates the positional encoding vector according to the total dimension and position index of the source code embedding vector, and adds the positional encoding vector to the source code embedding vector to obtain a source code embedding vector containing positional information, including: The source code embedding vector is input into the positional encoding layer integrated into the Keras model. The positional encoding layer integrated into the Keras model calculates the sine and cosine values of each position in each dimension according to the total dimension and position index of the source code embedding vector; For each position, the calculated sine and cosine values are interleaved and combined into a positional encoding vector with the same dimension as the total dimension of the source code embedding vector; The positional encoding vector is added to the source code embedding vector to obtain a source code embedding vector containing positional information.
4. The method according to claim 2, wherein The source code embedding vector containing positional information is passed to the lightweight Transformer block to output a source code feature depth-enhanced fusion vector, including: The source code embedding vector containing positional information is input into the multi-head attention layer to output a multi-head attention fusion vector; The multi-head attention fusion vector is subjected to residual connection with the source code embedding vector containing positional information and layer normalization to output a self-attention position-enhanced vector; The self-attention position-enhanced vector is passed into a feed-forward network composed of a fully connected layer, a dropout layer, and a relu activation function to output a feed-forward network enhanced vector; The feed-forward network enhanced vector is subjected to residual connection and layer normalization to output a source code feature depth-enhanced fusion vector.
5. A source code vulnerability prediction device, characterized in that, The device includes: A module for obtaining data to be predicted, which is used to obtain a dataset of source code to be predicted; A first tokenization processing module, which is used to perform text tokenization processing on the dataset of source code to be predicted to obtain multiple source code token units; A first word-frequency sorted vocabulary building module, which is used to build a word-frequency sorted vocabulary according to the source code text in the source code token units; A first one-hot encoding module, which is used to convert the source code token units into source code feature vectors in one-hot encoding format according to the word-frequency sorted vocabulary; A vulnerability prediction module, which is used to input the source code feature vector in one-hot encoding format into a source code vulnerability prediction model and output a vulnerability category prediction result; The device further includes: The source code public dataset acquisition module is used to acquire the source code public dataset; the source code public dataset includes data with CWE vulnerability labels and data without CWE vulnerability labels; The undersampling module is used to identify the data without CWE vulnerability labels in the source code public dataset, randomly extract a preset number of data from the data without CWE vulnerability labels, and recombine the extracted preset number of data with the data with CWE vulnerability labels to generate an undersampled source code subset; The dataset splitting module is used to split the undersampled source code subset into a training dataset, a validation dataset, and a test dataset according to a preset ratio; The second tokenization processing module is used to perform text tokenization processing on the training dataset, the validation dataset, and the test dataset to obtain corresponding multiple source code token units; The second word frequency sorted vocabulary building module is used to build a corresponding word frequency sorted vocabulary according to the source code text in the source code token units; The equal-length matrix construction module is used to convert the training dataset, the validation dataset, and the test dataset into equal-length input matrices according to the source code token units; The second one-hot encoding module is used to convert the input matrices of the training dataset, the validation dataset, and the test dataset into source code feature vectors in one-hot encoding format according to the word frequency sorted vocabulary; The model network construction module is used to construct a source code vulnerability prediction model network; The model training module is used to train the source code vulnerability prediction model network with the training dataset represented as source code feature vectors in one-hot encoding format, and perform validation and testing with the validation dataset and the test dataset represented as source code feature vectors in one-hot encoding format to obtain the source code vulnerability prediction model; The model training module includes: The model network parameter optimization sub-module is used to train the source code vulnerability prediction model network with the training dataset represented as source code feature vectors in one-hot encoding format, calculate the total loss using a loss function based on dynamic weights, and update the source code vulnerability prediction model network parameters based on the total loss; The source code vulnerability prediction model generation sub-module is used to record a state change during the training iteration when the change trend of the loss function on the validation dataset changes from decreasing in the previous iteration compared to the next iteration to increasing. If the recorded number exceeds a preset early stopping threshold, stop the training and save the model parameters when the loss function showed a decreasing trend for the last time during the current training process to obtain the source code vulnerability prediction model; Loss function L based on dynamic weights t is defined as follows: Among them, L i is the loss function for each task, and the cross-entropy loss of each task at the end of each iteration is M i , N is the number of tasks, S is a parameter used to control the smoothness of the weights, and S ∈ (0, 5].
6. The device according to claim 5, characterized in that, The source code vulnerability prediction model includes a positional encoding layer, a lightweight Transformer block, a convolutional layer, a global average pooling layer, and a fully connected layer integrated into the Keras model; the vulnerability prediction module includes: The embedding vector generation sub-module is used to map and convert the source code feature vectors in one-hot encoding format into source code embedding vectors; The position index generation sub-module is used to generate a corresponding position index for each source code embedding vector; A position encoding sub-module for inputting the source code embedding vector into the position encoding layer integrated into the Keras model. The position encoding layer integrated into the Keras model calculates the position encoding vector according to the total dimension of the source code embedding vector and the position index, and adds the position encoding vector to the source code embedding vector to obtain a source code embedding vector containing position information; A self-attention fusion sub-module for passing the source code embedding vector containing position information to the lightweight Transformer block and outputting a source code feature deeply enhanced fusion vector; A vulnerability category probability prediction sub-module for passing the source code feature deeply enhanced fusion vector through a convolutional layer and a global average pooling layer and then to multiple fully connected layers, and converting the outputs of the multiple fully connected layers into the prediction probabilities of each vulnerability category through the softmax function respectively.
7. An electronic device, characterized in that, It includes a processor, a communication interface, a memory, and a communication bus. Among them, the processor, the communication interface, and the memory complete communication with each other through the communication bus; The memory is used to store computer programs; When the processor is used to execute the program stored on the memory, it implements the source code vulnerability prediction method as described in any one of claims 1-4.
8. One or more computer-readable media, on which instructions are stored. When executed by one or more processors, the instructions cause the processor to execute the source code vulnerability prediction method as described in any one of claims 1-4.