Source code vulnerability detection method based on multi-modal comparative learning and low-rank adaptation

Abstract

The invention provides a source code vulnerability detection method based on multi-modal comparative learning and low-rank adaptation, belongs to the technical field of computers, and solves the technical problems that in a traditional vulnerability detection method, multi-modal information cannot be fully utilized, and the bottleneck of calculation efficiency is faced. According to the technical scheme, the method comprises the following steps that S1, a data set is preprocessed and segmented; s2, extracting a preliminary representation of each mode through a pre-training model GraphCodeBERT; s3, performing fine tuning on the representation of each mode by applying low-rank adaptation; s4, aligning representations of different modes by using comparative learning; s5, in the training process, multi-modal representation of the model is optimized; and S6, inputting the test set into the trained model for vulnerability detection, and outputting a classification result of whether the codes have vulnerabilities or not. The method has the beneficial effects that the code vulnerability identification capability of the model is improved, and the calculation cost is reduced.

Description

Technical Field

[0001] The present invention relates to the field of computer technology, and in particular to a source code vulnerability detection method based on multimodal contrast learning and low-rank adaptation. Background Art

[0002] In the process of software development and maintenance, the detection of source code vulnerabilities is a long-standing and challenging problem. These vulnerabilities may not only cause the application to crash, but also cause serious security issues such as information leakage and data tampering. Therefore, timely and effective vulnerability detection is crucial to ensure software security.

[0003] In recent years, with the continuous development of deep learning technology, vulnerability detection methods based on deep learning have gradually become the focus of research. These methods use pre-trained language models to learn semantic representations from a large amount of source code, overcoming the defects of traditional machine learning methods that rely on manually designed features. For example, the CodeBERT model can effectively learn the basic semantics of source code by converting code snippets into token sequences and training on this basis. However, these methods ignore the structural information of the code when processing source code, which limits the effectiveness of vulnerability detection. Due to the complexity and diversity of the code, traditional sequence-based representation methods have difficulty capturing potential structural vulnerabilities in the code, thus affecting the detection accuracy.

[0004] How to solve the above problems has become the subject faced by the present invention. Summary of the invention

[0005] The purpose of the present invention is to provide a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation. The present invention enhances the learning ability of different modalities such as source code, abstract syntax tree and code comments in the fine-tuning stage by introducing multimodal contrastive learning and low-rank adaptation technology; specifically, the present invention performs independent low-rank adaptation adjustment on each modality so that the features of each modality can be fully extracted and optimized, thereby solving the limitations of single modality methods in vulnerability detection tasks. At the same time, with the help of contrastive learning, the present invention effectively aligns and distinguishes the embedded representations of multimodal, improves the robustness and distinguishing ability of the model in identifying similar codes, improves the accuracy and efficiency of source code vulnerability detection, and reduces training resources and computational burden.

[0006] The idea of ​​the present invention is: by introducing multimodal contrastive learning, the representation of source code, AST and code comments are aligned and optimized in the fine-tuning stage, and the similarity difference between different modalities is calculated by contrastive learning, so as to enhance the ability of the model to distinguish similar codes. At the same time, low-rank adaptation technology is used to optimize the features of each modality, reduce the computational consumption in the model training process, and improve the training efficiency of the model. It can maintain low computing resource consumption and high training efficiency while improving the accuracy of vulnerability detection.

[0007] The present invention is achieved by the following measures: a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation, which includes the following steps:

[0008] 1.1: Preprocess and split the data set 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 and generate a preliminary embedding representation h for each modality code 、h ast and h comment ;

[0011] 1.4: Apply low-rank adaptation to fine-tune the representation of each modality obtained in step 1.3 to optimize the feature extraction of each modality and reduce the 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 discrimination ability 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 results of whether each code snippet has a vulnerability.

[0014] In step 1.2, source code, abstract syntax tree (AST) and code comments are extracted from the training set D1 divided in step 1.1 as multimodal input, which specifically includes the following steps:

[0015] 2.1: Generate missing or incomplete code comments in the training set D1 through GPT-4o-mini to enhance the semantic information of the code comment modality;

[0016] 2.2: Generate the AST tree of the source code in the training set D1 through the Tree-Sitter tool to represent the semantic information of the code AST modality;

[0017] 2.3: Take source code, AST and code comments as multi-modal input and extract feature representations of each modality respectively.

