C ode R etriever: A Large Scale Contrastive Pre-Training Method for Code Search (2022)

Abstract: In this paper, we propose the CodeRetriever model, which learns the function-level code semantic representations through large-scale code-text contrastive pre-training. We adopt two contrastive learning schemes in CodeRetriever: unimodal contrastive learning and bimodal contrastive learning. For unimodal contrastive learning, we design an unsupervised learning approach to build semantic-related code pairs based on the documentation and function name. For bimodal contrastive learning, we leverage the documentation and in-line comments of code to build code-text pairs. Both contrastive objectives can fully leverage large-scale code corpus for pre-training. Extensive experimental results show that CodeRetriever achieves new state-of-the-art with significant improvement over existing code pre-trained models, on eleven domain/language-specific code search tasks with six programming languages in different code granularity (function-level, snippet-level and statement-level).These results demonstrate the effectiveness and robustness of CodeRetriever.The codes and resources are available at https://github.com/microsoft/AR2/tree/main/CodeRetriever .

## introduction code search aims to retrieve functionally relevant code given a natural language query to boost @xcite . recently, it has been shown that code pre-training techniques, such as code-bert @xcite and @xcite , could significantly improve code search performance via self-supervised pre-training using large-scale code corpus @xcite . doc: get the fibonacci number. code: def fibonacci_number(index): @xmath0,index+1): cache[i] = \ cache[i-1] + cache[i-2] return cache[index] (a) fibonacci doc: sort the input array into ascending order. code: def bubblesort(arr): @xmath1: for j in range(0, n-i-1): # if adjacent elements appear in # descending order, swap them. if arr[j] > arr[j+1]: arr[j], arr[j+1] = arr[j+1], arr[j] training objective which targets on learning to predict (masked) tokens in a given code context @xcite @xcite . however, this token-based approach generally results in poor code semantic representations due to two reasons. the first one is the anisotropy representation issue. as discussed in @xcite , the token-level self-training approach causes the embeddings of high-frequency tokens clustered and dominate the representation space, which greatly limits the expressiveness of long-tailed low-frequency tokens in pre-trained models. thus, the anisotropic representation space induces poor function-level code semantic representation @xcite . in programming language, the problem of token imbalance is even more severe than that of natural language. for example, common keywords and operators such as "=", "{", and "}" appear almost everywhere in java code. the second one is the cross-language representation issue. the widely used codesearch-net corpus @xcite ) contains codes from six different programming languages such as python, java, etc. since the code with mixed programming languages can hardly appear within the same context, it is challenging for the pre-trained model to learn a unified semantic representation of the code with the same functionality but using different programming languages. to address these limitations, we propose the coderetriever model, focusing on learning the function-level code representations, specifically for code search scenarios. the coderetriever model consists of a text encoder and a code encoder, which encodes text/code into separate dense vectors. the semantic relevance between code and text (or code and code) is measured by the similarity between dense vectors @xcite @xcite . in the training of coderetriever, the code/text encoders are optimized by minimizing two types of contrastive losses: 1.unimodal contrastive loss, encourages the model to push codes with similar functionality closer in representation space. to estimate whether two codes are semantically close, the model needs to reason based on the given code and understand its semantics. 2.bimodal contrastive loss, helps model the relevance between code and text. since the document or comment contains rich semantic information of the code, it can encourage the model to learn better code representation from natural language. in this work, we adopt the commonly used code-searchnet corpus @xcite for training the coderetreiver. codesearchnet mainly contains paired dataset (a function paired with a document) and unpaired dataset (only a function). the paired dataset could be directly used for bimodal contrastive learning. for unimodal contrastive learning in coderetriver, we build positive codecode pairs by an unsupervised semantic-guided approach. figure 1 (a) shows a code-code example: two implementations of the fibonacci number algorithm. moreover, the generated code-code pairs can be with different programming languages, which can mitigate the cross-language representation issues. to further take advantage of the largescale code in unpaired data and paired data, we extract the code and in-line comment pairs to enhance the bimodal contrastive learning in coderetriever. figure 1( @xcite shows an example to indicate that the in-line comment (comment shortly) can also reflect the code's semantics and internal logic. specifically, the underlying logic of "if adjacent elements appear in descending order, swap them" corresponds to sorting the input array into ascending order and such fine-grained semantic information can also help learn better code representation. through contrasting these unimodal and bimodal pairs, coderetriever can 1. learn better the function-level code semantic representation, which could alleviate the anisotropy representation issue @xcite ; 2. explicitly model the relevance of codes with different programming languages and treat unified natural language as a fulcrum to mitigate cross-language representation issue. we evaluate coderetriever on eleven code search datasets covering six programming languages, real-world scenarios and codes with different granularity (function-level, snippetlevel and statement-level), and the results show that coderetriever achieves a new state-of-the-art performance. ## preliminary: code search codesearchnet corpus @xcite is the largest publicly available code dataset. the corpus is collected from open-source non-fork github repositories, which contains 2.1m paired data (a function paired with a document) and 6.4m unpaired data (only functions). in the literature, code-search approaches @xcite @xcite ) make use of the paired codedocument dataset in codesearchnet corpus to train a siamese encoder model for language to code retrieval. however, rich unlabeled code corpus is either simply abandoned or severed as code pretraining corpus @xcite . we argue that token-level code pre-training objectives do not explicitly learn the function-level code representation. thus existing code pre-training models @xcite @xcite are sub-optimal for code search. in this work, we propose the coderetriever to learn the function-level code semantic representation. coderetriever is initialized with the code pre-trained model (i.e., graphcodebert). it takes code-doc and code-comment paired data for bimodal contrastive learning, and code-code paired data for unimodal contrastive learning. after coderetriever's pre-training, it can serve for downstream domain/language specified datasets. ## approach in this section, we present the model architecture and training objective of coderetriever. coderetriever adopts a siamese code/text encoder architecture to represent code/text as dense vectors. let e code (•; θ) and e text (•; ϕ) denote code and text encoders, respectively. the semantic similarities between code-code pair (c, c + ), and textcode pair (t, c + ) are calculated as: ## experiment for fair comparison, coderetriever adopts the same model architecture as previous works @xcite . coderetriever shares parameters of code encoder and text encoder. it contains 12 layers transformer with hidden size of 768 and attention heads of 12. to accelerate the training process, we initilaize coderetriever with the released parameters of @xcite . we show more details in appendix e. ## conclusion in this paper, we introduce coderetriever that combines unimodal and bimodal contrastive learning as pre-training tasks for code search. for unimodal contrastive learning, we propose a semanticguided method to build positive code pairs. for bimodal contrastive learning, we utilize the document and in-line comment to build positive textcode pairs. extensive experimental results on several publicly available benchmarks show that the proposed coderetriever brings significant improvement and achieves new state-of-the-art on all benchmarks. further analysis results show that coderetriever is also powerful on low resource and crosslanguage code search tasks, and demonstrate the effectiveness of unimodal and bimodal contrastive learning.