A malicious traffic classification method based on an improved federated learning algorithm

Abstract

The present invention belongs to the field of machine learning, and particularly relates to a malicious traffic classification method based on an improved federated learning algorithm. Aiming at the "data island" problem existing in the current malicious traffic classification methods and the data heterogeneity problem existing in the original federated learning algorithm, the present invention designs and implements a malicious traffic classification method based on an improved federated learning algorithm. Compared with the traditional classification methods, the present invention does not require centralized training of data, but only needs to aggregate the model parameters trained by each participant using local data through a central server for collaborative training, thus maintaining the data security of network security manufacturers. Compared with the original federated learning method, the present invention divides the participants with approximate data distributions into a cluster, so that the data distributions in each cluster are independently and identically distributed, and then the method of serial training for each cluster enables the federated learning algorithm to also show good effects on non-independent and identically distributed data. Experiments prove that the present invention can achieve good effects in malicious traffic classification.

Description

Technical Field

[0001] The present invention belongs to the field of machine learning, and particularly relates to a malicious traffic classification method based on an improved federated learning algorithm. Background Art

[0002] The rapid development of the Internet has brought great changes to people's lives. The number of Internet users has shown an explosive growth, and the average Internet access time of Internet users has also shown a continuous growth trend. At the same time, the Internet environment has suffered more and more frequent network attacks, and the attack methods are also emerging in an endless stream. These network attacks have brought great potential hazards to people's information property security and network space security. How to effectively prevent the harm caused by network security problems has become increasingly important. Among them, classifying malicious traffic is an important step in effectively preventing network security problems.

[0003] Currently, each network security vendor and relevant national agencies have their own malicious traffic classification systems. However, due to the increasing emphasis on data privacy protection by people, these organizations can only analyze based on their own data, which forms the situation of "data islands", that is, government departments, scientific research departments, and Internet companies cannot share data due to privacy protection. For example, 360 Company can only build its own malicious traffic classification system based on its own network security data. Due to privacy protection, it cannot obtain the network security data of network security vendors such as Tencent and Great Wall, which results in the malicious traffic classification model trained by it being not comprehensive enough and having many deficiencies.

[0004] Federated learning is a distributed machine learning with privacy protection functions. Under the coordination of a central server, multiple participants collaborate to solve the problem of machine learning model training, and the original private data of each participant is stored locally, which well solves the above limitations of the current malicious traffic classification technology. However, the existing federated learning algorithms have data heterogeneity, and when the data is non-independent and identically distributed, federated learning cannot achieve good results, while the data in the real world is all non-independent and identically distributed. Summary of the Invention

[0005] The object of the present invention is to address the above problems and propose a malicious traffic classification method that shares network traffic data under the condition of privacy protection, and fully utilizes the maximum value of data by improving the original federated learning algorithm.

[0006] To achieve the objectives of the present invention, the following technical solutions will be adopted: To enable network security vendors to collaborate in training, we introduce the federated learning algorithm into the malicious traffic classification model. Non-independent and identically distributed data causes the local model parameters in federated learning to diverge, making it impossible for the central server to aggregate a good model; however, for independent and identically distributed data, federated learning can obtain satisfactory results. Accordingly, we improve the original federated learning algorithm, divide the participants according to the data distribution of each participant, and group the participants with high data distribution similarity into one cluster; thus, the data distribution in this cluster is similar to independent and identically distributed data. Then, to traverse all participants, mimicking the training method in centralized data, each cluster is regarded as a batch-size for serial training. Thus, the adverse effects of non-independent and identically distributed data in federated learning are eliminated, and the accuracy of malicious traffic classification is improved. This method includes the following steps:

[0007] S1. Each network security vendor calculates its own data attribute I and sends it to the central server.

[0008] S2. The central server uses a clustering algorithm to aggregate all network security vendors into N clusters S N in.

[0009] S3. The central server initializes the global model parameters w0.

[0010] S4. The central server sends the initialized global model parameters w0 to the participants in the first cluster S1.

[0011] S5. The participants in the first cluster S1 use w0 as the starting point of the model, train the model using local private network security data, and upload the trained local model parameters to the central server.

[0012] S6. After the central server receives the local model parameters sent by the clients in the first cluster S1 , it performs weighted aggregation on these parameters to obtain the global model parameters w1. Then, it sends the global model parameters w1 to the clients in the second cluster S2.

[0013] S7. Repeat S5 - S6 until the global model parameters w N are obtained, which constitutes one round of communication. Then, the global model parameters w N are sent to the participants in the first cluster S1 for a new round of communication.

