Multi-task learning network flow classification method

Abstract

The invention discloses a multi-task learning network traffic classification method, and the method comprises the steps: S1, collecting a traffic data set, and carrying out the preprocessing of the traffic data set; s2, constructing a network architecture; s3, training a universal pre-training model according to the network architecture; and S4, carrying out fine adjustment on the pre-training model. And the learning ability of the model on the internal structure and mode of data is enhanced.

Description

Technical Field

[0001] The present invention relates to the field of network security and management, and in particular to a multi-task learning network traffic classification method. Background Art

[0002] Network traffic classification technology is an important network management tool. It classifies traffic into different applications or services by analyzing the characteristics and behaviors of network traffic. This technology can help network administrators implement more fine-grained control policies, thereby providing users with better service quality. With the widespread application of traffic encryption technology, protecting user privacy and anonymity has become an important issue. However, malicious traffic and cybercriminals often use encryption technologies such as Tor and VPN to evade monitoring systems, which poses a huge challenge to traffic classification. The traditional method of capturing patterns and keywords in data packets from the payload, namely deep packet inspection (DPI), cannot be applied to encrypted traffic. In addition, in the face of growing network complexity, the model must be able to handle multiple classification tasks simultaneously. Therefore, how to capture implicit and robust patterns in diverse encrypted traffic and support accurate multi-task traffic classification has become the key to ensuring network security and effective management.

[0003] In recent years, pre-trained models such as BERT, GPT, and ResNet have made breakthrough progress in fields such as natural language processing (NLP) and computer vision (CV). The reason why these models are powerful is mainly because they can use large-scale data resources, thereby reducing the dependence on labeled data and improving the robustness and migration ability of the model. In the research field of network traffic classification, features are effectively extracted through multi-head attention mechanisms, and the contextual information of data packets is used to enhance the representation and generalization capabilities of the model, but there is a lack of pre-training tasks designed to match the characteristics of traffic data. This is because network traffic data has its unique structure and pattern, and standard pre-training tasks may not be able to fully capture these characteristics.

[0004] The existing technologies have problems such as traffic classification relying on a large amount of labeled data, performance degradation in complex scenarios, and high computing power requirements. Summary of the invention

[0005] In view of the above problems, the present invention is proposed to provide a multi-task learning network traffic classification method that overcomes the above problems or at least partially solves the above problems.

[0006] According to one aspect of the present invention, a multi-task learning network traffic classification method is provided, the classification method comprising:

[0007] Step S1: collecting a flow data set and preprocessing the flow data set;

[0008] Step S2: constructing a network architecture;

[0009] Step S3: training a general pre-training model according to the network architecture;

[0010] Step S4: fine-tune the pre-trained model.

[0011] Optionally, the step S1: collecting a traffic data set and preprocessing the traffic data set specifically includes:

[0012] According to the design of the payload bytes, it is revealed in the initial packet of the data flow in the real network environment;

[0013] Traffic in multiple session flows is captured according to IP addresses, port numbers, and protocol types, and time-adjacent payload data is obtained from each session flow as an input unit.

[0014] Optionally, the traffic data set uses an ISCXVPN-nonVPN real data set.

[0015] Design a trace sequence for a capture file in It is a session flow consisting of packets from source to destination;

[0016] Identified by a five-tuple (IPsrc:PORTsrc,IPdst:PORTdst,Protocol), S-payload is defined as:

[0017]

[0018] Where m and n represent the maximum number of one-way packets from source to destination and from destination to source respectively.

[0019] Optionally, the step S2: constructing a network architecture specifically includes:

[0020] A multi-layer bidirectional Transformer is used, each layer consists of a multi-head self-attention mechanism to share information when performing traffic feature classification and application identification tasks.

[0021] Optionally, the step S2: constructing a network architecture specifically includes:

[0022] Use the self-attention mechanism to reconstruct the representation of the target word based on the context S-paload to capture the subtle differences in the data stream;

