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

TL;DR

NetMamba+ introduces an efficient, multimodal traffic classification framework using advanced attention mechanisms and label-aware fine-tuning, significantly improving accuracy, speed, and few-shot learning in real-world network environments.

Contribution

It is the first to adapt Mamba architecture for network traffic classification, combining efficiency, multimodal representation, and label-aware strategies for superior performance.

Findings

Up to 6.44% improvement in F1 score over baselines

1.7x higher inference throughput than state-of-the-art

Enhanced few-shot learning capabilities

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.