Test-Time Scaling with Reflective Generative Model

著者 Zixiao Wang、Yuxin Wang、Xiaorui Wang、Mengting Xing、Jie Gao、Jianjun Xu、Guangcan Liu、Chenhui Jin、Zhuo Wang、Shengzhuo Zhang、Hongtao Xie
所属 MetaStone-AI、USTC (University of Science and Technology of China)
投稿日 2025年7月2日
カテゴリ cs.LG

Test-Time Scaling with Reflective Generative Model

基本情報

  • arXiv ID: 2507.01951v1 (https://arxiv.org/abs/2507.01951)
  • 著者: Zixiao Wang、Yuxin Wang、Xiaorui Wang、Mengting Xing、Jie Gao、Jianjun Xu、Guangcan Liu、Chenhui Jin、Zhuo Wang、Shengzhuo Zhang、Hongtao Xie
  • 所属: MetaStone-AI、USTC (University of Science and Technology of China)
  • 投稿日: 2025年7月2日
  • カテゴリ: cs.LG

簡単に説明すると

この論文は、OpenAI o3と同等の性能を持つ言語モデル「MetaStone-S1」を提案しています。
このモデルの特徴は、通常は別々に学習される「推論するモデル」と「推論の正しさを評価するモデル」を1つに統合したことです。

従来の方法では、言語モデルが複雑な問題を解く際に、まず複数の解答を生成し、その後別の大規模な評価モデルで最良の答えを選ぶ必要がありました。
しかしMetaStone-S1では、自己教師あり学習により1つのモデルで両方の機能を実現しています。
これにより、パラメータ数を99%以上削減しながら、32Bという比較的小さなモデルサイズでOpenAI o3-miniと同等の性能を達成しています。

1. 研究概要

1.1 背景と動機

大規模言語モデル(LLMs)の分野は過去2年間で急速に発展し、GPT-4、Gemini、LLaMA、Qwen、DeepSeek R1などの高度なモデルが登場しています。
特に最近の分析では、OpenAI o3モデルがTest-Time Scaling(TTS)技術を活用して高度な推論能力を実現していることが示されています。

TTSアプローチは内部TTSと外部TTSの2種類に分類されます。
内部TTSは長いChain-of-Thought(CoT)を使用して詳細な思考プロセスを生成します。
外部TTSは、Process Reward Model(PRM)を使用して複数の候補から最良の解答を選択します。

既存の外部TTS手法の主な課題は以下の通りです。

  • 別途大規模なPRM(7B-72Bパラメータ)が必要で、訓練と推論のコストが高い
  • プロセスレベルのアノテーションが必要で、データ収集コストが高い
  • ポリシーモデルとPRMが別々に訓練されるため、分布のミスマッチが発生する

1.2 主要な貢献

本研究では、以下の3つの主要な貢献をしています。

  • Reflective Generative Formという新しいフレームワークを提案する。
    これはポリシーモデルとPRMを統一的なインターフェースで実現する
  • Self-supervised Process Reward Model(SPRM)を導入する。
    これは最終的な答えの正誤のみを使用してプロセスレベルの評価を学習する
  • MetaStone-S1モデルを開発する。
    制御可能な思考長を持つ3つの推論モード(low、medium、high)を提供する

2. 提案手法

2.1 手法の概要

Self-supervised Process Reward Model(SPRM)は、ポリシーモデル自身が解答の正しさを判断できるように訓練されます。
SPRMはポリシーモデルとバックボーンネットワークを共有し、次トークン予測とプロセススコアリングのための2つのタスク固有ヘッドを使用します。

この統一的な設計により、推論と評価を単一のモデルで実行でき、別途PRMを必要としません。
また、既存の推論モデルから簡単にコールドスタートできる利点があります。

2.2 技術的詳細

SPRMの核心的な技術要素は以下の通りです。

ステップ分割: 思考プロセスを'.\n\n'トークンに基づいて複数のステップに動的に分割します。
特別なトークンを追加する必要がなく、既存のトークナイザーをそのまま使用できます。

プロセススコア予測: 軽量な二値分類器(SPRMヘッド)を使用して各ステップのスコアを予測します。
SPRMヘッドは2つの線形層とドロップアウト層のみで構成され、非常に軽量です。
最終スコアは全ステップのスコアの幾何平均として計算されます。

自己教師あり最適化: Self-supervised Process Reward Loss(SPR Loss)を提案します。
これは最終答えの正誤ラベルのみを使用してプロセスレベルの識別能力を学習します。
ノイズの多い監督信号を緩和するため、動的な重み付けメカニズムを導入しています。

2.3 新規性

既存手法との主な違いは以下の点にあります。

従来のPRMは別個の大規模モデルとして訓練され、プロセスレベルのアノテーションが必要でした。
SPRMは、ポリシーモデルとパラメータを共有し、最終答えのラベルのみで学習可能です。
これにより、追加パラメータを99%以上削減し、オンポリシー最適化を実現しています。

また、"aha moment"と呼ばれる興味深い現象を発見しました。
訓練の特定の段階で、モデルが正しい推論と誤った推論を突然区別し始める瞬間があります。

3. 実験結果

3.1 実験設定

実験は3つのモデルサイズで実施されました:MetaStone-S1-1.5B、7B、32Bです。
これらはDeepSeek-R1-DistillとQWQ-32Bから初期化され、継続的な強化学習で訓練されました。

評価は以下のベンチマークで行われました。

  • AIME2024/2025:高度な数学推論能力を評価
  • LiveCodeBench:コーディング能力を評価
  • C-EVAL:中国語の知識と推論能力を評価

推論時には、k=2、8、32の3つの推論努力レベルを設定しました。

3.2 主要な結果

MetaStone-S1は全てのベンチマークで優れた性能を示しました。

小規模モデル(1.5B/7B)では以下の結果が得られました。

  • MetaStone-S1-1.5B-highはAIME24で57.9%を達成し、7B/8Bモデルを上回る
  • MetaStone-S1-7B-highはAIME24で70.2%を達成

大規模モデル(32B)では以下の結果が得られました。

  • MetaStone-S1-32B-highはAIME24で85.2%を達成し、OpenAI o3-miniの79.6%を上回る
  • AIME25では73.6%を達成し、o3-miniの74.8%に匹敵

SPRMは約26Mの追加パラメータのみで、72BのPRMを上回る性能を実現しました。

3.3 既存手法との比較

SPRMと他のPRM手法を比較した結果、以下が明らかになりました。

パラメータ効率に関しては、SPRMは72B PRMと比較して、追加パラメータを99%以上削減しています。
それにも関わらず、全てのベンチマークで優れた性能を示しました。

汎化能力に関しては、数学データのみで訓練されたにも関わらず、コーディングや一般的な推論タスクでも性能向上を示しました。
これは、SPRMがドメインに依存しない推論パターンを学習していることを示唆しています。

スケーリング則に関しては、総思考計算量とTTS性能の間に明確なスケーリング則を確立しました。
性能は計算予算の対数に正の相関を示します。

4. 実用性評価

4.1 実装の容易性

MetaStone-S1は既存の推論モデルから簡単にコールドスタート可能です。
SPRMヘッドは非常に軽量(2層のMLP)で、実装が簡単です。
特別なトークンや複雑な前処理は不要で、既存のインフラストラクチャに統合しやすい設計です。

オープンソースとして公開されており、GitHubから利用可能です。

4.2 計算効率

SPRMの追加パラメータは約5M-26Mと非常に少なく、メモリ効率が高いです。
推論時には、ポリシーモデルとSPRMが特徴抽出を共有するため、計算効率が向上します。

従来の別個のPRM(72B)と比較して、推論時の計算コストを約90%削減できます。
Best-of-N戦略での候補評価が高速化され、実用的なTTSが可能になります。

4.3 応用可能性

MetaStone-S1は数学、コーディング、一般的な推論タスクで優れた性能を示しています。
3つの推論モード(low、medium、high)により、精度と計算コストのトレードオフを調整可能です。

将来的には、Monte Carlo Tree Search(MCTS)などの探索ベースのTTS手法との統合も期待されます。
実験では、MCTSとの組み合わせでさらなる性能向上が確認されています。

5. まとめと所感

5.1 論文の意義

この研究は、大規模言語モデルの推論能力向上において重要な進歩を示しています。
特に、推論と評価を統一的に扱うReflective Generative Formは、TTSにおける計算効率を最大90%向上させる新しいパラダイムを提示しています。

99%以上のパラメータ削減を達成しながら、最新のベンチマークで高い性能を実現した点は特筆すべきです。
また、"aha moment"の発見は、モデルが推論能力を獲得するメカニズムの理解に貢献する可能性があります。

5.2 今後の展望

今後の発展として以下が期待されます。

ステップレベルの探索ベースTTS手法との統合により、さらなる性能向上が期待されます。
リアルタイム推論強化への応用も検討されており、動的な推論能力の向上が可能になるでしょう。

現在の制限として、自己教師あり学習を使用しているため、初期段階では正確なプロセス評価が困難です。
しかし、"aha moment"後は急速に改善することが示されています。
この現象の理解と制御が今後の重要な研究課題となるでしょう。