Software vulnerability detection method based on Transform architecture pre-training model

Abstract

The invention discloses a software vulnerability detection method based on a Transform architecture pre-training model, and the method comprises the steps: adding an adapter module on a self-attention matrix based on a base model; in the forward propagation process of the model, original parameters of the pre-training model are frozen, only parameters of an adapter are finely adjusted, and a loss function of the model is minimized; calculating a disturbance matrix in the SPSA method by using a reinforcement learning method so as to minimize a loss function of the detection model; and the parameters A and B are iteratively updated, so that the model learns finer representation. According to the method, an efficient parameter updating strategy is designed on the basis of the adapter, the adapted model is expected to show powerful performance in the field of vulnerability detection, and a reliable tool is provided for security engineers to efficiently and accurately identify and classify various potential risks in software.

Description

Technical Field

[0001] The present invention relates to the fields of network security and NLP technology, and mainly to a software vulnerability detection method based on a Transformer architecture pre-training model. Background Art

[0002] Software vulnerability detection is a key technology in the field of information security. It aims to identify weaknesses or defects in software that may be exploited by attackers to undermine system security. In this process, various tools and methods have been developed to discover potential security threats in an automated or semi-automated manner, thereby protecting the software and its data from damage.

[0003] As computer software systems become more complex, the technology behind vulnerability detection is also constantly improving. Initial methods include static analysis and dynamic analysis techniques. Static analysis is the analysis of source code or compiled code to find potential vulnerabilities, such as buffer overflows, SQL injections, etc., without running the program. Dynamic analysis involves actually running the program and monitoring its behavior to identify anomalies, and is usually used to find vulnerabilities that only appear when the program is running.

[0004] In recent years, the introduction of deep learning technology has brought new breakthroughs in software vulnerability detection. Using deep learning models, researchers are able to learn and extract features from a large number of code samples and train models to automatically identify and classify vulnerabilities. In particular, convolutional neural networks (CNN) and recurrent neural networks (RNN) perform well in processing the grammatical and semantic features of code. At the same time, the Transformer model, with its self-attention mechanism, has set a new standard in the semantic understanding of source code.

[0005] Large-scale pre-trained models based on Transformer, such as BERT, GPT, etc., have demonstrated their powerful capabilities in many fields. These models have not only revolutionized the landscape of natural language processing, but have also begun to penetrate into many other fields, such as computer vision, audio processing, and cross-modal (such as image and text) tasks, demonstrating their excellent versatility and flexibility.

[0006] On the other hand, despite the limitations of data quality and quantity, deep learning methods still achieve superior performance in vulnerability detection. Traditional vulnerability databases, such as the National Vulnerability Database (NVD) and the Open Source Vulnerability Database (OSS), provide rich historical vulnerability report information. After proper cleaning and processing, these data can be used to train more powerful and accurate vulnerability detection models.

[0007] Each vulnerability type, often represented by a unique CWE-ID, provides in-depth explanations for the detected vulnerable code and assists security engineers in understanding, classifying, and responding to these vulnerabilities. Software Vulnerability Classification (SVC) thus becomes a critical task that assists the SVP model by providing richer context for detected vulnerabilities, allowing end users to better understand potential security risks.

[0008] The Transformer-based pre-trained large-scale model shows amazing capabilities, and it is very meaningful to propose a software vulnerability detection method. Summary of the invention

[0009] In view of this, the present invention provides a software vulnerability detection method based on a Transformer architecture pre-training model to efficiently and accurately identify and classify various potential risks in the software.

[0010] In order to achieve the above object, the present invention provides a software vulnerability detection method based on a Transformer architecture pre-training model, comprising:

[0011] Based on the base model, add an adapter module on the self-attention matrix;

[0012] During the forward propagation of the model, the original parameters of the pre-trained model are frozen, and only the parameters of the adapter are fine-tuned to minimize the loss function of the model;

[0013] Reinforcement learning is used to calculate the perturbation matrix in the SPSA method to minimize the loss function of the detection model. Reinforcement learning is performed, and the parameters A and B of the adapter module are used as the input of the proxy model, and the perturbation matrices zA and zB are used as outputs. The loss function and F1 score are converted and used as a reward signal to guide the training process of the proxy model.

