A code pre-training model training method based on program dependency graph prediction

Abstract

The application discloses a code pre-training model training method based on program dependency graph prediction, which predicts the program dependency graph of code in the pre-training stage, helps the pre-training model to perceive and capture the semantic structure features of the code, and improves the performance of downstream tasks sensitive to the semantic structure of the code. The method first extracts the program dependency graph of the code by using a static program analysis tool and constructs pre-training supervision information according to the program dependency graph, then encodes the code into a real value feature vector sequence by using a neural network encoder, and predicts the program dependency graph of the code based on the feature vector sequence, and finally updates the weight of the model according to the difference between the prediction result and the actual program dependency graph. The pre-trained model can be directly used for program dependency analysis, and the massive code can be quickly analyzed by using a parallel processor such as a GPU, and the pre-trained model can also be used for downstream tasks sensitive to the semantic structure of the code after fine tuning, such as vulnerability analysis related tasks and the like.

Description

Technical Field

[0001] This invention relates to the fields of deep learning and software engineering, and in particular to a method for training a code pre-trained model based on program dependency graph prediction. Background Technology

[0002] In recent years, pre-trained models have significantly improved the performance of various Natural Language Processing (NLP) tasks, exhibiting strong generalization and versatility. Pre-trained models typically learn general knowledge on massive amounts of unsupervised data using multiple pre-training tasks, and then fine-tune them using supervised data from downstream tasks to quickly adapt them to those tasks. Given the superior performance of pre-trained models on NLP tasks, a number of models pre-trained on massive amounts of code data have been proposed, such as CodeBERT and GraphCodeBERT, and have been used for various code-related tasks, such as code search, code clone detection, code completion, and code summarization. These models are called code pre-trained models.

[0003] The performance of a pre-trained model is closely related to the pre-training task. An effective pre-training task can guide the model to learn general knowledge from massive amounts of data that is helpful for downstream tasks. Some code pre-training models directly adopt pre-training tasks designed for natural language processing, such as Replacement Word Prediction (RTD). However, unlike natural language, code has a strong structure. Therefore, some code pre-training models have proposed pre-training tasks that consider code structure information to improve model performance, such as Fine-Grained Type Inference, AST node type prediction, and data flow edge prediction. However, these pre-training models either only consider the syntactic structure of the code, ignoring its semantic structure (e.g., control dependencies and data dependencies), thus limiting their performance on downstream tasks; or they require explicit input of structural information extracted by code analysis tools, rather than just source code, limiting their generality and making them unsuitable for analyzing code structure information. Currently, there is no code pre-training model that considers the semantic structure of the code during the pre-training stage and only uses source code as input. This invention, by learning the semantic structure knowledge of code, aims to improve the performance of pre-trained models for downstream tasks sensitive to code semantic structure. Summary of the Invention

[0004] The purpose of this invention is to address the shortcomings of existing technologies by providing a code pre-training model training method based on program dependency graph prediction.

[0005] The objective of this invention is achieved through the following technical solution: a method for training a code pre-training model based on program dependency graph prediction, comprising the following steps:

[0006] (1) Extracting the program dependency graph of code data as pre-training supervision information, wherein the program dependency graph includes a control dependency graph and a data dependency graph;

[0007] (2) The code is segmented into a symbol sequence T by a word segmenter, where T = [t cls ,t1,...,t n ], t cls Let n be the symbol inserted at the very beginning of the sequence, and n be the length of the symbol sequence T.

[0008] (3) Based on the symbol sequence T obtained in step (2), the program dependency graph is converted into a statement-level control dependency graph and a symbol-level data dependency graph;

[0009] (4) Encode the segmented symbol sequence T into a symbol feature vector sequence H using an encoder. t ,in The symbol t inserted at the very beginning of the sequence cls The corresponding vector, where n is the symbol sequence H t Length;

[0010] (5) Input the symbolic feature vector sequence into the statement-level control dependency graph prediction task, the symbolic data dependency graph prediction task, and the masked language model task, and summarize the loss values ​​of each task to obtain the overall goal of the pre-training stage.

[0011] (6) Use the optimizer to optimize the overall objective. And update the parameters of the encoder and neural network components.

[0012] Further, step (1) specifically involves: using a static program analysis tool to analyze the code data used for pre-training and obtaining the program dependency graph of each code segment; obtaining the control dependency graph of the code by extracting all AST nodes and edges indicating control dependencies in the program dependency graph; and obtaining the data dependency graph of the code by extracting all AST nodes and edges indicating data dependencies in the program dependency graph.

