FedEdge AI-TC: A Semi-supervised Traffic Classification Method based on Trusted Federated Deep Learning for Mobile Edge Computing

Deep learning, also referred to as deep structured learning or hierarchical learning, is achieved by acquiring the representation of data. In contrast to traditional machine learning algorithms, deep learning can automatically extract features without human intervention, rendering it an ideal approach for traffic classification. The application of deep learning techniques to Network TC[8, 9, 10, 11, 12] involves three steps: firstly, characterizing the input data by defining and designing the model input using data packets, PCAP files, or traffic statistics vectors as features; secondly, selecting suitable models and algorithms based on classifier objectives and model characteristics; finally extracting traffic features automatically through training a DL-based classifier and associating input data with corresponding category labels.

Recent research has demonstrated deep learning methods’ superiority in traffic classification. For example, CNNs[13, 14] are widely used in traffic classification. They can automatically extract features from raw network traffic data and have end-to-end learning capabilities during training. In addition, RNNs and LSTMs[15, 16] are used to process traffic sequence data and can capture the temporal dependencies in the data. These types of networks are often used in traffic classification to identify persistent attacks such as DDoS (Distributed Denial of Service) attacks. GNN has proven to be a novel information representation method for DL, which has been applied in TC or IDS[17].

Semi-supervised traffic classification[18, 19, 20, 21] is an approach that utilizes a small set of labeled data along with a substantial amount of unlabeled data to distinguish various network traffic types. There are four primary methods for semi-supervised traffic classification: cluster-based methods, generative models, GANs (Generative Adversarial Networks), and discriminative models. Cluster-based methods[22, 23] have low computational complexity but can be influenced by data distribution and may exhibit instability in practical usage. Generative model-based methods[24] are effective for unknown or dynamically changing network applications; however, they require prior knowledge to select appropriate statistical features and clustering parameters which might limit the generalization of classification results. GAN-based methods[25, 26] can enhance dataset diversity and quality, thereby improving the model’s generalization performance. Nevertheless, GAN models have complex structures and numerous parameters that pose challenges in training and render them impractical for deployment on edge devices. Therefore, this study primarily adopts a discriminative model-based approach by directly learning the mapping function from feature space to class space. Model parameters are optimized by minimizing the classification error of labeled data while incorporating a regularization term for unlabeled data. Subsequently, the unlabeled data is classified based on predicted results.

Federated Learning (FL) is a distributed ML/DL framework that focuses on decentralized training data, aiming to obtain ML/DL models by distributing the data across numerous nodes while ensuring privacy and security. FL allows local clients to retain their data, sharing only the model parameters with a central server, thereby reducing communication overhead and preserving client data privacy[7, 27]. Recently, two main approaches have emerged for traffic classification tasks by combining federated learning with deep learning techniques. The first approach[28, 29, 30, 31] empowers child nodes to annotate the data through various means. The second approach[32, 33, 34, 35, 36] involves transforming the model structure and training objectives so that sub-nodes can train the model using unlabeled data; subsequently, fine-tuning is performed by the server using labeled data to achieve semi-supervised traffic classification. We propose a semi-supervised traffic classification model based on VAE-CNN that incorporates a federated learning paradigm enabling edge devices’ semi-supervised training of the model.

Smart home networks face four major challenges: terminal heterogeneity, application diversity, high privacy, and rapid evolution. A network TC system must continuously learn through long-term iterative optimization to overcome these challenges. The process follows the full life cycle of federated learning, as shown in Fig. 2. The edge-side AI-TC classification system workflow based on federated learning includes initialization, broadcast, training, parameter uploading, model aggregation and evaluation, edge deployment, and model monitoring. Here are the steps for implementing an AI-powered network traffic classification system:

1. 1.

   Initialization: Provide the client node with an initialization model for efficient local/global model training.

2. 2.

   Broadcast: The centralized server broadcasts the initialization model to all the client nodes like 5G CPE/HGU.

3. 3.

   Local Training: The client node performs feature engineering, model construction, and local training.

   1. (a)

      Feature Engineering: Extract, select, represent, and compress network traffic features to build an optimal feature subset for the AI-TC classification system.

   2. (b)

      Model Construction and Local Training: This step involves selecting what kind of learning methods (supervised/semi-supervised/unsupervised/weakly supervised), training methods (centralized/distributed training, federated learning), whether to pre-train, whether to use classical models for transfer learning, to form a local model (Local Model) on the local client node.

