A lightweight model design method for small sample malicious traffic classification

Abstract

The present application belongs to the technical field of network security, and particularly relates to a lightweight model design method for small sample malicious traffic classification. Unlike the traditional transfer method based on pre-training fine-tuning, the present application performs model transfer and fine-tuning according to different feature extraction targets of shallow networks and deep networks, so as to improve the fitting of different transfer network layers to the target classification task. For the lightweight task, a one-dimensional fully convolutional neural network suitable for traffic classification is proposed as a source network, and redundant neurons are captured and removed through step-by-step pruning, so as to reduce the model complexity. In addition, in order to solve the problem of insufficient model generalization ability caused by overfitting of the model to small sample data, the present application finds a basic representation of the data to improve the model generalization ability when training the transfer model.

Description

Technical Field

[0001] The present invention belongs to the field of network security technology, and specifically relates to a lightweight model design method for classifying small sample malicious traffic. Background Art

[0002] The advent of the big data era has been accompanied by a surge in network traffic, with malicious traffic generated by malicious programs intermingling with it and posing a threat to network security. Malicious traffic classification, which can associate malicious traffic with the malicious family of its generating program, has attracted widespread attention as the first step in detecting malicious network resources. With the widespread use of encryption technology, traditional traffic classification methods, including port-based methods, payload inspection techniques, and machine learning-based methods, have been ineffective due to the limitations of prior knowledge. In recent years, deep learning methods have been widely studied in various fields. Some end-to-end deep learning models can integrate feature engineering and model training into a whole, learning features directly from raw input, making them more suitable for processing encrypted data. Therefore, deep learning-based traffic classification technology has continued to develop and become an important solution to the problem of malicious traffic classification.

[0003] However, two issues hinder the application of deep learning in malicious traffic classification. First, deep learning-based methods rely on large amounts of data. However, malicious traffic is highly cryptic and updates rapidly, resulting in a small volume of malicious traffic that can be captured and difficulty accurately labeling data. This is known in academic circles as the small sample problem, and the small sample problem often leads to insufficient training of deep learning models and low classification accuracy. Second, current deep learning models are becoming increasingly large, requiring ample storage and computing resources to support training. For example, the classic traffic classification model DeepPacket has over 30 million parameters and requires up to 2.1 billion computations. This places high demands on the deployment environment and equipment for deep learning-based training. However, some edge network devices, such as routers, have limited available computing resources, making it difficult to meet the requirements of deep learning. Furthermore, the computational overhead of large models, which is disproportionate to the target task, conflicts with the requirements of near-real-time classification (e.g., for latency-sensitive traffic). Therefore, research on lightweight, small-sample malicious traffic classification models is of great significance.

[0004] In fact, research on any of the above issues is accompanied by an issue that cannot be ignored, namely the problem of model generalization. The abundant amount of training data allows the model to focus on the diversity of features. However, when using small sample data to train the model, due to the high penalty coefficient for error terms, it is easy to derive an overfitting model, which is inaccurate when applied to other tasks. The lightweight model aims to remove as much model structure as possible without affecting the accuracy. This results in the removal of some neurons that extract basic features, making the lightweight model structure only applicable to the current task. The above two target characteristics make solving the problem of lightweight model generalization in small sample tasks a challenge.

[0005] Small sample traffic classification based on deep transfer learning: Since the transferability of deep networks has been confirmed, inductive transfer learning has become a promising framework for small sample classification tasks. It uses a pre-trained source model (usually a large model trained with a large data set) to migrate some of its network parameters that extract common features of the source and target data sets, and uses the target data set to fine-tune the remaining parameters to obtain a model suitable for the target data set. The method reduces the demand for the target data set by migrating the source model parameters and fine-tuning them. Researchers continue to use deep transfer learning to solve the problem of small sample traffic classification based on the characteristics of traffic data. Idriss et al. (Accelerating the update of a DL-based IDS for IoT using deep transfer learning [J]. Indonesia. J. Electr. Eng. Comput. Sci., 2021, 23 (2): 1059-1067) applied the transfer learning method to the intrusion detection system and updated the solution of the deep learning-based intrusion detection system (DL-IDS). This method achieved good results in multiple indicators such as detection rate. Eva et al. (Transfer-Learning-Based Intrusion Detection Framework in IoT Networks. Sensors. 22, 5621) proposed an efficient intrusion detection framework based on transfer learning, knowledge transfer, and model refinement. This framework is applicable to 5G IoT scenarios with imbalanced and scarce labeled datasets. Experiments demonstrate that the approach achieves high accuracy and low false positive rates. Compared to any previous deep learning-based intrusion detection system, the proposed solution achieves better detection rates for a diverse range of known and zero-day attacks.

