PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Tongze Wang, Xiaohui Xie, Wenduo Wang, Chuyi Wang, Youjian Zhao, Yong Cui

arXiv:2405.11449v4cs.LGcs.NI

TL;DR

Network traffic models face challenges from limited labeled data and unreliable traffic representations. NetMamba combines a linear-time state space model with a traffic representation scheme; across classification tasks, it reports strong accuracy, faster inference, and few-shot gains.

  • Problem

    Deep-learning traffic methods require extensive labeled datasets and can be susceptible to bias and less adaptable to novel data distributions.

    Source: §I · p. 1

  • Method

    NetMamba uses a linear-time state space model and a traffic representation scheme, with stride-based inputs and masked-autoencoder pre-training followed by fine-tuning.

    Source: §I · pp. 1–2

  • Results

    Across classification tasks, NetMamba reports accuracy above 90%, sometimes above 99%, inference speedups up to 60 times, and better few-shot performance with fewer labels.

    Source: §I · p. 2

  • Takeaways & Limitations

    The authors state that NetMamba’s representation scheme and model design can address broader network-domain tasks, including quality-of-service and network-performance prediction.

    Source: §VIII · p. 10

    The current implementation depends on specialized GPU hardware.

    Source: §VIII · p. 10

Abstract

from arXiv · show

Network traffic classification is a crucial research area aiming to enhance service quality, streamline network management, and bolster cybersecurity. To address the growing complexity of transmission encryption techniques, various machine learning and deep learning methods have been proposed. However, existing approaches face two main challenges. Firstly, they struggle with model inefficiency due to the quadratic complexity of the widely used Transformer architecture. Secondly, they suffer from inadequate traffic representation because of discarding important byte information while retaining unwanted biases. To address these challenges, we propose NetMamba, an efficient linear-time state space model equipped with a comprehensive traffic representation scheme. We adopt a specially selected and improved unidirectional Mamba architecture for the networking field, instead of the Transformer, to address efficiency issues. In addition, we design a traffic representation scheme to extract valid information from massive traffic data while removing biased information. Evaluation experiments on six public datasets encompassing three main classification tasks showcase NetMamba's superior classification performance compared to state-of-the-art baselines. It achieves an accuracy rate of nearly 99% (some over 99%) in all tasks. Additionally, NetMamba demonstrates excellent efficiency, improving inference speed by up to 60 times while maintaining comparably low memory usage. Furthermore, NetMamba exhibits superior few-shot learning abilities, achieving better classification performance with fewer labeled data. To the best of our knowledge, NetMamba is the first model to tailor the Mamba architecture for networking.

I. INTRODUCTION

Network traffic classification is increasingly important but harder as encryption and anonymity grow. NetMamba addresses efficiency and representation gaps with a linear-time state space model and a traffic representation scheme, and reports strong classification, speed, and few-shot results.

A. Transformer-based Traffic Classification

  • Transformer-based traffic models incur quadratic self-attention costs on long sequences, while deep-learning approaches also require extensive labels and can be vulnerable to bias and distribution shifts.The introduction frames these as efficiency and traffic-representation challenges.Source: §I · p. 1

I. INTRODUCTION

  • NetMamba is presented as a linear-time state space model with a comprehensive traffic representation scheme, intended to improve inference speed and reduce memory use.The authors describe it as the first state space model specifically designed for network traffic classification.Source: §I · pp. 1–2
  • Its representation scheme retains useful header and payload content while removing unwanted biases through packet anonymization, byte-allocation balancing, and stride-based cutting.The authors report extensive evaluations across traffic-classification tasks.Source: §I · p. 2
  • Across classification tasks, NetMamba reports accuracy above 90%, sometimes above 99%, inference speedups up to 60 times, and better few-shot performance with fewer labels.The comparison is against existing baselines; memory usage is described as low.Source: §I · p. 2

B. Mamba-based Representation Learning

The paper situates NetMamba amid domain-specific Mamba models and identifies network traffic as an unaddressed application. It motivates a representation that preserves raw-traffic information and introduces state space models as a sequence-modeling foundation.

