Tackling Long Code Search with Splitting, Encoding, and Aggregating

Early studies Nie et al. (2016); Yang and Huang (2017); Rosario (2000); Hill et al. (2011); Satter and Sakib (2016); Lv et al. (2015); Van Nguyen et al. (2017) in code search mainly applied information retrieval (IR) techniques directly, treating code search as a text matching task. Both queries and code snippets were considered plain text, and traditional text matching algorithms such as bag-of-words (BOW) Schütze et al. (2008), Jaccard Jaccard (1901), term frequency-inverse document frequency (TF-IDF) Robertson and Jones (1976), BM25 (an improved version of TF-IDF) Robertson and Zaragoza (2009), and the extended boolean model Lv et al. (2015) were employed. Since code length has minimal impact on modeling complexity, these methods could encode long code without truncation.

Following the introduction of the large-scale pre-training model BERT Devlin et al. (2019), CodeBERT was proposed by Feng et al. (2020). CodeBERT is a model pre-trained on unlabeled source code and comments, which achieved impressive performance in text-based code search through fine-tuning on text-code paired datasets. Huang et al. (2021) introduced CoCLR, a contrastive learning method that enhances query-code matching. Sun et al. (2022) developed a context-aware code translation technique that translates code snippets into natural language descriptions. Gu et al. (2022) utilized deep hashing and code classification to accelerate code search, while Chai et al. (2022) adapted few-shot meta-learning to code search. Guo et al. (2021) proposed GraphCodeBERT, incorporating structure-aware pre-training tasks to improve code understanding and performance. Recently, Hu et al. (2023) utilized a two-stage fusion code search framework that combines bi-encoders and cross-encoders to enhance performance. However, the computational complexity of Transformers and limited GPU memory often lead to the truncation of long code snippets.

Recently, there have been notable advancements in neural code representation methods that leverage code structure, particularly Abstract Syntax Trees (AST), yielding impressive performance Alon et al. (2020); Sun et al. (2020); Bui et al. (2021); Kim et al. (2021); Peng et al. (2021); Hellendoorn et al. (2019); Allamanis et al. (2021); Georgiev et al. (2022); Ma et al. (2023); Du and Yu (2023). MMAN Wan et al. (2019) incorporates a multi-modal attention fusion layer to combine AST and Control Flow Graph (CFG) representations. ASTNN Zhang et al. (2019a) and CAST Shi et al. (2021) segment large ASTs into sequences of smaller statement trees, encoding them into vectors by capturing the lexical and syntactical information of each statement. TBCAA Chen et al. (2019) employs a tree-based convolution network over API-enhanced ASTs. UniXcoder Guo et al. (2022) leverages both AST and code comments to enrich code representation. GraphCodeBERT Guo et al. (2021) incorporates variable relations extracted from ASTs in its pre-training tasks. In our work, we specifically aim to capture and model the structural information present in long code snippets.

The application of Transformer models for long text can be broadly divided into two categories: scaling up attention and enhancing the original Transformer model, and aggregation methods. The first category includes four main approaches: sparse attention Child et al. (2019); Correia et al. (2019); Beltagy et al. (2020); Kitaev et al. (2019); Roy et al. (2021); Ainslie et al. (2020); Jiang et al. (2020); Günther et al. (2023), recurrence Dai et al. (2019), hierarchical mechanisms Zhang et al. (2019b); Gao and Callan (2022), and compressed attention Ye et al. (2019); Guo et al. (2019). Sparse attention restricts each token to attend to only a subset of other tokens. Recurrence integrates recurrent neural network elements into Transformer models to extend their attention span. Hierarchical mechanisms model long input text hierarchically, from sentences to paragraphs. Compressed attention selectively compresses specific parts of the input.

