[논문 리뷰] Group Distributionally Robust Optimization-Driven Reinforcement Learning for LLM Reasoning

Kishan Panaganti, Zhenwen Liang외 3명|arXiv (Cornell University)|2026. 01. 27.

Artificial Intelligence in Healthcare and Education인용 수 0

한 줄 요약

논문은 온라인 난이도에 따라 프롬프트를 동적으로 분할하고 그룹 간 롤아웃을 할당하여 LLM의 추론을 개선하는 다-적대적 GDRO 프레임워크를 제시하며, GRPO 대비 주목할 만한 이득을 달성한다.

ABSTRACT

Recent progress in Large Language Model (LLM) reasoning is increasingly driven by the refinement of post-training loss functions and alignment strategies. However, standard Reinforcement Learning (RL) paradigms like Group Relative Policy Optimization (GRPO) remain constrained by static uniformity: uniform prompt sampling and a fixed number of rollouts per prompt. For heterogeneous, heavy-tailed reasoning data, this creates structural inefficiencies that waste compute on already-solved patterns while under-training the long tail of hard problems. To address this, we propose Multi-Adversary Group Distributionally Robust Optimization (GDRO), an optimization-first framework that moves beyond uniform reasoning models by dynamically adapting the training distribution. We introduce an Online Difficulty Classifier that partitions prompts into dynamic pass@k difficulty groups. We then propose two independent GDRO games for post-training: (1) Prompt-GDRO, which employs an EMA-debiased multiplicative-weights bandit sampler to target the intensive difficulty margin and upweight persistently hard groups without frequency bias; and (2) Rollout-GDRO, which uses a shadow-price controller to reallocate rollouts across groups, maximizing gradient variance reduction on hard tasks under a fixed mean budget (compute-neutral). We provide no-regret guarantees for both controllers and additionally a variance-proxy analysis motivating a square-root optimal rollout allocation for Rollout-GDRO. We validate our framework on the DAPO 14.1k dataset using Qwen3-Base models. Prompt-GDRO and Rollout-GDRO achieve average relative gains of +10.6% and +10.1%, respectively, in pass@8 accuracy across 1.7B, 4B, and 8B scales compared to the GRPO baseline. Qualitative analysis shows an emergent curriculum: the adversaries shift resources to the evolving reasoning frontier, enhancing the reasoning model's performance.

연구 동기 및 목표

• 난이도 분포가 과중하게 꼬여 있는 경우 추론 과제에 대한 비균일 학습의 동기를 제시한다.
• 데이터-의존적이지 않은 온라인 난이도 분류기를 제안하여 프롬프트를 동적으로 그룹으로 분할한다.
• 두 개의 독립적인 GDRO 기반 적대자(Prompt-GDRO와 Rollout-GDRO)를 개발하여 샘플링과 계산 배치를 최적화한다.
• 엔트로피-정규화된 GDRO 및 분산 프록시 분석과의 이론적 연결을 제공한다.
• 다양한 모델 규모에서 DAPO 14.1k를 대상으로 경험적 개선을 보여준다.

제안 방법

• 동적 pass@k 기반 빈으로 프롬프트를 분할하는 온라인 난이도 분류기 정의
• EMA-편향 EXP3P를 사용하여 BIN 난이도에 따라 GRPO 업데이트의 가중치를 재조정하는 Prompt-GDRO 구현
• 평균 예산 제약 하에서 BIN 간 롤아웃을 할당하는 계산적 적대자로 Rollout-GDRO 구현
• UMA 보정(EMA 점수)을 사용하여 집중 손실(평균 손실)을 추적하고 빈도 편향을 피한다.
• 그림자 가격 μ를 사용하여 그래디언트 분산 감소를 최대화하는 제약 최적화로 Rollout-GDRO를 공식화
• 소프트-최악그룹 목적과 무-후회 보장을 보이는 엔트로피 GDRO 해석을 제공

Figure 1: Beyond Uniform Reasoning—A Multi-Adversary Post-Training Framework. Plots on the right represent training steps tail averages ( $\geq$ 60th percentile) capturing the curriculum. (Left) Our framework significantly outperforms the standard GRPO baseline across mathematical reasoning benchmar

실험 결과

연구 질문

• RQ1 동적이고 데이터에 의존하지 않는 난이도 그룹화가 정적 균일 샘플링에 비해 LLM 추론 학습 신호를 개선할 수 있는가?
• RQ2 두 개의 독립적인 GDRO 적대자(샘플링 및 롤아웃 예산)가 추론 과제를 위한 GRPO 대비 후처리에서 가산 이득을 가져오는가?
• RQ3 EMA-편향 및 분산 인식 할당이 최악의 그룹 강건성과 그래디언트 분산에 어떤 영향을 미치는가?
• RQ4 제안된 적대적 프레임워크를 뒷받침하는 이론적 보장이나 해석(엔트로피-정규화 GDRO 및 분산 프록시)이 있는가?
• RQ5 제안된 방법이 DAPO 추론 데이터셋의 서로 다른 모델 규모에서 측정 가능한 개선을 보이는가?

주요 결과

• Prompt-GDRO는 GRPO 대비 1.7B, 4B, 8B Qwen3-Base 모델에서 pass@8을 약 9.74%에서 13.13% 향상시킨다.
• Rollout-GDRO는 동일한 모델 규모에서 GRPO 대비 pass@8을 약 10.64%에서 9.20% 향상시킨다.
• 프레임워크는 진화하는 추론 프런티어로 자원이 전개되는 커리큘럼의 자발적 형성을 유도한다.
• EMA-편향 점수화는 빈도 편향을 피하고 다양한 활성 난이도 그룹을 유지한다.
• 이론적 기초는 Prompt-GDRO를 엔트로피-정규화 GDRO 대리 목표 및 무-후회 해석과 연결한다.
• 제곱근 법칙은 컴퓨트 중립 예산 하에서 분산 최적 롤아웃 배치를 정당화한다.

Figure 2: Conceptual Illustration: Static Uniformity vs. Multi-Adversary GDRO (Dynamic). (Left) Standard GRPO samples prompts uniformly ( $q=1/B$ ) and assigns a fixed number of rollouts (schematically $N=16$ ), causing it to overfit easy tasks while under-exploring the frontier. (Right) Our framewo