[0014] S8. Repeat S5 - S7 until the predetermined number of communication rounds is reached or the global model converges. Then, use the trained global model for malicious traffic classification.

[0015] The beneficial effects of the present invention are as follows:

[0016] (1) By combining federated learning with malicious traffic classification technology, the "data island" problem caused by the inability of network security vendors to share data due to data privacy protection is solved. Under the condition that the private data is stored locally, cooperation training can be carried out with other network security vendors to improve the accuracy of malicious traffic classification;

[0017] (2) The federated learning algorithm is improved to solve the problem that the training effect of the federated learning algorithm is not ideal when the data distribution shows non-independent and identically distributed, and further improve the accuracy of malicious traffic classification;

[0018] (3) The improved federated learning algorithm provided by the present invention has a wide range of applications and has excellent effects on different data distributions and different numbers of participants. BRIEF DESCRIPTION OF THE DRAWINGS

[0019] Figure 1 It is a block diagram of the malicious traffic classification method based on the improved federated learning algorithm of the present invention.

[0020] Figure 2 It is a detailed block diagram of the improved federated learning algorithm of the present invention.

[0021] Figure 3 It is a comparison diagram of the training processes of the improved federated learning algorithm and the original federated learning.

[0022] Figure 4 It is the confusion matrix after training malicious traffic using the improved federated learning algorithm.

[0023] Figure 5 It is the confusion matrix after training malicious traffic using the original federated learning algorithm. DETAILED DESCRIPTION OF THE INVENTION

[0024] The technical solution of the present invention will be further described below with reference to the accompanying drawings.

[0025] As Figure 1 and Figure 2 shown, the present invention includes the following steps:

[0026] Step 1: Each network security vendor calculates its own data attribute I based on the local private data and sends it to the central server.

[0027] Step 2: The central server uses a clustering algorithm to aggregate all network security vendors into N clusters S N .

[0028] If the data distribution is independent and identically distributed, then after local training, the local model parameters of each network security vendor are very similar. At this time, aggregating the model parameters is beneficial to improving the performance of the global model and enabling the global model to have stronger generalization ability for unknown data. However, if the data is non-independent and identically distributed, then the data of each network security vendor is very different, and the knowledge learned by each local model after training with local data is also very different, and the model parameters are also very different. If the local model parameters are simply averaged at this time, then the knowledge learned by each local model will be lost, resulting in a deterioration of the global model effect. Therefore, the present invention divides each network security vendor according to the data distribution, and divides the network security vendors with high data distribution similarity into a cluster; in this way, the data distribution in the cluster is similar to independent and identically distributed. At this time, performing federated optimization within each cluster can greatly improve the performance of the model within each cluster and improve efficiency.

[0029] Step 3: The central server initializes the global model parameter w0.

[0030] Step 4: The central server sends the initialized model parameter w0 to the participants in the first cluster S1.

[0031] Step 5: The clients in the first cluster S1 use w0 as the starting point of the model and perform model training using the local dataset, and upload the trained model parameters to the central server.

[0032] Step 6: After the central server receives the model parameters sent by the participants in the first cluster S1 it performs weighted aggregation on these parameters according to the data sample size of the participants to obtain the global model parameter w1. Then it sends the global model parameter w1 to the clients in the second cluster S2. Among them

[0033] Step 7: Repeat Step 5 to Step 6 until the global model parameter w N is obtained, which constitutes one round of communication. Then the global model parameter w N is sent to the participants in the first cluster S1 for a new round of communication.

[0034] There is a significant gap between federated learning and centralized machine learning in terms of the method of updating model parameters. The present invention reduces the impact of data imbalance on federated learning by changing the training process of federated learning to make it closer to the data training process of centralized machine learning. The present invention regards the global model obtained by the central server aggregating local models in the cluster as the model generated by centralized machine learning from batch data. Then, the central server transmits the aggregated global model to another cluster for local training and then aggregates the parameters to obtain a new global model, which corresponds to centralized machine learning using another batch of data to continue training the model.

[0035] Step 8, repeat S5 - S7 until the predetermined number of communication rounds is reached or the global model converges. Then use the trained global model for malicious traffic classification.

[0036] In step 1, the data attributes are defined as: n is the total number of samples of a certain participant, n i is the number of samples of the i-th type of data it owns.

[0037] In step 2, the clustering algorithm is as follows:

[0038] (1) Consider each network security vendor as a separate cluster.

[0039] (2) Calculate the distance between two clusters and merge the two clusters with the minimum distance. The distance between two clusters is defined as

[0040] D p,q = max(d ij = ||I i - I j ||2 | i ∈ p, j ∈ q) (1)

