Vulnerability detection method based on Transform interlayer similarity weighted fusion

Abstract

The invention belongs to the field of vulnerability detection methods, and relates to a vulnerability detection method based on Transform interlayer similarity weighted fusion, which comprises the following steps of: calculating the similarity between layers in a Transform model, enhancing the understanding ability of the model on a code structure and a mode, and specifically, calculating the similarity between a current layer and an upper layer in each layer, so as to improve the understanding ability of the model on the code structure and the mode; and adding to the output of the current layer, and giving a proper weight. According to the method, the complex dependency relationship in the code can be captured, and the change of the code structure can be identified, so that the accuracy of vulnerability detection is improved.

Description

Technical Field

[0001] The patent of this invention relates to the field of vulnerability detection methods, and in particular to a vulnerability detection method based on weighted fusion of Transformer inter-layer similarity. Background Art

[0002] As software applications grow in size and complexity, code bases become increasingly large, and it becomes increasingly difficult to find and fix specific software defects. Vulnerability detection is a key technology in the field of software development and maintenance. Its goal is to identify and fix security vulnerabilities in the code to ensure the security and reliability of the software. Therefore, the development of vulnerability detection technology has become a key factor in software quality assurance and security improvement.

[0003] Initial methods of vulnerability detection included manual inspection and static analysis, but over time, more automated tools and algorithms have been developed to improve efficiency and reduce labor costs. The basic idea is to identify potential security vulnerabilities by analyzing the structure and behavior of the code and provide repair suggestions. In recent years, with the rapid development of artificial intelligence technology, machine learning and deep learning methods have been introduced into the field of vulnerability detection. These methods use data from large-scale code repositories for learning and automatically generate detection strategies in order to find security vulnerabilities in the code in a shorter time.

[0004] In particular, the self-attention mechanism of the Transformer-based model greatly enhances the model's ability to understand code semantics, providing a new perspective for vulnerability detection. Through the multi-layer self-attention mechanism, the Transformer model can capture complex dependencies and patterns in the code, thereby improving the accuracy and efficiency of vulnerability detection.

[0005] Although Transformer-based vulnerability detection technology has improved performance to a certain extent, it still faces problems such as insufficient capture of contextual information, limited understanding of code structure, insufficient training data, and poor model interpretability. Summary of the invention

[0006] In order to overcome the technical problems existing in the prior art, the present invention proposes a vulnerability detection method based on weighted fusion of Transformer inter-layer similarity, which is used to improve the accuracy and efficiency of vulnerability detection. By introducing inter-layer similarity in the Transformer model, the code structure and pattern can be better captured, thereby improving the model's ability to detect code vulnerabilities.

[0007] The technical solution adopted by the present invention is as follows: a vulnerability detection method based on weighted fusion of Transformer inter-layer similarity, the steps of which are:

[0008] 1) Dataset preparation: After collecting code snippets containing vulnerabilities and corresponding labels and preprocessing them, they are divided into training set, validation set and test set; preprocessing includes removing comments and formatting code. After preprocessing, the data set is divided into training set, validation set and test set according to the ratio of 70%, 20% and 10%.

[0009] 2) Build the model:

[0010] The Transformer architecture is used, including an encoder and a decoder. The similarity between the current layer and the previous layer is calculated in each layer and appended to the output of the current layer, where the input embedding is:

[0011] X=[x1,x2,…,x n ] (1)

[0012] Among them, x i Represents the i-th input vector; after the projection matrix W q ,W k ,W v ,get:

[0013] Q=XW q (2)

[0014] K=XW k (3)

[0015] V=XW v (4)

[0016] Self-Attention Mechanism:

[0017]

[0018] Where Q, K, and V represent query, key, and value matrices, respectively. k is the dimension of the key, K T is the transpose of K; multi-head attention:

[0019] MultiHead(Q,K,V)=Concat(head1,…,head h )W O (6)

[0020] Among them, head i =Attention(QW i Q ,KW i K ,VW i V ), W i Q , W i K , Wi V Respectively represent the Q, K, and V weight matrices of the i-th head; residual link:

[0021] x=LayerNorm(X+MultiHead(Q,K,V)) (7)

[0022] Where LayerNorm stands for layer normalization, x is the output; layer normalization:

[0023] x=LayerNorm(X+FFN(x)) (8)

[0024] Among them, FFN is a feedforward neural network, and the output of the lth layer is:

[0025] H (l) =max(0,xW1+b1)W2+b2 (9)

[0026] W1 and W2 are the weights of the two feedforward neural networks, and b1 and b2 are bias terms.

[0027] Calculate the similarity between the output of layer l and the output of the previous layer k:

[0028]

[0029] Among them, H (l) represents the output of the lth layer, H (k) represents the output of the kth layer, · represents the dot product of the vector, ∥·∥ represents the modulus of the vector;

[0030] Record the similarity vector between the output of the lth layer and all previous layers:

[0031] Sim (l) = {Sim(H (l) ,H (k) )|k=1,2,...,l-1} (11)

[0032] Normalize the similarity vector to obtain the normalized similarity vector:

[0033]

[0034] Use the normalized similarity vector as the weight to perform weighted fusion on the outputs of the previous layers:

[0035]

[0036] The balance between the original output and the fused output is controlled by the weight α, where 0≤α≤1. The new l-th layer output is expressed as:

[0037]

[0038] 3) Training the model: using the cross entropy loss function, forward propagation and back propagation are performed on the training set to update the model parameters;

[0039]

[0040] Among them, y i is the true label, is the predicted probability.

[0041] 4) Preprocess the code snippet to be tested and convert it into the model input format. Perform forward propagation through the trained Transformer model, calculate the output of each layer and the similarity between layers, and identify potential vulnerabilities in the code based on the probability distribution of the model output.

[0042] The beneficial effects created by the present invention are as follows: the present invention provides a vulnerability detection method based on weighted fusion of Transformer inter-layer similarity, which enhances the model's ability to capture code context relationships and improves the understanding of code structure by calculating the similarity between Transformer layers, thereby facilitating more accurate detection of complex vulnerabilities. In addition, the introduction of inter-layer similarity can also alleviate the problem of insufficient training data to a certain extent and improve data utilization efficiency. At the same time, by analyzing the inter-layer similarity, we can more easily understand the decision-making process of the model, enhance the interpretability of the model, and enable developers to better understand and fix vulnerabilities. Therefore, the improved method based on Transformer inter-layer similarity has significant advantages in vulnerability detection and is expected to bring new breakthroughs to the development of vulnerability detection technology. DETAILED DESCRIPTION

[0043] The present invention will be described in full below. Obviously, the described embodiments are only a part of the embodiments of the present invention, rather than all the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0044] The vulnerability detection method based on weighted fusion of Transformer inter-layer similarity is as follows:

[0045] 1) Dataset preparation: After collecting code snippets containing vulnerabilities and corresponding labels and preprocessing them, they are divided into training set, validation set and test set; preprocessing includes removing comments and formatting code. After preprocessing, the data set is divided into training set, validation set and test set according to the ratio of 70%, 20% and 10%.

[0046] 2) Build the model:

[0047] The Transformer architecture is used, including an encoder and a decoder. The similarity between the current layer and the previous layer is calculated in each layer and appended to the output of the current layer, where the input embedding is:

[0048] X=[x1,x2,…,x n ] (1)

[0049] Among them, x i Represents the i-th input vector; after the projection matrix W q ,W k ,W v ,get:

[0050] Q=XW q (2)

[0051] K=XW k (3)

[0052] V=XW v (4)

[0053] Self-Attention Mechanism:

[0054]

[0055] Where Q, K, and V represent query, key, and value matrices, respectively. k is the dimension of the key, K T is the transpose of K; multi-head attention:

[0056] MultiHead(Q,K,V)=Concat(head1,…,head h )W O (6)

[0057] Among them, head i =Attention(QW i Q ,KW i K ,VW i V ), W i Q , W i K , W i V Respectively represent the Q, K, and V weight matrices of the i-th head; residual link:

[0058] x=LayerNorm(X+MultiHead(Q,K,V)) (7)

[0059] Where LayerNorm stands for layer normalization, x is the output; layer normalization:

[0060] x=LayerNorm(X+FFN(x)) (8)

[0061] Among them, FFN is a feedforward neural network, and the output of the lth layer is:

[0062] H (l) =max(0,xW1+b1)W2+b2 (9)

[0063] W1 and W2 are the weights of the two feedforward neural networks, and b1 and b2 are bias terms.

[0064] Calculate the similarity between the output of layer l and the output of the previous layer k:

[0065]

[0066] Among them, H (l) represents the output of the lth layer, H (k) represents the output of the kth layer, · represents the dot product of the vector, ∥·∥ represents the modulus of the vector;

[0067] Record the similarity vector between the output of the lth layer and all previous layers:

[0068] Sim (l) = {Sim(H (l) ,H (k) )|k=1,2,...,l-1} (11)

[0069] Normalize the similarity vector to obtain the normalized similarity vector:

[0070]

[0071] Use the normalized similarity vector as the weight to perform weighted fusion on the outputs of the previous layers:

[0072]

