Malicious software classification method based on cooperation of pseudo tag optimization and main and auxiliary models

Abstract

The invention provides a false label optimization and main and auxiliary model collaborative malicious software classification method, which comprises the following steps of: constructing and pre-training a malicious software classification model based on a DAN-OfficientNet network as a main model and a false label generation model based on a DAN-MobileNet network as an auxiliary model; a new unknown malicious software sample is subjected to pseudo-label generation through a pseudo-label generation model, and then is screened through a Bayesian controller to be used for retraining of a malicious software classification model; and finally, carrying out final updating training on the retrained malicious software classification model and the pre-trained pseudo label generation model by utilizing the real label of the new malicious software sample. According to the method, the main model is helped to be retrained through the main and auxiliary model collaboration framework, so that the influence of new malicious software on the model performance is resisted; and quality control is carried out on the pseudo labels through a Bayesian controller, so that the low-quality pseudo labels are prevented from influencing the retraining performance of the master model.

Description

Technical Field

[0001] The present invention relates to the technical field of malware classification, and in particular to a malware classification method combining pseudo-label optimization with primary and secondary models. Background Art

[0002] Image-based malware classification uses modern computer vision and deep learning techniques to convert malware binaries into images and perform feature extraction and analysis. The images are passed through deep learning models, such as convolutional neural networks (CNNs), to obtain the final classification output, and the model training requires a large number of labeled samples.

[0003] As time goes by, a large number of new unknown malware appear. Facing new malware, the performance of the model will continue to decline without training. At this time, retraining with new malware samples can maintain the high performance of the model.

[0004] However, the existing classifier uses machine learning RF to classify malware, which is difficult to capture complex nonlinear features, performs poorly on high-dimensional data, and is difficult to expand to large-scale data. In addition, the pseudo-labels of the existing technology need to be generated quickly to meet the needs of retraining the main model, but the quality of pseudo-label generation is not controlled, which seriously affects the performance of the model. Summary of the invention

[0005] In view of the deficiencies in the prior art, the present invention provides a malware classification method that optimizes pseudo labels and collaborates with primary and secondary models. The present invention trains the primary model through screened pseudo labels to resist the impact of the emergence of new malware on the performance of the primary model.

[0006] The technical solution of the present invention is: a malware classification method based on pseudo-label optimization and collaboration between primary and secondary models, comprising the following steps:

[0007] S1), build and pre-train a malware classification model based on the DAN-EfficientNet network as the main model; use the malware classification model to perform high-precision classification of known malware image samples;

[0008] S2) Build and pre-train a pseudo-label generation model based on the DAN-MobileNet network as a secondary model; when new malware appears, generate pseudo-labels for unknown malware samples through the pseudo-label generation model;

[0009] S3), using a Bayesian controller to perform uncertainty assessment and screening on the quality of pseudo labels generated by the pseudo label generation model, and optimize the credibility of the pseudo labels;

[0010] S4), using the high-quality pseudo labels selected in step S3) to retrain the malware classification model to form a preliminary updated detection model;

[0011] S5), when the real label of new malware is detected, the pseudo label is replaced and the malware classification model retrained in step S4) and the pseudo label generation model pre-trained in step S2) are finally updated and trained using the real label of the new malware sample to obtain the final malware classification model and pseudo label generation model.

[0012] Preferably, the malware classification model and pseudo-label generation model introduce a distribution alignment normalization mechanism DAN, which improves the robustness of the malware classification model and pseudo-label generation model when new malware appears by dynamically adjusting and normalizing the distribution of input data.

[0013] Preferably, the distribution alignment normalization mechanism DAN specifically comprises the following steps:

[0014] A1) Input feature standardization

[0015] The input feature map A is standardized according to the mean μ A and standard deviation σ A Calculate the standardized feature map A′, that is:

[0016]

[0017] A2) Feature Statistics Recalibration

[0018] The standard deviation and mean are dynamically adjusted to obtain the recalibrated A″, that is:

[0019] A″=σ A ′ A′+μ ′ A

[0020] In the formula, σ A ′ , μ ′ A are the dynamically adjusted standard deviation and mean, respectively.