[0041] where d ij represents the distance between network security vendor i in cluster p and network security vendor j in cluster q.

[0042] (3) Repeat step (2) until the number of clusters meets the requirements.

[0043] In step 6, under the coordination of the central server, all vendors participate in federated learning collaborative training. Assume that there are C network security vendors in a cluster as participants in federated learning, and each vendor has its own private local dataset is the feature vector, is the corresponding label, n c represents the size of the dataset of the c-th user participating in federated learning. For the learning objective of non-convex neural network models, such as the formula:

[0044]

[0045] In a cluster, c manufacturers each have their local private dataset |D i | = n c , where n represents the size of all datasets in the cluster. Therefore So Equation 2 can be expressed as:

[0046]

[0047] When each manufacturer in a cluster participates in federated learning collaborative training, starting from the global model parameters sent by the central server, with a fixed learning rate η, calculate the average loss gradient on its own private dataset

[0048] These manufacturers synchronously update their local models and send the updated local model parameters to the central server for aggregation to update the global model parameters:

[0049]

[0050]

[0051]

[0052]

[0053] For each manufacturer c, Combined with Equation (7), we get:

[0054]

[0055] More samples can train a better model. Therefore, by allocating weights according to the amount of private samples owned by each manufacturer, a better global model can be formed.

[0056] Next, combined with simulation examples and Figure 3 、 Figure 4 and Figure 5 , prove the effectiveness of the present invention:

[0057] Use the CICIDS2017 network security dataset as the experimental dataset. To simulate the non-independent and identically distributed data distribution form in reality, the dataset is allocated to 60 manufacturers, and each manufacturer only has one type of malicious traffic. The network used is a multi-layer perceptron with two hidden layers. The first layer has 100 units and the second layer has 50 units.

[0058] From Figure 3It can be seen that although the accuracy of the improved federated learning algorithm for classifying malicious traffic is lower than that of the original federated learning algorithm at the beginning, as the number of communications increases, the accuracy of the improved federated learning algorithm exceeds that of the original federated learning algorithm and gradually increases. In the improved federated learning algorithm, the data of network security vendors in the same cluster are almost independently and identically distributed, and the parameters of each local model are not very different, so a global model with better performance can be aggregated. In the original federated learning algorithm, the local models of different network security vendors vary greatly, and the performance of the aggregated model is also poor. From Figure 4 and Figure 5 It can be seen that the recognition rate of the improved federated learning algorithm for each type of malicious traffic is close to 100%, while the recognition rate of the original federated learning algorithm reaches 100% only for several types of malicious traffic. This shows that the improved federated learning algorithm can fully learn the characteristics of various types of malicious traffic, so as to better classify. Therefore, each participating network security vendor can obtain a better malicious traffic recognition model by using the improved federated learning algorithm, which fully verifies the feasibility of the present invention in the field of malicious traffic classification.

Claims

1. A malicious traffic classification method based on an improved federated learning algorithm, characterized in that, It includes the following steps: S1. Each network security vendor participating in federated learning calculates its own data attribute I and sends it to the central server; where the data attribute is defined as: n is the number of samples owned by the participant, n i is the number of samples of the i-th type of data owned; S2. The central server uses a clustering algorithm to aggregate all network security vendors into N clusters S N where the clustering algorithm is as follows: (1) Treat each network security vendor as a separate cluster; (2) Calculate the distance between two clusters and merge the two clusters with the minimum distance. The distance between two clusters is defined as D p,q = max(d ij = ||I i - I j ||2 | i ∈ p, j ∈ q) (1) where d ij represents the distance between cybersecurity vendor i in cluster p and cybersecurity vendor j in cluster q; (3) Repeat step (2) until the number of clusters meets the set requirements; S3. The central server initializes the global model parameter w0; S4. The central server sends the initialized global model parameter w0 to the vendors in the first cluster S1; In S5, the participants in the first cluster S1 use w0 as the model starting point, train the model using their local data, and upload the trained local model parameters to the central server; S6. After the central server receives the local model parameters sent by the clients in the first cluster S1 it performs weighted aggregation on these parameters to obtain the global model parameter w1; then it sends the global model parameter w1 to the clients in the second cluster S2, where S7. Repeat S5 - S6 until the global model parameter w is obtained N ; This constitutes one round of communication. Then the central server transmits the global model parameter w N to the manufacturers in the first cluster S1 for a new round of communication; S8. Repeat S5 - S7 until the predetermined number of communication rounds is reached or the global model converges, and then use the obtained global model for malicious traffic classification.

Images