[0023] The self-attention mechanism calculates each input x through a similarity function i With all his inputs x1-x kcompatibility between them, and calculate x i The weighted sum of , the self-attention mechanism is calculated as follows:

[0024]

[0025] x i and x j The similarity between them is represented by a scaled dot product operator, where d k It is k j The dimension of , Z is the normalization factor;

[0026] Feature extraction on the data stream will be performed for each x i Perform M linear projections to generate multiple attention vectors [att i,1 ,att i,2 ,...,att i,M ];

[0027] Use a concatenation operator to get the final attention vector:

[0028]

[0029] Multi-head self-attention can better capture the subtle differences in data.

[0030] Calculate the similarity scores between multiple datagrams to help the model understand and integrate the information between the previous and next datagrams of the entire input sequence;

[0031] By processing information in parallel, it can capture features in multiple dimensions.

[0032] Optionally, the step S3: training a general pre-training model according to the network architecture specifically includes:

[0033] Special tokens [CLS], [SEP], [PAD], and [MASK] are introduced into the token sequence for classification tasks, pre-training tasks, insufficient padding sequences, and token masks;

[0034] Self-supervised learning is used on unlabeled network traffic data to train a general pre-trained model through the network architecture.

[0035] Optionally, the step S3: training a general pre-training model according to the network architecture further includes:

[0036] The two pre-training tasks capture the contextual relationship between traffic bytes by predicting mask tokens, and the correct transmission order by predicting S-payload;

[0037] When dealing with non-semantic data such as network traffic, the dynamic masking task forces the model to predict the masked data parts, and the model learns specific patterns in the packet sequence;

[0038] During pre-training, at the beginning of each epoch, the model randomly selects tokens to mask;

[0039] Each marker has a 15% chance of being selected;

[0040] The model attempts to predict the original value of the masked token by replacing it with [MASK] with 80% probability, or choosing a random token to replace it with 10% probability;

[0041] In an epoch, the above process is repeated, and different markers are masked each time;

[0042] We use negative log-likelihood as the loss function and formally define it as:

[0043]

[0044] Among them, θ represents the set of trainable parameters in the model, k is the total number of tags in the sequence, and MASK i is the i-th mask token, S-Payload{token}_i is the actual token at the corresponding position, X is the given before and after payload, and theta is the parameter of the model;

[0045] The formula evaluates the overall prediction performance of the model by summing the logarithm of the predicted probability of each mask position;

[0046] The input is Payload, separated by [SEP]. The first Payload is preceded by [CLS], and the second one is also followed by [SEP].

[0047] The output of the SPP task is a binary classification, which determines whether the two payloads are arranged in the order of the original text;

[0048] If yes, output IsPayload, otherwise, output NotPayload;

[0049]

[0050] Optionally, the step S4: fine-tuning the pre-trained model specifically includes:

[0051] In the fine-tuning stage, the model parameters of the pre-trained model trained on task-specific data are further adjusted using labeled traffic data to capture task-related features.

[0052] The present invention provides a multi-task learning network traffic classification method, the classification method comprising: step S1: collecting traffic data sets and preprocessing the traffic data sets; step S2: building a network architecture; step S3: training a general pre-trained model according to the network architecture; step S4: fine-tuning the pre-trained model to enhance the model's ability to learn the intrinsic structure and pattern of the data.

[0053] The above description is only an overview of the technical solution of the present invention. In order to more clearly understand the technical means of the present invention, it can be implemented according to the contents of the specification. In order to make the above and other purposes, features and advantages of the present invention more obvious and easy to understand, the specific implementation methods of the present invention are listed below. BRIEF DESCRIPTION OF THE DRAWINGS

[0054] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the accompanying drawings required for use in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other accompanying drawings can be obtained based on these accompanying drawings without paying creative work.

[0055] Figure 1 A flowchart of a multi-task learning network traffic classification method provided by an embodiment of the present invention;

