A malicious code dataset expansion method

Abstract

The application provides a malicious code data set expansion method, including: obtaining a malicious code data set to be expanded, and generating corresponding gray images by visualizing malicious codes in the malicious code data set; randomly initializing each gray image in the malicious code data set; in the forward process of the diffusion model, gradually introducing Gaussian noise until reaching a preset diffusion degree; in the reverse process of the diffusion model, adjusting sample distribution through the back propagation of the diffusion model to generate a large number of new samples, and expanding the original malicious code data set. The application learns the conversion process from the original data to the noise, and restores the noise samples to new similar samples by using the reverse process, realizes the expansion of the data set, generates new samples similar to but different from the original data set, alleviates the influence of the imbalance of part of the family data on the classification detection task, and improves the generalization ability of the model.

Description

Technical Field

[0001] This invention relates to the field of malicious code detection technology, and in particular to a method for expanding a malicious code dataset. Background Technology

[0002] With the popularization of the internet and the advent of the digital age, malware has brought enormous challenges to cybersecurity. To improve the accuracy of malware detection, various malware analysis methods and technologies have been proposed. Malware analysis technologies are categorized into dynamic analysis and static analysis technologies based on whether the file is executable. Dynamic analysis refers to the practice of running executable files in sandboxes, emulators, and virtual machines and monitoring and analyzing application behavior through system calls. Static analysis methods extract static features of malware to identify the malicious behavior of samples. Static analysis methods are superior to dynamic analysis in terms of speed and effectiveness because they can capture information related to structural characteristics. Traditional malware detection methods rely on signature-based model matching, requiring manual extraction of malware signatures based on expert knowledge and comparison with known signatures in a database. With the development of malware obfuscation and packing techniques, malware has spawned numerous variants. Traditional detection methods are inefficient and struggle to effectively detect and identify malware variants; therefore, how to accurately and efficiently detect and classify malware and its variants has become a research hotspot in this field.

[0003] To address the challenges of traditional malware detection methods and more effectively detect malware variants that have been obfuscated or packed, visualization-based malware detection methods have emerged. This method first maps malware to images. Based on the similarity of texture features within the same malware family and the differences in texture features between different malware families, it extracts and classifies the texture features of the malware images. This method has been proven to effectively detect malware variants, and its detection speed is 4000 times faster than dynamic detection techniques. Since its inception, numerous experts and scholars have conducted research on this method. Nataraj et al. fused image and signal features to describe malware and used KNN (K-Nearest Neighbor) as a classifier to identify malware. Kancherla et al. fused Gabor features, wavelet features, and intensity features as the total features to enhance feature diversity and trained an SVM (Support Vector Machines) classifier to classify malware. Liu Yashu et al. constructed anti-obfuscation features by fusing GIST features and LBP (Local Binary Pattern) features of malware images to address the problem of decreased classification performance of the model among similar malware images. To reduce computation time, Naeem et al. proposed an LGMP feature descriptor that integrates local and global features of malware images. The aforementioned existing techniques apply machine learning to visualization-based malware detection methods, focusing on feature extraction and classification models, but without addressing the quality of the dataset itself.

[0004] The Malimg dataset, a widely used dataset containing a large number of grayscale images of malware, provides an important resource for malware research. However, the limited number of samples from individual malware families in the Malimg dataset cannot cover all possible malware variants. This can lead to trained models failing to accurately identify the latest malware types. In summary, while deep learning models have demonstrated exceptional capabilities in fields such as computer vision and natural language processing, the limited number of training samples for some malware families, which cannot fully describe the distribution characteristics of the entire dataset, can cause deep learning models to overfit to these scarce data points, resulting in poor model performance. Summary of the Invention

[0005] To address the shortcomings of the prior art, this invention proposes a method for expanding a malicious code dataset. This method generates new samples that are similar to but different from the original dataset based on a diffusion model, thereby expanding the dataset.

[0006] The technical solution adopted in this invention, a method for expanding a malicious code dataset, includes the following steps: Step 1: Obtain the malware dataset to be expanded, and visualize the malware in the dataset to generate grayscale images of the corresponding samples; Step 2: Randomly initialize each grayscale image in the malicious code dataset. During the forward propagation of the diffusion model, Gaussian noise is gradually introduced until the preset diffusion level is reached, generating a pure noise image. Step 3: Perform reverse denoising based on the pure noise image. Adjust the sample distribution through backpropagation of the diffusion model to generate a large number of new samples and expand the original malicious code dataset.