[0021] Preferably, in step S1), the malware classification model based on the DAN-EfficientNet network includes multiple convolutional layers, a DAN module, multiple stacked first modules, second modules, third modules, and a pooling layer, a fully connected layer, and a global pooling layer.

[0022] Preferably, in step S1), the malware classification model based on the DAN-EfficientNet network takes the malware image as input, and then performs a DAN operation through the DAN module after passing through a convolution layer with a convolution kernel of 3x3, and then extracts features through multiple stacked first modules, second modules, and third modules after processing with a Swish activation function. The extracted features are processed through a convolution layer, a pooling layer, a fully connected layer, a global pooling layer and a Softmax function to obtain a classification result.

[0023] Preferably, in step S1), the first module includes a depthwise separable convolutional layer, a DAN module, a SE module, and a convolutional layer; the spatial and channel features are extracted through a depthwise separable convolutional layer with a convolution kernel of 3x3, and then the DAN module performs distribution alignment normalization through the DAN mechanism, and then the gradient flow is improved using the Swish activation function; then the correlation between channels is enhanced through the SE module, and the representation of the feature map is dynamically adjusted by learning the channel weights, and then the 1x1 convolutional layer and the DAN mechanism are used again to obtain the initial output, and the initial output is jump-connected with the input of the first module to obtain the final output.

[0024] Preferably, in step S1), the second module and the third module are similar in structure to the first module, and the input of the second module and the third module must first pass through a convolution layer with a convolution kernel of 1x1 and the convolution kernel of the depth-separable convolution layer of the third module is 5x5.

[0025] Preferably, in step S2), the pseudo-label generation model based on the DAN-MobileNet network includes a convolutional layer, a DAN module, a backbone network, and an average pooling layer; the input enters the backbone network after passing through a 3x3 convolutional layer, a DAN module, and an H-Swish activation function in sequence; the output of the backbone network passes through a 1×1 convolutional layer, a DAN module, an H-Swish activation function, an average pooling layer, a 1×1 convolutional layer, an H-Swish activation function, a 1×1 convolutional layer, and enters the Softmax activation function to output the final classification probability distribution.

[0026] Preferably, in step S2), the backbone network is composed of several Bneck modules stacked together, and each of the Bneck modules is composed of a 1×1 convolution layer, a depthwise separable convolution layer, and a DAN module; some of the Bneck modules also include an SE module, in which a 1×1 convolution layer is first used to perform a dimensionality increase operation, and then a depthwise separable convolution layer is used to perform a depthwise separable convolution operation, and feature extraction and feature combination are performed separately, first using a 3×3 depthwise separable convolution layer to perform feature extraction by channel, and then activated by the DAN module and the nonlinear activation function NL, and the SE module is used to enable the network to focus on more important channel features, and the obtained weights are multiplied by the features after activation by the nonlinear activation function NL, and finally residual jump connections are performed after being processed by a 1×1 convolution kernel.

[0027] Preferably, in step S1), the SE module includes an average pooling layer and two fully connected layers, a RELU function is connected between the average pooling layer and the first fully connected layer, and an output of the second fully connected layer is connected to an H-Swish activation function.

[0028] Preferably, in step S3), the quality of the pseudo-labels generated by the pseudo-label generation model is evaluated for uncertainty using a Bayesian controller, specifically:

[0029] S31), pseudo labels generated by pseudo label generation model It is expressed as:

[0030]

[0031] In the formula, f dan-mobilenet represents the pseudo-label generation model; x represents a malware sample;

[0032] S32) Use Bayesian reasoning to calculate the credibility of each pseudo label, based on the sample x and the pseudo label Calculate the posterior probability So as to evaluate the credibility of the pseudo-label, namely:

[0033]

[0034] In the formula, Indicates that the sample is a pseudo label of x Indicates that the pseudo label is The probability of sample x appearing under the condition of ; represents the prior probability of the category pseudo label; P(x) represents the marginal probability;

[0035] S33), preset quality threshold τ, when the posterior probability Greater than the preset quality threshold τ, the pseudo label If it is credible, otherwise, the pseudo label is discarded.

[0036] The beneficial effects of the present invention are:

[0037] 1. This paper proposes an image-based malware classification main and secondary model collaborative framework to help malware classification models be retrained to resist the impact of new malware on model performance;

