A Software Vulnerability Type Identification Method Based on Long-Tail Incremental Learning
By using a long-tail incremental learning approach, combining source code and vulnerability descriptions, and employing a hybrid hint template and mapping mechanism, the problem of difficulty in integrating source code and vulnerability descriptions in existing technologies is solved, achieving more efficient and accurate identification of software vulnerability types and alleviating the problems of large data volume and forgetting.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-11-27
- Publication Date
- 2026-03-10
AI Technical Summary
Existing technologies struggle to effectively integrate source code and vulnerability descriptions, leading to challenges for software vulnerability type identification models in learning specific knowledge from samples. Furthermore, these models face issues such as high data requirements and the risk of catastrophic forgetting.
We employ a long-tail incremental learning approach, designing a hybrid prompt template and mapping mechanism. We use source code and vulnerability descriptions as bimodal inputs, train the code using the pre-trained model CodeT5, and use Focal Loss, Label Smooth CE Loss, and EWC regularization terms to mitigate class imbalance and catastrophic forgetting.
It achieves more accurate and efficient identification of software vulnerability types, significantly improves adaptability and robustness to long-tail distributions, and reduces the model's forgetting rate in real-world projects.
Smart Images

Figure CN119646822B_ABST
Abstract
Description
TECHNICAL FIELD
[0001] The application relates to the technical field of software engineering, and in particular to a software vulnerability type identification method based on long-tail incremental learning. BACKGROUND
[0002] Software vulnerability type identification is of great significance in ensuring the security and reliability of computer systems, and can effectively reduce potential economic losses caused by vulnerabilities. With the continuous emergence of new technologies, the complexity of the Internet has gradually increased, and the frequency of software vulnerabilities has significantly increased. However, due to the complexity of software and the limited nature of repair resources, it is difficult to complete the repair of all vulnerabilities in the short term. Therefore, it is particularly important to carry out accurate identification of software vulnerability types. By quickly locating high-risk vulnerabilities, it can provide a scientific basis for vulnerability repair and resource allocation, and significantly improve the security protection capability of the system.
[0003] Software vulnerability type identification is a key task aimed at automatically classifying vulnerability types (such as using CWE-ID). Its core goal is to help developers and security analysts quickly identify high-risk vulnerabilities and develop targeted mitigation strategies. The importance of software vulnerability type identification (SVTI) lies in its ability to significantly enhance software security while improving classification efficiency. Traditional manual classification methods are not only time-consuming, but also highly dependent on expert knowledge, making it difficult to meet the growing demand. For example, as of 2023, among the 28,902 newly added vulnerabilities in the NVD database, 4,113 have not been classified, highlighting the urgent need for automated methods. The following sections will discuss the development of code-based and code description-based vulnerability type identification methods in detail.
[0004] Vulnerabilities are usually classified by Common Weakness Enumeration (CWE) identifiers, which help to unify communication and understanding of security issues. A typical case is the Log4j vulnerability, which exploits a weakness in a popular logging library, leading to widespread attacks on various applications. By automating vulnerability type identification, security practitioners can prioritize the repair of high-risk vulnerabilities, thereby optimizing response strategies and comprehensively improving software security. This proactive approach is crucial for dealing with the increasingly complex and frequent problems of vulnerabilities in modern software systems.
[0005] To solve this problem, researchers have built a vulnerability type identification model based on available information to predict vulnerability types. Compared with the fine-tuning method based on pre-trained language models, the performance of traditional models has been insufficient. However, directly using pre-trained language models to process downstream tasks still faces two major problems. On the one hand, pre-trained language models have certain challenges in learning sample-specific knowledge; on the other hand, the fine-tuning method usually requires a large amount of data to achieve good performance. The emergence of prompt learning alleviates these problems. However, existing researches mostly focus on predicting based on a single information source (such as source code or vulnerability description), which cannot fully exploit vulnerability-related information. Therefore, there is an urgent need to develop an advanced method that integrates source code and vulnerability description to better solve the software vulnerability type identification problem.
[0006] How to solve the above technical problems becomes a problem faced by the present application. SUMMARY
[0007] The purpose of the present application is to provide a software vulnerability type identification method based on long-tail incremental learning, which can predict the type of software vulnerability.
[0008] The idea of the present application is that the present application proposes a software vulnerability type identification method based on long-tail incremental learning, the core idea of which is to combine source code and vulnerability description to construct a dual-modal input through a designed mixed prompt template; in this method, a mapping mechanism is designed to associate categories with mapping words, so that the model can input dual-modal information into the pre-trained model CodeT5 for training; through this mapping mechanism, the natural language output of the model can be accurately mapped to specific categories; at the same time, in view of the dynamic nature and long-tail distribution of software vulnerability type data in real-world projects, incremental learning and long-tail learning strategies are introduced; by comparing and analyzing the performance changes and combining the early stopping strategy, a software vulnerability type identification model with optimal performance is constructed, thereby realizing more accurate and efficient software vulnerability type identification.
[0009] The present application is realized by the following measures: a software vulnerability type identification method based on long-tail incremental learning, comprising the following steps:
[0010] (1) By analyzing the CVE vulnerability database and the actual development projects based on the Git software project hosting platform, collecting vulnerability-related information including CWE ID, vulnerability code, vulnerability description and submission time. The data processing flow first extracts the CWE ID, vulnerability code and vulnerability description from the CVE data, then retrieves the code changes associated with the vulnerability in the open source Git project and obtains the source code before modification. Then, pre-processing operations such as deleting comments, blank lines and line beginning comments are performed on the source code, and finally a vulnerability type prediction dataset D is constructed;
[0011] (2) The vulnerability type prediction dataset D is task-divided according to the submission time of the vulnerability. Then, the dataset of each task is divided into a training set, a validation set, and a test set in a ratio of 8:1:1. During the division process, a stratified sampling strategy is adopted to ensure that the distribution of vulnerability types in different datasets remains consistent;
[0012] (3) For the current task that needs to identify the type of vulnerability, a hybrid replay strategy is used to select typical vulnerability samples from the dataset of the old task. Before each model training, the tail data in the old task (i.e., vulnerability samples of rare types) is preferentially selected, and then vulnerability samples with high uncertainty are selected according to the Mahalanobis distance of the samples. After merging the vulnerability samples selected by the hybrid replay with the vulnerability samples of the current task, they are re-divided into a training set, a validation set, and a test set in a ratio of 8:1:1. During the division process, a stratified sampling method is used to ensure the consistency of the distribution of vulnerability types in the two types of data sets (i.e., the head sample set and the tail sample set), which includes the following steps:
[0013] (3-1) In order to identify the tail category, we set a threshold τ, for example, less than 5% of the total dataset, and classify it as a tail category. Specifically, the tail category can be defined as follows:
[0014]
[0015] where τ is set to 5% to define the tail data.
[0016] (3-2) Select samples with high uncertainty, which represent the most challenging instances that the model considers, usually located at the edge of the decision boundary. To quantify the uncertainty of the samples, we use the Mahalanobis distance between the feature vector of each sample and the mean of the feature vectors of all samples. Specifically, step 1, use the feature extraction layer of the model to obtain the feature vector x i of each sample. Step 2, calculate the mean μ and covariance matrix ∑ of the feature vectors of all samples. Step 3, calculate the Mahalanobis distance d M (x i ), the formula for each sample is as follows:
[0017]
[0018] where x i is the feature vector of sample i, μ is the mean of the feature vectors, ∑ -1 is the inverse of the covariance matrix.
[0019] (3-3) Combine the tail data samples and the samples with high uncertainty in a balanced proportion (i.e., mixed at 1:1). This mixed replay method can ensure that the model not only retains the knowledge of the tail data, but also improves the accuracy of the prediction by learning from the uncertain samples. This method improves the adaptability of the model to the long-tail distribution and enhances its robustness to catastrophic forgetting.
[0020] (4) For the data set processed in step (3), extract the source code and vulnerability description, and design a prompt template. In the prompt template, the source code and the vulnerability description are distinguished by hard prompts; in the type identification, a [SOFT] tag is used to create a soft prompt, and a [MASK] tag is used to predict or fill in the missing information. Then a mapping method is designed to map the natural language output generated by the model to specific categories. Finally, the constructed dual-modal information is input into the pre-trained language model CodeT5, and the model and the prompt template are trained through prompt fine-tuning, which includes the following steps:
[0021] (4-1) Construct a prompt template. Specifically, use hard prompts to distinguish source code and vulnerability description, i.e., add "The code snippet:" and "The vulnerability description:". For type classification, use a [SOFT] tag to create a soft prompt, i.e., initialize it to "Identify the vulnerability type:" during training, and use a [MASK] tag to predict or fill in the missing information;
[0022] (4-2) Construct a mapping method: map the natural language output of the model to specific categories. According to each CWE-ID, map it to its corresponding software vulnerability type, and construct the approximate words of this vulnerability type, finally get the mapping method used in the method of the application. For example, the vulnerability type corresponding to CWE-125 is "Out-of-bounds Read", and its corresponding two approximate words are "Buffer Overread", "Invalid Memory Access", and finally the mapping of CWE-125 is constructed as CWE-125: ["Out-of-bounds Read", "Buffer Overread", "Invalid Memory Access"];
[0023] (4-3) Combine the dual-modal information of the vulnerability sample to be identified (i.e., the vulnerability code and the vulnerability description) with the pre-defined mixed prompt template P, input the pre-trained model CodeT5, and generate natural language output results through prompt fine-tuning of the model and the prompt template training;
[0024] (4-4) Convert the natural language results generated by the model into corresponding software vulnerability types and return the final prediction results.
[0025] (5) During model training, Focal Loss and Label Smooth CE Loss are used to focus on tail data to alleviate the class imbalance problem. For new task learning, the EWC regularization term is used to limit the model parameter weight to be close to the optimal parameter of the previous task, to ensure that the training for the new task does not significantly change the parameters critical to the old task, thereby effectively alleviating the catastrophic forgetting phenomenon.2. The software vulnerability type identification method based on long-tail incremental learning according to claim 1, wherein the step (3) comprises the following steps:
[0026] (5-1) Use the traditional stochastic gradient descent method to train the neural network model, train the initial task model, and during the training process, the data will be randomly shuffled and processed in small batches. After completing the training of the initial task, the trained model parameters are saved;
[0027] (5-2) Use Focal Loss and Label Smooth CE Loss in combination to ensure that the model can focus on tail data during training. Focal Loss introduces a scaling factor in the standard CE Loss, allowing the model to focus more on difficult-to-classify samples, which are often tail samples. It achieves this by applying a focusing parameter γ (to increase the weight of misclassified samples) and a balancing factor α (to address the class imbalance problem). The final loss function is defined as:
[0028]
[0029] Label Smooth CE Loss softens the target label by reducing overconfidence in predictions to improve generalization, especially for non-representative classes. Its calculation method is defined as follows:
[0030]
[0031] where y k =(1-∈) is the true class, ∈ is the smoothing parameter, and K is the total number of classes. Next, use the weight w to fit the two loss functions, and the specific calculation method is defined as follows:
[0032]
[0033] (5-3) Restrict the model parameter weight by EWC regularization term, so that it is close to the optimal parameter of the previous task. Specifically, the Fisher information matrix is calculated to quantify the importance of each parameter to the initial task. Assuming F t represents the Fisher matrix of the current task t, represents the cumulative Fisher matrix. After each task is completed, will be updated as follows:
[0034]
[0035] where β is the decay factor, which attempts to assign higher weights to recent tasks while preserving the knowledge of early tasks. Using the cumulative Fisher matrix, the regularization term of EWC can be defined as:
[0036]
[0037] where λ is the penalty coefficient, is the i-th element of the cumulative Fisher matrix, θ i is the current parameter, is the parameter value learned from the previous task. When training a new task, we use the loss function with the regularization term to ensure that the model parameters do not deviate too much from the parameters of the initial task.(5-4) Calculate the F1 index change of the model on the validation set, when the F1 value does not exceed the current best performance for 10 consecutive rounds, stop training, return the best vulnerability type recognition model and mixed prompt template.
[0038] Compared with the prior art, the beneficial effects of the present application are: the software vulnerability type recognition method based on long-tail incremental learning proposed by the present application takes the dual modal information fusion of vulnerability source code and vulnerability description as input, uses an advanced pre-trained language model and performs prompt fine-tuning, and considers the use scenario of incremental learning, better mines vulnerability information, can more accurately and reliably classify vulnerability types, and to some extent, alleviates the problem of catastrophic forgetting caused by the increase of data over time under real projects. BRIEF DESCRIPTION OF DRAWINGS
[0039] The accompanying drawings are included to provide a further understanding of the present application, and constitute a part of the specification, together with the embodiments of the present application, to explain the present application, and do not constitute a limitation of the present application.
[0040] Figure 1 A system framework diagram of the software vulnerability type recognition method based on long-tail incremental learning provided by the present application. DETAILED DESCRIPTION
[0041] In order to make the purpose, technical scheme and advantages of the present application clearer, the present application will be further described in detail below in combination with the drawings and examples. Of course, the specific examples described herein are only used to explain the present application and not used to limit the present application.
[0042] Example 1
[0043] The technical scheme of the present embodiment is a software vulnerability type identification method based on long-tail incremental learning. As shown in Figure 1 , the present embodiment specifically includes the following contents:
[0044] (1) By analyzing the CVE vulnerability database and the actual development projects based on the Git software project hosting platform, collect vulnerability related information, including CVE ID, CWE ID, vulnerability code, vulnerability description and submission time. The data processing process first extracts the CWE ID, code and description of the vulnerability from the CVE data, then retrieves the code changes associated with the vulnerability in the open source Git project, and obtains the source code before modification. Then, the source code is preprocessed, such as deleting comments, blank lines and line beginning comments, and finally the vulnerability dataset D is constructed, including the following steps:
[0045] (1) By analyzing the CVE vulnerability database and the actual development projects based on the Git software project hosting platform, collect vulnerability related information, including CWE ID, vulnerability code, vulnerability description and submission time. The data processing process first extracts the CWE ID, code and description of the vulnerability from the CVE data, then retrieves the code changes associated with the vulnerability in the open source Git project, and obtains the source code before modification. Then, the source code is preprocessed, such as deleting comments, blank lines and line beginning comments, and finally the vulnerability type prediction dataset D is constructed;
[0046] (2) Divide the vulnerability type prediction dataset D according to the submission time of the vulnerability. Then, divide the dataset of each task into training set, validation set and test set according to the ratio of 8:1:1. In the division process, stratified sampling strategy is adopted to ensure that the distribution of vulnerability types in different datasets remains consistent;
[0047] (3) For the current task that needs to identify the type of vulnerability, a hybrid replay strategy is used to select typical vulnerability samples from the dataset of the old task. Before each model training, the tail data in the old task (i.e. vulnerability samples of rare types) is preferentially selected, and then vulnerability samples with high uncertainty are selected according to the Mahalanobis distance of the samples. After the vulnerability samples selected by the hybrid replay are combined with the vulnerability samples of the current task, they are re-divided into training set, validation set and test set according to the ratio of 8:1:1. In the division process, the stratified sampling method is used to ensure the consistency of the vulnerability type distribution in the two data sets (i.e. the head sample set and the tail sample set), which includes the following steps:
[0048] (3-1) In order to identify the tail category, we set a threshold τ, for example, less than 5% of the total dataset, and classify it as a tail category. Specifically, the tail category can be defined as follows:
[0049]
[0050] Where τ is set to 5% to define the tail data.
[0051] (3-2) Select samples with high uncertainty, which represent the most challenging instances that the model considers, usually on the edge of the decision boundary. In order to quantify the uncertainty of the sample, we use the Mahalanobis distance between the feature vector of each sample and the mean of the feature vectors of all samples. Specifically, step 1, use the feature extraction layer of the model to obtain the feature vector x i of each sample. Step 2, calculate the mean μ and covariance matrix ∑ of the feature vectors of all samples. Step 3, calculate the Mahalanobis distance d M (x i ), the calculation formula for each sample is as follows:
[0052]
[0053] Where x i is the feature vector of sample i, μ is the mean of the feature vector, ∑ -1 is the inverse of the covariance matrix.
[0054] (3-3) Combine the tail data samples and the samples with high uncertainty in a balanced ratio (i.e. 1:1 mix). This hybrid replay method can ensure that the model not only retains the knowledge of the tail data, but also improves the accuracy of prediction by learning from uncertain samples. This method improves the adaptability of the model to long-tailed distribution and enhances its robustness to catastrophic forgetting.
[0055] (4) For the data set processed in step (3), extract the source code and vulnerability description, and design a prompt template. In the prompt template, distinguish the source code and vulnerability description by hard prompts; in type identification, create soft prompts using the [SOFT] tag, and use the [MASK] tag to predict or fill in the missing information. Then design a mapping method to map the natural language output generated by the model to specific categories. Finally, input the constructed dual-modal information into the pre-trained language model CodeT5, and train the model and the prompt template through prompt fine-tuning, which includes the following steps:
[0056] (4-1) Construct a prompt template. Specifically, use hard prompts to distinguish source code and vulnerability description, that is, add “The code snippet:” and “The vulnerability description:” respectively. For type classification, create soft prompts using the [SOFT] tag, that is, initialize it to “Identify the vulnerability type:” during training, and use the [MASK] tag to predict or fill in the missing information;
[0057] (4-2) Construct a mapping method: map the natural language output of the model to specific categories. According to each CWE-ID, map it to its corresponding software vulnerability type, and construct the approximate words of this vulnerability type, and finally obtain the mapping method used in this embodiment. For example, the vulnerability type corresponding to CWE-125 is “Out-of-bounds Read”, and its corresponding two approximate words are “Buffer Overread” and “Invalid Memory Access”. Finally, the mapping of CWE-125 is constructed as CWE-125: [“Out-of-bounds Read”, “Buffer Overread”, “Invalid Memory Access”];
[0058] (4-3) Combine the dual-modal information of the vulnerability sample to be identified (i.e., the vulnerability code and the vulnerability description) with the pre-defined mixed prompt template P, input it into the pre-trained model CodeT5, and generate natural language output results through prompt fine-tuning of the model and the prompt template;
[0059] (4-4) Convert the natural language results generated by the model into corresponding software vulnerability types, and return the final prediction results.
[0060] (5) In the model training process, first, Focal Loss and Label Smooth CE Loss are used to focus on the tail data to alleviate the class imbalance problem. For the learning of new tasks, the EWC regularization term is used to limit the model parameter weight to be close to the optimal parameter of the previous task, so as to ensure that the training for the new task will not significantly change the parameters critical to the old task, thereby effectively alleviating the catastrophic forgetting phenomenon.2. The software vulnerability type identification method based on long-tail incremental learning according to claim 1, characterized in that the step (3) comprises the following steps:
[0061] (5-1) Use the traditional stochastic gradient descent method to train the neural network model, train the initial task model, and in the training process, the data will be randomly shuffled and processed in small batches. After completing the training of the initial task, the trained model parameters are saved;
[0062] (5-2) Use Focal Loss and Label Smooth CE Loss in combination to ensure that the model can focus on the tail data during training. Focal Loss introduces a scaling factor in the standard CE Loss, allowing the model to focus more on difficult-to-classify samples, which are often tail samples. It achieves this by applying a focusing parameter γ (to increase the weight of misclassified samples) and a balancing factor α (to solve the class imbalance problem). The final loss function is defined as:
[0063]
[0064] Label Smooth CE Loss softens the target label, reducing overconfidence in predictions to improve generalization, especially for non-representative classes. Its calculation method is defined as follows:
[0065]
[0066] where y k =(1-∈) is the true class, ∈ is the smoothing parameter, and K is the total number of classes. Next, use the weight w to fit the two loss functions, and the specific calculation method is defined as follows:
[0067]
[0068] (5-3) Use the EWC regularization term to limit the model parameter weight to be close to the optimal parameter of the previous task. Specifically, calculate the Fisher information matrix to quantify the importance of each parameter to the initial task. Assume that F t represents the Fisher matrix of the current task t, denotes the accumulated Fisher matrix. After each task is completed, is updated as follows:
[0069]
[0070] where β is a decay factor that attempts to assign higher weights to recent tasks while preserving the knowledge of early tasks. Using the accumulated Fisher matrix, the regularization term of EWC can be defined as:
[0071]
[0072] where λ is a penalty coefficient, is the i-th element of the accumulated Fisher matrix, θ i is the current parameter, is the parameter value learned from previous tasks. When training a new task, we use a loss function with a regularization term to ensure that the model parameters do not deviate too much from the initial task parameters.
[0073] (5-4) Calculate the F1 index change of the model on the validation set, when the F1 value does not exceed the current best performance for 10 consecutive rounds, stop training, return the best vulnerability type recognition model and mixed prompt template.
[0074] (6) Compare the method of the embodiment with the existing software vulnerability type recognition method on the same data set, use five performance indicators from the field of software vulnerability task research to automatically evaluate the quality of the model:
[0075] Table 1 Comparison of results of the method of the embodiment and the remaining methods
[0076]
[0077] Experiments show that the software vulnerability type recognition method based on long-tail incremental learning proposed in the embodiment is superior to the baseline method in all indicators, and has more reliable and accurate software vulnerability type recognition ability. Specifically, compared with the baseline method, the accuracy of the embodiment is improved from 125.81% to 438.46%, the precision is improved from 130.77% to 757.14%, the recall is improved from 144% to 1120%, the F1 score is improved from 166.67% to 522.22%, and the MCC is relatively improved by 168% to 737.50%. The method of the embodiment combines the incremental learning strategy based on mixed replay and elastic weight consolidation (EWC), and fuses source code and vulnerability description information, which not only can more comprehensively mine vulnerability features, but also effectively reduces the forgetting rate of old task data. Compared with the baseline method, the method significantly improves the recognition ability of long-tail category vulnerabilities, and shows strong advancement and practical value.
[0078] Example 2
[0079] On the basis of Example 1, different incremental learning methods are used to compare the method of the present example and the method using different incremental learning strategies, and five performance indicators from the field of software vulnerability task research are used to automatically evaluate the quality of the model:
[0080] Table 2 Comparison table of results of the method of the present example and the remaining methods
[0081]
[0082] Experiments show that the software vulnerability type recognition method based on long-tail incremental learning proposed in the present example has an accuracy rate of 2.94% to 9.37% higher than that of using other incremental learning methods, an accuracy rate of 3.45% to 13.21% higher, an F1 score of 1.82% to 7.69% higher, and an MCC of 1.52% to 8.06% higher. The results show that the method of the present example has higher performance and stronger robustness when dealing with incremental learning tasks, and can perform more reliable and accurate software vulnerability type recognition.
[0083] Example 3
[0084] On the basis of Example 1, the pre-trained model used is used to compare the method of the present example and the method using other pre-trained models, and five performance indicators from the field of software vulnerability task research are used to automatically evaluate the quality of the model:
[0085] Table 3 Comparison table of results of the method of the present example and the remaining methods
[0086]
[0087]
[0088] Experiments show that the software vulnerability type recognition method based on long-tail incremental learning proposed in the present example has all the highest indicators compared to the method using other pre-trained models, and thus can perform more reliable and accurate software vulnerability type recognition. Specifically, compared with UnixCoder which has the second best performance, the present example has an accuracy rate of 4.48% higher, an accuracy rate of 11.11% higher, an F1 score of 1.82% higher, and an MCC of 3.08% higher. These results show the competitiveness of the method of the present example.
[0089] Example 4
[0090] Based on Example 1, the type of prompt in this example is changed. The method of this example is compared with methods using different prompts (i.e., soft prompts and hard prompts). Five performance metrics from the field of software vulnerability task research are used to automatically evaluate the quality of the model:
[0091] Table 4 Comparison of Results between the Method in This Embodiment and Other Methods
[0092]
[0093] Experiments show that the software vulnerability type identification method based on long-tail incremental learning proposed in this embodiment, employing hybrid hints, achieves the highest scores across all metrics compared to methods using other types of hints, thus enabling more reliable and accurate software vulnerability type identification. Compared to hard hints, this embodiment improves accuracy by 1.45%, precision by 3.45%, recall by 6.67%, F1 score by 3.70%, and MCC by 3.08%. Simultaneously, compared to soft hints, this embodiment improves accuracy by 1.45%, precision by 7.14%, recall by 8.47%, F1 score by 5.66%, and MCC by 1.52%. These results highlight the effectiveness of this embodiment in combining the advantages of hard and soft hints, thereby achieving superior performance in software vulnerability type identification tasks.
[0094] Example 5
[0095] Based on Example 1, the number of mappings in this example is changed, and the method in this example is compared with methods using different numbers of mappings. Five performance metrics from the field of software vulnerability task research are used to automatically evaluate the quality of the model:
[0096] Table 5. Comparison of results between the method in this embodiment and other methods.
[0097]
[0098]
[0099] Table 6 Mapping method of this embodiment
[0100]
[0101] Experiments show that the software vulnerability type identification method based on long-tail incremental learning proposed in this embodiment achieves the highest scores across all metrics compared to methods using other mapping quantities, thus enabling more reliable and accurate software vulnerability type identification. Specifically, for the F1 score, the method in this embodiment improves by up to 3% compared to configurations using other mapping quantities; for the MCC metric, the method in this embodiment improves by up to 5%. These results fully validate the competitiveness of the proposed method, demonstrating its significant advantages in optimizing performance and controlling complexity, thus possessing important practical application value in software vulnerability type identification tasks.
[0102] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the protection scope of the present invention.
Claims
1. A method for software vulnerability type identification based on long-tail incremental learning, characterized in that, The method comprises the following steps: (1) Collecting vulnerability-related information including CWE ID, vulnerability code, vulnerability description and submission time by analyzing CVE vulnerability database and actual development projects based on Git-based software project hosting platform, the data processing process firstly extracts the CWE ID, vulnerability code and vulnerability description of the vulnerability from the CVE data, then retrieves the code changes associated with the vulnerability in the open source Git project and obtains the source code before modification, then, the source code is preprocessed to delete comments, blank lines and line comments, and finally a vulnerability type prediction dataset D is constructed; (2) Dividing the vulnerability type prediction dataset D according to the submission time of the vulnerability, then dividing the dataset of each task into training set, validation set and test set according to the ratio of 8:1:1, and adopting stratified sampling strategy in the division process; (3) For the current task requiring vulnerability type identification, a hybrid replay strategy is adopted to select typical vulnerability samples from the dataset of the old task, and before each model training, the tail data in the old task is preferentially selected, then according to the Mahalanobis distance of the samples, the vulnerability samples with high uncertainty are selected, the vulnerability samples selected by the hybrid replay are combined with the vulnerability samples of the current task, and then the training set, the validation set and the test set are divided according to the ratio of 8:1:1, and in the division process, stratified sampling method is used to ensure the consistency of the vulnerability type distribution in the two types of data sets; (4) For the dataset processed in step (3), extract the source code and vulnerability description, and design a prompt template, in which the source code and vulnerability description are distinguished by hard prompt; In type identification, [SOFT] label is used to create soft prompt, and [MASK] label is used to predict or fill in missing information, then a mapping method is designed to map the natural language output generated by the model to specific categories, finally, the built dual-modal information is input into the pre-trained language model CodeT5, and the model and prompt template are trained through prompt fine-tuning; (5) In the model training process, firstly, Focal Loss and Label Smooth CE Loss are used to focus on tail data to alleviate the problem of class imbalance, and for the learning of new tasks, EWC regularization term is used to limit the model parameter weight to be close to the optimal parameter of the previous task; Focal Loss and Label Smooth CE Loss are used in combination to ensure that the model focuses on tail data during training, Focal Loss introduces a scaling factor in the standard CE Loss to make the model focus on difficult-to-classify samples, which are often tail samples, and it achieves this by applying a focusing parameter γ and a balancing factor α, and finally the loss function is defined as: Label Smooth CE Loss softens the target label to improve the generalization ability by reducing the excessive confidence in prediction, especially for the representative classes, and its calculation method is defined as: where y k = (1 - e) is the true class, e is the smoothing parameter, K is the total number of classes, and the weights w are used to fit the two loss functions, which are calculated as follows:
2. The software vulnerability type identification method based on long tail incremental learning according to claim 1, characterized in that, In step (3), the following steps are included: (2-1) In order to identify the tail category, a threshold τ is set, which accounts for less than 5% of the total dataset, and is classified as a tail category, and the tail category is defined as follows: Where τ is set to 5% to define the tail data; (2-2) Select samples with high uncertainty, in order to quantify the uncertainty of the sample, use the Mahalanobis distance between the feature vector of each sample and the mean of the feature vector of all samples, including the following steps: Step 1, use the feature extraction layer of the model to obtain the feature vector x of each sample i ; Step 2, calculate the mean μ and covariance matrix ∑ of the feature vectors of all samples; Step 3, calculate Mahalanobis distance d M (x i ), the calculation formula for each sample is as follows: where x i is the feature vector of sample i, μ is the mean of the feature vectors, ∑ -1 is the inverse of the covariance matrix; (2-3) Mix and combine tail data samples and samples with high uncertainty in a balanced ratio of 1:
1.
3. The software vulnerability type identification method based on long tail incremental learning according to claim 1, characterized in that, The step (4) comprises the following steps: (3-1) Construct a prompt template, use hard prompts to distinguish source code and vulnerability description, add "The codesnippet:" and "The vulnerability description:" respectively, for type classification, use [SOFT] tags to create soft prompts, that is, initialize it to "Identify the vulnerability type:" during training, and use [MASK] tags to predict or fill in missing information; (3-2) Build mapping method: map the natural language output of the model to the specific category, according to each CWE-ID, map it to the corresponding software vulnerability type, and build the approximate words of this vulnerability type, and finally get the mapping method used by this method; (3-3) Combine the dual-modal information of the vulnerability sample to be identified, that is, the vulnerability code and the vulnerability description, with the pre-defined mixed prompt template P, input the pre-trained model CodeT5, and train the model and the prompt template through prompt fine-tuning to generate the output result in natural language form; (3-4) Convert the natural language result generated by the model into the corresponding software vulnerability type, and return the final prediction result.
4. The software vulnerability type identification method based on long tail incremental learning according to claim 1, characterized in that, The step (5) comprises the following steps: (4-1) Train the neural network model using the traditional stochastic gradient descent method, train the initial task model, and during the training process, the data will be randomly shuffled and processed in small batches, after completing the training of the initial task, save the trained model parameters; (4-2) Restrict the model parameter weight by EWC regularization term, make it close to the optimal parameter of the previous task, calculate the Fisher information matrix to quantify the importance of each parameter to the initial task, assume F t Fisher matrix of the current task t, Cumulative Fisher matrix, after each task is completed, Update as follows: Where β is the decay factor, which attempts to assign higher weights to recent tasks while preserving the knowledge of early tasks, using the cumulative Fisher matrix, the regularization term of EWC is defined as: where λ is a penalty coefficient, is the i-th element of the cumulative Fisher matrix, θ i is the current parameter, is the parameter value learned from the previous task, and a regularization term is used in the loss function when training the new task to ensure that the model parameters do not deviate too much from the parameters of the initial task. (4-3) Calculate the F1 index change of the model on the validation set, when the F1 value does not exceed the current best performance for 10 consecutive rounds, stop training, return the best vulnerability type identification model and mixed prompt template.
Citation Information
Patent Citations
Multi-type vulnerability detection method and system based on multi-sample comparison and fusion
CN114756861A
Software vulnerability assessment method based on continuous learning
CN119026138A