SNAKE: A Sustainable and Multi-functional Traffic Analysis System utilizing Specialized Large-Scale Models with a Mixture of Experts Architecture

Traditional traffic classification techniques, such as port-based  [25] and payload-based  [49] methods, rely on plain-text information or fixed rules in unencrypted traffic. However, dynamic port mapping and the widespread use of encrypted transmission protocols like TLS/SSL have rendered these methods largely ineffective. In response, researchers have turned to machine learning, using labeled traffic datasets and supervised learning models for automated classification. For instance, T. Bujlow and colleagues applied the C5.0 algorithm  [8], which combines multiple decision trees to enhance accuracy. Similarly, S. Huang and others utilized the k-Nearest Neighbors (kNN) algorithm  [22] to classify internet traffic based on statistical features, bypassing the need for a training phase. Other methods, such as Support Vector Machine (SVM) and Naive Bayes (NB)  [5], have also been explored. However, traditional machine learning techniques often struggle with over-fitting and lack the capability for fine-grained classification in an open world [21]. This has led to the exploration of deep learning approaches, which offer improved performance and adaptability in complex network scenarios.

With the rapid advancement of deep learning technology, neural network architectures such as residual neural networks have enabled deeper exploration of features related to data packets. For example, Wei et al.  [46]transformed the first 784 bytes of traffic packets from session groupings into image pixels to construct a CNN-based traffic classifier. In 2022, Dai et al. [12] applied the Transformer model architecture, known for its success in natural language processing, to traffic classification with impressive results. More recently, Wang et al. have established proofs for various mathematical theoretical foundations of deep learning in the field of traffic classification  [47]. Additionally, some deep learning methods for low-quality traffic data have also emerged  [39].

Despite these advancements, the goals of encrypted traffic classification are diverse, including application classification, VPN  [20] and Tor  [41] detection, service type identification, and monitoring of malicious traffic such as malware  [7], attacks in IPv6  [11], DDoS attacks  [24], and even attacks in IoT environments  [14]. However, these models are not well-integrated, posing challenges for multi-objective traffic analysis.

Multi-task learning approach: In the early stages of deep learning, multi-task learning was introduced to tackle various classification tasks within a single model  [52]. This approach connects multiple fully connected layers, each with its own loss function, to a shared neural network, allowing simultaneous training for different tasks. Some research has applied this method to traffic classification; however, our practical experience suggests limited feasibility for two main reasons:
       1) Different tasks require distinct feature sets. For instance, identifying unknown protocols demands analyzing local byte entropy changes, while VPN traffic classification relies on detecting packet re-encryption. In contrast, application and service type classification focuses on payload length sequences, and detecting malicious behaviors like DDoS attacks requires attention to statistical traffic features. These diverse features are challenging to integrate into a single model.
       2) Traffic classification attributes can exhibit subordinate and parallel relationships, creating complex task domains  [38]. For example, traffic may simultaneously belong to a specific protocol, application, and be classified as an attack or associated with VPN/Tor routing. When multi-task learning directly connects these classification targets, task divergence or entanglement can impede effective model training.

Large Model Fine-tuning Approach:Since last year large language models has advanced rapidly, with some of its latest innovations being applied in the field of network  [53]. However, there has yet to be relevant research that has developed a traffic analysis model capable of achieving multiple classification objectives simultaneously. Current studies typically take one of two approaches: ①Utilize established open-source large language models, such as Llama and Grok, and develop a model specifically tailored for traffic classification using fine-tuning techniques. ② Drawing inspiration from the Masked AutoEncoder (MAE) concept found in BERT  [31] and GPT  [51], training a large model head specifically for traffic classification, followed by fine-tuning the terminal layers to perform designated classification tasks  [13]. The first approach, which leverages the parameters of large language models, undoubtedly possesses strong classification capabilities. However, since language models are generally designed as broad conversational tools, most of their parameters are actually irrelevant to traffic classification. This vast number of parameters can result in high computational costs and delays in output during task execution, rendering it less suitable for high-throughput traffic classification scenarios. The second approach appears to be more advantageous, but it primarily aids in training downstream models. Ultimately, these models still need to be executed separately to produce classification results, which does not address the challenges highlighted in the first subsection.