B. Mamba-based Representation Learning

  • Prior Mamba variants tailor sequence modeling to vision, language, graphs, and other domains, but the paper reports no prior successful application to network traffic classification.The authors position their work as addressing that gap.Source: B. Mamba-based Representation Learning · p. 2

C. Traffic Representation Schemes

  • Traffic representations based on compressed statistics can lose raw-datagram information, while deep-learning schemes may omit header bytes or introduce byte-balance and splitting biases.The proposed scheme is described as preserving hierarchical traffic information while eliminating biases.Source: C. Traffic Representation Schemes · p. 3

III. PRELIMINARIES

  • The preliminaries identify the Mamba block as NetMamba’s foundation and describe state space models as sequence models connected to recurrent and convolutional networks.The section introduces their definitions and components.Source: §III · p. 3, A. State Space Models · p. 3

A. State Space Models

  • A state space model maps an input sequence to an output sequence through an intermediate latent state, using linear ordinary differential equations.The passage identifies A as the evolution parameter and B and C as projection parameters.Source: A. State Space Models · p. 3

B. Discretization

This section connects state space model computation to efficient training and describes NetMamba’s traffic-input pipeline. The framework converts raw traffic into stride representations before pre-training and downstream use.

B. Discretization

  • Discretization converts continuous-time state space models for discrete real-world data; the resulting recurrence has linear-time inference but is not parallelizable during training.The convolutional formulation addresses that training-parallelization difficulty.Source: B. Discretization · p. 3

C. Selection Mechanism

  • Mamba makes the state space parameters input-dependent, enabling dynamic selection of relevant information, and uses a GPU-friendly implementation to reduce memory I/O and intermediate-state storage.This selection mechanism addresses the time-invariant parameters of standard state space models.Source: C. Selection Mechanism · p. 3

IV. NETMAMBA FRAMEWORK

  • NetMamba’s framework converts raw binary traffic into stride-based representations, then uses self-supervised pre-training followed by supervised fine-tuning for downstream traffic understanding.The training approach is described as dual-stage.Source: §IV · p. 3

1) Traffic Representation Phase:

  • Traffic is segmented into flows, and bytes from initial packets are aggregated into an array before non-overlapping strides are formed to preserve adjacent-byte relationships.The array integrates information across byte, packet, and flow levels while avoiding patch-splitting and tokenization issues described in the passage.Source: 1) Traffic Representation Phase · pp. 3–4

2) Pre-training Phase:

NetMamba learns traffic representations from unlabeled data through masked-stride reconstruction, then fine-tunes on labeled data for classification. The section also outlines flow formation and packet parsing for its traffic inputs.

2) Pre-training Phase:

  • During pre-training, NetMamba embeds and positionally encodes flow strides, randomly masks some, and trains an encoder-decoder to reconstruct the masked strides.The encoder processes visible strides; reconstruction loss is computed on masked strides.Source: 2) Pre-training Phase · p. 4

3) Fine-tuning Phase:

  • For fine-tuning, the decoder is replaced with an MLP classifier, all strides are visible, and the trailing class token aggregates traffic features for classification.The encoder is fine-tuned with labeled traffic data, and the passage notes adaptability with limited labels for downstream tasks.Source: 3) Fine-tuning Phase · p. 4

V. TRAFFIC REPRESENTATION

  • The traffic-representation section identifies Table II as the location of NetMamba’s key hyperparameters.Source: §V · p. 4, Table II · p. 4

1) Flow Splitting:

  • A flow groups packets belonging to a protocol and transmitted between two hosts through two ports.The passage describes flow-level packets as carrying interaction information between the hosts.Source: 1) Flow Splitting · pp. 4–5

2) Packet Parsing:

  • Packet parsing excludes non-IP packets, retains selected header and payload information, and removes Ethernet headers to anonymize identifying information.The retained fields include packet information such as total length and payload text.Source: 2) Packet Parsing · p. 5

3) Packet Cropping & Padding, and Concatenating:

NetMamba standardizes the first packets of each flow, concatenates their bytes, and cuts the resulting array into sequential strides. The stride sequence is then projected and position-embedded for the model.

