[論文レビュー] Prompt Augmentation Scales up GRPO Training on Mathematical Reasoning

Wenquan Lu, Hai Huang+1 more|arXiv (Cornell University)|Feb 3, 2026

Topic Modeling被引用数 0

ひとこと要約

要約: 本論文はGRPO型強化学習のプロンプト拡張を導入し、推論フォーマットを多様化し、低エントロピー領域での学習を安定化させ、標準ベンチマークでの数学推論の最先端性能を達成する。

ABSTRACT

Reinforcement learning algorithms such as group-relative policy optimization (GRPO) have demonstrated strong potential for improving the mathematical reasoning capabilities of large language models. However, prior work has consistently observed an entropy collapse phenomenon during reinforcement post-training, characterized by a monotonic decrease in policy entropy that ultimately leads to training instability and collapse. As a result, most existing approaches restrict training to short horizons (typically 5-20 epochs), limiting sustained exploration and hindering further policy improvement. In addition, nearly all prior work relies on a single, fixed reasoning prompt or template during training. In this work, we introduce prompt augmentation, a training strategy that instructs the model to generate reasoning traces under diverse templates and formats, thereby increasing rollout diversity. We show that, without a KL regularization term, prompt augmentation enables stable scaling of training duration under a fixed dataset and allows the model to tolerate low-entropy regimes without premature collapse. Empirically, a Qwen2.5-Math-1.5B model trained with prompt augmentation on the MATH Level 3-5 dataset achieves state-of-the-art performance, reaching 45.2 per-benchmark accuracy and 51.8 per-question accuracy on standard mathematical reasoning benchmarks, including AIME24, AMC, MATH500, Minerva, and OlympiadBench. The code and model checkpoints are available at https://github.com/wenquanlu/prompt-augmentation-GRPO.

研究の動機と目的

• LLMの強化学習による数学的推論の改善を動機付ける。
• 拡張後の学習中のエントロピー崩壊とトレーニング不安定性に対処する。
• 推論テンプレートとフォーマットを多様化するためのプロンプト拡張を導入する。
• 複数の推論プロンプトがより長く効果的なトレーニング・ホライズンを可能にすることを示す。

提案手法

• テンプレート固有の報酬を用いた多様な推論テンプレートをサンプリングすることでプロンプト拡張を定義する。
• デカップリングされたクリッピングとエントロピー考慮を含むトークンレベルGRPO損失を採用する。
• 推論を多様化するためにDeepSeek、free-form、reflection、explicit CoTの4カテゴリで13のテンプレートをキュレーションする。
• 各推論フォーマットの順守を強制するテンプレート固有のフォーマット報酬を使用する。
• 標準ベンチマークで評価する際には固定データセット（MATH Level 3–5）とQwen2.5-Math-1.5Bベースモデルを使用する。

実験結果

リサーチクエスチョン

• RQ1 GRPOトレーニングにおいて、プロンプト拡張は単一テンプレートを超えて推論軌道を多様化できるか？
• RQ2 低エントロピー領域でのRL後の安定化と長期のホライズンを可能にするか？
• RQ3 テンプレート固有の報酬は形式への遵守と全体的な数学的推論精度を改善するか？
• RQ4 標準的な数学ベンチマークにおける最先端性能への影響はどの程度か？

主な発見

Method	AIME24	AMC	MATH500	Minerva	Olympiad	AVG benchmark	AVG question
Qwen2.5-Math-1.5B ∗	16.7	43.4	61.8	15.1	28.4	33.1	37.4
GRPO †	23.3	49.9	75.4	26.3	38.3	42.6	48.4
Dr. GRPO ∗	20.0	53.0	74.2	25.7	37.6	42.1	47.7
SEED GRPO ∗	23.3	50.6	75.4	26.8	41.3	43.5	49.8
GMPO ∗	20.0	53.0	77.6	30.1	38.7	43.9	50.1
DAPO †	23.3	54.9	76.9	26.0	39.4	44.1	49.6
DAPO w/ Prompt Aug † (Step 2720)	23.3	53.5	78.4	31.3	39.4	45.2	50.9
DAPO w/ Prompt Aug † (Step 2480)	20.0	49.4	80.4	29.7	41.2	44.1	51.8

• プロンプト拡張は1つのモデルと1回のトレーニング実行の下で多様な推論軌道を生み出す。
• 低エントロピー領域でのRL後の安定化を促し、崩壊なしで長めのトレーニング（約50エポック）を可能にする。
• MATH Level 3–5でQwen2.5-Math-1.5Bを用いると、特定のステップでベンチマークごとに45.2、質問ごとに51.8といった最先端の結果を達成する。
• DAPO with prompt augmentationは5つのベンチマーク（AIME24、AMC、MATH500、Minerva、OlympiadBench）でベースラインのGRPOおよびDAPOを上回る。
• このアプローチは高性能領域を広げ、定性的な推論軌跡で示される多様な推論フォーマットを生み出す。