The Mixture of Experts (MoE) architecture  [16] is designed to facilitate the horizontal expansion of models by dynamically routing inputs to different sub-model modules based on actual demands. In neural network models, transforming specific layer outputs into MoE layers allows for the identification of the most suitable expert modules for processing current inputs. By distributing tasks among one or more expert models, the system can efficiently handle a wide array of tasks while maintaining a sparse parameter distribution. This design enables the selective activation of parameters during both training and inference, thereby reducing unnecessary computational power consumption. When applied to traffic analysis, this approach can significantly enhance the speed of result generation.

Additionally, the MoE architecture provides strong horizontal scalability, allowing for the seamless integration of new functional modules into an existing trained model. This capability is especially advantageous in the context of rapidly evolving network traffic scenarios. Advanced variants such as the Multi-gate Mixture of Experts (MMoE)  [36] further refine this approach by employing multiple gating networks to allocate inputs to the most relevant experts, enhancing model performance in multi-task learning environments. Other techniques, such as Hierarchical MoE  [29] and Switch Transformers  [17], provide additional flexibility and efficiency in routing and processing, demonstrating the versatility of MoE architectures in adapting to complex and dynamic classification tasks.

As discussed, the current landscape of traffic classification tasks is exceedingly diverse, with many tasks playing a crucial role in network security and management. While the deployment of classification systems aims to provide a comprehensive understanding of data flows, there is a lack of a cohesive framework to integrate techniques from various domains. A single model often struggles to meet the complex demands of different scenarios. For instance, intrusion detection systems must address malicious data from various environments, including mobile devices, IoT, even in devices using microphone [33], as well as diverse attack patterns such as DDoS [30] and malicious DoH .

Moreover, common user traffic can experience concept drift due to changes in internet protocols and emerging services, impacting the overall classification effectiveness of the model. The significant disparities among different tasks further impede multi-task learning architectures designed to accommodate multiple objectives. This challenge stems from the direct sharing of underlying model parameters  [16], which restricts their ability to meet the varied requirements of distinct tasks.

The objective of the SNAKE system is to facilitate rapid model fusion and multi-target identification in the complex domain of network traffic analysis. Our final implementation enables the architecture trained by the SNAKE system to classify various application types, service categories, VPN and Tor usage, as well as identify four major categories of Android malware (including 42 specific attack tools), three types of DoH attacks, and ten types of Trojans. Achieving these functionalities requires the SNAKE model to exhibit strong scalability and specific handling mechanisms for compatibility across different task patterns.

To construct a highly scalable traffic classification system compatible with multiple identification targets, SNAKE draws on the MMoE architecture to build its classification model. At the model hierarchy level, SNAKE consists of four important components. First, all traffic classification tasks share a unified data preprocessing system for feature extraction. To accommodate different tasks, we have referenced the work of Distiller  [3], extracting information from traffic headers and payloads to form a common input for our model.

The second component is the combination of expert networks, which are tailored for different traffic classification tasks. The SNAKE system can set up the same network structure to directly read pre-trained model parameters. Certainly, the model architecture also supports the option to train a specific classification task from the ground up; however, due to the complexity of the overall network structure, we do not recommend adopting this approach.

The third component comprises an intelligent multi-gating structure featuring configurable modes, which is employed to regulate the outputs of multiple expert networks tailored for specific tasks. We have developed three distinct gating schemes corresponding to three different task fusion scenarios, with detailed descriptions to be provided in subsequent sections.

Finally, the architecture includes several classification tower networks designed for various models, which typically comprises fully connected networks. Following the model fusion process, only a limited number of fine-tuning iterations will be conducted on this segment of the model parameters. The overall architecture of the SNAKE system is depicted in Figure 2.

