NetMamba: Efficient Network Traffic Classification via Pre-training Unidirectional Mamba

Given the stride array, we initially perform a linear projection on each stride $\mathbf{s}_{i}$ to a vector with size $\mathtt{D}_{\text{enc}}$ and incorporate position embeddings $\mathbf{E}_{\text{enc}}^{\text{pos}}\in\mathbb{R}^{N_{s}\times\mathtt{D}_{\text{enc}}}$ as shown below:

where $\mathbf{W}\in\mathbb{R}^{L_{s}\times\mathtt{D}_{\text{enc}}}$ represents the learnable projection matrix. Inspired by ViT [29] and BERT [7], we introduce a class token to represent the entire stride sequence, denoted as $\mathbf{x}_{\text{cls}}$. Since the unidirectional Mamba processes sequence information from front to back, we opt to append the class token to the end of the sequence for enhanced information aggregation.

Recently, several variants of Mamba have been proposed to accommodate domain-specific data formats and task requirements. For instance, Vim [12] incorporates bidirectional Mamba blocks for spatial-aware understanding of vision tasks, Graph-Mamba [16] introduces a graph-dependent selection mechanism for graph learning, while MiM-ISTD [30] customizes a cascading Mamba structure for extracting hierarchical visual information. We argue that the original unidirectional Mamba design[14], tailored for sequence modeling, is well-suited for representation learning in sequential network traffic, offering increased efficiency through the elimination of omnidirectional scans and redundant blocks. We carefully test different Mamba variants, demonstrating that the selected unidirectional Mamba is more suitable for processing network traffic. Please refer to the ablation studies for more details.

Hence, we implement the NetMamba encoder and decoder using unidirectional Mamba blocks, as illustrated in Figure 2. The operational process of the NetMamba block forward pass is outlined in Algorithm 1. For a given input token sequence $\mathbf{X}_{t-1}$ with a batch size $\mathtt{B}$ and sequence length $\mathtt{L}$ from the $(t-1)$-th NetMamba block, we begin by normalizing it and then projecting it linearly into $\mathbf{x}$ and $\mathbf{z}$, both with dimension size of $\mathtt{E}$. We subsequently apply causal 1-D convolution to $\mathbf{x}$, resulting in $\mathbf{x}^{\prime}$. Based on $\mathbf{x}^{\prime}$, we compute the input-dependent step size $\mathbf{\Delta}$, as well as the projection parameters $\mathbf{B}$ and $\mathbf{C}$ having a dimension size of $\mathtt{N}$. We then discretize $\overline{\mathbf{A}}$ and $\overline{\mathbf{B}}$ using $\mathbf{\Delta}$. Following this, we calculate $\mathbf{y}$ employing a hardware-aware SSM. Finally, $\mathbf{y}$ is gated by $\mathbf{z}$ and added residually to $\mathbf{X}_{t-1}$, resulting in the output token sequence $\mathbf{X}_{t}$ for the $t$-th NetMamba block.

Given the embedded stride tokens $\mathbf{X}_{0}\in\mathbb{R}^{\mathtt{L}\times\mathtt{D}_{\text{enc}}}$, a portion of strides is randomly sampled while the remaining ones are removed. For a predefined masking ratio $r\in(0,1)$, the length of visible tokens is determined as $\mathtt{L}_{\text{vis}}=\lceil(1-r)\mathtt{L}\rceil$. The visible tokens are then sampled as follows:

where the $\mathbf{Shuffle}$ operation permutes the token sequence randomly. Notably, we ensure that the trailing class token remains unmasked throughout this process since its role in aggregating overall sequence information necessitates its preservation at all times.

The primary objective behind random masking is the elimination of redundancy. This approach creates a challenging task that resists straightforward solutions through extrapolation from neighboring strides alone. Additionally, the reduction in input length diminishes computational and memory costs, offering an opportunity for more efficient model training.

The NetMamba encoder is tasked with capturing latent inter-stride relationships using the visible tokens, whereas the NetMamba decoder’s objective is to reconstruct masked strides utilizing both the encoder output tokens and mask tokens. Each mask token represents a shared, trainable vector indicating the presence of a missing stride. Additionally, new positional embeddings are added to provide location information to the mask tokens.