The second category, aggregation methods, involves aggregating multiple passage scores or representations for a long document. For instance, Wang et al. (2019) proposed a multi-passage BERT model to globally normalize answer scores across all passages in the question answer task. In the context of document ranking, SMITH Yang et al. (2020) learns a document representation through hierarchical sentence representation aggregation. PARADE Li et al. (2020) employs Max, CNN, Attention, and Transformer to aggregate the passage representations. Tsujimura et al. (2023) uses a sliding window method to manage long input sequences in the context of medical Named Entity Recognition tasks.

However, these methods may not be entirely suitable for highly structured code. In well-designed programs, code within the same module, such as a function, is closely interconnected, while interactions between different modules are loosely coupled, adhering to the principle of high cohesion and low coupling. Conversely, long text in natural language tends to exhibit coherence. In this paper, we investigate the applicability of long text methods in the field of code search and propose a new baseline SEA for long code search.

Code search aims to find the most relevant code snippet $C$ from a given codebase that matches a query $Q$. For a current deep-learning model, we first transform query $Q$ and the code snippets $C$ to query and code tokens with the $\mathbf{tokenizer}$ such as BPE Sennrich et al. (2016). Then we transform the token ids of the query $Q$ and the code snippets $C$ to vector representations $\mathbf{e_{q}}$ and $\mathbf{e_{c}}$ by neural network encoders, and calculate the similarity (or distance) measures in Euclidean space such as Cosine similarity or Euclidean distance to obtain the cross-modal similarity score $s$. The calculation can be formalized as follows:

$$\begin{cases}\mathbf{e_{q}}=\Gamma(\mathbf{tokenizer}(Q))\\ \mathbf{e_{c}}=\Gamma^{\prime}(\mathbf{tokenizer}(C)),C\in Codebase\\ s=sim(\mathbf{e_{q}},\mathbf{e_{c}})\end{cases}$$

(1)

where $\Gamma$ and $\Gamma^{\prime}$ are two well-trained neural network encoders learned from labeled paired data.

To control memory and computation costs in training stage, it is common practice to truncate long code. For example, GraphCodeBERT typically takes the first 256 code tokens by default. To investigate whether this truncation method results in information loss, we conducted token length statistics on CodeSearchNet. As shown in Table 2, we found that snippets with a token length less than 256 accounted for only 14.1%, while 53.5% of code snippets exceeded the maximum encoding length of 512 tokens for Transformers. This indicates that truncation leads to information loss for snippets with a token length greater than 256.

To examine the search performance difference of GraphCodeBERT across query subsets with varying ground truth (GT) code lengths, we divided the python test subset of CodeSearchNet (CSN) into 5 distinct query sets based on different GT code token lengths. We calculated the Mean Reciprocal Rank (MRR) of GraphCodeBERT for various code truncation lengths, as shown in Table 1. Notably, we observed a downward trend in search performance as the ground-truth code token length increased (from top to bottom) for code token lengths surpassing 256 tokens, indicating that long code snippets pose challenges for GraphCodeBERT. Moreover, as the code truncation length extended from left to right, we observed a relatively consistent search performance when the truncation length exceeded the token length. And there emerged an upward trend in the search performance for code snippets with the token length surpassing the truncation length. This suggests that simply truncating long code may result in the loss of valuable information.

We introduce our SEA in this section. The overall pipeline is illustrated in Figure 2. Given a code snippet $C$, our objective is to derive a code representation $e_{c}$. To achieve this, we employ a multi-step approach. We first split the code snippet into a code piece set:

$$P=\mathbf{Split}(C)=\{p_{1},p_{2},\dots,p_{n}\}.$$

(2)

Then we use the sliding window method to obtain a partially overlapping code block set:

$$B=\mathbf{SlidingWindow}(P)=\{b_{1},b_{2},\dots,b_{k}\}.$$

(3)

Assuming the window size is $w$ and the step is $s$, then the code block number is $k=\lfloor\frac{n-w}{s}+1\rfloor$, where $\lfloor\cdot\rfloor$ refers to round down. Next, we utilize a code encoder, such as GraphCodeBERT, to obtain embeddings for each of the $k$ code blocks:

$$e_{B}=\{e_{b_{1}},e_{b_{2}},\dots,e_{b_{k}}\}.$$

(4)

Finally, an aggregation method is applied to combine the $k$ embeddings into the code representation $e_{c}$:

$$e_{c}=\mathbf{Aggregation}(e_{B})$$

(5)

To obtain the code piece set, we explore four splitting methods, namely space-based splitting, token-based splitting, line-based splitting, and AST-based splitting. Space-based splitting is simply splitting by space, resulting in splitting a string like “def read_image_file” is divided into {‘def’, ‘read_image_file’}. Similarly, token-based splitting and line-based splitting entail splitting based on tokens and lines, respectively.

An Abstract Syntax Tree (AST) is a tree representation of the syntactic structure of source code written in a programming language. Each node in the AST corresponds to a specific construct in the code, such as expressions, statements, or declarations. The hierarchical structure of ASTs reflects the syntax of programming languages, abstracting away certain syntactic details to focus on the core structure.

For AST-based splitting, our goal is to devise a method that is both straightforward and applicable to various programming languages. Inspired by CAST Shi et al. (2021), we parse a source code into an Abstract Syntax Tree with tree_sitter¹¹1https://github.com/tree-sitter/py-tree-sitter, and visit this AST by preorder traversal. In the case of composite structures (i.e. for, if, def, etc.), as depicted in Figure 2(a), we define the set of AST nodes {head_block, body}, where head_block is responsible for splitting the header and body of nested statements such as if and While statements, while body corresponds the method declarations. When encountering a composite structure, we insert a splitting mark before and after the head_block, effectively dividing a large AST into a sequence of non-overlapping subtrees. Subsequently, based on the AST splitting, we construct the code piece set $P$ by splitting the original code accordingly.

Meanpooling / Maxpooling. A straightforward approach to aggregate the embeddings of $k$ code blocks is to calculate the mean or maximum of their embeddings:

$$e_{c}=\mathbf{Mean/Max}(\{e_{b_{1}},e_{b_{2}},\dots,e_{b_{k}}\}).$$

(6)

However, a limitation of meanpooling is that each code block contributes equally to the final representation, regardless of their individual qualities. Similarly, maxpooling gives prominence to the block with the highest value. To address these limitations and enhance the aggregation process, we propose the incorporation of weighted embedding methods.

Attention-based aggregation. Recognizing that not all code blocks hold equal importance in representing long code snippets, we introduce self-adaptive weights $\alpha$ for each block embedding during aggregation:

$$e_{c}=\sum_{i}^{k}\alpha_{i}e_{b_{i}}.$$

(7)

Inspired by attention-based Multi-Instance Learning Li et al. (2021) and Lightweight Attentional Feature Fusion Hu et al. (2022), we compute the weights $\{\alpha_{1},\dots,\alpha_{k}\}$ as follows:

$$\{a_{1},\ldots,a_{k}\}=softmax(Linear(\{e_{b_{1}},\ldots,e_{b_{k}}\})).$$

(8)

For one layer attention, $Linear$ refers to a fully connected layer that transforms the dimension to 1. For two layer attention, $Linear$ refers to two fully connected layers that first transform the dimension to 128 and then transform the dimension to 1. Furthermore, as illustrated in Figure 3(a) and Figure 3(b), we explore the combination of attention with meanpooling / maxpooling methods:

$$e_{c}=\sum_{i}^{k}(\alpha_{i}e_{b_{i}})+\mathbf{Mean/Max}(\{e_{b_{1}},e_{b_{2}},\dots,e_{b_{k}}\}).$$

(9)

For computation cost analysis, SEA employs the sliding window method to significantly reduce complexity to $1/k$. The original complexity of GraphCodeBERT is given by $O(n^{2}\cdot d\cdot l)$, where $n,d,l$ represent sequence length, representation dimension, and layer number, respectively. By using the sliding window method, the complexity for each window becomes $O(w^{2}\cdot d\cdot l)$, where $w$ denotes the window size. Setting the step $s=w$, the total number of code blocks becomes $k=\frac{n}{w}$, leading to the window size $w=\frac{n}{k}$. Consequently, the overall complexity is simplified to:

$$O(k\cdot w^{2}\cdot d\cdot l)=O(k\cdot{(\frac{n}{k})}^{2}\cdot d\cdot l)=O(\frac{n^{2}}{k}\cdot d\cdot l).$$

(10)

This remarkable reduction in complexity to $\frac{1}{k}$ allows SEA to encode long code with less memory and computation costs.

Furthermore, as shown in Table 3, we observe that compared to GraphCodeBERT, SEA incorporating one-layer attention Aggregation introduces only $d$ additional learnable parameters. Despite this modest increase in parameter count, it plays a pivotal role in enhancing the effectiveness of the aggregation stage, as our experiments will provide the evidence in Section 6.1.

To enhance inference efficiency on large datasets, it is necessary to devise a batch processing method capable of encoding multiple long code snippets simultaneously. As outlined in Section 4.1, we obtain multiple code blocks from each long code snippet. However, due to the varying number of corresponding code blocks for different long code snippets, devising a general batch processing approach poses a challenge.

To address this issue, we introduce the combine-divide method. As illustrated in Figure 4, assuming a batch size of 3 (comprising three code snippets), the corresponding number of code blocks for each snippet is 2, 3, and 1, respectively. We begin by combining these six code blocks into a block batch and establish a mapping $M$ that links the code index to the block index. Subsequently, we input this block batch into the code encoder in parallel to obtain block embeddings. Finally, leveraging the information from mapping $M$, we segregate the embeddings into three groups and input them into the aggregation module to obtain distinct code representations.

We conduct experiments on the widely used CodeSearchNet Husain et al. (2019) dataset, comprising six programming languages, i.e. , Ruby, JavaScript, Go, Python, Java, and PHP. Following the approach in Guo et al. (2021), we apply filtering to eliminate low-quality queries and expand the retrieval set to encompass the entire code corpus.

In our evaluation, we use two popular automatic criteria: MRR (Mean Reciprocal Ranking) and R@k (top-k accuracy, k=1, 5, 10, 100). They are commonly used for in previous code search studies Lv et al. (2015); Gu et al. (2018); Sachdev et al. (2018); Husain et al. (2019); Feng et al. (2020); Huang et al. (2021); Guo et al. (2021). In addition, we report the number of parameter and inference time as the efficiency measure.

Our baseline is GraphCodeBERT. The parameters of code and natural language encoders are initialized by GraphCodeBERT. For training, we randomly select 6 code blocks from the divided code blocks of one long code. The training batch size is 32. For evaluation, we use all divided code blocks of one long code. The evaluated batch size is 256. All experiments are conducted on a machine with Intel Xeon E5-2698v4 2.2Ghz 20-Core CPU and two Tesla V100 32GB GPUs.

To identify the optimal configuration for SEA, we conducted experiments by varying our architecture using different code splitting methods and aggregation methods, while measuring the resulting changes in search performance. Given that the CodeSearchNet Ruby dataset is relatively small, we focused on conducting experiments on the ruby subset, and we present the results in Table 4.

In Table 4 rows SpaceSplitting, we experimented with various aggregation methods as described in Section 4.3. Our findings showed that using any single aggregation method in isolation did not yield significant performance improvements compared to the GraphCodeBERT Baseline. However, upon fusing the attention method with meanpooling, we observed substantial performance enhancement. Specifically, the Attention (one layer) + Mean aggregation method improved MRR and R@1 by 7.9% and 11.5%, respectively. Consequently, for subsequent experiments, we opted to use the Attention (one layer) + Mean aggregation method.

