A lightweight malicious traffic identification method based on ViT improvement

Abstract

This invention provides a lightweight malicious traffic identification method based on ViT improvement. First, a grayscale image constructed from the original traffic data is generated using a grayscale image conversion method for data packets in a session. The session length is limited, the length of each data packet in the session is standardized, and the resolution of the grayscale image is unified. The grayscale image is then converted into a two-dimensional matrix according to the order of the data packets in the session. Next, based on a simplified attention mechanism of the hybrid Vision Transformer model, the MobileNetV2 model is used to extract and downsample the grayscale image, ignoring adjacent pixels to generate two sequences. These sequences are then converted back to the shape of the original feature image by a dimensionality-reducing Transformer Encoder module, and finally passed through a fully connected layer to obtain the malicious traffic identification result. This invention provides a high-precision, lightweight malicious traffic identification method that reduces the number of parameters and computational load required by the model while ensuring the accuracy of malicious traffic identification.

Description

Technical Field

[0001] This invention belongs to the field of malicious traffic identification for Internet of Things (IoT) devices, and specifically relates to a lightweight malicious traffic identification method based on ViT improvement. Background Technology

[0002] Over the past few years, the application of IoT devices has seen continuous growth in various fields, such as industry, healthcare, automation, education, smart homes, and smart cities. While the interconnectedness of everything proposed by IoT technology can drive economic development, it also brings new security challenges; it is not as secure as it seems. In fact, IoT devices are typically allocated limited resources, such as small storage media and processors with limited computing power. Attackers often manipulate controlled IoT devices to launch malicious attacks. Therefore, the faster malicious traffic received by IoT devices can be identified, the more effectively malicious attacks can be prevented, reducing the impact of network attacks.

[0003] Most malicious traffic identification technologies rely on known attack signature databases, comparing the traffic to be identified against a predefined malicious traffic signature database to identify potential malicious attacks. This method is clearly very accurate and effective in identifying known attacks, but its main drawbacks are twofold: First, the technology requires significant resources and overhead, and necessitates manual intervention to update attack signatures, such as maintaining and updating the malicious traffic signature database; second, if new types of malicious traffic attacks emerge rapidly, these new types of malicious traffic cannot be updated in the signature database in a timely manner.

[0004] Deep learning-based malicious traffic identification methods can accurately identify malicious traffic by automatically learning from large amounts of raw malicious traffic data, without requiring manual design of features for malicious traffic identification. Due to the robust automatic learning capabilities of deep learning, this topic has received widespread attention from the academic community. Meanwhile, the sheer number of IoT devices provides a vast amount of traffic data for training deep learning models, improving identification accuracy. Therefore, researching deep learning-based malicious traffic identification methods is an inevitable trend, especially applicable to the resource-constrained problems involved in IoT devices. Designing a high-precision, lightweight malicious traffic identification method is of great significance for ensuring the security of IoT devices. Summary of the Invention

[0005] To address the aforementioned issues, this invention proposes a lightweight malicious traffic identification method based on ViT improvement. To ensure the accuracy of malicious traffic identification, grayscale images converted from data packets in a session are used as input to the neural network model, and malicious traffic identification is performed using a simplified hybrid Vision Transformer model.

[0006] A lightweight malicious traffic identification method based on ViT improvement includes the following steps:

[0007] Step 1: Obtain the malicious traffic dataset and perform session segmentation and parsing processing on it;

[0008] Step 1.1: Obtain publicly available datasets of IoT network traffic from the internet;

[0009] Step 1.2: Divide the raw traffic files stored in the dataset into several session files based on TCP or UDP streams;

[0010] Step 1.3: Parse the session file, identify the format of the data packets stored in the file, delete the format information, and retain the data packet information specified by the TCP / IP protocol;

[0011] Step 2: Represent the data packets of the session parsed in Step 1 into a graphical representation and generate a grayscale image;

[0012] Step 2.1: Limit the session length by setting the session length to L. session Its range is from 576 bytes to 50176 bytes. Any excess will be discarded, and any insufficient part will be padded. The padding value is selected as the padding value, and its range is an integer from 0 to 255.

[0013] Step 2.2: Standardize the length of data packets in the session, requiring that the length of each data packet in the session is L. packet Insufficient L packet Fill in the part;

[0014]

[0015] Among them, L session N represents the length of the session. patch This represents the number of blocks.

[0016] Step 2.3: Represent the data packets in the session as grayscale images. Convert the data packets in the session into two-dimensional matrices in sequence to generate grayscale images.

[0017] Step 3: Use a sliding data packet-based approach to augment the data for the categories with smaller sample sizes in the dataset, resulting in a dataset with balanced samples across all categories;

