技術要約: STEP-KTODER

問題提起

中間的な推論ステップに対してフィードバックを提供するプロセス・スーパービジョン（過程監督）は、大規模言語モデル（LLM）の数学的推論を大幅に向上させてきましたが、コード生成への適用については未だ十分に探索されていません。数学の領域では、ステップは自然に「思考の連鎖（Chain of Thought）」として定義されます。しかし、コード生成においては、「ステップ」の標準的な概念が存在しません。監督対象となるターゲットは、行、推論トレース、あるいはプログラムの状態など多岐にわたり、何をラベル付けし、何を最適化すべきかという点において曖昧さが生じます。

DPO（Direct Preference Optimization）やKTO（Kahneman–Tversky Optimization）のような現在の選好最適化手法は、通常、アウトカム（結果）レベルで動作し、プログラム全体を単一のフィードバック単位として扱います。この粗いシグナルは、部分的に成功しているプログラム内にある、正解と不正解の構成要素を区別することができません。例えば、大部分は正しいものの、特定の狭いエッジケースで失敗するプログラムは、全体として「望ましくない」とラベル付けされます。これにより、どの関数が成功に寄与し、どの関数が失敗に寄与したのかという詳細が隠されてしまいます。このような粒度の欠如は、局所的に正しいコードを強化しつつ、特定の誤ったコンポーネントを罰するという能力を制限します。

手法: STEP-KTODER

著者らは、分解されたマルチファンクション・プログラムにおけるモジュールレベルの関数を「ステップ」として定義する、コード選好最適化のためのフレームワークであるSTEP-KTODERを提案しています。この手法は、関数レベルのプロセス・スーパービジョンと、実行ベースのラベリング・パイプラインを用いたアウトカムレベルのフィードバックを組み合わせたものです。

1. データ構築パイプライン

本フレームワークは、以下の4つのステージからなる自動データ構築プロセスに依存しています。

• 参照用分解（Reference Decomposition）: 強力な指示チューニング済みコードモデルを用いて、グラウンドトゥルース（正解）のソリューションを、元のインターフェースと挙動を維持したまま、分解されたマルチファンクション・プログラムへと書き換えます。書き換えられたプログラムのうち、元のテストスイートをパスしたもののみが保持されます。
• 関数レベルのユニットテスト生成: 各分解された関数に対して、その関数のインターフェースとエッジケースをカバーするターゲットを絞ったユニットテストをモデルが生成します。これらのテストは、参照実装に対して実行することで検証され、テストを通過し、かつ非自明なアサーションを含むもののみが保持されます。
• オンポリシー候補生成: ターゲットモデルは、分解された参照から派生した関数スケルトン（関数の本体がpassに置き換えられたもの）を埋めることで、候補となる補完コードを生成します。
• ラベリング:
  • アウトカム・ラベル ($o$): フル候補プログラムがデータセット提供のテストスイートを通過するかどうかによって決定されます。
    イルステップ・ラベル ($z$): 生成されたユニットテストに対して各候補関数を実行することによって決定されます。ラベル $1$は局所的な正当性を示し、$0$は失敗を示し、$\emptyset$ は有効なテストが存在しないことを示します。

2. 目的関数

STEP-KTODERは、アウトカムレベルとステップレベルのスーパービジョンを組み合わせた結合目的関数を導入することで、KTOを拡張しています。

• アウトカムレベルのKTO: プログラム全体の望ましさに基づく標準的なKTO損失 ($L_{out}$) を使用します。
• ステップレベルのKTO: 個々の関数のトークンに対してステップごとの損失 ($L_{step}$) を定義します。各関数 $s_i$ について、ローカルな対数比（log-ratio）が計算され、ステップ・ラベル $z_i$ に基づいて価値関数が適用されます。
• 結合目的関数: 最終的な損失は $L_{STEP-KTODER} = L_{out} + \lambda_{step} L_{step}$ となります。

極めて重要な点として、本フレームワークは**コンフリクト保持型ラベル（conflict-preserving labels）**を採用しています。アウトカム・ラベル ($o$) とステップ・ラベル ($z$) が一致しない例（例：グローバルなテストスイートには合格するが、局所的に失敗している関数を含むプログラム）を、学習例として保持します。これらのコンフリクトは、アウトカムのみの手法では捉えられない独自の学習シグナルを提供し、モデルが「全体としては失敗していても、特定のコンポーネントは正しい可能性がある」こと、あるいはその逆を学習することを可能にします。

