Transform-Augmented GRPO Improves Pass@k

📄 arXiv: 2601.22478v1 📥 PDF

作者: Khiem Le, Youssef Mroueh, Phuc Nguyen, Chi-Heng Lin, Shangqian Gao, Ting Hua, Nitesh V. Chawla

分类: cs.LG

发布日期: 2026-01-30

💡 一句话要点

提出TA-GRPO，通过转换增强提升GRPO在数学推理中的Pass@k指标

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 数学推理 语言模型 强化学习 策略优化 数据增强 泛化能力 GRPO Pass@k

📋 核心要点

  1. 现有基于GRPO的语言模型在数学推理中存在多样性崩溃和梯度消失问题，限制了模型的泛化能力。
  2. TA-GRPO通过生成语义等价的问题变体，并汇集奖励来计算优势，从而缓解梯度消失并促进多样性。
  3. 实验表明，TA-GRPO在数学推理基准测试中显著提高了Pass@k指标，尤其是在竞赛数学和分布外科学推理任务上。

📝 摘要（中文）

大型语言模型通过预测下一个token进行训练，本质上是模式匹配器，对表面的措辞变化非常敏感，即使底层问题相同。Group Relative Policy Optimization (GRPO) 旨在提高推理能力，但实际上通过两种失效模式加剧了这种情况：多样性崩溃，即训练放大了单一解决方案策略，而忽略了梯度信号的其他选择；梯度消失，即大部分问题产生零梯度，因为所有rollout都获得相同的奖励。我们提出了TA-GRPO（Transform-Augmented GRPO），它生成每个问题的语义等价转换变体（通过释义、变量重命名和格式更改），并通过汇集整个组的奖励来计算优势。这种汇集计算确保了混合奖励，即使原始问题太容易或太难，而对不同措辞的训练促进了多种解决方案策略。我们提供了理论证明，表明TA-GRPO降低了零梯度概率，并通过减少训练-测试分布偏移来提高泛化能力。在数学推理基准上的实验表明，Pass@k指标得到了持续的改进，在竞赛数学（AMC12、AIME24）上获得了高达9.84个点的增益，在分布外科学推理（GPQA-Diamond）上获得了5.05个点的增益。

🔬 方法详解

问题定义：论文旨在解决大型语言模型在数学推理中，由于对表面措辞的敏感性以及GRPO训练过程中的多样性崩溃和梯度消失问题，导致的泛化能力不足的问题。现有方法容易陷入单一解题策略，且对简单或困难的问题无法有效学习。

核心思路：核心思路是通过转换增强来提升GRPO的性能。具体来说，对每个问题生成语义等价的变体，并在这些变体上进行训练。通过这种方式，模型可以接触到更多样化的表达方式，从而提高其泛化能力。同时，通过汇集所有变体的奖励来计算优势，可以避免梯度消失的问题。

技术框架：TA-GRPO的整体框架是在GRPO的基础上，增加了一个转换增强模块。该模块负责生成原始问题的语义等价变体，包括释义、变量重命名和格式更改等。然后，模型在原始问题及其变体上进行训练，并使用汇集的奖励来计算优势。训练完成后，模型可以用于解决新的数学推理问题。

关键创新：最重要的技术创新点在于转换增强策略和汇集奖励的优势计算方法。转换增强策略使得模型能够接触到更多样化的表达方式，从而提高其泛化能力。汇集奖励的优势计算方法可以避免梯度消失的问题，从而提高模型的学习效率。与现有方法相比，TA-GRPO能够更好地利用数据，并学习到更鲁棒的解题策略。

关键设计：在转换增强模块中，使用了多种转换方法，包括释义、变量重命名和格式更改等。这些转换方法旨在生成语义等价但表达方式不同的问题变体。在汇集奖励的优势计算方法中，使用了平均奖励或最大奖励等策略。这些策略旨在确保即使原始问题太容易或太难，模型也能获得有意义的梯度信号。具体的参数设置和损失函数等细节在论文中进行了详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TA-GRPO在数学推理基准测试中显著提高了Pass@k指标。在竞赛数学（AMC12、AIME24）上，TA-GRPO获得了高达9.84个点的增益。在分布外科学推理（GPQA-Diamond）上，TA-GRPO获得了5.05个点的增益。这些结果表明，TA-GRPO能够有效地提高语言模型在数学推理中的性能。

🎯 应用场景

TA-GRPO具有广泛的应用前景，可以应用于各种需要数学推理能力的场景，例如自动解题、科学研究、金融分析等。该方法可以提高语言模型在这些领域的性能，并促进相关领域的发展。此外，TA-GRPO的转换增强策略也可以应用于其他任务，例如自然语言理解和机器翻译等。

📄 摘要（原文）

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).