[0006] Model lightweighting: With the continuous advancement of deep learning, the bulky structures of neural networks are increasingly incompatible with edge device resources, raising the need for lightweight, small-sample classification models. In recent years, researchers have proposed numerous methods to lightweight models, including low-rank decomposition, knowledge extraction, quantization, and pruning. For lightweight migration models, it is important to identify redundant model structures that arise during the migration process for the target task and selectively "slim down" the model rather than compressing it. A precise model structure is crucial for the classification accuracy it can achieve. Therefore, pruning methods offer valuable insights for lightweight migration models. Model pruning uses a method to select neurons that contribute most to the classification task, pruning away redundant neurons to create a new model structure. This lightweighting can then be achieved by fine-tuning the model using the target dataset. Currently, research on pruning-based model lightweighting methods in the field of malicious traffic classification is limited. Liu et al. ("Learning Efficient Convolutional Networks through Network Slimming," 2017 IEEE International Conference on Computer Vision (ICCV), Venice, Italy, 2017, pp. 2755-2763) pioneered a lightweight method called network slimming, which achieves this by enforcing channel-level sparsity in a simple and effective way. This method was experimentally demonstrated on several state-of-the-art CNN models for image processing, demonstrating its effectiveness. For VGGNet, it reduced the model size by approximately 20 times and the number of computational operations by 5 times. Subsequently, Liao et al. ("Convolution Filter Pruning for Transfer Learning on Small Dataset," 2020 International Computer Symposium (ICS), 2020, pp. 79-84) proposed a scheme to reduce the size of pre-trained full-size models for small-sample classification models. This scheme combines model compression and transfer learning to construct a simplified custom model. Experiments show that using this scheme reduces the accuracy of the Flowers-102 image dataset by less than 2% when pruning 60% of the network.

[0007] Most current classification methods based on transfer learning assume that the data differences between the source and target datasets used to train the transfer model are small. Under this assumption, datasets with high similarity generally share the same basic features, and directly copying the shallow network parameters used by the transfer model to extract these basic features is beneficial. However, methods for determining similarity between traffic datasets remain to be explored. If the datasets are poorly correlated, directly migrating some parameters of the source model can easily produce negative transfer effects, and the correlation between the source and target datasets becomes a factor limiting the model's accuracy. Therefore, for the task of classifying small-sample malicious traffic based on transfer learning, more effective methods are needed to select parameters that should be retained during the model migration process.

[0008] Secondly, the traditional lottery ticket hypothesis generally holds that a network with a large number of parameters is more likely to contain well-defined sparse subnetworks. For small-sample classification tasks, transfer learning achieves good results based on the assumption that the fine-tuned transfer network contains well-defined sparse subnetworks suitable for classifying the target small-sample data. However, the complexity of small-sample datasets and the transfer model is often highly mismatched. Therefore, pruning redundant structures in the transfer model that are specific to the target small-sample task can improve model performance.

[0009] The effectiveness of pruning relies on accurately identifying redundant neurons, which requires that the weights of the neurons in the trained model be positively correlated with the target classification task. In traditional transfer methods, the parameters of the shallow network are directly copied from the source network, while only the deeper network is fine-tuned using the target dataset. Pruning shallow network neurons based on the training results of the source network can result in a decrease in the model's classification accuracy for the target task after removing so-called "redundant neurons." Therefore, a more rational method is needed to rank the importance of neurons in the transferred network.

[0010] Secondly, most current traffic classification models consist of convolutional neural networks and fully connected layers. Pruning typically involves slimming down the convolutional layer neurons, splicing the optimized convolutional layer network with the original fully connected layer. This approach still retains the large number of parameters in the fully connected layer. This is because each neuron in the fully connected layer is connected to all neurons in the upper layer, using global information for samples. Its fully connected nature makes it difficult to restructure the layer. Therefore, optimizing the fully connected layer structure has become a challenging issue for lightweight models. Summary of the Invention

[0011] To address the current demand for accurate, lightweight, and generalizable models in the task of classifying small-sample malicious traffic, this paper proposes a lightweight model design method for small-sample malicious traffic classification, which includes the following steps:

