Towards Learning (Dis)-Similarity of Source Code from Program Contrasts (2022)

Abstract: 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.

## introduction understanding the functional similarity/dissimilarity of source code is at the core of several code modeling tasks such as software vulnerability and code clone detection, which are important for software maintenance @xcite . existing pre-trained transformer models @xcite @xcite show promises for understanding code syntax (i.e., tokens and structures). however, they still get confused when trying to identify functional (dis)-similarities. for instance, syntax-based models can embed two code fragments with identical functionality but very different tokens and structures as distinct vectors and fail to identify them as semantically similar. likewise, these models cannot distinguish between two code fragments that differ in functionalities but share a close syntactic resemblance. for example, consider an if statement if(len(buf) < n) checking buffer length before accessing the buffer. keeping the rest of the program the same, if we simply replace the token '<' with '≤,' the modification can potentially trigger security vulnerability, e.g., buffer overflow bug foot_0 . it is challenging for existing pre-training techniques to tell apart such subtle differences in the functionalities. in addition, existing pre-training techniques rely on a huge volume of training corpus that is randomly selected. for fine-tuning tasks like code clone detection or vulnerability detection, such random selection of training data is never tailored to teach the model about code functionalities. to address these limitations, we present disco, a self-supervised pre-trained model that jointly learns the general representations of source code and specific functional features for identifying source code similarity/dis-similarity. similar to state-of-the-art pre-trained transformer models @xcite , we apply the standard masked language model (mlm) to capture the token features of source code. to learn about the structural code properties, we propose a new auxiliary pre-training task that consumes additional inputs of local tree-based contexts (e.g., parent or sibling nodes in abstract syntax trees) and embeds such structural context, together with the token-based contexts, into each token representation. on top of such well-learned general code representations, we further incorporate prior knowledge of code clones and vulnerable programs into the pre-training to help the model learn the functional (dis)-similarity. we design structure-guided code transformation heuristics to automatically augment each training sample with one synthetic code clone (i.e., positive samples) that is structurally different yet functionally identical and one vulnerable contrast (i.e., hard negative samples) that is syntactically similar but injected with security bugs. during the pre-training, disco learns to bring similar programs closer in the vector space and differentiate the benign code from its vulnerable contrast, using a contrastive learning objective. since we augment the dataset in a more targeted way than existing works and the model explicitly learns to reason about a code w.r.t. its functional equivalent and different counterparts during pre-training, disco can learn sufficient knowledge for downstream applications from a limited amount of data, consequently saving computing resources. in particular, we evaluate disco for clone detection and vulnerability detection, as the knowledge of similar/dissimilar code fragments is at the core of these tasks. to this end, we pre-train disco on a small dataset, with only 865 mb of c code and 992 mb java code from 100 most popular github repositories, and evaluate the model on four different datasets for vulnerability and code clone detection. experiments show that our small models outperform baselines that are pre-trained on 20× larger datasets. the ablation study ( §5.4) also reveals that pre-training our model with 10× larger datasets further improves the performance up to 8.2%, outperforming state-of-the-art models by 1% for identifying code clones and up to 9.6% for vulnerability detection, even if our dataset is still smaller. in summary, our contributions are: 1) we design structure-guided code transformation heuristics to automatically augment training data to integrate prior knowledge of vulnerability and clone detection without human labels. 2) we propose a new pre-training task to embed structural context to each token embedding. 3) we develop disco, a self-supervised pre-training technique that jointly and efficiently learns the textual, structural, and functional properties of code. even though pre-trained with significantly less data, disco matches or outperforms the state-of-the-art models on code clone and vulnerability detection. ## data augmentation without human labels our pre-training aims to identify similar programs that can be structurally different (positive sample) and differentiate the buggy programs (negative sample) that share structural resemblances with the benign ones. thus, we need a labeled positive and a negative example for each original sample. manually collecting them is expensive, especially at the scale of pre-training. to this end, we design code transformation heuristics to automatically generate such positive and negative samples so that the transformation can be applied to any amount of programs without human efforts. we first represent a code sample as abstract syntax tree (ast), and build a control/data flow graph from the ast. the code transformation heuristics are then applied to this graph. for every original code sample (x), we apply semantic preserving transformation heuristics ( §3.2) to generate a positive sample (x + ) and a bug injection heuristics ( §3.1) to generate a hard-negative code example (x -). we design the heuristics in a way that makes x + be the functional equivalent or semantic clone of x and x -be the buggy/noisy version of x. noted that not all heuristics are applicable to all code samples; we decide on applicable heuristics based on the flow graph of the original code. figure 1 shows an example of the code transformation. ## disco this section presents the model architecture, input representation, and pre-training tasks. disco uses a 12-layered transformer encoder model similar to bert. we feed the model with both source code text and structure (ast) information ( §4.1). we pre-train disco using three different pre-training tasks ( §4.2). figure 2 depicts an example workflow of disco. we randomly select tokens in the original sample, mask them and their node types, and then use the embedding of these masks to predict them back. we further extract the sequence embeddings within a minibatch and contrast them based on the code functionality. figure 2 : an illustration of disco pre-training with a minibatch of three. the original code and its node types will be randomly masked with [m ask], and the final representation of masked tokens will be used to recover their source tokens and node types. the original code, say x, will also be transformed to build (x, x + , x -). then the triplet will be fed into the same transformer encoder and get the embedding of each sequence with [cls] tokens for contrastive learning. ## conclusion in this work, we present disco, a self-supervised contrastive learning framework to both learn the general representations of source code and specific characteristics of vulnerability and code clone detections. our evaluation reveals that disco pretrained with smaller dataset can still outperform the large models' performance and thus prove the effectiveness of our design.