Tackling Long Code Search with Splitting, Encoding, and Aggregating

TL;DR

This paper introduces SEA, a method that splits long code into blocks, encodes, and aggregates them to improve Transformer-based code search without retraining, significantly enhancing performance on long code retrieval tasks.

Contribution

SEA provides a simple yet effective baseline for long code search by splitting, encoding, and aggregating code, compatible with existing pretrained models without re-pretraining.

Findings

SEA outperforms existing models on CodeSearchNet benchmark

Achieves 10.1% higher mean reciprocal rank than GraphCodeBERT

Enables effective long code modeling without changing internal model structures

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. Our source code and experimental data are available at: https://github.com/fly-dragon211/SEA.