[0056] Figure 2 A schematic diagram of S-payload data integration provided by an embodiment of the present invention;

[0057] Figure 3 A schematic diagram of data input provided by an embodiment of the present invention;

[0058] Figure 4 A schematic diagram of the overall network structure provided by an embodiment of the present invention. DETAILED DESCRIPTION

[0059] The exemplary embodiments of the present disclosure will be described in more detail below with reference to the accompanying drawings. Although the exemplary embodiments of the present disclosure are shown in the accompanying drawings, it should be understood that the present disclosure can be implemented in various forms and should not be limited by the embodiments set forth herein. On the contrary, these embodiments are provided to enable a more thorough understanding of the present disclosure and to fully convey the scope of the present disclosure to those skilled in the art.

[0060] The terms "comprises" and "having" and any variations thereof in the description embodiments, claims and drawings of the present invention are intended to cover non-exclusive inclusions, for example, including a series of steps or units.

[0061] The technical solution of the present invention is further described in detail below in conjunction with the accompanying drawings and embodiments.

[0062] like Figure 1 As shown, a multi-task learning network traffic classification method includes:

[0063] Step S1: Collecting traffic data sets and preprocessing the traffic data sets;

[0064] Step S2: constructing a network architecture;

[0065] Step S3: training a general pre-trained model according to the network architecture;

[0066] Step S4: Fine-tune the pre-trained model.

[0067] Step S1, prepare the data set, preprocess the general traffic data set, and design the payload bytes to reveal visible information, especially in the first few packets of the data flow in the real network environment. The traffic is captured into different session flows according to the IP address, port number and protocol type, and the temporally adjacent payload data is obtained from each session flow as the input unit.

[0068] In step S1, the ISCXVPN-nonVPN real data set is used. In order to obtain the same-source payload data, we designed a capture file tracking sequence Trace = {flow i ,i∈N +}, where flow = {p j ,j∈N +} is a session flow consisting of source to destination packets, identified by a five-tuple (IPsrc:PORTsrc,IPdst:PORTdst,Protocol), and S-payload is defined as:

[0069]

[0070] Where m and n represent the maximum number of one-way packets from source to destination and from destination to source respectively.

[0071] In step S2, the network architecture is constructed, which uses a multi-layer bidirectional Transformer. Each layer consists of a multi-head self-attention mechanism. This design enables the model to share information when performing traffic feature classification and application identification tasks, thereby improving the accuracy and efficiency of classification.

[0072] In step S2, the self-attention mechanism is used to reconstruct the representation of the target word according to the context S-payload, so as to better capture the subtle differences in the data stream. Generally speaking, the self-attention mechanism is to calculate each input x through a similarity function i With all his inputs x1-x kThe compatibility between them, and further calculate x i The weighted sum of , which means its overall data flow information. Our self-attention mechanism is specifically calculated as follows:

[0073]

[0074] x i and x j The similarity between them is represented by a scaled dot product operator, where d k It is k j The dimension of , Z is the normalization factor. Feature extraction on the data stream will be for each x i Perform M linear projections to generate multiple attention vectors [att i,1 ,att i,2 ,...,att i,M ]. Then, a concatenation operator is used to obtain the final attention vector:

[0075]

[0076] This multi-head self-attention can better capture the subtle differences in datagrams. It helps the model understand and integrate the information between the previous and next datagrams of the entire input sequence by calculating the similarity scores between different datagrams. By processing information in parallel, it can capture features of different dimensions, thereby providing richer representation information when dealing with complex tasks.

[0077] In step S3, the pre-training phase, special tokens [CLS], [SEP], [PAD], and [MASK] are introduced into the token sequence of step S2 for classification tasks, pre-training tasks, padding sequences of insufficient length, and token masks. A general pre-training model is trained on unlabeled network traffic data through the network architecture of step S3 through self-supervised learning.

