Towards Learning (Dis)-Similarity of Source Code from Program Contrasts
This work addresses code modeling challenges for software developers and security analysts, offering a more data-efficient approach compared to prior methods.
The paper tackles the problem of understanding functional similarity and dissimilarity in source code for tasks like vulnerability and clone detection, by introducing DISCO, a self-supervised model that uses targeted data augmentation and pre-training with program contrasts, achieving state-of-the-art performance while using only 5% of the training data of existing models.
Understanding the functional (dis)-similarity of source code is significant for code modeling tasks such as software vulnerability and code clone detection. We present DISCO(DIS-similarity of COde), a novel self-supervised model focusing on identifying (dis)similar functionalities of source code. Different from existing works, our approach does not require a huge amount of randomly collected datasets. Rather, we design structure-guided code transformation algorithms to generate synthetic code clones and inject real-world security bugs, augmenting the collected datasets in a targeted way. We propose to pre-train the Transformer model with such automatically generated program contrasts to better identify similar code in the wild and differentiate vulnerable programs from benign ones. To better capture the structural features of source code, we propose a new cloze objective to encode the local tree-based context (e.g., parents or sibling nodes). We pre-train our model with a much smaller dataset, the size of which is only 5% of the state-of-the-art models' training datasets, to illustrate the effectiveness of our data augmentation and the pre-training approach. The evaluation shows that, even with much less data, DISCO can still outperform the state-of-the-art models in vulnerability and code clone detection tasks.