Applying Contrastive Learning to Code Vulnerability Type Classification (2024)

Abstract: Vulnerability classification is a crucial task in software security analysis, essential for identifying and mitigating potential security risks. Learning-based methods often perform poorly due to the long-tail distribution of vulnerability classification datasets. Recent approaches try to address the problem but treat each CWE class in isolation, ignoring their relationships. This results in non-scalable code vector representations, causing significant performance drops when handling complex real-world vulnerabilities. We propose a hierarchical contrastive learning framework for code vulnerability type classification to bring vector representations of related CWEs closer together. To address the issue of class collapse and enhance model robustness, we mix self-supervised contrastive learning loss into our loss function. Additionally, we employ max-pooling to enable the model to handle longer vulnerability code inputs. Extensive experiments demonstrate that our proposed framework outperforms state-of-the-art methods by 2.97% - 17.90% on accuracy and 0.98% - 22.27% on weighted-F1, with even better performance on higher-quality datasets. We also utilize an ablation study to prove each component’s contribution. These findings underscore the potential and advantages of our approach in the multi-class vulnerability classification task.

## introduction nowadays, software is ubiquitous in people's lives, permeating nearly every aspect of daily activities. however, as software systems continue to grow in scale and complexity, the types of vulnerabilities are becoming more diverse @xcite . common weakness enumeration (cwe) is a comprehensive vocabulary for describing and classifying weaknesses in software. maintained by mitre foot_0 , its purpose is to aid security experts in identifying and addressing software security vulnerabilities. for example, the standardized terminology provided by cwe enables security experts to accurately report discovered vulnerabilities, offer specific repair recommendations, and determine the severity of vulnerabilities, which helps prioritize and address high-risk vulnerabilities first. classifying vulnerability code according to cwe is a challenging task that requires security experts to manually analyze the code and identify the specific types of vulnerabilities. while manual classification provides detailed analysis, it also has several significant drawbacks. first, it relies on the knowledge and experience of security experts. second, manual analysis is typically slow, making it difficult to meet the demands of large-scale code reviews. for example, among the 28,902 common vulnerabilities and exposures (cve) entries newly published in 2023 in the national vulnerability database (nvd) foot_1 , 4,113 cases still have unidentified types. this underscores the need for automated tools to quickly classify potential vulnerabilities. existing approaches. many methods have been proposed to classify software vulnerabilities. traditional methods include static analysis @xcite or dynamic analysis @xcite @xcite . however, they still face issues such as manual feature design, high false positives, and incomplete coverage. given the significant success of deep learning methods in image recognition and natural language processing(nlp) tasks @xcite , particularly in terms of automated feature extraction, handling complex patterns, and understanding code semantics, researchers are now exploring their application to the task of classifying code vulnerability types. some researchers applied recently, several studies have been proposed @xcite @xcite to address this issue by leveraging learning-based methods to predict the types of vulnerabilities (i.e., cwe-id). according to figure 1 ## methodology in this section, we present the details of our novel framework that uses hierarchical supervised contrastive learning to enhance the performance of a transformer-based model in vulnerability classification. the overview of the proposed method is illustrated in figure 3 and the model architecture is detailed in section 3.3. ## conclusion we propose a multi-class vulnerability type classification framework based on hierarchical contrastive learning to bring the vector representations of highly related cwes closer together. additionally, we employ geometric spread to address the class collapse problem and max-pooling to receive a longer input length, thereby enhancing the model's robustness. the results demonstrate that our hierarchical contrastive learning framework outperforms current baseline methods across various metrics. this improvement is primarily due to advancements in code representation provided by our method and these better quality vectors can be used in subsequent works.