[논문 리뷰] Transform-Augmented GRPO Improves Pass@k

Khiem Le, Youssef Mroueh외 5명|arXiv (Cornell University)|2026. 01. 30.

Topic Modeling인용 수 0

한 줄 요약

TA-GRPO는 각 질문의 의미적으로 동등한 변형 버전에 학습하고 그룹 전체에서 보상을 풀링하여 기울기 감소 문제를 줄이고 Pass@k를 개선하며, 특히 더 큰 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).

연구 동기 및 목표

• 질문-당 패턴 매칭이 추론 과제에서 갖는 한계를 제시한다.
• GRPO의 기울기 감소 및 다양성 붕괴 문제를 다룬다.
• 변형된 문제 변형들 간의 이점을 풀링하기 위해 TA-GRPO를 도입한다.
• 학습-테스트 격차 감소 및 비제로 기울기에 대한 이론적 근거를 제시한다.
• 수학 및 과학 추론 벤치마크에서 Pass@k 이득을 실증적으로 검증한다.

제안 방법

• 질문당 N개의 의미 보존 변환을 도입한다(의역, 변수 재명명, 형식 변경).
• 각 질문을 변형들과 함께 그룹화하고 그룹 전체에서 풀링된 이점을 계산한다.
• 풀링된 화이트닝을 사용: A = (R - mu_group) / (sigma_group + epsilon) 모든 (N+1) 변형 및 롤아웃에 걸쳐.
• Bernoulli 분산 및 Pinsker-KL에 의한 0 기울기 확률 감소 및 일반화 경계에 대한 이론적 근거를 제시하고, 풀링된 목적함수를 도출한다.
• Qwen-3.1B 및 Qwen-3.4B 모델에서 AMC12, AIME24, AIME25, OlympiadBench, Minerva, GPQA-Diamond에 걸쳐 실증적으로 평가한다.
• 풀링이 필수적임을 보여주고, 애블레이션에서 데이터 증강만으로는 풀링된 이점 없이는 일반적으로 GRPO보다 못함을 보인다.

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 풀링된 이점 목표가 이론적으로 타당하고 실제로 유익한가?

주요 결과

• TA-GRPO는 일관된 Pass@k 향상을 제공하며, 1.7B 모델의 경우 AMC12에서 9.84포인트까지, GPQA-Diamond에서 5.05포인트까지의 이득을 달성했다.
• 4B 모델의 경우 벤치마크 전반에서 Pass@32에서 GRPO 대비 최대 3.54포인트 이득을 달성한다.
• TA-GRPO는 학습 중 0 기울기 문제를 12–16 퍼센트 포인트 감소시킨다.
• 애블레이션은 이점 풀링이 필수적임을 보여준다; 풀링 없이 데이터 증강만으로는 일부 벤치마크에서 GRPO보다 성능이 낮은 경우가 많다.
• TA-GRPO는 분포 외 과제에 더 잘 일반화하며 GPQA-Diamond를 5.05(1.7B) 및 3.53(4B) 개선한다.
• 변형된 버전들의 다양성은 다수의 해결 전략을 유지하게 하여 더 높은 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.