[0007] Furthermore, the specific steps of step 2 include: Step 2.1: Obtain the visualized grayscale image; Step 2.2: Gradually add Gaussian noise to each grayscale image. Gaussian noise added in each step Prediction is performed using a noise prediction model. After T diffusion processes, all original grayscale images are transformed into pure noise images. These T diffusion processes can be represented as follows: ; ; in, Represents constant hyperparameters. This represents the sampling data at step t. This represents the sampled data at step t-1. Represents the original sample data; Step 2.3: Assumptions , Establish sample data at any time. Compared with the original sample data Relationship : ; Will Further reparameterization is as follows: ; in, Represents a grayscale image. This represents the grayscale image at time step t. This represents the probability that the grayscale image remains unchanged at time step t. This represents the probability that a grayscale image remains unchanged from time step 1 to time step t. Represents the identity matrix.

[0008] Furthermore, step 3 includes the following specific steps: Step 3.1: Process the pure noise image obtained from the forward diffusion process. To perform denoising, the denoising process at time step t is represented as: ; ; in, This represents the mean of the Gaussian distribution during the reverse process; Step 3.2: Calculate the posterior probability distribution : ; in, ; Step 3.3: Based on the loss function of the prediction noise model, using the data generated during the forward diffusion process as training data, adjust the parameters of the prediction noise model. Train the model to obtain a well-trained predictive noise model. Step 3.4: Extract the pure noise image The input is fed into a trained noise prediction model to predict the mean. and variance Sampling is performed according to the following formula. The process is repeated T times to finally obtain the generated grayscale image. The sampling formula is: ; in, This indicates that given the current noise data x t Given the original data x0, the mean of the denoised data predicted by the model.

[0009] Furthermore, the predicted noise model is an improved U-Net model, which is divided into a contraction path and a diffusion path. The contraction path includes a downsampling module, a convolutional unit, and a residual block, while the diffusion path includes an upsampling module, a convolutional unit, and a residual block, wherein: The downsampling module adopts a 4-layer downsampling structure. Each downsampling layer uses a convolutional layer with a kernel size of 3×3, a stride of 2, and padding of 1, which reduces the size of the feature map by half each time. The upsampling module adopts a 4-layer upsampling structure. Each upsampling layer uses the nearest neighbor interpolation method for calculation. Each time, the feature map size is doubled and then concatenated with the corresponding downsampled feature map. Convolutional units are used for feature extraction in the contraction path and gradually reduce the spatial dimension through pooling layers or convolutional layers with a stride greater than 1, while increasing the number of feature channels, so that the network can extract deeper feature representations from the input image. In the expansion path, they are used to upsample the feature map, gradually increase the spatial dimension, and reduce the number of feature channels, and finally output a feature map with fine spatial resolution. The residual block consists of two modules and an embedding layer. Each module includes group normalization, a ReLU function, and a convolutional layer that ensures the feature map size remains constant. The time difference between the residual block and the processing time after the embedding layer is calculated. When the number of input and output channels are equal, a direct short circuit is used to connect the two channels. When the number of input and output channels are different, a 1×1 convolutional layer is used to scale the number of channels before connecting them.

[0010] Furthermore, the loss function for the noise prediction model is: ; in, Indicates noise The predicted value.

[0011] Therefore, the above-mentioned method for expanding a malicious code dataset has the following beneficial effects: First, this invention expands existing malware datasets by introducing a diffusion model, effectively addressing the problem of insufficient sample size. By introducing a diffusion process, the model can learn the inherent patterns and characteristics of data from limited data samples, thereby enabling better prediction and classification. Secondly, this invention effectively addresses the data imbalance problem by introducing a diffusion model. By adjusting the diffusion process, the model can better handle data with different distributions, thus improving its performance when dealing with original imbalanced data.

[0012] In summary, this invention applies a diffusion model to expand the existing malicious code dataset, solving the problems of insufficient sample size and data imbalance. Experimental results show that the diffusion model, by generating more synthetic samples, not only enriches the training set but also improves the model's generalization ability and recognition performance.

[0013] The technical solution of the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. Attached Figure Description

[0014] Figure 1 A flowchart for visualizing malicious code.

[0015] Figure 2 Malicious images visualized for different malware families.

[0016] Figure 3 These are the two processes in the diffusion model.