[0013] Furthermore, step (3) includes the following sub-steps:

[0014] (3.1) For the control dependency graph, for each edge in the control dependency graph output in step (1), extract the specific positions of the code elements corresponding to the start and end AST nodes. Then, use the code line number corresponding to the AST node as the code statement number to map the AST node to the corresponding code statement. This transforms the control dependency relationship of the original AST node into the control dependency relationship between code statements, thereby constructing a statement-level control dependency graph and storing it as a control dependency graph matrix.

[0015] (3.2) For the data dependency graph, for each edge in the data dependency graph output in step (1), extract the identifiers corresponding to its start and end AST nodes; then, by checking the overlap between the code character range corresponding to the identifier and the code character range corresponding to the tokenized symbol, map the AST node to the symbol in the symbol sequence T; finally, based on the above data dependency relationships between symbols, construct a symbol-level data dependency graph and store it as a data dependency graph matrix.

[0016] Furthermore, in step (5), the overall goal of pre-training is... From the symbolic eigenvector sequence H t The calculated loss is defined as a weighted sum of the losses from the statement-level control dependency graph prediction task, the symbol-level data dependency graph prediction task, and the masked language model task, and is expressed as follows:

[0017]

[0018] in and denoted as a1, a2, and a3 respectively, representing the loss values ​​generated by the statement-level control dependency graph prediction task, the symbolic data dependency graph prediction task, and the masked language model task, respectively. The weights in the equation.

[0019] Furthermore, step (5) includes the following sub-steps:

[0020] (5.1) Calculation of loss value for statement-level control dependency graph prediction task: First, based on the code symbol feature vector sequence H t Calculate each node q in the statement-level control dependency graph k eigenvectors Then, bilinear layer prediction is used for node q. i Control depends on another node q j probability Finally, the cross-entropy loss function is used to measure the predicted probability. With statement-level control dependency graph matrix G c The label value at the corresponding position in the middle The differences are analyzed, and a statement-level code control dependency graph is generated to predict the loss value of the task. The following expression exists:

[0021]

[0022] (5.2) Calculation of loss value for symbol-level data dependency graph prediction task: First, based on the code symbol feature vector sequence H tCompute each node t in the symbolic data dependency graph i Features Then, bilinear layer prediction symbol t is used. i The data depends on another symbol t j probability Finally, the masked cross-entropy loss function is used to measure the predicted probability. Label values ​​at corresponding positions in the symbol-level data dependency graph The differences are analyzed, and a symbolic data dependency graph prediction task loss value is generated. The following expression exists:

[0023]

[0024] Where each symbol t i Corresponding mask value The mask value m of code symbols belonging to the identifier type is calculated based on the type analysis results of the parser. i d =1, and the mask value m of symbols that do not belong to the identifier type. i d =0;

[0025] (5.3) Calculation of loss value for masked language model task: First, some symbols in the input are randomly replaced with special mask symbols [MASK]. Then, the encoder encodes the input of random mask into a feature vector sequence. Finally, based on the feature vectors at the mask positions, the masked symbols in the input are predicted. A random sampling ratio of r is taken from the segmented symbol sequence T. sample The symbols, among these extracted symbols, r mask The % symbol is replaced with the [MASK] symbol, r replace The % symbol is replaced with a random symbol from the vocabulary; then the masked symbol sequence is input into the encoder to obtain the symbol feature vector sequence H. t Based on the symbolic feature vector sequence H t A multilayer perceptron model is used to predict the original symbol before the mask, and the cross-entropy loss is calculated based on the difference between the predicted result and the original symbol as the loss value for the masked language model task. The following expression exists:

[0026]

[0027]

[0028] Where V is the vocabulary of the model. The symbol t i The probability of predicting correctly, M tThis represents the set of indices that are randomly sampled.

[0029] The beneficial effects of this invention are as follows: During the pre-training stage, this invention guides the model to learn knowledge related to code semantic structure by explicitly predicting the program dependency graph of the code, thereby improving the performance of downstream tasks that are sensitive to code semantic structure. Simultaneously, when transferring this pre-trained model to downstream tasks, it does not require code analysis tools to extract code structural information, avoiding the impact of the generality issues of code analysis tools on the generality of the pre-trained model. After pre-training is complete, the model can not only be fine-tuned to adapt to downstream tasks, but can also be directly used as a high-efficiency, neural network-based program dependency analysis tool, relying on parallel processors such as GPUs to quickly perform program dependency analysis on massive amounts of code. Attached Figure Description