[0073] The balance between the original output and the fused output is controlled by the weight α, where 0≤α≤1. The new l-th layer output is expressed as:

[0074]

[0075] 3) Training the model: using the cross entropy loss function, forward propagation and back propagation are performed on the training set to update the model parameters;

[0076]

[0077] Among them, y i is the true label, is the predicted probability.

[0078] 4) Preprocess the code snippet to be tested and convert it into the model input format. Perform forward propagation through the trained Transformer model, calculate the output of each layer and the similarity between layers, and identify potential vulnerabilities in the code based on the probability distribution of the model output.

[0079] 5) Experiment:

[0080] In the experiments, we used multiple challenging software datasets that contain code bases of various types and sizes, as well as known security vulnerabilities. Through comparative experiments, we conducted an in-depth comparison between the layer-by-layer similarity method based on the Transformer model and traditional vulnerability detection methods.

[0081] Experimental results show that traditional methods often have difficulty in fully capturing potential vulnerabilities when dealing with large and complex code bases, especially when dealing with complex dependencies and structural changes between codes. In contrast, our proposed method significantly improves the model's ability to recognize code structure and patterns by calculating inter-layer similarity. This not only enables the model to locate potential vulnerabilities more accurately, but also predicts the type and severity of vulnerabilities to a certain extent.

[0082] Table 1: Time result comparison table

[0083] method Accuracy (percentage) F1 score Devign 51.09 0.49 RXGD 57.33 0.56 CodeBert 63.19 0.43 CodeGPT 63.06 0.61 GraphCodeBERT 62.28 0.62 BAGS 58.91 0.57 LFME 61.57 0.61 Inter-layer similarity method 65.76 0.65

Claims

1. A vulnerability detection method based on weighted fusion of similarity between Transformer layers, characterized in that: The steps are: 1) Dataset preparation: Collect code snippets containing vulnerabilities and their corresponding labels, preprocess them, and divide them into training set, validation set, and test set; 2) Build the model: The Transformer architecture is used, including an encoder and a decoder. The similarity between the current layer and the previous layer is calculated in each layer and appended to the output of the current layer, where the input embedding is: X=[x1,x2,…,x n ] (1) Among them, x i Represents the i-th input vector; after the projection matrix W q ,W k ,W v ,get: Q=XW q (2) K=XW k (3) V=XW v (4) Self-Attention Mechanism: Where Q, K, and V represent query, key, and value matrices, respectively. k is the dimension of the key, K T is the transpose of K; multi-head attention: MultiHead(Q,K,V)=Concat(head1,…,head h )W O (6) Among them, head i =Attention(QW i Q ,KW i K ,VW i V ), W i Q , W i K , W i V Respectively represent the Q, K, and V weight matrices of the i-th head; residual link: x=LayerNorm(X+MultiHead(Q,K,V)) (7) Where LayerNorm stands for layer normalization, x is the output; layer normalization: x=LayerNorm(X+FFN(x)) (8) Among them, FFN is a feedforward neural network, and the output of the lth layer is: H (l) =max(0,xW1+b1)W2+b2 (9) W1 and W2 are the weights of the two feedforward neural networks, and b1 and b2 are bias terms. Calculate the similarity between the output of layer l and the output of the previous layer k: Among them, H (l) represents the output of the lth layer, H (k) represents the output of the kth layer, · represents the dot product of the vector, ∥·∥ represents the modulus of the vector; Record the similarity vector between the output of the lth layer and all previous layers: Sim (l) ={Sim(H (l) ,H (k) )|k=1,2,...,l-1} (11) Normalize the similarity vector to obtain the normalized similarity vector: Use the normalized similarity vector as the weight to perform weighted fusion on the outputs of the previous layers: The balance between the original output and the fused output is controlled by the weight α, where 0≤α≤1. The new l-th layer output is expressed as: 3) Training model; 4) Preprocess the code snippet to be tested and convert it into the model input format. Perform forward propagation through the trained Transformer model, calculate the output of each layer and the similarity between layers, and identify potential vulnerabilities in the code based on the probability distribution of the model output.

2. According to claim 1, the vulnerability detection method based on weighted fusion of Transformer inter-layer similarity is characterized in that: In the above 1), the preprocessing includes removing comments and formatting codes. The preprocessed data is divided into a training set, a validation set and a test set according to the ratio of 70%, 20% and 10%.

3. The vulnerability detection method based on weighted fusion of Transformer inter-layer similarity according to claim 1 is characterized in that: In the above 3), Using the cross entropy loss function, forward propagation and back propagation are performed on the training set to update the model parameters; Among them, y i is the true label, is the predicted probability.

Images