In Table 4 rows SpaceSplitting, TokenSplitting, LineSplitting, ASTSplitting, we explored different code split methods, as detailed in Section 4.2. For space and token-based splitting methods, we set the window size from 64 to 256 due to the finer granularity of division. Conversely, for line and AST-based split methods, we set the window size from 16 to 64. Notably, we observed that the AST-based split method displayed outstanding performance, achieving the highest MRR and R@1 with a window size of 32. As a result, in subsequent experiments, SEA refers to SEA-ASTSplitting with a window size of 32, step size of 16 and the Attention (one layer) + Mean aggregation method.

In this section, we conduct a comparison between SEA and three sparse Transformers, BIGBIRD Zaheer et al. (2020), Longformer Beltagy et al. (2020), and LongCoder Guo et al. (2023). BIGBIRD and Longformer are two well-known long document-oriented Transformers. LongCoder employs a sliding window mechanism to handle long code input for code completion. Specifically, we leverage the bigbird-roberta-base²²2https://huggingface.co/google/bigbird-roberta-base, longformer-base-4096³³3https://huggingface.co/allenai/longformer-base-4096 and longcoder-base⁴⁴4https://huggingface.co/microsoft/longcoder-base models, with a token length of 1024. Due to BIGBIRD and Longformer not being pretrained on the code dataset, we also conducted experiments to initialize BIGBIRD and Longformer with the parameters of GraphCodeBERT. The results are presented in Table 5. Comparing the results before and after initializing BIGBIRD and Longformer with the parameters of GraphCodeBERT, we found that MRR results improved from 0.2952 and 0.5016 to 0.6121 and 0.6595, respectively. We attribute this performance gap to the need for re-pretraining models that were originally pretrained on natural language datasets. We observed that LongCoder’s MRR was 0.4718, which represents a significant decrease compared to GraphCodeBERT, suggesting that LongCoder may be primarily suited for Code Completion tasks. We also conducted t-tests between our SEA and other baselines, and the results demonstrate that SEA significantly outperforms all sparse Transformer baselines ($p<0.01$), highlighting its superior performance in the domain of code search.

In terms of model parameters and search efficiency, SEA stands out as it boasts a lower parameter count and shorter inference time compared to BIGBIRD, Longformer and LongCoder. It’s worth noting that SEA’s parameter count is closely aligned with that of GraphCodeBERT, differing only by the addition of a single attention layer. However, this minor change results in a significant boost in search performance. We also present experimental results without employing the combine-divide method in Table 5. We observed that while the search performance stays stable, the inference time increases by more than threefold. It highlights the considerable improvement in inference time brought about by the combine-divide method, thereby confirming its effectiveness in accelerating the model’s inference process.

To explore the improvement of the proposed SEA for code snippets with varying lengths, we present the search performance comparison between the baseline method GraphCodeBERT and SEA under different ground-truth code token lengths. The results are depicted in Figure 5.

Notably, the retrieval performance of each query subset exhibits noticeable enhancements, particularly for long code retrieval results. We attribute this improvement to two crucial factors. Firstly, the aggregation module of SEA adaptively captures and incorporates information from diverse segments of the long code, leading to a more comprehensive and informative code representation. Secondly, the code splitting method employed by SEA can be viewed as a form of data augmentation, providing additional context and variation that aids in strengthening the code representation. In summary, SEA yields a more robust code representation, significantly enhancing the overall retrieval performance.

To ensure a fair and reproducible comparison, we carefully selected pretraining-based baselines that meet the following three criteria: 1) The source code is publicly available; 2) The overall model is adaptable to all the six programming languages on the CodeSearchNet dataset; 3) The paper is peer-reviewed if it is published as a research paper. Consequently, we select four deep end-to-end approaches: RoBERTa  Liu et al. (2019), UniXcoder Guo et al. (2022), CodeBERT Feng et al. (2020), and GraphCodeBERT Guo et al. (2021).

In Table 6, we present the MRR results, demonstrating that SEA outperforms all methods across all six programming languages. Notably, this conclusion remains consistent for the recall metric and another variant of SEA, the results of which can be found in our replication package. These findings reinforce the superiority of SEA as compared to the pretraining-based baselines across diverse programming languages.