[0017] Figure 4 This is a diagram of the improved U-Net architecture.

[0018] Figure 5 This is a diagram of the residual block structure.

[0019] Figure 6 This is the confusion matrix for a multi-class classification problem.

[0020] Figure 7 Experimental results were selected to expand the sample size.

[0021] Figure 8 a- Figure 8 c represents the data generated after training for the three families: VB.AT, Fakerean, and C2LOP.P.

[0022] Figure 9 a- Figure 9 b represents the curves showing the changes in accuracy and loss rate of the training and test sets as a function of training batches during the training process of this model.

[0023] Figure 10 a- Figure 10 b represents the confusion matrix before and after dataset expansion, respectively. Detailed Implementation

[0024] In the description of this invention, it should also be noted that, unless otherwise expressly specified and limited, these embodiments are for illustrative purposes only and are not intended to limit the scope of the invention. Furthermore, it should be understood that after reading the teachings of this invention, those skilled in the art can make various alterations or modifications to the invention, and these equivalent forms also fall within the scope defined by the appended claims.

[0025] This invention proposes a method for expanding a malware dataset by introducing a diffusion model. The relevant background is introduced below: 1. Malicious code visualization Malicious code visualization is the process of converting malicious code binary files into grayscale images, including steps such as file loading, file processing, data conversion, resizing, and normalization. Figure 1 As shown in the diagram. First, the given malicious code binary file is read in groups of 8-bit unsigned integers. Then, each group of binary numbers is converted to a decimal integer. Next, the line width is determined based on the size of the PE file (malicious code binary file) and converted into a two-dimensional array. The correspondence between the line width and the file size is shown in Table 1. Finally, each element in the two-dimensional array is used as the grayscale value of the image, mapping the two-dimensional array into a grayscale image. Some of the converted malicious family samples (Adialer.C, Instantaccess, Fakerean, Malex.gen!, Skintrim.N, and Yuner.A) are shown below. Figure 2 As shown.

[0026] Table 1. Correspondence between line width and malicious file size ; 2. Diffusion Model The diffusion model expands the dataset by simulating the diffusion process and using the backpropagation algorithm. Its principle is based on the gradual expansion and simulation of the data distribution.

[0027] During the diffusion process, the model also calculates the differences between each sample and the original samples, and backpropagates these differences back to the original samples. This means that as the diffusion process continues, the original samples gradually expand into the new sample distribution. This expansion process is based on the backpropagation algorithm, which adjusts the position and distribution of the original samples according to the differences, so that the newly generated samples are closer to the true distribution.

[0028] By repeatedly performing the above steps, the preset diffusion level is reached. A large number of new samples are generated by simulating the diffusion process of the data distribution and adjusting the sample distribution using backpropagation. Since the diffusion model introduces noise and perturbations during the sample expansion process, the generated samples may differ from the original samples. These differences can be seen as an extension and supplement to the original data distribution. This expansion method not only increases the number of samples in the dataset but also improves the model's generalization ability and robustness.

[0029] Next, based on the diffusion model, this invention constructs a malicious code diffusion model to expand the malicious code dataset.

[0030] 3. Malicious code propagation model construction

[0031] This invention expands upon malware datasets using a diffusion model. Specifically for the Malimg dataset, the diffusion model first randomly initializes each sample (the resulting grayscale image after visualization) in the Malimg dataset. This provides the model with an initial input containing random noise during training, simulating the data generation process in the real world. By randomly initializing the grayscale images, the model can learn to gradually recover the original data distribution from noise during training, thereby improving its ability to generate high-quality samples. This random initialization process starts with a completely random noisy image (typically sampled from a Gaussian distribution) and gradually adds noise until the image becomes entirely noisy.

[0032] Then, after random initialization, the diffusion process is simulated by gradually introducing Gaussian noise. In each iteration, the model calculates the Gaussian noise to be added to the current sample based on its location and diffusion parameters. This noise gradually increases until a preset diffusion level is reached. Once the preset diffusion level is reached, the model simulates the diffusion process of the data distribution and uses backpropagation to adjust the sample distribution, generating a large number of new samples to expand the original Malimg dataset, thereby increasing the number of samples in the dataset.