[0012] Step 1: Source model construction and training: Design a model based on a fully convolutional neural network as the source model for transfer learning. Use a large dataset to train the source model to obtain good model parameters for the classification task.

[0013] Step 2: Migrate the source model in steps: Migrate and fine-tune the shallow and deep networks of the source model as public and private feature extractors, respectively. The fine-tuning goal of the public feature extractor is to generate common features that can extract the source and target datasets. At the same time, this feature serves as a useful feature for distinguishing the target small sample task, and the source and target datasets are used to guide fine-tuning. The private feature extractor aims to extract high-level features for the target task and is fine-tuned using the target dataset.

[0014] Step 3: Migrate the model and prune it: After the public and private feature extractors are migrated, sort the neuron importance and perform model pruning on the fine-tuned model. Use the target dataset to reshape the new model structure to obtain the final lightweight model for small-sample malicious traffic classification.

[0015] Furthermore, in step one, the designed source model is an end-to-end fully convolutional neural network model, using one-dimensional CNN as the neural network backbone architecture. The one-dimensional CNN can capture the dependencies between bytes in network data packets and find identification methods for different categories. The source model consists of multiple layers of one-dimensional convolutional layers and multiple layers of maximum pooling layers. A BN layer is added after each convolutional layer, and the features processed by the BN layer will be input into the RELU activation function layer; at the same time, two convolutional layers are used instead of the fully connected layer. The neurons in the convolutional layer are only connected to a local area in the input data and share parameters with the neurons in the convolution column.

[0016] Furthermore, the one-dimensional convolution layer is 10 layers, the maximum pooling layer is 7 layers, the maximum pooling layer step size is set to 2, and the convolution layer uses 200 and 100 neurons in the shallow and deep layers to process the input data, respectively, where the size of each neuron is [6, 1], and each neuron moves 1 step after one convolution operation.

[0017] Furthermore, in step 2,

[0018] For the public feature extractor, the shallow network of the pre-trained network is used to extract features. The domain classification of the sample is used as the classification task 1 of the network. The gradient reversal layer is introduced. For the shallow network, the feature function f(x) is extracted. x is the traffic sample. The gradient reversal layer is formally expressed as:

[0019]

[0020] At the same time, we construct classification task 2, which aims to predict the classification labels of samples in the source and target domains. The final training goal of the network is to reduce the following losses:

[0021]

[0022] in, is a parameter that controls the interaction of loss terms, is the loss function, is the loss term for classification task 1, This is the loss term for classification task 2. Under the action of the two loss terms, the loss function outputs the common loss of the current network parameters for domain classification and sample category classification, performs backpropagation, and guides the network to adjust parameters to extract useful common features.

[0023] For the private feature extractor: Freeze the fine-tuned public feature extractor parameters so that it can directly output public features, use the target dataset to fine-tune the remaining layers of the migration network as the private feature extractor, and minimize the difference between each target dataset sample y and the predicted value. The loss function is defined as:

[0024]

[0025] Where D t is the target data.

[0026] Furthermore, in step three, it specifically includes:

[0027] For the public feature extractor, obtain the γ parameter of the migrated network BN layer, prune it according to the threshold, and obtain the optimal network structure that can extract effective public features;

[0028] For the private feature extractor, obtain the γ parameter of the migrated network BN layer and prune it according to the threshold to obtain the optimal network structure that can extract effective private features;

[0029] After integrating the two optimal network structures to obtain a new network, the new network structure is reshaped using the target data set.

[0030] Furthermore, in step three, it specifically includes: using the γ parameter in the BN layer as an indicator for evaluating neurons, statistically sorting the γ values ​​of all BN layers, obtaining the specified number of retained BN layers, that is, obtaining the weight threshold after sorting, traversing the BN layer γ value weights in the model, making masks for each layer, obtaining the index of non-zero values, retaining the neuron weights with a value of 1, building a new model structure according to the mask, loading the pruned model, assigning the pruned weights to the new model, and reshaping the model.

