An improved method of ALBERT model applied to encrypted traffic classification

Abstract

The application discloses an improved method of an ALBERT model applied to encrypted traffic classification, and belongs to the field of network encrypted traffic classification. Firstly, a network traffic packet is captured, and after pretreatment, the network traffic packet is converted into a standard input format required by the ALBERT model. The Self-Attention layer and the FFN layer in the ALBERT model are improved to obtain an improved ALBERT model. By using the standard input of the ALBERT model, the improved ALBERT model is unsupervised pre-trained through an MLM method; when the unsupervised pre-training is completed, the model is fine-tuned by using category label data, and a mask shielding layer is removed, and the final ALBERT model is saved. The application can effectively reduce resource loss, maintain stable accuracy, and has a wider application prospect.

Description

Technical Field

[0001] This invention belongs to the field of network encrypted traffic classification, and specifically relates to an improved method for the ALBERT model applied to encrypted traffic classification. Background Technology

[0002] Network traffic classification is a technology used to identify network traffic categories. To protect the rights of internet users on the open network, traffic encryption technology is increasingly needed. However, this also provides opportunities for malicious actors to implant malware. Therefore, the demand for network traffic classification is growing daily to build a healthy network environment.

[0003] Network encrypted traffic classification employs both traditional machine learning methods and deep learning-based approaches. Due to the complexity of dynamic network environments, traditional traffic classification methods are no longer sufficient. In contrast, deep learning has achieved better results. However, its high resource consumption makes it difficult to apply to smartphones or ordinary computers. Summary of the Invention

[0004] The purpose of this invention is to reduce resource consumption without compromising model accuracy, thereby making it easier to implement. It provides an improved method for ALBERT models applied to encrypted traffic classification, using fewer layers of ALBERT to handle encrypted traffic classification problems and identify malicious traffic data.

[0005] An improved method for applying the ALBERT model to encrypted traffic classification, with the following specific steps:

[0006] Step 1: Capture a network traffic data packet, preprocess it, and convert it into the standard input format required by the ALBERT model.

[0007] The specific process of preprocessing is as follows:

[0008] First, parse network traffic data packets to extract SSL handshake data and transmission data;

[0009] Then, remove the Ethernet header, IP and TCP port information from the SSL handshake data and transmission data, and convert each bit value into the corresponding double character in hexadecimal to generate two strings, denoted as A and B respectively.

[0010] Finally, A and B are each divided into segments of 4 characters from left to right, generating two character sequences. These sequences are then combined. The length of the combined character sequence is checked to see if it meets the standard input length s of ALBERT. If the length is less than s, zeros are added to the end. If the length is greater than s, the end strings of the longer sequence in A and B are gradually removed until the condition is met.

[0011] The two character sequences A and B are respectively:

[0012] ,

[0013] The transformed ALBERT standard input is as follows:

[0014]

[0015] in .

[0016] Step 2: Improve the Self-Attention layer and FFN layer in the ALBERT model to obtain the improved ALBERT model;

[0017] The self-attention layer is a mutual self-attention layer between tokens. Its improvement is the addition of a mask layer to the attention matrix. The mask layer's function is to mask the SSL handshake or file data transmission part, so that other parts, except for the masked part, pay attention to the global token.

[0018] The improvement to the FFN layer involves reducing the number of output nodes in a certain layer of the original FFN neural network, and then adding a new neural network. This new neural network takes the reduced number of output nodes as input and outputs nodes equal to the reduced number of nodes. Finally, the input and output nodes of this new neural network are concatenated to form an output of the same size as the original FFN. The reduction ratio of the network output nodes varies depending on the task.

[0019] Step 3: Use the masked language modeling (MLM) method to perform unsupervised pre-training on the improved ALBERT model;

[0020] The specific process of unsupervised pre-training is as follows:

[0021] Step 301: Input the data in the standard input format obtained in Step 1 into the improved ALBERT model;