[0018] In step 1.3, the GraphCodeBERT model is used to learn the feature representation of the code snippets of each modality and generate a preliminary embedding representation h for each modality. code 、h ast and h comment , specifically including the following steps:

[0019] 3.1: Input the embedded representations of source code, AST, and code comments into the GraphCodeBERT model to generate preliminary feature representations 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 by applying low-rank adaptation to optimize the feature extraction of each modality and reduce the computational overhead, which specifically includes the following steps:

[0021] 4.1: Apply low-rank adaptation to the representation of the source code modality to generate a source code-specific low-rank adaptation matrix A code and B code , to enhance its structural information;

[0022] 4.2: Apply low-rank adaptation to the representation of AST modality to generate AST-specific low-rank adaptation matrix A ast and B ast , to enhance the tree structure and grammatical dependency information;

[0023] 4.3: Apply low-rank adaptation to the representation of the code comment modality 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 calculate the contrast loss To enhance the robustness and discrimination ability of the model, the specific steps include:

[0025] 5.1: Calculating contrast loss between different modalities Including the contrast loss between code and AST, the contrast loss between code and comments, and the contrast loss between AST and comments. By maximizing the similarity between samples of the same type and minimizing the similarity between samples of different types, the model's ability to distinguish similar code changes is improved. The specific loss function is as follows:

[0026]

[0027] Among them, sim(h m ,h n ) represents the similarity between the embedding representation and , τ is the temperature coefficient, h m and h n is a positive sample, and h is used for negative samples a express.

[0028] 5.2: Weighted sum of each contrast loss to get the total contrast loss And calculate the total loss together with the classification loss To optimize the parameters of the model. The final total loss function is as follows:

[0029]

[0030] in, represents the cross entropy loss, represents the total contrastive loss, and 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 a classification result of whether each code snippet has a vulnerability is output.

[0032] Compared with the prior art, the present invention has the following beneficial effects:

[0033] 1. Through multimodal comparative learning, the model can effectively enhance the alignment and differentiation capabilities between the three modes of source code, AST, and code comments, thereby improving the sensitivity and processing capabilities of similar code changes. At the same time, low-rank adaptation technology effectively reduces the number of trainable parameters, improves training efficiency, and reduces computational overhead. These innovations significantly improve the accuracy and robustness of vulnerability detection tasks.

[0034] 2. The present invention enhances the ability to understand code context through multimodal feature fusion. Traditional methods usually rely on single modal input for vulnerability detection. The present 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. The present invention introduces low-rank adaptation technology to improve training efficiency and reduce computing overhead by reducing the trainable parameters of the model. Traditional deep learning models require a large number of trainable parameters, which leads to long-term training and large consumption of computing resources. In order to solve this problem, the present invention introduces a low-rank matrix, so that the number of parameters of each modality is effectively compressed while maintaining the efficiency and high performance of the model. In this way, the present invention can greatly improve the training efficiency and reduce the consumption of computing resources while maintaining a high vulnerability detection accuracy.

[0036] 4. The present invention introduces multimodal contrastive learning to solve the limitations of existing pre-trained models when dealing with similar code changes. The present invention enhances the model's ability to distinguish similar code changes by applying contrastive learning methods under multimodal input. This significantly improves the accuracy of the model in vulnerability detection, especially when dealing with complex code modifications. BRIEF DESCRIPTION OF THE DRAWINGS

[0037] 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.

[0038] Figure 1 A system framework diagram of a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation provided by the present invention. DETAILED DESCRIPTION

[0039] 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.

[0040] Example 1

[0041] See also Figure 1 As shown, this embodiment provides a source code vulnerability detection method based on multimodal contrast learning and low-rank adaptation, comprising the following steps:

[0042] (1-1) Preprocess and split the data set into training set D1 and test set D2;

[0043] (1-2) extracting 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 the feature representation of the code snippets of each modality and generate a preliminary embedding representation h for each modality. code 、h ast and hcomment ;

[0045] (1-4) Apply low-rank adaptation to fine-tune the representation of each modality obtained in step (1-3) to optimize the feature extraction of each modality and reduce the 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 discrimination ability 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 the step (1-2), source code, abstract syntax tree AST and code comments are extracted from the training set D1 divided in step (1-1) as multimodal input:

[0049] (2-1) Generate missing or incomplete code comments in the training set D1 through GPT-4o-mini to enhance the semantic information of the code comment modality;

[0050] (2-2) Generate the AST tree of the source code in the training set D1 through the Tree-Sitter tool to represent the semantic information of the code AST modality;

[0051] (2-3) Source code, AST and code comments are taken as multimodal inputs, and feature representations of each modality are extracted separately.

[0052] In the steps (1-3), the GraphCodeBERT model is used to learn the feature representation of the code snippet of each modality, and a preliminary embedding representation h of each modality is generated. code 、h ast and h comment , specifically including the following steps:

[0053] (3-1) Input the embedded representations of source code, AST, and code comments into the GraphCodeBERT model to generate a preliminary feature representation h for each modality code 、h ast and h comment .

[0054] In the step (1-4), the representation of each modality obtained in the step (1-3) is fine-tuned by applying low-rank adaptation to optimize the feature extraction of each modality and reduce the computational overhead, which specifically includes the following steps:

[0055] (4-1) Apply low-rank adaptation to the representation of the source code modality 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 the AST modality to generate an AST-specific low-rank adaptation matrix A ast and B ast , to enhance the tree structure and grammatical dependency information;

[0057] (4-3) Apply low-rank adaptation to the representation of the code comment modality to generate a code comment-specific low-rank adaptation matrix A comment and B comment , to enhance its semantic information.

[0058] In the steps (1-5), contrastive learning is applied to align the embedding representations of each modality and calculate the contrast loss To enhance the robustness and discrimination ability of the model, the specific steps include:

[0059] (5-1) Calculate the contrast loss between different modalities Including the contrast loss between code and AST, the contrast loss between code and comments, and the contrast loss between AST and comments. By maximizing the similarity between samples of the same type and minimizing the similarity between samples of different types, the model's ability to distinguish similar code changes is improved. The specific loss function is as follows:

[0060]

[0061] Among them, sim(h m ,h n ) represents the similarity between the embedding representation and , τ is the temperature coefficient, h m and h n is a positive sample, and h is used for negative samples a express.

[0062] (5-2) The weighted sum of each contrast loss is obtained to obtain the total contrast loss And calculate the total loss together with the classification loss To optimize the parameters of the model. The final total loss function is as follows:

[0063]

[0064] in, represents the cross entropy loss, represents the total contrastive loss, and x is a hyperparameter that controls the weight of the structural loss.

[0065] In the steps (1-6), the trained model is used to perform vulnerability detection on the data in the test set D2, and a classification result of whether each code snippet has a vulnerability is output.

[0066] (7) The method of this embodiment and the existing vulnerability detection methods are evaluated on the same dataset Devign, and four performance indicators (i.e., Accuracy, Recall, Precision, and F1-score) from the vulnerability detection research field are used to automatically evaluate the quality of the model.

[0067] Table 1 Comparison of results between the method of this embodiment and other methods

[0068]

[0069] Experiments have shown that the source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation proposed in this embodiment can perform more accurate vulnerability detection compared to the baseline method. Specifically, the method of this embodiment combines multimodal input with contrastive learning and low-rank adaptation, so that the model can better distinguish similar code changes, further improve the accuracy of vulnerability detection, and can surpass these baseline methods; among them, for Accuracy, the method of this embodiment can at least relatively improve the performance by 4.71%; for Recall, the method of this embodiment can at least relatively improve the performance by 54.52%; for F1-score, the method of this embodiment can at least relatively improve the performance by 12.85%; these results show the competitiveness of the method proposed in this embodiment.

[0070] Example 2

[0071] Based on Example 1, the method of this embodiment evaluates the method of this embodiment and existing vulnerability detection methods on the same dataset Reveal, and uses four performance indicators (i.e., Accuracy, Recall, Precision, and F1-score) from the field of vulnerability detection research to automatically evaluate the quality of the model.

[0072] Table 2 Comparison of results between the method of this embodiment and other methods

[0073]

[0074] Experiments have shown that the source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation proposed in this embodiment can perform more accurate vulnerability detection compared to the baseline method. Specifically, the method of this embodiment improves the accuracy of vulnerability detection by combining multimodal input with contrastive learning and low-rank adaptation, and can surpass these baseline methods; among them, for Accuracy, the method of this embodiment can at least improve the performance by 2.20%; for Precision, the method of this embodiment can at least improve the performance by 26.06%; for F1-score, the method of this embodiment can at least improve the performance by 13.67%; these results show the competitiveness of the method proposed in this embodiment.