[0018] Step 3.1: Introduce a variable to record the current data packet position, i.e., the starting position of the window;

[0019] Step 3.2: Determine the difference between the number and location of data packets in the current raw traffic file;

[0020] Step 3.3: After the data packet at the current position in the session has been processed by the data packet visualization method, the position variable points to the next data packet;

[0021] Step 4: Assign labels to the grayscale images converted by the packet visualization method according to the label file provided in the dataset, and divide the training dataset and the test dataset in a 6:1 ratio, including two cases: binary classification and multi-class classification of malicious traffic;

[0022] Step 5: Construct a simplified hybrid Vision Transformer model, SHViT, for malicious traffic identification;

[0023] Step 5.1: Use the MobileNetV2 model to downsample and generate feature maps of data packets, which will then be used as input to the Transformer Encoder;

[0024] Step 5.2: Construct the SHViT Block in the simplified hybrid Vision Transformer model;

[0025] Step 5.2.1: Use the feature map as the input to this module, and use a convolutional layer to change the number of channels in the feature map;

[0026] Step 5.2.2: Improve the multi-head self-attention mechanism of the Transformer model, that is, ignore the attention calculation of neighboring pixels;

[0027]

[0028]

[0029] Where n represents the length of the input sequence to the Transformer Encoder, and k represents the factor by which the length of the input sequence is reduced;

[0030] Step 5.2.3: After concatenating the input feature map and the output of the Transformer model through residual connection, the number of channels is kept consistent with that of the input feature map;

[0031] Step 5.3: Pass the output of SHViTBlock through a global pooling layer and a linear layer to obtain the complete SHViT model;

[0032] Step 5.4: Configure training parameters and input training dataset to improve the accuracy of malicious traffic identification;

[0033] Step 5.5: Input the grayscale image converted from the processed raw traffic file using the packet visualization method into the model to obtain the malicious traffic identification result;

[0034] Step 5.6: Calculate the number of parameters and computational cost of the model, which are the metrics for evaluating the model's lightweightness; the parameter and computational cost of convolutional layers and fully connected layers are defined as follows:

[0035] parameters conv =(k w ×k h ×C in +1)×C out (4)

[0036] parameters Fc =(N in +1)×N out (5)

[0037] FLOPs conv =[(C in ×k w ×k h )+(C in ×k w ×k h -1)+1]×C Out ×w×h (6)

[0038] FLOPS Fc =[N in +(N in -1)+1]×N out (7)

[0039] Where, k w kh represents the size of the convolution kernel, C in and C out N represents the number of input and output channels of the convolutional layer. in and N out represent the size of the input and output dimensions in the fully connected layer, respectively; w and h represent the size of the feature map; and 1 in the formula represents the bias in the convolutional layer.

[0040] Beneficial technical effects of the present invention:

[0041] This invention proposes a lightweight malicious traffic identification method based on ViT improvements. To ensure lightweight and high-precision malicious traffic identification, a grayscale image conversion method based on data packets in a session is employed. A data augmentation method is proposed to address the imbalanced sample problem in the dataset, ensuring that the number of samples in each category is roughly equal. Data from data packets in the session is converted into grayscale images and used as input to the neural network model, improving the accuracy of malicious traffic identification. Simultaneously, a simplified hybrid VisionTransformer model is used. This model reduces the number of parameters and computational cost by simplifying the multi-head self-attention mechanism in the Transformer model. Attached Figure Description

[0042] Figure 1 This is a flowchart of a lightweight malicious traffic identification method based on ViT improvement according to the present invention.

[0043] Figure 2 This is a structural diagram of the simplified hybrid Vision Transformer model in this invention;

[0044] Figure 3 This is a structural diagram of the SHViT Block model in this invention;

[0045] Figure 4 This presents the classification accuracy in the binary and multi-class scenarios of the IOT-23 dataset in this invention. Detailed Implementation

[0046] The invention will be further explained below with reference to the accompanying drawings and specific implementation examples.

[0047] This invention uses the IoT-23 dataset for IoT traffic validation. The malicious traffic identification experiment can be understood as a classification task; therefore, classification accuracy is used as the model's evaluation metric to compare the training performance of the model across different rounds. The malicious traffic identification process is as follows: Figure 1 As shown, it specifically includes:

[0048] Step 1: Obtain the Aposemat IoT-23 public dataset of IoT network traffic provided by Avast AIC Labs from the network; as shown in Table 1; and perform session segmentation and parsing processing on it;

[0049] Table 1. Summary data of network traffic in ten IoT scenarios;

[0050]