[0031] Unlike traditional migration methods based on pre-training and fine-tuning, the present invention performs model migration and fine-tuning according to the different feature extraction targets of shallow networks and deep networks, so as to improve the fit of different migration network layers to the target classification task. For lightweight tasks, a one-dimensional fully convolutional neural network suitable for traffic classification is proposed as the source network. Redundant neurons are captured and eliminated through step-by-step pruning, thereby reducing the complexity of the model. In addition, in order to address the problem of insufficient model generalization ability due to overfitting of the model to small sample data, the present invention improves the model generalization ability by finding a basic representation of the data when training the migration model.

[0032] The present invention has the following beneficial effects:

[0033] 1. The present invention proposes a lightweight model method for classifying small sample malicious traffic based on step-by-step migration and pruning, so that the migrated model parameters and model structure better match the target small sample classification task.

[0034] 2. The present invention designs a fully convolutional network as the source network for migration, and performs step-by-step pruning on this basis to make the model more lightweight.

[0035] 3. An adversarial learning idea is introduced into the training of the model's shallow network to find the invariant features between the dataset of the training source model and the target dataset, so that the model can have better generalization between different tasks.

[0036] 4. This method can maintain a classification accuracy of over 96% on multiple small-sample datasets, exceeding the current mainstream small-sample traffic classification methods. With an accuracy drop of no more than 1%, this method reduces model parameters by over 80% and is more generalizable than other lightweight methods. BRIEF DESCRIPTION OF THE DRAWINGS

[0037] Figure 1 Flow chart of the method of the present invention;

[0038] Figure 2 is the source model structure diagram;

[0039] Figure 3 Flowchart for public feature extractor training. DETAILED DESCRIPTION

[0040] To better understand the objectives, technical solutions, and functions of the present invention, the present invention is described in further detail below with reference to the accompanying drawings. However, the present invention may be implemented in a variety of different ways as defined and covered by the claims. The accompanying drawings, which constitute a part of this application, are intended to provide a further understanding of the present invention. The exemplary embodiments of the present invention and their descriptions are intended to explain the present invention and do not constitute undue limitations of the present invention.

[0041] like Figure 1 As shown, the present invention proposes a lightweight model design method for classifying small sample malicious traffic, including step one, source model construction and training: first, a model based on a fully convolutional neural network is designed as a source model for transfer learning, and a large data set is used to train the source model to obtain good model parameters for the classification task.

[0042] Step 2: Migrate the Source Model Step by Step: This paper migrates and fine-tunes the shallow and deep networks of the source model as public and private feature extractors, respectively. The public feature extractor is fine-tuned to generate common features that can be used to extract features from both the source and target datasets. This feature should also be useful for distinguishing the target small-sample task. Therefore, both the source and target datasets are used to guide fine-tuning. The private feature extractor, with the goal of extracting high-level features for the target task, is fine-tuned using the target dataset.

[0043] Step 3: Migrate the model and prune it: After the two feature extractors are migrated, sort the importance of neurons on the fine-tuned model and perform model pruning. Use the target dataset to reshape the new model structure to obtain the final lightweight model for small-sample malicious traffic classification.

[0044] Problem Definition

[0045] The present invention introduces symbols and annotations to define the optimization goal of the method. Assume that W s is a pre-trained model, s indicates that the model is optimized for the source domain, and the present invention uses D s and D t Represents source data and target data respectively. is the loss function used to optimize the network. In the pre-training phase, the optimization objective can be expressed as follows:

[0046]

[0047] In the transfer learning scenario, the present invention uses the source data and the target data set to jointly guide the public feature extractor E c Optimization, using the target dataset to fine-tune the private feature extractor E p , the optimization objective of the migration model is defined as:

[0048]

[0049] In the pruning phase, assuming W f is the migration model, and f is the number of all neurons in the pre-trained model. The goal of this invention is to find the optimal Minimize the loss of the model in the target domain, as shown below:

[0050]

[0051] In the following sections, the implementation method of each step will be described in detail through specific embodiments.

[0052] Step 1: Source model construction and training.

[0053] This paper designs a model based on a fully convolutional neural network as the source model for transfer learning, avoiding the problem of insufficient pruning caused by the inability to optimize the structure of the fully connected layer during the subsequent pruning process. With the exception of the final classification layer, all layers of the fully convolutional neural network can rank and select the importance of the neurons they contain, further optimizing the network structure.

