[論文レビュー] 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}.

研究の動機と目的

• 既存のコード LLM がインポートや親クラス、ファイル関係といったリポジトリレベルのコンテキストを理解できないという制限を解消し、不正確なコード補完を回避すること。
• 推論時にリポジトリ全体の関連する複数のコードスニペットを統合できるようにモデルを学習させることで、コード補完の正確性を向上させること。
• リポジトリコンテキストを用いて学習させた小さなモデルが、より大きなモデルを上回ることを示し、モデル規模が性能を決定づけるという前提に疑問を呈すること。
• 今後の研究を支援するため、200の許可付きライセンスの Java リポジトリと3種類のリポジトリコンテキストを含む新しいデータセット Stack-Repo を公開すること。
• リtrieバル手法に依存しない再利用可能なフレームワークを提供し、さまざまなコード生成タスクに応用可能にする。

提案手法

• RepoFusion は、Fusion-in-Decoder (FiD) アーキテクチャを用い、各コンテキストを別々にエンコードし、その後、デコーダー内でそれらを統合的に参照することで、複数のリポジトリコンテキストを統合的に扱う。
• 各リポジトリコンテキストは、関連するファイル（例：インポート済みファイル、サブファイル、親クラス）から取得され、文法的・意味的コンテキストを保持するために、固定長の周囲コードを付加して拡張される。
• モデルは、ターゲットの穴とそれに対応するマルチソースのリポジトリコンテキストを組み合わせて、次トークン予測の目的関数に基づいて微調整される。
• フレームワークは、インポート、サブファイル、親クラスなど、複数のコンテキストタイプをサポートしており、コンテキストの数や長さは、アブレーションスタディにより最適化されたハイパーパramータである。
• RepoFusion はエンドツーエンドで学習され、下位のリトリーブメカニズムに依存しない設計となっており、さまざまなリトリーブ戦略との統合が可能である。
• この方法により、モデルはリポジトリ内の多様なコンテキスト信号を優先的に統合する能力を学習し、予測の正確性が向上する。

実験結果

リサーチクエスチョン

• RQ1 リポジトリレベルのコンテキストを活用して学習した小さなコードモデルが、そのようなコンテキストを持たないより大きなモデルと比較して、著しくコード補完性能が向上するか？
• RQ2 インポート済みファイルやサブファイルなどの異なる種類のリポジトリコンテキストが、モデル性能にどのように寄与するか？
• RQ3 モデル性能を最大化するための最適なコンテキストの数と長さは何か？
• RQ4 初期化戦略と学習目的が、マルチソースコンテキストを統合する能力にどのように影響するか？
• RQ5 リポジトリコンテキストを用いて学習した小さなモデルが、そのようなコンテキストを用いないより大きなモデルと同等またはそれ以上の性能を発揮できるか？

主な発見

• 220Mパラメータの RepoFusion は、16Bパラメータの CodeGen-16B-multi よりも、単一のコード補完タスクで優れた性能を示しており、73倍も小さいにもかかわらずである。
• RepoFusion は、より複雑な Fill-in-the-Middle 目的で学習された 15.5Bパラメータの StarCoderBase モデルと、ほぼ同等の性能を達成している。
• アブレーションスタディの結果、インポートやサブファイルといったリポジトリ内多様なソースからの情報を活用することが、性能向上に不可欠であることが示された。
• コンテキストの数や長さは性能に顕著な影響を及ぼし、体系的なアブレーションにより最適な設定が特定された。
• モデルの効果は、コンテキスト長さやコンテキスト数そのものではなく、統合する能力を学習したコンテキスト信号の質と多様性に起因している。
• 200の許可付きライセンスの Java リポジトリと3種類のリポジトリコンテキストを含む Stack-Repo の公開により、リtrieバル拡張型コードモデル分野における再現性と今後の研究が可能になった。