🤖 AI

Taming Data Challenges in ML-based Security Tasks Using Generative AI

This paper demonstrates that augmenting training datasets with synthetic data generated by Generative AI, including a novel method called Nimai, can significantly improve the performance and adaptability of machine learning-based security classifiers, particularly in data-constrained settings, while also identifying specific data characteristics that may hinder such improvements.

Original authors: Shravya Kanchi, Neal Mangaokar, Aravind Cheruvu, Sifat Muhammad Abdullah, Shirin Nilizadeh, Atul Prakash, Bimal Viswanath

Published 2026-05-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Shravya Kanchi, Neal Mangaokar, Aravind Cheruvu, Sifat Muhammad Abdullah, Shirin Nilizadeh, Atul Prakash, Bimal Viswanath

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a security guard (a computer program) how to spot a thief in a crowded city. Usually, the best way to teach this guard is to show them thousands of photos of real thieves and real citizens. But in the world of cybersecurity, there's a big problem: we don't have enough photos of the thieves.

Thieves are rare, they hide well, and sometimes the photos we do have are blurry or labeled incorrectly. This makes the security guard bad at their job.

This paper asks a simple question: Can we use "Generative AI" (GenAI) to create fake photos of thieves to help train the guard?

Here is the breakdown of their findings, using simple analogies:

1. The Problem: The "Empty Photo Album"

The researchers looked at 35 different cybersecurity studies and found that almost all of them struggled with the same data problems:

Not enough examples: There are very few examples of attacks compared to normal activity.
Biased views: The few attack examples we have might only show one specific type of thief, missing others.
Noisy photos: Sometimes the labels are wrong (a photo of a citizen is labeled as a thief).
Changing faces: Thieves change their appearance over time (this is called "concept drift"), so the guard's old training becomes useless.

2. The Solution: The "AI Artist"

The team tried using advanced AI tools (GenAI) to act like an artist. Instead of just showing the guard real photos, they asked the AI to paint new, realistic-looking photos of thieves based on the few real ones it had seen. They then added these fake photos to the training album to see if the guard got smarter.

They tested this on 7 different security tasks, ranging from spotting malware (computer viruses) to detecting hijacked internet routes.

3. The Results: When the Artist Shines

The results were a mix of "Wow!" and "Not so fast."

The Success Stories (The "Wow!" moments):

The "Super-Training" Effect: In tasks where the data was scarce (like the BGP hijacking task, which had only about 180 training samples), the AI-generated data was a game-changer. One specific method they invented, called Nimai, improved the security guard's accuracy by 32.6%. That's like going from a guard who catches 6 out of 10 thieves to one who catches 9.
Fixing the "Blurry" Photos: They found that the AI didn't just copy-paste; it actually helped smooth out the "noise" in the data. It made the patterns of the thieves clearer, helping the guard generalize better.
The "Shape-Shifter" Problem: When thieves changed their tactics (concept drift), the AI could quickly generate new "fake" examples of the new tactics. This allowed the security system to adapt almost instantly with very little new human labeling.

The Failures (The "Not so fast" moments):

The "Overcrowded Room": In some tasks, the "thieves" and "citizens" looked so similar (overlapping data) that the AI couldn't tell the difference. It just made more confusing photos, and the guard didn't get any better.
The "Too Big to Fit" Problem: Some of the existing AI tools were like giant trucks trying to drive down a tiny alley. They were too complex for the specific security data and crashed or failed to train.
The "Noisy Label" Trap: If the original data had too many mistakes (e.g., labeling a good file as bad), the AI just learned to make more mistakes. It couldn't fix a broken foundation.

4. The New Tool: "Nimai"

The researchers realized that existing AI tools were a bit like a blindfolded painter—they could only paint based on a broad category (e.g., "paint a thief"). They couldn't say, "Paint a thief that looks exactly like this specific person I'm holding."

So, they built a new tool called Nimai.

