o1-Coder: an o1 Replication for Coding

3.2節で説明した推論データ合成タスクを完了した後、我々はデータセット内の各完全な推論解を用いて方策モデル$\pi_{\theta}$を初期化する。このステップは、$\pi_{\theta}$がタスク要件をより良く理解し、期待される行動パターンに従うことを助け、その後の反復訓練のための最適な出発点を提供することを目的としている。

質問$Q_{i}$が与えられた場合、ステップ$j$において方策モデル$\pi_{\theta}$によって生成される特定の推論ステップの内容は$\pi_{\theta}(S_{i}^{j}\mid Q_{i},S_{i}^{1:j-1})$と表現できる。ここで、$S_{i}^{j}=(w_{1},w_{2},\dots,w_{k})$である。$S_{i}^{j}$は推論ステップの内容を表し、特定の区切り文字で区切られており、$w$は各デコーディングステップで$\pi_{\theta}$によって生成されるトークンを表す。$S_{i}^{1:j-1}$は前の推論ステップの出力によって形成されるコンテキストを表す。

その後、方策モデル$\pi_{\theta}$は検証済みの正しい推論解のセット$\mathcal{D}^{+}_{\text{process}}$を用いて初期化される。この初期化は、以下の訓練目的を最適化することによって実行される：

ここで、$\circ$は推論ステップ$S_{i}^{1:m}$と最終的なコード$C_{i}^{\prime}$の連結を表す。初期化された方策モデル$\pi_{\theta}^{\text{SFT}}$は、その後の訓練段階の基礎として機能する。

問題$Q_{i}$と現在の状態に対応する解答の接頭辞が与えられた場合、プロセス報酬モデル（PRM）（$Q\times S\rightarrow\mathbb{R}^{+}$と表記）は、現在のステップ$S_{i}^{j}$に報酬値を割り当て、最終的な解答への貢献度を推定する。3.2節のデータ合成中に使用されたツリー探索アプローチに基づき、プロセス報酬モデルの訓練には、ポイントワイズとペアワイズと呼ばれる2つの形式のデータ編成を使用できる。以下に詳細を述べる。

ポイントワイズ この形式では、探索ツリーから収集されたデータは$D=\{(Q_{i},S_{i}^{1:j-1},S_{i}^{j},v_{i}^{j})\mid i=1,2,\ldots,N\}$として編成される。ここで、$N$はサンプル数であり、$v_{i}^{j}$はツリー探索プロセス中にステップ$S_{i}^{j}$に割り当てられた値ラベルを表す。処理方法に応じて、このラベルはハードまたはソフトな推定値を導出するために使用できる。(Wang et al., 2024)のアプローチに従い、PRMは以下の目的関数を用いて訓練される：

ここで、$r(Q_{i},S_{i}^{1:j})$はPRMによって割り当てられた正規化された予測スコアである。

ペアワイズ ペアワイズ形式では、探索ツリーの深さ$d$にあるノード$n^{d}$について、その子ノードを$\sum_{i}n_{i}^{d+1}$として表現し、選好ペアデータを$D_{\text{pair}}=\{(Q_{i},S_{i}^{1:j-1},S_{i}^{j_{\text{win}}},S_{i}^{j_{\text{lose}}})\mid i=1,2,\ldots,N\}$として編成する。ここで、$S_{i}^{j_{\text{win}}}$はツリー探索中に$S_{i}^{j_{\text{lose}}}$と比較してより高い値の推定を達成した推論ステップを表す。

Bradley-Terryモデル(Bradley & Terry, 1952)に従い、PRMは以下の目的関数を用いて訓練される：

ここで、$\sigma(x)$はシグモイド関数を表す。ポイントワイズの設定とは異なり、ここでのスコア$r$は正規化されない。これにより、モデルは絶対的な値の予測ではなく、行動間の相対的な選好の学習に焦点を当てることができる。

