代码补全相关论文 ReACC: A Retrieval-Augmented Code Completion Framework

一些术语：https://developer.baidu.com/article/details/2686964

Intro

当前对于代码补全的方法仅仅关注文件或者项目内部的代码上下文，但是研究表明程序员倾向于复制来重用现有的代码，所以文章中将输入范围（information scope？）从仅输入文件扩大到了大型的代码库。作者提出 ReACC – a Retrieval-Augmented Code Completion framework，并推测使用具有相同语义的代码作为辅助信息有利于预测后续的code tokens。

在作者的框架下，代码补全任务可以被重新制定（reformulated）为：给定一个用于搜索的源代码库和一份未完成的代码片段，利用未完成的代码片段作为query在源代码库中检索相似的代码片段，然后利用检索到的代码片段来预测后续的code tokens。

ReACC包含两个核心的组件：

用于代码-代码检索的双编码器模型（dual-encoder model）
作为代码补全生成器的自回归语言模型。ReACC采用阶段性训练策略（Stage-Wise Training Strategy）

最简单的的代码检索方式就是构建稀疏向量检索器，例如TF-IDF，BM25，这两者都是基于关键词匹配算法。

稀疏检索器能够捕获词汇信息，并且对于代码标识符是敏感的。相反，密集检索器通过将代码片段转换成密集向量，能够捕获语法和语义的信息。

在代码补全任务中，代码检索器需要理解源代码的意图来检索语义上相似的代码。另一方面，考虑到程序员是很容易去复制粘贴已有的代码，检索器也应该能够评估词汇的相似性。为此，作者采用了混合式的检索器，结合了稀疏和密集检索器。由于交叉编码器具有更高的计算复杂度，作者使用双编码器模型架构作为密集检索器。同时为了达到更好的理解能力，作者使用GraphCodeBERT来初始化密集检索器，这是一个基于BERT的预训练的编程语言理解模型。然后继续利用对比学习来预训练检索器以增强句子嵌入（sentence embedding）。As the labeled data containing similar code pairs is rare, we utilize various transformations to generate programs with similar functionality for data augmentation.（这一句没想好怎么翻译理解）

作者使用decoder-only transformer model来实现了生成器。为了合并来自检索到的相似代码的外部信息，将获取到的代码和代码上下文拼接起来作为输入。生成器由CodeGPT-adapted初始化，这是一个在代码语料库上预训练的GPT2的领域自适应模型（domain-adaptation model）。

作者在两个基准数据集上评估了ReACC框架，CodeXGLUE 和 CodeNet ，采用Python和Java语言。ReACC 在两个数据集上都获得了最先进的性能。实验结果表明，检索器检索到的外部源代码对于自动完成部分代码是有用的。

作者的主要贡献可以总结为

作者提出了一种检索增强方法协助代码自动完成任务。
为了适应代码补全场景，其中检索查询是未完成的代码片段，我们提出部分code-to-code搜索任务并创建用于评估的数据集。
我们采用语义保留转换进行数据增强来预训练代码检索模型。

ps：周三晚上听完会议就记得听到了很多BM25，终于看到是用来干什么的了

Approach

Task Formulation

假设我们有一个包含大量软件库的源代码数据库，其由D个源文件$f_1,f_2…f_D$构成。根据Dense Passage Retriever（DPR）模型，我们将每个文件分为相同长度的代码片段作为基础的检索单元。这样的分割方法不仅能带来更好的检索效果，而且支持了极长的代码文件，其中文件的每个部分都代表了不同的语义。

因此，我们得到M个代码片段作为数据库$C={c_1,c_2,…,c_M}$，让$X={x_1,x_2,…,x_k}$作为提前写好的未完成代码，一个检索器$R: (X,C)\rightarrow C$检索出C中最相近的代码片段$C_s$。然后基于上下文合检索到的代码，生成器$G$预测接下来的code tokens $Y={x_{k+1},…,x_{k+n}}$，其中token级别的代码补全任务中n=1。
$$
P(Y) = \prod_{i=1}^{n}{P(x_{k+i}|c_s, x_{1:k+i-1})}
$$

Retriever

ReACC的检索模型应当能够在给定不完整的代码的情况下，检索出语义上相同的代码。作者采用结合稀疏合密集检索器的得分（score）的混合检索器。

其中使用的稀疏检索器是基于ElasticSearch实现的BM25。作为term-based的检索方法，BM25将每个代码片段视为code token序列，并且采用词袋表示。由BM25计算的匹配得分指出了查询和文档（document）之间的词汇上的相似性。

