技術概要：強化学習による LLM コード生成のためのプロンプト最適化

問題提起

大規模言語モデル（LLM）は、自然言語からのコード生成において特に自動化されたソフトウェア工学を進展させてきたが、その出力には機能的および論理的な欠陥が頻繁に含まれている。これらの誤りは実用的な展開を制限し、手動でのデバッグを必要とする。静的なプロンプトエンジニアリング手法（例：Chain-of-Thought）は一般的な問題を緩和できるが、中間フィードバックに基づく反復的な改善が求められる場合、適応性に欠ける。進化アルゴリズム（例：EPiC）や強化学習（RL）フレームワーク（例：Reflexion）といった既存の適応的アプローチは、しばしば語彙的および意味的なプロンプト変更を個別に扱ったり、二値の正誤信号に依存したりしており、部分的な正しさから得られる豊富な情報を活用できず、複数の変換次元を統合的に結合できていない。

手法

著者らは、プロンプトの改善を逐次意思決定問題としてモデル化する強化学習ベースのフレームワークを提案する。手法の中核的構成要素は以下の通りである：

1. フレームワークアーキテクチャ

このシステムは、プロンプトを反復的に改善するために Proximal Policy Optimization（PPO）エージェントを利用する。プロセスは、カスタムの Gymnasium 互換環境内で動作する：

• 状態: 現在のプロンプトは、MiniLM 文書トランスフォーマーを用いて 384 次元の意味的埋め込みとして符号化される。注目すべきは、状態表現がプロンプトのみに限定されており、生成されたコード、ランタイム例外、または実行トレースバックを符号化していない点である。
• 行動空間: エージェントは、3 つの戦略からなるハイブリッド行動空間から選択する：
  1. 直接生成: 現在のプロンプトを維持する（デコードの確率性を活用）。
  2. 語彙的突然変異: テストケースの性能に導かれて、トークンレベルの挿入、削除、置換、および交差を含む、遺伝的アルゴリズムに着想を得たアプローチ。
  3. 意味的書き換え: 自己反省的パラダイムに従い、事前学習済みモデルを用いてプロンプトのより高レベルな再構成と明確化を行う。
• 環境ループ: エージェントはプロンプトを観察し、行動をサンプリングし、変換を適用し、新しいプロンプトを凍結されたコード生成モデルに渡し、出力を単体テストに対して評価する。

2. 報酬構造

主要な革新は、二値の成功/失敗信号のみに依存するのではなく、単体テストのフィードバックから導出された形状化報酬関数の使用である。報酬 $R_{i,t}$ は以下のように定義される：

• 通過率 $\rho_{i,t} = 1.0$（すべてのテストに合格）の場合、$1.0$。
• $0 < \rho_{i,t} < 1.0$（部分的な正しさ）の場合、$\rho_{i,t}$。
• $\rho_{i,t} = 0$（どのテストも不合格）の場合、$-1$。
• 実行失敗（構文エラーまたはランタイム例外）の場合、$-2$。

この設計は、より密な学習信号を提供し、エージェントが漸進的な改善に対してクレジットを受け取り、構文上無効なコードを抑制するために実行失敗をより重く罰則化することを可能にする。

3. 評価指標

このフレームワークは、2 つの指標を用いて評価される：

• Strict Pass@1: 主要指標であり、あるステップでいずれかの解決策がすべてのテストに合格した場合にのみタスクを成功とみなす。
• SoftPass@1: 改善ステップ全体にわたるテスト合格の累積確率を測定する補助指標であり、中間的な進捗を分析するために使用される。

主要な貢献

本論文は、4 つの主要な貢献を概説する：

1. 逐次定式化: コード生成のためのプロンプト最適化を逐次意思決定問題として定式化し、PPO ベースの RL 環境として実装した。
2. ハイブリッド行動空間: 適応的な多段階改善のための、直接生成、遺伝的語彙的突然変異、および意味的書き換えを統合した結合行動空間の定義。
3. 形状化報酬: 二値の成功のみよりも密な学習信号を提供するために、部分的なテストケースの正しさを利用する報酬関数の導入。
4. 包括的評価: CodeT5+、CodeLLaMA、DeepSeek-Coder を凍結されたバックボーンとして使用し、MBPP+、HumanEval+、APPS ベンチマークにおける広範な評価。これらは、直接生成、EPiC、Reflexion、および Random-Hybrid ベースラインと比較された。

実験結果

このフレームワークは、MBPP+ テストセットの 500 タスク、ならびに HumanEval+ および APPS 上でテストされた。

• 性能向上: MBPP+ において、PPO エージェントは Strict Pass@1 スコアで 57.58%（CodeT5+）、64.80%（CodeLLaMA）、85.50%（DeepSeek-Coder）を達成した。これらの結果は、すべてのバックボーンモデルにおいて EPiC、Reflexion、および Random-Hybrid ベースラインを上回った。
• ベースラインとの比較:
  • PPO エージェントは、Random-Hybrid ベースライン（同じ行動空間を使用するが、行動を均一にランダムに選択するもの）を大幅に上回った。例えば、CodeT5+ を用いた MBPP+ において、PPO は Random-Hybrid よりも Strict Pass@1 を 26.46 パーセントポイント改善した。これは、性能向上が単に複数の変換タイプの可用性から生じたのではなく、学習された行動スケジューリングに由来することを示している。
  • PPO は EPiC および Reflexion も凌駕し、学習されたポリシーを介して語彙的および意味的変換を調整することが、個別の進化探索や反省ベースの手法よりも効果的であることを示した。
• 報酬形状化の影響: 疎な二値報酬と形状化報酬を比較する実験により、後者がより速い収束と高い安定性をもたらすことが示された。二値報酬は学習が遅く不安定であり、部分的な正しさの信号が効果的な多段階改善に不可欠であることを確認した。
• 統計的有意性: MBPP+ テストセットにおける McNemar 検定は、Strict Pass@1 において EPiC および Reflexion に対して統計的に有意な改善（$p < 0.001$）を確認し、中程度の効果量を示した。

意義と主張

本論文は、その主要な意義が、形状化されたテスト駆動型報酬を用いた強化学習が、LLM ベースのコード生成の機能的正しさを改善できることを実証した点にあると主張する。

著者らが強調する主な教訓は以下の通りである：

• 適応的調整: エージェントが中間フィードバックに基づいて、語彙的突然変異（局所探索）と意味的書き換え（大域再構成）の間を動的に調整する能力が、成功にとって決定的である。
• 部分的信号の価値: 多段階改善タスクにおけるエージェントの訓練において、二値の成功/失敗よりも部分的な正しさを学習信号として利用する方が優れている。
• 頑健性: このフレームワークは、異なるバックボーンモデルおよび変化するタスクの難易度（MBPP+ からより複雑な APPS まで）において一貫した改善を示しており、アーキテクチャの差異に対する頑健性を示唆している。

著者らは謙虚な立場を維持し、コードやトレースバックを除外したプロンプトのみの状態表現、凍結モデルへの依存、および反復的な生成と実行の計算コストといった限界を認めている。将来の作業としては、トレースバックを考慮した状態表現や、実用的な開発ツールへの統合が提案されている。