Lightweight model design method for small sample malicious traffic classification

Abstract

The invention belongs to the technical field of network security, and particularly relates to a lightweight model design method for small sample malicious traffic classification, which is different from a traditional migration method based on pre-training fine tuning. Therefore, the fitting of different migration network layers to the target classification task is improved. For a lightweight task, a one-dimensional full convolutional neural network suitable for traffic classification is provided as a source network, and redundant neurons are captured and eliminated through step-by-step pruning, so that the model complexity is reduced. Besides, aiming at the problem that the model generalization ability is insufficient due to overfitting of the model for small sample data, when the migration model is trained, the model generalization ability is improved by searching a basic representation for the data.

Description

Technical Field

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

[0002] The advent of the big data era is accompanied by a spurt in the growth of network traffic. Malicious traffic generated by malicious programs is mixed in, endangering network security. Malicious traffic classification can associate malicious traffic with the malicious family of its generating program. As the first step in the task of detecting network malicious resources, it has always been widely concerned. With the widespread use of encryption technology, traditional traffic classification methods include port-based methods, payload inspection technology, and machine learning-based methods, which are not effective 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 combine feature engineering and model training into a whole, learn features directly from the original input, and are more suitable for processing encrypted data. Therefore, traffic classification technology based on deep learning has also continued to develop into an important solution to the problem of malicious traffic classification.

[0003] However, there are two problems that hinder the application of deep learning in the field of malicious traffic classification. First, deep learning-based methods rely on a large amount of data-driven, while malicious traffic is highly concealed and updates quickly, resulting in a small amount of malicious traffic that can be captured and data labels that are difficult to accurately label. This is called the small sample problem in the academic field. The small sample problem often leads to insufficient training of deep learning models and low classification accuracy. Secondly, current deep learning models are gradually tending towards large models, requiring sufficient storage and computing resources to support training. For example, the classic traffic classification model deeppacket has more than 30 million parameters and requires up to 2.1 billion calculations, which makes the model based on deep learning training have relatively high requirements for the deployment environment and equipment. However, some edge network devices (such as routers) have limited available computing resources and it is difficult to meet the resources required for deep learning. At the same time, due to the large model The amount of computing that is disproportionate to the target task conflicts with the requirements of near-real-time classification (such as for delay-sensitive traffic). Therefore, the research on the design method of lightweight small sample malicious traffic classification model is of great significance.

[0004] In fact, the research on any of the above problems is accompanied by an issue that cannot be ignored, namely the model generalization problem. 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 has poor accuracy when applied to other tasks. The lightweight model aims to remove model structures that are irrelevant to the current task to the maximum extent 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 it challenging to solve the problem of generalization of lightweight models in small sample tasks.

[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 dataset) to migrate some of its network parameters that extract common features of the source dataset and the target dataset, and uses the target dataset to fine-tune the remaining parameters to obtain a model suitable for the target dataset. The method reduces the demand for the target dataset 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 IoTNetworks.Sensors.22,5621) proposed an efficient intrusion detection framework based on transfer learning, knowledge transfer, and model refinement. The framework is suitable for 5G IoT scenarios with unbalanced and scarce labeled datasets. Experiments show that the method achieves high accuracy and low false prediction rate. Compared with any previous deep learning-based intrusion detection system, the proposed solution has better detection rate for different series of known and zero-day attacks.

[0006] Model lightweighting: With the continuous development of deep learning, the huge structure of neural networks is gradually unable to match the resources of edge devices, and the demand for lightweight small sample classification models has been proposed. In recent years, researchers have proposed many methods to lightweight models, including low-rank decomposition, knowledge extraction, quantization, and pruning. For lightweight migration models, we should find redundant model structures that appear in the migration process for the target task, and selectively "slim down" the model instead of compressing the model. The precise model structure is also of certain significance for the classification accuracy that the model can provide. Therefore, for lightweight migration models, pruning methods have a certain reference value. Model pruning uses a certain method to select neurons that contribute highly to the classification task, and obtains a new model structure by pruning redundant neurons. The lightweight purpose can be achieved by fine-tuning the model using the target data set. Currently, there are few studies on model lightweighting methods based on pruning in the field of malicious traffic classification. Liu et al. ("Learning Efficient Convolutional Networks through Network Slimming," 2017 IEEE International Conference on Computer Vision (ICCV), Venice, Italy, 2017, pp. 2755-2763) first proposed a lightweight method called network slimming, which is implemented by enforcing channel-level sparsity in a simple and effective way. The method was experimentally tested on several state-of-the-art CNN models in the image field and proved its effectiveness. For VGGNet, the method reduced the model size by about 20 times and the 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 the pre-trained full-size model for small sample classification models. The scheme combines model compression and transfer learning to construct a simplified custom model. Experiments show that when 60% of the network is pruned, the accuracy of the picture dataset Flowers-102 drops by less than 2% using this scheme.

