Enhancing long-tailed software vulnerability type classification via adaptive data augmentation and prompt tuning

Abstract

Software vulnerability type classification (SVTC) is essential for efficient and targeted remediation of vulnerabilities. With the rapid increase in software vulnerabilities, the demand for automated SVTC approaches is becoming increasingly critical. However, the SVTC is significantly affected by the long-tailed issues, where the distribution of vulnerability types is highly unbalanced. Specifically, a small number of head classes contain a large volume of samples, while a substantial portion of tail classes consists of only a limited number of samples. This imbalance poses a significant challenge to the classification accuracy of existing approaches. To alleviate these challenges, we propose an innovative approach VulTC-LTPF, which integrates prompt tuning with long-tailed learning to enhance the effectiveness of SVTC. Within VulTC-LTPF, an adaptive error-rate-based data augmentation strategy is developed. This strategy allows the SVTC model to dynamically augment data for tail classes types with limited sample size during training, thereby mitigating the impact of the long-tailed problem. Furthermore, VulTC-LTPF employs a hybrid prompt tuning strategy, aligning the training process more closely with pre-training, which enhances adaptability to downstream tasks. Unlike existing approaches that rely solely on either vulnerability description or source code, VulTC-LTPF leverages both sources of information. By incorporating a combination of hard and soft prompts, it facilitates a more comprehensive and effective classification strategy. Experimental results demonstrate that VulTC-LTPF achieves substantial performance improvements over four state-of-the-art SVTC baselines, with gains ranging from 26.1% to 55.1% in MCC. Ablation studies further validate the effectiveness of the adaptive data augmentation, prompt tuning, the integration of two types of vulnerability information, and the use of hybrid prompts. These findings highlight that VulTC-LTPF represents a promising advancement in the field of SVTC, offering significant potential for further progress in addressing software vulnerability type classification challenges.