Feature Distribution Shift Mitigation with Contrastive Pretraining for Intrusion Detection

TL;DR

We address feature distribution shift in network intrusion detection and propose SwapCon, a pretrained model that uses contrastive pretraining to learn shift-invariant representations and refines them during finetuning. The approach is evaluated on Kyoto2006+ with IID, NEAR, and FAR splits, showing that properly sized pretraining improves robustness by over 8% and that numerical embedding strategies further boost performance. SwapCon outperforms XGBoost and KNN under distribution shift, and results highlight the critical role of embedding choices and model size in pretraining gains.

Abstract

In recent years, there has been a growing interest in using Machine Learning (ML), especially Deep Learning (DL) to solve Network Intrusion Detection (NID) problems. However, the feature distribution shift problem remains a difficulty, because the change in features' distributions over time negatively impacts the model's performance. As one promising solution, model pretraining has emerged as a novel training paradigm, which brings robustness against feature distribution shift and has proven to be successful in Computer Vision (CV) and Natural Language Processing (NLP). To verify whether this paradigm is beneficial for NID problem, we propose SwapCon, a ML model in the context of NID, which compresses shift-invariant feature information during the pretraining stage and refines during the finetuning stage. We exemplify the evidence of feature distribution shift using the Kyoto2006+ dataset. We demonstrate how pretraining a model with the proper size can increase robustness against feature distribution shifts by over 8%. Moreover, we show how an adequate numerical embedding strategy also enhances the performance of pretrained models. Further experiments show that the proposed SwapCon model also outperforms eXtreme Gradient Boosting (XGBoost) and K-Nearest Neighbor (KNN) based models by a large margin.

Figures (8)

• Figure 1: ML model training pipeline with pretraining.
• Figure 2: Visualization of the shapes of a feature distribution shift over time in the Kyoto2006+ dataset kyoto2006. In each year, the horizontal expansion of the feature plot shows its probability density. The Y-axis means the percentage value that ranges from 0 to 100 which relates to the feature values.
• Figure 3: Splits of IID, NEAR, and FAR based on 10 years of data. Note that the training set is sampled from the IID split.
• Figure 4: SwapCon pretraining pipeline. Colors represent different features.
• Figure 5: The finetuning pipeline of SwapCon.