A malware robust identification method based on contrastive learning
By training the model using a contrastive learning method, it generates similar representations for adversarial examples and clean samples, solving the problems of high feature engineering costs and vulnerability in existing technologies, and achieving efficient and robust identification of malware organizations.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- JINAN UNIVERSITY
- Filing Date
- 2022-07-08
- Publication Date
- 2026-06-05
AI Technical Summary
Existing technologies for identifying malware organizations require high human resources for feature engineering and are vulnerable to adversarial attacks, lacking research on adversarial attacks and model robustness.
A contrastive learning-based approach is adopted, which trains the model through an adversarial training module, a contrastive learning module, and a KL divergence module, enabling it to produce similar representations of clean samples and adversarial samples, thereby improving the model's robustness.
It significantly improves the accuracy of adversarial sample identification, enhances the model's ability to resist malware organizations, and has stronger anti-attack capabilities.
Smart Images
Figure CN115310083B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of software security technology, specifically to a robust malware identification method based on contrastive learning. Background Technology
[0002] With the proliferation of malware tools, Advanced Persistent Threats (APTs) are a particularly challenging and sophisticated attack method, representing one of the most significant advanced threats to cybersecurity. Global APT groups are focusing their attacks on industries such as medicine, cybersecurity, energy, and trade, with the scale and power of their cyber weapons continuously increasing. After detecting malware attacks, victims often seek to trace the organizations behind them and classify the malware families. Malware organization or family identification technologies are crucial for analyzing and tracing malicious attacks, defending against them, and deterring attackers, making them of significant research value.
[0003] There are currently some studies on the identification of organizations or families based on malware samples, but these studies have the following two shortcomings: (1) Feature engineering is required: Most existing organization identification technologies are based on machine learning (including deep learning) models, which require feature engineering such as extraction and representation of the organizational features of malware. Such feature engineering requires analysts to have expert knowledge such as reverse binary code analysis, resulting in high manpower and time costs; (2) Adversarial attack issues are not fully considered: Most existing organization identification methods are based on machine learning models, which are susceptible to adversarial attacks. However, there is currently no research on sample escape, adversarial attacks, and model robustness related to malware organization identification. Studying the adversarial attack and defense issues of end-to-end organization identification of APT malware can provide new ideas and methods for tracing the source of APT organizations, which is conducive to maintaining my country's cyberspace security and has important research significance and application value.
[0004] Therefore, a robust malware identification method based on contrastive learning is invented. Summary of the Invention
[0005] In view of the problems existing in the above and / or existing methods for robust identification of malware based on contrastive learning, the present invention is proposed.
[0006] Therefore, the purpose of this invention is to provide a robust malware identification method based on contrastive learning. This method trains a model that can generate similar representations for clean samples and adversarial samples. Adversarial samples are generated by performing adversarial transformations on clean samples. The purpose of these transformations is to deceive the model and cause it to make misjudgments. This invention utilizes three main modules—an adversarial training module, a contrastive learning module, and a KL divergence module—to perform adversarial training based on contrastive learning, thereby solving the aforementioned existing problems.
[0007] To address the aforementioned technical problems, according to one aspect of the present invention, the present invention provides the following technical solution:
[0008] A robust malware identification method based on contrastive learning, comprising a training phase and a testing phase;
[0009] The specific steps of the training phase are as follows:
[0010] Step 1: Preprocess the malware binary sample;
[0011] Step 2: Map the input vector to an embedding vector;
[0012] Step 3: Feed the embedding vector into two one-dimensional convolutional layers respectively, and perform a dot product operation on the output results of the two convolutional layers;
[0013] Step 4: Feed the computation result into max pooling, then through ReLU layer and fully connected layer to obtain the representation h;
[0014] Step 5: Feed the representation h into the fully connected layer to obtain the representation z;
[0015] Step Six: For the adversarial sample of malware binary, repeat steps one through five to obtain z. adv and h adv ;
[0016] Step 7: Combine h and h adv The results are fed into the fully connected layer, and the adversarial training loss is calculated using the total loss function formula. This step is also known as the adversarial training module.
[0017] Step 8: Combine z and z adv The data are fed into the projection head, and the output of the projection head is used to calculate the contrastive learning loss using the adversarial training loss function formula. This step is also known as the contrastive learning module.
[0018] Step 9: Combine h and h adv The results are fed into the fully connected layer, and the KL divergence loss is calculated using the contrastive learning loss function formula. This step is also known as the KL divergence loss module.
[0019] Step 10: Sum the adversarial training loss, contrastive learning loss, and KL divergence loss using the KL divergence loss function formula and backpropagate to update the model parameter gradients;
[0020] The specific process for the testing phase is as follows:
[0021] Step 1: Preprocess the malware binary sample;
[0022] Step 2: Map the input vector to an embedding vector;
[0023] Step 3: Feed the embedding vector into two one-dimensional convolutional layers respectively, and perform a dot product operation on the output results of the two convolutional layers;
[0024] Step 4: The computation result is fed into max pooling, passing through a ReLU layer and a fully connected layer to obtain the representation h;
[0025] Step 5: Feed the representation h into the fully connected layer to obtain the classification results of the samples.
[0026] As a preferred embodiment of the robust malware identification method based on contrastive learning described in this invention, in step one, the preprocessing of the malware binary samples is as follows: by padding the end of the file with 0s, the length of each file is transformed to L bytes, where L is an integer greater than the maximum actual length of the input files in all the datasets of this embodiment. Through the above preprocessing, each input vector has a fixed length.
[0027] As a preferred embodiment of the robust malware identification method based on contrastive learning described in this invention, in step two, the process of embedding vectors is as follows: the word embedding layer of the model uses an embedding layer word vector matrix of size [256, 8] to map each input vector to an embedding vector. The word embedding operation is equivalent to a matrix lookup operation. By retrieving the embedding layer word vector matrix, each byte in the input vector is mapped to an 8-dimensional vector.
[0028] As a preferred embodiment of the robust malware identification method based on contrastive learning described in this invention, in step three, the convolution and dot product operations of the embedded vector are as follows: the embedded vector is fed into two one-dimensional convolutional layers respectively, one of which learns the representation of bytes, and the other convolutional layer combines with the sigmoid activation function to form an attention layer, and the representation is filtered by performing a dot product calculation with the representation learned by the other convolutional layer.
[0029] As a preferred embodiment of the robust malware identification method based on contrastive learning described in this invention, the total loss function formula is:
[0030] .
[0031] As a preferred embodiment of the robust malware identification method based on contrastive learning described in this invention, the adversarial training loss function formula is:
[0032] .
[0033] As a preferred embodiment of the robust malware identification method based on contrastive learning described in this invention, the contrastive learning loss function formula is:
[0034] .
[0035] As a preferred embodiment of the robust malware identification method based on contrastive learning described in this invention, the KL divergence loss function formula is:
[0036] ;
[0037] in, , and The values are 1, 0.3, and 0.3 respectively.
[0038] Compared with existing technologies:
[0039] 1. To illustrate the effectiveness and superiority of the present invention, the following evaluation metrics used in the prior art are adopted: Standard Accuracy (SA), used to verify the ability to identify clean sample tissues, is obtained by weighting and summing the metrics according to the proportion of each tissue sample in the total number of clean samples; Robust Accuracy (RA), used to verify the ability to identify adversarial sample tissues, is obtained by calculating the accuracy for all tissue categories on the adversarial sample test set and then weighting and summing the results, and the calculation method is consistent with the calculation method of SA.
[0040] 2. Experiments were conducted on two malware datasets: the dataset named "small" contains 1759 samples from 9 organizations, and the dataset named "big" contains 5133 samples from 12 organizations. The SA value obtained for the "small" test set was 92.9%, and the SA value obtained for the "big" test set was 92.0%, demonstrating the excellent effectiveness of the present invention in identifying APTs on two clean sample sets. In this experiment, the RA value obtained by the present invention for the "small" test set was 89.7%, and the RA value obtained for the "big" adversarial dataset was 86.7%. In contrast, the APT organization identification model trained only with MConv only achieved RA values of only 39.4% and 26.5% for the two adversarial datasets, respectively. The present invention achieved significant improvements in RA metrics of 50.3% and 60.2%, respectively, demonstrating that the present invention has a significantly better defense capability against adversarial samples.
[0041] 3. This invention includes a training phase and a testing phase. In the training phase, the convolutional neural network model MConv is used, combined with three modules (adversarial training module, contrastive learning module, and KL divergence module) to perform adversarial training on MConv based on contrastive learning. In the testing phase, MConv and fully connected layers are used for sample identification. This invention combines the advantages of deep learning convolutional neural network models with adversarial training, contrastive learning, and KL divergence applications to achieve robust end-to-end identification of malware organizations (or families). At the same time, this invention fully considers the possibility of identifying malware evasion organizations and targets the possibility of adversarial attacks aimed at evasion, which is more in line with the real needs of cyberspace and has more practical application significance. Moreover, the method shows good identification and defense against adversarial attacks on various evaluation indicators. Attached Figure Description
[0042] Figure 1 This is the MConv model of the present invention;
[0043] Figure 2 This is the training phase of the end-to-end robust identification method RMI proposed in this invention;
[0044] Figure 3 This is the testing phase of the end-to-end robust identification method RMI proposed in this invention. Detailed Implementation
[0045] To make the objectives, technical solutions, and advantages of the present invention clearer, the embodiments of the present invention will be described in further detail below with reference to the accompanying drawings.
[0046] This invention provides a robust malware identification method based on contrastive learning. Please refer to [link / reference]. Figures 1-3This includes the training phase and the testing phase;
[0047] like Figure 2 As shown, the specific steps of the training phase are as follows:
[0048] Step 1: Preprocess the malware binary sample;
[0049] The preprocessing process for the binary samples of malware is as follows: by padding the end of the file with 0s, the length of each file is transformed to L bytes, where L is an integer greater than the maximum actual length of the input files in all the datasets of this embodiment. Through the above preprocessing, each input vector has a fixed length.
[0050] Step 2: Map the input vector to an embedding vector;
[0051] The process of embedding vectors is as follows: The word embedding layer of the model uses an embedding layer word vector matrix of size [256, 8] to map each input vector (i.e. a byte sequence of length L) to an embedding vector. The word embedding operation is equivalent to a matrix lookup operation. By searching the embedding layer word vector matrix, each byte in the input vector is mapped to an 8-dimensional vector.
[0052] Step 3: Feed the embedding vector into two one-dimensional convolutional layers respectively, and perform a dot product operation on the output results of the two convolutional layers;
[0053] The convolution and dot product operations of the embedded vector are as follows: the embedded vector is fed into two one-dimensional convolutional layers. One convolutional layer learns the representation of the byte, and the other convolutional layer combines the sigmoid activation function to form an attention layer. The representation is filtered by performing a dot product calculation with the representation learned by the other convolutional layer.
[0054] Step 4: Feed the computation result into max pooling, then through ReLU layer and fully connected layer to obtain the representation h;
[0055] Step 5: Feed the representation h into the fully connected layer to obtain the representation z;
[0056] Step Six: For the adversarial sample of malware binary, repeat steps one through five to obtain z. adv and h adv ;
[0057] Step 7: Combine h and h adv The results are fed into the fully connected layer, and the adversarial training loss is calculated using the total loss function formula. This step is also known as the adversarial training module.
[0058] The formula for the total loss function is as follows:
[0059] ;
[0060] Step 8: Combine z and z adv The data are fed into the projection head, and the output of the projection head is used to calculate the contrastive learning loss using the adversarial training loss function formula. This step is also known as the contrastive learning module.
[0061] The formula for the adversarial training loss function is as follows:
[0062] ;
[0063] Step 9: Combine h and h adv The results are fed into the fully connected layer, and the KL divergence loss is calculated using the contrastive learning loss function formula. This step is also known as the KL divergence loss module.
[0064] The formula for the contrastive learning loss function is as follows:
[0065] ;
[0066] Step 10: Sum the adversarial training loss, contrastive learning loss, and KL divergence loss using the KL divergence loss function formula and backpropagate to update the model parameter gradients;
[0067] The KL divergence loss function formula is as follows:
[0068] ;
[0069] in, , and The values are 1, 0.3, and 0.3 respectively.
[0070] Steps one, two, three, and four represent inputting the sample as follows: Figure 1 The MConv shown.
[0071] like Figure 3 As shown, the specific process of the testing phase is as follows:
[0072] Step 1: Preprocess the malware binary sample;
[0073] Step 2: Map the input vector to an embedding vector;
[0074] Step 3: Feed the embedding vector into two one-dimensional convolutional layers respectively, and perform a dot product operation on the output results of the two convolutional layers;
[0075] Step 4: The computation result is fed into max pooling, passing through a ReLU layer and a fully connected layer to obtain the representation h;
[0076] Step 5: Feed the representation h into the fully connected layer to obtain the classification results of the samples.
[0077] Working Principle: This invention consists of a training phase and a testing phase. In the training phase, the samples are first preprocessed and then input into MConv (steps one, two, three, and four). Then, three main modules (adversarial training module, contrastive learning module, and KL divergence module) are used to perform adversarial training on the model based on contrastive learning. Each module uses a different loss function. The sum of the losses from the three modules is used as the total loss for backpropagation to update the model parameters. The training phase proceeds as follows: Figure 2 As shown;
[0078] The coefficients before the loss functions of the three modules are hyperparameters, mainly used to adjust the weight of the loss of each module;
[0079] The steps of the adversarial training module are as follows: take clean samples and adversarial samples and input them into MConv and the fully connected layer respectively to obtain confidence scores, and use cross-entropy to calculate the adversarial training loss;
[0080] The steps of the contrastive learning module are as follows: The representations h obtained from the adversarial training model, which represent clean samples and adversarial samples respectively, are fed into the projection head. The results are used to calculate the contrastive loss using the NT-Xent loss function. The projection head also updates its parameters during the training process.
[0081] The steps of the KL divergence loss module are as follows: directly calculate the KL divergence loss on the representations obtained from the adversarial training model, which represent clean samples and adversarial samples respectively.
[0082] experiment:
[0083] This invention was tested on two malware organization datasets: the "small" dataset contains 1759 samples from 9 organizations, and the "big" dataset contains 5133 samples from 12 organizations. The SA value obtained for the "small" test set was 92.9%, and the SA value obtained for the "big" test set was 92.0%, demonstrating the excellent identification effectiveness of this invention for APTs on two clean sample sets. In this experiment, the RA value obtained by this invention for the adversarial sample set generated from the "small" test set was 89.7%, and the RA value obtained for the adversarial sample set generated from the "big" test set was 86.7%. In contrast, the organization identification model trained only with MConv only achieved RA values of only 39.4% and 26.5% for the two adversarial sample sets, respectively. This invention achieved significant improvements in RA metrics of 50.3% and 60.2%, respectively, demonstrating that this invention has a significantly better defense capability against adversarial samples.
[0084] Although the present invention has been described above with reference to embodiments, various modifications can be made and components can be replaced with equivalents without departing from the scope of the invention. In particular, as long as there is no structural conflict, the features in the disclosed embodiments can be combined with each other in any manner. The lack of an exhaustive description of these combinations in this specification is merely for the sake of brevity and resource conservation. Therefore, the present invention is not limited to the specific embodiments disclosed herein, but includes all technical solutions falling within the scope of the claims.
Claims
1. A robust malware identification method based on contrastive learning, characterized in that, Includes training and testing phases; The specific steps of the training phase are as follows: Step 1: Preprocess the malware binary sample; Step 2: Map the input vector to an embedding vector; Step 3: Feed the embedding vector into two one-dimensional convolutional layers respectively, and perform a dot product operation on the output results of the two convolutional layers; Step 4: Feed the computation result into max pooling, then through ReLU layer and fully connected layer to obtain the representation h; Step 5: Feed the representation h into the fully connected layer to obtain the representation z; Step Six: For the adversarial sample of malware binary, repeat steps one through five to obtain z. adv and h adv ; Step 7: Combine h and h adv The results are fed into the fully connected layer, and the adversarial training loss is calculated using the total loss function formula. This step is also known as the adversarial training module. Step 8: Combine z and z adv The data are fed into the projection head, and the output of the projection head is used to calculate the contrastive learning loss using the adversarial training loss function formula. This step is also known as the contrastive learning module. Step 9: Combine h and h adv The results are fed into the fully connected layer, and the KL divergence loss is calculated using the contrastive learning loss function formula. This step is also known as the KL divergence loss module. Step 10: Sum the adversarial training loss, contrastive learning loss, and KL divergence loss using the KL divergence loss function formula and backpropagate to update the model parameter gradients; The specific process for the testing phase is as follows: Step 1: Preprocess the malware binary sample; Step 2: Map the input vector to an embedding vector; Step 3: Feed the embedding vector into two one-dimensional convolutional layers respectively, and perform a dot product operation on the output results of the two convolutional layers; Step 4: The computation result is fed into max pooling, passing through a ReLU layer and a fully connected layer to obtain the representation h; Step 5: Feed the representation h into the fully connected layer to obtain the classification result of the sample; The formula for the total loss function is: ； The formula for the adversarial training loss function is: ； The formula for the contrastive learning loss function is: ； The formula for the KL divergence loss function is: ； in, , and The values are 1, 0.3, and 0.3 respectively.
2. The robust malware identification method based on contrastive learning according to claim 1, characterized in that, In step one, the preprocessing of the malicious binary sample is as follows: by padding the end of the file with 0s, the length of each file is transformed to L bytes, where L is an integer greater than the maximum actual length of the input files in all datasets. Through the above preprocessing, each input vector has a fixed length.
3. The robust malware identification method based on contrastive learning according to claim 1, characterized in that, In step two, the process of embedding vectors is as follows: the word embedding layer of the model uses an embedding layer word vector matrix of size [256, 8] to map each input vector to an embedding vector. The word embedding operation is equivalent to a matrix lookup operation. By searching the embedding layer word vector matrix, each byte in the input vector is mapped to an 8-dimensional vector.
4. The robust malware identification method based on contrastive learning according to claim 1, characterized in that, In step three, the convolution and dot product operations of the embedded vector are as follows: the embedded vector is fed into two one-dimensional convolutional layers respectively. One convolutional layer learns the representation of the byte, and the other convolutional layer combines the sigmoid activation function to form an attention layer. The representation is filtered by performing a dot product calculation with the representation learned by the other convolutional layer.