CMAE-MTC: A Contextual Masked AutoEncoder Based Multi-Level Traffic Classifier
Zecheng Yuan, Min Luo, Cong Peng, Qin Liu · 2024
Traffic classification is vital for network management, ensuring efficient resource allocation, network security, and quality of service. Due to the increasing complexity and anonymity of network traffic, traditional deep learning methods of traffic classification have exposed the following limitations on this critical task. First, traditional methods tend to consider the whole raw packet data as input of the model, ignoring the importance of a well-formed presentation. Second, direct application of the simple models without targeted improvement cannot deeply capture the feature of the traffic flow, especially those encrypted. Last but not least, supervised learning of traditional methods requires a much higher cost to learn for specific scenarios, resulting from the heavy dependence on labels. To break above limitations, we propose a classifier called CMAE-MTC, which involves a well-designed multi-level presentation matrix of traffic flows, reflecting the association between traffic flows and packets, headers and payloads. Meanwhile, our method introduces an improved masked autoencoder paradigm with a latent contextual regressor for self-supervised learning instead of supervised learning. At last, we replace the naive Vision Transformer in the fine-tuning stage with a multi-level attention module, forcing the model to capture the features from not only the small patches of headers and payloads but also the overall packets and flows. We validate the performance of our model on four real-world available encrypted traffic datasets, ISCXVPN, ISCXTor, USTC-TFC, and CICIoT. Our experimental results demonstrate that our proposed method outperforms state-of-the-art methods for traffic classification tasks.