[0014] By iteratively updating parameters A and B, the model learns more refined representations.

[0015] Preferably, the base model is pre-trained on a large-scale code dataset.

[0016] Preferably, the software vulnerability detection method comprises:

[0017] Step 1: Select the CodeGeex2 model as the base model and add the self-attention matrix W Q , W K , W V Add an adapter module to freeze the original parameters of the pre-trained model during the forward propagation of the base model, and only train the parameters of the LoRA adapter;

[0018] Step 2: Apply the SPSA gradient estimation method to train the parameters of the LoRA adapter, and update the perturbation matrix through the reinforcement learning method, and finally update the adapter matrix only through forward propagation; obtain the detection model, as shown in formula (1);

[0019] P=W0H+ΔWH=W0H+BAH (1)

[0020] Where W0 represents the parameters of the pre-trained model, ΔW represents the adapter parameters, and H represents the input;

[0021] Step 3: Merge the frozen parameters of the LoRA adapter and pre-trained model;

[0022] Step 4: Input code H, through the projection matrix W q ,W k ,W v ,get:

[0023] Q = HW q (2)

[0024] K = HW k (3)

[0025] V = HW v (4)

[0026] Then through self-attention

[0027]

[0028] Through the encoder:

[0029] A n = LN(Multi(Attention(Out enc n-1 )))+Out enc n-1 (6)

[0030] Out enc n =LN(FeedForward(A n )+A n ) (7)

[0031] Among them, Out enc n-1 It is the output of the n-1 layer encoder, which passes through the multi-head self-attention Multi and layer normalization LN, and then passes through the feedforward network FeedForward and residual link, and then layer normalization again;

[0032] After the decoder:

[0033]

[0034] Among them, Q cross is the decoder self-attention layer output, Q dec ,K dec ,V dec The decoder input sequence is transformed into a hidden dimension, and CrossAttention calculates the encoder output layer and Q cross Cross-attention, Out dec n-1 It is the self-attention output of the n-1 layer decoder, which is subjected to multi-head self-attention calculation and layer normalization for residual connection, and finally to layer normalization after passing through the feedforward network and residual connection;

[0035] Finally, after average pooling and SoftMax function, the vulnerability category C is output;

[0036] C=Softmax(AveragePolling(Out dec n )) (12).

[0037] Preferably, step 2 specifically includes:

[0038] Step 21: The B and A matrices in the adapter are initialized with random Gaussian for A and zero for B using the LoRA method;

[0039] Step 22: Modify the loss function of the model. In vulnerability prediction, the vulnerability prediction task is regarded as a classification task. The loss function considers both the cross entropy loss and the F1 score. The cross entropy loss is used to measure the difference between the true label distribution and the predicted label distribution. The F1 score is the harmonic mean of the precision and recall. Finally, the loss function is:

[0040]

[0041]

[0042] Among them, TP is a true positive example and FP is a false positive example;

[0043]

[0044] Where FN is a false negative example;

[0045] Step 23: Use the zero-order gradient descent algorithm RL-ZO so that only forward propagation is performed when training the model; introduce reinforcement learning into the zero-order gradient descent algorithm RL-ZO, and the reward function is shown in formula (17):

[0046]

[0047] Subtract the F1 score from the cross entropy loss and introduce a weight factor to merge the two indicators into a unified dimension. Then, set appropriate priorities for the two indicators by adjusting the values ​​of γ and λ.

[0048] Step 24: The reinforcement learning agent model uses a layer of transformer model. The transformer uses its self-attention mechanism to focus on the relationship between different parameters. Through the encoder and decoder structure, the adapter parameters θ of the original model are input into the agent model, and the corresponding perturbation matrix z is output.

[0049] After the proxy model outputs the perturbation matrix z, it calculates the SPSA gradient estimate through two forward propagations, updates the A and B matrices, and performs model fine-tuning and reinforcement learning model training alternately. During the reinforcement learning model training, the DDPG algorithm is used to update the parameters.

