Vulnerability repair model construction method based on abstract syntax tree position coding

Through the positional encoding method based on the abstract syntax tree, the Clang word segmentation position is mapped to the BPE word segmentation position, and the rotational position encoding RoPE is expanded, which solves the problem of insufficient token position recognition in the Transformer model and improves the code structure understanding and repair accuracy of the vulnerability repair model.

CN120688058APending Publication Date: 2025-09-23LIAONING UNIVERSITY
View PDF 0 Cites 1 Cited by

Patent Information

Application Number
CN202510783424.8
Authority / Receiving Office
CN · China
Patent Type
Applications(China)
Current Assignee / Owner
Filing Date
2025-06-12
Publication Date
2025-09-23

Smart Images

  • Figure BDA0005446409280000031
    Figure BDA0005446409280000031
  • Figure BDA0005446409280000032
    Figure BDA0005446409280000032
  • Figure BDA0005446409280000041
    Figure BDA0005446409280000041
Patent Text Reader

Abstract

The invention discloses a vulnerability repair model construction method based on abstract syntax tree position coding. The method comprises the following steps: step 1, obtaining a 2D position sequence; and step 2, expanding rotation position encoding (Rotation Position Encoding, RoPE) to adapt to the 2D position sequence, so that a position fusion module is constructed. And step 3, inserting the position fusion module into a transformer-based model to form a final vulnerability repair model. By means of the method, the understanding ability of the model for the code structure is enhanced, and irrelevant attention calculation is reduced.
Need to check novelty before this filing date? Find Prior Art

Description

Technical Field

[0001] The present invention provides a program vulnerability repair technology, specifically a vulnerability repair model construction method based on abstract syntax tree position coding, and belongs to the field of software security. Background Art

[0002] While research in the field of automatic vulnerability repair isn't overwhelming, it does occasionally appear in academic research. From traditional methods based on dynamic and static analysis to deep learning-based approaches, researchers are constantly innovating to safeguard software security. Consequently, with the rapid growth of software vulnerabilities, vulnerability repair technology is also evolving and advancing rapidly. Vulnerability repairs are often intertwined with bug repairs, and sometimes even interchangeable. In some cases, the distinction between the two is clear, while in others, it's more ambiguous. Faced with this situation, some researchers are combining the two, especially with the development of deep learning. Bug repair data is more readily available than high-quality vulnerability repair data, and a large amount of bug repair data is crucial for training deep learning models. Using the concept of transfer learning, knowledge learned in the source domain can be applied to the target domain. Using transfer learning, Chen et al. collated a dataset of vulnerability repair and bug repair data and trained the Transformer model VRepair. This approach has become the foundation of many subsequent automatic vulnerability repair methods, including this one, which utilize transfer learning when training their repair models and incorporate VRepair into baseline comparisons.

[0003] Many previous research approaches have incorporated code structure, such as textual sequences of ASTs or DFGs, or vector representations passed through intermediate models, into repair models. After the advent of the Transformer, research on adding code structure continued this approach to make the model aware of code structure. However, one of the Transformer's most significant advantages over previous CNNs, RNNs, and other approaches is its attention mechanism. This attention mechanism allows different tokens in a sequence to pay attention to each other. However, it is worth noting that this mechanism, which relies on inner products between vectors, is inherently order-insensitive. This means that in the absence of additional information, the attention mechanism cannot distinguish the relative or absolute position of tokens in the sequence. Therefore, to address this limitation and enable the model to perceive positional information in the sequence, positional encoding is introduced. Positional encoding provides a positional identifier for each token in the sequence. This allows the attention mechanism to indirectly perceive their positional relationships, even if the attention mechanism does not directly consider their order when processing them. Summary of the Invention

[0004] The purpose of this invention is to provide a vulnerability repair model construction method based on abstract syntax tree position encoding, which expands the one-dimensional Transformer sequence position to 2D position. The AST sequence position is replaced by the parent node position of the simplified AST after parent representation. To solve the problem of position mismatch between the code after the Clang word segmenter and the BPE word segmenter, a mapping algorithm is designed. The encoder and decoder position fusion module is realized by expanding the Rotational Position Encoding (RoPE).

[0005] The technical solution adopted by the present invention is as follows: A method for constructing a vulnerability repair model based on abstract syntax tree position encoding:

[0006] Step 1: Get a 2D position sequence based on the abstract syntax tree;

[0007] 1.1: Prepare source code snippets containing vulnerabilities;