[0054] The designed source model is an end-to-end fully convolutional neural network model. Based on the characteristics of traffic data, a one-dimensional convolutional neural network (CNN) is used as the neural network backbone. One-dimensional CNNs can capture inter-byte dependencies in network packets and identify different categories. The source model consists of ten one-dimensional convolutional layers and seven max pooling layers. The max pooling layer has a stride of 2. The convolutional layers process the input data using 200 and 100 neurons in the shallow and deep layers, respectively. Each neuron has a size of [6, 1], and each neuron shifts by 1 step after each convolution operation. A batch normalization (BN) layer is added after each convolutional layer. In addition to preventing overfitting and accelerating model convergence, the BN layer also uses a scaling factor to assess the importance of neurons for subsequent pruning. The features processed by the BN layer are then fed into the RELU activation function layer.

[0055] In order to alleviate the problem of a large number of parameters in the fully connected layer, unlike the traditional classification model, the proposed model uses two convolutional layers instead of the fully connected layer. The neurons in the convolutional layer are only connected to a local area in the input data, and share parameters with the neurons in the convolution column. This feature is more in line with the subsequent requirements for model pruning. The neurons in the fully connected layer and the convolutional layer are both dot product calculations, which are the same in functional form, so it is feasible to replace the fully connected layer with the convolutional layer. The number of input channels of the last convolutional layer is set to the number of classification categories. After multiple convolutions and pooling, the output of the last layer of the model is the most important high-dimensional feature map of the input sample. The feature map is used as the classification prediction score for each class of the sample, and forward propagation is performed to complete the training of the model. The present invention splits the data set into 50% training sets and 50% test sets. After 5-fold cross-validation on the training set, the hyperparameters are set based on the grid search of the best average accuracy. The detailed structure and parameters of the model are as follows. Figure 2 shown.

[0056] Step 2: Migrate the source model step by step.

[0057] In order to solve the problem that small sample data sets cannot drive neural network training, the present invention uses transfer learning. Based on a pre-trained source model, a classification model adapted to the target task is obtained after the source model is migrated and fine-tuned. The migration source model is often limited by the difference between the source data set and the target data set. When the difference between the two data sets is large, the traditional migration method, that is, directly copying the shallow network parameters of the source model and fine-tuning the deep network parameters, is no longer applicable. Therefore, inspired by the model's shallow network and deep network extracting public features and private features respectively, the present invention reconstructs the migration source model into a public feature extractor and a private feature extractor for migration and fine-tuning respectively.

[0058] To encourage both networks to develop better representations for the target task, the two extractors are fine-tuned differently, with separate loss functions constructed. After the step-by-step migration, the two feature extractors are combined into a single model with a common representation similar to both datasets and a private representation specific to the target small-sample dataset. By migrating these models separately, the classifier fine-tuned on the common representation generalizes better across domains, as its input is unaffected by the domain-specific representations. Furthermore, since the small-sample dataset only bears the burden of fine-tuning the private representation, the pressure on the small-sample dataset is reduced, resulting in a classifier with improved performance.

[0059] Public feature extractor: Since shallow networks are often copied directly from one model to another during the migration process, the features extracted by this part of the network are expected to be common features of the two models, that is, public features. However, due to the differences between the data sets, the migration network parameters often have deviations, which is allowed in the large model as long as the large model contains a sub-model for the target small sample task. However, since the goal of the present invention is to generate a lightweight small sample malicious traffic classification model, the parameters to be migrated are expected to be more accurate and effective, that is, a neuron parameter weight that is positively correlated with the target task is required. Therefore, the present invention proposes the concept of a public feature extractor, which is used to find the invariant representation between the source data set and the target data set without losing the classification ability of the target data set. By mining the common features between data sets, the features extracted by the model are also beneficial when performing other tasks.

[0060] Based on the above analysis, the present invention introduces an adversarial learning concept to conduct the confrontation and reconciliation of the two target tasks. Based on experience, the present invention uses the shallow network of the pre-trained network to extract features. The present invention takes the domain classification of the sample as the classification task 1 of the network. If favorable common features are extracted, the loss of the classifier will be very large. Therefore, the gradient reversal layer (GRL) is introduced. For the shallow network to extract the feature function f(x), the gradient reversal layer is formally expressed as

[0061]

[0062] x is a traffic sample, which makes the back propagation have the opposite effect, further increasing the loss of the classifier and promoting the extraction of common features. However, only extracting features that cannot be classified into domains as a method of determining common features can easily lead the network to extract meaningless features. Therefore, the present invention constructs classification task 2, the goal of which is to obtain the prediction results of the corresponding classification labels of samples in the source domain and the target domain, so as to combat the meaningless behavior of the common feature extractor during training. The ultimate training goal of the network is to reduce the following losses:

[0063]

[0064] in It is a parameter that controls the interaction of loss terms. Under the action of the two loss terms, the loss function outputs the common loss of the network's current parameters for domain classification and sample category classification, performs back propagation, and guides the network to adjust parameters to extract useful common features. The common feature extractor training process is as follows: Figure 3 shown.

[0065] Private feature extractor: The goal of the private feature extractor is to find high-level features that distinguish the target small sample dataset from the public features. Therefore, the present invention freezes the fine-tuned public feature extractor parameters so that it can directly output the public features. Only the remaining layers of the target dataset fine-tuned migration network are used as the private feature extractor. The loss function encourages the extractor to extract features that can distinguish different categories, hoping to minimize the difference between each target dataset sample y and the predicted value. The gap between , so the loss function is defined as:

[0066]

[0067] Step 3: Migrate model pruning.

[0068] The main problem with pruning a small sample classification model is that the sample size is not enough to support the reshaping work after the network is pruned. In addition, model pruning evaluates the contribution of neurons in the trained network through a certain evaluation method, which requires that the weight training of neurons must be effective during model training. The present invention solves this problem by using the reconstruction model to migrate separately during the migration phase. At the same time, the traditional pruning method directly performs a unified sorting and pruning of all neuron parameters of the migration model. Since the migration network only uses the target data set to adjust the parameters of the deep network, the unified sorting may cause the pruned redundant neurons to be only beneficial to the classification task of the source data set, but not to the target task after migration. To solve this problem, during the pruning process, the present invention is based on a step-by-step migration work, and divides the model into a public feature extractor and a private feature extractor for pruning and reshaping respectively, ensuring that the pruned redundant parameters are beneficial for the small sample classification task. Since the parameters have been retrained according to the different tasks of the two extractors during the step-by-step migration, it is only necessary to determine the parameter weight sorting in the pruning work and prune them in sequence.

[0069] The source model proposed in this paper consists only of convolutional layers, so only the redundancy of neurons in the convolutional layers needs to be considered. This reduces the complexity of the pruning phase and allows for a global pruning method based on the scaling factor of the batch normalization layer. The batch normalization layer, as a data normalization method before the activation layer in the neural network, performs the following transformations:

[0070]

[0071] z in and z out is the input and output of the BN layer, B is the current training batch, μ and σ are the mean and standard deviation values ​​of the input activation, γ and β are trainable affine transformation parameters (scale and shift), which provide the possibility of linearly transforming the normalized activation back to any scale. The present invention uses the γ parameter in the BN layer as an indicator to evaluate neurons, statistically sorts the γ values ​​of all BN layers, obtains the specified number of retained BN layers, and obtains the weight threshold after sorting. Traverse the BN layer γ value weights in the model, make masks for each layer (weight>threshold is 1, weight<threshold is 0), obtain the index of non-zero values, and retain the neuron weights with a value of 1. Construct a new model structure according to the mask, load the pruned model, assign the pruned weights to the new model, and reshape the model.

[0072] Based on the above pruning method, the model is pruned in steps. First, for the public feature extractor, the γ parameter of the migrated network's BN layer is obtained and pruned according to a threshold, aiming to search for the optimal network structure that can extract effective public features. The private feature extractor is pruned using the same method. After pruning, it is integrated with the optimal structure found by the public feature extractor to obtain a new network. The target dataset is then used to reshape the new network structure. The pruning algorithm is shown in Algorithm 1.

[0073] Algorithm 1

[0074]

[0075] The present invention proposes a design method for a lightweight model for classifying small-sample malicious traffic. The method designs an end-to-end full-convolution classification model, which is used as the source model for step-by-step migration and pruning to obtain a lightweight small-sample malicious traffic classification model. This method solves four challenges: the similarity between the source data set and the target data set limits the classification accuracy of the migration model; the lightweight model ignores the attention to the fully connected layer; the pruning of the migration network lacks effective sorting of the importance of neurons; and the model has poor generalization. Experimental results show that the method proposed in the present invention has higher accuracy than the current small-sample traffic classification method and lightweight method, saves more computing resources and is more generalizable. The work of the present invention can provide a new perspective for the study of lightweight models for small-sample malicious traffic classification.