[0050] The software vulnerability detection method based on the Transformer architecture pre-trained model provided by the present invention designs an efficient parameter update strategy based on the adapter, aiming to customize these powerful models for specific areas of software vulnerability detection. Through such transformation, the original large model-pre-trained model will be adapted to downstream tasks to receive suspicious code snippets as input and produce accurate vulnerability classification results.

[0051] The design of the adapter takes into account the complexity of the pre-trained model architecture and the special needs of downstream tasks. By introducing the adapter, the number of parameters that need to be updated can be significantly reduced, thereby improving overall efficiency and reducing resource consumption. At the same time, in order to better adapt to vulnerability detection tasks, a new parameter update method is adopted to enable fast parameter updates on non-differentiable targets. The model after such adaptation is expected to show strong performance in the field of vulnerability detection and provide security engineers with a reliable tool to efficiently and accurately identify and classify various potential risks in software. BRIEF DESCRIPTION OF THE DRAWINGS

[0052] In order to more clearly illustrate the technical solutions in the embodiments of the present invention, the following briefly introduces the drawings required for describing the embodiments. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative work:

[0053] Figure 1 A schematic diagram of the structure of a software vulnerability detection system based on a Transformer architecture pre-training model provided by the present invention. DETAILED DESCRIPTION

[0054] The present invention will be further explained below in conjunction with specific implementation schemes, but the present invention is not limited thereto.

[0055] The present invention provides a method based on adapter and reinforcement learning gradient optimization to adapt the pre-trained model of Transformer architecture to the code vulnerability detection method. Figure 1 As shown, the entire system framework is based on a pre-trained base model. The main steps of the method include:

[0056] Step 1: Select a pre-trained model of a suitable Transformer architecture as the base model, preferably the CodeGeex2 model. Then, in its self-attention matrix W Q , W K , W V Add the adapter module. During the forward propagation of the model, freeze the original parameters of the pre-trained model and only train the parameters of the LoRA adapter.

[0057] Step 2: Apply the SPSA (Simultaneous Perturbation Stochastic Approximation) gradient estimation method to train the parameters of the LoRA adapter, and update the perturbation matrix through the reinforcement learning method, and finally update the adapter matrix only through forward propagation.

[0058] Detection model, as shown in formula (1);

[0059] P=W0H+ΔWH=W0H+BAH (1)

[0060] Where W0 represents the parameters of the pre-trained model, ΔW represents the adapter parameters, and H represents the input;

[0061] Step 3: Merge the frozen parameters of the LoRA adapter and the model. In the inference phase, ΔW is merged with the original frozen W0 without increasing the inference latency.

[0062] Input code, represented as H, passes through the projection matrix W q ,W k ,W v ,get:

[0063] Q = HW q (2)

[0064] K = HW k (3)

[0065] V = HW v (4)

[0066] Then through self-attention

[0067]

[0068] Through the encoder:

[0069] A n = LN(Multi(Attention(Out enc n-1 )))+Out enc n-1 (6)

[0070] Out enc n =LN(FeedForward(A n )+A n ) (7)

[0071] Among them, Out enc n-1 It is the output of the n-1 layer encoder, which passes through the multi-head self-attention Multi and layer normalization LN, and then passes through the feedforward network FeedForward and the residual link, and then is normalized again.

[0072] After the decoder:

[0073]

[0074] Among them, Q cross is the output of the decoder self-attention layer, Q dec ,K dec ,V dec The decoder input sequence is transformed into a hidden dimension, and CrossAttention calculates the encoder output layer and Q cross Cross-attention, Out dec n-1 It is the self-attention output of the n-1 layer decoder, which is subjected to multi-head self-attention calculation and layer normalization for residual connection, and finally to layer normalization after passing through the feedforward network and residual connection.

[0075] Finally, after average pooling and SoftMax function, the vulnerability category C is output.

[0076] C=Softmax(AveragePolling(Out dec n )) (12)

[0077] The detailed training process in step 2 is described as follows:

[0078] Step 1: For the B and A matrices in the adapter, use the LoRA method to initialize A with random Gaussian and initialize B with zero.