4. 4.

   Parameters Uploading: Upload encrypted parameter information obtained from local training to the centralized server.

5. 5.

   Aggregation: The centralized server performs ’secure aggregation’ on the parameter information uploaded by each client node (such as using the FedAvg algorithm) and performs global training.

6. 6.

   Model Evaluation: Evaluate the global model obtained from the centralized server’s global training. If the training process converges, it will enter the model deployment step; otherwise, it will inform the client node to continue training and iterative optimization. In addition, the model evaluation also needs to consider its computational complexity, time complexity, and the computing resources and time required for training/inference.

7. 7.

   Edge deployment: Deploy the model on the edge or terminal side using the pull/push/subscribe model deploy method and update strategy. Model compression and interpretation are two important tasks in this step.

   1. (a)

      Model Compression: Compress the inference/classification model small enough to meet the fast classification under limited computing power.

   2. (b)

      XAI(Model Interpretation/Explanation): Solve the ”black box” problem of the AI-TC model to make the classification model users trust the model.

8. 8.

   Model Monitoring: Monitor the status of the classification system, including model, and real-time network flow, to report some key issues such as classification system failures and model degradation.

9. 9.

   Continuous Learning: Initiate iterative optimizations and continuous learning from initialization to maintain high adaptability, robustness, and reliability of the classification system.

The following sections of this article will focus on initialization, model training, compression, and explanation.

The overall architecture of FedEdge AI-TC is illustrated in Fig. 3, which is divided into the client node and central server/central aggregator based on our previous work[37]. HGU, functioning as the client node, performs local training and inference classification tasks[38]. On the other hand, the central aggregator acts as a centralized server responsible for aggregated training, model evaluation, compression, interpretation, and deployment of inference models. The workflow can be summarized as follows: Initially, the central aggregator broadcasts the initial classification model to HGU. Subsequently, HGU collects real-time packets and performs redundant/invalid packet filtering, network flow attribute calculation, and normalization to formulate a FAM (Flow Attribute Matrix), depicted in Fig.4. This matrix is then utilized for training an initialization model locally. Afterward, encrypted gradient information, loss, and other parameter details are uploaded to the central aggregator for aggregate training. Since each HGU exhibits similar flow characteristics, Horizontal Federated Learning (HFL) is employed for aggregate training by standard secure aggregation algorithms like FedAvg. Then the model performance metrics will be evaluated, including accuracy, precision, recall, and F1-Score for assessing model convergence. If convergence occurs successfully, the aggregated global model undergoes compression to meet computing resource constraints of HGU (including CPU/memory/Flash). Once an available inference model is obtained, XAI-based methods are applied to provide both global and local explanations of the model. Finally, the resulting inference model will be distributed across all HGUs. Otherwise, if convergence fails, the global model will be issued to each HGU, initiating a new epoch of iterative training and optimization until it converges.

The initial model plays a crucial role in determining the convergence speed and performance of federated aggregation training, which serves as the initial stage of federated learning. While randomly initialized models can be used within the FL framework, it is essential to construct an initial model that enhances local/global model training efficiency. Unlike traditional AI domains like images and text, network traffic classification lacks pre-trained models, necessitating the development of the initial model by ourselves. In this study, we utilize three benchmarking datasets including (ISCIX[39], UNSW-NB15[40], and MIRAGE[41]) as baseline datasets for constructing our initial model. We adopt CNN as the supervised training algorithm, and further details about the initial model can be found in our previous work[37].

We conducted the experimental evaluation on two datasets. One is a benchmarking public dataset, ISCX-VPN2016. The other is a private dataset built by ourselves. The latter comprises six popular applications and background traffic from terminals collected in the campus network scenario, including Bilibili, QQ music, Honor of Kings, Teamfight Tactics, and Game for peace. These apps cover the five popular app categories of video, music, Moba games, First Person Shooting (FPS) and Role-playing game (RPG). To collect the data, we used a semi-automatic web traffic generation tool. We leveraged an automated traffic generator to collect traffic from Bilibili and QQ music, with 13,314 flows from Bilibili and 20703 flows from QQ music. For the other interactive games, we chose to collect them manually. We used PCAPDroid to mark the network traffic as it occurred at the endpoint and also at the router, where the two were compared and filtered. The network flow features were calculated by CICFlowMeter for each application’s PCAP files. In addition, this experiment also provides the statistics of the background traffic information, which contains information about network location services, security components, and syslog services. Table IV shows the exact number of flows.

