DKI-GMDM: domain-knowledge integrated and imbalance-guided diffusion model for network intrusion detection

Abstract

The deployment of learning-based Network Intrusion Detection Systems (NIDS) is severely hindered by the class imbalance commonly observed in network traffic, where malicious activities account for only a small fraction of the data. Although generative data augmentation provides a promising way to alleviate this problem, existing methods such as generative adversarial network (GAN)-based oversampling and standard diffusion-based tabular generators may fail to sufficiently capture minority-attack distributions or may generate samples that violate important network-domain constraints. To address these challenges, we propose the Domain-Knowledge Integrated and Imbalance-Guided Diffusion Model (DKI-GMDM), a diffusion-based augmentation framework designed for imbalanced NIDS data. Specifically, we introduce a domain-knowledge constraint (DKC) module that incorporates protocol-related validity rules and statistical dependencies into the denoising process, thereby improving the semantic validity of generated traffic samples under a defined constraint set. We further develop an imbalance-aware guidance mechanism that adaptively adjusts the sampling trajectory according to class frequency, encouraging the generation of more class-specific samples for minority attack categories. Extensive experiments on selected class subsets of the UNSW-NB15 and CICIDS2017 benchmarks show that DKI-GMDM improves downstream detection performance over the selected baselines under the evaluated subset-based settings. On the selected UNSW-NB15 subset, DKI-GMDM improves the F1-scores of the highly imbalanced Shellcode and Worms classes by 7.71 and 6.52 percentage points, respectively, compared with Modelling Tabular Data with Diffusion Models (TabDDPM). Overall, these findings demonstrate the effectiveness of DKI-GMDM in alleviating the impact of class imbalance in NIDS training under the evaluated benchmark settings.