Abstract:To perform fine-grained vulnerability detection, an ideal model must determine whether software contains vulnerabilities and identify the type of vulnerability (i.e., perform vulnerability classification). A series of deep learning models have demonstrated strong overall performance in vulnerability classification tasks. However, a severe data imbalance exists across different vulnerability types. Many vulnerability types are represented by only a small number of samples (referred to as few-shot types in this study), resulting in poor classification performance and generalization for these few-shot types. To enhance classification performance for these types, VulFewShot is proposed. This contrastive learning-based vulnerability classification framework assigns more weight to few-shot types by bringing samples of the same type closer together while keeping samples from different types further apart. Experimental results show that VulFewShot improves classification performance across all vulnerability types. The smaller the number of samples for a given type, the more significant the improvement. Therefore, VulFewShot improves classification performance for vulnerabilities with limited samples and mitigates the impact of sample size on the learning process.