[0008] 1.2: The source code will be parsed into a series of tokens after lexical analysis, recorded as C code ={c1, c2, ..., c n}, C code Represents the sequence under Clang word segmentation. Each token has its own absolute position in the order, denoted by P original ={1, 2, ..., n}, where n is the length of the sequence after Clang segmentation;

[0009] 1.3: Generate AST based on lexical analysis. At this time, the AST node contains two kinds of information: node type and position information;

[0010] 1.4: Replace the token represented by each node of AST with the absolute position P of the token in the source code text sequence original After the above steps, the AST is simplified into a simplified abstract syntax tree (SAST) in which all nodes are represented by numbers only. Then, the simplified abstract syntax tree is represented by the parent representation to form a 2D position sequence C(P) under Clang word segmentation. original , P ast ), C represents Clang word segmentation, P original Represents the original position information of each node in the Clang word segmentation sequence, P ast Indicates the original position information of each node's parent node in the Clang word segmentation sequence.

[0011] 1.5 In order to solve the vocabulary overflow problem, the BPE word segmenter is used based on the transformer, so it is necessary to implement the conversion from the Clang position to the BPE position; the requirements for using BPE word segmentation are: scan the code sequence from left to right, and select the longest matching subword unit for segmentation each time; for the i-th token c under the Clang word segmentation i , after being processed by the BPE word segmenter, it will be split into multiple tokens of smaller lengths {b j+1 , b j+2 ,...}, the entire sequence C originally represented under Clang word segmentation code ={c1, c2, ..., c n}, which is represented as B under BPE segmentation code ={b1, b2, ..., b m}; For locators that locate vulnerable areas <SAST_StartBug>、<SAST_EndBug> and patch locators<SAST_ModStart> 、<SAST ModEnd> And the space retainer for text prompts in the code<SAST blank> , are all set as full tokens under the BPE segmentation and are not further decomposed into sub-words; a special token "CWE-XXX" is added before the vulnerable code to indicate which CWE category the vulnerability belongs to. For unknown vulnerabilities, "CWE-000" is used to indicate "generic fix"; "CWE-XXX" is set as the full token under the BPE segmentation;

[0012] 1.6 Link the respective sequence positions of Clang and BPE word segmentation with the positions of the original code string, so as to achieve original , P ast ) to B(P original , P ast ) mapping; the specific algorithm is implemented by the proposed algorithm 1, B(P original , P ast ) original is the sequential position, P ast The results of Algorithm 1 show that the input of Algorithm 1 is source code and Clang word segmentation sequence C code and BPE segmentation B code ;

[0013] Algorithm 1 is as follows: mapClangToBPE first calculates the c in the source code i (where i = 1, ..., n) the start and end indexes are expressed by formula 1; then, calculate b j(where j = 1, ..., m) The starting and ending indexes in the source code are expressed by formula 2; then initialize hasmapBtoC, where the key corresponds to P obtained by BPE segmentation original , the value corresponds to P obtained by Clang original ; hasmapCtoB is used as an inverse map: its keys correspond to Clang positions and its values ​​correspond to BPE positions; if a string is split into multiple subwords by BPE, the starting index of the string must be less than or equal to the starting index of all its subwords, and the ending index must be greater than or equal to the ending index of all its subwords; by comparing O Clang and O BPE , establish two types of participle P original The corresponding relationship between them is shown in Formula 3. When a string is divided into multiple subwords, only the position of the first subword is recorded in hasmapCtoB. Finally, the algorithm returns two hasmaps. Using these two hashmaps, the final B(P original , P ast ), B(P original , P ast ) is the final 2D position sequence;

[0014] O Clang (C, T Clang )={(s c1 ,e c1 ), s c2 , e c2 ),...,(s cn , e cn )) (1)

[0015] O BPE (C, T BPE )={(s B1 , e B1 ), (s B2 , e B2 ),…,s Bm ,e Bm )} (2)

[0016]

[0017] Step 2: Extend Rotary Position Embedding (RoPE) to adapt to the 2D position sequence to build a position fusion module;

[0018] The position information is added to the Q and K of the attention mechanism. Different position fusion modules are proposed for the encoder and decoder because they process word vectors slightly differently. The sequence processed by the encoder has a fixed length and the input code is complete. The complete AST tree is parsed by the Clang analyzer, and then the 2D position encoding B(P) under BPE word segmentation is formed. original , P ast ), use B(P original , P ast ) is recorded as the Original-AST (OA) module; the decoder's position fusion module is designed to use 2D position encoding, but both dimensions are in the form of absolute position B (P original , P original ), use B(P original , P original )’s position fusion module is recorded as Original-Original (OO) module.