[0078] In step S3, the two pre-training tasks we proposed capture the contextual relationship between traffic bytes by predicting mask tokens, and the correct transmission order by predicting S-payload. When dealing with non-semantic data such as network traffic, the dynamic masking task forces the model to predict the masked data parts. The model can learn specific patterns in the packet sequence, which may indicate certain characteristics of network behavior. During the pre-training process, at the beginning of each epoch, the model randomly selects some tokens to mask. Usually, each token has a 15% probability of being selected. We replace it with [MASK] with 80% probability, or choose a random token to replace it, or keep it unchanged with 10% probability. The model then tries to predict the original values ​​of these masked tokens, which helps the model learn the deep structure of the language. In the next epoch, the above process is repeated, but different tokens are masked each time. We use negative log-likelihood as the loss function and formally define it as:

[0079]

[0080] θ represents the set of trainable parameters in the model, (k) is the total number of tokens in the sequence, and MASK i is the i-th mask token, S-Payload{token}_i is the actual token at that position, X is the given previous and next payloads, and theta is the model parameter. The formula evaluates the overall prediction performance of the model by summing the logarithms of the predicted probabilities for each mask position. The purpose of the S-payload prediction task is to allow the model to learn the coherence between payloads, such as whether they are arranged in the order of traffic data transmission. First, the input is the payload, separated by [SEP]. The first payload is preceded by [CLS], and the second sentence is also followed by [SEP]. The output of the SPP task is a binary classification that determines whether the two payloads are arranged in the order of the original text. If so, output IsPayload, if not, output NotPayload.

[0081]

[0082] Step S4, the fine-tuning stage, uses a small amount of labeled traffic data on the data of a specific task to further adjust the model parameters of the pre-trained model trained in step S4 so that it can better capture the features related to the task.

[0083] In step S4, fine-tuning is used to allow the model to serve downstream classification tasks. First, fine-tuning can utilize the high-level features that the pre-trained model has learned. Apply any category of traffic representation; second, the input of the pre-trained model is at the datagram byte level, so the downstream task of classifying packets and flows is converted into corresponding datagram byte tokens and classified by the model; finally, through fine-tuning, the features learned by the pre-trained model are used to explain the model's classification decision on the input data, thereby enhancing the interpretability of the model and making it easier to understand and trust. The special [CLS] token output by the pre-trained model models the representation of the entire input traffic and is directly used for classification.

[0084] (1) S-payload data integration

[0085] like Figure 2 As shown in the real network environment, the massive traffic contains various traffic of different categories (such as different applications, protocols or services), and it is difficult to learn a stable and discriminative traffic representation. Therefore, before representing the traffic, the flows with the same IP, port and protocol are first separated from the trace. Therefore, each separated flow comes from the same traffic category containing the complete flow session. Specifically, a trace sequence of capture files is designed. in It is a session flow consisting of source to destination packets, identified by a five-tuple (IPsrc:PORTsrc,IPdst:PORTdst,Protocol). S-payload is defined as:

[0086]

[0087] Where m and n represent the maximum number of one-way packets from source to destination and from destination to source respectively.

[0088] (2) Schematic diagram of the input layer

[0089] Each token in the payload is represented by three types of embeddings: token embedding, position embedding, and segment embedding. Figure 3 The visual representation of the structure is shown. Token embedding is the process of mapping each token into a high-dimensional vector space. Using a vector of size d model The embedding matrix to represent the embedding of all tags. For each tag x in the input sequence i , and its corresponding token embedding vector is expressed as By formula e i =E[i] calculation, position embedding is to encode the position information of each token in the sequence into a vector representation. A special position encoding method is used, usually using sine and cosine functions, as shown in the formula

[0090]

