Prompt Augmentation Scales up GRPO Training on Mathematical Reasoning

TL;DR

Prompt augmentation diversifies the reasoning templates used during GRPO post-training for mathematical reasoning, mitigating entropy-collapse and enabling longer, more stable optimization horizons. By sampling multiple templates and applying template-specific rewards, it elicits diverse reasoning trajectories and improves learning efficiency without relying on strict KL regularization. On the Qwen2.5-Math-1.5B model trained on MATH Level $3$–$5$, the approach achieves state-of-the-art performance across major benchmarks, including $AIME24$, $AMC$, $MATH500$, $Minerva$, and $OlympiadBench$, with $44.5$ per-benchmark and $51.3$ per-question accuracies at representative checkpoints. The work suggests that simple, inexpensive prompt augmentation can unlock deeper policy improvement and points to inference-time ensembles across prompts as a promising future direction.

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 44.5 per-benchmark accuracy and 51.3 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.

Figures (7)

• Figure 1: Accuracy curves of different training methods based on the performance on five test datasets (i.e., AIME24, AMC, MATH500, Minerva, and OlympiadBench). The accuracies are evaluated every 20 steps. DAPO with prompt augmentation outperforms both DAPO and GRPO baselines in both per-benchmark and per-question average accuracy.
• Figure 2: Examples of prompt templates used in our experiments. The templates can be generally divided into 4 main categories: DeepSeek style, freeform generation, reflection-based and explicit CoT prompting. Complete set of templates can be found in Appendix \ref{['sec:all_prompts']}.
• Figure 3: Distribution of template categories used in our training. Fewer reflection-based templates are included to improve training efficiency, as reflection typically requires longer rollouts.
• Figure 4: Visualization of token-level policy entropy during training of different methods. DAPO with prompt augmentation can sustain stable training at low entropy regimes ($\mathcal{H} < 0.05$) for an extended duration (step 1000 to 3000).
• Figure 5: Reasoning trajectories produced by the model trained with prompt augmentation (checkpoint at step 2820) when answering the same question under different prompt templates. The model follows template-specific instructions to generate diverse reasoning formats.