通过选择性检索增强代码库级代码补全
自我监督方法学习何时从代码库中检索上下文信息，使代码补全速度提高70%，同时提升准确性。
研究背景
大型代码语言模型虽然在基于上下文的任意程序函数代码补全方面表现优异，但在处理大型软件开发项目时面临挑战——正确的代码补全可能依赖于代码库中其他位置定义的API调用或函数。
检索增强生成（RAG）通过从代码库中获取相关上下文来解决这个问题，丰富模型的理解并改进其输出。但检索需要时间并会减慢生成速度：这始终是最佳选择吗？
在ICML 2024会议上发表的研究表明，实际上80%的情况下检索并不会提高代码生成质量。
方法创新
为了应对这种低效问题，研究人员对大型语言模型进行微调，以确定检索是否可能有所帮助，并根据答案输出两个特殊标记之一。
数据集构建
通过从开源许可代码库中采样代码，随机掩码代码行，并从代码库其他位置检索相关代码来构建训练数据集。然后比较LLM在有和没有附加上下文情况下对掩码代码的重建效果，并根据检索是否改进生成来标记示例。
检索方法比较
实验了多种从代码库检索上下文信息的方法，包括：
- UniXCoder：使用基于Transformer的语义嵌入匹配代码序列
- CodeBLEU：使用n-gram数据、语法树和代码流语义
最终发现更高效的Jaccard相似度（两个符号序列交集与并集的比率）表现最佳，因此大多数实验采用此方法。
技术实现
微调机制
采用"填充中间"机制：将掩码代码从代码序列中切除，前后部分用特殊标记标识。训练目标由输入字符串组成，掩码代码附加在字符串末尾，同样用特殊标记标识。这种方法允许模型利用掩码代码前后的上下文信息。
双重训练目标
微调过程中设置两个训练目标：
- 正确重建缺失代码
- 准确评估检索信息何时有助于重建
性能评估
准确性提升
与StarCoder等现有模型相比，该方法（称为Repoformer）在RepoEval和CrossCodeEval等多个基准测试中提高了准确性并降低了推理延迟。
延迟优化
在实际"在线服务"场景中展示延迟降低能力。假设工作代码库已建立索引，给定包含当前文件的代码补全请求时，系统同时启动三个过程：
- 使用Repoformer做出检索决策
- 使用代码LMM生成无跨文件上下文的代码补全
- 检索跨文件上下文并使用其生成代码补全
实验结果
选择性检索能够在各种固定选择阈值下同时提高准确性和推理速度。Repoformer作为即插即用的策略模型，能够减少RAG中各种强大代码LLM作为生成模型的推理延迟。
检索决策准确率超过85%，确保仅在增加价值时使用上下文检索。进一步分析表明，该策略提高了Repoformer对检索的鲁棒性，减少了有害检索，增加了通过检索改进的实例。
技术特点
- 自监督方法：所有数据集创建步骤可自动化，无需人工标注
- 选择性检索：仅在必要时进行检索，平衡准确性与速度
- 通用性强：适用于不同模型和代码补全任务
- 效率优化：相比始终检索的方法，推理速度提升70%
该方法为代码补全领域提供了新的技术思路，通过智能检索决策机制实现了性能与效率的最佳平衡。