Mean Masked Autoencoder with Flow-Mixing for Encrypted Traffic Classification
Authors: Xiao Liu, Xiaowei Fu, Fuxiang Huang, Lei Zhang
Organizations: School of Microelectronics and Communication Engineering, Chongqing University, Chongqing 400044, China · School of Data Science, Lingnan University, Hong Kong, China
Network traffic classification using self-supervised pre-training models based on Masked Autoencoders (MAE) has demonstrated a huge potential. However, existing methods are confined to isolated byte-level reconstruction of individual flows, lacking adequate perception of the multi-granularity contextual relationship in traffic. To address this limitation, we propose Mean MAE (MMAE), a teacher-student MAE paradigm with flow mixing strategy for building encrypted traffic pre-training model. MMAE employs a self-distillation mechanism for teacher-student interaction, where the teacher provides unmasked flow-level semantic supervision to advance the student from local byte reconstruction to multi-granularity comprehension. To break the information bottleneck in individual flows, we introduce a dynamic Flow Mixing (FlowMix) strategy to replace traditional random masking mechanism. By constructing challenging cross-flow mixed samples with interferences, it compels the model to learn discriminative representations from distorted tokens. Furthermore, we design a Packet-importance aware Mask Predictor (PMP) equipped with an attention bias mechanism that leverages packet-level side-channel statistics to dynamically mask tokens with high semantic density. Numerous experiments on a number of datasets covering encrypted applications, malware, and attack traffic demonstrate that MMAE achieves state-of-the-art performance. The code is available at https://github.com/lx6c78/MMAE
Figures & tables
Fig. 1: Comparison of reconstruction loss during pre-training between MAE and its variant with Self-Distillation strategy (i.e., MAE+SD). By introducing flow-level semantics via SD, the loss is consistently lower.
Fig. 2: Comparison of pre-training paradigms. (a) Standard MAE with a random masking strategy in an isolated single flow. (b) MMAE (ours): A novel teacher-student pre-training architecture with a flow-mixing strategy, where the teacher is a copy of student and interacted via self-distillation strategy under the unmasked flow-level semantic supervision from the teacher.
Fig. 3: Comparison of hierarchical semantic extraction capabilities between MAE and the proposed MMAE. MAE relies solely on a single flow and random masking for reconstruction. Our MMAE incorporates cross-flow mixing and packet-level side-channel statistical features. Within the self-distillation teacher-student architecture, flow-level reference features are used as supervision for reconstruction and mask prediction tasks.
Fig. 4: Flowchar of the proposed MMAE, which mainly includes FlowMix and teacher-student based twin MAEs. MMAE incorporates a traffic pre-processing unit, a statistics-based flow matcher and a cross-flow mixing unit (i.e., FlowMix). Within the self-distillation teacher-student twin MAE architecture, flow-level reference features are exploited as supervision for reconstruction and mask prediction tasks.
Category
Feature Name
Description
Time (6 dims)
F-Duration
Total time of a flow from start to end
F-Time
Timestamp of the first packet in a flow
P-inter-min
Shortest time interval between packets
P-inter-max
longest time interval between packets
P-inter-avg
Average deviation of packet intervals
P-inter-std
Standard deviation of packet intervals
TABLE I: Summary of Extracted Side-Channel Statistical Features
Fig. 5: Architecture of the Packet-importance aware Mask Predictor (PMP). PMP leverages packet-level side-channel priors to generate a low-rank attention bias via a token-aware gating mechanism. This bias dynamically modulates self-attention to identify and mask challenging, semantically dense regions.
Method
CrossPlatform(Android)
CrossPlatform(iOS)
CICIoT2022
AC
PR
RC
F1
AC
PR
RC
F1
AC
PR
RC
F1
AppScanner
0.1626
0.1646
0.1456
0.1413
0.1718
0.1400
0.1440
0.1283
0.7556
0.8093
0.7244
0.6938
FlowPrint
0.8739
0.8941
0.8739
0.8700
0.8712
0.8687
0.8712
0.8603
0.5820
0.4164
0.5820
0.4643
FS-Net
0.0147
0.0023
0.0147
0.0034
0.0293
0.0014
0.0293
0.0025
0.5747
0.3800
0.5747
0.4216
ET-BERT
0.9386
0.9451
0.9386
0.9401
0.9105
0.8809
0.9105
0.8850
0.9937
0.9938
0.9937
0.9937
YaTC
0.9042
0.9081
0.9042
0.9042
0.9310
0.9307
0.9310
0.9295
0.9959
0.9959
0.9959
0.9959
TABLE II: Comparisons with previous SoTA methods on CrossPlatform(Android), CrossPlatform(iOS), and CICIoT2022 datasets
Method
ISCXTor2016
ISCXVPN2016
USTC-TFC2016
AC
PR
RC
F1
AC
PR
RC
F1
AC
PR
RC
F1
AppScanner
0.4034
0.2850
0.2149
0.2113
0.7643
0.8047
0.7045
0.7256
0.6998
0.8591
0.6062
0.6633
FlowPrint
0.1316
0.0173
0.1316
0.0306
0.9666
0.9733
0.9666
0.9681
0.7992
0.7745
0.7992
0.7755
FS-Net
0.7020
0.7010
0.7020
0.6999
0.7023
0.7487
0.7023
0.6660
0.4381
0.2011
0.4381
0.2672
ET-BERT
0.9980
0.9981
0.9980
0.9980
0.9566
0.9566
0.9566
0.9565
0.9910
0.9911
0.9910
0.9910
YaTC
0.9959
0.9959
0.9959
0.9959
0.9819
0.9820
0.9819
0.9819
0.9947
0.9749
0.9747
0.9734
TABLE III: Comparison with previous SoTA methods on ISCXTor2016, ISCXVPN2016, and USTC-TFC2016 datasets
Method
CSTNET-TLS
AC
PR
RC
F1
AppScanner
0.7441
0.7232
0.6963
0.7023
FS-Net
0.7814
0.7670
0.7316
0.7311
ET-BERT
0.7993
0.7832
0.7689
0.7700
YaTC
0.8391
0.8364
0.8101
0.8140
TrafficFormer
0.7982
0.7883
0.7736
0.7704
TABLE IV: Comparison with SoTA methods on the CSTNET-TLS dataset
Method
CrossPlatform(Android)
CrossPlatform(iOS)
CSTNET-TLS
AC
PR
RC
F1
AC
PR
RC
F1
AC
PR
RC
F1
Baseline
0.9852
0.9855
0.9852
0.9850
0.9838
0.9836
0.9836
0.9829
0.9271
0.9284
0.9271
0.9267
+ SD
0.9856
0.9858
0.9856
0.9854
0.9843
0.9855
0.9848
0.9840
0.9284
0.9302
0.9284
0.9280
+ DMM
0.9878
0.9876
0.9878
0.9873
0.9881
0.9872
0.9881
0.9866
0.9336
0.9356
0.9336
0.9332
+ SFM
0.9882
0.9882
0.9882
0.9880
0.9883
0.9857
0.9883
0.9866
0.9353
0.9361
0.9353
0.9346
+ PMP
0.9897
0.9898
0.9897
0.9893
0.9900
0.9872
0.9900
0.9881
0.9435
0.9459
0.9435
0.9434
TABLE V: Ablation experiments of each component in MMAE
Fig. 6: Comparison between MAE and MAE+SD on three datasets. The charts show 10 classes exhibiting prominent F1-score enhancements.
Fig. 7: t-SNE visualization of hidden state representations across different training stages on the CSTNET-TLS 1.3 dataset. Distinct colors denote different traffic categories. MMAE effectively learns highly discriminative representations, exhibiting significantly clearer class boundaries compared to the baseline.
Fig. 8: Efficiency evaluation of MMAE and baseline models. (a) Inference throughput, (b) GPU memory consumption across varying batch sizes, and (c) Comprehensive cost-effectiveness comparison at a batch size of 64.
Phase
GPUs
Time
Unit/Granularity
Params
GPU Memory
Pre-training
1
24.66 h
5.918 s/100 steps
8.65 M
18GB
Fine-tuning
1
4.6 h
137.87 s / epoch
7.72 M
3.6GB
Inference
1
-
5343 samples/sec
-
787MB
TABLE VI: Computational Efficiency Across Different Phases