Our method adopts a traffic preprocessing identifying individual flows as the fundamental classification unit, uniquely defined by their five-tuple. The input data provided to the classifier comprises two types: (a) the first Nb bytes of the transport-layer payload (PAY) of the Traffic Classification object  [48]; and (b) informative protocol header fields (HDR) from the first Np packets  [34]. In the first case, the input is represented in binary format, organized in a byte-wise manner, and normalized within the range of (0, 1). The second type of input data includes: (i) the number of bytes in the transport-layer payload, (ii) the TCP window size (which is set to zero for UDP packets), (iii) inter-arrival time, and (iv) packet direction $\in{0,1}$ derived from the first $N_{p}$ packets.

It is important to note that, in both scenarios, longer instances are truncated while shorter ones are padded with zeros to conform to the specified lengths of bytes ($N_{b}$) or packets ($N_{p}$). This specific input data selection is driven by the need to mitigate biased inputs—a common pitfall identified in related works—which may arise from factors such as PCAP metadata, data-link layer information, and certain transport-layer header fields (e.g., source and destination ports). Such biases can lead to inflated performance metrics and hinder generalization capabilities.

This processing approach enables us to accommodate both flow-level and packet-level perspectives, thereby supporting a range of tasks with varying emphases. For instance, while tasks related to DDoS attacks or port scanning require a flow-centric view, application classification tasks necessitate a detailed analysis of the payload data.

The expert classification networks served as the primary entity for executing a variety of tasks. In this paper, we adopt a fine-tuning approach for the integration of expert sub-models, wherein we first independently train usable models on their respective labeled datasets. Subsequently, the SNAKE system retrieves and freezes the parameters of the core components from these models, relying solely on the gating mechanisms and tower networks in the latter stages for fine-tuning. It is acknowledged that there exists a plethora of model architectures tailored for different classification tasks, with extensive research conducted in the field. However, the diverse input formats of these models are evidently detrimental to the model fusion process. In light of the data preprocessing methods discussed in the previous section, this paper proposes a unified architecture for the expert sub-models.

In the original Distiller work, two modalities of input were analyzed using GRU and convolutional neural networks, which were effective for capturing temporal dependencies and texture features of payload data, respectively. However, with the advent of the transformer architecture, we have identified that this framework can effectively accommodate both modalities, thereby simplifying the model structure and avoiding the complexities arising from the use of heterogeneous architectures. The transformer architecture employs a multi-head attention mechanism composed of queries (Q), keys (K), and values (V), enabling it to extract sequential dependencies across varying spans, thus fulfilling the analytical requirements for both flow-level and packet-level modalities. In this paper, we uniformly adopt an expert model structure that flattens and concatenates the data from both modalities as input, establishing two attention heads for feature extraction. The specific model structure parameters and hyper-parameters utilized for training each traffic classification task are presented in Table 1.

During the model fusion process within the SNAKE system, only the parameters of the core transformer model components are read, excluding the final fully connected layers. We would like to emphasize that there may exist more optimal data processing or model configuration strategies for each specific task. However, this paper primarily focuses on evaluating the model fusion and expansion capabilities of SNAKE; hence, we have selected a relatively generic architecture for this purpose. The system itself supports customization of the structure, parameters, and hyper-parameters for each expert network.

In this paper, we adopt a supervised learning paradigm for traffic classification. In the field of network traffic research, labeled sample data suitable for various tasks, such as malicious attacks and traffic application types, can be collected using constructed scripts and some automated tools. Thus, dataset can be described as $D_{k}(x_{i},y_{i,k})$. The subscript $k$ denotes the classification task for the $k_{th}$ attribute, where $x_{i}$ represents the $i_{th}$ traffic sample and $y_{i,k}$ is the corresponding label. We build the SNAKE system to provide multiple attribute categories of traffic at once, and this system employs a MMoE architecture, establishing an independent gating layer for each task. Therefore, for the known classification task of $k$ attributes, we need to establish $k$ Gates and corresponding Tower classification networks for the dataset. We assume that $n$ expert sub-models $f_{j}(x)$ have been used previously, and we need to define a concatenation operation to represent it as follows:

After the $k_{th}$ gate $g_{k}$ receives the overall output $X$ , it needs to be integrated and input into the corresponding Tower $h_{k}$. Specific operation can be represented as follows:

In the conventional MMoE architecture, the gate is a learnable linear transformations of the input with a softmax layer. However, in the field of traffic classification, whether the outputs of different expert models are mixed depends on the correlation between the corresponding task attributes. Therefore, the gating we use is merely a vector of the same dimension as the traditional MMoE method, with a length equal to the number of expert models $n$, like $g_{i}(\delta_{1i},\delta_{2i},\dots,\delta_{ni})$. For the specific value settings, we have the following three modes:

Default Mode: In this mode, if there is only one expert model connected to the gate, the task and the expert are in a one-to-one correspondence. Consequently, the gating mechanism is configured to output only the result from the corresponding expert, as expressed below:

Top-K Mode: In this mode, multiple expert models are connected to the gate, and these experts share the same classification granularity. For instance, Expert 1 may classify $app_{1}$, $app_{2}$, …, $app_{n}$, while Expert 2 classifies $app_{n+1}$, …, $app_{m}$. Here, the task and experts are in a one-to-many relationship, and the gating mechanism is designed to equally consider the results from those experts (which could be represent by a set $S_{j}\subseteq\{1,2,\ldots,n\}$). The setting of the gating vector in the top-k mode is as follows:

Trainable Mode: In this mode, multiple expert models are connected to the gating layer, and these experts operate at different classification granularities. For example, Expert Model 1 may judge whether this traffic is malicious, while Expert Model 2 focuses on specific tools. In this case, there exists a many-to-many relationship between the tasks and the experts. If the gating layer were to equally consider the outputs of both models, it could interfere with their respective tasks. Therefore, these gating layers should be set to an updatable state, allowing for timely updates during fine-tuning based on the gradients returned by the loss function. Specifically, the gating vector can be updated using the following formula:

where $\omega$ is the trainable weight calculated through a linear layer and softmax function, ensuring that the gating layer can flexibly adapt to the outputs of different expert models, thereby optimizing overall classification performance.

Each gating layer connects to the final tower network for the task, which is responsible for producing the ultimate output. For this purpose, we uniformly configure a fully connected layer network. Additionally, to address potential over-fitting during fine-tuning, we have incorporated dropout layers. Specifically, tower networks are configured as a two-layer fully connected network, utilizing a dropout rate of 0.2, a learning rate of 0.001, a batch size of 128, and the ReLU activation function.

It is important to note that the architecture, parameters, and hyper-parameters of each tower can also be customized. Through experimentation, we have found that adding more fully connected layers does not significantly impact model performance; thus, we have opted for a simplified configuration in this context.

The pseudo-code for the system’s operation is presented in Algorithm 1.

In this paper, aside from the corresponding comparative experiments, the default expert network structure and parameters, tower model parameters, dropout settings, and other hyper-parameter configurations are consistent with those outlined in Section 3. Some brief information about the datasets used in our experiments can be found in Table 3. Due to the use of multiple datasets, a detailed introduction is provided in Appendix Section B for the sake of narrative clarity.

During the model training, we performed a uniform split of the dataset, allocating $75\%$ to the training set, $10\%$ to the validation set, and $20\%$ to the test set. Whether it is the training of the expert model or the fine-tuning and integration operations conducted by the SNAKE system, the training samples are consistently drawn from the training set. Therefore, the experimental results will not be misjudged due to any overlap in the datasets. We conducted experiments with ten random data splits, so all the following experiments present the results of these ten repeated trials.

In our experiments, we utilized deep learning models implemented with PyTorch 2.0.0 and CUDA 11.7. Data pre-processing and post-processing were primarily conducted using the NumPy and Scapy libraries. For graphical data representation, we employed Matplotlib and MATLAB. All experiments were performed on a PC with the following hardware specifications: a 13th Gen Intel® Core™ i7-13700KF processor running at 3.40 GHz, 32 GB of RAM, and an NVIDIA GeForce RTX 4080 GPU. The operating system used was Windows 11.

In this subsection, we aimed to compare the recognition performance of expert networks using different model architectures. Here, we compare five traffic classification models that utilize data processing schemes similar to ours. We evaluated the detection results of these models on the ISCX-VPN-2016 dataset for three tasks: encapsulation, service type, and application, as shown in Table 3.