[0030] Figure 1 This is a flowchart of the method of the present invention;

[0031] Figure 2 This is an example diagram illustrating how the program dependency graph of the code is linked to the code symbol sequence obtained from word segmentation in this invention;

[0032] Figure 3 This is a schematic diagram of the pre-training task used in this invention;

[0033] Figure 4 This is an example diagram of the symbol mask values ​​generated by the parser in this invention to alleviate the sparsity problem of the symbol-level data dependency matrix;

[0034] Figure 5 This is a table showing the results of the intrinsic evaluation in the experimental evaluation of this invention;

[0035] Figure 6 This is a table showing the results of the external evaluation in the experimental evaluation of this invention. Detailed Implementation

[0036] This invention extracts the program dependency graph from the code using static program analysis tools, predicts the program dependency graph during the pre-training phase, calculates the difference between the output program dependency graph and the actual program dependency graph, obtains the loss value for the program dependency graph prediction task, and finally summarizes the loss values ​​of all pre-training tasks to obtain the total pre-training loss value. (i.e., the overall goal) And use the optimizer to minimize the total pre-training loss value. To update the model parameters.

[0037] A program dependency graph displays the control and data dependencies between code elements, reflecting the semantic structure of the code. It contains control dependency edges and data dependency edges. Control dependency edges show the control flow relationship between conditional assertions and code statements, and can be used to infer the control flow conditions for code statement execution; data dependency edges show the definition-use relationships in the code.

[0038] The present invention will be further described below with reference to the accompanying drawings and examples.

[0039] This invention provides a code pre-training model training method based on program dependency graph prediction, see [link to relevant documentation]. Figure 1 This includes the following steps:

[0040] (1) Extracting the Program Dependence Graph from the Code Data as Supervision Information for the Program Dependence Graph Prediction Task in the Pre-training Stage. This invention first utilizes a static program analysis tool to extract the program dependency graph from the code in the pre-training code dataset as labels, providing supervision information for the program dependency graph prediction task in the pre-training stage. The output of the program analysis tool is a graph, where the nodes are Abstract Syntax Tree (AST) nodes, and the edges in the graph have different types. By extracting all nodes and control dependency edges from the program dependency graph, the control dependency graph defined on the AST nodes can be obtained. Since data dependencies refer to the existence of definition-reference (def-use) relationships between variables, by extracting all nodes and merging the edges of type "def" and "use" in the graph, the data dependency graph defined on the AST nodes can be obtained.

[0041] (2) Code segmentation. The code input to the model is essentially a piece of text; in order for the neural network to process the code, it first needs to be converted into a sequence of symbols. This invention uses the BPE (Byte-Pair Encoding) algorithm to construct the vocabulary and segmenter of the pre-trained model, and uses this segmenter to convert the input code into a sequence of symbols to eliminate the "out-of-vocabulary problem," that is, to eliminate symbols that exist in the test data but not in the training data. The symbol sequence obtained after segmentation is defined as T = [t cls ,t1,...,t n ], where t cls A symbol is a special symbol inserted at the very beginning of the sequence. Its feature vector, obtained after processing by the model, is often used to represent the entire code, where n is the length of the code symbol sequence. Furthermore, since word segmentation does not remove characters from the code, each symbol corresponds to a character range within the code.

[0042] (3) Construct the program dependency graph matrix. See [link / reference] Figure 3 This invention uses the prediction of the program dependency graph based on the symbol sequence T as a pre-training task. However, the program dependency graph output in step (1) is defined on the AST node and is not associated with the symbol sequence T input to the neural network model. Therefore, this step needs to reconstruct the control dependency graph and the data dependency graph based on the symbol sequence output in step (2) based on the output of step (1). This step is mainly divided into the following two sub-steps:

[0043] (3.1) For the control dependency graph, since the control dependencies of the code are usually defined between code statements, the nodes q in the reconstructed control dependency graph will be... k Defined as code statements. First, for each edge in the control dependency graph output in step (1), extract the specific positions (code lines and corresponding character ranges) of the code elements corresponding to the start and end AST nodes. Then, map the AST nodes to the corresponding code statements (code lines) using the code line number corresponding to the AST node as the code statement sequence number, thereby transforming the control dependency relationship of the original AST nodes into the control dependency relationship between code statements, and thus constructing a statement-level control dependency graph, which is stored as a control dependency graph matrix. in This indicates that there is a control dependency between the i-th code statement and the j-th code statement. This indicates that there is no control dependency, m c To control the number of nodes in the dependency graph, finally, G... c As can be seen from the definition, each row and each column This represents the control dependency between a code statement and other statements; therefore, for each symbol t in the symbol sequence T... i By examining the line of code in which it appears, it is mapped to a specific statement, thus establishing a relationship between the symbol sequence T and the control dependency graph matrix G. c The connection.

[0044] (3.2) For data dependency graphs, since data dependencies usually exist between variables, defining data dependencies directly between code statements is too coarse-grained, which may cause the model to fail to accurately capture and learn data dependencies in the code. Therefore, this invention defines code data dependencies at the symbol level and constructs a symbol-level data dependency graph defined on a sequence of symbols. Specifically, firstly, for each edge in the data dependency graph output in step (1), the identifiers corresponding to its start and end AST nodes are extracted. Then, by checking the overlap between the code character range corresponding to the identifier and the code character range corresponding to the segmented symbols, the AST nodes are mapped to symbols in the symbol sequence T. Since the BPE algorithm may segment an identifier into multiple symbols, the AST node corresponding to an identifier may be mapped to multiple (a group of) symbols. Figure 2 For example, the first "temp_flag" identifier in the code corresponds to the character range [4, 12]. By checking the overlap of the character ranges, it is mapped to three characters, i.e. The characters "_" and "flag" correspond to the code ranges [3,7], [8,8], and [9,12], respectively. Therefore, the data dependencies between AST nodes are transformed into data dependencies between one set of symbols and another set of symbols. To avoid generating a large number of meaningless data dependency edges in this "many-to-many" relationship, this invention only retains the first of the two sets of symbols corresponding to the starting and ending points of the data dependency edges, discarding the others. (Continuing from the foregoing...) Figure 2 In the example, the data dependency between the two "temp_flag" identifiers is simplified to the first symbol of the two identifiers. The data dependencies between the symbols (with indices 1 and 13 respectively) are established. Finally, based on the above data dependencies between symbols, a symbol-level data dependency graph is constructed and stored as a matrix. in Represents t i With t j There are data dependencies between them, and vice versa. The m represents non-existence. d is the length of the symbol sequence.