3. ラベルの信頼性

論文では、実行ベースのラベルが不可欠であることを強調しています。著者らは、実行ベースのステップ・ラベルを「LLM-as-a-judge」（GPT-5.4 miniを使用）によるアノテーションに置き換えてテストを行いました。LLMは全体として中程度の合致度（73.2%）を示しましたが、系統的なバイアスを露呈しました。具体的には、失敗する関数を特定する能力は94.1%と高い一方で、合格している関数を正しく判定できる割合はわずか52.2%であり、失敗を過剰に予測する傾向がありました。このバイアスは、学習に必要なポジティブなステップ・ラベルを損ない、特に困難なベンチマークにおいて性能を低下させる結果となりました。

主な貢献

1. STEP-KTODER フレームワーク: ステップを独立してテスト可能な関数として定義することで、プロセス・スーパービジョンとコード生成の間の溝を埋める、コード特化型のステップワイズKTOの実装。
2. 自動データパイプライン: ソリューションを分解し、関数レベルのユニットテストを生成し、実行ベースのステップ・ラベルを導出するパイプラインにより、手動のアノテーションなしで微細な（ファイングレインな）スーパービジョンを可能にする。
3. 実証的検証: 実行ベースのラベルがステップ・スーパービジョンにおいてLLMによる判断よりも優れていること、および局所的/グローバルなラベルのコンフリクトを保持することが性能維持に不可欠であることを示す実験結果。

実験結果

本手法は、Qwen2.5-Coder および DeepSeek-Coder を用いて、HumanEval(+), MBPP(+), BigCodeBench, LiveCodeBench で評価されました。

• 性能向上: STEP-KTODER は、アウトカムのみの KTO や DPO を一貫して上回り、最も困難なベンチマークで顕著な改善が見られました。
  • Qwen2.5-Coder-1.5B において、BigCodeBench Hard で KTO より +11.5%、LiveCodeBench で +7.7% 向上しました。
  • Qwen2.5-Coder-3B において、BigCodeBench Hard で +11.3%、LiveCodeBench で +9.5% の利得を得ました。
  • DeepSeek-Coder-6.7B において、BigCodeBench Hard で KTO に対して +4.9% の利得を達成しました。
• Target-DPO との比較: Target-DPO も局所的な失敗をターゲットとしていますが、STEP-KTODER は、より少ない訓練サンプル（約11,400対 vs 59,000対）を用いながら、最も困難なベンチマーク（例：3Bモデルにおいて BigCodeBench Hard で Target-DPO に対し +18.5% の改善）において優れた結果を達成しました。
• アブレーション研究:
  • コンフリクト保持: コンフリクト（$o \neq z$ となる例）をマスク（除去）すると、性能が KTO レベルまで低下しました。これは、これらの不一致が不可欠な学習シグナルを提供していることを裏付けています。
  • ラベルソース: 実行ベースのラベルを LLM-as-a-judge によるアノテーションに置き換えると、性能が大幅に低下しました（例：BigCodeBench Hard で -25.8%）。これにより、ステップ・スーパービジョンにおける実行ベースのフィードバックの必要性が検証されました。

意義と主張

本論文は、プログラムの自然なモジュール性を活用することで、STEP-KTODER がコード生成におけるプロセス・スーパービジョンへの実用的な道筋を提供すると主張しています。ステップを関数として定義し、実行を通じて監督することで、本手法は正しいコンポーネントを強化し、特定の欠陥を罰するという、アウトカムのみのスーパービジョンには欠けている局所的なフィードバックを提供します。

著者らは、本手法の有効性が主に以下の2つの要因に依存していると強調しています。

1. 実行ベースの信頼性: 自動化されたユニットテストがステップ・ラベルのための必要なグラウンドトゥルースを提供します。一方、LLMによる判定は系統的なバイアスを導入し、学習を阻害します。
2. コンフリクトの保持: 局所的な正当性とグローバルな成功が乖離する例から学習できる能力が、標準的な選好最適化との主要な差別化要因です。

本研究は、関数レベルの実行フィードバックが、特に粗い粒度のフィードバックでは不十分な複雑なタスクにおいて、コードモデルを整列（アライメント）させるための、実行可能かつ効果的なメカニズムであることを示唆しています。また、解決策がテスト可能な関数に分解可能である必要があることや、新しいモデルに対してオンポリシーの候補を生成するためのエンジニアリングコストといった限界についても述べています。