Software vulnerability type identification method based on long-tail incremental learning
By combining source code and vulnerability description, a hybrid prompt template is designed and trained using the pre-trained language model CodeT5, combined with incremental learning and long-tail learning strategies, the problem of difficult to automatically identify software vulnerability types in the existing technology is solved, and more accurate and efficient vulnerability type recognition is achieved.
Patent Information
- Application Number
- CN202411714608.0
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-11-27
- Publication Date
- 2025-03-18
- Estimated Expiration
- 2044-11-27
AI Technical Summary
The prior art is difficult to effectively automate the identification of software vulnerability types, especially when dealing with complex modern software systems. Traditional models have challenges in learning sample-specific knowledge and handling downstream tasks, and cannot fully tap vulnerability-related information.
A software vulnerability type recognition method based on long-tail incremental learning is proposed. By combining source code and vulnerability description, mixed prompt templates are designed, dual-modal input is constructed, and pre-trained language model CodeT5 is used for training, and incremental learning and long-tail learning strategies are introduced to build a software vulnerability type recognition model with the best performance.
It realizes more accurate and efficient software vulnerability type identification, can better mine vulnerability information, reduce catastrophic forgetting, and significantly improves the ability to identify long-tail vulnerabilities.
Smart Images

Figure CN119646822A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the technical field of software engineering, and in particular to a method for identifying software vulnerability types based on long-tail incremental learning. Background Art
[0002] Identification of software vulnerability types 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 increased significantly. However, due to the complexity of software and the limited resources for repair, it is difficult to complete the repair of all vulnerabilities in a short period of time. Therefore, it is particularly important to carry out accurate identification of software vulnerability types. By quickly locating high-risk vulnerabilities, a scientific basis can be provided for vulnerability repair and resource allocation, significantly improving the security protection capabilities of the system.
[0003] Software vulnerability type identification is a critical task that aims to automatically classify vulnerability types (e.g., 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) is reflected in its ability to significantly enhance the security of software while improving classification efficiency. Traditional manual classification methods are not only time-consuming, but also highly dependent on expert knowledge, and it is difficult to meet the growing demand. For example, as of 2023, 4,113 of the 28,902 new vulnerabilities in the NVD database remain unclassified, highlighting the urgent need for automated methods. The following sections will explore the development of code-based and code description-based vulnerability type identification methods in detail.
[0004] Vulnerabilities are often classified by Common Weakness Enumeration (CWE) identifiers, which helps unify communication and understanding of security issues. A typical case is the Log4j vulnerability, which exploited a weakness in a popular logging library and led to widespread attacks on various applications. By automating the identification of vulnerability types, security practitioners are able to prioritize the remediation of high-risk vulnerabilities, thereby optimizing response strategies and improving the security of the software overall. This proactive approach is critical to addressing the increasing complexity and frequency of vulnerabilities in modern software systems.
[0005] To address this problem, researchers have built a vulnerability type recognition model based on available information to predict vulnerability types. Compared with fine-tuning methods based on pre-trained language models, the performance of traditional models has been insufficient. However, there are still two major problems in directly using pre-trained language models to handle downstream tasks. On the one hand, pre-trained language models are challenging to learn sample-specific knowledge; on the other hand, fine-tuning methods usually require a lot of data to achieve good performance. The emergence of prompt learning has alleviated these problems. However, existing studies mostly focus on a single information source (such as source code or vulnerability description) for prediction, which cannot fully mine vulnerability-related information. Therefore, it is urgent to study an advanced method that integrates source code and vulnerability description to better solve the problem of software vulnerability type identification.
[0006] How to solve the above technical problems becomes the subject faced by the present invention. Summary of the invention
[0007] The purpose of the present invention is to provide a method for identifying software vulnerability types based on long-tail incremental learning, which can predict the types of software vulnerabilities.
[0008] The idea of the present invention is: the present invention proposes a software vulnerability type identification method based on long-tail incremental learning, the core idea of which is to combine source code with vulnerability description, and construct a bimodal input through a designed hybrid prompt template; in this method, a mapping mechanism is designed to associate categories with mapping words, so that the model can input bimodal information into a 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 and long-tail distribution problems of software vulnerability type data in real-world projects, incremental learning and long-tail learning strategies are introduced; through comparative analysis of performance changes and combined with an early stopping strategy, a software vulnerability type identification model with optimal performance is constructed, thereby achieving more accurate and efficient software vulnerability type identification.
[0009] The present invention is implemented by the following measures: a method for identifying software vulnerability types based on long-tail incremental learning, which includes the following steps:
[0010] (1) By analyzing the CVE vulnerability database and actual development projects based on the Git-based software project hosting platform, we collect vulnerability-related information, including CWE ID, vulnerability code, vulnerability description, and submission time. The data processing process first 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. Subsequently, the source code is preprocessed, such as deleting comments, blank lines, and comments at the beginning of the line, and finally constructs a vulnerability type prediction dataset D;
[0011] (2) The vulnerability type prediction dataset D is divided into tasks according to the vulnerability submission time. Subsequently, 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. A stratified sampling strategy is used in the division process to ensure that the distribution of vulnerability types in different datasets is consistent;
[0012] (3) For the current task that requires vulnerability type identification, a hybrid replay strategy is used to select typical vulnerability samples from the data set of the old task. Before each model training, the tail data in the old task (i.e., rare vulnerability samples) is preferentially selected, and then vulnerability samples with higher uncertainty are selected based on 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 training set, validation set, and 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 vulnerability type distribution in the two data sets (i.e., the head sample set and the tail sample set), which specifically includes the following steps:
[0013] (3-1) In order to identify the tail category, we set a threshold τ, such as the category that accounts for less than 5% of the total data set, and classify it as the tail category. Specifically, the tail category can be defined as follows:
[0014]
[0015] Here, τ is set to 5% to define it as tail data.
[0016] (3-2) Select samples with high uncertainty. These samples represent the most challenging instances that the model considers, and are 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, in step 1, use the feature extraction layer of the model to obtain the feature vector x of each sample. i Step 2: Calculate the eigenvector mean μ and covariance matrix ∑ of all sample features. Step 3: Calculate the Mahalanobis distance d M (x i ), the calculation formula for each sample is as follows:
[0017]
[0018] where x i is the eigenvector of sample i, μ is the mean of the eigenvector, ∑ -1 is the inverse of the covariance matrix.
[0019] (3-3) Combine tail data samples and samples with high uncertainty in a balanced ratio (i.e., 1:1 mixing). This mixed replay method ensures that the model not only retains the knowledge of the tail data, but also improves the accuracy of predictions by learning from uncertain samples. This method improves the model's adaptability to long-tail distributions 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, hard prompts are used to distinguish between the two types of information, source code and vulnerability description; in type recognition, the [SOFT] tag is used to create soft prompts, and the [MASK] tag 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 constructed bimodal information is input into the pre-trained language model CodeT5, and the model and prompt template are trained through prompt fine-tuning, which specifically includes the following steps:
[0021] (4-1) Construct a prompt template. Specifically, use hard prompts to distinguish source code and vulnerability descriptions, that is, add "The code snippet:" and "The vulnerability description:" respectively. For type classification, use the [SOFT] tag to create soft prompts, that is, initialize it to "Identify the vulnerability type:" during training, and use the [MASK] tag to predict or fill in missing information;
[0022] (4-2) Construct a mapping method: map the natural language output of the model to a specific category. According to each CWE-ID, map it to the corresponding software vulnerability type, and construct an approximate word for the vulnerability type, and finally obtain the mapping method adopted by the method of the present invention. For example, the vulnerability type corresponding to CWE-125 is "Out-of-bounds Read", and its two corresponding approximate words are "Buffer Overread" and "Invalid Memory Access". The final mapping of CWE-125 is, CWE-125: ["Out-of-bounds Read", "Buffer Overread", "Invalid Memory Access"];
[0023] (4-3) The bimodal information (i.e., vulnerability code and vulnerability description) of the vulnerability sample to be identified is combined with the predefined hybrid prompt template P, input into the pre-trained model CodeT5, and the model and prompt template are trained through prompt fine-tuning to generate output results in natural language form;
[0024] (4-4) Convert the natural language results generated by the model into the corresponding software vulnerability type and return the final prediction result.
[0025] (5) During the model training process, Focal Loss and Label Smooth CE Loss are first used to focus on the tail data to alleviate the problem of category imbalance. For the learning of new tasks, the EWC regularization term is used to limit the model parameter weights to make them close to the optimal parameters of the previous task to ensure that the training for the new task will not significantly change the parameters that are critical to the old task, thereby effectively alleviating the catastrophic forgetting phenomenon. 2. A software vulnerability type identification method based on long-tail incremental learning according to claim 1, characterized in that the step (3) specifically includes the following steps:
[0026] (5-1) Use the traditional stochastic gradient descent method to train the neural network model and the initial task model. 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) Focal Loss and Label Smooth CE Loss are used together to ensure that the model can focus on the tail data during training. Focal Loss introduces a scaling factor in the standard CE Loss, which enables the model to pay more attention to difficult-to-classify samples, which are often tail samples. It achieves this by applying a focusing parameter γ (used to increase the weight of misclassified samples) and a balancing factor α (used to address the problem of class imbalance). The final loss function is defined as:
[0028]
[0029] Label Smooth CE Loss softens the target label to improve generalization ability by reducing overconfidence in predictions, especially for unrepresentative categories. Its calculation method is defined as follows:
[0030]
[0031] Among them, y k =(1-∈) is the true category, ∈ is the smoothing parameter, and K is the total number of categories. Next, the weight w is used to fit the two loss functions. The specific calculation method is defined as follows:
[0032]
[0033] (5-3) The EWC regularization term is used to constrain the model parameter weights to be close to the optimal parameters of the previous task. Specifically, the Fisher information matrix is calculated to quantify the importance of each parameter to the initial task. Assume that 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 a decay factor that attempts to assign higher weights to recent tasks while retaining knowledge from earlier 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 a loss function with a 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 and return the best vulnerability type recognition model and mixed prompt template.
[0038] Compared with the prior art, the beneficial effects of the present invention are as follows: a software vulnerability type identification method based on long-tail incremental learning proposed in the present invention takes the bimodal 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 usage scenarios of incremental learning, so as to better mine vulnerability information and classify vulnerability types more accurately and reliably, while alleviating to a certain extent the catastrophic forgetting problem of the model caused by the increase of data changing over time in real projects. BRIEF DESCRIPTION OF THE DRAWINGS
[0039] The accompanying drawings are used to provide further understanding of the present invention and constitute a part of the specification. They are used to explain the present invention together with the embodiments of the present invention and do not constitute a limitation of the present invention.
[0040] Figure 1 A system framework diagram of a software vulnerability type identification method based on long-tail incremental learning provided by the present invention. DETAILED DESCRIPTION
[0041] In order to make the purpose, technical solution and advantages of the present invention more clearly understood, the present invention is further described in detail below in conjunction with the accompanying drawings and embodiments. Of course, the specific embodiments described here are only used to explain the present invention and are not used to limit the present invention.
[0042] Example 1
[0043] This embodiment provides a technical solution, which is a method for identifying software vulnerability types based on long-tail incremental learning. Figure 1 As shown, specifically including the following:
[0044] (1) By analyzing the CVE vulnerability database and actual development projects based on the Git-based software project hosting platform, we 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. Subsequently, the source code is preprocessed, such as deleting comments, blank lines, and comments at the beginning of the line, and finally constructs a vulnerability dataset D, which includes the following steps:
[0045] (1) By analyzing the CVE vulnerability database and actual development projects based on the Git-based software project hosting platform, we collect vulnerability-related information, including CWE ID, vulnerability code, vulnerability description, and submission time. The data processing process first 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. Subsequently, the source code is preprocessed, such as deleting comments, blank lines, and comments at the beginning of the line, and finally constructs a vulnerability type prediction dataset D;
[0046] (2) The vulnerability type prediction dataset D is divided into tasks according to the vulnerability submission time. Subsequently, 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. A stratified sampling strategy is used in the division process to ensure that the distribution of vulnerability types in different datasets is consistent;
[0047] (3) For the current task that requires vulnerability type identification, a hybrid replay strategy is used to select typical vulnerability samples from the data set of the old task. Before each model training, the tail data in the old task (i.e., rare vulnerability samples) is preferentially selected, and then vulnerability samples with higher uncertainty are selected based on 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 training set, validation set, and 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 vulnerability type distribution in the two data sets (i.e., the head sample set and the tail sample set), which specifically includes the following steps:
[0048] (3-1) In order to identify the tail category, we set a threshold τ, such as the category that accounts for less than 5% of the total data set, and classify it as the tail category. Specifically, the tail category can be defined as follows:
[0049]
[0050] Here, τ is set to 5% to define it as tail data.
[0051] (3-2) Select samples with high uncertainty. These samples represent the most challenging instances that the model considers, and are 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, in step 1, use the feature extraction layer of the model to obtain the feature vector x of each sample. i Step 2: Calculate the eigenvector mean μ and covariance matrix ∑ of all sample features. 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 eigenvector of sample i, μ is the mean of the eigenvector, ∑ -1 is the inverse of the covariance matrix.
[0054] (3-3) Combine tail data samples and samples with high uncertainty in a balanced ratio (i.e., 1:1 mixing). This mixed replay method ensures that the model not only retains the knowledge of the tail data, but also improves the accuracy of predictions by learning from uncertain samples. This method improves the model's adaptability to long-tail distributions 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, hard prompts are used to distinguish between the two types of information, source code and vulnerability description; in type recognition, the [SOFT] tag is used to create soft prompts, and the [MASK] tag 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 constructed bimodal information is input into the pre-trained language model CodeT5, and the model and prompt template are trained through prompt fine-tuning, which specifically includes the following steps:
[0056] (4-1) Construct a prompt template. Specifically, use hard prompts to distinguish source code and vulnerability descriptions, that is, add "The code snippet:" and "The vulnerability description:" respectively. For type classification, use the [SOFT] tag to create soft prompts, that is, initialize it to "Identify the vulnerability type:" during training, and use the [MASK] tag to predict or fill in 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 the corresponding software vulnerability type, and construct an approximate word for the vulnerability type, and finally obtain the mapping method adopted by the method of this embodiment. For example, the vulnerability type corresponding to CWE-125 is "Out-of-bounds Read", and its two corresponding approximate words are "Buffer Overread" and "Invalid Memory Access". The final mapping of CWE-125 is, CWE-125: ["Out-of-bounds Read", "Buffer Overread", "Invalid Memory Access"];
[0058] (4-3) The bimodal information (i.e., vulnerability code and vulnerability description) of the vulnerability sample to be identified is combined with the predefined hybrid prompt template P, input into the pre-trained model CodeT5, and the model and prompt template are trained through prompt fine-tuning to generate output results in natural language form;
[0059] (4-4) Convert the natural language results generated by the model into the corresponding software vulnerability type and return the final prediction result.
[0060] (5) During the model training process, Focal Loss and Label Smooth CE Loss are first used to focus on the tail data to alleviate the problem of category imbalance. For the learning of new tasks, the EWC regularization term is used to limit the model parameter weights to make them close to the optimal parameters of the previous task to ensure that the training for the new task will not significantly change the parameters that are critical to the old task, thereby effectively alleviating the catastrophic forgetting phenomenon. 2. A software vulnerability type identification method based on long-tail incremental learning according to claim 1, characterized in that the step (3) specifically includes the following steps:
[0061] (5-1) Use the traditional stochastic gradient descent method to train the neural network model and the initial task model. 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.
[0062] (5-2) Focal Loss and Label Smooth CE Loss are used together to ensure that the model can focus on the tail data during training. Focal Loss introduces a scaling factor in the standard CE Loss, which enables the model to pay more attention to difficult-to-classify samples, which are often tail samples. It achieves this by applying a focusing parameter γ (used to increase the weight of misclassified samples) and a balancing factor α (used to address the problem of class imbalance). The final loss function is defined as:
[0063]
[0064] Label Smooth CE Loss softens the target label to improve generalization ability by reducing overconfidence in predictions, especially for unrepresentative categories. Its calculation method is defined as follows:
[0065]
[0066] Among them, y k =(1-∈) is the true category, ∈ is the smoothing parameter, and K is the total number of categories. Next, the weight w is used to fit the two loss functions. The specific calculation method is defined as follows:
[0067]
[0068] (5-3) The EWC regularization term is used to constrain the model parameter weights to be close to the optimal parameters of the previous task. Specifically, the Fisher information matrix is calculated to quantify the importance of each parameter to the initial task. Assume that 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:
[0069]
[0070] Where β is a decay factor that attempts to assign higher weights to recent tasks while retaining knowledge from earlier tasks. Using the cumulative Fisher matrix, the regularization term of EWC can be defined as:
[0071]
[0072] Where λ is the penalty coefficient, is the i-th element of the cumulative Fisher matrix, θ i is the current parameter, are the parameter values learned from the previous task. 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 parameters of the initial task.
[0073] (5-4) Calculate the change in the F1 index of the model on the validation set. When the F1 value does not exceed the current best performance for 10 consecutive rounds, stop training and return the best vulnerability type recognition model and mixed prompt template.
[0074] (6) The method of this embodiment is compared with the existing software vulnerability type identification method on the same data set, and five performance indicators from the field of software vulnerability task research are used to automatically evaluate the quality of the model:
[0075] Table 1 Comparison of results between the method of this embodiment and other methods
[0076]
[0077] Experiments show that the software vulnerability type identification method based on long-tail incremental learning proposed in this embodiment is superior to the baseline method in all indicators, showing a more reliable and accurate software vulnerability type identification ability. Specifically, compared with the baseline method, the accuracy of this embodiment is improved from 125.81% to 438.46%, the precision is improved from 130.77% to 757.14%, the recall rate 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 this embodiment combines the incremental learning strategy based on hybrid playback and elastic weight consolidation (EWC), integrates source code and vulnerability description information, and can not only mine vulnerability features more comprehensively, but also effectively reduce the forgetting rate of old task data. Compared with the baseline method, this method significantly improves the ability to identify long-tail category vulnerabilities, showing 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 this embodiment with methods 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 of results between the method of this embodiment and other methods
[0081]
[0082] Experiments show that the method for identifying software vulnerability types based on long-tail incremental learning proposed in this embodiment improves the accuracy by 2.94% to 9.37%, the precision by 3.45% to 13.21%, the F1 score by 1.82% to 7.69%, and the MCC by 1.52% to 8.06% compared to other incremental learning methods. The results show that the method in this embodiment has higher performance and stronger robustness when processing incremental learning tasks, and can perform more reliable and accurate software vulnerability type identification.
[0083] Example 3
[0084] Based on the pre-trained model used in Example 1, the method of this embodiment is compared with 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 of results between the method of this embodiment and other methods
[0086]
[0087]
[0088] Experiments show that the method for identifying software vulnerability types based on long-tail incremental learning proposed in this embodiment has the highest indicators compared to other methods using pre-trained models, so it can perform more reliable and accurate software vulnerability type identification. Specifically, compared with UnixCoder, which performed second best, this embodiment improved the accuracy by 4.48%, the precision by 11.11%, and the F1 score by 1.82%. In terms of the MCC indicator, the performance of this embodiment improved by 3.08%. These results show the competitiveness of the method proposed in this embodiment.
[0089] Example 4
[0090] On the basis of Example 1, the type of prompt in this example is changed, and the method in this example is compared with methods using different prompts (i.e., soft prompts and hard prompts). Five performance indicators 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 of this embodiment and other methods
[0092]
[0093] Experiments have shown that the method for identifying software vulnerability types based on long-tail incremental learning proposed in this embodiment uses hybrid prompts. Compared with methods using other types of prompts, all indicators are the highest, so more reliable and accurate software vulnerability type identification can be performed. Compared with hard prompts, this embodiment has improved accuracy by 1.45%, precision by 3.45%, recall by 6.67%, F1 score by 3.70%, and MCC by 3.08%. At the same time, compared with soft prompts, this embodiment has improved 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 prompts and soft prompts, thereby achieving better performance in the task of software vulnerability type identification.
[0094] Example 5
[0095] On the basis of Example 1, the number of corresponding mappings of this example is changed, and the method of this example is compared with the method using different numbers of mappings. Five performance indicators 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 of this embodiment and other methods
[0097]
[0098]
[0099] Table 6 Mapping method of the method of this embodiment
[0100]
[0101] Experiments have shown that the method for identifying software vulnerability types based on long-tail incremental learning proposed in this embodiment has the highest indicators in all aspects compared to methods using other mapping numbers, so it can perform more reliable and accurate software vulnerability type identification. Specifically, for the F1 score, the method of this embodiment has a maximum improvement of 3% compared to other configurations of the number of mappings; for the MCC indicator, the method of this embodiment has a maximum improvement of 5%. These results fully verify the competitiveness of the method of this embodiment, and reflect its significant advantages in optimizing performance and controlling complexity, thus having important practical application value in the task of identifying software vulnerability types.
[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 principle of the present invention should be included in the protection scope of the present invention.
Claims
1. A method for identifying software vulnerability types based on long-tail incremental learning, characterized in that: The following steps are involved: (1) By analyzing the CVE vulnerability database and actual development projects based on the Git-based software project hosting platform, we collect vulnerability-related information, including CWE ID, vulnerability code, vulnerability description, and submission time. The data processing process first 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. Subsequently, the source code is preprocessed to delete comments, blank lines, and beginning comments, and finally a vulnerability type prediction dataset D is constructed. (2) The vulnerability type prediction dataset D is divided into tasks according to the vulnerability submission time. Then, the dataset of each task is divided into training set, validation set and test set in a ratio of 8:1:
1. The stratified sampling strategy is adopted in the division process. (3) For the current task that requires vulnerability type identification, a hybrid replay strategy is used to select typical vulnerability samples from the data set of the old task. Before each model training, the tail data in the old task is selected first, and then vulnerability samples with higher uncertainty are selected based on the Mahalanobis distance of the samples. After the vulnerability samples selected by the hybrid replay are merged with the vulnerability samples of the current task, they are re-divided into training set, validation set and 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; (4) extracting source code and vulnerability description from the data set processed in step (3), and designing a prompt template, in which hard prompts are used to distinguish between source code and vulnerability description; In type recognition, the [SOFT] tag is used to create soft prompts, and the [MASK] tag 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 constructed bimodal information is input into the pre-trained language model CodeT5, and the model and prompt template are trained through prompt fine-tuning. (5) During the model training process, Focal Loss and Label Smooth CE Loss are first used to focus on the tail data to alleviate the problem of category imbalance. For the learning of new tasks, the EWC regularization term is used to limit the model parameter weights to make them close to the optimal parameters of the previous task.
2. According to claim 1, a method for identifying software vulnerability types based on long-tail incremental learning is characterized in that: The step (3) comprises the following steps: (2-1) In order to identify the tail category, a threshold τ is set. Categories that account for less than 5% of the total data set are classified as tail categories. The tail category is defined as follows: Here, τ is set to 5% to define it as tail data; (2-2) Select samples with high uncertainty. To quantify the uncertainty of samples, the Mahalanobis distance between the feature vector of each sample and the mean of the feature vectors of all samples is used, including the following steps: Step 1: Use the model’s feature extraction layer to obtain the feature vector x for each sample i ; Step 2: Calculate the eigenvector mean μ and covariance matrix ∑ of all sample features; Step 3: Calculate the Mahalanobis distance d M (x i ), the calculation formula for each sample is as follows: where x i is the eigenvector of sample i, μ is the mean of the eigenvector, ∑ -1 is the inverse of the covariance matrix; (2-3) Mix in a balanced ratio of 1:1 to combine the tail data samples and samples with high uncertainty.
3. According to claim 1, a method for identifying software vulnerability types based on long-tail incremental learning is characterized in that: The step (4) comprises the following steps: (3-1) Construct a prompt template and use hard prompts to distinguish source code and vulnerability descriptions, adding "The codesnippet:" and "The vulnerability description:" respectively. For type classification, use the [SOFT] tag to create soft prompts, that is, initialize it to "Identify the vulnerability type:" during training, and use the [MASK] tag to predict or fill in missing information; (3-2) Constructing a mapping method: Mapping the natural language output of the model to specific categories, mapping each CWE-ID to its corresponding software vulnerability type, and constructing an approximate word for the vulnerability type, and finally obtaining the mapping method adopted by the method; (3-3) The bimodal information of the vulnerability sample to be identified, i.e., the vulnerability code and the vulnerability description, is combined with the predefined hybrid prompt template P, input into the pre-trained model CodeT5, and the model and prompt template are trained through prompt fine-tuning to generate output results in natural language form; (3-4) Convert the natural language results generated by the model into the corresponding software vulnerability type and return the final prediction result.
4. According to claim 1, a method for identifying software vulnerability types based on long-tail incremental learning is characterized in that: The step (5) comprises the following steps: (4-1) Use the traditional stochastic gradient descent method to train the neural network model and the initial task model. 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. (4-2) Focal Loss and Label Smooth CE Loss are used together to ensure that the model focuses on the tail data during training. Focal Loss introduces a scaling factor in the standard CE Loss to make the model focus on samples that are difficult to classify. These samples are often tail samples. It achieves this by applying the focus parameter γ and the balance factor α. The final loss function is defined as: Label Smooth CE Loss softens the target label to improve generalization ability by reducing overconfidence in predictions, especially for unrepresentative categories. Its calculation method is defined as follows: Among them, y k =(1-∈) is the true category, ∈ is the smoothing parameter, K is the total number of categories, and then the weight w is used to fit the two loss functions. The calculation method is defined as follows: (4-3) The EWC regularization term is used to limit the model parameter weights to make them close to the optimal parameters of the previous task. The Fisher information matrix is calculated to quantify the importance of each parameter to the initial task. Assume that F t represents the Fisher matrix of the current task t, Represents the cumulative Fisher matrix. After each task is completed, Update as follows: Where β is a decay factor that attempts to assign higher weights to recent tasks while retaining knowledge from earlier tasks. Using the cumulative Fisher matrix, the regularization term of EWC is defined as: Among them, λ is the penalty coefficient, is the i-th element of the cumulative Fisher matrix, θ i is the current parameter, are the parameter values learned from the previous task. When training a new task, a loss function with a regularization term is used to ensure that the model parameters do not deviate too much from the parameters of the initial task. (4-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 and return the best vulnerability type recognition model and mixed prompt template.
Citation Information
Patent Citations
Information security vulnerability classification method based on improved Bert model
CN113961933A
Multi-type vulnerability detection method and system based on multi-sample comparison and fusion
CN114756861A
Evaluation data generation method and device and storage medium
CN118779880A
Software vulnerability assessment method based on continuous learning
CN119026138A
System and methods for reducing an organization's cybersecurity risk by determining the function and seniority of employees
US11640470B1
Cited By
Statement-level code vulnerability detection method based on fusion feature prompt tuning
CN119475345A
Statement-level code vulnerability detection method based on fusion feature prompt tuning
CN119475345B
Double-task vulnerability intelligent evaluation method and system
CN121765734A