[0051] The IoT-23 dataset consists of network traffic captured from 23 different IoT scenarios. These scenarios are divided into 20 network traffic capture files from IoT devices infected with malware and 3 network traffic capture files from real IoT devices. This invention selects only 10 of these scenarios as the dataset. Table 1 shows some data from IoT botnet-infected scenarios and benign scenarios. In the scenario names, HC stands for Honeypot-Capture, and MC stands for IoT-Malware-Capture.

[0052] Step 1.1: Since the label file provided by the dataset gives the labels of the entire TCP or UDP stream, rather than giving labels to each packet, and each scenario dataset provides a large raw traffic file, it is impossible to distinguish each TCP or UDP stream. At this time, it is necessary to process the raw traffic file and split it into individual small files of TCP or UDP streams.

[0053] Step 1.2: After obtaining the traffic files of individual TCP or UDP streams, the network traffic capture file needs to be parsed, and the data packets in the stream are stored in memory as a list for easy access at any time;

[0054] Step 2: Visual representation of data packets in the session of the original dataset;

[0055] Step 2.1: Limit the length of the session. This invention chooses a length of 50176 bytes to ensure that the resolution of the data packet after being converted into a grayscale image is exactly 224×224, which is convenient for subsequent model training.

[0056] Step 2.2: Standardize the length of data packets in the session, requiring that the length of each data packet in the session is a multiple of 256 bytes, and padding is used for any part that is less than 256 bytes;

[0057] Step 2.3: Grayscale representation of data packets in the session. As mentioned above, the size of each data packet in the session is a multiple of 256 bytes. This method requires converting the first 256 bytes into the first 16×16 patch in the order of a one-dimensional matrix during the conversion process. Then, every 256 bytes is converted into a 16×16 patch, for a total of 14×14 patches, finally forming a 224×224 two-dimensional matrix.

[0058] Step 3: Data augmentation of the original dataset;

[0059] Step 3.1: Introduce a variable to record the current position, i.e., the starting position of the window;

[0060] Step 3.2: Determine the difference between the number and location of data packets in the current raw traffic file;

[0061] Step 3.3: After a packet is processed in a session, the position variable points to the next packet.

[0062] Step 4: Assign labels to grayscale images and divide the training dataset and test dataset, including two cases: binary classification and multi-class classification of malicious traffic;

[0063] Step 4.1: Assign labels to grayscale images. Read the first data packet in the session in sequence, obtain the timestamp information of the data packet, and check in the label file whether there is a session record corresponding to the timestamp. If the corresponding session is found, use the label category corresponding to the session as part of the name of the generated grayscale image.

[0064] Step 4.2: Dataset partitioning. The dataset used in this invention consists of approximately 70,000 grayscale images after the data packets in the session are visualized. The dataset is divided into training and test sets in a 6:1 ratio.

[0065] Step 5: Construct a simplified hybrid Vision Transformer model (SHViT) for malicious traffic identification, such as... Figure 2 As shown;

[0066] Step 5.1: Use the MobileNetV2 model to downsample and generate feature maps of data packets, which will then be used as input to the Transformer Encoder;

[0067] Step 5.2: Construct the SHViT Block in the simplified Hybrid Vision Transformer model, such as... Figure 3 As shown;

[0068] Step 5.2.1: Use the feature map as the input to this module, and use a convolutional layer to change the number of channels in the feature map;

[0069] Step 5.2.2: Improve the multi-head self-attention mechanism of the Transformer model, that is, ignore the attention calculation of neighboring pixels;

[0070]

[0071]

[0072] Step 5.2.3: After concatenating the input feature map and the output of the Transformer model through residual connection, the number of channels is kept consistent with that of the input feature map;

[0073] Step 5.3: Pass the output of the SHViT Block through a global pooling layer and a linear layer to obtain the complete SHViT model;

[0074] Step 5.4: Configure training parameters as shown in Table 2, input the training dataset, and improve the accuracy of malicious traffic identification;

[0075] Table 2 Experimental Parameters;

[0076]

[0077] Step 5.5: Input the test dataset and calculate the classification accuracy obtained in each round of training, such as... Figure 4 As shown.

[0078] Accuracy is defined as follows:

[0079]

[0080] In this model, TP represents positive samples predicted as positive, TN represents negative samples predicted as negative, FP represents negative samples predicted as positive, and FN represents positive samples predicted as negative. TP and TN represent the number of correctly classified samples. Therefore, accuracy indicates the proportion of correctly classified samples out of all samples. The higher the accuracy, the better the model's training performance.

[0081] Step 5.6: Calculate the number of parameters and computational cost of the model, i.e., the metrics for evaluating the model's lightweightness; the parameter and computational cost of convolutional and fully connected layers are defined as follows:

[0082] parameters conv =(k w ×k h ×C in +1)×C out (4)

[0083] parameters Fc =(N in +1)×N out (5)

