Transformation-Augmented GRPO for Enhancing Exploration in Reasoning of Large Language Models
Khiem Le, Phuc Nguyen, Youssef Mroueh, Chi-Heng Lin, Shangqian Gao, Ting Hua, Nitesh V. Chawla
原文摘要
Group Relative Policy Optimization (GRPO) has become the dominant reinforcement learning method for eliciting reasoning capabilities in LLMs using verifiable rewards. Still, it suffers from two critical issues that limit models' exploration: \textbf{Gradient vanishing} and \textbf{Diversity collapse}. When a training question is too easy or too hard, all sampled responses receive identical rewards, yielding no gradient signal. Meanwhile, the model tends to collapse its responses toward a single reasoning pattern rather than exploring diverse strategies. We propose Transformation-Augmented GRPO (TA-GRPO), a simple yet effective method that addresses both issues via question transformation. For each training question, we automatically create multiple problem-equivalent rephrased questions that alter wording, format, and information order while preserving the underlying intent. As these rephrasings shift the model's perceived difficulty, responses across the original question and its rephrasings receive mixed rewards and encode more diverse reasoning strategies within it. TA-GRPO jointly computes advantages over this expanded response group and aligns all importance sampling ratios to the original question, enabling the model to learn from a richer set of attempts. Experiments with four LLMs (Qwen3-1.7B/4B-Base and Llama-3.2-1B/3B-Instruct) show that TA-GRPO consistently improves pass@$k$ on competition-level mathematics benchmarks (AMC, OlympiadBench, AIME24, AIME25) and out-of-distribution benchmarks (Minerva, GPQA-Diamond). Notably, it improves the average pass@32 of Qwen3-1.7B-Base and Qwen3-4B-Base by \textbf{4.97}\% and \textbf{4.35}\%, respectively, and matches the pass@32 of baselines trained with up to 2.5x more data.
cs.LG