The formal forward process of NetMamba pre-training can be outlined as follows:

where the $\mathbf{Unshuffle}$ operation restores the original sequence order, and $\mathbf{E}_{\text{dec}}^{\text{pos}}\in\mathbb{R}^{\mathtt{L}\times\mathtt{D}_{\text{dec}}}$ represents decoder-specific positional embeddings. Subsequently, the mean square error (MSE) loss for self-supervised reconstruction is calculated as shown below:

where $\mathbf{y}_{\text{real}}$ represents the ground-truth mask tokens, and $\mathbf{y}_{\text{rec}}$ signifies the predicted ones.

To assess the effectiveness and generalization abilities of NetMamba, we conducted experiments using six publicly available real-world traffic datasets encompassing three main classification tasks.

1. 1.

   Encrypted Application Classification: This task aims to classify application traffic under various encryption protocols. Specifically, the CrossPlatform (Android) [31] and CrossPlatform (iOS) [31] contain 254 and 253 applications respectively. Additionally, we use Tor traffic data from 8 communication categories in ISCXTor2016 [32] and VPN traffic data from 7 communication categories in ISCXVPN2016 [33].

2. 2.

   Attack Traffic Classification: This task aims to identify potential attack traffic, such as Denial of Service (DoS) attacks and brute force attacks. We construct 6 data categories using CICIoT2022[34].

3. 3.

   Malware Traffic Classification: This task aims to distinguish between traffic generated by malware and benign traffic. We use all 20 data categories from the USTC-TFC2016 dataset [35].

We observed an imbalance in flow counts across traffic categories, which adversely impacts model performance. To address this, we set upper and lower flow limits for each category. Categories below the lower limit are discarded, while those above the upper limit are randomly sampled.

To comprehensively evaluate NetMamba, we conducted comparisons with various open-source baselines and state-of-the-art techniques, as outlined below:

1. 1.

   Classical machine learning methods such as AppScanner [3] and FlowPrint [2] that rely on statistical features for traffic classification.

2. 2.

   Deep learning approaches like FS-Net [4] and TFE-GNN [6] that utilize packet lengths or raw bytes to perform traffic analysis in a supervised manner.

3. 3.

   Transformer-based models such as ET-BERT [9] and YaTC [10] that capture traffic representations during pre-training and subsequently fine-tune for specific tasks with limited labeled data. In particular, we implement YaTC(OF) by substituting packet-level and flow-level attention with a global attention module, which expedites model inference while removing its original memory optimization.

4. 4.

   Transformer variants within the NetMamba backbone, including NT-Vanilla and NT-Linear. The former replaces Mamba blocks in NetMamba with vanilla Transformer blocks [36] featuring quadratic complexity, while the latter adopts Linear Transformer blocks [37] with linear complexity.

At the pre-training stage, we set the batch size to $\mathtt{B}=128$ and train models for 150,000 steps. The initial learning rate is set to $1.0\times 10^{-3}$ with the AdamW optimizer, alongside a linear learning rate scaling policy. Additionally, a masking ratio of $r=0.9$ is employed for randomly masking strides.

For fine-tuning, we adjust the batch size to $\mathtt{B}=64$ and set the learning rate to $2.0\times 10^{-3}$. Each dataset is partitioned into training, validation, and test sets following an 8:1:1 ratio. All models are trained for 120 epochs on the training data, with checkpoints saving the best accuracy on the validation set, subsequently evaluated on the test set.

The NetMamba architecture features an encoder composed of 4 Mamba blocks and a decoder composed of 2 Mamba blocks. More hyper-parameter details can be found in Table III.

The proposed model is implemented using PyTorch 2.1.1, with all experiments conducted on a Ubuntu 22.04 server equipped with CPU of Intel(R) Xeon(R) Gold 6240C CPU @ 2.60GHz, GPU of NVIDIA A100 (40GB $\times$ 4).

We assess the performance of NetMamba using four typical metrics: Accuracy(AC), Precision(PR), Recall(RC), and weighted F1 Score(F1).

