Transform-Augmented GRPO Improves Pass@k

TL;DR

This paper tackles the pattern-matching shortcomings of LLMs in reasoning tasks by addressing two failure modes of GRPO: diversity collapse and gradient diminishing. It introduces Transform-Augmented GRPO (TA-GRPO), which creates semantically equivalent transforms of each question and pools advantages across the whole group to provide non-zero gradients even for easy or hard problems, while promoting multiple solution strategies. The authors provide theoretical justification via non-vacuous initialization, gradient analysis, and generalization bounds, and demonstrate consistent Pass@k improvements on math benchmarks (up to +9.84 points at Pass@32) and out-of-distribution science reasoning (up to +5.05). Empirical results, including ablations, show that pooling advantages—not mere data augmentation—drive the gains, and that TA-GRPO increases rollout diversity, explaining its effectiveness at larger k. Overall, TA-GRPO offers a practical, theoretically grounded path to more robust reasoning in LLMs with broader applicability beyond mathematics.

Abstract

Large language models trained via next-token prediction are fundamentally pattern-matchers: sensitive to superficial phrasing variations even when the underlying problem is identical. Group Relative Policy Optimization (GRPO) was designed to improve reasoning, but in fact it worsens this situation through two failure modes: diversity collapse, where training amplifies a single solution strategy while ignoring alternatives of gradient signal, and gradient diminishing, where a large portion of questions yield zero gradients because all rollouts receive identical rewards. We propose TA-GRPO (Transform-Augmented GRPO), which generates semantically equivalent transformed variants of each question (via paraphrasing, variable renaming, and format changes) and computes advantages by pooling rewards across the entire group. This pooled computation ensures mixed rewards even when the original question is too easy or too hard, while training on diverse phrasings promotes multiple solution strategies. We provide theoretical justification showing that TA-GRPO reduces zero-gradient probability and improves generalization via reduced train-test distribution shift. Experiments on mathematical reasoning benchmarks show consistent Pass@k improvements, with gains up to 9.84 points on competition math (AMC12, AIME24) and 5.05 points on out-of-distribution scientific reasoning (GPQA-Diamond).

Figures (4)

• Figure 1: Percentage of zero-gradient questions throughout training (Qwen3-1.7B). Questions that are "too easy" (all rollouts correct) or "too hard" (all rollouts incorrect) yield zero gradients and contribute nothing to learning: lower is better. TA-GRPO consistently reduces zero-gradient questions compared to GRPO, maintaining a 12--16 percentage point advantage throughout training (e.g., 75% vs 86% at convergence).
• Figure 2: Pass@k curves across all benchmarks for Qwen3-1.7B (top) and Qwen3-4B (bottom). TA-GRPO (red) consistently outperforms GRPO (blue), with the gap widening as $k$ increases.
• Figure 3: Distribution of rollout diversity (average pairwise cosine distance in embedding space) throughout training for Qwen3-1.7B. TA-GRPO (red) maintains higher diversity than GRPO (blue) at all checkpoints, explaining the Pass@k improvements.
• Figure 4: Ablation study on Pass@32 (Qwen3-1.7B). "TA-GRPO w/o pooling" trains on transforms but computes advantages per-variant rather than pooled. Data augmentation alone does not guarantee improvement; pooled advantages are essential.

Theorems & Definitions (15)

• Definition 3.1: Pass Rate and Pass@k
• Proposition 5.4: Non-Vacuous Initialization
• Theorem 5.5: Reduced Zero-Gradient Questions
• proof
• Remark 5.6: Empirical Validation
• Proposition 5.7: Train-Test Gap via Pinsker's Inequality
• Corollary 5.8: Non-Vacuous Test Bound
• Lemma 5.9: Bernoulli Variance of Pooled Reward
• Corollary 5.10: Augmentation Improves Pooled Success
• Remark 5.11: Why Pooled Optimization