[論文レビュー] Planning with Large Language Models for Code Generation

Shun Zhang, Zhenfang Chen+4 more|arXiv (Cornell University)|Mar 9, 2023

Software Testing and Debugging Techniques被引用数 13

ひとこと要約

Planning-Guided Transformer Decoding (PG-TD) は、プランナーとテストケース報酬を用いた先読みを行う計画型デコード法で、事前学習済みのコード生成トランスフォーマーと協調して、より高品質なプログラムを生成します。これにより、ファインチューニングなしで制御可能なコード生成を実現し、パス率の向上を達成します。

ABSTRACT

Existing large language model-based code generation pipelines typically use beam search or sampling algorithms during the decoding process. Although the programs they generate achieve high token-matching-based scores, they often fail to compile or generate incorrect outputs. The main reason is that conventional Transformer decoding algorithms may not be the best choice for code generation. In this work, we propose a novel Transformer decoding algorithm, Planning-Guided Transformer Decoding (PG-TD), that uses a planning algorithm to do lookahead search and guide the Transformer to generate better programs. Specifically, instead of simply optimizing the likelihood of the generated sequences, the Transformer makes use of a planner to generate candidate programs and test them on public test cases. The Transformer can therefore make more informed decisions and generate tokens that will eventually lead to higher-quality programs. We also design a mechanism that shares information between the Transformer and the planner to make our algorithm computationally efficient. We empirically evaluate our framework with several large language models as backbones on public coding challenge benchmarks, showing that 1) it can generate programs that consistently achieve higher performance compared with competing baseline methods; 2) it enables controllable code generation, such as concise codes and highly-commented codes by optimizing modified objective.

研究の動機と目的

• 従来の Transformer デコード（ビーム探索/サンプリング）によるコード生成の制限を動機づける。高トークン一致であってもコンパイルされない、あるいは誤ったプログラムとなるケースがある。
• PG-TD を提案する。プランナーを用いてテストケース報酬に導かれた候補プログラムを探索する計画誘導デコードフレームワーク。
• ファインチューニングや追加サンプル収集を必要とせず、標準的なコード生成トランスフォーマーとモデルに依存しない統合を可能にする。
• PG-TD によりパス率が向上し、報酈設計を通じて要約やコメント追加などの制御可能な目的をサポートできることを示す。

提案手法

• コード生成を、状態が問題記述と部分プログラムを連結したMDP、アクションが次のトークンであるという形で定式化する。
• トランスフォーマーが提供する次のトークン確率を指針として、MCTS風のプランナー（P-UCB選択付き）を適用し先読みを行う。
• 上位k個の高確率の次トークンでノードを展開し、候補継続を形成する。
• 部分状態を、トランスフォーマーのビーム探索で完成させ、公開テストケースで評価して報酬を得て木に伝播させる。
• トランスフォーマーとプランナー間で情報を共有するため、木構造と生成シーケンスのキャッシュ機構を導入し、効率を向上させる。
• APPSと CodeContests のデータセットで、GPT-2 および GPT-Neo のバックボーンでの適用によりモデル依存性の高さを示す。
• プランナー報酬を変更することで（例：コードを短く、コメントを増やす）制御可能な目的を探求し、PG-TDで生成されたサンプルを用いたファインチューニングの可能性を示す。

実験結果

リサーチクエスチョン

• RQ1 計画アルゴリズムを事前学習済みのコード生成トランスフォーマーと統合することで、標準的なデコードやサンプリングベースのベースラインよりも良いプログラムを得られるか。
• RQ2 PG-TD は異なるトランスフォーマーのバックボーン（モデル非依存性）で効果的か。
• RQ3 情報共有とキャッシュによりサンプル効率と計算時間が改善されるか。
• RQ4 ファインチューニングせずにPG-TD が制御可能なコード生成の目的をサポートできるか。
• RQ5 PG-TD 生成サンプルを用いてトランスフォーマーをファインチューニングして性能をさらに向上させられるか。

主な発見

• PG-TD は APPS および CodeContests において、パス率の点でビーム探索およびサンプリング＋フィルタリングのベースラインを一貫して上回る。
• PG-TD は GPT-2 および GPT-Neo のバックボーンに適用した場合にもモデル非依存の改善を示す。
• 計画ベースのデコードは有効な先読みを提供し、Sequential Monte-Carlo-Guided TD (SMCG-TD) を上回り、同等の計算予算下でパス率に匹敵するか上回る。
• キャッシュ機構（木構造およびシーケンスのキャッシュ）は計算時間とトランスフォーマーの呼び出しの冗長性を大幅に削減する。
• PG-TD はプランナー報酬を変更することで制御可能なコード生成を可能にし（例：短いコード、より多くのコメント）、生成されたサンプルを用いたファインチューニングで性能をさらに向上させることができる。