[0038] 2. The present invention uses DAN to improve EfficientNet and MobileNet, reducing the impact of model performance when data distribution changes; improving the classification ability and label generation ability of malware classification model and pseudo-label generation model;

[0039] 3. The Bayesian controller of the present invention performs quality control on pseudo labels, further optimizing the credibility of pseudo labels and preventing low-quality pseudo labels from affecting the retraining performance of the model. BRIEF DESCRIPTION OF THE DRAWINGS

[0040] Figure 1 It is a schematic diagram of the process of the present invention;

[0041] Figure 2 It is a framework diagram of the malware classification model based on the DAN-EfficientNet network of the present invention;

[0042] Figure 3 It is a structural framework diagram of the first module, the second module and the third module of the present invention;

[0043] Figure 4 This is a structural framework diagram of the pseudo-label generation model based on the DAN-MobileNet network of the present invention;

[0044] Figure 5 It is a framework diagram of the Bneck module of the present invention;

[0045] Figure 6 It is a structural schematic diagram of the SE module of the present invention. DETAILED DESCRIPTION

[0046] The specific implementation of the present invention will be further described below in conjunction with the accompanying drawings:

[0047] like Figure 1 As shown, the present embodiment provides a malware classification method that combines pseudo-label optimization with primary and secondary models, including the following steps:

[0048] S1), build and pre-train a malware classification model based on the DAN-EfficientNet network as the main model; use the malware classification model to perform high-precision classification of known malware image samples;

[0049] S2) Build and pre-train a pseudo-label generation model based on the DAN-MobileNet network as a secondary model; when new malware appears, generate pseudo-labels for unknown malware samples through the pseudo-label generation model;

[0050] S3), using a Bayesian controller to perform uncertainty assessment and screening on the quality of pseudo labels generated by the pseudo label generation model, and optimize the credibility of the pseudo labels;

[0051] S4), using the high-quality pseudo labels selected in step S3) to retrain the malware classification model to form a preliminary updated detection model;

[0052] S5), when the real label of new malware is detected, the pseudo label is replaced and the malware classification model retrained in step S4) and the pseudo label generation model pre-trained in step S2) are finally updated and trained using the real label of the new malware sample to obtain the final malware classification model and pseudo label generation model.

[0053] In this embodiment, in the initial stage, the malware classification model and the pseudo-label generation model are pre-trained for the first time using a labeled malware dataset. When new malware appears, a pseudo-label is generated using the pseudo-label generation model, the pseudo-label is quality screened using a Bayesian controller, and the screened pseudo-label malware classification model is retrained. When a sandbox or manually detects a real label of new malware, the pseudo-label is replaced and the malware classification model and the pseudo-label generation model are updated and trained using new malware samples.

[0054] As a preferred embodiment of this embodiment, the malware classification model and pseudo-label generation model introduce a distribution alignment normalization mechanism DAN to cope with the performance fluctuation problem caused by different data distributions.

[0055] The distribution alignment normalization mechanism DAN improves the robustness of the malware classification model and the pseudo-label generation model in the presence of new malware by dynamically adjusting and normalizing the distribution of input data, thereby slowing down the performance degradation of the malware classification model and improving the quality of pseudo-labels generated by the pseudo-label generation model. Specifically, the following steps are included:

[0056] A1) Input feature standardization

[0057] Standardize the input feature map A and calculate the mean μ A and standard deviation σ A , and get the standardized feature map A′, that is:

[0058]

[0059] Where H and W are the height and width of feature map A respectively; A h,w It represents the pixel value or feature value at the hth row and wth column in the input feature map A; h and w represent the index of the height and width of the feature map, that is, the row number and column number in the two-dimensional coordinate system;

[0060] A2) Feature Statistics Recalibration

[0061] The standard deviation and mean are dynamically adjusted to obtain the recalibrated A″, that is:

[0062] A″=σ′ A A′+μ′ A

[0063] In the formula, σ′ A , μ′ A are the dynamically adjusted standard deviation and mean, respectively.

[0064] in,

[0065] μ′ A =f(μ A ,σ A )μ A

[0066] σ′ A =g(μ A ,σ A )σ A

