Saved in:
Bibliographic Details
Main Authors: Le, Khiem, Nguyen, Phuc, Mroueh, Youssef, Lin, Chi-Heng, Gao, Shangqian, Hua, Ting, Chawla, Nitesh V.
Format: Preprint
Published: 2026
Subjects:
Online Access:https://arxiv.org/abs/2601.22478
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866909054677286912
author Le, Khiem
Nguyen, Phuc
Mroueh, Youssef
Lin, Chi-Heng
Gao, Shangqian
Hua, Ting
Chawla, Nitesh V.
author_facet Le, Khiem
Nguyen, Phuc
Mroueh, Youssef
Lin, Chi-Heng
Gao, Shangqian
Hua, Ting
Chawla, Nitesh V.
contents Group Relative Policy Optimization (GRPO) has become the dominant method for reinforcement learning with verifiable rewards in large language models, but it suffers from two critical limitations: gradient vanishing and diversity collapse. When training questions are too easy or too hard, all sampled responses receive identical rewards, yielding zero gradients. 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 but effective method that addresses both issues via question rephrasing. For each training question, we automatically generate multiple problem-equivalent rephrasings that alter wording, format, and information order while preserving the underlying meaning. Because these rephrasings shift the model's perceived difficulty, pooling responses across the original and its rephrasings yields mixed rewards and more diverse reasoning paths. TA-GRPO jointly computes advantages over this expanded response set and aligns all importance ratios to the original question, enabling the model to learn from a richer set of solution attempts. Experiments on four LLMs (Qwen3-1.7B, Qwen3-4B, Llama-3.2-1B, Llama-3.2-3B) show that TA-GRPO consistently improves pass@$k$ on competition-level benchmarks (AMC, OlympiadBench, AIME24, AIME25) and out-of-distribution benchmarks (Minerva, GPQA-Diamond). Notably, it improves the average pass@32 of Qwen3-1.7B and Qwen3-4B by \textbf{4.97} and \textbf{4.34} points, respectively, and matches the exploration quality of baselines trained on up to 2.5$\times$ more data.
format Preprint
id arxiv_https___arxiv_org_abs_2601_22478
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Transformation-Augmented GRPO for Enhancing Exploration in Reasoning of Large Language Models
Le, Khiem
Nguyen, Phuc
Mroueh, Youssef
Lin, Chi-Heng
Gao, Shangqian
Hua, Ting
Chawla, Nitesh V.
Machine Learning
Group Relative Policy Optimization (GRPO) has become the dominant method for reinforcement learning with verifiable rewards in large language models, but it suffers from two critical limitations: gradient vanishing and diversity collapse. When training questions are too easy or too hard, all sampled responses receive identical rewards, yielding zero gradients. 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 but effective method that addresses both issues via question rephrasing. For each training question, we automatically generate multiple problem-equivalent rephrasings that alter wording, format, and information order while preserving the underlying meaning. Because these rephrasings shift the model's perceived difficulty, pooling responses across the original and its rephrasings yields mixed rewards and more diverse reasoning paths. TA-GRPO jointly computes advantages over this expanded response set and aligns all importance ratios to the original question, enabling the model to learn from a richer set of solution attempts. Experiments on four LLMs (Qwen3-1.7B, Qwen3-4B, Llama-3.2-1B, Llama-3.2-3B) show that TA-GRPO consistently improves pass@$k$ on competition-level benchmarks (AMC, OlympiadBench, AIME24, AIME25) and out-of-distribution benchmarks (Minerva, GPQA-Diamond). Notably, it improves the average pass@32 of Qwen3-1.7B and Qwen3-4B by \textbf{4.97} and \textbf{4.34} points, respectively, and matches the exploration quality of baselines trained on up to 2.5$\times$ more data.
title Transformation-Augmented GRPO for Enhancing Exploration in Reasoning of Large Language Models
topic Machine Learning
url https://arxiv.org/abs/2601.22478