The Analogy: Imagine a sculptor who doesn't just make a generic statue of a person. Instead, they take a real person, look at them, and create a new statue that is very similar but has slight, controlled variations. This allows the security guard to see every tiny nuance of the threat.
The Result: Nimai was the most successful tool, especially in the hardest, most data-starved situations.

5. The Big Takeaway

The paper concludes that Generative AI is a powerful tool for cybersecurity, but it's not magic.

It works wonders when you have very little data or when the data is unbalanced. It can fill in the gaps and help the system learn faster.
It struggles when the data is too messy, the categories overlap too much, or the AI tools are too heavy for the job.

The authors suggest that future security tools need to be smarter about how they generate data, using controlled methods like Nimai, rather than just churning out random fake data. They also note that this approach could help security systems adapt quickly when hackers change their tactics, without needing to hire armies of humans to label new data.

In short: They taught the security guard by having an AI paint new pictures of the bad guys. It worked brilliantly when the pictures were scarce, but it got confused when the bad guys and good guys looked too much alike.

Technical Summary: Taming Data Challenges in ML-based Security Tasks Using Generative AI

1. Problem Statement

Machine learning (ML) classifiers are widely deployed for cybersecurity tasks, such as malware detection, network intrusion detection, and privacy classification. While algorithmic advancements have been the primary focus of research, the data challenges inherent to security domains have received limited attention, often resulting in performance ceilings. The authors identify five critical data challenges:

Limited or biased training data: Including severe class imbalance, under-represented attack patterns, and insufficient sample sizes due to the high cost of labeling and domain expertise.
High dimensionality: Features often exceed thousands of dimensions, exacerbating the "curse of dimensionality."
Concept drift: Attackers adapt over time, causing the test distribution to diverge from the training distribution post-deployment.
Noisy labels and overlapping class distributions: Common in security data where ground truth is difficult to establish or classes are inherently similar.
Adversarial examples: Though noted, this paper focuses primarily on the first four challenges.

The core research question is whether Generative AI (GenAI) can address these data challenges to improve classifier performance, specifically by augmenting training datasets with synthetic data.

2. Methodology

2.1 Experimental Setup

The study evaluates GenAI across 7 diverse security tasks (3 binary, 4 multi-class) using real-world threat data. These tasks include BGP hijacking detection, Tor-based malware traffic detection, IoT cryptojacking detection, web cookie privacy classification, Windows malware classification (MS-Malware), PE malware classification (BODMAS), and OS fingerprinting (nPrintML).

The authors compare 6 State-of-the-Art (SoTA) GenAI models designed for tabular data against traditional augmentation methods:

GenAI Models: TVAE (VAE), CTAB-GAN+ (cGAN), TabDDPM (Diffusion), TabSyn (VAE+Diffusion), GReaT (LLM), and REaLTabFormer (LLM).
Traditional Baselines: SMOTE (linear interpolation) and MC-CCR (cleaning and re-positioning).

2.2 The Nimai Framework

Recognizing that existing GenAI tools primarily offer class-conditional generation (creating samples based on a label) but lack control over specific regions of the data manifold, the authors propose Nimai.

Architecture: Nimai is a VAE-based model utilizing a discrete latent space (Vector Quantized VAE) combined with a Masked Token Model (MTM) inspired by BERT.
Mechanism: Unlike standard VAEs that use continuous latent spaces (prone to posterior collapse), Nimai discretizes latent tokens. It employs a masking strategy where a subset of latent features is "blanked out" and predicted based on context.
Control Modes:
- Nimai-S (Sample-Conditioned): Generates samples in the vicinity of existing real samples by masking and reconstructing features of a specific input. This targets underrepresented regions within a class.
- Nimai-C (Class-Conditioned): Generates samples based solely on a class label, similar to existing tools.
- Nimai-Hybrid: Combines both approaches for post-deployment drift recovery.

2.3 Evaluation Strategy

The evaluation focuses on data augmentation before deployment (pre-training) and rapid recovery after deployment (post-training).