[0019] 2.1OA Position Fusion Module

[0020] In the encoder fusion OA module, assuming that vectors q and k come from Q and K respectively, the vector dimension is d, and the corresponding 2D positions in Q and K are (m o , m A ) and (n o , n A ), the position information (m o , m A ) and (n o , n A ) is fused into vectors q and k, as shown in formulas 4 and 5:

[0021]

[0022] where q f represents the first d / 2 dimensions (q0, q1, ..., q d / 2-2 ,q d / 2-1 ) T ,q l Represents the last d / 2 dimensions (q d / 2 ,q d / 2+1 ,...,q d-2 ,q d-1 ) T , the rotation matrix and As shown in formulas 6 and 7 respectively, and Similarly, we can get, where θ=10000 -2i / d ;

[0023]

[0024] As shown in formulas 8, 9, and 10, in the process of calculating the attention score of q and k, n is dynamically included. o -m o Such relative position information;

[0025]

[0026] According to the original formula 3-10 of RoPE, R n-m is an orthogonal matrix

[0027]

[0028] It can be seen that 3-11 is an orthogonal matrix, and is an orthogonal matrix

[0029]

[0030] use as well as The sparsity of is used to provide a more computationally efficient implementation method than directly applying Formula 5, as shown in Formula 11.

[0031] After the fusion of the above position information, the initial query Q and key K contain absolute position information. The original position is integrated into the first d / 2 dimensions, and the AST sequence information is integrated into the last d / 2 dimensions. In the subsequent self-attention calculation process, the result contains relative position information through the inner product operation.

[0032]

[0033] Using the same method, we can come up with a more computationally efficient implementation method than directly applying Formula 4, as shown in Formula 12.

[0034] For the other calculations of the encoder, the formula is as follows:

[0035]

[0036] Formula 12 represents the output of the previous encoder Output of multi-head self-attention after layer normalization (LN) Formula 13 shows After the feedforward neural network FeedForward, residual connection and layer normalization are performed again.

[0037] 2.2OA Position Fusion Module

[0038] The cross attention mechanism in the decoder allows one sequence to focus on the relevant part of another sequence during the generation process. The query Q comes from the self-attention in the decoder, while the key K and value V come from the last layer output of the encoder. The OA position fusion module is applied to the cross attention K and V in the decoder according to the encoder. For the self-attention mechanism of the decoder, a 2D position encoding B (P original ,P original ), the length automatically increases as the input sequence increases, and the process of position fusion using RoPE is the same as the method in 2.1;

[0039] For the other calculations of the decoder, the formula is as follows:

[0040]

[0041] Formula 14 represents the self-attention of the current decoder The final output of the encoder Do the cross attention calculation, formula 15 represents the output of the previous decoder Output of multi-head self-attention after layer normalization (LN) Formula 16 shows After the feedforward neural network FeedForward, residual connection and layer normalization are performed again.

[0042] The model passes the decoder's final output through the linear layer Converted to a vector z whose dimension is equal to the size of the vocabulary V, that is, z∈R V , each element z i represents the unnormalized score logits of the i-th token in the vocabulary, and finally the Softmax function converts the vector z into a probability distribution for each token:

[0043]

[0044] in, is the exponential function value of the i-th logits, It is the sum of the exponential function values ​​of all logits and is used for normalization;

[0045] The input of the vulnerability repair model is the code sequence B code And a vulnerability type token added at the beginning, where the 2D position encoding schemes in the position fusion module of the encoder and decoder are (P original ,P ast ) and (P original ,P original ); Vulnerability repair model generates repair patch Rep={r1,r2,…,rn}, r i Represents the token in the patch, and the task of model training is to learn the vulnerability patch pair {B code ,Rep} features to generate patch predictions, as shown in Formula 3-18:

[0046]

[0047] Among them, Rep is the reference repair, Rep t Represents the token generated at time t, Rep <t Indicates the repair token generated before time t, B code Indicates the token sequence under BPE segmentation.

[0048] Step 3: Insert the position fusion module into the transformer-based model to form the final vulnerability repair model;

[0049] 3.1: Encoder Position Fusion Module: Inserted into encoder self-attention and decoder cross-attention; Insert the OA position fusion module into any transformer-based model, and the OA position fusion module is inserted after the encoder self-attention Q, K, and after the decoder cross-attention K, V.

