Exploring Representation-level Augmentation for Code Search (2022)

Abstract: Code search, which aims at retrieving the most relevant code fragment for a given natural language query, is a common activity in software development practice. Recently, contrastive learning is widely used in code search research, where many data augmentation approaches for source code (e.g., semantic-preserving program transformation) are proposed to learn better representations. However, these augmentations are at the raw-data level, which requires additional code analysis in the preprocessing stage and additional training cost in the training stage. In this paper, we explore augmentation methods that augment data (both code and query) at representation level which does not require additional data processing and training, and based on this we propose a general format of representation-level augmentation that unifies existing methods. Then, we propose three new augmentation methods (linear extrapolation, binary interpolation, and Gaussian scaling) based on the general format. Furthermore, we theoretically analyze the advantages of the proposed augmentation methods over traditional contrastive learning methods on code search. We experimentally evaluate the proposed representation-level augmentation methods with state-of-the-art code search models on a large-scale public dataset consisting of six programming languages. The experimental results show that our approach can consistently boost the performance of the studied code search models.

## introduction in software development, developers often search and reuse commonly used functionalities to improve their productivity @xcite . with the growing size of large-scale codebases such as github, retrieving semantically relevant code fragments accurately becomes increasingly important in this field @xcite . traditional approaches @xcite @xcite @xcite @xcite leverage information retrieval techniques to treat code snippets as natural language text and match certain terms in code with queries, hence suffering from the vocabulary mismatch problem @xcite . deep siamese neural networks first embed queries and code fragments into a joint embedding space, then measure similarity by calculating dot product or cosine distance @xcite @xcite . recently, with the popularity of large scale pre-training techniques, some big models for source code @xcite @xcite @xcite with various pre-training tasks are proposed and significantly outperform previous models. contrastive learning is widely adopted by the above-mentioned models. it is suitable for code search because the learning objective aims to push apart negative query-code pairs and pull together positive pairs at the same time. in contrastive learning, negative pairs are usually generated by in-batch augmentation @xcite . for positive pairs, besides labeled ones, some researchers proposed augmentation approaches to generate more positive pairs @xcite @xcite @xcite . the main hypothesis behind these approaches is that augmentations do not change the original semantics. however, these approaches are resource-consuming @xcite . models have to embed the data again for the augmented data. to solve this problem, some researchers proposed representation-level augmentation, which augments the representations of the original data. for example, linear interpolation, a representationlevel augmentation method, is adopted by many classification tasks in nlp @xcite @xcite . the augmented representation captures the structure of the data manifold and hence could force model to learn better features, as argued by @xcite . these augmentation approaches are also considered to be semantic-preserving. the representation-level augmentation methods are not investigated on the code search task before. to the best of our knowledge, @xcite is the only work to bring representation-level augmentation approaches to a retrieval task. besides linear interpolation, it also proposes another approach called stochastic perturbation for document retrieval. although these augmentation methods bring improvements to model performance, they are not yet fully investigated. the relationships between the existing methods and how they affect model performance remain to be explored. in this work, we first unify linear interpolation and stochastic perturbation into a general format of representation-level augmentation. we further propose three augmentation methods (linear extrapolation, binary interpolation, and gaussian scaling) based on the general format. then we theoretically analyze the advantages of the proposed augmentation methods based on the most commonly used infonce loss @xcite . as optimizing infonce loss equals to maximizing the mutual information between positive pairs, applying representation-level augmentation leads to tighter lower bounds of mutual information. we evaluate representation-level augmentation on several siamese networks across several large-scale datasets. experimental results show the effectiveness of the representation-level augmentation methods in boosting the performance of these code search models. to verify the generalization ability of our method to other tasks, we also conduct experiments on the paragraph retrieval task, and the results show that our method can also improve the performance of several paragraph retrieval models. in summary, our contributions of this work are as follows: the rest of the paper is organized as follows. we introduce related work of code search and data augmentation in section 2. section 3 introduces the main part, including the general format of representation-level augmentation, new augmentation methods, and their application on code search. in section 4, we analyze the theoretical lower bounds of mutual information and study why our approach works. in section 5 and section 6, we conduct extensive experiments to show the effectiveness of our approach. then we discuss the generality of our approach in section 7 and section 8 concludes this paper. ## approach in this section, we unify the linear interpolation and stochastic perturbation into a general format. based on it, we propose three other augmentation methods for the code retrieval task, linear extrapolation, binary interpolation and gaussian scaling. then, we explain how to apply representation-level augmentation with infonce loss in code retrieval. ## theoretical analysis in this section, we mathematically analyze the effect of infonce loss with or without representationlevel augmentation and prove that optimizing in-fonce loss with representation-level augmentation leads to mutual information with tighter lower bounds between positive pairs. here we take linear interpolation as an example to demonstrate the benefits. other forms of representation-level augmentation are left to future work. the mutual information between a query q and a code fragment c is: theorem 1 optimizing infonce loss l n improves lower bounds of mutual information i(q, c) for a positive pair: proof of theorem 1 is presented in appendix a.1. this is proved in the original paper of in-fonce loss @xcite ). here, to better extend the proof of theorem 2, we prove it in another way. theorem 2 optimizing infonce loss l n with representation-level augmentation improves a tighter lower bounds of mutual information i(q, c) for a positive pair: proof of theorem 2 is presented in appendix a.2. since we interpolate q with other samples q - in the batch (c with c -), the mutual information between (q, c -), (q -, c) and (q -, c -) are also incorporated into the optimizing process. as defined in eq.1, β is a small value that close to 0. according to eq.4, for negative pairs p(c|q) can be expressed as p(c) the optimal mutual information between negative pairs is also 0. note that we interpolate q and c with different samples to make sure that (q -, c -) is a negative pair. thus, the last three terms in eq.6 can be ignored. comparing 1 α 2 (log(n b) -l n ) with log(b) -l n , we can see that representationlevel augmentation improves the lower bounds of mutual information. ## experimental setup in this section, we describe datasets, baselines, and implementation details. ## results in this section, we first show the overall performance on code search when applying representation-level augmentation, and then individually analyze each augmentation method. then, we demonstrate the relationship between loss and mrr and the effect of augmentation on vector distribution. finally, we take an ablation study to analyze the impact of augmentation times. ## discussion according to the proof of theorem 2, the advantage of applying representation-level augmentation should be task-agnostic. to show its generalization ability, we evaluate our proposed approaches on two passage retrieval benchmark datasets, nfcorpus @xcite and fiqa-2018 foot_0 . the difference between code search and passage retrieval is that the retrieved items are changed from code snippets written in programming language to passages written in english. we take distilbert @xcite and roberta @xcite for experiments. we implement our approach based on an open-sourced framework beir @xcite . the two models are fine-tuned on two datasets for 20 epochs, respectively. the settings of augmentation are the same as those in the code search task. for other hyper-parameters, we follow the settings that are provided by the framework. results are shown in ## conclusion in this work, we unify existing approaches to propose a general format of representation-level augmentation in code search. based on the general format, we propose three other augmentation methods. we further theoretically analyze the effect of representation-level augmentation by proving that it helps optimize a tighter lower bound of mutual information between positive pairs. we evaluate our approach on several models and datasets and the results demonstrate the effectiveness of the proposed approach.