LongCoder: A Long-Range Pre-trained Language Model for Code Completion (2306.14893v1)

Abstract: In this paper, we introduce a new task for code completion that focuses on handling long code input and propose a sparse Transformer model, called LongCoder, to address this task. LongCoder employs a sliding window mechanism for self-attention and introduces two types of globally accessible tokens - bridge tokens and memory tokens - to improve performance and efficiency. Bridge tokens are inserted throughout the input sequence to aggregate local information and facilitate global interaction, while memory tokens are included to highlight important statements that may be invoked later and need to be memorized, such as package imports and definitions of classes, functions, or structures. We conduct experiments on a newly constructed dataset that contains longer code context and the publicly available CodeXGLUE benchmark. Experimental results demonstrate that LongCoder achieves superior performance on code completion tasks compared to previous models while maintaining comparable efficiency in terms of computational resources during inference. All the codes and data are available at https://github.com/microsoft/CodeBERT.

• The paper introduces a long-range Transformer model employing sparse attention to efficiently handle lengthy code sequences.
• It leverages window, bridge, and global attention mechanisms to reduce computational overhead and improve code completion accuracy.
• Experimental results on the LCC dataset and CodeXGLUE demonstrate superior performance over dense and sparse Transformer models.

Overview of LongCoder: A Long-Range Pre-trained LLM for Code Completion

The paper introduces LongCoder, a sparse Transformer model designed for efficient code completion tasks, particularly with long code sequences. Code completion aids software developers by suggesting code snippets based on context. However, existing Transformer models face computational inefficiencies with lengthy inputs due to the quadratic complexity of self-attention mechanisms. LongCoder addresses this by leveraging a sparse attention approach, which reduces this complexity to linear.

Key Components and Mechanisms

LongCoder's architecture is distinguished by three core attention mechanisms:

1. Window Attention: This mechanism applies a sliding window approach, focusing on local contexts. By limiting each token to attend only within a fixed window, the model reduces computational overhead, enhancing efficiency.
2. Bridge Attention: Involves inserting bridge tokens throughout the input sequence. These tokens aggregate local information, facilitating global interactions across distant sections of the code, thus allowing efficient context assimilation.
3. Global Attention: This introduces memory tokens, which provide global access to crucial code elements such as imports and function definitions. These tokens enable the model to remember significant elements that have wider scope and potential long-term impact.

Experimental Setup and Results

LongCoder was evaluated on a newly curated dataset, Long Code Completion (LCC), which consists of longer code contexts from Python, Java, and C# repositories. The model outperformed previous solutions on the LCC dataset as well as on the CodeXGLUE benchmark. Noteworthy improvements were observed with an enhanced Exact Match (EM) and Edit Similarity metrics, demonstrating LongCoder's superiority over both dense and sparse Transformer models.

Implications and Future Directions

The proposed LongCoder model is not only efficient in handling extensive code sequences but also illustrates significant advancements in long-range dependency modeling. The integration of code-specific heuristics into sparse attention mechanisms holds promise for developing models capable of cross-file and repository-level code completions. This work encourages further exploration into the scaling of sparse Transformers with larger datasets and model sizes.

Limitations and Considerations

Despite its strengths, LongCoder's pre-training was limited to the CodeSearchNet dataset, less extensive than those available to larger models like OpenAI Codex. The evaluation datasets primarily source their samples from GitHub, raising potential concerns regarding data leakage and fairness in assessments.

In conclusion, LongCoder presents a significant step forward in handling long-range dependencies in code completion tasks, promoting both efficiency and practicality. The methodology and results invite further exploration into leveraging sparse attention models, potentially influencing future developments in AI-driven code generation and analysis tools.