Vulnerability detection method based on mixed text and code pre-training coding

Abstract

The invention relates to the technical field of vulnerability detection, and discloses a mixed text and code pre-training coding-based vulnerability detection method, which comprises the following steps of: coding an input code through a pre-trained CodeBERT model to obtain a code feature vector; encoding the input cue word through a pre-trained codeGPT model to obtain a cue word feature vector; after the code feature vector and the cue word feature vector are mapped to the same dimension, splicing is carried out, and a spliced feature vector is obtained; performing fine tuning on the large language model by using a low-rank adaptation method; and inputting the spliced feature vector into the fine-tuned large language model for reasoning to obtain a vulnerability detection result, namely whether the code has a corresponding vulnerability or not. According to the method, the advantages of the CodeBERT model and the codeGPT model are creatively combined, the codes and the prompt words are subjected to feature coding, and the deep structures of the codes and semantic information of the prompt words are effectively fused.

Description

Technical Field

[0001] The present invention relates to the technical field of vulnerability detection, and in particular to a vulnerability detection method based on mixed text and code pre-training encoding. Background Art

[0002] At present, with the continuous increase in the scale and complexity of software development, code vulnerability detection has become a key link in ensuring software security. Existing code vulnerability detection technologies mainly include static analysis, dynamic analysis, and hybrid analysis. Static analysis mainly identifies potential vulnerabilities through syntax parsing and symbolic execution of code text, but it is easily affected by code complexity and has false positives and false negatives. Dynamic analysis discovers vulnerabilities by actually running the program to observe its behavior. Although it can capture runtime errors well, it may not cover all potential risk scenarios due to the limited test coverage. Hybrid analysis combines the advantages of both, but the implementation cost is high and the efficiency is low.

[0003] In recent years, the successful application of deep learning technology in the field of natural language processing has inspired its exploration in the field of code analysis. In particular, the application of large language models (such as GPT, BERT, etc.) has been attempted to be used for code vulnerability detection due to its powerful language understanding and generation capabilities. However, the direct application of large language models for such tasks faces the following challenges:

[0004] Context length limit: Frequent function calls and complex nested structures in the code may cause the context length of the code and the prompt word to far exceed the input limit of the large language model. Excessively long context may cause the model performance to degrade and affect the accuracy of vulnerability detection.

[0005] Cross-language migration problem: Traditional natural language processing-based models are often trained for a single programming language and are difficult to migrate between different programming languages, limiting their widespread application in multiple programming environments. Summary of the invention

[0006] In order to solve the above technical problems, the present invention provides a vulnerability detection method based on mixed text and code pre-training encoding.

[0007] In order to solve the above technical problems, the present invention adopts the following technical solutions:

[0008] A vulnerability detection method based on mixed text and code pre-training encoding, comprising:

[0009] The input code is encoded through the pre-trained CodeBERT model to obtain the code feature vector e; the input prompt word is encoded through the pre-trained codeGPT model to obtain the prompt word feature vector EP;

[0010] After mapping the code feature vector e and the prompt word feature vector EP to the same dimension, they are concatenated to obtain a concatenated feature vector E;

[0011] The large language model is fine-tuned using a low-rank adaptation method. The concatenated feature vector E is input into the fine-tuned large language model for inference to obtain the vulnerability detection result, that is, whether there is a vulnerability in the code.

[0012] Furthermore, the pre-trained CodeBERT model is used to encode the input code to obtain the code feature vector e, which specifically includes:

[0013] Code = {c1, c2, c3...c m-1 , c m}, m is the length of the code Code, c m is the mth character of Code;

[0014] Input the code Code into the pre-trained CodeBERT model and get the code feature vector e = {e 1 , e 2 , e 3 ...e h-1 , e h}, h is the dimension of the code feature vector e, which represents the hidden layer dimension used by the CodeBERT model in the encoding process.

[0015] Furthermore, the input prompt word is encoded by the pre-trained codeGPT model to obtain the prompt word feature vector EP;

[0016] Prompt word P = {p1, p2, p3...p n -1, p n}, n is the length of the prompt word P, p n is the nth word of the prompt word P; the prompt word is processed by the Tokenizer tool, and the obtained sequence P′ consisting of integer indices is input into the pre-trained codeGPT model; the codeGPT model maps each element in the sequence P′ to the internal hidden state space and generates the corresponding hidden state vector. The set of hidden state vectors constitutes the prompt word feature vector EP = {ep1, ep2, ep3...ep n -1,ep n}, ep n is the nth hidden state vector in the prompt word feature vector EP.

[0017] Compared with the prior art, the beneficial technical effects of the present invention are:

[0018] 1. Efficient processing of long contexts: The present invention solves the problem of performance degradation of large language models caused by inputting too long contexts through code characterization and vector concatenation, ensuring high-precision detection when processing complex codes.

[0019] 2. Cross-programming language detection: The present invention has the ability to detect code vulnerabilities across programming languages, breaking the traditional model’s reliance on a single programming language and enhancing versatility and scope of application.

[0020] 3. Accurate detection and strong generalization capabilities: The present invention adopts the LoRA fine-tuning strategy to optimize the model parameters, making the model more accurate in code vulnerability detection tasks while maintaining good generalization capabilities for unknown code types and new vulnerability types. BRIEF DESCRIPTION OF THE DRAWINGS

[0021] Figure 1 It is a schematic diagram of the overall structure of the present invention. DETAILED DESCRIPTION

[0022] A preferred embodiment of the present invention is described in detail below with reference to the accompanying drawings.

[0023] When using a large language model for code vulnerability detection, the function calls and other contents of the code plus the prompt words of the large language model may cause the context length to be too long, and the excessively long context will affect the performance of the large language model, which may cause the large language model to be less effective in the code vulnerability detection task; at the same time, the large language model trained by traditional natural language processing methods can only demonstrate code vulnerability detection capabilities in one programming language and lacks the ability to migrate between programming languages.

[0024] The present invention proposes a vulnerability detection method based on mixed text and code pre-training encoding. The overall architecture of the method is as shown in the attached Figure 1 As shown in the figure, it mainly consists of two modules, the characterization module and the large model question answering module.

[0025] The input of the large language model consists of two parts: the prompt word P and the code Code. Assume that the prompt word P is {p1, p2, p3...p n -1, p n}, the code is {c1, c2, c3...c m -1, c m}, where n and m represent the length of the prompt word and code respectively.

[0026] 1. The characterization module is divided into two parts, including prompt word feature encoding and code feature encoding.

[0027] When performing deep learning processing on the code to extract its intrinsic features, the present invention selects a pre-trained model designed specifically for programming languages, the CodeBERT model. The CodeBERT model is a bidirectional encoder based on the Transformer architecture. It is pre-trained on a large-scale source code dataset and can capture the semantics, syntax and structural information of the code, providing a powerful feature representation for subsequent code understanding and generation tasks. For a given code Code, the present invention first uses the CodeBERT model to encode it. Specifically, the code Code is passed as input to the CodeBERT model, and the CodeBERT model performs context-related deep encoding on each token in the code through its multi-level self-attention mechanism. In this process, CodeBERT not only considers the meaning of each token itself, but also fully considers its positional relationship in the overall structure of the code and the interaction information with other tokens, thereby generating a feature vector obtained by encoding that can represent an in-depth and comprehensive understanding of the code. After encoding by the CodeBERT model, the present invention obtains a vector representing the features of the entire code, denoted as the code feature vector e={e 1 , e 2 , e 3 ...e h-1 , e h}, the dimension of the vector is recorded as h, where h represents the hidden layer dimension used by the CodeBERT model in the encoding process. The hidden dimension h is usually a pre-set value (such as 768, 1024, etc.), reflecting the complexity and sophistication of the model in representing code information internally. The code feature vector e is regarded as a compact and rich feature representation of the code Code, which contains many key features such as its grammatical structure, semantic logic, potential programming patterns, etc., and can be used for various downstream tasks such as subsequent code similarity calculation, code retrieval, code defect detection, etc.

[0028] At the same time, for the encoding of prompt words, the present invention chooses to use the Tokenizer tool that is compatible with the large language model (codeGPT) for processing. Tokenizer is responsible for converting the original natural language prompt word sequence into an input format that the model can understand, including word segmentation, adding special tags (such as start symbols, end symbols), and mapping words to predefined vocabulary indexes. After being processed by the Tokenizer, the prompt word is converted into a sequence P′ consisting of integer indexes. Furthermore, the present invention inputs the sequence P′ into the codeGPT model for encoding. As a generative model based on Transformer, the encoding part of the codeGPT model also has the ability to capture long-distance dependencies between sequence elements. For each element corresponding to index i (i∈I, I is the set of all indexes in sequence P) in sequence P, codeGPT maps it to the hidden state space inside the model and generates the corresponding hidden state vector. The set of these hidden state vectors constitutes the encoding result of the prompt word, which is recorded as the prompt word feature vector EP={ep1, ep2, ep3...ep n -1,ep n}. EP and P have the same set size, that is, each prompt word index i corresponds to a hidden state vector, and the dimension of each hidden state vector is k, where Here, k is the hidden dimension of the codeGPT model, which may be the same as or different from the hidden dimension h of the CodeBERT model, depending on the specific configuration of each model.