The experimental environment is AMD Ryzen 3600, 16GB RAM, NVIDIA GTX 1660, CUDA 7.5, CDNN10.5. In this paper, Python3 is the primary programming language. The following is a description of the evaluation metrics: Precision, Recall, F1, Accuracy, and AUC.

Time-complexity-related metrics about training like training time are not included in this paper because we think those are highly dependent on the hardware resources.

The training process of the VAE part in the VAE+CNN (E-CNN) model is divided into two parts. Firstly, the labels were removed from the datasets, which can be acted as the unlabeled flows for the unsupervised learning training. After training, one can save the trained encoder in the VAE model for further semi-supervised learning. The training process for the CNN part mainly aim to convert the labeled data into digital encoding for fine-tuning.

In the single CNN model, we process the dataset and divide it into training and testing sets according to ratio, which is then trained and evaluated in the CNN.

For the E-CNN and single CNN models, we obtained the following two diagrams by adjusting the ratio of the dataset to the training and testing sets. The horizontal axis in the diagram represents the partitioning ratio. For example, 0.2 means that 80$\%$ of the entire dataset is allocated to the training set, and the remaining 20$\%$ is allocated to the testing set. The vertical axis represents the accuracy of the model training results under this partition. Fig. 9 displays the model results of E-CNN and single CNN in real-life scenarios, while Fig. 10 shows the results in the public dataset.

According to the figures, as the partition ratio increases, meaning that the training set data decreases, the accuracy of both models tends to decrease and becomes similar, both at around 0.7. Within the partition ratio interval [0.2, 0.8], E-CNN achieves good results with higher accuracy than a single CNN.

We have set the partition ratio to 0.45. After training and testing, we obtained the respective Confusion Matrix and Classification Report for E-CNN and single CNN using this ratio. The Confusion matrix is shown in Fig. 11 and Fig. 12. In the figures, the x-axis represents the prediction labels, the y-axis represents the actual labels, and the color intensity indicates the count of correct and incorrect predictions. According to the figure, under the partition with this ratio, E-CNN shows relatively accurate predictions for ’QQmusic,’ ’Teamfight Tactics,’ and ’Bilibili.’ On the other hand, compared to CNN, E-CNN’s predictions for ’QQmusic,’ ’IQiyi,’ and ’Teamfight Tactics’ are relatively accurate, but there are also several errors.

The classification report is presented in Table V and Table VI. It is evident from the tables that, under this partition ratio, E-CNN demonstrates relatively accurate predictions for ’Teamfight Tactics’ and ’Bilibili,’ with prediction accuracy rates above 0.9. The prediction accuracy for other applications is also around 0.7, resulting in a total classification accuracy of 0.7422. In contrast, CNN’s performance is notably inferior to that of E-CNN. The lowest prediction accuracy is approximately 0.5, and the highest prediction accuracy is only around 0.8. Its overall classification accuracy is merely 0.6989.

Based on the experimental results presented above, we can conclude that when the ratio of labeled data to total data falls within the range of [0.2, 0.8], the VAE+CNN method proposed in this paper can be employed for classifying and recognizing network application traffic. At this range, the recognition outcomes outperform those achieved by solely training with CNN, and the classification accuracy is commendable.

Fig. 13 visualizes the importance of ten features for the classification results. These ten features are the most significant ones contributing to the classification outcomes. For instance, in the case of the applications Bilibili and IQiyi, the most crucial feature for the classification result is s_idleMax, representing the maximum idle time of any neighboring packet in a stream. Notably, video streams exhibit distinct characteristics on this feature due to their continuous packet sending and shorter inter-packet idle time, which differentiates them from game streams.

The XAI-Pruning model effectively filters out unimportant weights and parameters, substantially reducing the number of parameters and the overall model size. As a result, the model becomes more compact, enabling significantly faster inference and quicker predictions. On the other hand, the baseline model, which includes all weights and parameters, generally exhibits higher prediction accuracy due to its capacity to capture more intricate model details and complexity.

As depicted in Table VII, the baseline model size is approximately three times larger than the XAI-Pruning model, and the inference time is doubled. However, despite these changes, both models still exhibit comparable prediction accuracy without significant differences. Consequently, the pruning method proposed in this paper effectively reduces the model’s compression size and inference time to some extent while having a relatively minor impact on prediction accuracy.