3) Packet Cropping & Padding, and Concatenating:

4) Stride Cutting:

  • One-dimensional, non-overlapping strides divide the byte array while aligning with traffic order and retaining inter-byte correlations.The number of strides is Ns = Lb/Ls.Source: 4) Stride Cutting · p. 5

1) Stride Embedding:

  • Each stride is linearly projected to dimension Denc, and positional embeddings are added to form the encoder input.A class token is appended to the stride sequence.Source: 1) Stride Embedding · p. 5

2) NetMamba Block :

NetMamba uses unidirectional Mamba blocks for sequential traffic and pre-trains an encoder-decoder with randomly masked stride tokens. The encoder processes visible strides, while the decoder reconstructs masked ones.

2) NetMamba Block :

  • The selected unidirectional Mamba design processes traffic sequences without omnidirectional scans or redundant blocks.The authors report testing variants and selecting this design as more suitable for network traffic.Source: 2) NetMamba Block · p. 5

1) Random Masking:

  • Random masking reduces input length and computational and memory costs while challenging the model to learn beyond neighboring-stride extrapolation.The trailing class token remains unmasked to aggregate sequence information.Source: 1) Random Masking · p. 6

2) Masked Pre-training:

  • During masked pre-training, the encoder processes visible strides and the decoder reconstructs masked strides using encoder outputs and mask tokens.Decoder-specific positional embeddings provide location information, and training minimizes masked-stride reconstruction MSE.Source: 2) Masked Pre-training · p. 6

C. NetMamba Fine-tuning

For downstream classification, NetMamba loads its pre-trained encoder, replaces the decoder with an MLP head, and fine-tunes on labeled traffic.

  • Fine-tuning loads the pre-trained encoder and replaces the decoder with an MLP head for supervised classification with all stride tokens visible.Source: C. NetMamba Fine-tuning · p. 6
  • The trailing class token represents the sequence for prediction, and classification minimizes cross-entropy against the ground-truth label.Source: C. NetMamba Fine-tuning · p. 6
  • The authors report that the pre-trained encoder adapts to downstream tasks including application classification and attack detection using limited labeled data.Source: C. NetMamba Fine-tuning · p. 6

VII. EVALUATION

The evaluation uses six public traffic datasets spanning encrypted application, attack, and malware classification. It compares NetMamba with classical, deep-learning, pre-trained Transformer, and Transformer-backbone variants.

1) Datasets:

  • Six public datasets cover encrypted application, attack, and malware classification, including CrossPlatform Android and iOS, Tor, VPN, CICIoT2022, and USTC-TFC2016 traffic.Source: 1) Datasets · pp. 6–7

2) Comparison Methods:

  • Baselines include statistical-feature methods, supervised deep-learning models, pre-trained Transformer models, and vanilla or linear Transformer replacements for NetMamba blocks.Source: 2) Comparison Methods · p. 7

3) Implementation Details:

  • Pre-training runs for 150,000 steps with masking ratio 0.9; fine-tuning uses an 8:1:1 train-validation-test split and 120 epochs.The described architecture has four encoder Mamba blocks and two decoder Mamba blocks.Source: 3) Implementation Details · p. 7

4) Evaluation Metrics:

Across six datasets, NetMamba compares favorably with baselines, while its inference evaluation reports higher speed and comparably low memory use. The section also describes the evaluation metrics and task-specific comparisons.

4) Evaluation Metrics:

1) CrossPlatform (Android & iOS) and ISCXTor2016:

  • On CrossPlatform Android, CrossPlatform iOS, and ISCXTor2016, NetMamba improves F1 scores by 0.18% to 0.31%.The passage attributes the comparison to differences in traffic representation, including ET-BERT's payload-only focus and YaTC's two-dimensional splitting.Source: 1) CrossPlatform (Android & iOS) and… · p. 7

B. Overall Evaluation

  • Across six datasets, NetMamba outperforms all baselines on three, ranks second on two, and trails slightly on CICIoT2022 by at most 0.72 percentage points in accuracy and F1.The reported accuracy range is 0.9094–0.9986, and the F1 range is 0.9096–0.9986.Source: B. Overall Evaluation · p. 7, Table IV · p. 8, Table V · p. 8

