LocAgent: Graph-Guided LLM Agents for Code Localization
- GitHub issueなどの自然言語の説明から、修正すべきファイル、クラス、関数を探す code localization の論文である。
- LocAgent はコードベースを有向の異種グラフとして索引化し、LLM agent が検索、グラフ探索、実体取得の三つの道具で対象箇所をたどる。
- SWE-Bench Lite と新しい LOC-BENCH で評価し、特にファインチューニングした Qwen2.5-Coder-32B が低い推論費用で高い局所化精度を示した。
Abstract(日本語訳)
コード局所化、すなわちコードベースのどこを変更すべきかを正確に特定することは、ソフトウェア保守における基本的でありながら難しい課題である。既存手法は、関連するコード片を特定する際に、複雑なコードベースを効率よく探索することに苦労している。難しさは、自然言語で書かれた問題説明と対象となるコード要素を結びつける点にあり、多くの場合、階層構造や複数の依存関係をまたいだ推論を要する。本論文では、グラフに導かれる agent によってコード局所化に取り組むフレームワーク LocAgent を提案する。LocAgent はコードベースを有向の異種グラフへ解析し、コード構造とその依存関係を捉える軽量な表現を作る。これにより、LLM agent は強力な multi-hop 推論を通じて、関連する実体を効果的に検索し、位置を特定できる。実世界のベンチマークでの実験結果は、本手法がコード局所化の精度を大きく高めることを示している。特に、ファインチューニングした Qwen-2.5-Coder-Instruct-32B モデルを用いた本手法は、SOTA のプロプライエタリモデルに匹敵する結果を、はるかに低い費用、およそ 86% の削減で達成した。ファイルレベルの局所化では最大 92.7% の accuracy に達し、下流の GitHub issue 解決成功率を複数試行、すなわち Pass@10 で 12% 改善した。コードは https://github.com/gersteinlab/LocAgent で公開されている。
論文の面白いところ
この論文の主眼は、LLM に大きなリポジトリをそのまま読ませるのではなく、LLM がたどりやすい形にコードベースを作り替える点にある。コードはディレクトリの木だけでなく、import、関数呼び出し、継承によって結びついている。バグ報告や機能要望は、しばしば表面に出た症状だけを述べ、実際に直すべき関数名を含まない。LocAgent はこのずれを、コードグラフ上の探索問題として扱う。既存の embedding 検索では、クエリとコード片の語彙的・意味的な近さが主な手がかりになるが、修正箇所が呼び出し先や親クラスの側にある場合には届きにくい。Graph-guided agent は、検索で見つけた手がかりから周辺の依存関係をたどれるため、この種の距離を埋めやすい。さらに、論文は LOC-BENCH という新しいベンチマークも作っている。SWE-Bench Lite がバグ修正に偏るのに対し、LOC-BENCH はバグ報告、機能要望、セキュリティ、性能改善を含む。コード修正 agent の実用性を考えるなら、この広がりは重要である。単に「LLM がコードを読む」話ではなく、LLM に読ませる単位と順序を設計する研究として読める。
問題設定
Code localization は、自然言語の issue description を入力として、修正に関係するコード位置を返す課題である。対象位置はファイル、クラス、関数、行などの粒度で評価される。この課題は、プログラム修復の前段にあたる。修正内容を生成する前に、どの場所を見るべきかを誤ると、後段の編集モデルが正しい patch を作ることは難しい。通常の検索課題と違い、クエリはコード中の識別子をそのまま含むとは限らない。たとえば「ユーザープロファイルに XSS 脆弱性がある」という報告は、表示処理ではなく共通の validation utility を直す必要があるかもしれない。リポジトリ全体を LLM の context window に入れる方法も、現実のコードベースでは費用と長さの面で無理がある。Dense retrieval は大規模で変化の速いリポジトリに対して embedding の更新が重くなる。既存の agent 型手法は探索できるが、grep やファイル閲覧に近い道具だけでは、依存関係を数 hop 先まで安定して追うのが難しい。したがって本論文は、軽量に構築でき、LLM agent が扱いやすい構造付き索引を作ることを問題の中心に置いている。
提案手法
LocAgent は、Python リポジトリを有向の異種グラフとして表す。ノードの型は directory、file、class、function であり、エッジの型は contain、import、invoke、inherit である。ファイルやディレクトリの包含関係だけでなく、関数呼び出しやクラス継承を明示的に持つため、ディレクトリ上は離れているが実行上は近い要素も結びつく。各ノードは entity として扱われ、完全修飾名、entity 名、BM25 による ID 検索、コード内容に対する BM25 検索の階層的な sparse index が作られる。実行時の agent には SearchEntity、TraverseGraph、RetrieveEntity の三つの道具が与えられる。SearchEntity はキーワードから関連 entity を探し、件数や長さに応じて fold、preview、full code の形式で返す。TraverseGraph は指定した entity を起点に、方向、hop 数、ノード型、エッジ型を制限しながら幅優先探索を行う。RetrieveEntity は指定 entity の完全なコード、ファイルパス、行番号などを返す。Agent は issue からキーワードを抽出し、それをコード entity に結びつけ、障害の表れから原因へ向かう論理の流れを作り、最後に疑わしい entity を順位付きで返す。さらに、複数回の推論で同じ場所が上位に出るほど信頼度を高くする consistency-based な集約も用いている。プロプライエタリ API への依存を減らすため、Claude-3.5 とファインチューニング済み Qwen2.5-32B の成功 trajectory を用い、Qwen2.5-Coder の 32B と 7B を LoRA による supervised fine-tuning で学習している。
結果
SWE-Bench Lite では、LocAgent はファイル、モジュール、関数の各粒度で強い結果を示した。Claude-3.5 を使う設定では、ファイルレベル Acc@5 が 94.16%、関数レベル Acc@10 が 77.37% であった。ファインチューニングした Qwen2.5-32B でも、ファイルレベル Acc@5 が 92.70%、関数レベル Acc@10 が 77.01% に達し、Claude-3.5 に近い性能を示した。費用面では、Qwen2.5-32B の平均コストは 1 例あたり 0.09 ドルで、LocAgent with Claude-3.5 の 0.66 ドルより低い。論文の abstract では、およそ 86% の費用削減と記されている。Ablation では SearchEntity を外すとファイルレベル Acc@5 が 88.32% から 68.98% に落ち、索引検索が中核であることが分かる。TraverseGraph を外した場合や contain 関係だけにした場合も、モジュール・関数レベルで性能が下がる。これは、呼び出し、import、継承を含むグラフ構造が局所化に効いていることを示す。LOC-BENCH では、LocAgent with Claude-3.5 がファイルレベル Acc@10 で 86.07%、関数レベル Acc@15 で 60.71% を得た。Qwen2.5-7B のファインチューニング版も、より高価な Claude-3.5 を使う SWE-agent と近い水準を保った。ただし、バグ報告以外の機能要望、セキュリティ、性能改善では性能低下が見られ、訓練データの偏りが残ることも示されている。下流の GitHub issue 解決では、Agentless の Pass@10 が 33.58% であるのに対し、LocAgent の局所化を使うと Qwen2.5-32B で 36.13%、Claude-3.5 で 37.59% まで上がった。
具体例
Django の issue として、「三つ以上の media object を merge すると不要な MediaOrderConflictWarning が出る」という報告が来たとする。入力はこの自然言語の説明と、Django のリポジトリ全体である。LocAgent はまず “media”, “merge”, “MediaOrderConflictWarning” のような語を取り出し、SearchEntity で関係しそうなファイル、クラス、関数を探す。検索で django/forms/widgets.py:Media.merge が見つかれば、TraverseGraph によってその関数がどのクラスに含まれ、どこから呼ばれ、どの関連処理とつながっているかをたどる。必要に応じて RetrieveEntity が関数本体を取り出し、agent は warning がどの条件で発生するかを読む。期待される出力は、修正候補として django/forms/widgets.py、Media クラス、Media.merge 関数などを順位付きで示すことである。間違えやすい点は、issue の語に “warning” が含まれるため、単純な検索では警告クラスの定義やログ出力だけを拾ってしまうことである。実際に直すべき箇所は、警告を出す処理そのものではなく、media object の順序を統合するロジックにあるかもしれない。LocAgent のグラフ探索は、語の一致から一歩進んで、呼び出し関係や包含関係を通じてそのロジックへ近づくことを狙っている。