Tighter Clusters, Safer Code? Improving Vulnerability Detection with Enhanced Contrastive Loss (2025)

Abstract: Distinguishing vulnerable code from non-vulnerable code is challenging due to high inter-class similarity. Supervised contrastive learning (SCL) improves embedding separation but struggles with intra-class clustering, especially when variations within the same class are subtle. We propose Cluster-Enhanced Supervised Contrastive Loss (CESCL), an extension of SCL with a distance-based regularization term that tightens intra-class clustering while maintaining inter-class separation. Evaluating on CodeBERT and GraphCodeBERT with Binary Cross Entropy (BCE), BCE + SCL, and BCE + CESCL, our method improves F1 score by 1.76% on CodeBERT and 4.1% on GraphCodeBERT, demonstrating its effectiveness in code vulnerability detection and broader applicability to high-similarity classification tasks.

## introduction code vulnerability detection is a cornerstone of software security, particularly as the world undergoes rapid digitization. in domains like finance, healthcare, and government, software vulnerabilities exploited by malicious actors could have catastrophic consequences. the ability to detect such weaknesses efficiently is essential for safeguarding the trust that underpins modern technological systems. beyond protecting sensitive data, robust vulnerability detection forms the backbone of a resilient digital society, ensuring confidence in the software solutions we rely on daily. recent years have witnessed a paradigm shift in this field with the rise of deep learning models, which have revolutionized how vulnerabilities are identified. these models far outperform traditional approaches such as static analysis tools and manual code reviews, which are labor-intensive, error-prone, and unable to keep pace with the accelerating rate of software development. leveraging neural networks has enabled researchers to automate vulnerability detection, improving scalability and accuracy by identifying subtle patterns in code that signal potential flaws. most deep learning models rely on loss functions such as binary cross-entropy to learn from labeled datasets of vulnerable and non-vulnerable code. however, despite these advancements, the field remains fraught with challenges. one of the most significant bottlenecks in existing models is the high semantic and structural similarity between vulnerable and non-vulnerable code samples. this similarity often causes embeddings of the two classes to overlap in highdimensional space, resulting in increased false positives and false negatives, thereby undermining the reliability of the models. to better quantify this challenge, we conducted a focused analysis of the embedding space. using a simple, yet effective method, we computed the average cosine similarity between embeddings of samples with opposite labels across both code and general text datasets. the results, as seen in the figure below 1, revealed that code datasets exhibit significantly higher similarity across labels than general text data, underscoring the unique difficulty of separating vulnerable from non-vulnerable code. this inherent overlap in the embedding space presents a key challenge in ensuring accurate and reliable vulnerability detection. contrastive learning has emerged as a promising technique to address this challenge. by structuring the embedding space to maximize separation between samples with opposite labels and encouraging tighter clustering of samples within the same class, contrastive learning reduces overlap and enhances the discriminative power of embeddings. this is particularly critical in the context of code vulnerability detection, where subtle differences between classes demand a highly optimized embedding space. however, existing contrastive learning methods, such as supervised contrastive loss (scl), exhibit limitations in this domain. while scl effectively prioritizes inter-class separation, it often fails to enforce sufficient intra-class cohesion. this can result in loosely clustered embeddings within each class, increasing the likelihood of misclassifications. consequently, scl struggles to handle the high similarity between vulnerable and nonvulnerable samples, limiting its effectiveness in real-world applications. to overcome these limitations, we propose a novel loss function, cluster enhanced supervised contrastive loss (cescl). cescl builds on the foundation of scl by introducing additional regularization techniques aimed at simultaneously minimizing intra-class separation and penalizing high cosine similarity between embeddings of vulnerable and non-vulnerable code snippets. this dual objective ensures tighter clustering within the same class while amplifying the dissimilarity between different classes, resulting in a well-structured embedding space optimized for classification. cescl achieves this by incorporating penalties for misaligned embeddings and emphasizing the structural and semantic nuances that distinguish vulnerable from non-vulnerable code. by fostering tighter intra-class cohesion and greater inter-class separation, cescl reduces embedding overlap, enabling models to better generalize across diverse and unseen code patterns. this results in lower false positive and false negative rates, addressing key reliability concerns in existing systems. in summary, this research introduces cescl as a targeted solution to the embedding challenges in code vulnerability detection. by addressing the shortcomings of existing loss functions, cescl provides a more robust and generalizable embedding space, significantly improving classification accuracy. our work also highlights the unique challenges of this domain through a quantitative analysis of embedding similarity, offering a new perspective on the limitations of current approaches. as software vulnerabilities continue to rise alongside the pace of digitization, the need for reliable and efficient detection methods has become more urgent than ever. cescl represents a step forward in building secure and trustworthy software systems, offering a foundation for future advancements in vulnerability detection. by bridging the gap between the limitations of scl and the demands of real-world applications, this research provides both a theoretical and practical contribu-tion to the field, paving the way for more secure digital ecosystems. codebert is a pretrained model tailor made for both programming and natural languages, capturing both syntactic and semantic features from a large corpus of code. it has been adopted widely for vulnerability detection because of its ability to handle tasks like code summarization, generation, and classification. graphcodebert extends codebert by incorporating data flow information within the model's architecture. this approach enhances the model's understanding of dependencies and control structures, allowing it to detect vulnerabilities that rely on intricate code flows, an area where traditional transformer models tend to disappoint. such advancements highlight the potential of transformer based models in advancing the field of code vulnerability detection. the devign dataset @xcite has been crucial in evaluating and benchmarking vulnerability detection models. devign contains over 21,000 labeled c/c++ code snippets drawn from opensource projects, with each snippet classified as vulnerable or non-vulnerable. the dataset presents unique challenges due to the incorporation of a wide variety of vulnerabilities. the dataset also constitutes complex vulnerabilities making it hard for models to effectively generalise across samples. devign also provides a rich structural context, including abstract syntax trees (asts) and control flow graphs (cfgs), which has proven useful for models designed to capture graph-based relationships in code, as demonstrated in @xcite original work making use of graph neural networks. supervised contrastive learning (scl), introduced by @xcite , is an emerging and powerful technique that focuses class separation by leveraging both positive and negative samples, making it immensely suitable for tasks where closely related samples are to be differentiated. in vulnerability detection, where samples of vulnerable and non vulnerable can appear notoriously similar, scl has shown potential @xcite by promoting embedding spaces that separate vulnerable and non-vulnerable samples. various regularization techniques have been applied to these losses to improve robustness in highsimilarity domains such as the one tackled in this paper. for instance, @xcite explored regularizing for invariance to data augmentation, improving the ability of models to handle difficult samples. in summary, this study builds on the strengths of transformer models like codebert and graph-codebert, evaluates performance on the devign dataset, and explores advanced contrastive learning techniques to enhance code vulnerability detection. by combining supervised contrastive learning with regularization strategies, we aim to improve the model's capability in embedding separation, thus enhancing overall classification performance. ## methodology to ameliorate the effectiveness of code vulnerability detection, this study builds a classification model on top of the both codebert and graph-codebert models, with additional dropout and batch normalization layers. these layers reduce overfitting and ensure stable training by normalizing activations, contributing to more reliable model performance. the central novelty in this approach is the novel loss function, which integrates supervised contrastive learning with a distance-based regularization term to improve embedding separation between classes. ## results and analysis in this study, three models were trained to assess the impact of different loss functions on code vulnerability detection. each model uses the same architecture, a classifier built on top of a pretrained codebert or graphcodebert model with dropout and batch normalization. the only difference between models being the loss function utilized during training: to evaluate performance, f1 scores were calculated on the test set for each model. these scores provide a comparison between each model's precision and recall, informing how well the different loss functions contribute to the model's accuracy and embedding separation. on top of this, the silhouette score was calculated as a measure of embedding separation. the silhouette score is a widely-used metric to evaluate clustering quality. it ranges from -1 to 1, where a value near 1 indicates that samples are well-separated and closely grouped within their respective clusters, and a value near -1 suggests significant overlap between clusters. in the context of our study, a higher silhouette score implies that code snippets belonging to the same class (vulnerable or non-vulnerable) are more similar to each other than to those in the opposing class, thereby indicating effective embedding separation. ## conclusion in this work, we introduced cluster-enhanced supervised contrastive loss (cescl), a novel loss the experimental results, as shown in tables 1 and 2 , demonstrate that cescl consistently outperforms both bce and bce + scl across the models tested. notably, codebert + cescl achieved the highest f1 score of 0.625 among codebert models and the best silhouette score of 0.056, highlighting its ability to generate wellclustered and discriminative embeddings for vulnerability classification. similarly, graphcodebert + cescl showed significant improvements, achieving a 4.1% increase in f1 score and a favorable silhouette score of 0.050 compared to the other configurations. it is worth noting that, although the incorporation of scl in isolation sometimes led to a reduction in the silhouette score, the overall improvement in f1 score indicates that the cescl framework effectively optimizes the embedding space for classification. this discrepancy suggests that while the silhouette score is a useful measure of clustering quality, it may not fully capture the nuances that contribute to enhanced detection performance in this context. these results indicate that cescl effectively improves intra-class compactness and inter-class separation, thereby enhancing the performance of the model in detecting vulnerabilities in code. future work could explore further optimizations to the cescl framework and test it on additional code-related tasks to fully realize its potential in improving the robustness and reliability of code classification models.