我々は、コード生成タスクを言語拡張マルコフ決定過程（MDP）としてモデル化し、形式的に $\mathcal{M}=(\mathcal{V},\mathcal{S},\mathcal{A},\mathcal{T},\mathcal{R},\phi)$ と表現する (Team, 2024; Carta et al., 2023)。 このフレームワークにおいて、$\mathcal{V}$ は語彙を表し、$w\in\mathcal{V}$ はモデルによって生成される個々のトークンを表す。 行動空間 $\mathcal{A}\subseteq\mathcal{V}^{N}$ と状態空間 $\mathcal{S}\subseteq\mathcal{V}^{N}$ はトークン列の集合であり、行動と状態の両方がトークン列であることを意味する。 このフレームワークでは、$s_{0}$ が質問を表し、行動 $a_{i}$ は推論ステップ（アルゴリズム 1 の $S_{i}$ を参照）とみなされ、行動の種類とそれに対応する思考の連鎖の両方で構成される。 状態遷移関数 $\mathcal{T}:\mathcal{S}\times\mathcal{A}\to\mathcal{S}$ は、行動 $a_{t}\in\mathcal{A}$ が取られたときに現在の状態 $s_{t}\in\mathcal{S}$ がどのように変化するかを定義する。具体的には、行動 $a_{t}$ が現在の状態にトークンを追加し、新しい状態 $s_{t+1}=\mathcal{T}(s_{t},a_{t})$ を形成する。このプロセスは、モデルが最終的な解を生成するまで続く。報酬関数 $\mathcal{R}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}^{+}$ は、推論プロセスや生成されたコードフラグメントなどの中間ステップの質を評価する。関数 $\phi$ は、プロセスベースの報酬と結果ベースの報酬を組み合わせて最終的な報酬信号を生成する。

各ステップで、モデルは行動 $a_{t}\in\mathcal{A}$ を選択し、システムを新しい状態 $s_{t+1}=\mathcal{T}(s_{t},a_{t})$ に遷移させる。行動を実行した後、モデルはPRMからプロセス報酬 $r^{t}=\rho_{\text{PRM}}(s_{t-1},a_{t})$ を受け取る。このプロセスは、モデルが最終的なコードを生成するか、事前に定義された最大深度に達するまで繰り返される。

モデルが最終的なコードを生成するか、探索プロセスを完了すると、結果報酬 $R_{i}$ が一連のテストケースに対して生成されたコードをテストすることで評価される。我々は、時間依存の重みと割引因子の両方を組み込んだ報酬集約関数を提案する：

ここで、$\alpha(t)$ は最終報酬 $R_{i}$ と累積中間報酬 $r_{i}^{1:m}$ のバランスを時間とともに調整する時間変動因子である。例えば、$\alpha(t)$ は時間とともに減少し、モデルが解を洗練するにつれて中間報酬の重みを徐々に増やし、モデルが最適な方策に近づくにつれて最終報酬の重みを減らすかもしれない。$r_{i}^{1:m}$ であり、$\alpha(t)$ は通常、線形または対数的減衰などのスケジュールに従う。パラメータ $\gamma\in[0,1]$ は割引因子であり、将来の報酬の重要性を即時の報酬に対して決定する。集約された報酬信号は、通常、PPO (Ziegler et al., 2019) や反復的DPO(Rafailov et al., 2024)などの強化学習アルゴリズムの実装を通じて、モデルの方策を洗練するために使用される。

この設定により、我々はコード生成タスクに特化した強化学習環境を定義する。モデルの行動は、中間的な推論ステップを促すプロセスベースの報酬と、最終的なコードの正確さを反映する結果ベースの報酬の両方によって駆動される。この二重報酬構造は、モデルが時間とともにコード生成能力を向上させるのに役立つ。

ステップ6では、更新されたポリシーモデル$\pi_{\theta}$を使用して、$\mathcal{D}^{\prime}_{\text{process}}$と表される新しい推論データを生成する。このデータは、新しい問題インスタンス$Q_{i}$を通じて推論し、段階的な推論パス$\{S_{i}^{1},S_{i}^{2},\dots,S_{i}^{m}\}$を生成することで作成される。各パスは最終的なコード出力$C_{i}^{\prime}$で culminate する。推論ステップは反復的に生成され、各ステップ$S_{i}^{j}$は前のステップに条件付けられる。

新しい推論データが生成されると、それは既存のデータセット$\mathcal{D}_{\text{process}}$に追加され、更新されたデータセット$\mathcal{D}_{\text{process}}\leftarrow\mathcal{D}_{\text{process}}\cup\mathcal{D}^{\prime}_{\text{process}}$を形成する。この更新により、推論例の多様性と品質が向上し、後続のステップのためのより包括的な訓練材料が提供される。

この新しいデータ生成プロセスにより、反復的な自己対戦訓練ループが可能となる。新しい推論データを追加した後、モデルはさらなる微調整を受け、第4ステップで説明したPRMの更新から始まる。PRMは、次に第5ステップで説明したRLを用いてポリシーモデルを調整する。データ生成、報酬モデルの更新、ポリシーの改善というこの反復的なサイクルにより、システムの推論能力の持続的な向上が確保される。