[0050] 3.2: Decoder Position Fusion Module: Inserted into the decoder self-attention, insert the OO position fusion module into any transformer-based model, the OO position fusion module is inserted after the decoder self-attention Q and K, and after inserting the OA and OO position modules, the vulnerability repair model is built and then trained using conventional training methods.

[0051] The beneficial effects created by the present invention are: through in-depth analysis of the attention mechanism in the Transformer architecture, which captures the absolute and relative positional relationships between tokens through positional encoding, this study inspires the use of positional sequences to annotate the structural hierarchy of the input code sequence. In terms of code structure analysis, the present invention adopts the abstract syntax tree AST and appropriately simplifies it for digital positional representation. In order to solve the vocabulary overflow OOV problem, the present invention designs an algorithm to map the positional sequence based on Clang word segmentation to the positional sequence based on byte pair encoding BPE word segmentation, and forms a 2D positional encoding. In addition, in order to effectively integrate the 2D positional encoding of the encoder and decoder into the model, this study extends the rotational position encoding RoPE to a block matrix form, constructs the OA positional module for the encoder and the OO positional module for the decoder, thereby enhancing the model's ability to understand the code structure without increasing the length of the additional input sequence, and reducing irrelevant attention calculations. BRIEF DESCRIPTION OF THE DRAWINGS

[0052] Figure 1 To simplify the abstract syntax tree process diagram;

[0053] Figure 2 This is an example diagram of 2D position encoding;

[0054] Figure 3 This is the flowchart of the mapClangToBPE algorithm;

[0055] Figure 4 This is an example diagram of the insertion position fusion module;

[0056] Figure 5 This is the overall framework diagram of the model construction method. DETAILED DESCRIPTION

[0057] Step 1: Get a 2D position sequence based on the abstract syntax tree, such as Figure 1 As shown:

[0058] 1.1: Prepare source code snippets containing vulnerabilities;

[0059] 1.2: The source code will be parsed into a series of tokens after lexical analysis, recorded as C code ={c1, c2, ..., c n}, C code Represents the sequence under Clang word segmentation. Each token has its own absolute position in the order, denoted by P original ={1, 2, ..., n}, where n is the length of the sequence after Clang segmentation;

[0060] 1.3: Based on the lexical analysis, AST is further generated. At this time, the AST node usually contains two kinds of information: node type (such as declaration, identifier, function name, etc.) and location information (such as line number and column number in the source code);

[0061] 1.4: Replace the token represented by each node of AST with the absolute position P of the token in the source code text sequence original After the above steps, the AST is simplified into a simplified abstract syntax tree (SAST) in which all nodes are represented by numbers only. Then, the simplified abstract syntax tree is represented by the parent representation to form a 2D position sequence C (P original , P ast ), C represents Clang word segmentation, P original Represents the original position information of each node in the Clang word segmentation sequence, P astIndicates the original position information of each node's parent node in the Clang word segmentation sequence.

[0062] 1.5 In order to solve the vocabulary overflow problem, the BPE word segmenter is used based on the transformer, so it is necessary to implement the conversion from the Clang position to the BPE position; the requirements for using BPE word segmentation are: scan the code sequence from left to right, and select the longest matching subword unit for segmentation each time; for the i-th token c under the Clang word segmentation i , after being processed by the BPE word segmenter, it will be split into multiple tokens of smaller lengths {b j+1 , b j+2 ,...}, the entire sequence C originally represented under Clang word segmentation code ={c1, c2, ..., c n}, which is represented as B under BPE segmentation code ={b1, b2, ..., b m}; For locators that locate vulnerable areas <SAST_StartBug>、<SAST_EndBug> and patch locators<SAST_ModStart> 、<SAST ModEnd> And the space retainer for text prompts in the code<SAST blank> , are all set as full tokens under the BPE segmentation and are not further decomposed into sub-words; a special token "CWE-XXX" is added before the vulnerable code to indicate which CWE category the vulnerability belongs to. For unknown vulnerabilities, "CWE-000" is used to indicate "generic fix"; "CWE-XXX" is set as the full token under the BPE segmentation;

[0063] 1.6 Link the respective sequence positions of Clang and BPE word segmentation with the positions of the original code string, so as to achieve original , P ast ) to B(P original , P ast ) mapping; the specific algorithm is implemented by the proposed algorithm 1, B(P original , P ast ) original is the sequential position, P ast The results of Algorithm 1 show that the input of Algorithm 1 is source code and Clang word segmentation sequence C code and BPE segmentation B code ;

