A Novel Self-Supervised Framework Based on Masked Autoencoder for Traffic Classification

• Published in: IEEE/ACM transactions on networking Vol. 32; no. 3; pp. 2012 - 2025
• Main Authors: Zhao, Ruijie, Zhan, Mingwei, Deng, Xianwen, Li, Fangqi, Wang, Yanhao, Wang, Yijun, Gui, Guan, Xue, Zhi
• Format: Journal Article
• Language: English
• Published: New York IEEE 01.06.2024; The Institute of Electrical and Electronics Engineers, Inc. (IEEE)
• Subjects: Algorithms; Analytical models; Classification; Classifiers; Data models; Deep learning; Feature extraction; Machine learning; Network security; Representations; Self-supervised learning; Task analysis; traffic classification; Traffic information; Training; Transformers
• ISSN: 1063-6692; 1558-2566
• DOI: 10.1109/TNET.2023.3335253

Traffic classification is a critical task in network security and management. Recent research has demonstrated the effectiveness of the deep learning-based traffic classification method. However, the following limitations remain: (1) the traffic representation is simply generated from raw packet bytes, resulting in the absence of important information; (2) the model structure of directly applying deep learning algorithms does not take traffic characteristics into account; and (3) scenario-specific classifier training usually requires a labor-intensive and time-consuming process to label data. In this paper, we introduce a masked autoencoder (MAE) based traffic transformer with multi-level flow representation to tackle these problems. To model raw traffic data, we design a formatted traffic representation matrix with hierarchical flow information. After that, we develop an efficient Traffic Transformer, in which packet-level and flow-level attention mechanisms implement more efficient feature extraction with lower complexity. At last, we utilize MAE paradigm to pre-train our classifier with a large amount of unlabeled data, and perform fine-tuning with a few labeled data for a series of traffic classification tasks. Experiment findings reveal that our method outperforms state-of-the-art methods on five real-world traffic datasets by a large margin. The code is available at https://github.com/NSSL-SJTU/YaTC .