通过选择性检索增强仓库级代码补全
一种自监督方法，用于学习何时从代码仓库中检索上下文信息，可将代码补全速度提升70%，同时提高准确性。
作者：Di Wu
2024年10月17日
4分钟阅读
用于代码的大语言模型是在源代码而非自然语言文本上预训练的模型。它们在仅基于上下文完成任意程序函数的代码方面表现出色。然而，在处理新的、大型的软件开发项目时，它们会遇到困难，因为正确的代码补全可能依赖于代码仓库中其他地方定义的API调用或函数。检索增强生成通过从仓库中获取相关上下文来解决此问题，从而丰富模型的理解并改善其输出。但执行检索需要时间并减慢生成速度：它是否总是最佳选择？
在今年的国际机器学习大会上发表的一篇论文中，我们研究了这个问题，并发现实际上，在80%的情况下，检索并不能提高代码生成的质量。
![上下文检索对模型性能的影响。橙色条表示无变化。]（图片占位符）
为了解决这种低效问题，我们微调了一个大语言模型，使其能够判断检索是否可能有帮助，并根据答案输出两个特殊令牌之一。
![有无上下文检索的代码补全对比。]（图片占位符）
对于微调，我们使用了一个数据集，该数据集通过从开放许可仓库中采样代码、随机屏蔽代码行以及从仓库其他地方检索相关代码构建而成。然后，我们比较了大语言模型在有额外上下文和无额外上下文的情况下对屏蔽代码的重建结果。随后根据检索是否改善了生成来标记这些示例。
在实验中，我们发现，在代码补全任务上，使用我们的数据集进行微调的代码大语言模型甚至比始终执行检索的模型表现得更好——并且由于选择性检索，推理速度提高了70%。在论文中，我们还报告了广泛的实验，旨在证明我们的方法能够很好地泛化到不同的模型和不同的代码补全任务。
方法
创建数据集的所有步骤——采样和屏蔽代码、检索相关代码以及在有和无检索上下文的情况下进行代码生成——都可以自动化，这使得我们的方法是自监督的：它不需要人工标注，并且可以扩展到任意大的数据集。
![相关内容]（图片占位符）
我们尝试了多种从仓库检索上下文信息的方法，包括使用基于Transformer的语义嵌入来匹配代码序列的UniXCoder，以及使用n-gram数据、语法树和代码流语义的CodeBLEU。然而，这两种方法的性能都没有超过效率高得多的Jaccard相似度，即两个符号序列的交集与并集之比。因此，在我们的大部分实验中，我们使用Jaccard相似度进行检索。我们推测，通过使用具有结构感知分块而非固定行分块的语义检索，可以获得更好的性能。我们将此留作未来的工作。
对于模型微调，我们使用了“中间填充”机制，在该机制中，被屏蔽的代码从代码序列中被移除，并且其前后部分用特殊令牌标识。训练目标由输入字符串组成，该字符串的末尾附加了被屏蔽的代码，同样用特殊令牌标识。这使得模型能够利用被屏蔽代码前后的上下文信息；已有研究表明，这比训练模型将被生成的代码插入到前后部分之间能产生更好的结果。
在微调过程中，我们有两个训练目标：正确重建缺失的代码，以及准确评估检索到的信息何时有助于重建。
准确性评估
与现有模型（如StarCoder）相比，我们的方法（我们称之为Repoformer）在各种基准测试中提高了准确性并减少了推理延迟，这些基准测试包括RepoEval和CrossCodeEval（一个针对长形式代码补全的新基准）。
![模型性能，根据精确匹配、编辑相似度和单元测试通过率衡量。SelectiveG（其中“G”代表“贪心”）如果最可能的下一个令牌来自仓库的其他地方则执行检索；SelectiveT仅在可能性超过某个阈值时执行检索。]（图片占位符）
延迟评估
我们在一个现实的“在线服务”环境中展示了Repoformer减少延迟的能力。我们假设工作仓库已经被索引。给定一个包含当前文件的代码补全请求，系统会同时启动三个进程：
- 使用Repoformer做出检索决策；
- 使用一个代码大语言模型在无跨文件上下文的情况下生成代码补全；
- 检索跨文件上下文并使用它生成代码补全。
在一系列固定的选择阈值中，Repoformer的选择性检索能够同时提高准确性和推理速度。其性能在广泛的阈值设置下也能保持。
![十亿参数Repoformer模型的自选择性RAG的延迟-准确性权衡。]（图片占位符）
![当十亿参数的Repoformer作为选择性RAG的策略模型时，更大规模代码大语言模型的准确性和延迟。“SU”代表“加速比”（相对于始终检索）。]（图片占位符）
![对Repoformer放弃检索实例的分析。深蓝色表示模型在无RAG的情况下生成了正确输出；浅蓝色表示模型生成了错误输出，但RAG并未改善性能；红色表示模型生成了错误输出，而RAG本可以帮助。]（图片占位符）
更有趣的是，Repoformer能够作为一个即插即用的策略模型，在RAG中作为生成模型，降低各种强大的代码大语言模型的推理延迟。
Repoformer在检索决策中的准确率超过85%，确保了上下文检索仅在增加价值时才被使用。
进一步的分析表明，所提出的策略提高了Repoformer对检索的鲁棒性，减少了有害的检索，并增加了通过检索得到改善的实例。
致谢
我们非常感谢Wasi Uddin Ahmad和Dejiao Zhang作为本项目导师所做的贡献。从项目构思到定期会议中的诸多出色建议，他们的指导带来了巨大的改变。我们还要感谢其他合著者和匿名的ICML审稿人提供的宝贵反馈，这些反馈确实帮助我们改进和完善了这项工作。
研究领域 对话式人工智能
标签 代码生成，检索增强生成
会议 ICML 2024
相关出版物 REPOFORMER: Selective retrieval for repository-level code completion
关于作者 Di Wu是加州大学洛杉矶分校计算机科学专业的博士生。他在完成这项工作期间是某机构的实习生。FINISHED
