[论文解读] RepoFusion: Training Code Models to Understand Your Repository

Disha Shrivastava, Denis Kocetkov+3 more|arXiv (Cornell University)|Jun 19, 2023

Software Engineering Research被引用 6

一句话总结

RepoFusion 是一种训练框架，通过使用解码器中的融合（Fusion-in-Decoder）架构，使代码模型能够联合关注编码后的多个仓库级上下文片段（如导入文件、同级文件和父类），从而实现显著提升的单行代码补全性能。尽管参数量仅为 220M，其性能仍优于参数量大 73 倍的 CodeGen-16B-multi 模型，并与参数量为 15.5B 的 StarCoderBase 模型性能相当。

ABSTRACT

Despite the huge success of Large Language Models (LLMs) in coding assistants like GitHub Copilot, these models struggle to understand the context present in the repository (e.g., imports, parent classes, files with similar names, etc.), thereby producing inaccurate code completions. This effect is more pronounced when using these assistants for repositories that the model has not seen during training, such as proprietary software or work-in-progress code projects. Recent work has shown the promise of using context from the repository during inference. In this work, we extend this idea and propose RepoFusion, a framework to train models to incorporate relevant repository context. Experiments on single-line code completion show that our models trained with repository context significantly outperform much larger code models as CodeGen-16B-multi ($\sim73 imes$ larger) and closely match the performance of the $\sim 70 imes$ larger StarCoderBase model that was trained with the Fill-in-the-Middle objective. We find these results to be a novel and compelling demonstration of the gains that training with repository context can bring. We carry out extensive ablation studies to investigate the impact of design choices such as context type, number of contexts, context length, and initialization within our framework. Lastly, we release Stack-Repo, a dataset of 200 Java repositories with permissive licenses and near-deduplicated files that are augmented with three types of repository contexts. Additionally, we are making available the code and trained checkpoints for our work. Our released resources can be found at \url{https://huggingface.co/RepoFusion}.

研究动机与目标

• 为解决现有代码大语言模型在理解仓库级上下文（如导入、父类和文件关系）方面的局限性，这些局限性会导致代码补全不准确。
• 通过在推理过程中训练模型整合来自仓库中多个相关代码片段，以提升代码补全的准确性。
• 证明当使用仓库上下文进行训练时，小型模型可超越大型模型，从而挑战仅凭模型规模决定性能的假设。
• 发布一个新数据集 Stack-Repo，包含 200 个宽松许可的 Java 仓库及三种类型的仓库上下文，以支持未来研究。
• 提供一个与检索机制无关、可复用于多种代码生成任务的可重用框架。

提出的方法

• RepoFusion 使用解码器中的融合（Fusion-in-Decoder, FiD）架构，通过分别编码每个上下文及其周围上下文（不包括目标补全位置），并在解码器中联合关注这些编码后的上下文片段，实现多仓库上下文的融合。
• 每个仓库上下文从相关文件（如导入文件、同级文件或父类文件）中检索，并通过添加固定长度的周围代码进行增强，以保留语法和语义上下文。
• 模型通过结合目标补全位置及其对应的多源仓库上下文，在下一项标记预测目标上进行微调。
• 该框架支持多种上下文类型（如导入、同级、父类），其数量和长度为通过消融实验调优的超参数。
• RepoFusion 采用端到端训练，且对底层检索机制保持无关性，可与多种检索策略无缝集成。
• 该方法使模型能够学习优先处理并融合来自仓库的多样化上下文信号，从而提升预测准确性。

实验结果

研究问题

• RQ1 与未使用此类上下文的更大模型相比，使用仓库级上下文训练的小型代码模型是否能显著提升代码补全性能？
• RQ2 不同类型的仓库上下文（如导入文件、同级文件）对模型性能的贡献如何？
• RQ3 为最大化模型性能，仓库上下文的最佳数量和长度是什么？
• RQ4 初始化策略和训练目标如何影响模型融合多源上下文的能力？
• RQ5 在使用仓库上下文进行训练时，小型模型是否能与未使用此类上下文训练的大型模型性能相当甚至更优？

主要发现

• 参数量为 220M 的 RepoFusion 模型在单行代码补全任务上优于参数量为 16B 的 CodeGen-16B-multi 模型，尽管其参数量仅为后者的 1/73。
• RepoFusion 的性能与参数量为 15.5B 的 StarCoderBase 模型非常接近，后者是使用更复杂的“填空式”（Fill-in-the-Middle）目标训练的。
• 消融实验表明，利用仓库内部多样化来源的信息（如导入文件和同级文件）对性能提升至关重要。
• 仓库上下文的数量和长度显著影响性能，最优配置通过系统性消融实验确定。
• 模型的有效性并非仅源于上下文长度或上下文数量，而在于其学习融合高质量、多样化上下文信号的能力。
• 发布 Stack-Repo 数据集，包含 200 个宽松许可的 Java 仓库及三种类型的仓库上下文，为检索增强型代码模型的研究提供了可复现性与未来研究支持。