過去10年間、AI分野は計算能力-知能変換効率の最大化という中心線に沿って発展してきた。これは、絶えず増加する計算能力を効率的により高度な知能レベルに変換することを意味する。この流れに沿って、図6の上部に示されているように、初期の進歩はモデル側の改善を優先した：SVMからDNN、そしてTransformerへと、計算能力を最大限に活用できるスケーラブルなモデルアーキテクチャが設計された。

近年、焦点はデータ側にシフトしている。事前学習における半教師あり学習（SSL）や事後学習における強化学習（RL）などの技術は、データをより効果的に活用することを目的としている。o1モデルはこの流れを継続している。高品質な教師あり学習データを活用するSFTから、理論上無限のデータにアクセスできる環境フィードバックを利用するRLHFへ、そして最終的にo1の革新的なアプローチである、生成された推論プロセス自体から導出された報酬信号を通じて生成プロセスを監督する方法へと移行している。

この進展は、Transformerアーキテクチャが現在、膨大な量のデータを処理し、十分な規模のモデルを訓練できるようになったことを考えると、残された唯一の課題は適切なデータの獲得に収束することを示唆している。一つのアプローチは、システム2能力のための推論データや、身体化された知能のための物理世界の軌跡など、不足しているデータを求めることである。もう一つのアプローチは、人間の世界にはまだ存在しないデータタイプを探索することであり、これにはRLやSelf-Playなどの技術のさらなる探求が必要である。

LLMに対する一般的な批判は、既存の人間が記録したデータに依存していることであり、これが本質的にその潜在能力を制限している。ウィトゲンシュタインが述べたように、「私の言語の限界が私の世界の限界を意味する」。人間の言語記録の有限な範囲と深さが、LLMの認知能力を制約している。しかし、o1の成功は、我々が強化学習を通じてこれらの記録されたデータの背後にある思考プロセスを探求できるようになったことを示している。この進歩は、AI開発における重要な転換点を示しており、単なる人間の言語の模倣から、新しい認知プロセスの自律的な生成へと移行している。

さらに興味深いことに、これらの思考プロセスデータは必ずしも自然言語に限定される必要はない。最近のNature論文で強調されているように(Fedorenko et al., 2024)、「言語は思考の本質というよりも、主にコミュニケーションのための道具として機能する」。我々の観察では、o1によって生成された思考の連鎖の一部に無意味なテキストが含まれており、思考トークンが離散的な自然言語の単語に対応していない可能性を示唆している。モデルが思考のためのより効率的な内部表現形式を自ら開発したとすれば、これは思考プロセスと問題解決メカニズムの効率を大幅に向上させ、人間の言語データによって課された制限を超えるだけでなく、モデルの能力の可能性をさらに解き放つことになるだろう。

自己対戦+強化学習フレームワークは、基礎となるデータを探索するための実行可能な解決策を提供し、これまでシステム1の能力に依存していた多くのタスクにおいてシステム2の解決策を探索する可能性を開いている。タスク実行により思慮深く段階的なプロセスを統合することで、この手法が幅広い領域で好ましい結果をもたらすと我々は考えている(Kant et al., 2024; Ganapini et al., 2021; Valmeekam et al., 2024; Lowe, 2024)。報酬モデリング(Mahan et al., 2024)、機械翻訳(Zhao et al., 2024)、検索拡張生成（RAG）(Li et al., 2024)、マルチモーダルQA(Islam et al., 2024)など、従来システム1の能力を用いて解決されていたタスクは、すでにシステム2の思考によって可能となったより深い推論能力の恩恵を受けている。

o1モデルのシステムカードは、モデルの安全性における顕著な改善を示している。これに触発され、我々は最近システム2アライメントの概念を探求している。これは、入力を徹底的に評価し、潜在的なリスクを考慮し、推論におけるバイアスを修正するようモデルを導くことを含む(Wang & Sang, 2024)。我々はシステム2アライメントを実現するために、プロンプトエンジニアリング、教師あり微調整、プロセス監視付き強化学習という3つの方法を導入した。本稿で提示した自己対戦+強化学習フレームワークをシステム2アライメントに適用し、複雑なシナリオにおいてモデルが慎重に思考する能力をさらに向上させ、脆弱性を減少させることを目指す。