至于密集检索器，它将每个代码片段映射为一个d纬度的密集向量。作者在这篇文章中基于DPR模型来构建了密集检索器。

Dense Retriever

下图展示了ReACC的密集检索器的训练过程

该密集检索器包含两个基于双向transformer的编码器$E_C$和$E_Q$。$E_C$对于检索数据库$C$中的每个代码片段编码并创建索引。而查询由$E_Q$编码。用[CLS] token作为输出，并以$sim(q,c) = E_C©^TE_Q(q)$来计算相似度。

由于$E_C$和$E_Q$都将源代码作为输入且他们唯一的区别在于他们是否是部分的（partial？如何理解），双编码器在ReACC中是共享权重的。

在训练阶段，根据DPR，作者采用in-batch negatives策略来计算InfoNCE的对比损失。
$$
L(q,c^+,c_1-,c_2^-,…,c_m-) = -log \frac{e^{sim(q,c+)}}{e^{sim(q,c+)}+\sum_{i=1}^{m}e{sim(q,c^+)}}
$$
然而和DPR不同，作者并没有使用从BM25检索到的难分负样本（hard negatives）。因为程序员趋向于直接复制token，所以具有不同语义但是大量词汇相似性的代码能更好的帮助代码补全。

Data Augmentation

ReACC密集检索器的对比学习的目的是学习如何表示代码片段，使得其更靠近相似代码或者是等效语义，远离（far apart）不相似的代码。这要求大量的正负代码对。然而，基于未标记的语料库来识别相似的程序是困难的，特别是从GitHub仓库中挖掘的大量使用的数据集。

搜索语义相等的代码需要额外的代码编译和运行开销，这在大型数据库中是不现实的。除了搜索以外，==代替方法是创建具有相同功能的代码片段来进行数据增强。==为了实现这种方法，作者对源代码进行了一些语义保留的转换来构建了一系列变体。

作者主要采用标识符重命名和死代码（无法访问或未使用的代码）插入的方法。下图展示了对python代码实施这种转换。

Identifier renaming

是将一个标识符重命名为另一个的方法。只重命名变量和方法名称，因为其他标识符不能像内置类型或 API 调用一样任意更改。和以前的工作不一样的是，作者在修改名称时，保留了部分的词汇信息，这是基于标识符名称通常传达含义以及词汇相似性对检索有很大贡献的考虑。作者标记了程序中的所有标识符，利用GraphCodeBERT来预测。选择前10个预测作为重命名的参考集。

Dead code insertion

将dead code插入到代码片段的正确的位置。

死代码是永远无法到达的代码片段或者可以到达但其结果永远不能用于任何其他计算。

在软件工程中，dead code插入是最常见的代码混淆技术之一，其目标是修改代码使其难以理解，但保留其功能，这与我们的目标类似。作者首先随机挑选程序中没有出现过的变量名，然后使用它们从一组预定义的死代码中形成语句，例如赋值、方法调用、循环语句、条件语句等。遍历 AST 并识别所有语句，然后随机选择一个语句并在其后面插入死代码，从而在 AST 中生成一个新的子树。

Input Format

将代码token序列和API调用序列集成作为输入，API调用序列和代码片段的功能高度相关。为了提高代码表示，提取API序列并将它附加到源代码序列中。

最后，在训练过程中使用原始代码的随机截断作为查询，并使用整个创建的程序作为正例来解决如何基于不完整语义进行检索的问题。

Generator