[0033] (1) Model building The diffusion model is a generative model that uses variational inference to train a parameterized Markov chain, progressively generating latent variables from a simple distribution. Generating data in complex distributions ,in Let represent the total diffusion step size. It is each diffusion time step As a result, and Having the same dimension The diffusion model of Gaussian noise is a stochastic process that can be divided into two parts: the forward process (also known as the forward diffusion process, the noise addition process, etc.) and the reverse process (also known as the reverse diffusion process, the reverse denoising process, etc.).

[0034] like Figure 3 As shown, the forward process is a gradual addition of noise. At each step of the forward process, diagonal Gaussian noise is added to the sample data distribution. By continuously adding Gaussian noise, the original sample data distribution is transformed into a simple standard Gaussian distribution. The reverse process is a denoising process. It samples from the standard Gaussian distribution, removing a small amount of Gaussian noise at each step, gradually approximating the true data distribution, and thus obtaining samples from the true data distribution, thereby achieving the purpose of generating data.

[0035] Traditional diffusion models typically employ fixed noise addition and removal steps when processing general image data. However, grayscale images after malware visualization possess unique texture and structural features, which may be more easily disrupted or blurred by noise. To address this issue, this invention optimizes the forward and backward processes of traditional diffusion models, particularly in the noise addition and removal strategies. This invention proposes an adaptive noise addition strategy that dynamically adjusts the amount of noise added based on the local features of the malware grayscale image. This allows for more effective noise introduction while maintaining the overall image structure, thereby enhancing the model's generalization ability. Simultaneously, the noise removal steps are adjusted accordingly in the backward process to ensure that more malware feature information is preserved when reconstructing the image.

[0036] The two processes will be described in detail below: Forward process The forward process is the process of spreading the original data towards a noise distribution. In this process, Gaussian noise is gradually added to the original data until it becomes entirely noisy data. middle, This represents the learnable components of the conditional noise predictor and the low-resolution encoder in the diffusion model. Posterior The diffusion process will distribute Transform into latent variables And fixed as a Markov chain, which is based on the variance table Gradually add Gaussian noise to the data .

[0037] (1) (2) In the formula, It is a very small normal number and can be regarded as a constant hyperparameter.

[0038] Let coefficients , The diffusion process allows for any time step sampling Its form is a closed form: (3) Further reparameterization to: (4) in, Represents a grayscale image. This represents the grayscale image at time step t. This represents the probability that the grayscale image remains unchanged at time step t. This represents the probability that a grayscale image remains unchanged from time step 1 to time step t. Represents the identity matrix; The above equation describes the process at each time step where noise is gradually added to the original data and diffused into the noise distribution.

[0039] This invention targets grayscale images generated after visualization of malicious code, specifically the image generation problem. Correspondingly, the noise addition in the forward process is to convert the original image into a noisy image.

[0040] Reverse process The backward process is the process of diffusing noisy data towards the original data distribution. In this process, the model needs to learn a parameterized function that predicts the noise added at each step of the forward process, and then progressively subtracts this noise in the backward process to approximate the original data distribution. This process is typically achieved by minimizing a log-likelihood loss, which distributes the latent variables... Converted to Parameterized data distribution It is made by Defined as a Markov chain with Gaussian transitions starting with .

[0041] (5) (6) in, This represents the mean of the Gaussian distribution during the reverse process. Indicates the variance of the prediction; During the training phase of the noise prediction model, the maximum variational lower bound (ELBO) approximation and the maximum log-likelihood function are adopted to indirectly minimize the negative log-likelihood function, and KL divergence and variance reduction are introduced.

[0042] (7) in, This represents the expected value of the negative log-likelihood. Indicates the variational lower bound. Denotes KL divergence, , , These represent the prior matching term, the denoised matching term, and the reconstructed term, respectively. This transformation requires The corresponding diffusion process posterior Direct comparison between them. Let: (8) in, This indicates that given the current noise data x t Given the original data x0, the mean of the denoised data predicted by the model; We can obtain: (9) in, Does it represent the variance of a Gaussian distribution? This represents a hyperparameter that gradually increases with t. Represents the input sample Identity matrices with the same dimensions; This represents a random noise vector sampled from a standard Gaussian distribution; when season , Since C is a constant, we can obtain: (10) in, This indicates that in the reverse process of the diffusion model, the first... The loss function for each step; Correspondingly, the denoising process in the reverse process of the present invention is to restore the noisy image to the original image.