In this experiment, all models were set with a batch size of 16 and trained for 50 epochs. The dataset was randomly split with 75$\%$ for the training set, 10$\%$ for the validation set, and 15$\%$ for the test set, and this process was run ten times. The table shows the accuracy, precision, and F1-score values of five model architectures. The CNN-RNN-2a model exhibited significant oscillations, possibly because the small batch size and number of epochs were not well-suited for training RNN models. Therefore, its range of variation is presented in the table, while the results of the other models are shown as average results and range performance. It is evident that the expert model architecture adopted in this paper performed the best in this three tasks, outperforming the other optimal results by approximately 0.5$\%$ to 2.5$\%$ on average. According to the experimental results, the transformer method we employed demonstrates significantly superior overall performance across the three tasks. Therefore, we used this model architecture for all expert networks in the classification tasks.

To explore the optimal gating configuration pattern and to lay the groundwork for the SNAKE system to integrate complex traffic classification task sets, we set up experiments in three traffic classification task fusion scenarios to test the effectiveness of different gating configurations. The schematic diagram of the scenarios is shown in Figure 4.

In this section of the experiment, we aimed to verify whether the settings of certain hyper-parameters in SNAKE would affect the overall performance of the model. First, we analyzed Task Attribute-Independent Scenarios and utilized three classification from the ISCX-VPN-2016 dataset to examine the impact of different batch sizes, the number of layers in the tower, and the dropout rates in the tower layers on the model’s recognition performance. The corresponding trends in model accuracy are illustrated in Figure 8.

From the figure, it is not difficult to see that the batch size has a relatively minor average impact on the model accuracy in this scenario. However, a smaller batch size can lead to bolder updates in the direction of optimization, potentially resulting in slightly better performance. That said, when fusing models, if there are too many label types across two or more tasks, a smaller batch size might result in some labels not being represented during a single training iteration, which could negatively affect the training process. Therefore, we believe that setting the batch size to 128 is more reasonable. In terms of other hyper-parameters, setting dropout for the Tower network results in a slight improvement in the model’s average performance, although this is not very noticeable. On the other hand, increasing the depth of the Tower and not freezing the parameters of the expert models can lead to a significant enhancement in model performance. However, both of these approaches may slightly increase the computational resource requirements, so whether to implement these settings should depend on the specific needs of the user.

For the scenario in Experiment 5.3.2, the selection of batch size is similar to the situation mentioned above, so we will not elaborate further here. In this case, the base model did not use dropout, the learning rate was set to 0.001, and the parameters of the expert networks were frozen. We conducted experiments with a smaller learning rate, deployed dropout layers in the Tower layer, and allowed the parameters of the expert networks to be updated, observing the impact of these parameter settings on the experimental accuracy. The model’s classification accuracy for the IPTAS and ISCX-VPN-2016 datasets under these configurations is shown in Figure 9.

From the figure, we can conclude that in this scenario, model fusion indeed requires a slightly higher learning rate to quickly update the mapping relationships. Additionally, configuring dropout layers in the Tower network helps improve the model’s accuracy. While allowing the expert models to be updated can enhance the performance of the model fusion, it requires more computational resources and needs to be allocated as needed. Furthermore, as we observed in Section 5.3.2, after configuring the dropout layers, the model’s recognition performance has fully reached the recognition accuracy of the independent expert models. This indicates that as long as the hyper-parameters are set appropriately, the SNAKE system can effectively approximate the capabilities of individual models across each classification dimension.

In this final experiment, we aim to showcase the model expansion that the SNAKE system has achieved so far, as well as its recognition performance on related tasks. While there is potential to further expand upon a variety of additional models and tasks, this paper will limit its scope to the current findings. Figure 10 illustrates the detection performance on various task test sets following model fusion, in comparison to the original individual models. Due to constraints on the length of this manuscript, the training processes and detailed results of the individual models are provided in Appendix D .

Based on the figure, it can be observed that the performance of the model fusion for each task closely aligns with that of the individual models, indicating that the SNAKE system has preliminarily developed a relatively universal large-scale model in the field of traffic classification.