检索器的输出为检索代码$C_S$，考虑到$C_S$是由代码上下文$x$查询到的，然而目标是$x$后的代码，所以作者提出了片段对齐——使用同一文件中得$C_S$的下一个片段$C_S^{{'}$（在前文，为了检索已经将每个文件分割为了代码片段）来对$x$的下一个片段进行补全。因此，生成器的输入是$C_S}{‘}$和$x$的连接，即$x^{'}=C_S{’}\bigoplus x$

ReACC的生成器模块支持任意的能完成代码补全任务的模型架构。作者采用的是CodeGPT-adapted，这是一个通过因果语言模型（causal language model）使用来自CodeSearchNet的Python和Java数据集预训练的（单or仅）编码器的转换模型。在两个广泛使用的代码补全数据集上，CodeGPT-adapted在CodeXGLUE 基准上的代码补全任务上达到了预期的效果。

Experiments： Code Clone Detection

为了评估ReACC中code-to-code检索模型的效率，实施代码克隆检测任务，目的是检索语义等价的程序。

Dataset

数据集包括大量从Online Judge网站上获得的程序。分别从CodeNet创建Python和Java的0样本设置的代码克隆检测评估数据集。

对于相同问题的解决方法可以认为是语义相等的。

Retrieval Training Set

ReACC的密集检索器在CodeSearchNet数据集，一个从GitHub仓库中获取的大规模源代码数据集，上进行的预训练。作者使用了其中的1.6M Java方法和1.2M Python方法。

Baseline Methods

CodeBERT

CodeBERT是一个针对编程语言预训练的模型，它在来自 CodeSearchNet 数据集的NL-PL对上进行训练，涵盖了六种编程语言。

GraphCodeBERT

GraphCodeBERT也是在 CodeSearchNet 的NL-PL对上进行预训练的，并且考虑了代码的内在结构，即数据流。

Experiment Setup

检索编码器使用 GraphCode BERT 进行初始化。它通过MLM（Masked language model）客观和对比学习进行持续预训练。

使用批内负样本，批次大小为256。使用学习率为5e-5，我们分别对Python和Java的检索器进行训练，每个训练30个周期。

作者通过部分搜索的方式完成代码克隆检测实验，这种方式理想地适用于代码补全场景，因为它接受部分程序作为查询，同时保持相同的目标。

Results

下图显示了 CodeNet 数据集上零样本代码克隆检测任务的结果，采用部分搜索设置。

模型通过 MAP@K（即前K个的平均精度）进行评估，这是 CodeXGLUE 克隆检测任务中的评估指标，同时也通过1处的精度进行评估，因为在代码补全中我们只关心最相似的代码。

对比其他基于transformer的编码器，可以发现CodeBERT 和 GraphCodeBERT 很难检索到等价的代码。而ReACC的模型显著优于它们，这表明即使查询的语义不完整，我们的模型也能够检索到语义等价的代码。

作者也发现 BM25 在此任务中表现出色，这与开放领域 QA、代码摘要等其他任务的表现有很大不同。这些发现表明语义相关的代码很可能是词汇相似的，这会导致词汇相似对检索的贡献更多，这使得使用基于术语的检索方法进行code-to-code搜索比text-to-code或question-to-passage搜索更容易。

Experiments：Code Completion

这部分评估的是端到端的代码补全。

Dataset

CodeXGLUE

这是一个基准数据集，包含用于10个多样化代码智能任务的14个数据集。作者在代码补全任务中使用了其中的 PY150 数据集和 Java 的 GitHub Java Corpus 数据集。

Baseline Methods

CodeGPT/CodeGPT-adapted

CodeGPT/CodeGPT-adapted都在 CodeSearchNet 的 Python 和 Java 数据集上进行了预训练。CodeGPT 是从零开始训练的，而 CodeGPT-adapted 是一种领域适应模型，初始化自 GPT-2。

PLBART

PLBART基于 BART 架构，采用去噪序列到序列（Seq2Seq）预训练，并在编程语言（PL）和自然语言（NL）的未标注数据上进行预训练。

CodeT5

CodeT5也是一个编码器-解码器预训练模型，它采用了T5架构，并考虑了代码中识别符感知的标记类型信息。

X-CodeGPT

X-CodeGPT 是 CodeGPT 的一个变体，它将 eWASH 适配到 CodeGPT。Clement 等 (2021) 提出了 eWASH，一种利用源代码语法层次结构的方法，以扩大模型在文件中的视野，并在 CodeXGLUE 代码补全任务上实现了新的最先进性能。作者重现了他们的方法，并通过将 eWASH 适配到 CodeGPT-adapted 开发了 X-CodeGPT。

Experiment Setup

Fine-tune - 微调

作者分别在PY150和GitHub Java Corpus数据集上对CodeGPT-adated进行微调，并将它作为ReACC的生成器。训练PY150的epoch数（epoch数是一个超参数，它定义了学习算法在整个训练数据集中的工作次数）为30，Java语料库为10，batch size（批大小，是指在每一次模型权重更新时所使用的样本数量）为96，学习率为2e-5。

除了X-CodeGPT，所有其他基线模型都使用相同的设置进行微调。

对于 X-CodeGPT，我们使用从 CodeSearchNet 中提取的 eWASH 格式的训练集对其进行预训练，其中每个示例都是一个函数体及其相应的扩展上下文。由于 eWASH 需要将代码解析为 AST，但 CodeXGLUE 中的代码已被标记化且无法解析，因此我们从 PY150 构建了一个新数据集，以在 CodeXGLUE 上微调 X-CodeGPT。因此，作者下载 PY150 中的原始文件并创建一个保留训练/有效/测试拆分的新数据集

Evaluation

作者进行了两种代码完成场景，即token级和line级完成，以衡量模型预测一个或多个token的能力。困惑度是 token-level补全的评估指标，而精确匹配准确度 (EM) 和编辑相似度则用于line-level补全。对于token-level的补全，基于效率的考虑，不是在每一步都进行检索，而是在预测前100个token后根据当前上下文检索相似的代码，并利用它进行进一步的预测。

Retrieval Database

使用PY150的训练集和Java语料库作为检索数据库进行测试。我们不使用对比预训练语料库（即 CodeSearchNet），以避免 CodeXGLUE 和 CodeSearch Net 之间的重复，因为它们都是从 GitHub 中提取的。

Hybrid Retriever

BM25得分和密集检索器的线性组合形成了一个混合检索器。具体来说，可以用$sim(q,c)+\alpha *BM25(q,c)$来计算分数，并且基于 PY150 和 Java Corpus 数据集的dev set结果将$\alpha$设为0.9。

Results

下面的两张表比较了不同的baseline model在CodeXGLUE Python和Java数据集上的代码补全任务。使用混合检索器的 ReACC 框架在所有数据集上始终优于其他基线，这证明了我们的猜想：“外部”上下文有利于代码补全任务。

Table4中，和X-CodeGPT进行对比表明了利用外部上下文可能比充分利用当前的代码文件更加有效。在 ReACC 检索器的三种配置中，混合检索器在几乎所有指标上都表现最好，除了 PY150 新测试集中的精确匹配分数。

Table3中，可以观察到，比较两个数据集，PY150数据集的增强大于Java Corpus数据集。原因是 Java 的检索数据库（即训练集）要小得多。 CodeXGLUE Java 语料库数据集仅包含 12,934 个用于训练的文件，因此从它们中检索类似代码更加困难。

另一个发现是 BM25 显示了与密集检索器相当的结果，甚至在困惑度和精确匹配指标方面表现更好。研究结果表明，代码完成任务可以从语义和词汇上相似的代码中受益。

Analysis

ReACC in specific domain

PY150和Java corpus数据集都是从分布在广泛领域的 GitHub 存储库中提取的。由于有些人经常在更具体的领域编写代码，例如 Kaggle 用户的数据挖掘/模式识别领域、ACM 社区的算法领域等。为了在特定代码域中评估 ReACC，我们从 CodeNet 构建了一个代码补全 Python 数据集，可以在算法域中考虑。下表显示，ReACC 在编辑相似性和精确匹配方面明显优于 CodeNet 中采用的 CodeGPT，分别提高了 10% 和 18% 的绝对值。

研究结果表明，ReACC 在特定领域更有效。我们还注意到，在 CodeNet 中，具有密集检索器的 ReACC 表现优于 BM25。这可以通过以下事实来解释：在算法领域，语义相似的代码可能比词法相似的代码更有价值。

Ablation study

通过逐步移除系统的一部分来评估该系统的贡献

为了进一步了解我们的训练选项如何影响模型性能，我们进行了消融实验。如下表所示，当消除检索器或生成器中的数据增强和训练策略时，指标会下降。其中最重要的因素是查询截断。比较两种语义保留转换，标识符重命名比死代码插入贡献更大。当从生成器中删除片段对齐时，即使用检索到的代码片段本身作为生成器时，性能会略有下降。

ReACC vs GitHub Copilot

GitHub Copilot 是一种强大的代码补全技术，它使用 OpenAI Codex作为模型后端。我们在 VSCode 中使用其扩展运行了一些定性示例，如附录 B 所示。值得注意的是，Codex 比 CodeGPT 更强大，因为它是一个大规模预训练模型，基于 GPT3 在 GitHub 中的所有源代码上进行了训练。然而，在某些情况下，使用 CodeGPT 作为生成器的 ReACC 的性能优于 Copilot。并且在 Figure 6 中，Copilot 本身在利用 ReACC 的检索器时也可以从 ReACC 中受益，这表明检索增强方法对于强生成模型的有效性。

Conclusion

作者提出了 ReACC，这是一种检索增强的代码补全框架，它通过从现有代码库中检索语义和词汇上相似的代码，利用“外部”上下文来完成代码补全任务。我们预训练一个双编码器作为部分代码搜索的检索器，它检索给定部分代码的代码片段。我们的方法可以采用任何能够执行代码完成的架构作为生成器。在 CodeXGLUE 基准测试中，ReACC 在代码完成任务中实现了最先进的性能。