現在リリースされているo1-previewとo1-miniには、OpenAIが完全版に含まれると主張しているマルチモーダル機能や関数呼び出し機能が欠けている。マルチモーダルと関数呼び出しを超えて、o1のような推論モデルの改善に不可欠なもう一つの特徴は、推論時間の最適化である。これには、推論効率の向上—単位時間あたりのより高いパフォーマンスの達成—と、適応的な推論時間調整の実現が含まれる。具体的には、タスクの複雑さに基づいてシステム2の推論プロセスを動的に調整し、より人間に近いシステム1とシステム2の推論モード間のシームレスな切り替え能力を実現することである。

o1のような推論モデルをより広範な実世界のアプリケーションに展開するためには、両方とも強化学習環境に関連する2つの主要な課題に取り組む必要がある。第一の課題は報酬関数の一般化に関するものである。これはすでにコミュニティで議論されている。例えば、推論モデルの高レベルの自然言語指示を理解する能力の向上を活用し、Constitutional AI (Bai et al., 2022)のようなアプローチは、報酬関数を直接自然言語で定義する可能性がある。もう一つの戦略は、コーディング能力を向上させ、他のタスクをコーディング問題に変換して解決することに焦点を当てている。

あまり言及されていないもう一つの課題は、環境状態の更新に関するものである。Q学習のような古典的なモデルフリーの強化学習手法とは異なり、状態遷移が明示的にモデル化されていないo1のようなモデルは、行動シミュレーションと前方探索に依存しており、行動後の更新された状態の知識を必要とする。これにより、パラダイムはモデルベースの強化学習へとシフトする。プログラミング、数学、囲碁などの明確に定義された領域では、環境には多くの場合決定論的なルールがある。例えば、プログラミングはコンパイラで定義された言語仕様を使用し、数学は公理論理に従い、囲碁は固定されたゲームルールの下で動作する。これらの決定論的なフレームワークにより、特定の行動後の状態遷移確率$p(\text{state}_{i+1}\mid\text{state}_{i},\text{action}_{i})$を正確に計算することができる。

しかし、検索拡張生成（RAG）、デバイス使用、実体化エージェントなど、多くの実世界のアプリケーションでは、状態更新の取得には外部環境やシミュレータとの相互作用が必要である。これにより、大きな計算コストと時間コストが発生する。例えば、デバイス使用では、クリック、入力、スクロールなどの行動を、ページのレンダリング、状態更新、時にはネットワークリクエストのような複雑なバックエンド相互作用を含む方法でシミュレートする必要がある。さらに、o1のようなモデルは、推論中にオンラインの行動シミュレーションを実行できないという制限に直面しており、これにより、モデルが以前の状態に戻って行動を検証または修正することができなくなる。これは、バックトラックして決定を改善する能力の欠如につながる。

したがって、主要な方向性の一つは、状態遷移予測のためのワールドモデルを開発することにより、環境を明示的にモデル化することを試みることである。ワールドモデルは、現在および過去の状態と行動を入力として受け取り、次の状態を出力として生成する。これにより、モデルは実際の環境やシミュレータと直接相互作用するのではなく、内部のワールドモデルと相互作用することができる。我々は、このようなワールドモデルを構築する際の強化学習における継続的な課題の一つが、その精度を確保することであると認識している。その結果、ワールドモデルは通常、ダイナミクスが比較的単純で十分に理解されている環境に適用されてきた。良いニュースは、生成ゲーム(Sang, 2024)における最近の急速な進歩が、実世界のアプリケーションにおける推論モデルのためのより正確で実用的な環境モデリングを促進する可能性のある有望な進展を提供していることである。

展望。 o1モデルは明らかにAlphaGoの影響を受けている：AlphaGoは模倣学習を使用してポリシーネットワークを初期化し、強化学習を用いてポリシーを微調整し、価値ネットワークを学習し、MCTSをオンライン探索戦略として使用した。これはLLMの事前学習、事後学習、推論と並行している。AlphaGoZeroは、歴史的データに依存しないより高度なアプローチを取り、これは現在のLLM開発が事後学習段階をますます重視する傾向と正確に一致している。Alphaシリーズの進化に従えば、o1のような推論モデルでも同様の発展が予想される。当初、Alphaシリーズは一般化に向けて発展した：AlphaZeroは囲碁、チェス、将棋に適用され、MuZeroは57のAtariゲームで人間レベルのパフォーマンスを達成した。しかし、一般化以外のもう一つの目標は、これらのモデルをより複雑な実世界のタスクに適用することである。この進展は、AlphaFoldからAlphaCodeとAlphaGeometryへの飛躍、さらにはSIMAの3D仮想エージェントやRT-Xの具現化された知能など、AIの物理的環境への拡張に見られる。