[0064] Algorithm 1 is as follows: mapClangToBPE first calculates the start and end indexes of ci (where i = 1, ..., n) in the source code, expressed by formula 1; then, calculates b j(where j = 1, ..., m) The starting and ending indexes in the source code are expressed by formula 2; then initialize hasmapBtoC, where the key corresponds to P obtained by BPE segmentation original , the value corresponds to P obtained by Clang original ; hasmapCtoB is used as an inverse map: its keys correspond to Clang positions and its values ​​correspond to BPE positions; if a string is split into multiple subwords by BPE, the starting index of the string must be less than or equal to the starting index of all its subwords, and the ending index must be greater than or equal to the ending index of all its subwords; by comparing O Clang and O BPE , establish two types of participle P original The corresponding relationship between them is shown in Formula 3. When a string is divided into multiple subwords, only the position of the first subword is recorded in hasmapCtoB. Finally, the algorithm returns two hasmaps. Using these two hashmaps, the final B(P original , P ast ), B(P original , P ast ) is the final 2D position sequence;

[0065] O Clang (C, T Clang )={(s c1 ,e c1 ), s c2 , e c2 ),...,(s cn , e cn )) (1)

[0066] O BPE (C, T BPE )={(s B1 , e B1 ), (s B2 , e B2 ),…,s Bm ,e Bm )} (2)

[0067]

[0068] Step 2: Extend Rotary Position Embedding (RoPE) to adapt to the 2D position sequence to build a position fusion module;

[0069] The position information is added to the Q and K of the attention mechanism. Different position fusion modules are proposed for the encoder and decoder because they process word vectors slightly differently. The sequence processed by the encoder has a fixed length and the input code is complete. The complete AST tree is parsed by the Clang analyzer, and then the 2D position encoding B(P) under BPE word segmentation is formed. original , P ast ), use B(P original , P ast ) is recorded as the Original-AST (OA) module; the decoder's position fusion module is designed to use 2D position encoding, but both dimensions are in the form of absolute position B (P original , P original ), use B(P original , P original )’s position fusion module is recorded as Original-Original (OO) module.

[0070] 2.1OA Position Fusion Module

[0071] In the encoder fusion OA module, assuming that vectors q and k come from Q and K respectively, the vector dimension is d, and the corresponding 2D positions in Q and K are (m o , m A ) and (n o , n A ), the position information (m o ,m A ) and (n o ,n A ) is fused into vectors q and k, as shown in formulas 4 and 5:

[0072]

[0073] where q f Represents the first d / 2 dimensions (q0,q1,…,q d / 2-2 ,q d / 2-1 ) T ,q l Represents the last d / 2 dimensions (q d / 2 ,q d / 2+1 ,…,q d-2 ,q d-1 ) T , the rotation matrix and As shown in formulas 6 and 7 respectively, and Similarly, we can get, where θ=10000 -2i / d ;

[0074]

[0075] As shown in formulas 8, 9, and 10, in the process of calculating the attention score of q and k, n is dynamically included. o -m o Such relative position information;

[0076]

[0077] According to the original formula 3-10 of RoPE, R n-m is an orthogonal matrix

[0078]

[0079] It can be seen that 3-11 is an orthogonal matrix, and is an orthogonal matrix

[0080]

[0081] use as well as The sparsity of is used to provide a more computationally efficient implementation method than directly applying Formula 5, as shown in Formula 11.

[0082] After the fusion of the above position information, the initial query Q and key K contain absolute position information. The original position is integrated into the first d / 2 dimensions, and the AST sequence information is integrated into the last d / 2 dimensions. In the subsequent self-attention calculation process, the result contains relative position information through the inner product operation.

[0083]

[0084] Using the same method, we can come up with a more computationally efficient implementation method than directly applying Formula 4, as shown in Formula 12.

[0085] For the other calculations of the encoder, the formula is as follows:

[0086]

[0087] Formula 12 represents the output of the previous encoder Output of multi-head self-attention after layer normalization (LN) Formula 13 shows After the feedforward neural network FeedForward, residual connection and layer normalization are performed again.

[0088] 2.2OA Position Fusion Module

[0089] The cross attention mechanism in the decoder allows one sequence to focus on the relevant part of another sequence during the generation process. The query Q comes from the self-attention in the decoder, while the key K and value V come from the last layer output of the encoder. The OA position fusion module is applied to the cross attention K and V in the decoder according to the encoder. For the self-attention mechanism of the decoder, a 2D position encoding B (P original ,P original ), the length automatically increases as the input sequence increases, and the process of position fusion using RoPE is the same as the method in 2.1;