[0067] g(μ A ,σ A )=Sigmoid(BN(Conv1D(μ A ,σ A )))

[0068] f(μ A ,σ A )=Sigmoid(BN(Conv1D(μ A ,σ A )))

[0069] In the formula, f(μ A ,σ A )、g(μ A ,σ A ) represent the dynamic adjustment factors used to adjust the mean and standard deviation respectively; BN represents the BatchNorm1D normalization operation; Sigmoid represents an activation function that maps the input value to the (0,1) interval; Conv1D represents a one-dimensional convolution operation.

[0070] As preferred in this embodiment, Figure 2As shown, the malware classification model based on the DAN-EfficientNet network includes multiple convolutional layers, DAN modules, multiple stacked first modules, second modules, third modules, and pooling layers, fully connected layers, and global pooling layers. Specifically, in this embodiment, from the input end to the output end, there are sequentially connected a convolutional layer with a convolution kernel of 3x3, a DAN module, a Swish activation function, a first module, a second module, a third module, a second module, a third module, a third module, a second module, a convolutional layer with a convolution kernel of 1x1, a pooling layer, a fully connected layer, a global pooling layer, and a Softmax function.

[0071] In this embodiment, the malware classification model based on the DAN-EfficientNet network takes the malware image as input, then passes through the first 3x3 convolution layer and performs DAN operation through the DAN module, and then uses the Swish activation function to process and then passes through a plurality of stacked first modules, second modules, and third modules for feature extraction. The extracted features are processed through the convolution layer, the pooling layer, the fully connected layer, the global pooling layer and the Softmax function to obtain the classification result.

[0072] In this embodiment, the Swish activation function is defined as: swish(x)=x*sigmoid(x), where x is a malware image; this embodiment improves gradient flow through the swish activation function.

[0073] As preferred in this embodiment, Figure 3 As shown, the first module includes a depth-separable convolution layer, a DAN module, a SE module, and a convolution layer; the spatial and channel features are extracted through a depth-separable convolution layer with a convolution kernel size of 3x3, and then the DAN module performs distribution alignment normalization through the DAN mechanism, and then the Swish activation function is used to enhance the gradient flow; then the SE module is used to enhance the correlation between channels, and the representation of the feature map is dynamically adjusted by learning the channel weights, and then the 1x1 convolution and the DAN mechanism are used again to obtain the initial output, and the initial output is jump-connected with the input of the first module to obtain the final output.

[0074] In this embodiment, Figure 3 As shown, the second module and the third module are similar in structure to the first module. The input of the second module and the third module must first pass through a 1x1 convolution layer and the convolution kernel of the depth-separable convolution layer of the third module is 5x5.

[0075] As an embodiment preferably, Figure 4As shown, the pseudo-label generation model based on the DAN-MobileNet network includes a convolution layer, a DAN module, a backbone network, and an average pooling layer; the input new malware passes through a 3x3 convolution layer, a DAN module, and an H-Swish activation function in sequence before entering the backbone network; the output of the backbone network passes through a 1×1 convolution layer, a DAN module, an H-Swish activation function, an average pooling layer, a 1×1 convolution layer, an H-Swish activation function, a 1×1 convolution layer, and enters the Softmax activation function before outputting the final classification probability distribution.

[0076] In this embodiment, the backbone network is composed of 11 stacked Bneck modules, and a 1×1 convolution layer is connected after the Bneck module stack, and the important information in the global features is further strengthened through the second DAN module; the H-Swish activation function is used after the second DAN module to introduce nonlinear transformation to the features and enhance the classification effect. After H-Swish activation, average pooling is performed to compress the spatial features into channel-level feature vectors. After that, a 1×1 convolution layer is used for dimensionality increase, and an H-Swish activation function is used, and finally a 1×1 convolution layer is used to obtain the Logits of each category. The Logits of each category are input into Softmax to obtain the final probability distribution of the category.

[0077] In this embodiment, Figure 5 As shown, each of the Bneck modules is composed of a 1×1 convolution layer, a depth-separable convolution layer, and a DAN module. Some of the Bneck modules also include a SE module. In the Bneck module, a 1×1 convolution layer is first used to perform a dimensionality increase operation, and then a depth-separable convolution is performed to separate feature extraction and feature combination. A 3×3 depth-separable convolution layer is first used to extract features by channel, and then activated by the DAN module and the nonlinear activation function NL. The SE module enables the network to focus on more important channel features, and the obtained weights are multiplied by the features after activation by the nonlinear activation function NL, and finally processed by a 1×1 convolution kernel for residual jump connection.