[0084] FLOPs conv =[(C in ×k w ×k h )+(C in ×k w ×k h -1)+1]×C out ×w×h (6)

[0085] FLOPs FC =[N in +(Nin -1)+1]×N out (7)

[0086] k w kh represents the size of the convolution kernel, C in and C out N represents the number of input and output channels of the convolutional layer. in and N out represent the size of the input and output dimensions in the fully connected layer, respectively; w and h represent the size of the feature map; and 1 in the formula represents the bias in the convolutional layer.

Claims

1. A ViT improved lightweight malicious traffic identification method, characterized in that, Specifically comprising the following steps: Step 1: Obtain a malicious traffic dataset and perform session segmentation and parsing processing; Step 2: Perform packet image representation on the sessions parsed in step 1 to generate grayscale pictures; Step 3: Use a sliding packet-based method to perform data augmentation on the classes with less sample quantity in the dataset, obtaining a balanced dataset for each class; Step 4: Assign labels to the grayscale pictures converted by the packet image method according to the label file provided by the dataset, divide the training dataset and test dataset in a 6:1 ratio, including malicious traffic binary classification and multi-classification; Step 5: Construct a simplified hybrid Vision Transformer model SHViT for malicious traffic identification; Step 5.1: Use the mobileNetV2 model for downsampling to generate a feature map of the packet as input for the TransformerEncoder; Step 5.2: Construct the SHViT Block in the simplified hybrid Vision Transformer model; Step 5.3: Pass the output of the SHViT Block through a global pooling layer and a linear layer to obtain the complete SHViT model; Step 5.4: Configure the training parameters, input the training dataset, and improve the accuracy of malicious traffic identification; Step 5.5: Input the grayscale pictures converted from the processed original traffic files into the model to obtain the malicious traffic identification results; Step 5.6: Calculate the parameter quantity and computational complexity of the model, which are the indicators for evaluating the lightweight model; the parameter quantity and computational complexity of the convolutional layer and fully connected layer are defined as follows: parameters conv = (k w × k h × C in + 1) × C out (4) parameters FC = (N in + 1) x N out (5) FLOPs conv = [(C in × k w × k h ) + (C in × k w × k h - 1) + 1] x C out x w x h (6) FLOPs FC = [N in + (N in - 1) + 1] x N out (7) where k w and k h represent the size of the convolution kernel, C in and C out represent the number of input and output channels of the convolution layer, N in and N out represent the size of the input and output dimensions in the fully connected layer, w and h represent the size of the feature map, and 1 in the formula represents the bias in the convolution layer.

2. The method of claim 1, wherein the method is based on a ViT improved lightweight malicious traffic recognition method. Step 1 is specifically: Step 1.1: Obtain a public IoT network traffic dataset from the network; Step 1.2: Divide the original traffic files saved in the dataset into several session files according to TCP flow or UDP flow; Step 1.3: Parse the session file to identify the format of the data packets stored in the file, and delete the format information, retaining only the TCP / IP protocol specified data packet information.

3. The method of claim 1, wherein the method is based on a ViT improved lightweight malicious traffic recognition method. Step 2 is specifically: Step 2.1: Limit the length of the session, set the length of the session to L session The range of padding is 576 bytes to 50176 bytes, the exceeding part will be discarded, the insufficient part will be filled, and padding is selected as the value of the filling, which is an integer in the range of 0 to 255. Step 2.2: Pad the length of the data packet in the normative session, requiring the length of each data packet in the session to be L packet , with the remainder L packet being padded. wherein L session represents the length of the session, N patch represents the number of partitions; Step 2.3: Convert the data packet grayscale image representation in the session to a two-dimensional matrix in order, and generate a grayscale picture.

4. The method of claim 1, wherein the method is based on a ViT improved lightweight malicious traffic recognition method. Step 3 is specifically: Step 3.1: Introduce a variable to record the current packet position, i.e., the start position of the window; Step 3.2: Determine the difference between the number and position of the data packets in the current original traffic file; Step 3.3: After the data packet at the current position in the session is processed by the packet image method, the position variable points to the next data packet.

5. The method of claim 1, wherein the method is based on a ViT improved lightweight malicious traffic recognition method. Step 5.2 is specifically: Step 5.2.1: Use the feature map as the input of this module, and use the convolutional layer to change the number of channels of the feature map; Step 5.2.2: Improve the multi-head self-attention mechanism of the Transformer model, i.e., ignore the attention calculation of adjacent pixels; Where n represents the length of the input sequence of the Transformer Encoder, and k represents the multiple of the reduction of the input sequence length; Step 5.2.3: The input feature map and the output through the Transformer model are spliced by residual connection and kept consistent with the channel number of the input feature map.

Images