Tackling Long Code Search with Splitting, Encoding, and Aggregating (2024)

Abstract: Code search with natural language helps us reuse existing code snippets. Thanks to the Transformer-based pretraining models, the performance of code search has been improved significantly. However, due to the quadratic complexity of multi-head self-attention, there is a limit on the input token length. For efficient training on standard GPUs like V100, existing pretrained code models, including GraphCodeBERT, CodeBERT, RoBERTa (code), take the first 256 tokens by default, which makes them unable to represent the complete information of long code that is greater than 256 tokens. To tackle the long code problem, we propose a new baseline SEA (Split, Encode and Aggregate), which splits long code into code blocks, encodes these blocks into embeddings, and aggregates them to obtain a comprehensive long code representation. With SEA, we could directly use Transformer-based pretraining models to model long code without changing their internal structure and re-pretraining. We also compare SEA with sparse Trasnformer methods. With GraphCodeBERT as the encoder, SEA achieves an overall mean reciprocal ranking score of 0.785, which is 10.1% higher than GraphCodeBERT on the CodeSearchNet benchmark, justifying SEA as a strong baseline for long code search.

## tokens return a tensor containing the patches. graphcodebert truncates tokens beyond 256 tokens. key tokens are highlighted in yellow. as depicted in figure 2 , the process involves splitting the long code into a set of code pieces, followed by utilizing the sliding window method to generate a partially overlapping code block set. existing code encoders are then used to obtain embeddings for each code block. finally, these embeddings are aggregated to generate representations for the entire long code. through extensive experiments, we have found that the proposed astbased splitting method and attention-based aggregation method outperform other techniques for splitting and aggregation. due to the varying numbers of code blocks obtained from different code snippets, parallel operation becomes challenging. to address this problem, we have designed a combinedivide module for acceleration. it is important to note that sea is encoder-agnostic, meaning it can be used with different transformer-based encoders. when compared to various transformer-based encoder baselines, sea achieves a significant improvement in mean reciprocal ranking (mrr) performance, ranging from 7% to 10%. we also design a combine-divide module for acceleration.