[論文レビュー] rStar-Math: Small LLMs Can Master Math Reasoning with Self-Evolved Deep Thinking

Xinyu Guan, Li Lyna Zhang+6 more|arXiv (Cornell University)|Jan 8, 2025

Machine Learning and Data Classification被引用数 7

ひとこと要約

本論文は、小規模言語モデルが自己進化的な深い思考と Monte Carlo Tree Search、コードを組み込んだ CoT データ合成法、対ペアランキングで訓練されたプロセス好みモデルを用いることで、OpenAI o1-style の数学的推論に到達するか、あるいはそれを上回ることを示している。

ABSTRACT

We present rStar-Math to demonstrate that small language models (SLMs) can rival or even surpass the math reasoning capability of OpenAI o1, without distillation from superior models. rStar-Math achieves this by exercising "deep thinking" through Monte Carlo Tree Search (MCTS), where a math policy SLM performs test-time search guided by an SLM-based process reward model. rStar-Math introduces three innovations to tackle the challenges in training the two SLMs: (1) a novel code-augmented CoT data sythesis method, which performs extensive MCTS rollouts to generate step-by-step verified reasoning trajectories used to train the policy SLM; (2) a novel process reward model training method that avoids naïve step-level score annotation, yielding a more effective process preference model (PPM); (3) a self-evolution recipe in which the policy SLM and PPM are built from scratch and iteratively evolved to improve reasoning capabilities. Through 4 rounds of self-evolution with millions of synthesized solutions for 747k math problems, rStar-Math boosts SLMs' math reasoning to state-of-the-art levels. On the MATH benchmark, it improves Qwen2.5-Math-7B from 58.8% to 90.0% and Phi3-mini-3.8B from 41.4% to 86.4%, surpassing o1-preview by +4.5% and +0.9%. On the USA Math Olympiad (AIME), rStar-Math solves an average of 53.3% (8/15) of problems, ranking among the top 20% the brightest high school math students. Code and data will be available at https://github.com/microsoft/rStar.

研究の動機と目的

• 小型言語モデルが蒸留なしで最先端の数学的推論ベンチマークに匹敵するか、またはそれを上回ることを実証する。
• 数学的推論のためにポリシーモデルと報酬モデルの両方を反復的に改善する自己進化ワークフローを開発する。
• 信頼性の高い逐次ステップの軌跡を生成するためのコードを組み込んだチェーン・オブ・思考（CoT）データ合成法を提案する。
• ペアワイズランキングで訓練されたプロセス好みモデル（PPM）を導入し、密で信頼性のあるステップレベルの報酬を提供する。
• 複数データセットとモデルサイズで実証的な利得を示し、数学ベンチマークで大規模フロンティアモデルに近づく、あるいは超えることを示す。

提案手法

• ポリシーSLMとプロセス報酬モデル（PRM）を用いて、数学推論に深い思考を行うためのモンテカルロ木探索（MCTS）を使用する。
• 各ステップがPythonコードも生成するコード組み込みのCoT生成を導入する。コード実行済みの世代のみを保持して、有効な中間ステップを保証する。
• 大量のMCTSローアウトを介してステップ品質を注釈付けし、Q値を割り当て、終端およびPRM強化戦略でステップのスコアを洗練する。
• 高Qステップと低Qステップを用いたペアワイズ Bradley-Terry ランキングでプロセス好みモデル（PPM）を訓練し、ノイズの多い各ステップのスコアに直接依存することを回避する。
• シードデータセット 747k 問題を用いて、スクラッチからポリシーSLMとPPMを徐々に強化する4ラウンドの自己進化ループを実装する。
• MATH、AIME、AMC、Olympiad Bench などのベンチマークで 1.5B–7B の SLM を評価し、OpenAI o1 や他のベースラインと比較する。

実験結果

リサーチクエスチョン

• RQ1 小規模な LLM は大規模モデルからの蒸留なしでフロンティア級の数学推論を達成できるか？
• RQ2 ポリシーと報酬モデルの自己進化が、多段階の数学問題解決のギャップをどう埋めるか？
• RQ3 コード組み込みCoTデータ合成は推論ステップのトレース品質を向上させるか？
• RQ4 ペアワイズランキングで訓練されたプロセス好みモデルは、数学推論の信頼できるステップレベルの報酬を提供できるか？
• RQ5 MCTS の軌跡を増やすことが、さまざまな数学ベンチマークでの性能にどのような影響を与えるか？

主な発見

• rStar-Math は 7B規模のモデルを最先端レベルへと改善し、難解な数学ベンチマークでのスコアを OpenAI o1 相当以上へと引き上げる。
• MATH では、Qwen2.5-Math-7B が64軌道で 58.8% から 90.0% に改善され、o1-preview を上回り o1-mini に匹敵する。
• AIME 2024 では、rStar-Math が平均53.3%（8/15 問題）を達成し、上位20%の最も優秀な高校生数学者の中に位置する。
• 数百万の合成解を用いた4回の自己進化により、ポリシーモデルとPPMが次第に強化され、747k 問題のカバレッジを 90.25% へ拡大。
• Python 実行と MCTS ベースの Q 値注釈を伴うコード組み込みの CoT は中間ステップの誤りを減少させ、軌跡の品質を向上させる。
• ペアワイズランキングで訓練されたプロセス好みモデル（PPM）は信頼性の高いステップレベルの指針を提供し、アブレーション研究でベースラインの報酬モデルアプローチを上回る。