[0045] (4) The symbol sequence is encoded into a symbol feature vector sequence using an encoder. After word segmentation, the input code is converted into a symbol sequence T, which is then fed into the neural network encoder to obtain the feature vector of each symbol. Subsequent pre-training tasks are performed based on the symbol feature vector sequence. The encoder generally consists of an embedding layer and a feature extraction component. The embedding layer is responsible for mapping symbols to word embeddings (real-valued vectors) based on the vocabulary, while the feature extraction component performs information interaction between symbols based on the word embedding sequence to capture the feature vector of each symbol. Specifically, the symbol sequence T = [tcls ,t1,...,t n The sequence of symbolic feature vectors obtained after inputting into the encoder is denoted as... Each symbol is represented as a vector. d is the dimension of the feature vector. The symbol t inserted at the very beginning of the sequence cls The corresponding vector, where n is the symbol sequence H t The length of the encoder is not specified. This invention does not impose specific limitations on the structure of the encoder, but the encoder must at least be able to convert the symbol sequence T of code into a symbol feature vector sequence H. t The function.

[0046] (5) Based on the symbolic feature vector sequence, perform multiple pre-training tasks, including program dependency graph prediction, calculate the loss value of each pre-training task, and summarize them to obtain the total pre-training loss value (the overall goal of the pre-training stage). This step inputs the symbolic feature vector sequence into each pre-training task to generate the loss value of each pre-training task. The present invention uses three pre-training tasks: statement-level control dependency graph prediction, symbolic-level data dependency graph prediction, and masked language modeling. The loss values ​​generated by these three pre-training tasks are denoted as follows: and The total pre-training loss value is denoted as The following expression exists:

[0047]

[0048] Here, a1, a2, and a3 are hyperparameters used to adjust the weights of the loss values ​​for the statement-level control dependency graph prediction task, the symbol-level data dependency graph prediction task, and the masked language model task in the final pre-training loss value, respectively. The values ​​of this set of hyperparameters are selected using the validation set. The specific steps for calculating the loss values ​​for the three pre-training tasks are as follows:

[0049] (5.1) For the statement-level control dependency graph prediction task, it is first necessary to determine the symbol feature vector sequence H. t A feature representation of the code statements is constructed to predict control dependencies between statements. As mentioned in step (3.1), each symbol is mapped to a node q in the statement-level control dependency graph by retrieving its corresponding line of code. k Therefore, a node typically contains multiple symbols. This step involves retrieving and aggregating each node q. k The node (code statement) is constructed using the feature vectors of all symbols contained within it. k eigenvectors

[0050]

[0051]

[0052] Where |q k | is node q k The number of symbols contained in the code statements; MLP is a multilayer perceptron. MLP is mainly used for feature dimensionality reduction because the dimension of the symbol feature vector output by the encoder is usually very high (e.g., 768 in the code pre-trained model CodeBERT), therefore, it can be directly used... As q k The feature vectors will consume a significant amount of memory. Furthermore, although the formula uses the average of the feature vectors of all symbols within a node as the node's feature vector, this invention does not impose limitations here; any aggregation function for ordered or unordered sets can be used. After obtaining the feature vectors of the code statement, this invention utilizes a bilinear layer to predict the node (code statement) q. i Control depends on another node (code statement) q j probability

[0053]

[0054] Among them W c and b c The parameters are learnable, and σ is the sigmoid function. A bilinear layer is used here because it effectively models the relationship between the two inputs, and the two inputs of a bilinear layer do not satisfy the commutative law, which aligns with the directedness of the control dependency in the code. Finally, cross-entropy loss is used to measure the predicted probability. With the control dependency graph matrix G c The corresponding tags The differences were analyzed, and the results were averaged to generate a loss value.

[0055]

[0056] (5.2) For symbolic data dependency graph prediction tasks, since H t Since the sequence of feature vectors is already at the symbol level, there is no need to construct additional feature representations for the symbols as in step (5.1). Similarly, an MLP is used here to reduce the dimensionality of the symbol features, and then a bilinear layer is used to predict a symbol t based on the reduced symbol representation. i The data depends on another symbol t j probability value The following expression exists:

[0057]

[0058]

[0059] Among them W d and b d Here, σ represents the learnable parameters of the bilinear layer, and σ is the Sigmoid function. However, unlike other codebases, the number of symbols in a codebase is generally much greater than the number of statements. This results in a much larger number of nodes in the symbol-level data dependency graph compared to the statement-level control dependency graph, leading to a larger matrix to be predicted. For example, for a codebase containing 512 symbols and 30 statements, the statement-level control dependency graph to be predicted contains 30 × 30 = 900 elements, while the symbol-level data dependency graph contains 512 × 512 = 262144 elements. Worse still, many codes contain only a few dozen or even just a few data dependencies, resulting in a highly sparse symbol-level data dependency graph that is difficult to learn directly. To address this, this invention designs a symbol masking method to mask nodes in the symbol-level data dependency graph that are unlikely to have data dependencies, thus avoiding the prediction of symbol-level data dependencies at these locations and alleviating the data sparsity problem. Specifically, this invention first uses a program lexical analyzer to segment the code into a sequence of lexical symbols, and the lexical analyzer labels the type of each lexical symbol; then, it selects symbols of type identifier from this sequence to form a set of lexical symbols S'; finally, it checks the symbols t in the symbol sequence T based on the overlap of character intervals. i Whether a symbol t belongs to any lexical symbol in the set of lexical symbols S' is determined by the symbol t. i Generate a mask value m i If the symbol t i If it belongs to an identifier, then m i =1, otherwise m i =0. This invention uses cross-entropy loss to measure the predicted probability value. The label value at the corresponding position in the actual symbolic data dependency graph. The difference is calculated by averaging the values ​​of the unmasked elements to determine the symbolic data dependency loss. When, the following expression exists:

[0060]

[0061] Specifically, if the mask value of any symbol in a pair is 0 (that symbol is not an identifier), then the prediction loss value for the data dependency between this pair of symbols will not be included in the calculation. This reduces the size of the matrix that needs to be predicted. For example... Figure 4 As shown, this symbol masking method reduces the number of matrix elements that need to be predicted by 70%, greatly alleviating the problem of matrix sparsity.

[0062] (5.3) The masked language model task was first proposed by the BERT model and has since been widely used in pre-trained models. This task first randomly replaces some symbols in the input with special mask symbols [MASK]. Then, an encoder encodes the random masked input into a sequence of feature vectors. Finally, based on the feature vectors at the mask positions, the masked symbols in the input are predicted. This invention follows the settings of the masked language model in existing pre-trained models. Specifically, after word segmentation, this task randomly selects a proportion r from the segmented symbol sequence T. sample The symbols. Among these extracted symbols, r mask The % symbol is replaced with the [MASK] symbol, r replace The % symbol is replaced with a random symbol from the vocabulary. The masked symbol sequence is then input into the encoder to obtain the symbol feature vector sequence H. t Subsequently, this task is based on the symbolic feature vector sequence H t A multilayer perceptron model is used to predict the original symbol before the mask, and the cross-entropy loss is calculated based on the difference between the predicted result and the original symbol as the loss value for the masked language model task.

[0063]

[0064]

[0065] Where V is the vocabulary of the model. The symbol t i The probability of predicting correctly, M t This represents the set of indices that are randomly sampled.

[0066] (6) Update model parameters using the optimizer based on the total pre-training loss value. After obtaining the total loss value in step (5), update the parameters of all models end-to-end using the optimizer based on this loss value, including the parameters of the encoder model and the parameters of the neural network components (e.g., bilinear layers) of the three pre-trained tasks. After pre-training, only the encoder model will be transferred to the downstream task as a pre-trained model, while the task structure of the pre-trained task (e.g., bilinear layers) will not be transferred.

[0067] Experimental verification:

[0068] This invention evaluates the effectiveness of the proposed training method in two ways: first, it verifies the performance of the pre-trained model in predicting program dependency graphs, which is called intrinsic evaluation; then, it verifies the performance of the pre-trained model after fine-tuning on code-related downstream tasks, which is called extrinsic evaluation.

[0069] The specific experimental methods and settings for pre-training, internal evaluation, and external evaluation are described below.

[0070] Pre-training: The pre-training dataset used in the experiment contains over 2,280,000 C / C++ functions collected from open-source projects. This dataset is divided into three parts: a training set, a validation set, and a test set, containing approximately 1,900,000, 155,000, and 60,400 functions respectively. The training set is used for model pre-training, as described in steps (1)-(6) of the method described above; the validation set is used to monitor the model's generalization ability during training; and the test set is used to evaluate the model's performance in predicting program dependency graphs in the internal evaluation. In the evaluation experiment, the pre-trained model proposed in this invention adopts the model structure, tokenizer, and dictionary of the CodeBERT model, and initializes the model using CodeBERT weights. During pre-training, this invention uses Adam as the optimizer, with an initial learning rate set to 0.0001, and requires specifying three parameters: the maximum number of lines of code (the maximum number of nodes in the statement-level control dependency graph) m. c The maximum length m of the symbol sequence after word segmentation. d And the number of training cycles. In the experiment involved in this invention, m is set... c =50, m d =512, with 10 training cycles.

[0071] Internal evaluation: Experiments were conducted using a test set partitioned during the pre-training phase. The evaluation metric was the F-metric of the control dependencies. cdp F-metric value of data dependency edges ddp The F-metric value of the entire process and all dependent edges. overall The F-metric value is based on the formula:

[0072]

[0073]

[0074] Where Precision is the true measure of a statement and Recall is the true measure of a statement. For statement-level control dependency graph prediction tasks: TP (True Positive Instance) represents the number of statement pairs where the true label is 1 and the prediction result is also 1; FP (False Positive Instance) represents the number of statement pairs where the true label is 0 but the prediction result is 1; TN (True Negative Instance) represents the number of statement pairs where the true label is 0 and the prediction result is also 0; FN (False Negative Instance) represents the number of statement pairs where the true label is 1 but the prediction result is 1. For symbolic data dependency graph prediction tasks, TP (True Positive Instance) represents the number of symbol pairs where the true label is 1 and the prediction result is also 1; FP (False Positive Instance) represents the number of symbol pairs where the true label is 0 but the prediction result is 1; TN (True Negative Instance) represents the number of symbol pairs where the true label is 0 and the prediction result is also 0; FN (False Negative Instance) represents the number of symbol pairs where the true label is 1 but the prediction result is 1. For all program dependency edges as a whole, its TP is the sum of the TP values ​​of the statement-level control dependency graph prediction task and the symbol-level data dependency graph prediction task. The calculation methods for its TN, FP, and FN metrics are similar. The F metric ranges from 0 to 1; a larger value indicates better performance.

[0075] Figure 5 This is a table showing the results of the internal evaluation experiments of this invention. The experimental results show that the code model pre-trained using the training method proposed in this invention can achieve an F-metric score of over 99% when predicting statement-level control dependency graphs, and over 93% when predicting symbol-level data dependency graphs. Furthermore, the overall F-metric score for the program dependency graph exceeds 95%, indicating that the pre-training method proposed in this invention can predict the program dependency graph of the code relatively accurately.

[0076] External Evaluation: The external evaluation experiments involved in this invention selected vulnerability detection as the downstream task and used three existing publicly available benchmark datasets for validation: ReVeal, Devign, and Big-Vul, containing approximately 22,600, 272,000, and 188,600 C / C++ functions, respectively, with the proportion of vulnerable functions being 9.7%, 5.8%, and 45.6%, respectively. After pre-training, the model needs to be fine-tuned on the vulnerability detection dataset to adapt to the vulnerability detection task. The specific method of fine-tuning is to concatenate an MLP classifier after the code pre-trained model. This classifier will be represented by the first symbol output by the model (i.e., The model is used as input to predict the probability of a vulnerability in the code. Samples with a predicted probability greater than 0.5 are predicted as functions containing vulnerabilities. Then, the cross-entropy loss function is used to measure the difference between the output probability and the label value indicating whether the code actually contains a vulnerability, generating a fine-tuning loss. Finally, the Adam optimizer is used to jointly update the parameters of the pre-trained model and the classifier. The external evaluation metrics are vulnerability classification accuracy and the F-metric. The F-metric is described in the internal evaluation, where: TP (True Positives) represents the number of functions with vulnerabilities correctly predicted; FP (False Positives) represents the number of functions without vulnerabilities incorrectly predicted as having vulnerabilities; TN (True Negatives) represents the number of functions without vulnerabilities correctly predicted as not having vulnerabilities; and FN (False Negatives) represents the number of functions with vulnerabilities incorrectly predicted as not having vulnerabilities. Accuracy is based on the following formula:

[0077]

[0078] The code model pre-trained using the training method and pre-training settings proposed in this invention is called PDBERT because its pre-training task involves not only program dependency prediction but also the masked language model task proposed by the BERT model. To further verify the effectiveness of this invention, PDBERT is compared with mainstream vulnerability detection algorithms and code pre-trained models, including: five vulnerability detection models not based on pre-training, namely Bi-LSTM, Transformer, VulDeePecker, Devign, and ReVeal; and three code pre-trained models, namely CodeBERT, GraphCodeBERT, and DISCO. Figure 6 This is a table showing the external evaluation results of the present invention. Experimental results show that, compared to five untrained vulnerability detection models, PDBERT achieves the best performance after fine-tuning, with average performance improvements of 80.0%, 197.1%, and 24.7% on the corresponding metrics of the three benchmark datasets, respectively. Compared to three code-pretrained models, PDBERT achieves the best performance after fine-tuning, with average improvements of 7.0%, 9.5%, and 6.8% on the corresponding metrics of the three benchmark datasets, respectively. It can be seen that PDBERT significantly outperforms the other comparative models on all three vulnerability detection datasets, indicating that the pre-training method based on program dependency graph prediction proposed in this invention is effective.

[0079] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and are not intended to limit it. Although the present invention has been described in detail with reference to preferred embodiments, those skilled in the art should understand that modifications or equivalent substitutions can be made to the technical solutions of the present invention without departing from the spirit and scope of the technical solutions of the present invention, and all such modifications or substitutions should be covered within the scope of the claims of the present invention.

Claims

1. A method for training a code pre-training model based on program dependency graph prediction, characterized in that, Includes the following steps: (1) Extract the program dependency graph of the code data as pre-training supervision information, wherein the program dependency graph includes the control dependency graph and the data dependency graph; (2) The code is segmented into a sequence of symbols using a tokenizer. ,in , The symbol inserted at the very beginning of the sequence. For symbolic sequence Length; (3) Based on the symbol sequence obtained in step (2) Transform the program dependency graph into a statement-level control dependency graph and a symbol-level data dependency graph; (4) The token sequence after word segmentation is processed by the encoder. Encoded as a sequence of symbolic feature vectors ,in , The symbol inserted at the very beginning of the sequence The corresponding vector, For symbolic sequence Length; (5) Input the symbolic feature vector sequence into the statement-level control dependency graph prediction task, the symbolic data dependency graph prediction task, and the masked language model task, and summarize the loss values ​​of each task to obtain the overall goal of the pre-training stage. ; Pre-training overall goal From the sequence of symbolic eigenvectors The calculated loss is defined as a weighted sum of the losses from the statement-level control dependency graph prediction task, the symbol-level data dependency graph prediction task, and the masked language model task, and is expressed as follows: ； in , and These represent the loss values ​​generated by the statement-level control dependency graph prediction task, the symbol-level data dependency graph prediction task, and the masked language model task, respectively. , and The loss values ​​of the statement-level control dependency graph prediction task, the symbol-level data dependency graph prediction task, and the masked language model task are respectively in the overall objective. Weights in; (5.1) Calculation of loss value for statement-level control dependency graph prediction task: First, based on the code symbol feature vector sequence Calculate each node in the statement-level control dependency graph eigenvectors Then, a bilinear layer is used to predict nodes. Control depends on another node probability Finally, the cross-entropy loss function is used to measure the predicted probability. With statement-level control dependency graph matrix The label value at the corresponding position in the middle The differences are analyzed, and a statement-level code control dependency graph is generated to predict the loss value of the task. The expression is as follows: ； (5.2) Calculation of loss value for symbol-level data dependency graph prediction task: First, based on the code symbol feature vector sequence Compute each node in the symbolic data dependency graph Features Then, bilinear layer prediction symbols are used. Data depends on another symbol probability Finally, the masked cross-entropy loss function is used to measure the predicted probability. Label values ​​at corresponding positions in the symbol-level data dependency graph The differences are analyzed, and a symbolic data dependency graph prediction task loss value is generated. The expression is as follows: ； Each symbol Corresponding mask value The mask value of code symbols belonging to the identifier type is calculated based on the type analysis results of the parser. The mask value of a symbol that does not belong to the identifier type. ; (5.3) Calculation of loss value for masked language model task: First, some symbols in the input are randomly replaced with special mask symbols [MASK]. Then, the encoder encodes the input of the random mask into a feature vector sequence. Finally, based on the feature vectors at the mask positions, the masked symbols in the input are predicted; from the symbol sequence after word segmentation... The random sampling ratio is The symbols, among these extracted symbols, The % symbol has been replaced with the [MASK] symbol. The % symbol is replaced with a random symbol from the vocabulary; then the masked symbol sequence is input into the encoder to obtain the symbol feature vector sequence. Based on symbolic feature vector sequence A multilayer perceptron model is used to predict the original symbol before the mask, and the cross-entropy loss is calculated based on the difference between the predicted result and the original symbol as the loss value for the masked language model task. The expression is as follows: ； ； in, For the model's vocabulary, Symbols The probability of predicting correctly. Represents the set of indices to be randomly sampled; (6) Use the optimizer to optimize the overall objective. And update the parameters of the encoder and neural network components.

2. The code pre-training model training method based on program dependency graph prediction according to claim 1, characterized in that, The specific steps (1) are as follows: use static program analysis tools to analyze the code data used for pre-training and obtain the program dependency graph of each code segment; obtain the control dependency graph of the code by extracting all AST nodes and edges indicating control dependencies in the program dependency graph, and then obtain the data dependency graph of the code by extracting all AST nodes and edges indicating data dependencies in the program dependency graph.

3. The code pre-training model training method based on program dependency graph prediction according to claim 2, characterized in that, Step (3) includes the following sub-steps: (3.1) For the control dependency graph, for each edge in the control dependency graph output in step (1), extract the specific position of the code element corresponding to its start and end AST nodes. Then, use the code line number corresponding to the AST node as the code statement number to map the AST node to the corresponding code statement, thereby transforming the control dependency relationship of the original AST node into the control dependency relationship between code statements, and thus constructing a statement-level control dependency graph and storing it as a control dependency graph matrix. ; (3.2) For the data dependency graph, for each edge in the data dependency graph output in step (1), extract the identifiers corresponding to its start and end AST nodes; then, by checking the overlap between the code character range corresponding to the identifier and the code character range corresponding to the tokenized symbols, map the AST nodes to symbol sequences. The symbols in the diagram are then analyzed; finally, based on the data dependencies between the symbols, a symbol-level data dependency graph is constructed and stored as a data dependency graph matrix. .

Images