Metric: Relative gain ( $\Delta G$ ) in macro F-score (or malicious class F-score for specific tasks) compared to training on real data only.
Robustness: Experiments are run over 10 trials with different random seeds. Unreliable measurements (Coefficient of Variation $\ge$ 1) are flagged.

3. Key Results

3.1 Pre-Deployment Performance

Success in Data-Constrained Settings: GenAI significantly improved performance in 5 out of 7 tasks.
- BGP Hijacking: In an extremely data-constrained setting (only ~180 training samples, 17 malicious), Nimai-S achieved a 32.6% performance gain, outperforming all baselines.
- IoT & MS-Malware: Nimai-S and CTAB-GAN+ showed significant gains (up to 17.37% and 13.53% respectively) in highly imbalanced or high-dimensional scenarios.
- BODMAS (Concept Drift): For tasks involving drift, class-conditioned approaches (Nimai-C, TVAE, TabDDPM) generally outperformed sample-conditioned ones, achieving up to 59.42% gain in specific drift months.
Failure Cases:
- Tor, Cookie, nPrintML: GenAI schemes generally failed to improve performance. The authors attribute this to overlapping class distributions, noisy labels, and extreme feature sparsity (up to 98.66% zero values). In these cases, synthetic data often reinforced existing biases or failed to resolve class inseparability.
- Model Instantiation Failures: 5 out of 6 existing GenAI schemes failed to train or generate usable data on at least one task.
  - LLM-based models (GReaT, REaLTabFormer): Failed on 6/7 and 5/7 tasks respectively due to scalability issues with high-dimensional feature vectors (exceeding token limits or GPU memory).
  - TabDDPM: Suffered from mode collapse in the IoT task, failing to generate minority class samples.
  - TVAE/TabSyn: Struggled with high-dimensional or large-scale datasets (e.g., nPrintML, Cookie).

3.2 Post-Deployment (Concept Drift)

Rapid Recovery: The authors demonstrated that a Nimai-Hybrid approach could recover from concept drift with minimal labeling effort.
Efficiency: By using uncertainty sampling to select a small subset of drifted data (as few as 64 labeled samples) and applying sample-conditioned generation, the system achieved a 60.4% performance boost in the worst-affected months of the BODMAS task, outperforming retraining on the small subset alone (INSOMNIA baseline).

3.3 Mechanism of Improvement

Deep dive analysis using LIME and distribution metrics (entropy and skewness) revealed that successful augmentation (e.g., in IoT and BGP tasks) reduced feature skewness and increased entropy, effectively mitigating biases in the real dataset and promoting more symmetric feature distributions, which aids classifier generalization.

4. Key Contributions

Systematic Evaluation: The first study to systematically evaluate 6 SoTA GenAI models for tabular data across 7 diverse security tasks, identifying both opportunities and significant failure modes.
Nimai Framework: Introduction of a novel VAE-based scheme with a discrete latent space and masked token modeling, enabling highly controlled, sample-conditioned data synthesis to address in-class biases.
Drift Recovery: Demonstration of a hybrid GenAI strategy for rapid post-deployment adaptation to concept drift with minimal human labeling.
Failure Analysis: Identification of specific task characteristics (noisy labels, overlapping distributions, sparse features) that hinder GenAI performance, and the discovery that current LLM-based tabular generators are not yet scalable for high-dimensional security data.

5. Significance and Claims

The authors claim that their study provides a balanced perspective on the utility of GenAI in security. They argue that while GenAI is not a "one-size-fits-all" solution and struggles with specific data pathologies (like noise and overlap), it holds significant untapped potential for addressing data scarcity and bias in security tasks.

The paper posits that future progress in security classifiers may be stalled by data limitations rather than algorithmic ones. By developing GenAI tools that can effectively synthesize controlled, high-quality synthetic data, the community can "reignite progress" for security problems where algorithmic innovation has plateaued. The authors emphasize that their work is a step toward developing GenAI tools specifically tailored for the adversarial and data-constrained nature of cybersecurity, rather than simply applying off-the-shelf models.