发布时间：2024年03月15日
RAG 软件开发
随着 RAG 技术的最新突破，代码库级别的代码补全步入了一个崭新时代。然而，现行方法中固化的检索机制导致了效率与鲁棒性双重困境，许多检索到的上下文对 code LMs 来说并无裨益，甚至产生负面影响。为此，本论文提出一种创新的选择性 RAG 架构，能够在无需检索的情况下灵活规避这一过程。为支持这一架构的有效运作，我们开发了一种自我监督学习策略，让 code LM 能够精准自评检索对其输出质量的改善效果，并可靠地利用可能存在噪音的检索内容。利用这种兼具选择性检索策略与生成模型功能的 LM，我们的框架在RepoEval、CrossCodeEval等多个基准测试以及一个全新基准上均超越了基于固定检索方式的当前最优提示技术。此外，我们的选择性检索方案还带来了显著的效率提升，最高可达70%的推理加速，且性能不受损害。实验证明，无论是在何种生成模型、检索器还是编程语言环境下，我们的框架都能出色适应并发挥效用。以上成就标志着我们在构建更为精确高效代码库级别代码补全工具的道路上迈出了重要的一步。
答案是否定的，主要有两个原因：
首先，检索到的有效代码一般比较少。在各种代码完成任务中，我们发现标准的RAG方法所进行的检索中，高达80%的检索并没有提升强大代码LM（如StarCoder）的性能，而许多检索还因引入无关信息而降低了性能。
其次，始终进行检索会引入显著的低效性。检索索引的大小与存储库中的行数呈线性增长。对于中等规模的存储库来说，稀疏检索已经与使用3B LM进行完成一样耗时。这种低效性在密集检索、企业级规模的存储库和迭代检索方法中更为明显。
因此，有必要对“不变的检索”假设进行重新评估。这种重新评估不仅对于提高代码完成至关重要，而且还为使用领域特定结构化检索的广泛类似应用提供了宝贵的见解。
为了解决这些问题，我们提出了一个新颖的存储库级代码完成框架，其核心是一个选择性检索机制：系统使用策略主动避免不必要或潜在有害的检索。我们框架的核心是Repoformer，这是一个智能代码语言模型，通过自触发检索增强进行了精细调整，用于稳健的代码完成。
Repoformer的三个核心原则：
面向性能的自我评估。在观察当前文件后，Repoformer判断检索是否会提高预测的质量。
对检索上下文的稳健性。Repoformer学习使用检索到的上下文来提高其输出的质量，并避免由潜在嘈杂的检索信息引起的性能下降。
泛化能力。上述两个能力必须对任何完成粒度、语言和检索器进行泛化。此外，作为一种选择性检索策略，Repoformer应该能够作为其他黑盒代码语言模型的即插即用策略。
针对以上问题，作者设计了一个简单的自监督RAG模拟过程来创建训练数据。
对多样化的代码块和函数体进行抽样，并与检索到的跨文件上下文配对。然后，通过对比使用和不使用检索增强时代码语言模型输出质量，可以获取选择性检索的地面真实标签。
通过这个数据集，设计了一个多任务目标，共同教导代码语言模型自我评估的能力，并使用可选的检索增强稳健地完成代码。
为了支持自我选择的RAG，语言模型需要两个关键能力：准确的自我评估和对检索到的上下文的稳健性。我们设计了一种对比数据标注方案，从公共存储库中挖掘自监督信息，然后使用一种新颖的多任务目标进行微调。
数据构建
如何创建微调的训练数据，三个步骤：
- • 抽样目标行，这些行可以是：（1）不同长度的随机代码块或（2）函数体。
- • 使用当前文件检索上下文，有或没有检索。
- • 标记是否通过将当前文件扩展为使用检索上下文可以将代码语言模型 ℳ 的代码完成质量提高超过阈值 τ，通过编辑相似度与 ℳ 进行比较。
从Stack数据集中随机抽样了18k个Python存储库
- • （1）每个文件至少有三个 import
- • （2）每个文件至少有两个本地导入项
- • （3）至少有五个Python文件
这些标准确保了存在本地依赖项，RAG可以发挥作用。我们使用 ℳ = StarCoderBase-1B 和 τ = 0 标记了24万个代码块完成实例和12万个函数完成实例。我们将500个存储库保留用于验证，其余的用于训练。
使用最先进的单次迭代 prompt pipeline 和 StarCoder模型作为测试基准，因为与CodeGen系列相比，它们性能更强。
- • 无检索。仅在提示中向模型提供查询代码。
- • 固定检索。始终将查询代码和检索到的上下文进行增强。
- • 选择性检索。我们向Repoformer提供查询文本和目标代码，并根据两种选择性RAG策略可选地增强检索上下文：
- • 自我选择。
如果
是后最可能出现的标记，则进行检索。
- • P(
)。如果在后出现的概率大于阈值 τ，则为该实例执行检索增强。
与相同大小的StarCoderBase的无检索和固定检索相比，Repoformer的选择性检索策略表现出强大的性能改进。通过 P() 策略，Repoformer-3B 在大多数任务和度量标准上都能超过StarCoderBase-7B，甚至在API完成的EM方面也能超过，甚至在API和代码块完成的ES (edit similarity) 方面也能超过5倍大的StarCoder。还展示了我们方法的性能改进可以泛化到Python以外的三种语言。
就代码完成准确性而言，基于阈值控制的 P() 策略在所有任务上都优于自我选择策略。
基于作者的研究，未来可以通过将选择性RAG调整为在生成过程中动态调用多次检索，可以进一步提高当前的函数完成性能。通过利用代码库的独特特征来定制检索策略，可以提高我们选择性检索框架的有效性，开发更智能、更高效的基于RAG的代码补全解决方案。
Arxiv[1]
引用链接
[1] Arxiv: https://arxiv.org/abs/2403.10059