[0022] Step 302: With a 10% probability, randomly replace some tokens in the model input with any token, and use the original tokens as the prediction labels;

[0023] Step 303: Use the MLM method to perform feedforward computation and use cross-entropy to calculate the difference between the model output and the predicted label of the replaced token;

[0024] Step 304: Based on the difference between the model output and the predicted label, optimize the model parameters through feedback adjustment, and determine whether the training stopping condition is met. If not, return to step 302 to continue training until the training stopping condition is met; if so, the unsupervised pre-training of the model is complete.

[0025] The training stop condition is:

[0026] Once all training data is input into ALBERT sequentially, one round of model training is completed, and the count is incremented by 1 (initially 0). When the count value is greater than the preset batch size, the stopping condition is met, and training is terminated.

[0027] Step 4: After unsupervised pre-training is complete, fine-tune the model using the class label data, remove the mask layer, and save the final ALBERT model.

[0028] The fine-tuning model is specifically as follows:

[0029] Step 401: Input the standard input and class label data obtained in Step 1 into the pre-trained ALBERT model.

[0030] Step 402: Extract the output corresponding to [CLS] through feedforward operation, and then pass it through a trainable fully connected neural network to obtain an output of the same size as the category label. Calculate the difference between the output and the real category label data using cross-entropy.

[0031] Step 403: Based on the difference between the model output and the real category label data, optimize the model parameters through feedback adjustment, and determine whether the stopping condition is met. If not, return to step 402 until the stopping condition is met; if yes, execute step 404.

[0032] The stopping condition is determined by comparing whether the count of the model training is greater than a preset batch number.

[0033] Step 404: Remove the mask layer added to the Self-Attention layer in the model, cut off the input tensor of the masked part, and save the model.

[0034] The advantages and beneficial effects of this invention are as follows:

[0035] The method for optimizing training and improving models provided by this invention makes the session data transmitted over the network more suitable for the ALBERT model, and the trained model also significantly reduces resource consumption without reducing the original accuracy. Attached Figure Description

[0036] Figure 1 This is a schematic diagram of data preprocessing according to the present invention;

[0037] Figure 2 This is a schematic diagram of the improved ALBERT model of the present invention;

[0038] Figure 3 This is a flowchart of the network encrypted traffic classification technology of the present invention; Detailed Implementation

[0039] To make the objectives, technical solutions, and advantages of this invention clearer, the technical solutions of this invention will be further described in detail below with reference to the accompanying drawings and implementation steps. This invention is described in three aspects.

[0040] This invention optimizes data preprocessing and model pretraining methods, and improves ALBERT, including Self-Attention layers and FFN layers. The overall process is as follows: Figure 3 As shown, the specific steps are as follows:

[0041] Step 1: Capture a continuous, complete, and clean network traffic data packet, and after preprocessing, convert it into the standard input format required by the ALBERT model.

[0042] like Figure 1 As shown, the specific preprocessing process is as follows:

[0043] First, parse network traffic data packets to extract SSL handshake data and transmission data;

[0044] Then, remove the Ethernet header, IP and TCP port information from the SSL handshake data and transmission data, and convert each bit value into the corresponding double character in hexadecimal to generate two strings, denoted as A and B respectively.

[0045] Finally, A and B are each divided into segments of 4 characters from left to right, generating two character sequences. These sequences are then combined. The length of the combined character sequence is checked to see if it meets the standard input length s of ALBERT. If the length is less than s, zeros are added to the end. If the length is greater than s, the end strings of the longer sequence in A and B are gradually removed until the condition is met.

[0046] The two character sequences A and B are respectively:

[0047] ,

[0048] The transformed ALBERT standard input is as follows:

[0049]

[0050] in .

[0051] Step 2: Improve the Self-Attention layer and FFN layer in the ALBERT model to obtain the improved ALBERT model;