[0079] Step 2: Modify the loss function of the model. In vulnerability prediction, the vulnerability prediction task is regarded as a classification task. The accuracy evaluation of the vulnerability is also an important reference standard. A loss function is designed that takes into account both the cross entropy loss and the F1 score. The goal is to ensure that the model has both good category probability predictions and a balanced precision and recall. Cross entropy loss is a commonly used loss function in classification problems. It measures the difference between the true label distribution and the predicted label distribution. The F1 score is the harmonic mean of precision and recall. The F1 score itself is not a loss function because its optimal value is 1, which the loss function hopes to minimize. Let 1 minus the F1 score be optimized as part of the loss, that is, LF1=1-F1Score. Combining these two, the overall loss function is:

[0080]

[0081]

[0082] Among them, TP is true positives and FP is false positives.

[0083]

[0084] Where FN is false negatives.

[0085] Step 3: The calculation of the F1 score involves non-differentiable operations (because it includes the counts of true positives, false positives, and false negatives), which makes it difficult to directly calculate the gradient of the loss function. The original memory-efficient zero-order gradient descent algorithm is replaced by a faster zero-order gradient descent algorithm RL-ZO (Algorithm 1), so that only forward propagation is required when training the model. The algorithm implementation is shown in Algorithm 1, which introduces reinforcement learning into the algorithm. The reward function is formula (6). By combining the F1 score and the cross entropy loss, a system containing rewards (when the F1 score is high) and penalties (when the cross entropy loss is high) is actually created. Subtracting the F1 score from the cross entropy loss and introducing a weight factor can merge the two indicators into a unified dimension, and set appropriate priorities for the two indicators by adjusting the values ​​of γ and λ.

[0086]

[0087] As shown in Algorithm 1, the reinforcement learning agent model uses a layer of transformer model. The transformer can pay attention to the connection between different parameters with its self-attention mechanism. The adapter parameter θ of the original model is input into the agent model, and the corresponding perturbation matrix z is output. We regard it as a sequence-to-sequence task, which can be well achieved through the encoder-decoder structure.

[0088] After the agent outputs the perturbation matrix z, it calculates the SPSA gradient estimate through two forward propagations, updates the A and B matrices, and performs large model fine-tuning and reinforcement learning model training alternately. During the reinforcement learning model training, the DDPG algorithm is used to update the parameters.

[0089]

[0090]

[0091] The detection method (RL-ZO) provided by the present invention was compared with other detection methods on a data set containing 8636 vulnerabilities, as shown in Table 1. Among all the compared methods, the method of the present invention showed the highest performance in both precision and F1 score, with a precision of 65.86%, which is 2.67%-14.77% higher than the precision of other methods, and an F1 score of 0.64. This shows that the RL-ZO method has high accuracy and stability in identifying software vulnerabilities. The reason why the RL-ZO method can achieve the best performance in both precision and F1 score is that we use the F1 score as part of the loss function and use reinforcement learning to solve the problem that the designed loss function cannot be back-propagated.

[0092] Table 1 Performance of all vulnerability detection methods

[0093] 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 RL-ZO (the present invention) 65.86 0.64

[0094] Those skilled in the art will readily appreciate other embodiments of the present invention after considering the specification and practicing the invention disclosed herein. This application is intended to cover any variations, uses or adaptations of the present invention that follow the general principles of the present invention and include common knowledge or customary techniques in the art that are not disclosed by the present invention. The specification and examples are to be considered exemplary only, and the true scope and spirit of the present invention are indicated by the following claims.

[0095] It should be understood that the present invention is not limited to what has been described above and that various modifications and changes may be made without departing from its scope. The scope of the present invention is limited only by the appended claims.

Claims