[0091] Among them, pos represents the position of the tag in the sequence, i represents the dimension index of the position embedding vector, and d model is the embedding dimension. The dimension of the position embedding matrix is ​​(n,d model ), each row corresponds to a position in the input sequence. Segment embedding is used to process the feature information of different parts of the input sequence. For each input token, it is necessary to specify which segment it belongs to. The segment information is represented by a binary vector. For each token, if it belongs to the first segment, the first element of the vector is 1, otherwise it is 0; similarly, if it belongs to the second segment, the second element of the vector is 1, otherwise it is 0. Multiply the binary segment vector by a size d model The specific formula is s i = S·SegmentEncoding(i). Finally, the token embedding, position embedding, and segment embedding are added together to get the final input representation: Input i =e i +PE i +S i .

[0092] (4) Pre-training and fine-tuning

[0093] like Figure 4 As shown in the figure, the two proposed pre-training tasks capture the contextual relationship between traffic bytes by predicting mask tokens, and capture the correct transmission order by predicting S-payload. When processing non-semantic data such as network traffic, the dynamic mask task forces the model to predict the masked data part, and the model can learn specific patterns in the packet sequence, indicating certain characteristics of network behavior.

[0094] During pre-training, at the beginning of each epoch, the model randomly selects some tokens to mask. Typically, each token has a 15% probability of being selected. It is replaced with [MASK] with 80% probability, or a random token is selected to replace it, or it remains unchanged with 10% probability. The model then tries to predict the original values ​​of these masked tokens, which helps the model learn the deep structure of the language. In the following epochs, the above process is repeated, but different tokens are masked each time. The negative log-likelihood is used as the loss function and is formally defined as:

[0095]

[0096] θ represents the set of trainable parameters in the model, (k) is the total number of tokens in the sequence, and MASK iis the i-th mask token, S-Payload{token}_i is the actual token at that position, X is the given previous and next payloads, and theta is the model parameter. The formula evaluates the overall prediction performance of the model by summing the logarithms of the predicted probabilities for each mask position. The purpose of the S-payload prediction task is to allow the model to learn the coherence between payloads, such as whether they are arranged in the order of traffic data transmission. First, the input is the payload, separated by [SEP]. The first payload is preceded by [CLS], and the second sentence is also followed by [SEP]. The output of the SPP task is a binary classification that determines whether the two payloads are arranged in the order of the original text. If so, output IsPayload, if not, output NotPayload.

[0097]

[0098] Fine-tuning allows the model to serve downstream classification tasks well. First, fine-tuning can utilize the high-level features that the pre-trained model has learned. Apply any category of traffic representation; second, the input of the pre-trained model is at the datagram byte level, so the downstream task of classifying packets and flows is converted into corresponding datagram byte tokens and classified by the model; finally, through fine-tuning, the features learned by the pre-trained model are used to explain the model's classification decision on the input data, enhancing the interpretability of the model and making it easier to understand and trust. The special [CLS] token output by the pre-trained model models the representation of the entire input traffic and is directly used for classification.

[0099] Beneficial effects: Reduced reliance on labeled data: By combining pre-training and fine-tuning, the need for a large amount of labeled data is reduced, and the trainability and generalization ability of the model are improved. Improved adaptability: The MTP method can classify traffic in complex scenarios, adapt to different task requirements and network environments, and has strong adaptability and flexibility. Reduced computing power requirements: Compared with a single model, the MTP method can reduce computing power requirements and improve the efficiency and performance of the model by sharing the underlying representation and multi-task learning strategy.

[0100] The above specific implementation methods further illustrate the objectives, technical solutions and beneficial effects of the present invention in detail. It should be understood that the above are only specific implementation methods of the present invention and are not intended to limit the scope of protection of the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principles of the present invention should be included in the scope of protection of the present invention.

Claims

1. A multi-task learning network traffic classification method, characterized in that: The classification method includes: Step S1: collecting a flow data set and preprocessing the flow data set; Step S2: constructing a network architecture; Step S3: training a general pre-training model according to the network architecture; Step S4: fine-tune the pre-trained model.

2. A multi-task learning network traffic classification method according to claim 1, characterized in that: The step S1: collecting a traffic data set and preprocessing the traffic data set specifically includes: According to the design of the payload bytes, it is revealed in the initial packet of the data flow in the real network environment; Traffic in multiple session flows is captured according to IP addresses, port numbers, and protocol types, and time-adjacent payload data is obtained from each session flow as an input unit.