The CrossPlatform Android and iOS datasets consist of encrypted traffic generated by the top 100 applications from the US, China, and India, encompassing over 200 categories in both cases. Additionally, ISCXTor2016 includes application traffic using the Onion Router (Tor) for encrypted communications.

As shown in Table IV, NetMamba demonstrates superior performance across these three datasets, with F1 score improvements ranging from 0.18% to 0.31%. In comparison to state-of-the-art pre-trained methods, ET-BERT exclusively focuses on learning traffic representations from packet payloads, overlooking the valuable information carried by packet headers. Moreover, beyond the variances in base model architecture, YaTC differs from NetMamba in the traffic representation scheme, employing a two-dimensional splitting technique. In contrast, NetMamba effectively models both header and payload characteristics and optimizes traffic data segmentation to mitigate biases. This leads to a more comprehensive analysis of traffic patterns.

The CICIoT2022 dataset consists of traffic collected from a laboratory network designed for profiling, behavioral analysis, and vulnerability testing of various IoT devices. The ISCXVPN2016 dataset includes encrypted communication traffic tunneled through Virtual Private Networks (VPN). The USTC-TFC2016 dataset comprises encrypted traffic from both malware and benign applications.

As depicted in Table VI, NetMamba performs only slightly worse than YaTC on the ISCXVPN2016 and USTC-TFC2016 datasets, while achieving comparable performance to other state-of-the-art methods on the CICIoT2022 dataset. Although TFE-GNN achieves the highest performance on the CICIoT2022 dataset, this non-pre-trained model falls significantly behind NetMamba on the other datasets, highlighting its unstable classification performance.

Initially, we replaced all unidirectional Mamba blocks in NetMamba with bidirectional ones used by [12]. The experimental results revealed a slight performance decline across all datasets except for CICIoT2022. This suggests that unidirectional Mamba is well-suited for processing network traffic data, given that packets are transmitted sequentially and earlier packets possess limited information about subsequent ones. Moreover, incorporating bidirectional or even omnidirectional Mamba blocks introduces additional computational and memory overheads due to extra scan passes, ultimately reducing efficiency. Thus, unidirectional Mamba stands out as the preferable choice.

Following [30], we substituted the original unidirectional Mamba block with a cascading structure where each inner block processes data of different granularity. We observed a notable drop in classification performance across all datasets, indicating that this complex structure is inferior to processing sequential traffic data.

To further investigate the expressiveness of the Mamba architecture, we evaluated two Transformer-based ablation variants, NT-Vanilla and NT-Linear. The results show that NT-Vanilla performs slightly worse than NetMamba across all datasets except for CICIoT2022. This indicates that the linear-time Mamba model is well suited for capturing sequential traffic data using our proposed representation scheme. Moreover, NT-Linear, due to information loss from the over-compression of its attention mechanisms, exhibits unstable performance and falls significantly behind NetMamba on three datasets.

While positional information is inherently preserved in sequence models such as Mamba, eliminating explicit positional embedding still results in a reduction in accuracy ranging from 0.03% to 2.17% across all datasets. This suggests that reinforced positional information aids the model in capturing correlations within sequential traffic data.

The pre-training process is designed to capture general traffic understanding from extensive unlabeled data. When compared to the non-pre-trained counterpart, pre-trained NetMamba demonstrates accuracy improvements ranging from 0.20% to 4.70%, affirming the effectiveness of our MAE-based pre-training task.

When all header bytes in a packet are omitted, classification performance declines significantly, with accuracy dropping by 15.51% to 48.75%. This highlights the critical role of key fields within packet headers—such as port number, protocol, and packet length—which have been proven effective in traffic classification [4, 38, 39].

Regarding packet payloads, the ablation results show accuracy drops ranging from 0.16% to 0.84% across four datasets. This underscores the contribution of potential plaintext and specific encrypted payloads to improved traffic understanding.

Likewise, the vertical bias information introduced by the 2-dimensional patch splitting results in a maximum accuracy decline of 1.88%, highlighting the importance of 1-dimensional stride cutting.