[0007] In the current classification methods based on transfer learning, most default training migration models assume that the data differences between the source dataset and the target dataset are not large. Under this assumption, datasets with high similarity generally have the same basic features, and it is beneficial to directly copy the shallow network parameters used by the migration model to extract basic features. However, the similarity discrimination method between traffic datasets needs to be explored. If the correlation of datasets is low, directly migrating some parameters of the source model is likely to produce negative transfer effects, and the correlation between the source dataset and the target dataset becomes a factor that limits the accuracy of the model. Therefore, for the small sample malicious traffic classification task based on transfer learning, a more effective method is needed to select the parameters that should be retained during the model migration process.

[0008] Secondly, the traditional lottery hypothesis generally believes that a network with a large number of parameters is more likely to contain good sparse subnetworks. For small sample classification tasks, transfer learning achieves good results based on the fact that the fine-tuned transfer network contains good sparse subnetworks suitable for the classification of target small sample data. However, the complexity of small sample data sets and transfer models is often highly mismatched, so pruning the redundant structure of the transfer model for the target small sample task can provide better model performance.

[0009] For pruning tasks, the effect of pruning depends on the accurate judgment of redundant neurons, which requires that the weights of the neurons in the trained model are positively correlated with the target classification task. In the traditional migration method, the parameters of the shallow network are directly copied from the source network, and only the deep network is fine-tuned through the target data set. Pruning the neurons of the shallow network according to the training results of the source network may result in a decrease in the classification accuracy of the model for the target task after the so-called "redundant neurons" are pruned. Therefore, a more reasonable method is needed to rank the importance of the neurons of the migration network.

[0010] Thirdly, most current traffic classification models are composed of convolutional neural networks and fully connected layers, and pruning is mostly aimed at "slimming down" the neurons in the convolutional layer. The convolutional layer network with optimized structure is spliced ​​with the original fully connected layer. This method still retains the fully connected layer with a large number of parameters. This is because each neuron in the fully connected layer is connected to all neurons in the upper layer, and global information is used for samples. Its fully connected characteristics make it difficult to reshape its structure. Therefore, the optimization of the fully connected layer structure has become one of the challenging issues for lightweight models. Summary of the invention

[0011] In order to solve the current demand for accurate, lightweight and generalizable models in the task of classifying malicious traffic with small samples, the present invention proposes a lightweight model design method for classifying malicious traffic with small samples, comprising 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, and use a large dataset to train the source model to obtain good model parameters for classification tasks;

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

[0014] Step 3: Migration model pruning: After the migration of the public feature extractor and the private feature extractor is completed, the importance of neurons is sorted and the model is pruned on the fine-tuned model. The new model structure is reshaped using the target dataset to obtain the final small-sample malicious traffic classification lightweight model.

[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, and 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 step size of the maximum pooling layer 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 a 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, and 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, where x is the flow sample. The gradient reversal layer is formally expressed as:

[0019]

[0020] At the same time, we construct 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. 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, 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 back propagation, and guides the network to adjust parameters to extract beneficial 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 value of each target dataset sample y and the predicted value The gap between , 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, perform pruning 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, perform pruning according to the threshold, and 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 all BN layer γ values, 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] Different from the traditional migration method based on pre-training and fine-tuning, the present invention performs model migration and fine-tuning according to the different feature extraction targets of the shallow network and the deep network respectively, 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 removed through step-by-step pruning, thereby reducing the complexity of the model. In addition, in order to solve the problem of insufficient model generalization ability due to overfitting of the model to small sample data, when training the migration model, the present invention improves the model generalization ability by finding a basic representation of the data.

