A Source Code Vulnerability Detection Method Based on Multimodal Contrast Learning and Low-Rank Adaptation
By employing multimodal contrastive learning and low-rank adaptation techniques, the multimodal feature extraction and alignment capabilities of the source code vulnerability detection model are enhanced, addressing the issue of structural information being ignored in existing methods and achieving more efficient vulnerability detection.
Patent Information
- Application Number
- CN202411922007.9
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-12-25
- Publication Date
- 2025-10-28
- Estimated Expiration
- 2044-12-25
AI Technical Summary
Existing deep learning-based source code vulnerability detection methods ignore the structural information of the code, resulting in insufficient detection accuracy and difficulty in effectively capturing potential structural vulnerabilities in the code.
Multimodal contrastive learning and low-rank adaptation techniques are introduced to enhance multimodal feature extraction and alignment by performing independent low-rank adaptation adjustments on source code, abstract syntax trees, and code comments, and to use contrastive learning to improve the robustness and discriminability of the model.
It improves the accuracy and efficiency of vulnerability detection, reduces computational resources and training burden, and significantly enhances the model's ability to distinguish between similar code changes.
Smart Images
Figure CN119760728B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of computer technology, and in particular to a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation. Background Technology
[0002] Detecting source code vulnerabilities is a long-standing and highly challenging problem in software development and maintenance. These vulnerabilities can not only cause application crashes but also lead to serious security issues such as information leaks and data tampering. Therefore, timely and effective vulnerability detection is crucial for ensuring software security.
[0003] In recent years, with the continuous development of deep learning technology, deep learning-based vulnerability detection methods have gradually become a research focus. These methods utilize pre-trained language models to learn semantic representations from large amounts of source code, overcoming the shortcomings of traditional machine learning methods that rely on manually designed features. For example, the CodeBERT model effectively learns the basic semantics of source code by transforming code fragments into token sequences and training on them. However, these methods ignore the structural information of the code when processing it, which limits the effectiveness of vulnerability detection. Due to the complexity and diversity of code, traditional sequence-based representation methods struggle to capture potential structural vulnerabilities in the code, thus affecting detection accuracy.
[0004] How to solve the above problems has become the challenge facing this invention. Summary of the Invention
[0005] The purpose of this invention is to provide a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation. This invention enhances the learning ability for different modalities such as source code, abstract syntax trees, and code comments during the fine-tuning stage by introducing multimodal contrastive learning and low-rank adaptation techniques. Specifically, this invention performs independent low-rank adaptation adjustments on each modality, enabling the full extraction and optimization of features for each modality. This overcomes the limitations of single-modal methods in vulnerability detection tasks. Simultaneously, through contrastive learning, this invention effectively aligns and distinguishes the embedding representations of multiple modalities, improving the model's robustness and discriminative ability in identifying similar code, thus increasing the accuracy and efficiency of source code vulnerability detection while reducing training resources and computational burden.
[0006] The core idea of this invention is to introduce multimodal contrastive learning to align and optimize the representations of source code, AST, and code comments during the fine-tuning phase. Contrastive learning is then used to calculate the similarity differences between different modalities, thereby enhancing the model's ability to distinguish similar code. Simultaneously, a low-rank adaptation technique is employed to specifically optimize the features of each modality, reducing computational costs during model training and improving training efficiency. This allows the model to improve vulnerability detection accuracy while maintaining low computational resource consumption and high training efficiency.
[0007] This invention is achieved through the following measures: a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation, comprising the following steps:
[0008] 1.1: Preprocess and split the dataset into training set D1 and test set D2;
[0009] 1.2: Extract source code, abstract syntax tree (AST), and code comments from the training set D1 divided in step 1.1 as multimodal input;
[0010] 1.3: Use the GraphCodeBERT model to learn feature representations for code snippets of each modality, generating an initial embedding representation h for each modality. code h ast and h comment ;
[0011] 1.4: Fine-tune the representation of each modality obtained in step 1.3 by applying low-rank adaptation to optimize feature extraction for each modality and reduce computational overhead;
[0012] 1.5: Apply contrastive learning to align the embedding representations of each modality and calculate the contrastive loss. To enhance the robustness and discriminative power of the model;
[0013] 1.6: Use the trained model to perform vulnerability detection on the data in the test set D2, and output the classification result of whether each code snippet has a vulnerability.
[0014] The step 1.2, which involves extracting source code, abstract syntax tree (AST), and code comments from the training set D1 partitioned in step 1.1 as multimodal input, specifically includes the following steps:
[0015] 2.1: Enhance the semantic information of code comment modalities by generating missing or incomplete code comments in training set D1 using GPT-4o-mini;
[0016] 2.2: Generate the AST tree of the source code in training set D1 using the Tree-Sitter tool to represent the semantic information of the code AST modalities;
[0017] 2.3: Use source code, AST and code comments as multimodal inputs, and extract feature representations for each modality respectively.
[0018] In step 1.3, the GraphCodeBERT model is used to learn feature representations for the code snippets of each modality, generating an initial embedding representation h for each modality. code h ast and h comment Specifically, it includes the following steps:
[0019] 3.1: Input the embedded representations of the source code, AST, and code comments into the GraphCodeBERT model to generate the preliminary feature representation h for each modality. code h ast and h comment .
[0020] In step 1.4, the representation of each modality obtained in step 1.3 is fine-tuned using low-rank adaptation to optimize feature extraction for each modality and reduce computational overhead. Specifically, this includes the following steps:
[0021] 4.1: Apply low-rank adaptation to the representation of source code modes to generate a source code-specific low-rank adaptation matrix A. code and B code To enhance its structural information;
[0022] 4.2: Applying low-rank adaptation to the representation of AST modes to generate an AST-specific low-rank adaptation matrix A ast and B ast To enhance the tree structure and syntactic dependency information;
[0023] 4.3: Apply low-rank adaptation to the representation of code comment modalities to generate a code comment-specific low-rank adaptation matrix A comment and B comment To enhance its semantic information.
[0024] In step 1.5, contrastive learning is applied to align the embedding representations of each modality, and the contrastive loss is calculated. To enhance the robustness and discriminative power of the model, the following steps are included:
[0025] 5.1: Calculate the contrast loss between different modes This includes contrast loss between code and AST, contrast loss between code and comments, and contrast loss between AST and comments. It improves the model's ability to distinguish similar code variations by maximizing the similarity between samples of the same class and minimizing the similarity between samples of different classes. The specific loss function... as follows:
[0026]
[0027] Among them, sim(h) m ,h n ) represents the similarity between embedding and , τ is the temperature coefficient, h m and h n For positive samples, use h; for negative samples, use h. a express.
[0028] 5.2: The total contrast loss is obtained by weighted summation of the various contrast losses. And calculate the total loss together with the classification loss. The model parameters are optimized. The final total loss function is as follows:
[0029]
[0030] in, Represents cross-entropy loss, This represents the total contrastive loss, while x is a hyperparameter that controls the weight of the structural loss.
[0031] In step 1.6, the trained model is used to perform vulnerability detection on the data in the test set D2, and the classification result of whether each code snippet has a vulnerability is output.
[0032] Compared with the prior art, the beneficial effects of the present invention are as follows:
[0033] 1. This invention utilizes multimodal contrastive learning, enabling the model to effectively enhance the alignment and differentiation capabilities among source code, AST, and code comments, thereby improving sensitivity and processing capabilities for similar code changes. Simultaneously, low-rank adaptation techniques effectively reduce the number of trainable parameters, improving training efficiency and lowering computational overhead. These innovations significantly improve the accuracy and robustness of vulnerability detection tasks.
[0034] 2. This invention enhances the understanding of code context through multimodal feature fusion. Traditional methods typically rely on single-modal input for vulnerability detection. This invention enhances the model's understanding of multidimensional code information by fusing features from three modalities, further improving the accuracy of vulnerability detection.
[0035] 3. This invention introduces a low-rank adaptation technique, which improves training efficiency and reduces computational overhead by reducing the number of trainable parameters in the model. Traditional deep learning models require a large number of trainable parameters, leading to long training times and high computational resource consumption. To address this issue, this invention introduces a low-rank matrix, effectively compressing the number of parameters for each modality while maintaining the model's efficiency and performance. Through this method, this invention can significantly improve training efficiency and reduce computational resource consumption while maintaining high vulnerability detection accuracy.
[0036] 4. This invention introduces multimodal contrastive learning, overcoming the limitations of existing pre-trained models in handling similar code changes. By applying contrastive learning methods under multimodal input, this invention enhances the model's ability to distinguish between similar code changes. This significantly improves the model's accuracy in vulnerability detection, especially when handling complex code modifications. Attached Figure Description
[0037] The accompanying drawings are provided to further illustrate the invention and form part of the specification. They are used together with the embodiments of the invention to explain the invention and do not constitute a limitation thereof.
[0038] Figure 1 This invention provides a system framework diagram for a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation. Detailed Implementation
[0039] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and embodiments. Of course, the specific embodiments described herein are merely illustrative and not intended to limit the invention.
[0040] Example 1
[0041] See Figure 1 As shown, this embodiment provides a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation, including the following steps:
[0042] (1-1) Preprocess and split the dataset into training set D1 and test set D2;
[0043] (1-2) Extract source code, abstract syntax tree (AST) and code comments from the training set D1 divided in step (1-1) as multimodal input;
[0044] (1-3) Use the GraphCodeBERT model to learn feature representations for code snippets of each modality, generating an initial embedding representation h for each modality. code h ast and hcomment ;
[0045] (1-4) Fine-tune the representation of each modality obtained in step (1-3) using low-rank adaptation to optimize feature extraction for each modality and reduce computational overhead.
[0046] (1-5) Apply contrastive learning to align the embedding representations of each modality and calculate the contrastive loss. To enhance the robustness and discriminative power of the model;
[0047] (1-6) Use the trained model to perform vulnerability detection on the data in the test set D2, and output the classification result of whether each code snippet has a vulnerability.
[0048] In step (1-2), the source code, abstract syntax tree (AST), and code comments are extracted from the training set D1 partitioned in step (1-1) as multimodal input:
[0049] (2-1) Use GPT-4o-mini to generate missing or incomplete code comments in training set D1 to enhance the semantic information of code comment modality;
[0050] (2-2) Generate the AST tree of the source code in the training set D1 using the Tree-Sitter tool to represent the semantic information of the code AST modality;
[0051] (2-3) Take the source code, AST and code comments as multimodal inputs and extract the feature representation of each modality respectively.
[0052] In steps (1-3), the GraphCodeBERT model is used to learn feature representations for the code snippets of each modality, generating an initial embedding representation h for each modality. code h ast and h comment Specifically, it includes the following steps:
[0053] (3-1) Input the embedded representations of the source code, AST, and code comments into the GraphCodeBERT model to generate the preliminary feature representation h for each modality. code h ast and h comment .
[0054] In steps (1-4), the representation of each modality obtained in steps (1-3) is fine-tuned using low-rank adaptation to optimize feature extraction for each modality and reduce computational overhead. Specifically, this includes the following steps:
[0055] (4-1) Apply low-rank adaptation to the representation of source code modes to generate a source code-specific low-rank adaptation matrix A. code and Bcode To enhance its structural information;
[0056] (4-2) Apply low-rank adaptation to the representation of AST modes to generate an AST-specific low-rank adaptation matrix A. ast and B ast To enhance the tree structure and syntactic dependency information;
[0057] (4-3) Apply low-rank adaptation to the representation of code comment modalities to generate a code comment-specific low-rank adaptation matrix A. comment and B comment To enhance its semantic information.
[0058] In steps (1-5), contrastive learning is applied to align the embedded representations of each modality, and the contrastive loss is calculated. To enhance the robustness and discriminative power of the model, the following steps are included:
[0059] (5-1) Calculate the contrast loss between different modes This includes contrast loss between code and AST, contrast loss between code and comments, and contrast loss between AST and comments. It improves the model's ability to distinguish similar code variations by maximizing the similarity between samples of the same class and minimizing the similarity between samples of different classes. The specific loss function... as follows:
[0060]
[0061] Among them, sim(h) m ,h n ) represents the similarity between embedding and , τ is the temperature coefficient, h m and h n For positive samples, use h; for negative samples, use h. a express.
[0062] (5-2) The total comparison loss is obtained by weighted summation of the individual comparison losses. And calculate the total loss together with the classification loss. The model parameters are optimized. The final total loss function is as follows:
[0063]
[0064] in, Represents cross-entropy loss, This represents the total contrastive loss, while x is a hyperparameter that controls the weight of the structural loss.
[0065] In steps (1-6), the trained model is used to perform vulnerability detection on the data in the test set D2, and the classification result of whether each code snippet has a vulnerability is output.
[0066] (7) The method of this embodiment and existing vulnerability detection methods are evaluated on the same dataset Devign. The quality of the model is automatically evaluated using four performance metrics from the vulnerability detection research field (i.e., Accuracy, Recall, Precision and F1-score).
[0067] Table 1. Comparison of results between the method in this embodiment and other methods.
[0068]
[0069] Experiments show that the source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation proposed in this embodiment can achieve more accurate vulnerability detection compared to baseline methods. Specifically, by combining multimodal input with contrastive learning and low-rank adaptation, the method in this embodiment enables the model to better distinguish similar code changes, further improving the accuracy of vulnerability detection and surpassing these baseline methods. Specifically, for accuracy, the method in this embodiment can achieve a relative performance improvement of at least 4.71%; for recall, the method in this embodiment can achieve a relative performance improvement of at least 54.52%; and for F1-score, the method in this embodiment can achieve a relative performance improvement of at least 12.85%. These results demonstrate the competitiveness of the proposed method in this embodiment.
[0070] Example 2
[0071] Based on Example 1, the method in this example is evaluated on the same dataset Reveal, along with existing vulnerability detection methods. Four performance metrics from the vulnerability detection research field (Accuracy, Recall, Precision, and F1-score) are used to automatically evaluate the quality of the model.
[0072] Table 2 Comparison of results between the method in this embodiment and other methods.
[0073]
[0074] Experiments show that the source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation proposed in this embodiment can achieve more accurate vulnerability detection compared to baseline methods. Specifically, by combining multimodal input with contrastive learning and low-rank adaptation, the method in this embodiment improves the accuracy of vulnerability detection and surpasses these baseline methods. Specifically, in terms of accuracy, the method in this embodiment can improve performance by at least 2.20%; in terms of precision, the method in this embodiment can improve performance by at least 26.06%; and in terms of F1-score, the method in this embodiment can improve performance by at least 13.67%. These results demonstrate the competitiveness of the proposed method in this embodiment.
[0075] Example 3
[0076] Based on Example 1, this example evaluates the method and existing vulnerability detection methods on the same dataset Big_Vul, using four performance metrics (Accuracy, Recall, Precision, and F1-score) from the vulnerability detection research field to automatically evaluate the quality of the model.
[0077] Table 3 Comparison of results between the method in this embodiment and other methods
[0078]
[0079]
[0080] Experiments show that the source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation proposed in this embodiment can achieve more accurate vulnerability detection compared to baseline methods. Specifically, by combining multimodal input with contrastive learning and low-rank adaptation, the method in this embodiment improves the accuracy of vulnerability detection and surpasses these baseline methods. Specifically, in terms of accuracy, the method in this embodiment can improve performance by at least 2.04%; in terms of precision, the method can improve performance by at least 35.34%; in terms of F1-score, the method can improve performance by at least 3.47%; and in terms of F1-score, the method can improve performance by at least 17.26%. These results demonstrate the competitiveness of the proposed method.
[0081] 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 source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation, characterized by comprising the following steps: 1.1: Preprocess and split the dataset into training set D1 and test set D2; 1.2: Extract source code, abstract syntax tree (AST), and code comments from the training set D1 partitioned in step 1.1 as multimodal input; 1.3: Use the GraphCodeBERT model to learn feature representations for code snippets of each modality, generating an initial embedding representation h for each modality. code h ast and h comment ; 1.4: Fine-tune the representation of each modality obtained in step 1.3 by applying low-rank adaptation to optimize feature extraction for each modality and reduce computational overhead; 1.5: Apply contrastive learning to align the embedding representations of each modality and calculate the contrastive loss. 1.6: Use the trained model to perform vulnerability detection on the data in the test set D2, and output the classification result of whether each code snippet has a vulnerability.
2. The source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation according to claim 1, characterized in that, In step 1.2, the source code, abstract syntax tree (AST), and code comments are extracted from the training set D1 divided in step 1.1 as multimodal input, including the following steps: 2.1: Enhance the semantic information of code comment modalities by generating missing or incomplete code comments in training set D1 using GPT-4o-mini; 2.2: Generate the AST tree of the source code in training set D1 using the Tree-Sitter tool to represent the semantic information of the code AST modalities; 2.3: Use source code, AST and code comments as multimodal inputs, and extract feature representations for each modality respectively.
3. The source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation according to claim 1, characterized in that, In step 1.3, the GraphCodeBERT model is used to learn feature representations for the code snippets of each modality, generating an initial embedding representation h for each modality. code h ast and h comment Specifically: 3.1: Input the embedded representations of the source code, AST, and code comments into the GraphCodeBERT model to generate the preliminary feature representation h for each modality. code h ast and h comment .
4. The source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation according to claim 1, characterized in that, In step 1.4, the representation of each modality obtained in step 1.3 is fine-tuned using low-rank adaptation to optimize feature extraction for each modality and reduce computational overhead, including the following steps: 4.1: Apply low-rank adaptation to the representation of source code modes to generate a source code-specific low-rank adaptation matrix A. code and B code To enhance its structural information; 4.2: Applying low-rank adaptation to the representation of AST modes to generate an AST-specific low-rank adaptation matrix A ast and B ast To enhance the tree structure and syntactic dependency information; 4.3: Apply low-rank adaptation to the representation of code comment modalities to generate a code comment-specific low-rank adaptation matrix A comment and B comment To enhance its semantic information.
5. The source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation according to claim 1, characterized in that, In step 1.5, contrastive learning is applied to align the embedding representations of each modality, and the contrastive loss is calculated. To enhance the robustness and discriminative power of the model, the following steps are included: 5.1: Calculate the contrast loss between different modes This includes contrast loss between code and AST, contrast loss between code and comments, and contrast loss between AST and comments. It improves the model's ability to distinguish similar code variations by maximizing the similarity between samples of the same class and minimizing the similarity between samples of different classes. The loss function... as follows: Among them, sim(h) m ,h n ) represents the similarity between embedding and , τ is the temperature coefficient, h m and h n For positive samples, use h; for negative samples, use h. a express; 5.2: The total contrast loss is obtained by weighted summation of the various contrast losses. And calculate the total loss together with the classification loss. After optimizing the model parameters, the final total loss function is as follows: in, Represents cross-entropy loss, This represents the total contrastive loss, while x is a hyperparameter that controls the weight of the structural loss.
6. The source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation according to claim 1, characterized in that, In step 1.6, the trained model is used to perform vulnerability detection on the data in the test set D2, and the classification result of whether each code snippet has a vulnerability is output.