Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

NetMamba+: A Framework of Pre-trained Models for Efficient and Accurate Network Traffic Classification

Tongze Wang, Xiaohui Xie, Wenduo Wang, Chuyi Wang, Jinzhou Liu, Boyan Huang, Yannan Hu, Youjian Zhao, Yong Cui

TL;DR

As the first framework to adapt Mamba architecture for network traffic classification, NetMamba+ opens new possibilities for efficient and accurate traffic analysis in complex network environments.

Abstract

With the rapid growth of encrypted network traffic, effective traffic classification has become essential for network security and quality of service management. Current machine learning and deep learning approaches for traffic classification face three critical challenges: computational inefficiency of Transformer architectures, inadequate traffic representations with loss of crucial byte-level features while retaining detrimental biases, and poor handling of long-tail distributions in real-world data. We propose NetMamba+, a framework that addresses these challenges through three key innovations: (1) an efficient architecture considering Mamba and Flash Attention mechanisms, (2) a multimodal traffic representation scheme that preserves essential traffic information while eliminating biases, and (3) a label distribution-aware fine-tuning strategy. Evaluation experiments on massive datasets encompassing four main classification tasks showcase NetMamba+'s superior classification performance compared to state-of-the-art baselines, with improvements of up to 6.44\% in F1 score. Moreover, NetMamba+ demonstrates excellent efficiency, achieving 1.7x higher inference throughput than the best baseline while maintaining comparably low memory usage. Furthermore, NetMamba+ exhibits superior few-shot learning abilities, achieving better classification performance with fewer labeled data. Additionally, we implement an online traffic classification system that demonstrates robust real-world performance with a throughput of 261.87 Mb/s. As the first framework to adapt Mamba architecture for network traffic classification, NetMamba+ opens new possibilities for efficient and accurate traffic analysis in complex network environments.

NetMamba+: A Framework of Pre-trained Models for Efficient and Accurate Network Traffic Classification

TL;DR

As the first framework to adapt Mamba architecture for network traffic classification, NetMamba+ opens new possibilities for efficient and accurate traffic analysis in complex network environments.

Abstract

With the rapid growth of encrypted network traffic, effective traffic classification has become essential for network security and quality of service management. Current machine learning and deep learning approaches for traffic classification face three critical challenges: computational inefficiency of Transformer architectures, inadequate traffic representations with loss of crucial byte-level features while retaining detrimental biases, and poor handling of long-tail distributions in real-world data. We propose NetMamba+, a framework that addresses these challenges through three key innovations: (1) an efficient architecture considering Mamba and Flash Attention mechanisms, (2) a multimodal traffic representation scheme that preserves essential traffic information while eliminating biases, and (3) a label distribution-aware fine-tuning strategy. Evaluation experiments on massive datasets encompassing four main classification tasks showcase NetMamba+'s superior classification performance compared to state-of-the-art baselines, with improvements of up to 6.44\% in F1 score. Moreover, NetMamba+ demonstrates excellent efficiency, achieving 1.7x higher inference throughput than the best baseline while maintaining comparably low memory usage. Furthermore, NetMamba+ exhibits superior few-shot learning abilities, achieving better classification performance with fewer labeled data. Additionally, we implement an online traffic classification system that demonstrates robust real-world performance with a throughput of 261.87 Mb/s. As the first framework to adapt Mamba architecture for network traffic classification, NetMamba+ opens new possibilities for efficient and accurate traffic analysis in complex network environments.
Paper Structure (59 sections, 20 equations, 15 figures, 9 tables, 1 algorithm)

This paper contains 59 sections, 20 equations, 15 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Transformer-based Traffic Classification
  4. Mamba-based Representation Learning
  5. Traffic Representation Schemes
  6. Deep Long-Tailed Learning
  7. Preliminaries
  8. State Space Models
  9. Discretization
  10. Selective Scan
  11. NetMamba+ Framework
  12. Traffic Representation Phase
  13. Pre-training Phase
  14. Fine-tuning Phase
  15. Traffic Representation
  16. ...and 44 more sections

Figures (15)

  • Figure 1: Overview of NetMamba+ Framework.
  • Figure 2: The Prototype of NetMamba-based Online Traffic Classification System.
  • Figure 3: Comparison of Pre-trained Models on Huawei-VPN Dataset
  • Figure 4: The Inference Speed and GPU Memory Comparison
  • Figure 5: The Inference Efficiency Comparison on Fine-tuning Batch Size
  • ...and 10 more figures