[0052] In this work, the network model is pre-set to have the same structure for each transform layer, with the word vector dimension set to 768.

[0053] (1) Improvement of Self-Attention layer:

[0054] The self-attention layer is a layer where tokens mutually self-attention each other; its improvement lies in adding a masking layer to the attention matrix. For example... Figure 2 As shown in (c), this means that the SSL handshake portion is masked. This allows the file data transfer portion and the token [CLS] to focus on the global token, but the content of the SSL handshake phase cannot. Conversely, masking the file data transfer portion can also produce a similar result. Finally, after fine-tuning, a slicing operation is added to split the input tensor, so that the masked portion no longer participates in the calculation of the self-attention matrix. This design is more suitable for the characteristics of encrypted network data, and it also saves a significant amount of computational power during model inference.

[0055] (2) Improvements to the FFN layer:

[0056] A schematic diagram of the original FFN layer is shown below. Figure 2 (a) The dimensions of the token vector at different parts are: 768 --> 3072 --> 768, where "-->" indicates a fully connected neural network, the tail of the arrow represents the input dimension, and the head represents the output dimension. The method used in this embodiment is as follows: Figure 2 As shown in (b), specifically, the number of output nodes in the last layer of the original FFN is reduced by 25%. Then, a new neural network is added, using this output as input, and its output node count is the number of nodes previously reduced. Finally, the input and output nodes of this new neural network are concatenated to form an output of the same size as the original FFN. This method reduces the number of parameters in the FFN layer by 10% and implicitly increases the model depth, extracting higher-level features. Of course, the reduction ratio of network output nodes can be varied depending on the task.

[0057] Step 3: Use the masked language modeling (MLM) method to perform unsupervised pre-training on the improved ALBERT model;

[0058] The specific process of unsupervised pre-training is as follows:

[0059] Step 301: Input the data in the standard input format obtained in Step 1 into the improved ALBERT model;

[0060] Step 302: With a 10% probability, randomly replace some tokens in the model input with any token, and use the original tokens as the prediction labels;

[0061] Step 303: Use the MLM method to perform feedforward computation and use cross-entropy to calculate the difference between the model output and the predicted label of the replaced token;

[0062] Step 304: Based on the difference between the model output and the predicted label, optimize the model parameters through feedback adjustment (AdamW optimizer), and determine whether the training stopping condition is met. If not, return to step 302 to continue training until the training stopping condition is met; if so, the unsupervised pre-training of the model is complete.

[0063] The training stop condition is:

[0064] Once all training data is input into ALBERT sequentially, one round of model training is completed, and the count is incremented by 1 (initially 0). When the count value is greater than the preset batch size, the stopping condition is met, and training is terminated.

[0065] Step 4: After unsupervised pre-training is complete, fine-tune the model using the class label data and save it.

[0066] The fine-tuning model is specifically as follows:

[0067] Step 401: Input the standard input and class label data obtained in Step 1 into the pre-trained ALBERT model.

[0068] Step 402: Extract the output corresponding to [CLS] through feedforward operation, and then pass it through a trainable fully connected neural network to obtain an output of the same size as the category label. Calculate the difference between the output and the real category label data using cross-entropy.

[0069] Step 403: Based on the difference between the model output and the real category label data, optimize the model parameters through feedback adjustment, and determine whether the stopping condition is met. If not, return to step 402 until the stopping condition is met; if yes, execute step 404.

[0070] The stopping condition is determined by comparing whether the count of the model training is greater than a preset batch number.

[0071] Step 404: Remove the mask layer added to the Self-Attention layer in the model, cut off the input tensor of the masked part, and save the model.

[0072] After fine-tuning, since the mask layer is fixed along with the model parameters, the self-attention calculation and output of the masked part of the tokens are meaningless, so they should be removed to simplify the calculation.

[0073] It should be noted and understood that various modifications and improvements can be made to the invention described in the detailed description above without departing from the spirit and scope of the invention as claimed in the appended claims. Therefore, the scope of the claimed technical solutions is not limited to any specific exemplary teachings given.