1. A software vulnerability detection method based on a Transformer architecture pre-training model, characterized in that: include: Based on the base model, add an adapter module on the self-attention matrix; During the forward propagation of the model, the original parameters of the pre-trained model are frozen, and only the parameters of the adapter are fine-tuned to minimize the loss function of the model; Reinforcement learning is used to calculate the perturbation matrix in the SPSA method to minimize the loss function of the detection model. Reinforcement learning is performed, and the parameters A and B of the adapter module are used as the input of the proxy model, and the perturbation matrices zA and zB are used as outputs. The loss function and F1 score are converted and used as a reward signal to guide the training process of the proxy model. By iteratively updating parameters A and B, the model learns more refined representations.

2. According to the software vulnerability detection method based on the Transformer architecture pre-training model according to claim 1, the base model is pre-trained on a large-scale code dataset.

3. The software vulnerability detection method based on the Transformer architecture pre-training model according to claim 1 is characterized in that: include: Step 1: Select the CodeGeex2 model as the base model and add the self-attention matrix W Q , W K , W V Add an adapter module to freeze the original parameters of the pre-trained model during the forward propagation of the base model, and only train the parameters of the LoRA adapter; Step 2: Apply the SPSA gradient estimation method to train the parameters of the LoRA adapter, and update the perturbation matrix through the reinforcement learning method, and finally update the adapter matrix only through forward propagation; The detection model is obtained as shown in formula (1); P=W0H+ΔWH=W0H+BAH (1) Where W0 represents the parameters of the pre-trained model, ΔW represents the adapter parameters, and H represents the input; Step 3: Merge the frozen parameters of the LoRA adapter and pre-trained model; Step 4: Input code H, through the projection matrix W q ,W k ,W v ,get: Q = HW q (2) K = HW k (3) V = HW v (4) Then through self-attention Through the encoder: A n = LN(Multi(Attention(Out enc n-1 )))+Out enc n-1 (6) Out enc n =LN(FeedForward(A n )+A n ) (7) Among them, Out enc n-1 It is the output of the n-1 layer encoder, which passes through the multi-head self-attention Multi and layer normalization LN, and then passes through the feedforward network FeedForward and residual link, and then layer normalization again; After the decoder: Among them, Q cross is the output of the decoder self-attention layer, Q dec ,K dec ,V dec The decoder input sequence is transformed into a hidden dimension, and CrossAttention calculates the encoder output layer and Q cross Cross-Attention, Out dec n-1 It is the self-attention output of the n-1 layer decoder, which is subjected to multi-head self-attention calculation and layer normalization for residual connection, and finally to layer normalization after passing through the feedforward network and residual connection; Finally, after average pooling and SoftMax function, the vulnerability category C is output; C=Softmax(AveragePolling(Out dec n )) (12)。 4. The software vulnerability detection method based on the Transformer architecture pre-training model according to claim 1 is characterized in that: Step 2 specifically includes: Step 21: The B and A matrices in the adapter are initialized with random Gaussian for A and zero for B using the LoRA method; Step 22: Modify the loss function of the model. In vulnerability prediction, the vulnerability prediction task is regarded as a classification task. The loss function considers both the cross entropy loss and the F1 score. The cross entropy loss is used to measure the difference between the true label distribution and the predicted label distribution. The F1 score is the harmonic mean of the precision and recall. Finally, the loss function is: Among them, TP is a true positive example and FP is a false positive example; Where FN is a false negative example; Step 23: Use the zero-order gradient descent algorithm RL-ZO so that only forward propagation is performed when training the model; introduce reinforcement learning into the zero-order gradient descent algorithm RL-ZO, and the reward function is shown in formula (17): Subtract the F1 score from the cross entropy loss and introduce a weight factor to merge the two indicators into a unified dimension. Then, set appropriate priorities for the two indicators by adjusting the values ​​of γ and λ. Step 24: The reinforcement learning agent model uses a layer of transformer model. The transformer uses its self-attention mechanism to focus on the relationship between different parameters. Through the encoder and decoder structure, the adapter parameters θ of the original model are input into the agent model, and the corresponding perturbation matrix z is output. After the proxy model outputs the perturbation matrix z, it calculates the SPSA gradient estimate through two forward propagations, updates the A and B matrices, and performs model fine-tuning and reinforcement learning model training alternately. During the reinforcement learning model training, the DDPG algorithm is used to update the parameters.

Images