🤖 AI

LLM-based Vulnerable Code Augmentation: Generate or Refactor?

This paper investigates whether LLM-based data augmentation for vulnerable code is more effective through the generation of new samples or the semantics-preserving refactoring of existing ones, ultimately finding that a hybrid approach best improves the performance of deep learning-based vulnerability classifiers.

Original authors: Dyna Soumhane Ouchebara, Stéphane Dupont

Published 2026-02-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Dyna Soumhane Ouchebara, Stéphane Dupont

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

The Problem: The "Rare Disease" Dilemma in Code

Imagine you are training a highly advanced medical AI to detect rare diseases. To do this well, the AI needs to see thousands of examples of each disease. But there’s a problem: some diseases are incredibly common (like the flu), while others are extremely rare. Because you only have a handful of examples for the rare ones, the AI becomes an expert at spotting the flu but remains "blind" to the rare diseases.

In the world of cybersecurity, software vulnerabilities are exactly like those rare diseases. Most code is "healthy" (safe), but certain types of "bugs" (vulnerabilities) are very rare in datasets. Because AI models learn by seeing patterns, if they don't see enough examples of a specific type of bug, they won't know how to catch it when it actually appears in the real world.

The Solution: The "Digital Mimic" Strategy

The researchers in this paper decided to solve this scarcity by using a powerful AI (an LLM called Qwen) to act as a "Digital Mimic." Instead of waiting years to collect more real-world examples, they used the AI to "manufacture" new ones.

They tested two different ways of doing this:

1. The "Improviser" (Generation-based Augmentation)

Think of this like a professional actor who has seen a few clips of a villain. You don't give them the actual villain; you just describe the villain's personality and show them a few clips for inspiration. Then, you say, "Now, go out and create five brand-new villains from scratch!"

How it works: The AI looks at a few real examples of a bug and then writes entirely new, original pieces of code that contain that same type of bug.

2. The "Makeover Artist" (Refactoring-based Augmentation)

Think of this like taking an existing photo of a person and using Photoshop to change their hair color, clothes, and background, but keeping their face exactly the same. The person is still the same individual, but the photo looks different.

How it works: The AI takes a piece of real, vulnerable code and "refactors" it. It renames variables, adds some useless extra lines of code, or rearranges the logic. The "bug" stays in the exact same spot, but the "look" of the code changes.

The Results: What Happened?

The researchers then took these "fake" (but realistic) examples and fed them to a "student" AI (a classifier called CodeBERT) to see if it got smarter.

The "Improviser" (Generation) was fast: It could whip up new code very quickly, but it was a bit hit-or-miss with the "labels" (sometimes the AI thought it created a bug, but it wasn't quite right).
The "Makeover Artist" (Refactoring) was slower: It took more time to carefully rewrite the code, but the results were very high quality.
The "Hybrid" (The Dream Team): The best results came when they used both methods together. By giving the student AI both brand-new "villains" and "makeovers" of old ones, the AI's ability to detect vulnerabilities jumped significantly—improving by up to 24% in some cases!

The "Aha!" Moment: The Length Factor

The researchers also discovered a "hidden rule" of coding AI: Complexity is the enemy.

They found that the AI is great at catching bugs in short, simple snippets of code, but as the code gets longer and more "wordy," the AI starts to get confused and its accuracy drops. It’s like trying to find a needle in a small box versus finding a needle in a massive haystack.

Summary in a Nutshell

Instead of complaining that there isn't enough "bad code" to train on, these researchers used a smart AI to hallucinate realistic bad code. By mixing brand-new creations with clever rewrites of old ones, they turned a "blind" security AI into a much sharper digital detective.

Technical Summary: LLM-based Vulnerable Code Augmentation: Generate or Refactor?

1. Problem Statement

Deep learning models for software vulnerability detection suffer from a significant data imbalance problem. Vulnerability datasets are often skewed, with a massive disparity between "safe" and "vulnerable" code, and even greater imbalances across different types of vulnerabilities (CWEs). Many critical vulnerability types are under-represented, limiting the ability of classifiers (like CodeBERT) to generalize. While Data Augmentation (DA) is a standard solution in NLP and Computer Vision, it is difficult in software engineering because code must adhere to strict syntactic and semantic constraints; minor edits can break compilation or inadvertently "fix" the vulnerability being studied.

2. Methodology

The authors investigate two distinct LLM-based augmentation strategies using Qwen2.5-Coder-32B to expand the SVEN dataset. They compare these against a baseline classifier, CodeBERT.

Strategy A: Generation-based Augmentation (Few-Shot)
- Process: The LLM is prompted to synthesize entirely new, independent vulnerable functions from scratch.
- Prompting: Uses a few-shot approach where real-world examples of a specific CWE are provided as guidance for naming conventions and complexity.
- Constraints: The model is instructed to produce "industrial-style" code (20–150 lines) that embeds the vulnerability within realistic logic rather than making it the sole purpose of the function.
- Verification: Syntactic correctness is checked via parsers, and label quality is verified using a high-reasoning LLM (GPT-5.1 Thinking).
Strategy B: Refactoring-based Augmentation (Zero-Shot)
- Process: The LLM takes an existing vulnerable function and produces $n$ variants by applying structural changes.
- Techniques: The authors define 18 specific, semantics-preserving refactoring techniques (e.g., renaming, dead code insertion, loop enhancement, and "plus zero" arithmetic).
- Constraints: The prompt strictly forbids the model from repairing the vulnerability. It must preserve the function signature and the underlying logical flaw.
- Verification: Syntactic correctness is checked via parsers, and refactoring quality is verified via LLM-based inspection.

3. Key Contributions

Comparative Study: The paper provides a systematic comparison between generating new samples versus refactoring existing ones for security-critical tasks.
Hybrid Strategy Discovery: The research identifies that combining both methods yields the best results.
Complexity Analysis: The authors provide a novel insight into the relationship between code length (token count) and classification performance, noting that longer code snippets are significantly harder for models to classify correctly.
Rigorous Prompt Engineering: The work provides highly detailed, constrained system prompts designed to maintain the "industrial" feel of code and prevent the accidental "patching" of bugs.

4. Results

Augmentation Efficiency: Both methods successfully increased the dataset size by 581%. Generation was faster (13.38s/sample) than refactoring (58.51s/sample).
Classifier Performance (Macro F1):
- Original Data: 0.62
- Generation Only: 0.64
- Refactoring Only: 0.65
- Hybrid (Both): 0.67 (The best performing approach).
Impact on Minority Classes: The augmentation significantly boosted performance for under-represented CWEs. For example, CWE-190 saw an 18% increase in F1 score when using generation.
Correlation Findings: The study confirmed that F1 scores decrease as the average length of the code snippet increases, suggesting that "verbosity" is a challenge for current vulnerability classifiers.

5. Significance

This research demonstrates that LLMs can be effectively used to mitigate the "long-tail" problem in software security datasets. By moving beyond simple rule-based transformations to sophisticated LLM-driven generation and refactoring, researchers can create more robust, balanced training sets. The finding that a hybrid approach is superior suggests that for optimal model training, one should provide the classifier with both entirely new logical contexts (Generation) and diverse structural variations of known patterns (Refactoring).