Addressing Leakage in Self-Supervised Contextualized Code Retrieval (2022)

Abstract: We address contextualized code retrieval, the search for code snippets helpful to fill gaps in a partial input program. Our approach facilitates a large-scale self-supervised contrastive training by splitting source code randomly into contexts and targets. To combat leakage between the two, we suggest a novel approach based on mutual identifier masking, dedentation, and the selection of syntax-aligned targets. Our second contribution is a new dataset for direct evaluation of contextualized code retrieval, based on a dataset of manually aligned subpassages of code clones. Our experiments demonstrate that the proposed approach improves retrieval substantially, and yields new state-of-the-art results for code clone and defect detection.

## introduction ai-supported software development has recently experienced growing interest @xcite , addressing various code understanding tasks such as code auto-completion @xcite , natural language code search @xcite , and code clone detection @xcite . our focus is on a related task called contextualized code search @xcite : given an incomplete piece of code and a certain position of interest (e.g., the current cursor position), a retriever searches for code fragments that are relevant for filling in the missing piece. this setting aligns well with programmers' workflow, and differs substantially from the three tasks mentioned above as follows: (1) in contrast to natural language code search, contextualized code search can exploit local code context. (2) while code generated by autocompletion systems such as github's codex @xcite is prone to subtle programming errors, contextualized search leaves the developer in charge, and the origin of a solution remains transparent. (3) in contrast to clone detection, contextualized code search is not targeted at semantically similar code but code that complements the query. a key challenge with contextualized code search is that supervised labels for relevant code passages are missing. therefore, we bootstrap a selfsupervised learning process by drawing inspiration from cloze tasks in natural language processing @xcite : given a large-scale dataset containing pieces of code in 16 programming languages, we erase random blocks. we refer to these blocks as targets, and to their surrounding as contexts. together, these pairs form samples for contrastive learning. unfortunately, as figure 1 shows, this approach suffers from leakage between context and target, as the two share (1) common identifiers, (2) a matching indentation level, and (3) in some languagesif dividing a syntactic primitive such as for-loopsmatching brackets. retrievers might exploit these effects and bypass semantic similarity. to this end, our first contribution is a novel approach towards self-supervised code retrieval, which avoids the above bias through de-leaking steps such as mutual identifier masking and dedentation. the second challenge we address is evaluation: so far, the focus of evaluating code retrieval systems has been on natural language queries (which can be bootstrapped from docstrings) @xcite . contextualized code retrieval has been evaluated only indirectly via infilling quality @xcite , which reflects the actual retrieval quality poorly. therefore, our second contribution is a rigorous evaluation of contextualized code retrieval on a manually curated dataset based on aligned code clones. we call this dataset cocos and make it available for future research. on cocos, we demonstrate that retrieval quality benefits substantially from our de-leaking approach. also, we achieve state-of-the-art results figure 1: our approach bootstraps code pairs for contrastive learning by removing target passages (green) from random code contexts (gray). to address leakage between the two -which can be due to matching identifiers, indentation, and brackets -we (a) select the target using the code's syntactic structure, (b) dedent the target (orange arrow), and (c) mutually mask identifiers. finally, we apply contrastive learning on the resulting code pairs (d). ## approach given a piece of code as a token sequence @xmath0, . . . , x m , our goal is to boostrap a contexttarget pair for contrastive learning. the target is a subsequence @xmath1, . . . , x i+l with 1≤@xmath2. by replacing this subsequence with a special mask token, we obtain the context to x ′ and y we prepend a programming-languagespecific cls token. to address the above leakages, we suggest three steps called tree-based span selection (ts), mutual identifier masking (im) and dedenting (de). ## dataset in this section, we first describe the large-scale data which our retriever is trained on. second, we outline cocos, a new bechmark we propose for contextualized code retrieval. pre-training dataset our self-supervised code retrieval model is pre-trained on 33m files in 16 programming languages (see appendix a). as code files tend to be large, we truncate them using treebased span selection (cmp. section 2): starting from a whole file, we randomly select sufficiently large spans of code (length between 150 and 800 tokens). we remove those segments from the original file and feed the shortened file as well as all individual segments as inputs x into the learning process described in section 2. a special identifier (similar to code folding in an ide) marks those positions in the original file where segments have been removed. cocos evaluating contextualized code retrieval models is hard because little or no suitable evaluation data is available to indicate which subblocks in the code implement the same functionality. to address this gap, we have created a new dataset based on bigclonebench @xcite , a java code clone dataset that provides pairs of semantically similar functions. given a function in bigclonebench, we manually select a sub-passage modeling a particular target functionality (e.g. extracting a zip file). we then label which lines in the function's clones match this functionality (see listings 1 -3 in the appendix). based on these targets and their surrounding contexts, we evaluate how well a model retrieves targets implementing the same functionality in code clones. we manually gather 606 context-target pairs implementing 31 randomly selected functionalities. finally, we add 10k non-relevant disctractor snippets by randomly sampling top-level statements from method bodies in codesearchnet @xcite zeroshot code retrieval results for different de-leaking steps as described in section 2: tree-based span selection (ts); mutual identifier masking (im); dedenting (de). we report non-neural results for bm25 @xcite using the elasticsearch standard tokenizer (standard) and a tokenizer that splits on camel case (camel). ## evaluation we report results for zero-shot code retrieval on cocos and for two similar code understanding tasks from codexglue @xcite , namely code clone detection and code defect detection. for all experiments, we report test results of the model with the highest mean reciprocal rank (mrr) on 30k held-out validation samples of the pre-training dataset. ## conclusion we have proposed a new approach towards unsupervised code retrieval, which reduces leakage between randomly drawn targets and their contexts. we also contribute a dataset cocos, on which we demonstrate via ablations that leakage reduction is crucial for an efficient training. our approach also yields competitive representations for related tasks, as demonstrated by new state-of-the-art results on clone and defect detection. an interesting future direction will be to combine our retriever with generators for a combined, unsupervised training.