[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 model parameters and model structure after migration are more compatible with 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 shallow network of the model to find the invariant features between the data set of the training source model and the target data set, so that the model can have better generalization between different tasks.

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

[0037] Figure 1 is a flow chart of the method of the present invention;

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

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

[0040] In order to better understand the purpose, technical solution and function of the present invention, the present invention is further described in detail below in conjunction with the accompanying drawings, but the present invention can be implemented in a variety of different ways as defined and covered by the claims. The accompanying drawings constituting a part of this application are used to provide a further understanding of the present invention, and the exemplary embodiments of the present invention and their descriptions are used to explain the present invention and do not constitute an improper limitation 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: The present invention migrates and fine-tunes the shallow network and deep network of the source model as public feature extractors and private feature extractors respectively. The fine-tuning goal of the public feature extractor is to generate common features that can extract the source dataset and the target dataset. At the same time, the feature should be used as a useful feature for distinguishing the target small sample task. Therefore, the source dataset and the target dataset are used to jointly guide the fine-tuning. The private feature extractor aims to extract high-level features of the target task and is fine-tuned using the target dataset.

[0043] Step 3: Migrate model and prune: After the migration of the two parts of the feature extractor is completed, sort the importance of neurons and perform model pruning on the fine-tuned model, and use the target dataset to reshape the new model structure to obtain the final small sample malicious traffic classification lightweight model.

[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 stage, 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, fine-tuning the private feature extractor E using the target dataset p , the optimization goal 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 the present 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] The present invention designs a model based on a fully convolutional neural network as the source model for transfer learning, which avoids the problem of insufficient pruning caused by the inability to optimize the structure of the fully connected layer during the subsequent pruning process. Except for the last network layer as the classification layer, the remaining layers of the fully convolutional neural network can rank and select the importance of the neurons they contain, which can further optimize the network structure.

[0054] The designed source model is an end-to-end fully convolutional neural network model. In view of the characteristics of traffic data, a one-dimensional convolutional neural network CNN is used as the backbone architecture of the neural network. 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 ten one-dimensional convolutional layers and seven maximum pooling layers. The step size of the maximum pooling layer is set to 2. The convolutional layer uses 200 and 100 neurons in the shallow and deep layers to process the input data, respectively. The size of each neuron is [6, 1], and each neuron moves 1 step after a convolution operation. A BN layer is added after each convolutional layer. In addition to the traditional function of preventing overfitting and accelerating model convergence, the scaling factor of the BN layer is also used to evaluate the importance of neurons for subsequent pruning. The features processed by the BN layer will be input into the RELU activation function layer.

[0055] In order to alleviate the problem of the 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 of 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 set and 50% test set. 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, to obtain a classification model that adapts to the target task 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] In order to encourage the two parts of the network to have a better representation of the target task, the two extractors are fine-tuned in different ways and the loss functions are constructed separately. After the step-by-step migration, the two feature extractors will be combined into a model with a common representation similar to the two datasets and a private representation of the target small sample dataset. By migrating separately in this way, the classifier fine-tuned on the common representation can better generalize across domains because its input is not affected by the representation aspects unique to each domain. At the same time, because the small sample dataset only bears the fine-tuning training of the private representation, the pressure on the small sample dataset is reduced, and a classifier with better performance can be trained.

[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, it is expected that the migrated parameters are 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 extracted features are also beneficial when the model performs other tasks.

[0060] Based on the above analysis, the present invention introduces an adversarial learning idea to perform 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 uses 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 reverse propagation play the opposite role, 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 classification labels corresponding to the samples in the source domain and the target domain, so as to counter 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 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 identify 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 stage. At the same time, the traditional pruning method directly performs unified sorting and pruning on 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, in the pruning process, the present invention is based on the 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 small sample classification tasks. 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 in order.

[0069] The source model proposed in the present invention is composed of only convolutional layers, so only the redundancy of neurons in the convolutional layers needs to be considered, which reduces the complexity of the pruning phase and can perform a global pruning method based on the BN layer scaling factor. The BN layer is a data normalization method before the activation layer in the neural network, performing the following conversion:

[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 for evaluating neurons, statistically sorts all the BN layer γ values, 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 step by step. First, for the public feature extractor, the γ parameters of the migrated network BN layer are obtained, and pruning is performed according to the threshold, aiming to search for the optimal network structure that can extract effective public features. For the private feature extractor, pruning is performed in the same way. After pruning, it is integrated with the optimal structure searched by the public feature extractor. After obtaining the new network, the target data set is 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 a source model for step-by-step migration and pruning to obtain a lightweight malicious traffic small sample 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, so all 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 in that: 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, and use a large dataset to train the source model to obtain good model parameters for classification tasks; Step 2: Migrate the source model in steps: Migrate and fine-tune the shallow network and deep network of the source model as the public feature extractor and private feature extractor respectively. The fine-tuning goal of the public feature extractor is to generate a common feature that can extract the source dataset and the target dataset. At the same time, this feature is used as a useful feature for distinguishing the target small sample task. The source dataset and the target dataset are used to jointly guide the 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: Migration model pruning: After the migration of the public feature extractor and the private feature extractor is completed, the importance of neurons is sorted and the model is pruned on the fine-tuned model. The new model structure is reshaped using the target dataset to obtain the final small-sample malicious traffic classification lightweight model.

2. According to claim 1, a lightweight model design method for small sample malicious traffic classification is characterized in that: 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.

3. According to claim 2, a lightweight model design method for small sample malicious traffic classification is characterized in that: The one-dimensional convolution layer has 10 layers, the maximum pooling layer has 7 layers, the step size of the maximum pooling layer 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 a convolution operation.

4. According to claim 2, a lightweight model design method for small sample malicious traffic classification is characterized in that: In step 2, For the public feature extractor, the shallow network of the pre-trained network is used to extract features, and 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, where x is the flow sample. The gradient reversal layer is formally expressed as: At the same time, we construct 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. 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, 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 back propagation, and guides the network to adjust parameters to extract beneficial 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 value of each target dataset sample y and the predicted value The gap between , the loss function is defined as: Where D t is the target data.

5. According to claim 4, a lightweight model design method for small sample malicious traffic classification is characterized in that: In step three, it specifically includes: For the public feature extractor, obtain the γ parameter of the migrated network BN layer, perform pruning 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, perform pruning according to the threshold, and 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 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.

Images