通过混合提示模板提升数学推理模型的鲁棒性与多样性。
PrAg-PO: Prompt Augmented Policy Optimization for Robust and Diverse Mathematical Reasoning
- 训练时动态混合多种提示模板，避免固定模板过拟合。
- 在8.5K道MATH Level 3-5题目上，准确率显著高于GRPO和DAPO。
- 适合追求高鲁棒性和多样化推理输出的研究者与开发者。
基于强化学习的算法如组相对策略优化（GRPO）在提升大语言模型的数学推理能力方面展现出巨大潜力。然而，现有方法大多使用单一固定推理提示模板，易导致提示特异性过拟合和训练不稳定。本文提出提示增强策略优化（PrAg-PO），在训练中混合不同提示模板，并结合模板专属格式奖励。该方法促使模型在多样指令与输出格式下生成推理过程，显著提升采样多样性与模型鲁棒性。在DeepSeek-R1-Distill-Qwen-1.5B、Qwen2.5-Math-1.5B和Qwen3-1.7B三个模型上的实验表明，仅用8.5K道固定MATH Level 3-5题目集，PrAg-PO始终优于强基线，性能媲美近期先进方法。代码与模型检查点已公开于https://github.com/wenquanlu/PrAg-PO。
原文摘要 · Abstract (English)
Reinforcement learning algorithms such as group-relative policy optimization (GRPO) have shown strong potential for improving the mathematical reasoning capabilities of large language models. While a growing body of work seeks to improve training entropy, rollout diversity, and exploration, most existing methods still train models with a single fixed reasoning prompt or template, which can encourage prompt-specific overfitting and unstable training dynamics. In this work, we introduce Prompt Augmented Policy Optimization (PrAg-PO), a simple policy optimization method that mixes prompt templates with template-specific format rewards during training. By encouraging models to generate reasoning traces under diverse instructions and output formats, PrAg-PO increases rollout diversity and improves robustness. Compared with GRPO and DAPO, PrAg-PO achieves significantly higher reasoning accuracy while mitigating premature training collapse. Empirically, experiments on DeepSeek-R1-Distill-Qwen-1.5B, Qwen2.5-Math-1.5B, and Qwen3-1.7B show that PrAg-PO consistently outperforms strong baselines and achieves competitive performance against recent methods on mathematics benchmarks, using only a fixed MATH Level 3-5 training set of 8.5K problems. The code and model checkpoints are available at https://github.com/wenquanlu/PrAg-PO.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。