[Paper Review] Group Distributionally Robust Optimization-Driven Reinforcement Learning for LLM Reasoning

Kishan Panaganti, Zhenwen Liang+3 more|arXiv (Cornell University)|Jan 27, 2026

Artificial Intelligence in Healthcare and Education0 citations

TL;DR

The paper introduces a Multi-Adversary GDRO framework that dynamically partitions prompts by online difficulty and allocates rollouts across groups to improve reasoning in LLMs, achieving notable gains over 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.

Motivation & Objective

• Motivate non-uniform training for reasoning tasks due to heavy-tailed difficulty distribution.
• Propose a data-agnostic online difficulty classifier to partition prompts into dynamic groups.
• Develop two independent GDRO-based adversaries (Prompt-GDRO and Rollout-GDRO) to optimize sampling and compute allocation.
• Provide theoretical connections to entropy-regularized GDRO and variance-proxy analysis.
• Demonstrate empirical improvements on DAPO 14.1k across multiple model scales.

Proposed method

• Define an Online Difficulty Classifier that partitions prompts into dynamic pass@k-based bins.
• Implement Prompt-GDRO with EMA-debiased EXP3P to reweight GRPO updates by bin difficulty.
• Implement Rollout-GDRO as a compute-adversary allocating rollouts across bins under a mean-budget constraint.
• Use an EMA score to track intensive (mean) losses and avoid frequency bias.
• Formulate Rollout-GDRO as a constrained optimization with a shadow price mu to maximize gradient variance reduction.
• Provide entropic GDRO interpretation showing a soft-worst-group objective and no-regret guarantees.

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

Experimental results

Research questions

• RQ1 Can dynamic, data-agnostic difficulty grouping improve learning signals in LLM reasoning compared to static uniform sampling?
• RQ2 Do two independent GDRO adversaries (sampling and rollout budgeting) yield additive gains over GRPO in post-training for reasoning tasks?
• RQ3 How do EMA-debiased and variance-aware allocations affect worst-group robustness and gradient variance?
• RQ4 What theoretical guarantees or interpretations support the proposed adversarial framework (entropy-regularized GDRO and variance proxy)?
• RQ5 Do the proposed methods yield measurable improvements across different model scales on the DAPO reasoning dataset?

Key findings

• Prompt-GDRO improves pass@8 by about 9.74% to 13.13% across 1.7B, 4B, and 8B Qwen3-Base models, compared to GRPO.
• Rollout-GDRO improves pass@8 by about 10.64% to 9.20% across the same model scales, relative to GRPO.
• The framework yields an emergent curriculum with resources shifting toward the evolving reasoning frontier.
• The EMA-debiased scoring avoids frequency bias and maintains diverse active difficulty groups.
• Theoretical grounding connects Prompt-GDRO to an entropy-regularized GDRO surrogate with no-regret interpretations.
• A square-root law motivates the variance-optimal rollout allocation under a compute-neutral budget.

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