[0043] (2) Loss function Training the mean using the U-net neural network At that time, for the quantity that needs to be predicted There are three ways to obtain the training results. These three ways are: (1) directly predict the mean of Gaussian noise at each step of the inverse process. (2) Predicting initial values Then, using formula (8), we can obtain the mean. (3) First predict the noise z, that is, first sample one Then sample according to equation (8). Then calculate according to formula (11). : (11) Where θ represents the learnable components of the conditional noise predictor and the low-resolution encoder in the discrete model. The predicted value represents Gaussian noise; This invention uses a third method for prediction, and its loss function is: (12) Since the ultimate goal of U-net neural network optimization is to maximize the final result of the inverse process... To obtain the generation result that best matches the sample, the variational lower bound (VLB) is used to optimize the negative log-likelihood function associated with ELBO: (13) The prior art discloses that equation (12) can be regarded as a reweighted form of the variational lower bound loss LVLB, and that optimizing equation (12) will produce better grayscale image quality than directly optimizing LVLB. Therefore, in the specific implementation of malicious code grayscale image generation, this invention uses the L1 loss function instead of the MSE loss function, and combines it with equation (4) to obtain the final target loss function for training the U-net neural network: (14) in, Indicates the expected value. Represents the cumulative noise product; (3) Predictive noise model The noise prediction model is a U-Net with the same input and output dimensions. The improved U-Net structure used in this invention is as follows: Figure 4 As shown, it performs multi-scale processing on the input image and outputs feature maps of multiple scales in the decoder section to adapt to the segmentation requirements of objects of different sizes. The direct input of this improved U-Net network is... , is a single-channel 128×128 tensor, time After encoding using embedding technology, the residual structures are merged and input into the network, and the network output is... The number of channels and their size are the same as the input. The network structure adopts a 4-layer downsampling structure. The downsampling of the network uses a convolution operation with a kernel size of 3×3, a stride of 2, and padding of 1, reducing the size of the feature map by half each time. The upsampling of the network uses nearest neighbor interpolation, doubling the size of the feature map each time, and concatenating it with the feature map of the corresponding shrinking path. The number of convolutional channels in the convolutional units are set to 32, 64, 128, 256, and 512, respectively. To prevent gradient vanishing and gradient exploding, a residual structure is used to complete the network propagation and channel number expansion and reduction. In the shrinking path, the main role of the convolutional unit is to extract features and gradually reduce the spatial dimension (usually through pooling layers or convolutional layers with a stride greater than 1), while increasing the number of feature channels, so that the network can extract deeper feature representations from the input image. In the expanding path, the main role of the convolutional unit is to upsample the feature map, gradually increase the spatial dimension, and reduce the number of feature channels, ultimately outputting a feature map with fine spatial resolution. Specifically: 1) Contraction path Convolution operation: Convolutional layers perform convolution operations with the input image through convolution kernels to extract local features of the image.

[0044] Spatial dimensionality reduction: By using pooling layers (such as max pooling layers) or convolutional layers with a stride greater than 1, the spatial dimensions (width and height) of the feature map are gradually reduced, thereby reducing computational complexity.

[0045] Increased feature channels: As the network progresses, the number of feature channels gradually increases, enabling the network to learn richer feature representations.

[0046] 2) Extended Path Upsampling operation: By using upsampling layers (such as transposed convolutional layers, bilinear interpolation, etc.), the spatial dimension of the feature map is gradually increased in order to perform more refined feature representation and localization.

[0047] Feature channel reduction: As the network moves closer to the output layer, the number of feature channels gradually decreases in order to generate a final feature map with fewer channels.

[0048] Feature fusion: In the expansion path, the upsampled feature map is usually fused with the feature map of the corresponding scale in the contraction path (such as by splicing or adding) to combine low-level and high-level feature information and improve the performance of the model.

[0049] The structure of the residual block is as follows: Figure 5 As shown, it consists of two modules (containing convolutional layers, batch normalization layers, and the main path and shortcut connections of activation functions) and an embedding layer. Each module includes group normalization, a ReLU function, and convolutional layers that ensure the feature map size remains constant. The processing time is compared with that of the embedding layer. When the number of input and output channels are equal, a direct short circuit is used to connect the two channels. Figure 5 a) When the number of input and output channels is different, use a 1×1 convolutional layer to scale the number of channels before connecting them. Figure 5 b). The final convolutional layer of the network converts the number of channels to 1, resulting in noise. Predicted value . Example

[0050] To verify the effectiveness of the extended method proposed in this invention, the following simulation experiment was conducted.

