[論文レビュー] Transform-Augmented GRPO Improves Pass@k

Khiem Le, Youssef Mroueh+5 more|arXiv (Cornell University)|Jan 30, 2026

Topic Modeling被引用数 0

ひとこと要約

TA-GRPOは各質問の意味的に同等な変換 variant で学習し、グループ全体で報酬を集約して勾配の減衰問題を抑制し、特に高い k における数学・科学的推論ベンチマークで Pass@k を改善する。

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

研究の動機と目的

• LLM の推論タスクにおけるパターンマッチングの限界を動機づける。
• GRPO における勾配減衰と多様性崩壊の問題に対処する。
• transformed-question variants に across をプールするための TA-GRPO を導入する。
• 訓練データとテストデータのギャップ低減および非ゼロ勾配の理論的根拠を提供する。
• 数学・科学的推論ベンチマークで Pass@k の向上を経験的に検証する。

提案手法

• 質問ごとに意味情報を保つ 変換を N 個導入（パラフレーズ、変数名変更、形式変更）。
• 各質問をその変換とともにグループ化し、グループ全体で集約された利得を計算する。
• 集約したホワイトニングを使用：A = (R - mu_group) / (sigma_group + epsilon) をすべての (N+1) バリアントとロールアウトで適用する。
• ベルヌーイ分散と Pinsker-KL を用いてゼロ勾配確率の低下と一般化境界を理論的に正当化し、集約目标を導出する。
• AMC12, AIME24, AIME25, OlympiadBench, Minerva, GPQA-Diamond にわたって Qwen-3.1B および Qwen-3.4B モデルで経験的評価を行う。
• 集約が不可欠であることを示す；データ拡張だけでは集約利得なしでは GRPO を超えられない Ablation を示す。

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 question

実験結果

リサーチクエスチョン

• RQ1 TA-GRPO は標準 GRPO と比較して Pass@k、特に大きい k において改善をもたらすか？
• RQ2 TA-GRPO は分布外の推論タスクへより良く一般化するか？
• RQ3 変換拡張は勾配減衰を減らし、多様な解法戦略を促進するか？
• RQ4 集約利得目的は理論的に正当化され、実践的にも有益か？

主な発見

モデル	AMC12	AIME24	AIME25	OlympiadBench	Minerva	GPQA-Diamond
Qwen3-1.7B Base	65.06	30.00	30.00	60.09	48.53	57.58
Qwen3-1.7B GRPO	69.88	41.31	30.00	66.62	50.37	68.69
Qwen3-1.7B TA-GRPO	79.72	50.00	33.33	68.84	52.94	73.74
+9.84	+8.69	+3.33	+2.23	+2.57	+5.05
Qwen3-4B Base	73.49	43.33	33.33	65.88	59.19	78.79
Qwen3-4B GRPO	84.34	60.00	46.67	70.33	59.56	78.79
Qwen3-4B TA-GRPO	87.95	66.67	53.33	75.07	61.03	82.32
+3.62	+6.67	+6.67	+4.75	+1.47	+3.54

• TA-GRPO は一貫した Pass@k の改善をもたらし、1.7B モデルで AMC12 に対して最大 9.84 ポイント、GPQA-Diamond に対して最大 5.05 ポイントの向上を達成。
• 4B モデルでは、ベンチマーク全体で Pass@32 において GRPO より最大で 3.54 ポイントの利得を達成。
• TA-GRPO は訓練中のゼロ勾配質問を 12–16 ポイント減少させる。
• アブレーションは利得の集約が必須であることを示す；集約なしのデータ拡張のみは一部ベンチマークで GRPO を下回ることが多い。
• TA-GRPO は分布外タスクへの一般化が向上し、GPQA-Diamond を 1.7B で 5.05、4B で 3.53 改善。
• 変換変種の多様性は複数の解法戦略を維持し、より高い Pass@k でより大きな利得を説明する。

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.