[论文解读] LongCoder: A Long-Range Pre-trained Language Model for Code Completion

Daya Guo, Canwen Xu+3 more|arXiv (Cornell University)|Jun 26, 2023

Software Engineering Research被引用 7

一句话总结

LongCoder 提出一种稀疏 Transformer 模型，通过引入桥接令牌（bridge tokens）和记忆令牌（memory tokens），高效处理长代码序列的代码补全任务，在保持线性计算复杂度的同时，实现了长上下文任务的最先进性能，且推理效率与更大模型相当。

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.

研究动机与目标

• 解决标准 Transformer 在长代码文件中进行代码补全时面临二次方计算成本的挑战。
• 通过实现对导入语句和函数定义等关键代码元素的全局访问，提升模型在长上下文代码补全任务中的性能。
• 设计一种高效的注意力机制，在保持线性复杂度的同时捕捉代码中的长距离依赖关系。
• 构建一个新的基准数据集（LCC），以支持长代码建模研究，并在真实、长上下文的代码示例上评估模型性能。
• 通过有效记忆和检索全局相关的代码结构，实现在整个文件或项目范围内的可扩展代码补全。

提出的方法

• 引入滑动窗口机制用于自注意力，以限制局部计算并降低复杂度。
• 在规则间隔处插入桥接令牌，聚合局部上下文信息，实现跨窗口的长距离信息流动。
• 引入记忆令牌，实现对高影响力代码元素（如导入语句、类和函数定义）的全局注意力。
• 设计一种稀疏注意力机制，结合局部窗口注意力、桥接令牌聚合和全局记忆注意力，实现高效的长上下文建模。
• 在代码语料库（CodeSearchNet）上预训练 LongCoder，采用掩码语言建模目标以学习上下文表征。
• 在代码补全基准（包括新构建的 LCC 数据集和 CodeXGLUE）上微调 LongCoder，以评估下游性能。

实验结果

研究问题

• RQ1 稀疏注意力机制是否能在保持线性计算复杂度的同时，有效建模长代码序列的代码补全任务？
• RQ2 与标准 Transformer 相比，引入桥接和记忆令牌在长上下文代码文件的代码补全准确率上有多大提升？
• RQ3 使用更长上下文在多大程度上提升了代码补全性能？这种提升是否能在不增加推理成本的前提下实现？
• RQ4 在标准上下文窗口之外的长上下文代码示例上评估时，LongCoder 与大型模型（如 Codex）相比表现如何？
• RQ5 通过架构创新，一个在较小预训练语料上训练的模型是否仍能在长代码补全任务上实现优异性能？

主要发现

• LongCoder 在 LCC 数据集上达到最先进性能，该数据集包含的代码示例平均长度是现有基准的五倍。
• 在 LCC 基准上，LongCoder 超越了标准模型和 Codex-512，证明即使模型规模小于 Codex，长上下文建模仍具显著优势。
• 仅 LongCoder 和 Codex-2048 能够正确补全依赖于文件顶部远处导入或定义的代码（如时间戳函数）。
• 在 Java 案例研究中，LongCoder 能够从局部窗口之外的类定义中正确预测方法调用，而 Codex-512 因上下文窗口限制而失败。
• 尽管处理长上下文，LongCoder 仍保持与较小模型相当的推理效率，这得益于结合桥接和记忆令牌的稀疏注意力机制。
• 通过记忆令牌实现对全局相关代码元素的记忆与检索，使模型在长距离依赖关系上的泛化能力优于仅依赖窗口的方法。