[0090] For the other calculations of the decoder, the formula is as follows:

[0091]

[0092] Formula 14 represents the self-attention of the current decoder The final output of the encoder Do the cross attention calculation, formula 15 represents the output of the previous decoder Output of multi-head self-attention after layer normalization (LN) Formula 16 shows After the feedforward neural network FeedForward, residual connection and layer normalization are performed again.

[0093] The model passes the decoder's final output through the linear layer Converted to a vector z whose dimension is equal to the size of the vocabulary V, that is, z∈R V , each element z i represents the unnormalized score logits of the i-th token in the vocabulary, and finally the Softmax function converts the vector z into a probability distribution for each token:

[0094]

[0095] Among them, e zi is the exponential function value of the i-th logits, It is the sum of the exponential function values ​​of all logits and is used for normalization;

[0096] The input of the vulnerability repair model is the code sequence B code And a vulnerability type token added at the beginning, where the 2D position encoding schemes in the position fusion module of the encoder and decoder are (P original ,P ast ) and (P original ,P original ); Vulnerability repair model generates repair patch Rep={r1,r2,…,rn}, r i Represents the token in the patch, and the task of model training is to learn the vulnerability patch pair {B code ,Rep} features to generate patch predictions, as shown in Formula 3-18:

[0097]

[0098] Among them, Rep is the reference repair, Rep t Represents the token generated at time t, Rep <t Indicates the repair token generated before time t, B code Indicates the token sequence under BPE segmentation.

[0099] Step 3: Insert the position fusion module into the transformer-based model to form the final vulnerability repair model;

[0100] 3.1: Encoder Position Fusion Module: Inserted into encoder self-attention and decoder cross-attention; Insert the OA position fusion module into any transformer-based model, and the OA position fusion module is inserted after the encoder self-attention Q, K, and after the decoder cross-attention K, V.

[0101] 3.2: Decoder Position Fusion Module: Inserted into the decoder self-attention, insert the OO position fusion module into any transformer-based model, the OO position fusion module is inserted after the decoder self-attention Q and K, and after inserting the OA and OO position modules, the vulnerability repair model is built and then trained using conventional training methods.

[0102] Example 1:

[0103] To measure the repair performance of the repair model, a position fusion module was inserted into the pre-trained CodeT5 model to form the final vulnerability repair model. Six repair models, including the present invention, were trained on the same vulnerability dataset. Two evaluation metrics were used: Perfect Predictions (%PP) and CodeBLEU (higher scores indicate better performance for both metrics). The model was then compared with five other automatic vulnerability repair models: VQM, VulMaster, VulRepair, VRepair, and GraphCodeBERT. They are described in detail below:

[0104] 1) VQM: Vulnerable areas are annotated by learning a model VM that models the relationship between vulnerable tokens and repair tokens. The CodeT5 model is used as the vulnerability repair model, and the VM is used to weight the attention output of the CodeT5 model so that the query VQ can focus on vulnerable areas.

[0105] 2) VulMaster: The code, AST, vulnerability type, and repair case are input into the decoder fusion (FiD) framework respectively, and the feature representations are obtained. These intermediate results are then merged and input into the decoder.

[0106] 3) VRepair: uses a traditional vanilla Transformer architecture for automatic vulnerability repair tasks.

[0107] 4) VulRepair: A CodeT5-based approach that learns vulnerability repairs directly from a repair corpus.

[0108] 5) GraphCodeBERT: An extended version of CodeBERT that takes data flow graphs into account during training and adapts to the vulnerability repair task by continuing to train on the vulnerability repair dataset.

[0109] The comparison results are shown in Table 1. Compared with other vulnerability repair models, the method used in the present invention achieves the best performance in both indicators when the beam search is 1, 3, and 5.

[0110] Table 1 Performance comparison of various repair methods

[0111]

Claims

1. A method for constructing a vulnerability repair model based on abstract syntax tree position encoding, characterized by: Step 1: Get a 2D position sequence based on the abstract syntax tree; Step 2: Extend the rotation position encoding to adapt to the 2D position sequence to build a position fusion module; Step 3: Insert the position fusion module into the transformer-based model to form the final vulnerability repair model; 3.1: Encoder position fusion module: inserted into the encoder self-attention and decoder cross-attention; 3.2: Decoder position fusion module: inserted into the decoder self-attention.

