Improving Long-Tail Vulnerability Detection Through Data Augmentation Based on Large Language Models
The ability of automatic vulnerability detection models largely depends on the dataset used for training. However, annotating these datasets is costly and time-consuming, leading to a scarcity of labeled samples, particularly for diverse CWE types. This scarcity results in a pronounced long-tail issue, where less common types are underrepresented, thus diminishing the model's effectiveness in detecting them. In this paper, we address this challenge by employing large language models' (LLMs) generative and reasoning capabilities to create the necessary training samples for detecting less common vulnerability types. Specifically, we use GPT-4 to generate targeted samples of these types. After a well-defined self-filtering process, these samples are incorporated into the training of detection models. Extensive experiments show that our approach significantly enhances vulnerability detection capabilities, especially for long-tail vulnerabilities, across a variety of detection models, including both traditional deep learning and modern LLM-based detection models. Further, our comparative tests on unseen projects indicate that models trained with our generated data can identify more real-world vulnerabilities than traditional methods, proving the practicality and generalizability of our approach in real settings. The code and the generated dataset are publicly available at https://github.com/LuckyDengXiao/LERT.
Authors
- Rui Xie (ORCID: https://orcid.org/0000-0002-1756-7746)
- Shikun Zhang (ORCID: https://orcid.org/0000-0002-6628-4881)
- Fuyao Duan
- Wei Ye (ORCID: https://orcid.org/0000-0002-9331-4716)
- Xiao Deng
Institutions
- Peking University (CN)
Publication Details
- Published
- 2024-10-06
- DOI
- https://doi.org/10.1109/icsme58944.2024.00033
- Citations
- 1
- Primary Topic
- Network Security and Intrusion Detection
- Type
- article
- Field-Weighted Citation Impact
- 0.35