[0051] 1. Dataset and Experimental Environment This embodiment uses the Malimg dataset published by the University of California Visual Research Lab for model training and classification. This dataset contains 9435 malware samples from 25 malware families. The model was trained on Ubuntu 22.04.2, running in CUDA 11.8, using the PyTorch 1.11.0 deep learning framework, and on an NVIDIA GeForce RTX 3080Ti hardware configuration.

[0052] The Malimg dataset was divided into a training set and a validation set in a 9:1 ratio. The training set was used to train the model, and the validation set was used to observe and evaluate the model's performance. The names of each family and the distribution of the number of samples in each family are shown in Table 3.

[0053] Noise timeline of diffusion model The function for determining the value of is the cosine schedule proposed in reference [1], and its function form is: , According to Obtain the entire noise timeline. The total number of steps T for the diffusion model is set to 1000.

[0054] Table 2. Sample distribution of the Malimg dataset ; 2. Performance Evaluation Indicators This invention uses four metrics—accuracy, precision, recall, and F1 score—to evaluate the model's performance. These metrics have been widely used in related research, and their formulas are as follows: (15) (16) (17) (18) in, This indicates that the actual positive malicious code sample was correctly predicted as a positive sample. This indicates that an actual negative malicious code sample was incorrectly predicted as a positive sample. Similarly, The actual negative malware samples were correctly predicted as negative samples. This indicates that an actual positive malicious code sample was incorrectly predicted as a negative sample. In multi-class classification problems, , , and Values ​​such as Figure 6 As shown. Among them, This indicates a family category of malicious code.

[0055] 3. Comprehensive Experiment to Verify the Validity of the Diffusion Model (1) Parameter selection experiment In the data processing section, this model uses a diffusion model to expand the malware samples, where the selection of the sample size n directly affects the model's classification performance. To obtain the optimal value of n, with all other conditions remaining the same, five different values ​​of n (200, 300, 500, 800, and 1000) were compared. The experimental results are shown in Table 7. Compared to other values ​​of n, the model's accuracy reaches 99.51% when n=500. When n is higher than 500, the model's inference speed increases rapidly, but the accuracy gradually decreases. The experimental results indicate that n=500 is the optimal value for N.

[0056] Based on the experimental results, for malware families with fewer than 500 samples, the training expanded the number of visualization images for these families to 500. For other malware families, 500 samples were randomly selected for visualization. Through these steps, after processing the sample supplementation module, all 25 malware families participating in the experiment had 500 samples.

[0057] (2) Validation experiment of diffusion model To verify the performance of the proposed diffusion model, all experiments were conducted and tested under the same environment. During model training, the batch size was set to 32 samples per training iteration, the number of epochs was set to 200, the Adam optimizer was used, the initial learning rate was set to 0.0001, the weight decay was 0.05, and the peak learning rate was 10. -3 The learning rate was decayed using cosine scheduling, and the loss function was the cross-entropy loss function. The results are shown in Table 3.

[0058] Table 3. Experiments to verify the effectiveness of the diffusion model ; Taking the training process of the VB.AT, Fakerean, and C2LOP.P families as an example, the generated data results are as follows: Figure 8 a- Figure 8 As shown in c. After a specified number of iterations of training, realistic grayscale images of malicious code can be generated. It can be seen that the generated results effectively reflect the feature information in the original dataset, such as the three shadows of the VB.AT family and the unique black borders at the top and bottom of the Fakerean family, which are also represented in the generated results. The texture of C2LOP.P is also well represented after the diffusion model.

[0059] Figure 9 This shows the performance of the training and test sets as a function of training batches during the training process of this model. Figure 9 'a' is the curve showing how accuracy changes with the number of training batches. Figure 9b is the curve showing the loss rate as a function of training batches. The red line represents the training performance on the test set, and the black line represents the training performance on the training set. As can be seen, the model has a fast convergence speed, converging quickly within the first 10 epochs. After training and testing, the model achieved an accuracy of 99.47%, a loss rate of 0.0027, a precision of 99.39%, a recall of 99.35%, and an F1 score of 99.48%.

