A malicious network traffic classification method based on sample enhancement and contrastive learning
By using sample enhancement and contrastive learning methods, a shallow neural network model was constructed, which solved the problem of insufficient network traffic samples, achieved efficient classification of new malicious traffic, and improved the model's classification performance and resource utilization efficiency.
Patent Information
- Application Number
- CN202311005429.5
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2023-08-10
- Publication Date
- 2025-09-26
- Estimated Expiration
- 2043-08-10
AI Technical Summary
When the number of network traffic samples is insufficient, existing machine learning models find it difficult to effectively classify new malicious traffic, resulting in poor classification results.
A method based on sample enhancement and contrastive learning is adopted. By building a shallow neural network model, enhancing data is constructed by randomly masking network traffic features, and training the model through a contrastive learning algorithm to improve the classification performance of the model.
In the case of small samples, it achieves efficient classification of new malicious network traffic, with low resource consumption and good classification performance, and can distinguish between known and unknown network traffic categories.
Smart Images
Figure CN117034112B_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to network security and deep learning technology, and in particular to a malicious network traffic classification method based on sample enhancement and contrastive learning. Background Art
[0002] With the increasing ubiquity of the internet in modern life and the increasing number of devices connected through networks, cybersecurity has garnered increasing attention. Traffic intrusion classification systems are used to effectively categorize various malicious attacks on networks and are a key component of maintaining cybersecurity. From a machine learning perspective, intrusion classification systems can be defined as systems that classify network traffic. Simply put, they distinguish between legitimate and malicious network traffic. With the advancement of machine learning technology, machine learning-based methods for classifying malicious network traffic have gained widespread attention.
[0003] Machine learning methods rely on sufficient training samples for intrusion classification systems to accurately classify test samples. However, today's network environment is constantly changing, with new types of malicious traffic constantly emerging. This makes it difficult to obtain sufficient samples to train models in a short period of time. This insufficient number of samples hinders the adequate training of machine learning models, which in turn affects the effectiveness of malicious traffic classification.
[0004] In response to the above problems, how to use deep learning to classify new types of malicious network traffic when the collected network traffic samples are small is an issue that needs to be solved urgently. Summary of the Invention
[0005] The purpose of this invention is to solve the problem that insufficient number of samples makes it difficult for machine learning models to be fully trained, which in turn affects the effectiveness of malicious traffic classification, and proposes a malicious network traffic classification method based on sample enhancement and contrastive learning.
[0006] To achieve the above objectives, the malicious network traffic classification method based on sample enhancement and contrastive learning includes:
[0007] Acquire N types of network traffic, each type of network traffic including multiple samples, extract multiple network traffic features of each sample to form a feature vector corresponding to the sample, and construct a network traffic sample set, wherein the network traffic sample set includes network traffic and corresponding network traffic features;
[0008] Preprocess the network traffic features to obtain an enhanced set for each sample;
[0009] Take the union of all enhanced sets of samples contained in each type of network traffic. For each sample in each union, traverse other samples in the same union to form a positive sample pair. Take samples from another union to traverse to form a negative sample pair. Set a label for each sample pair and use the resulting set of positive and negative sample pairs as the pre-training set.
[0010] Build a basic model and use the normalized feature vectors of the two samples in the pre-training set as input to the basic model to obtain the two processed feature vectors of the sample pair;
[0011] Calculate the similarity of the two processed feature vectors of the sample pair, identify the label of the sample pair to obtain the label judgment value, and calculate the loss value. Based on the loss value, calculate and update the parameters of the basic model to obtain the comparative learning malicious traffic classification model;
[0012] A contrastive learning malicious traffic classification model is used to classify malicious traffic in new tasks.
[0013] Furthermore, the network traffic features are preprocessed to obtain an enhanced set for each sample, including:
[0014] Train n network traffic classification models on the network traffic sample set, input the feature vector of the j-th type of network traffic into the i-th network traffic classification model, and obtain the classification accuracy acc(i,j);
[0015] Input the dth network traffic feature into the i-th network traffic classification model to obtain the importance weight weight(i,d), where the dth network traffic feature belongs to the j-th category of network traffic;
[0016] The importance weight I(j,d) is calculated based on the classification accuracy acc(i,j) and the importance weight weight(i,d), which can be expressed as follows:
[0017]
[0018] Where I(j,d) represents the importance weight of the d-th network traffic feature to the j-th type of network traffic;
[0019] The ath sample s in the network traffic sample set a The network traffic characteristics are represented as feature vector x a , the sample s a Belonging to the jth type of network traffic, calculate the masking probability of each network traffic feature in the jth type of network traffic, which can be expressed as follows:
[0020]
[0021] Where P(j,d) represents the masking probability of the dth network traffic feature of all samples of the jth type of network traffic;
[0022] According to the masking probability in the feature vector x a Randomly select L network traffic features to mask, and obtain sample s after masking a An enhanced sample, repeat the random selection and masking operation m times to obtain m enhanced samples, and transform the feature vector x a and a set of m enhanced samples as sample s a Enhancement set.
[0023] Furthermore, for all positive sample pairs formed by samples in the same union, there is at least one sample that is not repeated between any two positive sample pairs.
[0024] Furthermore, setting a label for each sample pair includes:
[0025] The label of the positive sample pair is set to 1, and the label of the negative sample pair is set to 0.
[0026] Furthermore, a basic model is constructed, and the feature vectors of the two samples in the sample pair in the pre-training set are normalized and used as the input of the basic model to obtain two processed feature vectors of the sample pair, including:
[0027] The basic model is constructed using a multi-layer perceptron, and the feature vector x of the normalized sample pair is p and the eigenvector x q As the input of the basic model, the basic model updates the parameters of each layer of the multilayer perceptron based on the following formula:
[0028] X (l+1) =σ(A (l) X (l) +b (l) )
[0029] Among them, A (l) is the trainable parameter matrix of the lth layer of the multilayer perceptron, b (l) is the parameter vector of the lth layer of the multilayer perceptron, X (l) is the output of the lth layer of the multilayer perceptron, X (l+1) is the output of the l+1th layer of the multilayer perceptron, σ(…) is the activation function;
[0030] Get the eigenvector x p With the eigenvector x q The feature vector after multi-layer perceptron processing and
[0031] Furthermore, the calculating of the similarity between the two processed feature vectors of the sample pair includes:
[0032] The cosine similarity function is used to compare the two processed feature vectors of the sample pair, which is expressed as follows:
[0033]
[0034] in, Represents the similarity value of two feature vectors, ranging from [-1, 1], represents the eigenvector x p The feature vector after multi-layer perceptron processing, represents the eigenvector x q Feature vector after processing by multi-layer perceptron;
[0035] The similarity value The range is scaled to [0,1] and is expressed as:
[0036]
[0037] in, Indicates the scaling value of the similarity value.
[0038] Furthermore, the calculation of the loss value, and the calculation and updating of the parameters of the basic model according to the loss value, include:
[0039] The loss value is calculated by the binary cross entropy loss function, which is expressed as follows:
[0040]
[0041] Among them, Loss_L2(x p ,x q ) represents the eigenvector x p and the eigenvector x q The loss value, α represents the regularization factor, W represents the sum of all weights of the basic model, and y represents the label judgment value;
[0042] Backpropagation is used to calculate the gradient, and gradient descent is used to calculate and update the parameters of the neural network.
[0043] Furthermore, the use of the contrastive learning malicious traffic classification model to classify malicious traffic in the new task includes:
[0044] Obtain the trained comparative learning malicious traffic classification model;
[0045] Collecting training samples in a new task, wherein at least one malicious traffic type in the training samples does not overlap with a malicious traffic type in a network traffic sample set;
[0046] The contrastive learning malicious traffic classification model is used to classify malicious traffic in the new task. The training samples are compared with the samples to be tested, and the label with the highest output probability is selected as the predicted label.
[0047] Compared with the existing technology, the present invention has the following significant advantages: 1. It adopts a basic model with a shallow neural network architecture that is light on resources, occupies few resources and has high operating efficiency. 2. It adopts a heuristic method of randomly masking network traffic features to construct a comparison task. This method can perform data enhancement and effectively retain the structure of the feature vector before masking, ensuring the effectiveness of the enhanced data and improving the training effect of the basic model. 3. Based on the architecture of contrastive learning, the contrastive learning malicious traffic classification model has good classification performance to distinguish known network traffic categories and can classify unknown small sample malicious network traffic. BRIEF DESCRIPTION OF THE DRAWINGS
[0048] Figure 1 This is a flow chart of a method for classifying small sample malicious network traffic based on sample enhancement and contrastive learning according to the present invention;
[0049] Figure 2 Flowchart for constructing a training set for network traffic of the present invention. DETAILED DESCRIPTION
[0050] In order to make the purpose, technical solutions and advantages of the present invention more clearly understood, the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention.
[0051] The present invention provides a small sample malicious network traffic classification method based on sample enhancement and contrastive learning, such as Figure 1 As shown, the following steps are included:
[0052] (1) Capture network traffic data and extract network traffic features from it. Finally, enhance the network traffic features and construct a pre-training set.
[0053] (1-1) Network traffic feature extraction: First, use a network traffic packet capture tool (such as tcpdump or wireshark) to capture and generate a network traffic file (such as a pcap file). Then, use a network traffic analysis tool (such as CICFlowMeter) to parse and perform statistical analysis on the network traffic file to generate network traffic features. Multiple network traffic features are extracted for each network traffic file (i.e., sample). If the network traffic feature types extracted for each network traffic file are the same, the number of extracted network traffic features is the same. If no network traffic feature is extracted for a certain network traffic feature type, the network traffic feature is set to 0. The number of network traffic features is denoted as D. The network traffic sample set S is obtained (containing N types of network traffic and corresponding network traffic features, where N types of network traffic are normal network traffic and N-1 types of malicious network traffic).
[0054] (1-2) Figure 2 As shown in the figure, the network traffic features are enhanced and a pre-training set is constructed. The specific steps are as follows:
[0055] (1-2-1) Calculate the importance weight of network traffic features: First, select n machine learning algorithms and train n network traffic classification models C1, C2, ..., C on S. n , C i The classification accuracy of the j-th type of network traffic is recorded as acc(i,j), where i∈[1,n], j∈[1,N]. i The importance weight of the calculated d-th network traffic feature is recorded as weight(i,d), and the d-th network traffic feature belongs to the j-th type of network traffic. Then, the importance weight I(j,d) of the d-th network traffic feature to the j-th type of network traffic is calculated according to formula (1).
[0056]
[0057] (1-2-2) Sample enhancement: First, for each sample s of each type of network traffic a , and its eigenvector is recorded as x a , assuming s a Belonging to the jth type of network traffic, the masking probability of each network traffic feature d is calculated according to formula (2). Then, according to the masking probability, a L network traffic features are randomly selected for masking. The default value of L is 10, which means the value of the selected network traffic feature is set to 0. The masked feature vector is s a An enhanced sample of . Repeat the above random selection and masking operation m times, then each network traffic sample s a Generate m enhanced samples and transform the feature vector x a The set of its m augmented samples is called sa The enhanced set is denoted as AS a .
[0058]
[0059] (1-2-3) Construct a pre-training set: 1) For one type of network traffic, take the union of the enhanced sets of all samples it contains, and record the union of the enhanced sets of the j-th type of network traffic as CAS j 2) CAS j A sample of jb , take CAS j Any other sample s jk , and s jb Form a positive sample pair (s jb ,s jk ), the label is set to 1; any union of the enhanced sets of other network traffic types CAS o Any sample s in (o≠j) ok , and s jb Form a negative sample pair (s jb ,s ok ), the label is set to 0. 3) Repeat step 2), for CAS j For each sample in , traverse CAS j The other samples in form positive sample pairs, traversing CAS o All samples in the set form negative sample pairs. 4) Repeat step 1) until all samples in the union of the enhanced sets of all network traffic types have been traversed. For the generation of positive sample pairs of samples in the same union, for positive sample pairs that are repeated with both samples in the existing positive sample pairs, they can be directly abandoned during traversal or deleted after traversal (for example, positive sample pairs (s j1 ,s j2 ) and (s j2 ,s j1 ) Both are duplicate sample pairs, so only one can be retained. The final set of positive and negative sample pairs is used as the pre-training set, denoted as PNS.
[0060] (2) Model training based on contrastive learning: A shallow neural network is used to build a basic model, following the contrastive learning method, and using regularization and other techniques to pre-train the model to obtain a contrastive learning model.
[0061] (2-1) Definition of neural network basic model: Multi-layer perceptron (MLP) is used to build the basic model. The input of the basic model is the normalized network traffic feature vector x of a positive and negative sample pair. p with x q , the basic model updates the parameters of each layer of network traffic based on formula (3), and finally obtains xp with x q Feature vector after processing by multi-layer perceptron.
[0062] X (l+1) =σ(A (l) X (l) +b (l) ) (3)
[0063] Among them, A (l) is the trainable parameter matrix of the lth layer of the multilayer perceptron, b (l) is the parameter vector of the lth layer of the multilayer perceptron, X (l) is the output of the lth layer of the multilayer perceptron, X (l+1) is the output of the l+1th layer of the multilayer perceptron, and σ(…) is the activation function.
[0064] (2-2) Use the cosine similarity function to compare the feature vectors of the two network flows, see formula (4).
[0065]
[0066] in, Represents the similarity value of two feature vectors, ranging from [-1, 1], Represents x p The feature vector after multi-layer perceptron processing, Represents x q Feature vector after processing by multi-layer perceptron.
[0067] (2-3) Then use formula (5) to convert The range of values is scaled to [0,1].
[0068]
[0069] in, Indicates the scaling value of the similarity value.
[0070] After obtaining the scaling value (2-4), the label of the input sample pair is identified to obtain the label determination value y. In this embodiment, the label determination value y is the label of the sample pair, 0 or 1. The loss value is calculated using the binary cross entropy loss function and L2 regularization is added, see formula (6). After obtaining the loss value, the gradient is calculated using backpropagation, and the parameters of the multilayer perceptron are calculated and updated using gradient descent.
[0071]
[0072] Among them, y represents the label judgment value, Loss_L2(x p ,x q) represents the loss value, α represents the regularization factor, and W represents the sum of all weights of the basic model.
[0073] (2-5) After training a large number of positive sample pairs and negative sample pairs, a comparative learning malicious traffic classification model is obtained.
[0074] (3) Classification of small samples of network traffic: For the target task, a contrastive learning malicious traffic classification model is used to perform small sample classification of malicious traffic in the small sample target task.
[0075] (3-1) Model initialization: Obtain the trained contrastive learning malicious traffic classification model.
[0076] (3-2) Data Input: Collect training samples for the new task (including a small number of training samples of new malicious traffic types, where at least one malicious traffic type in the training samples does not overlap with the malicious traffic types in the network traffic sample set). The logic for collecting training samples for the new task is the same as in steps (1-1) to (1-2-3).
[0077] (3-3) Model implementation: A contrastive learning malicious traffic classification model is used to classify malicious traffic in the new task. The training samples and the samples to be tested are contrasted and learned, and the label with the highest output probability is selected as the predicted label.
[0078] The small sample training process uses the contrastive learning malicious traffic classification model to perform malicious traffic classification in the new task. The logic is the same as that of steps (2-1) to (2-4), and the network traffic sample set is replaced by the training sample.
[0079] In the small sample classification process, when performing comparative learning on the training samples and the samples to be detected, samples in the training samples and samples in the samples to be detected are taken to form sample pairs, and the feature vectors of the sample pairs are input into the comparative learning malicious traffic classification model. The similarity of the two processed feature vectors of the sample pairs is calculated, and the label (type) of the training sample in the sample pair with the highest similarity (highest output probability) is taken as the type of the sample to be detected.
[0080] The above-described embodiments merely represent one or several implementations of the present invention. While the descriptions are relatively specific and detailed, they should not be construed as limiting the scope of the patent. It should be noted that a person skilled in the art would be able to make various modifications and improvements without departing from the spirit of the present invention, and these modifications and improvements fall within the scope of protection of the present invention. Therefore, the scope of protection of the patent for this invention shall be determined by the appended claims.
Claims
1. A malicious network traffic classification method based on sample enhancement and contrastive learning, characterized by: The malicious network traffic classification method based on sample enhancement and contrastive learning includes: Acquire N types of network traffic, each type of network traffic including multiple samples, extract multiple network traffic features of each sample to form a feature vector corresponding to the sample, and construct a network traffic sample set, wherein the network traffic sample set includes network traffic and corresponding network traffic features; Preprocess the network traffic features to obtain an enhanced set for each sample; Take the union of all enhanced sets of samples contained in each type of network traffic. For each sample in each union, traverse other samples in the same union to form a positive sample pair. Take samples from another union to traverse to form a negative sample pair. Set a label for each sample pair and use the resulting set of positive and negative sample pairs as the pre-training set. Build a basic model and use the normalized feature vectors of the two samples in the pre-training set as input to the basic model to obtain the two processed feature vectors of the sample pair; Calculate the similarity of the two processed feature vectors of the sample pair, identify the label of the sample pair to obtain the label judgment value, and calculate the loss value. Based on the loss value, calculate and update the parameters of the basic model to obtain the comparative learning malicious traffic classification model; Use contrastive learning malicious traffic classification model to classify malicious traffic in new tasks; in, The network traffic features are preprocessed to obtain an enhanced set for each sample, including: Train n network traffic classification models on the network traffic sample set, input the feature vector of the j-th type of network traffic into the i-th network traffic classification model, and obtain the classification accuracy acc(i,j); Input the dth network traffic feature into the i-th network traffic classification model to obtain the importance weight weight(i,d), where the dth network traffic feature belongs to the j-th category of network traffic; The importance weight I(j,d) is calculated based on the classification accuracy acc(i,j) and the importance weight weight(i,d), which can be expressed as follows: Where I(j,d) represents the importance weight of the d-th network traffic feature to the j-th type of network traffic; The ath sample s in the network traffic sample set a The network traffic characteristics are represented as feature vector x a , the sample s a Belonging to the jth type of network traffic, calculate the masking probability of each network traffic feature in the jth type of network traffic, which can be expressed as follows: Where P(j,d) represents the masking probability of the dth network traffic feature of all samples of the jth type of network traffic; According to the masking probability in the feature vector x a Randomly select L network traffic features to mask, and obtain sample s after masking a An enhanced sample, repeat the random selection and masking operation m times to obtain m enhanced samples, and transform the feature vector x a and a set of m enhanced samples as sample s a Enhancement set.
2. The malicious network traffic classification method based on sample enhancement and contrastive learning according to claim 1 is characterized in that: For all positive sample pairs formed by samples in the same union, there is at least one sample that is not repeated between any two positive sample pairs.
3. The malicious network traffic classification method based on sample enhancement and contrastive learning according to claim 1 is characterized in that: Setting a label for each sample pair includes: The label of the positive sample pair is set to 1, and the label of the negative sample pair is set to 0.
4. The malicious network traffic classification method based on sample enhancement and contrastive learning according to claim 1 is characterized in that: Build a basic model and use the normalized feature vectors of the two samples in the pre-training set as the input of the basic model to obtain two processed feature vectors of the sample pair, including: The basic model is constructed using a multi-layer perceptron, and the feature vector x of the normalized sample pair is p and the eigenvector x q As the input of the basic model, the basic model updates the parameters of each layer of the multilayer perceptron based on the following formula: X (l+1) =σ(A (l) X (l) +b (l) ) Among them, A (l) is the trainable parameter matrix of the lth layer of the multilayer perceptron, b (l) is the parameter vector of the lth layer of the multilayer perceptron, X (l) is the output of the lth layer of the multilayer perceptron, X (l+1) is the output of the l+1th layer of the multilayer perceptron, σ(…) is the activation function; Get the eigenvector x p With the eigenvector x q The feature vector after multi-layer perceptron processing and 5. The malicious network traffic classification method based on sample enhancement and contrastive learning according to claim 4 is characterized in that: Calculating the similarity of two processed feature vectors of the sample pair includes: The cosine similarity function is used to compare the two processed feature vectors of the sample pair, which is expressed as follows: in, Represents the similarity value of two feature vectors, ranging from [-1, 1], represents the eigenvector x p The feature vector after processing by the multi-layer perceptron is represents the eigenvector x q Feature vector after processing by multi-layer perceptron; The similarity value The range is scaled to [0,1] and is expressed as: in, Indicates the scaling value of the similarity value.
6. The malicious network traffic classification method based on sample enhancement and contrastive learning according to claim 5 is characterized in that: The calculation of the loss value, and the calculation and updating of the parameters of the basic model according to the loss value, include: The loss value is calculated by the binary cross entropy loss function, which is expressed as follows: Among them, Loss_L2(x p ,x q ) represents the eigenvector x p and the eigenvector x q The loss value, α represents the regularization factor, W represents the sum of all weights of the basic model, and y represents the label judgment value; Backpropagation is used to calculate the gradient, and gradient descent is used to calculate and update the parameters of the neural network.
7. The malicious network traffic classification method based on sample enhancement and contrastive learning according to claim 1 is characterized in that: The use of the contrastive learning malicious traffic classification model to classify malicious traffic in the new task includes: Obtain the trained comparative learning malicious traffic classification model; Collecting training samples in a new task, wherein at least one malicious traffic type in the training samples does not overlap with a malicious traffic type in a network traffic sample set; The contrastive learning malicious traffic classification model is used to classify malicious traffic in the new task. The training samples are compared with the samples to be tested, and the label with the highest output probability is selected as the predicted label.