Claims

1. An improved method for applying the ALBERT model to encrypted traffic classification, characterized in that, The specific steps are as follows: Step 1: Capture network traffic data packets, preprocess them, and convert them into the standard input format required by the ALBERT model to serve as the training set; The preprocessing process for network traffic data packets is as follows: First, parse network traffic data packets to extract SSL handshake data and transmission data; Then, remove the Ethernet header, IP and TCP port information from the SSL handshake data and transmission data, and convert each bit value into the corresponding double character in hexadecimal to generate two strings, denoted as A and B respectively. Finally, A and B are each divided into segments of 4 characters from left to right, generating two character sequences. These sequences are then combined. The length of the combined character sequence is checked to see if it meets the standard input length s of ALBERT. If the length is less than s, zeros are added to the end. If the length is greater than s, the end strings of the longer sequence in A and B are gradually removed until the condition is met. The two character sequences A and B are respectively: ， The transformed ALBERT standard input is as follows: in ; Step 2: Improve the Self-Attention layer and FFN layer in the ALBERT model to obtain the improved ALBERT model; The self-attention layer is a mutual self-attention layer between tokens. Its improvement is to add a mask layer to the attention matrix. The mask layer is used to mask the SSL handshake part or the file data transmission part, so that other parts except the masked part pay attention to the global token. The improvement of the FFN layer is to reduce the number of output nodes of a certain layer of the original FFN neural network, and then add a new neural network. The new neural network takes the reduced number of output nodes as input and outputs nodes equal to the reduced number of nodes. Finally, the input and output nodes of this new neural network are concatenated to form the same output size as the original FFN. Step 3: Using the MLM method, the improved ALBERT model is pre-trained in unsupervised mode using the training set. MLM stands for masked language modeling. The specific process of unsupervised pre-training is as follows: Step 301: Input the training set data in standard input format into the improved ALBERT model; Step 302: With a 10% probability, randomly replace some tokens in the model input with any token, and use the original tokens as the prediction labels; Step 303: Use the MLM method to perform feedforward computation and use cross-entropy to calculate the difference between the model output and the predicted label of the replaced token; Step 304: Based on the difference between the model output and the predicted label, optimize the model parameters through feedback adjustment, and determine whether the training stopping condition is met. If not, return to step 302 to continue training until the training stopping condition is met; if so, the unsupervised pre-training of the model is complete. The training stop condition is: Once all training data is input into ALBERT sequentially, one round of model training is completed, and the count is incremented by 1. When the count value is greater than the preset batch size, the stopping condition is met, and training is terminated. Step 4: After unsupervised pre-training is complete, fine-tune the model using the class label data, remove the mask layer, and save the final ALBERT model.

2. The improved method for applying the ALBERT model to encrypted traffic classification according to claim 1, characterized in that, In the improvement of the FFN layer in step two, the reduction ratio of network output nodes varies with the task.

3. An improved method for applying the ALBERT model to encrypted traffic classification according to claim 1, characterized in that, In step four, the specific process of fine-tuning the pre-trained ALBERT model is as follows: Step 401: Input the standard input and class label data obtained in Step 1 into the pre-trained ALBERT model; Step 402: Extract the output corresponding to [CLS] through feedforward operation, and then pass it through a trainable fully connected neural network to obtain an output of the same size as the category label. Calculate the difference between the output and the real category label data using cross-entropy. Step 403: Based on the difference between the model output and the real category label data, optimize the model parameters through feedback adjustment, and determine whether the stopping condition is met. If not, return to step 402 until the stopping condition is met; if yes, proceed to step 404. The stopping condition is determined by comparing whether the count of the model training is greater than a preset batch number. Step 404: Remove the mask layer added to the Self-Attention layer in the model, cut off the input tensor of the masked part, and save the model.

Images