[0060] To further validate the effectiveness of the diffusion model and clearly observe the classification details before and after dataset augmentation, thus providing a more intuitive verification of its effectiveness, the confusion matrix of the U-Net network before and after the transformation was plotted. Figure 10 a- Figure 10 As shown in b, the values ​​on the main diagonal of the confusion matrix represent the true positive rate for malware family classification, while other values ​​represent the false negative rate. It can be seen that although the left graph shows excellent performance in multiple family classifications, the recall rate is low in the Autorun.K family. The reason for this is that the Autorun.K family has a small sample size, with only 81 samples, and only 9 samples are allocated to the test set. The small data volume leads to learning difficulties, and each misclassified sample accounts for a higher percentage. The right graph, after the diffusion model, shows that the network has high recognition accuracy in all families. Experimental results indicate that the diffusion model effectively alleviates the impact of dataset imbalance on model classification performance and improves the model's recognition accuracy.

[0061] (3) Comparison experiment between the model and existing advanced malware classification technologies To demonstrate that the proposed method achieves satisfactory performance and further verify the effectiveness of the diffusion model in generating malware, the model is compared with existing malware detection methods based on visualization analysis using the Malimg dataset. Table 4 summarizes the performance metrics of various malware detection methods. As can be seen from the table, the proposed model outperforms existing visualization-based malware detection methods both domestically and internationally.

[0062] Table 4 Comparison of performance metrics for various malware classification methods ; References: Literature[1] "NicholAQ,DhariwalP.ImprovedDenoisingDiffusionProbabilisticModels[C]. Proceedings of Machine Learning Research;PMLR.2021:1-7.” Document [2] "Hamad Naeem, Guo Bing, Ullah Farhan, et al. A Cross-Platform Malware Variant Classification based on Image Representation[J]. KSII Transactions on Internet and Information Systems, 2019, 13: 3756-3777." Document [3] "Zhihua Cui, Du Lei, Wang Penghong, et al. Malicious code detection based on CNNs and multi-objective algorithm[J]. Journal of parallel and distributed computing, 2019, 129: 50-58." Document [4] "R Vinayakumar, Alazab Mamoun, Soman K-P, et al. Robust Intelligent Malware Detection Using Deep Learning[J]. IEEE Access, 2019, 7: 46717-46738." Document [5] "Sitalakshmi Venkatraman, Alazab Mamoun, Vinayakumar R. A hybrid deep learning image-based analysis for effective malware detection[J]. Journal of Information Security and Applications, 2019, 47: 377-389." Document [6] "Hamad Naeem, Guo Bing, Ullah Farhan, et al. A Cross-Platform Malware Variant Classification based on Image Representation[J]. KSII Transactions on Internet and Information Systems, 2019, 13: 3756-3777." Reference [7] "Verma Vinita, Sunil K-Muttoo, Singh V-B. Multiclass Malware Classification via First- and Second-Order Texture Statistics[J]. Computers & Security, 2020, 97: 101895." Reference [8] "S-Abijah Roseline, Geetha S, Kadry Seifedine, et al. Intelligent Vision-Based Malware Detection and Classification Using Deep Random Forest Paradigm[J]. IEEE Access, 2020, 8: 206303-206324." Reference [9] "Hamad Naeem, Ullah Farhan, Naeem Muhammad-Rashid, et al. Malware detection in the industrial Internet of Things based on hybrid image visualization and deep learning model[J]. Ad hoc networks, 2020, 105: 102154." Reference

[10] "Vasan Danish, Mamoun Alazab, Sobia Wassan, et al. IMCFN: Image-based malware classification using fine-tuned convolutional neural network architecture[J]. Computer Networks, 2020, 171: 107138." Reference

[11] "V Anandhi, Vinod P, Menon Varun-G. Malware visualization and detection using DenseNets[J]. Personal and Ubiquitous Computing, 2021." Literature

[12] "Changguang Wang, Zhao Ziqiu, Wang Fangwei, et al. A NovelMalware Detection and Family Classification Scheme for IoT Based on DEAM andDenseNet[J]. Security and Communication Networks, 2021, 20211-16." Literature

[13] "Vassilios Moussas, Andreatos Antonios. Malware DetectionBased on Code Visualization and Two-Level Classification[J]. Information, 2021, 12(3): 118" Literature

[14] "Sudhakar, Kumar Sushil. MCFT-CNN: Malware classification with fine-tuned convolution neural networks using traditional and transferlearning in the Internet of Things[J]. Future Generation Computer Systems, 2021, 125334-351." Literature

