Smart Cuts: Enhance Active Learning for Vulnerability Detection by Pruning Hard-to-Learn Data

Abstract

Vulnerability detection is crucial for identifying security weaknesses in software systems. However, training effective machine learning models for this task is often constrained by the high cost and expertise required for data annotation. Active learning is a promising approach to mitigate this challenge by intelligently selecting the most informative data points for labeling. This paper proposes a novel method to significantly enhance the active learning process by using dataset maps. Our approach systematically identifies samples that are hard-to-learn for a model and integrates this information to create a more sophisticated sample selection strategy. Unlike traditional active learning methods that focus primarily on model uncertainty, our strategy enriches the selection process by considering learning difficulty, allowing the active learner to more effectively pinpoint truly informative examples. The experimental results show that our approach can improve F1 score over random selection by 61.54% (DeepGini) and 45.91% (K-Means) and outperforms standard active learning by 8.23% (DeepGini) and 32.65% (K-Means) for CodeBERT on the Big-Vul dataset, demonstrating the effectiveness of integrating dataset maps for optimizing sample selection in vulnerability detection. Furthermore, our approach also enhances model robustness, improves sample selection by filtering hard-to-learn data, and stabilizes active learning performance across iterations. By analyzing the characteristics of these outliers, we provide insights for future improvements in dataset construction, making vulnerability detection more reliable and cost-effective.