[0078] In this embodiment, the second and third Bneck modules do not include SE modules, and the nonlinear activation functions NL of the first, second and third Bneck modules adopt Relu activation functions; the nonlinear activation functions NL of the remaining Bneck modules adopt H-Swish activation functions.

[0079] As an embodiment preferably, Figure 6As shown, the SE module includes an average pooling layer and two fully connected layers. A RELU function is connected between the average pooling layer and the first fully connected layer, and the output of the second fully connected layer is connected to an H-Swish activation function.

[0080] As a preferred embodiment of this invention, a Bayesian controller is used to perform uncertainty assessment and screening on the quality of pseudo labels generated by the pseudo label generation model, as follows:

[0081] S31), pseudo labels generated by pseudo label generation model It is expressed as:

[0082]

[0083] In the formula, f dan-mobilenet represents the pseudo-label generation model; x represents a malware image sample;

[0084] S32) Use Bayesian reasoning to calculate the credibility of each pseudo label, based on the sample x and the pseudo label Calculate the posterior probability So as to evaluate the credibility of the pseudo-label, namely:

[0085]

[0086] In the formula, Indicates that the sample is a pseudo label of x Indicates that the pseudo label is The probability of sample x appearing under the condition of ; represents the prior probability of the category pseudo label; P(x) represents the marginal probability;

[0087] In this embodiment, the marginal probability P(x) is calculated as:

[0088]

[0089] Where n is the number of categories, is the pseudo label for the i-th category The conditional probability of sample x under , is also approximated by the classification probability output by DAN-MobileNet’s Softmax. is the pseudo label of the i-th category The prior probability of is also estimated by counting the frequency of each category in the training data.

[0090] S33), preset quality threshold τ, when the posterior probability Greater than the preset quality threshold τ, the pseudo label If it is credible, it can be used to retrain the malware classification model; otherwise, the pseudo label is discarded.

[0091] The above embodiments and descriptions are only for illustrating the principles and best embodiments of the present invention. Without departing from the spirit and scope of the present invention, the present invention may be subject to various changes and improvements, all of which fall within the scope of the present invention to be protected.

Claims

1. A malware classification method based on pseudo-label optimization and collaboration between primary and secondary models, characterized in that: The following steps are involved: S1), build and pre-train a malware classification model based on the DAN-EfficientNet network as the main model; S2) Build and pre-train a pseudo-label generation model based on the DAN-MobileNet network as a secondary model; when new malware appears, generate pseudo-labels for unknown malware samples through the pseudo-label generation model; S3), using a Bayesian controller to perform uncertainty assessment and screening on the quality of pseudo labels generated by the pseudo label generation model, and optimize the credibility of the pseudo labels; S4), retraining the malware classification model using the pseudo labels selected in step S3); S5), when the real label of the new malware is detected, the malware classification model retrained in step S4) and the pseudo-label generation model pre-trained in step S2) are finally updated and trained using the real label of the new malware sample to obtain the final malware classification model and pseudo-label generation model.

2. According to claim 1, a malware classification method combining pseudo-label optimization and primary and secondary models, characterized in that: The malware classification model and pseudo-label generation model introduce a distribution alignment normalization mechanism DAN, and the distribution of input data is dynamically adjusted and normalized through the distribution alignment normalization mechanism DAN.

3. According to claim 2, a malware classification method combining pseudo-label optimization and primary and secondary models, characterized in that: The distribution alignment normalization mechanism DAN dynamically adjusts and normalizes the distribution of input data, specifically including the following steps: A1) Input feature standardization The input feature map A is standardized according to the mean μ A and standard deviation v A Calculate the standardized feature map A′, that is: A2) Feature Statistics Recalibration The standard deviation and mean are dynamically adjusted to obtain the recalibrated A″, that is: A″=σ′ A A′+μ′ A In the formula, σ′ A , μ′ A are the dynamically adjusted standard deviation and mean, respectively.