[0029] The present invention characterizes the code Code through the CodeBERT model to obtain a code feature vector e with a dimension of h, which fully reflects the deep structure and semantic information of the code. At the same time, with the help of the Tokenizer supporting the codeGPT model and the codeGPT model itself, the prompt words are effectively encoded to form a prompt word feature vector EP, in which each prompt word corresponds to a hidden state vector with a dimension of k, and the semantic features of the prompt word and its contextual relationship in the sequence are completely preserved. These two parts of work lay a solid foundation for the subsequent code and prompt word interaction analysis and code vulnerability detection tasks.

[0030] 2. Feature vector concatenation.

[0031] Since k and h are not necessarily the same, a linear layer is used to map e to ep before merging. i The same dimension, that is, ec = {ec1, ec2, ec3...ec k-1 ,ec k}, and then insert ec into the specified position in EP to obtain the final concatenated feature vector E = {ep1, ep2, ep3, ..., ec, ..., ep n-1 ,ep n}.

[0032] 3. Use large language models for reasoning.

[0033] The large language model needs to be fine-tuned. At this stage, efficient parameter update strategies such as LoRA (Low-Rank Adaptation) can be used for training. LoRA is a lightweight fine-tuning method that is particularly suitable for large pre-trained models. While maintaining the original knowledge of the model, only a few parameters are adjusted in a low-rank approximate manner to improve the performance of the model on specific tasks in a targeted manner. Through LoRA fine-tuning, the large language model can more accurately capture the features required for code detection tasks, thereby improving the accuracy and reliability of code vulnerability detection. The fine-tuned large language model infers based on E to obtain the final conclusion, that is, whether the code has the corresponding vulnerability.

[0034] It is obvious to those skilled in the art that the present invention is not limited to the details of the exemplary embodiments described above, and that the present invention can be implemented in other specific forms without departing from the spirit or essential features of the present invention. Therefore, from any point of view, the embodiments should be regarded as exemplary and non-limiting, and the scope of the present invention is defined by the appended claims rather than the above description, and it is intended that all changes falling within the meaning and scope of the equivalent elements of the claims are included in the present invention, and any reference numerals in the claims should not be regarded as limiting the claims involved.

[0035] In addition, it should be understood that although this specification is described in terms of implementation methods, not every implementation method contains only one independent technical solution. This narrative method of the specification is only for the sake of clarity. Those skilled in the art should regard the specification as a whole. The technical solutions in each embodiment may also be appropriately combined to form other implementation methods that can be understood by those skilled in the art.

Claims

1. A vulnerability detection method based on mixed text and code pre-training encoding, characterized in that: include: The input code is encoded through the pre-trained CodeBERT model to obtain the code feature vector e; the input prompt word is encoded through the pre-trained codeGPT model to obtain the prompt word feature vector EP; After mapping the code feature vector e and the prompt word feature vector EP to the same dimension, they are concatenated to obtain a concatenated feature vector E; The large language model is fine-tuned using a low-rank adaptation method. The concatenated feature vector E is input into the fine-tuned large language model for inference to obtain the vulnerability detection result, that is, whether there is a vulnerability in the code.

2. A vulnerability detection method based on mixed text and code pre-training encoding according to claim 1, characterized in that: The input code is encoded by the pre-trained CodeBERT model to obtain the code feature vector e, which specifically includes: Code = {c1,c2,c3...c m-1 ,c m }, m is the length of the code Code, c m is the mth character of Code; Input the code Code into the pre-trained CodeBERT model and get the code feature vector e = {e 1 ,e 2 ,e 3 ...e h -1 ,e h }, h is the dimension of the code feature vector e, which represents the hidden layer dimension used by the CodeBERT model in the encoding process.

3. A vulnerability detection method based on mixed text and code pre-training encoding according to claim 1, characterized in that: The pre-trained codeGPT model is used to encode the input prompt word to obtain the prompt word feature vector EP; Prompt word P = {p1, p2, p3...p n-1 ,p n }, n is the length of the prompt word P, p n is the nth word of the prompt word P; The prompt word is processed by the Tokenizer tool, and the resulting sequence P consisting of integer indices is ′ Input to the pre-trained codeGPT model; the codeGPT model converts the sequence P ′ Each element in is mapped to the internal hidden state space and generates the corresponding hidden state vector. The set of hidden state vectors constitutes the prompt word feature vector EP = {ep1, ep2, ep3...ep n-1 ,ep n }, ep n is the nth hidden state vector in the prompt word feature vector EP.

Images