EN-BERT: A Transformer-based Model for Encrypted Attack Traffic Detection with Pre-training and Fine-tuning Phases

Encrypted attacks represented by ransomware and APT are becoming increasingly complex, posing a huge threat to cyberspace. As a result, encrypted attack traffic detection is imperative. Traditional encrypted attack traffic detection methods face challenges which include feature extraction limitations, dataset imbalance, and poor generalization capabilities. To address these issues, this paper proposes EN-Bert, a Transformer-based model with both pre-training and fine-tuning phases. In the pre-training stage, the encrypted traffic dataset is first processed using Token serialization for traffic shunting, segmentation, and feature extraction. Then the model is pre-trained on two tasks: Masked Flow Model MFM and Same-origin Flow Prediction SFP to uncover the contextual relationships between traffic flows. In the fine-tuning stage, this paper addresses the imbalance issues through data enhancement. From the model’s perspective, weighted cross-entropy loss and K-L divergence are designed in this phase to optimize the model’s performance and enhance its generalization ability. In the experimental section, with the CIC-IDS-2017 dataset and a self-collected encrypted DOS attack traffic dataset, comparative and ablation experiments demonstrate that EN-Bert model proficiently addresses challenges related to dataset imbalance and poor model generalization, proving to be an effective and reliable approach for encrypted attack traffic detection.