[0075] Example 3

[0076] Based on Example 1, the method of this embodiment evaluates the method of this embodiment and existing vulnerability detection methods on the same dataset Big_Vul, and uses four performance indicators (i.e., Accuracy, Recall, Precision, and F1-score) from the field of vulnerability detection research to automatically evaluate the quality of the model.

[0077] Table 3 Comparison of results between the method of 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 perform more accurate vulnerability detection compared to the baseline method. Specifically, the method of this embodiment improves the accuracy of vulnerability detection by combining multimodal input with contrastive learning and low-rank adaptation, and can surpass these baseline methods; among them, for Accuracy, the method of this embodiment can at least improve the performance by 2.04%; for Precision, the method of this embodiment can at least improve the performance by 35.34%; for Precision, the method of this embodiment can at least improve the performance by 3.47%; for F1-score, the method of this embodiment can at least improve the performance by 17.26%; these results show the competitiveness of the method proposed in this embodiment.

[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 principle of the present invention should be included in 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 in that it includes the following steps: 1.1: Preprocess and split the data set 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 divided in step 1.1 as multimodal input; 1.3: Use the GraphCodeBERT model to learn feature representations for code snippets of each modality and generate a preliminary embedding representation h for each modality code 、h ast and h comment ; 1.4: Apply low-rank adaptation to fine-tune the representation of each modality obtained in step 1.3 to optimize the feature extraction of each modality and reduce the 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 results of whether each code snippet has a vulnerability.

2. According to claim 1, a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation is characterized in that: In the step 1.2, 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: Generate missing or incomplete code comments in the training set D1 through GPT-4o-mini to enhance the semantic information of the code comment modality; 2.2: Generate the AST tree of the source code in the training set D1 through the Tree-Sitter tool to represent the semantic information of the code AST modality; 2.3: Take source code, AST and code comments as multi-modal input and extract feature representations of each modality respectively.

3. According to claim 1, a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation is characterized in that: In step 1.3, the GraphCodeBERT model is used to learn the feature representation of the code snippets of each modality and generate a preliminary embedding representation h for each modality. code 、h ast and h comment , specifically: 3.1: Input the embedded representations of source code, AST, and code comments into the GraphCodeBERT model to generate preliminary feature representations h for each modality code 、h ast and h comment .

4. According to claim 1, a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation is characterized in that: In step 1.4, the representation of each modality obtained in step 1.3 is fine-tuned by applying low-rank adaptation to optimize the feature extraction of each modality and reduce the computational overhead, including the following steps: 4.1: Apply low-rank adaptation to the representation of the source code modality to generate a source code-specific low-rank adaptation matrix A code and B code , to enhance its structural information; 4.2: Apply low-rank adaptation to the representation of AST modality to generate AST-specific low-rank adaptation matrix A ast and B ast , to enhance the tree structure and grammatical dependency information; 4.3: Apply low-rank adaptation to the representation of the code comment modality to generate a code comment-specific low-rank adaptation matrix A comment and B comment , to enhance its semantic information.

5. According to claim 1, a source code vulnerability detection method based on multimodal contrastive learning and low-rank adaptation is characterized in that: In step 1.5, contrastive learning is applied to align the embedding representations of each modality and calculate the contrast loss To enhance the robustness and discrimination ability of the model, the following steps are included: 5.1: Calculating contrast loss between different modalities Including the contrast loss between code and AST, the contrast loss between code and comments, and the contrast loss between AST and comments. By maximizing the similarity between samples of the same type and minimizing the similarity between samples of different types, the model's ability to distinguish similar code changes is improved. The loss function as follows: Among them, sim(h m ,h n ) represents the similarity between the embedding representation and , τ is the temperature coefficient, h m and h n is a positive sample, and h is used for negative samples a express; 5.2: Weighted sum of each contrast loss to get the total contrast loss And calculate the total loss together with the classification loss To optimize the parameters of the model, the final total loss function is as follows: in, represents the cross entropy loss, represents the total contrastive loss, and 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 a classification result of whether each code snippet has a vulnerability is output.

Images