2. The method for constructing a vulnerability repair model based on abstract syntax tree position encoding according to claim 1, characterized in that: The specific method in step 1 is: 1.1: Prepare source code snippets containing vulnerabilities; 1.2: The source code will be parsed into a series of tokens after lexical analysis, recorded as C code ={c1,c2,…,c n }, C code Represents the sequence under Clang word segmentation. Each token has its own absolute position in the order, denoted by P original ={1,2,…,n}, where n is the length of the sequence after Clang segmentation; 1.3: Generate AST based on lexical analysis. At this time, the AST node contains two kinds of information: node type and position information; 1.4: Replace the token represented by each node of AST with the absolute position P of the token in the source code text sequence original After the above steps, the AST is simplified into a simplified abstract syntax tree (SAST) where all nodes are represented by numbers only. This simplified abstract syntax tree is then represented using the parent representation to form a 2D position sequence under Clang word segmentation. C represents Clang participle, P original Represents the original position information of each node in the Clang word segmentation sequence, P ast Represents the original position information of each node's parent node in the Clang word segmentation sequence; 1.5 In order to solve the vocabulary overflow problem, the BPE word segmenter is used based on the transformer, so it is necessary to implement the conversion from the Clang position to the BPE position; the requirements for using BPE word segmentation are: scan the code sequence from left to right, and select the longest matching subword unit for segmentation each time; for the i-th token c under the Clang word segmentation i , after being processed by the BPE word segmenter, it will be split into multiple tokens of smaller lengths {b j+1 ,b j+2 ,…}, the entire sequence C originally represented under Clang word segmentation code ={c1,c2,…,c n }, which is represented as B under BPE segmentation code ={b1,b2,…,b m }; For locators that locate vulnerable areas<SAST_StartBug> 、<SAST_EndBug> and patch locators<SAST_ModStart> 、<SAST_ModEnd> And the space retainer for text prompts in the code<SAST_blank> , are all set to the full token under the BPE word segmentation and are not further decomposed into subwords; a special token "CWE-XXX" is added before the vulnerable code to indicate which CWE category the vulnerability belongs to. For unknown vulnerabilities, "CWE-000" is used to indicate "generic fix"; "CWE-XXX" is set to the full token under the BPE word segmentation; 1.6 Link the respective sequence positions of Clang and BPE word segmentation with the positions of the original code string, so as to achieve original ,P ast ) to B(P original ,P ast )’s mapping; The specific algorithm is implemented by the proposed algorithm 1, B(P original ,P ast ) original is the sequential position, P ast The results of Algorithm 1 show that the input of Algorithm 1 is source code and Clang word segmentation sequence C code and BPE segmentation B code ; Algorithm 1 is as follows: mapClangToBPE first calculates the c in the source code i (where i = 1, ..., n) the start and end indexes are expressed by formula 1; then, calculate b j (where j = 1, ..., m) The starting and ending indexes in the source code are expressed by formula 2; then initialize hasmapBtoC, where the key corresponds to P obtained by BPE segmentation original , the value corresponds to P obtained by Clang original ; hasmapCtoB as an inverse map: its keys correspond to Clang locations and its values ​​correspond to BPE locations; If a string is split into multiple subwords by BPE, the starting index of the string must be less than or equal to the starting index of all its subwords, and the ending index must be greater than or equal to the ending index of all its subwords; by comparing O Clang and O BPE , establish two types of participle P original The corresponding relationship between them is shown in Formula 3. When a string is divided into multiple subwords, only the position of the first subword is recorded in hasmapCtoB. Finally, the algorithm returns two hasmaps. Using these two hashmaps, the final B(P original ,P ast ), B(P original ,P ast ) is the final 2D position sequence; O Clang (C,T Clang )={(s c1 ,e c1 ),(s c2 ,e c2 ),…,(s cn ,e cn )} (1) O BPE (C,T BPE )={(s B1 ,e B1 ),(s B2 ,e B2 ),…,(s Bm ,e Bm )} (2) 3. The method for constructing a vulnerability repair model based on abstract syntax tree position encoding according to claim 1, characterized in that: The specific method in step 2 is: The position information is added to the Q and K of the attention mechanism. Different position fusion modules are proposed for the encoder and decoder because they process word vectors slightly differently. The sequence processed by the encoder has a fixed length and the input code is complete. The complete AST tree is parsed by the Clang analyzer, and then the 2D position encoding B(P) under BPE word segmentation is formed. original ,P ast ), use B(P original ,P ast ) is recorded as the Original-AST (OA) module; the decoder's position fusion module is designed to use 2D position encoding, but both dimensions are in the form of absolute position B (P original ,P original ), use B(P original ,P original )’s position fusion module is recorded as Original-Original (OO) module; 2.1OA Position Fusion Module In the encoder fusion OA module, assuming that vectors q and k come from Q and K respectively, the vector dimension is d, and the corresponding 2D positions in Q and K are (m o ,m A ) and (n o ,n A ), the position information (m o ,m A ) and (n o ,n A ) is fused into vectors q and k, as shown in formulas 4 and 5: where q f Represents the first d / 2 dimensions (q0,q1,…,q d / 2-2 ,q d / 2-1 ) T ,q l Represents the last d / 2 dimensions (q d / 2 ,q d / 2+1 ,…,q d-2 ,q d-1 ) T , the rotation matrix and As shown in formulas 6 and 7 respectively, and Similarly, we can get, where θ=10000 -2i / d ; As shown in formulas 8, 9, and 10, in the process of calculating the attention score of q and k, n is dynamically included. o -m o Such relative position information; According to the original formula 3-10 of RoPE, R n-m is an orthogonal matrix It can be seen that 3-11 is an orthogonal matrix, and is an orthogonal matrix use as well as The sparsity of is used to provide a more computationally efficient implementation method than directly applying Formula 5, as shown in Formula 11. After the fusion of the above position information, the initial query Q and key K contain absolute position information. The original position is integrated into the first d / 2 dimensions, and the AST sequence information is integrated into the last d / 2 dimensions. In the subsequent self-attention calculation process, the result contains relative position information through the inner product operation. Using the same method, we can come up with a more computationally efficient implementation method than directly applying Formula 4, as shown in Formula 12. For the other calculations of the encoder, the formula is as follows: Formula 12 represents the output of the previous encoder Output of multi-head self-attention after layer normalization (LN) Formula 13 shows After the feedforward neural network FeedForward, residual connection and layer normalization are performed again; 2.2OA Position Fusion Module The cross attention mechanism in the decoder allows one sequence to focus on the relevant part of another sequence during the generation process. The query Q comes from the self-attention in the decoder, while the key K and value V come from the last layer output of the encoder. The OA position fusion module is applied to the cross attention K and V in the decoder according to the encoder. For the self-attention mechanism of the decoder, a 2D position encoding B (P original ,P original ), the length automatically increases as the input sequence increases, and the process of position fusion using RoPE is the same as the method in 2.1; For the other calculations of the decoder, the formula is as follows: Formula 14 represents the self-attention of the current decoder The final output of the encoder Do the cross attention calculation, formula 15 represents the output of the previous decoder Output of multi-head self-attention after layer normalization (LN) Formula 16 shows After the feedforward neural network FeedForward, residual connection and layer normalization are performed again; The model passes the decoder's final output through the linear layer Converted to a vector z whose dimension is equal to the size of the vocabulary V, that is, z∈R V , each element z i represents the unnormalized score logits of the i-th token in the vocabulary, and finally the Softmax function converts the vector z into a probability distribution for each token: in, is the exponential function value of the i-th logits, It is the sum of the exponential function values ​​of all logits and is used for normalization; The input of the vulnerability repair model is the code sequence V code And a vulnerability type token added at the beginning, where the 2D position encoding schemes in the position fusion module of the encoder and decoder are (P original ,P ast ) and (P original ,P original ); Vulnerability repair model generates repair patch Rep={r1,r2,…,r n }, r i Represents the token in the patch, and the task of model training is to learn the vulnerability patch pair {B code ,Rep} features to generate patch predictions, as shown in Formula 3-18: Among them, Rep is the reference repair, Rep t Represents the token generated at time t, Rep <t Indicates the repair token generated before time t, B code Indicates the token sequence under BPE segmentation.

4. The method for constructing a vulnerability repair model based on abstract syntax tree position encoding according to claim 1, characterized in that: The specific method in step 3 is: 3.1 Inserting the OA Position Fusion Module into Any Transformer-Based Model The OA position fusion module is inserted after the encoder self-attention Q, K, and after the decoder cross attention K, V; 3.2 Inserting the OO Position Fusion Module into Any Transformer-Based Model The OO position fusion module is inserted into the decoder after the self-attention Q and K, and after the OA and OO position modules are inserted, the vulnerability repair model is built and then trained using the conventional training method.

Citation Information

Cited By

  • Data processing method and device, equipment and medium

    CN121502759A