VDCRL: vulnerability detection with supervised contrastive code representation learning

Code vulnerability detection plays a pivotal role in ensuring software security, and recent deep learning-based approaches have demonstrated their effectiveness. However, these methods often exhibit strong detection capabilities only on the datasets used. When applied to unknown datasets, their performance deteriorates significantly, highlighting a limitation in generalization ability. To address the above issues, we propose VDCRL, a novel vulnerability detection framework based on supervised contrastive code representation learning. VDCRL employs both input-based and feature space-based code augmentation techniques to generate two augmented versions of each original sample. It then extracts source code features and assembly instruction features from these augmented samples. These features are integrated using a feature fusion encoder, SAFE, which captures and embeds both types of information. Finally, VDCRL leverages supervised contrastive learning and a Bidirectional Gated Recurrent Unit (BGRU) model to train and detect vulnerabilities effectively. We train VDCRL on a synthetic dataset and test its performance on two real-world datasets. Experimental results demonstrate that VDCRL significantly outperforms current state-of-the-art vulnerability detection methods, achieving superior generalization and detection performance.