VulScribeR: Exploring RAG-based Vulnerability Augmentation with LLMs

TL;DR

VulScribeR leverages large language models with retrieval-augmented generation to effectively augment vulnerable code datasets, significantly improving vulnerability detection performance across multiple datasets and models.

Contribution

The paper introduces VulScribeR, a novel LLM-based data augmentation method using prompt templates and three strategies to generate vulnerable code, outperforming state-of-the-art methods.

Findings

VulScribeR improves F1-score by up to 54% over SOTA methods.

The approach generates large-scale vulnerable datasets cost-effectively.

VulScribeR is effective across multiple datasets and vulnerability types.

Abstract

Detecting vulnerabilities is vital for software security, yet deep learning-based vulnerability detectors (DLVD) face a data shortage, which limits their effectiveness. Data augmentation can potentially alleviate the data shortage, but augmenting vulnerable code is challenging and requires a generative solution that maintains vulnerability. Previous works have only focused on generating samples that contain single statements or specific types of vulnerabilities. Recently, large language models (LLMs) have been used to solve various code generation and comprehension tasks with inspiring results, especially when fused with retrieval augmented generation (RAG). Therefore, we propose VulScribeR, a novel LLM-based solution that leverages carefully curated prompt templates to augment vulnerable datasets. More specifically, we explore three strategies to augment both single and multi-statement vulnerabilities, with LLMs, namely Mutation, Injection, and Extension. Our extensive evaluation across four vulnerability datasets and DLVD models, using three LLMs, show that our approach beats two SOTA methods Vulgen and VGX, and Random Oversampling (ROS) by 27.48%, 27.93%, and 15.41% in f1-score with 5K generated vulnerable samples on average, and 53.84%, 54.10%, 69.90%, and 40.93% with 15K generated vulnerable samples. Our approach demonstrates its feasibility for large-scale data augmentation by generating 1K samples at as cheap as US$ 1.88.