2) CICIoT2022, ISCXVPN2016 & USTC-TFC2016:

  • On ISCXVPN2016 and USTC-TFC2016, NetMamba performs slightly worse than YaTC, while on CICIoT2022 it is comparable to other state-of-the-art methods.TFE-GNN has the highest performance on CICIoT2022 but performs significantly worse than NetMamba on the other datasets.Source: 2) CICIoT2022, ISCXVPN2016 &… · p. 8

C. Inference Efficiency Evaluation

  • At fine-tuning batch size 64, NetMamba is 2.24 times faster than YaTC(OF), and it achieves the highest inference speeds among deep learning methods with comparably low GPU memory use.Across input batch sizes, its speed improvement ranges from 1.22 to 60.11 times; at large batch sizes, its GPU memory use is lower than most models, with FS-Net, YaTC, and NT-Linear as exceptions.Source: C. Inference Efficiency Evaluation · pp. 8–9, Figure 3 · p. 9, Figure 4 · p. 9

D. Ablation Study

The ablations compare alternative sequence models, positional information, and pre-training. They report benefits for unidirectional Mamba and pre-training, alongside performance declines for several alternatives and removals.

1) Model-level Ablation:

  • Replacing unidirectional Mamba blocks with bidirectional ones slightly reduces performance on every dataset except CICIoT2022, while adding scan passes and computational and memory overhead.The authors favor unidirectional blocks because traffic packets are sequential and earlier packets contain limited information about later ones.Source: 1) Model-level Ablation · p. 9
  • A cascading Mamba structure that processes different granularities causes a notable classification-performance drop across all datasets.The comparison substitutes a cascading structure for the original unidirectional Mamba block.Source: 1) Model-level Ablation · p. 9
  • NT-Vanilla performs slightly worse than NetMamba on all datasets except CICIoT2022, while NT-Linear is unstable and falls significantly behind on three datasets.The authors attribute NT-Linear's result to information loss from over-compressing its attention mechanisms.Source: 1) Model-level Ablation · p. 9
  • Removing explicit positional embeddings reduces accuracy by 0.03% to 2.17% across all datasets.The authors interpret the result as evidence that reinforced positional information aids capture of correlations in sequential traffic.Source: 1) Model-level Ablation · p. 9
  • Pre-trained NetMamba improves accuracy by 0.20% to 4.70% over its non-pre-trained counterpart.The pre-training process uses extensive unlabeled data to capture general traffic understanding.Source: 1) Model-level Ablation · p. 9

2) Data-level Ablation:

Data-level ablations report the effects of removing packet headers or payloads and changing stride cutting. The section also reports few-shot evaluation and notes a deployment boundary for the current implementation.

2) Data-level Ablation:

  • Omitting packet payloads drops accuracy by 0.16% to 0.84% across four datasets.The authors connect the contribution to potential plaintext and specific encrypted payloads.Source: 2) Data-level Ablation · p. 9
  • Replacing one-dimensional stride cutting with two-dimensional patch splitting causes a maximum accuracy decline of 1.88%.Source: 2) Data-level Ablation · pp. 9–10
  • Omitting all packet header bytes drops accuracy by 15.51% to 48.75%.The passage identifies port number, protocol, and packet length as key header fields.Source: 2) Data-level Ablation · p. 9, Table VI · p. 10

E. Few-Shot Evaluation

  • Across few-shot and leave-one-out settings, the pre-trained models generally outperform supervised methods; YaTC performs comparably to NetMamba across datasets.Fine-tuning uses 10%, 40%, 70%, or 100% of the full training set, and leave-one-out excludes the fine-tuning dataset from pre-training data.Source: E. Few-Shot Evaluation · p. 10

VIII. CONCLUSION AND FUTURE WORK

  • The conclusion identifies specialized GPU hardware dependence as a limitation on deploying the current NetMamba implementation on real-world network devices.It also names quality-of-service and network-performance prediction as broader potential tasks.Source: §VIII · p. 10

Related papers

Loading 2405.11449v4…