[0076] The above embodiments are only preferred implementations of the present invention. Therefore, any equivalent changes or modifications made according to the structures, features and principles described in the scope of the patent application of the present invention are included in the scope of the patent application of the present invention.

Claims

1. A lightweight model design method for small sample malicious traffic classification, characterized by: The steps include: Step 1: Source model construction and training: Design a model based on a fully convolutional neural network as the source model for transfer learning. Use a large dataset to train the source model to obtain good model parameters for the classification task. Step 2: Migrate the source model in steps: Migrate and fine-tune the shallow and deep networks of the source model as public and private feature extractors, respectively. The fine-tuning goal of the public feature extractor is to generate common features that can extract the source and target datasets. This feature is also used as a useful feature for the target small sample task. The source and target datasets are used together to guide fine-tuning. The private feature extractor aims to extract high-level features of the target task and is fine-tuned using the target dataset; Step 3: Migrate the model and prune it: After the public and private feature extractors are migrated, sort the neuron importance and perform model pruning on the fine-tuned model. Use the target dataset to reshape the new model structure to obtain the final lightweight model for small-sample malicious traffic classification.

2. A lightweight model design method for small sample malicious traffic classification according to claim 1, characterized in that: In step 1, the designed source model is an end-to-end fully convolutional neural network model, using one-dimensional CNN as the neural network backbone architecture. One-dimensional CNN can capture the dependencies between bytes in network data packets and find identification methods for different categories. The source model consists of multiple layers of one-dimensional convolutional layers and multiple layers of maximum pooling layers. A BN layer is added after each convolutional layer, and the features processed by the BN layer will be input into the RELU activation function layer; at the same time, two convolutional layers are used instead of the fully connected layer. The neurons in the convolutional layer are only connected to a local area in the input data and share parameters with the neurons in the convolutional layer.

3. A lightweight model design method for small sample malicious traffic classification according to claim 2, characterized in that: The one-dimensional convolution layer has 10 layers, the maximum pooling layer has 7 layers, the maximum pooling layer step size is set to 2, and the convolution layer uses 200 and 100 neurons in the shallow and deep layers to process the input data, respectively, where the size of each neuron is [6, 1], and each neuron moves 1 step after one convolution operation.

4. A lightweight model design method for small sample malicious traffic classification according to claim 2, characterized in that: In step 2, For the public feature extractor, the shallow network of the pre-trained network is used to extract features. The domain classification of the sample is used as the classification task 1 of the network. The gradient flip layer is introduced. For the shallow network, the feature function f(x) is extracted. x is the traffic sample. The gradient flip layer is formally expressed as: At the same time, we construct classification task 2, which aims to predict the classification labels of samples in the source and target domains. The final training goal of the network is to reduce the following losses: in, is a parameter that controls the interaction of loss terms, is the loss function, is the loss term for classification task 1, This is the loss term for classification task 2. Under the influence of the two loss terms, the loss function outputs the common loss of the network's current parameters for domain classification and sample category classification, performs backpropagation, and guides the network to adjust parameters to extract useful common features. For the private feature extractor: Freeze the fine-tuned public feature extractor parameters so that it can directly output public features, use the target dataset to fine-tune the remaining layers of the migration network as the private feature extractor, and minimize the difference between each target dataset sample y and the predicted value. The loss function is defined as: Where D t is the target data.

5. A lightweight model design method for small sample malicious traffic classification according to claim 4, characterized in that: In step three, specifically include: For the public feature extractor, obtain the γ parameter of the migrated network BN layer, prune it according to the threshold, and obtain the optimal network structure that can extract effective public features; For the private feature extractor, obtain the γ parameter of the migrated network BN layer and prune it according to the threshold to obtain the optimal network structure that can extract effective private features; After integrating the two optimal network structures to obtain a new network, the new network structure is reshaped using the target data set.

6. A lightweight model design method for small sample malicious traffic classification according to claim 5, characterized in that: In step three, it specifically includes: using the γ parameter in the BN layer as an indicator to evaluate neurons, statistically sorting the γ values ​​of all BN layers, obtaining the specified number of retained BN layers, that is, obtaining the weight threshold after sorting, traversing the BN layer γ value weights in the model, making masks for each layer, obtaining the index of non-zero values, retaining the neuron weights with a value of 1, building a new model structure according to the mask, loading the pruned model, assigning the pruned weights to the new model, and reshaping the model.

Images