3. A multi-task learning network traffic classification method according to claim 1, characterized in that: The traffic data set uses the ISCXVPN-nonVPN real data set. Design a trace sequence for a capture file in It is a session flow consisting of packets from source to destination; Identified by a five-tuple (IPsrc:PORTsrc,IPdst:PORTdst,Protocol), S-payload is defined as: Where m and n represent the maximum number of one-way packets from source to destination and from destination to source respectively.

4. A multi-task learning network traffic classification method according to claim 1, characterized in that: The step S2: constructing the network architecture specifically includes: A multi-layer bidirectional Transformer is used, each layer consists of a multi-head self-attention mechanism to share information when performing traffic feature classification and application identification tasks.

5. A multi-task learning network traffic classification method according to claim 1, characterized in that: The step S2: constructing the network architecture specifically includes: Use the self-attention mechanism to reconstruct the representation of the target word based on the context S-payload to capture the subtle differences in the data stream; The self-attention mechanism calculates each input x through a similarity function i With all his inputs x1-x k compatibility between them, and calculate x i The weighted sum of , the self-attention mechanism is calculated as follows: x i and x j The similarity between them is represented by a scaled dot product operator, where d k It is k j The dimension of , Z is the normalization factor; Feature extraction on the data stream will be performed for each x i Perform M linear projections to generate multiple attention vectors [att i,1 ,att i,2 ,...,att i,M ]; Use a concatenation operator to get the final attention vector: Multi-head self-attention can better capture the subtle differences in data. Calculate the similarity scores between multiple datagrams to help the model understand and integrate the information between the previous and next datagrams of the entire input sequence; By processing information in parallel, it can capture features in multiple dimensions.

6. A multi-task learning network traffic classification method according to claim 1, characterized in that: The step S3: training a general pre-training model according to the network architecture specifically includes: Special tokens [CLS], [SEP], [PAD], and [MASK] are introduced into the token sequence for classification tasks, pre-training tasks, insufficient padding sequences, and token masks; Self-supervised learning is used on unlabeled network traffic data to train a general pre-trained model through the network architecture.

7. A multi-task learning network traffic classification method according to claim 1, characterized in that: The step S3: training a general pre-training model according to the network architecture further comprises: The two pre-training tasks capture the contextual relationship between traffic bytes by predicting mask tokens, and the correct transmission order by predicting S-payload; When dealing with non-semantic data such as network traffic, the dynamic masking task forces the model to predict the masked data parts, and the model learns specific patterns in the packet sequence; During pre-training, at the beginning of each epoch, the model randomly selects tokens to mask; Each marker has a 15% chance of being selected; The model attempts to predict the original value of the masked token by replacing it with [MASK] with 80% probability, or choosing a random token to replace it with 10% probability; In an epoch, the above process is repeated, and different markers are masked each time; We use negative log-likelihood as the loss function and formally define it as: Among them, θ represents the set of trainable parameters in the model, k is the total number of tags in the sequence, and MASK i is the i-th mask token, S-Payload{token}_i is the actual token at the corresponding position, X is the given before and after payload, and theta is the parameter of the model; The formula evaluates the overall prediction performance of the model by summing the logarithm of the predicted probability of each mask position; The input is Payload, separated by [SEP]. The first Payload is preceded by [CLS], and the second one is also followed by [SEP]. The output of the SPP task is a binary classification, which determines whether the two payloads are arranged in the order of the original text; If yes, output IsPayload, otherwise, output NotPayload; 8. A multi-task learning network traffic classification method according to claim 1, characterized in that: The step S4: fine-tuning the pre-trained model specifically includes: In the fine-tuning stage, the model parameters of the pre-trained model trained on task-specific data are further adjusted using labeled traffic data to capture task-related features.

Images