我正在参加「掘金·启航计划」
导语
本文是RLPG（ICML 2023）论文的后续工作，本文通过结合不同的存储库上下文来提升基本的代码模型的补全能力，实验显示，作者通过基于CodeT5-base（220M）的RepoFusion模型实现了与StarCoderBase（15.5B）相当的性能表现。
- 会议：Arxiv 2023
- 链接：arxiv.org/pdf/2306.10…
摘要
大型语言模型在像GitHub Copilot这样的编码助手中取得了很大的成功，但是它们难以理解存储库中的上下文，并因此会产生不准确的代码完成。最近的研究表明，使用存储库上下文可以很好地解决这个问题。作者提出了一个叫做RepoFusion的框架，以在模型训练中整合相关的存储库上下文。实验表明，使用存储库上下文训练的模型性能明显优于其他较大的代码模型，作者发布了一个包含Java存储库的数据集，并提供了代码和训练好的检查点。
1 简介
现如今大型语言模型（LLMs）在编码助手中变得越来越受欢迎，但是它们往往难以理解代码存储库中的上下文，并因此会产生不准确的代码完成。为了解决这一问题，该论文提出了一个叫做RepoFusion的框架，以在模型训练中整合相关的存储库上下文。实验表明，使用存储库上下文训练的模型性能明显优于其他较大的代码模型。此外，该论文还发布了一个包含Java存储库的数据集，并提供了代码和训练好的检查点。在具体实现方面，论文使用了Fusion-in-Decoder架构来将存储库上下文（Repo Context，RC）组合起来，以生成更准确和上下文感知的代码完成。下图1展示了1个例子，其中补全任务就是根据Surrounding Context来预测Target Hole。
本文的贡献如下：
- 提出了RepoFusion框架，该框架通过学习组合存储库中相关的上下文提示来帮助代码模型进行更好的预测。
- 通过广泛的实验，RepoFusion（一个220M参数模型）明显优于多个基于下一个令牌预测目标进行训练的较大模型，例如CodeGen-16B [24]。
- 进行了彻底的消融研究，例如存储库上下文的性质、它们的长度、存储库上下文的数量和其他训练配置。其中一个关键发现是利用存储库内不同来源的信息是RepoFusion有效的关键。
- 创建和发布了Stack-Repo，这是一个包含200个Java存储库的数据集，具有宽松的许可证和近似去重的文件，并增加了三种存储库上下文。资源可以在以下网址找到：huggingface.co/RepoFusion。
2 使用存储库上下文训练
2.1 Fusion-in-decoder (FID)
FiD直接将检索回来的每个document都与question通过encoder分别编码，然后concat在一起输入decoder生成最终的输出。
2.2 存储库上下文
Shrivastava等人（这里的内容是之前一篇论文的工作www.yuque.com/glamour/xv5… ，建议可以先读一下再来读这篇论文）根据编程语言的语法和语义以及常见编程模式，提出了一组存储库级别的提示提案（Prompt Proposal，PP），利用存储库中跨文件的结构和相关上下文。PP是一个函数，它以目标空白的位置和关联存储库作为输入，并返回称为Prompt Proposal Context（PPC）的字符串作为输出。PPC是通过从特定类别的相关源文件（提示源）中提取特定类型的上下文（提示上下文类型）创建的。之前的工作一共总结了63个提示提案（有关详细信息，请参见Shrivastava等人之前的论文）。本文中，作者按照上下文位于target hole的相对位置，分为prior PPC（前面位置的）和post PPC（后面位置的）。 存储库级别的PP可以被看作是一个确定性的检索机制，它从存储库中返回相关的代码片段。除了PP，为了了解检索到的存储库上下文的作用，作者还考虑了另外两种机制来检索存储库级别的上下文：(a) BM25：使用基于BM25的相似性将存储库中每个文件的上下文进行评分，(b) RandomNN（也在Shrivastava等人中使用）：从一个随机选择的存储库块列表中，根据嵌入块与表示空间中嵌入周围上下文的相似性，选择前k个。
2.3 RepoFusion
RepoFusion的核心思想是训练一个代码模型，使其了解存储库的上下文，从而有助于生成准确的target hole预测。给定一组检索到的存储库上下文，RepoFusion学习使用FiD方法来组合这些上下文的相关部分。Surrounding Context与每个存储库上下文连接，然后独立编码（参见图1，底部）。作者将Surrounding Context附加到Repo Context（RC）的末尾。这与原始Fid不同。RepoFusion使用N个长度为l个token的RC（由于PPC有63种，所以需要进行排序决定优先选哪N个来作为RC，下面的策略中讨论了这个排序问题）。本文尝试了以下四种策略，根据PPC生成和排序RC（见图2）。
- Truncated-Ranked (T-Rank)。在这种设定下，每个PPC只截断后得到1个RC；他们的顺序按Shrivastava等人之前的工作；
- Truncated-Random (T-Rand)。与1类似，同样每个PPC只截断后得到1个RC，但顺序随机；
- Not Truncated-Ranked (NT-Rank)。该情况下不丢弃截断数据，而是把他们顺序的切割后都作为RC候选，这里1个PPC可能得到多个RC，顺序与1一致；
- Not Truncated-Prior-Last (NT-Prior-Last)。与NT-Rank相同，只是prior PPC始终按顺序放置在末尾。由于解码器按照输入的顺序关注存储库上下文的连接编码表示，因此这种策略有助于我们理解从prior PPC的编码表示继续代码生成的作用，因为它是解码器中最近关注的表示。请注意，根据N的值，可能需要删除某些排名靠前的PPC的部分块，以容纳位于末尾的prior PPC。
3 实验和结果
3.1-3.2 数据集构建与实验细节
作者从The stack V1.1数据集中选择了200个Java的存储库，其统计指标如下：
其中对于每个target hole，作者选择其前面两行和后面两行作为Surrounding Context。并选择了CodeT5-base作为RepoFusion的基础模型，并主要对比了以下Baseline：
- CodeT5(FT)：选择CodeT5-base、-large并使用不同的PPC类型和Context长度进行Fine-tune；
- BigCode Model：选择SantaCoder、StarCoderBase这两种使用Fill-in-the-middle(FiM)方式训练的模型进行对比；
- CodeGen：对比了CodeGen-2Bmulti, CodeGen-6B-multi, and CodeGen-16B-multi。
作者发现预训练的CodeT5模型不能很好地完成Java代码。因此，为了获得RepoFusion训练的基础模型，作者使用512的输入上下文长度，在Java存储库上使用Next Token Prediction目标对CodeT5-base进行微调。对于评估指标，作者采取了与之前工作一致的成功率Success Rate（SR）。
3.3 实验结果
实验结果如表2所示，作者主要强调了以下几点发现：
- Model Size和Context Length的增加可以提升Baseline性能。例如，CodeT5-large比CodeT5-base性能更好。采用4096的Context length比2048性能要好；
- RepoFusion是有效的。RepoFusion在相同的有效上下文长度下超越了其他更大的模型，即使被限制使用较少的存储库上下文也比其它模型表现更好。当提供额外的存储库上下文时，RepoFusion的性能还有进一步提升。与StarCoderBase模型相比，RepoFusion的成功率仅略低。
- Prompt Proposal很重要。图3右侧说明在使用T-Rank和NT-Rank时，RepoFusion使用Random-NN、BM25和PPC的SR。结果表明，使用PPC的RC效果最好。
- NT-Prior-Last最有效。图3的左侧说明了四种策略在两个不同的设置中的成功率：N = 32，l = 768和N = 63，l = 512。可以看到，有序的RC，尤其是NT-Prior-Last、NT-Rank和T-Rank比随机排序的RC（T-Rand）表现更好。此外，与T版本相比，NT版本的改进表明，提供来自顶部提示建议的完整上下文的价值，而不是显示更多提示建议的截断上下文。
- 更长的上下文更好。图4的左侧绘制了成功率随存储库上下文长度l变化的变化。结果表明，随着每个存储库上下文的大小增加，性能有所改善。然而，在两种情况下（N = 32，N = 63），随着l的值增加，性能达到饱和点。
- 使用更多的RC效果更好。为了了解多个存储库上下文的作用，作者评估了模型在不同的N值下的表现。从图4的中部可以看出，当l = 512时，RepoFusion的性能随着N的增加而增加，最高达到N = 63，而当存储库上下文更长，l = 768时，性能增加最高达到N = 32。在此之后，增加N的值不会再带来进一步的改进。
- 多样的PP对RepoFusion有帮助。作者研究了每个RC数量N对应的平均不同PPC的数量的成功率。请注意，PPC的数量小于N，因为通常一个PPC会产生多个RC。从图4的右侧可以看出，使用许多不同的PPC对于获得更好的RepoFusion性能至关重要。
- 对CodeT5使用Next Token Prediction进行Fine-tune很重要。表4展示了使用预训练的CodeT5-base模型初始化与使用微调版本进行初始化进行训练时，RepoFusion的评估结果。可以看到，使用RC进行训练可以提高基于预训练的CodeT5模型的性能（原始CodeT5性能参见表7）。作者发现，在所有情况下，使用微调的模型进行训练都具有明显的优势。
4 相关工作
Shrivastava等人提出了RLPG，一种基于target hole选择PP的分类器，并利用所选PP和先前上下文来提示Codex。类似地，RepoCoder通过注入LLM的先前预测以及检索到的上下文和输入提示中的先前上下文，迭代地细化目标空洞的预测。在这项工作中，本文借鉴了Shrivastava等人有关在推理过程中利用存储库级别信息的见解，将他们的发现扩展到涉及不同代码语言模型（LLMs）的各种配置中，考虑到不同的上下文长度和大小。此外，本文的框架是使用RC进行训练的，并学习有效地利用从存储库获取的多个相关上下文。
5 讨论
RepoFusion随着存储库上下文数量的增加而呈线性扩展，这可能导致推理时间较慢。生成的代码难以理解或调试，可能存在安全隐患。过度依赖这些模型可能会导致用户忽略代码中的错误或过于自信，引入错误。部署RepoFusion需要仔细考虑，并有可能需要使用优化技术以提高推理速度。虽然RepoFusion在单行代码自动完成方面表现出色，但在其他任务中的性能需要进一步研究。