[15] "Sanjeev Kumar, Janet B. DTMIC: Deep transfer learning formalware image classification[J]. Journal of Information Security andApplications, 2022, 64103063." Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and not to limit them. Although the present invention has been described in detail with reference to preferred embodiments, those skilled in the art should understand that modifications or equivalent substitutions can still be made to the technical solutions of the present invention, and these modifications or equivalent substitutions cannot cause the modified technical solutions to deviate from the spirit and scope of the technical solutions of the present invention.

Claims

1. A method for expanding a malicious code dataset, characterized in that, Includes the following steps: Step 1: Obtain the malware dataset to be expanded, and visualize the malware in the dataset to generate grayscale images of the corresponding samples; Step 2: Randomly initialize each grayscale image in the malicious code dataset. During the forward propagation of the diffusion model, Gaussian noise is gradually introduced until the preset diffusion level is reached, generating a pure noise image. Step 3: Perform reverse denoising based on the pure noise image. Adjust the sample distribution through backpropagation of the diffusion model to generate a large number of new samples and expand the original malicious code dataset.

2. The method for expanding a malicious code dataset as described in claim 1, characterized in that, Step 2 includes the following specific steps: Step 2.1: Obtain the visualized grayscale image; Step 2.2: Gradually add Gaussian noise to each grayscale image. Gaussian noise added in each step Prediction is performed using a noise prediction model. After T diffusion processes, all original grayscale images are transformed into pure noise images. These T diffusion processes can be represented as follows: ； ； in, Represents constant hyperparameters. This represents the sampling data at step t. This represents the sampled data at step t-1. Represents the original sample data; Step 2.3: Assumptions , Establish sample data at any time. Compared with the original sample data Relationship : ； Will Further reparameterization is as follows: ： in, Represents a grayscale image. This represents the grayscale image at time step t. This represents the probability that the grayscale image remains unchanged at time step t. This represents the probability that a grayscale image remains unchanged from time step 1 to time step t. Represents the identity matrix.

3. The method for expanding a malicious code dataset as described in claim 2, characterized in that, Step 3 includes the following specific steps: Step 3.1: Process the pure noise image obtained from the forward diffusion process. To perform denoising, the denoising process at time step t is expressed as: ； ； in, This represents the mean of the Gaussian distribution during the reverse process; Step 3.2: Calculate the posterior probability distribution : ； in, ; Step 3.3: Based on the loss function of the prediction noise model, using the data generated during the forward diffusion process as training data, adjust the parameters of the prediction noise model. Train the model to obtain a well-trained predictive noise model. Step 3.4: Extract the pure noise image The input is fed into a trained noise prediction model to predict the mean. and variance Sampling is performed according to the following formula. The process is repeated T times to finally obtain the generated grayscale image. The sampling formula is: ； in, This indicates that given the current noise data x t Given the original data x0, the mean of the denoised data predicted by the model.

4. The method for expanding a malicious code dataset as described in claim 3, characterized in that, The predicted noise model is an improved U-Net model, which consists of a contraction path and a diffusion path. The contraction path includes a downsampling module, a convolutional unit, and a residual block, while the diffusion path includes an upsampling module, a convolutional unit, and a residual block. The downsampling module adopts a 4-layer downsampling structure. Each downsampling layer uses a convolutional layer with a kernel size of 3×3, a stride of 2, and padding of 1, which reduces the size of the feature map by half each time. The upsampling module adopts a 4-layer upsampling structure. Each upsampling layer uses the nearest neighbor interpolation method for calculation. Each time, the feature map size is doubled and then concatenated with the corresponding downsampled feature map. Convolutional units are used for feature extraction in the contraction path and gradually reduce the spatial dimension through pooling layers or convolutional layers with a stride greater than 1, while increasing the number of feature channels, so that the network can extract deeper feature representations from the input image. In the expansion path, they are used to upsample the feature map, gradually increase the spatial dimension, and reduce the number of feature channels, and finally output a feature map with fine spatial resolution. The residual block consists of two modules and an embedding layer. Each module includes group normalization, a ReLU function, and a convolutional layer that ensures the feature map size remains constant. The time difference between the residual block and the processing time after the embedding layer is calculated. When the number of input and output channels are equal, a direct short circuit is used to connect the two channels. When the number of input and output channels are different, a 1×1 convolutional layer is used to scale the number of channels before connecting them.

5. The method for expanding a malicious code dataset as described in claim 3, characterized in that, The loss function for the noise prediction model is: ； in, Indicates the expected value. This represents the cumulative noise product.

Images