4. According to claim 1, a malware classification method combining pseudo-label optimization and primary and secondary models, characterized in that: In step S1), the malware classification model based on the DAN-EfficientNet network includes multiple convolutional layers, a DAN module, a plurality of stacked first modules, a second module, a third module, a pooling layer, a fully connected layer, and a global pooling layer.

5. According to claim 4, a malware classification method combining pseudo-label optimization and primary and secondary models, characterized in that: In step S1), the first module includes a depth-separable convolutional layer, a DAN module, a SE module, and a convolutional layer; the spatial and channel features are extracted through a depth-separable convolutional layer with a convolution kernel of 3x3, and then the DAN module performs distribution alignment normalization through the DAN mechanism, and then the Swish activation function is used to enhance the gradient flow; then the SE module is used to enhance the correlation between channels, and the representation of the feature map is dynamically adjusted by learning the channel weights, and then the 1x1 convolutional layer and the DAN mechanism are used again to obtain the initial output, and the initial output is jump-connected with the input of the first module to obtain the final output.

6. The malware classification method of pseudo-label optimization and primary and secondary model collaboration according to claim 5 is characterized by: In step S1), the second module and the third module have similar structures to the first module. The input of the second module and the third module must first pass through a convolution layer with a convolution kernel of 1x1 and the convolution kernel of the depth-separable convolution layer of the third module is 5x5.

7. The malware classification method of claim 1 that combines pseudo-label optimization with primary and secondary models is characterized by: In step S2), the pseudo-label generation model based on the DAN-MobileNet network includes a convolution layer, a DAN module, a backbone network, and an average pooling layer; the input new malware image sample enters the backbone network after passing through a 3x3 convolution layer, a DAN module, and an H-Swish activation function in sequence; the output of the backbone network passes through a 1×1 convolution layer, a DAN module, an H-Swish activation function, an average pooling layer, a 1×1 convolution layer, an H-Swish activation function, a 1×1 convolution layer, and a Softmax activation function in sequence to output the final classification probability distribution.

8. The malware classification method of pseudo-label optimization and primary and secondary model collaboration according to claim 7 is characterized by: In step S2), the backbone network is composed of several Bneck modules stacked together, and each of the Bneck modules is composed of a 1×1 convolution layer, a depthwise separable convolution layer, and a DAN module; some of the Bneck modules also include an SE module, in which a 1×1 convolution layer is first used to perform a dimensionality increase operation, and then a depthwise separable convolution layer is used to perform a depthwise separable convolution operation, and feature extraction and feature combination are performed separately. A 3×3 depthwise separable convolution layer is first used to perform feature extraction by channel, and then activated by the DAN module and the nonlinear activation function NL. The SE module enables the network to focus on more important channel features, and the obtained weights are multiplied by the features after activation by the nonlinear activation function NL, and finally residual jump connections are performed after being processed by a 1×1 convolution kernel.

9. The malware classification method of pseudo-label optimization and primary and secondary model collaboration according to claim 5 or 8, characterized in that: The SE module includes an average pooling layer and two fully connected layers. A RELU function is connected between the average pooling layer and the first fully connected layer, and an output of the second fully connected layer is connected to an H-Swish activation function.

10. The malware classification method of pseudo-label optimization and primary and secondary model collaboration according to claim 1 is characterized by: In step S3), the Bayesian controller is used to perform uncertainty assessment and screening on the quality of the pseudo labels generated by the pseudo label generation model, which specifically includes the following steps: S31), pseudo labels generated by pseudo label generation model It is expressed as: In the formula, f dan-mobilenet represents the pseudo-label generation model; x represents a malware image sample; S32) Use Bayesian reasoning to calculate the credibility of each pseudo label, based on the sample x and the pseudo label Calculate the posterior probability So as to evaluate the credibility of the pseudo-label, namely: In the formula, Indicates that the sample is a pseudo label of x Indicates that the pseudo label is The probability of sample x appearing under the condition of ; represents the prior probability of the category pseudo label; P(x) represents the marginal probability; S33), preset quality threshold τ, when the posterior probability Greater than the preset quality threshold τ, the pseudo label If it is credible, then the pseudo label is discarded.

Images