Group Distributionally Robust Optimization-Driven Reinforcement Learning for LLM Reasoning

TL;DR

This work tackles inefficiencies in reasoning post-training caused by static uniform sampling and fixed rollout budgets. It introduces a Multi-Adversary GDRO framework that decomposes the training loop into two independent adversaries: Prompt-GDRO, which dynamically reweights prompts by online pass@k difficulty to form a traveling learning frontier, and Rollout-GDRO, which reallocates rollout compute under a mean-budget constraint to minimize gradient variance in hard tasks. The authors provide theoretical interpretations—entropic GDRO, no-regret dynamics, and a square-root variance-optimal rollout law—alongside toy validation and rigorous empirical results on the DAPO 14.1k dataset with Qwen3-Base models, achieving up to ~13% gains in pass@8. The results show an emergent curriculum and adaptive compute behavior that better concentrates learning on the evolving frontiers of difficulty, suggesting practical improvements for scalable, robust reasoning in large language models.

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.

Figures (9)

• Figure 1: Beyond Uniform Reasoning—A Multi-Adversary Post-Training Framework. Plots on the right represent training steps tail averages ($\ge$60th percentile) capturing the curriculum. (Left) Our framework significantly outperforms the standard GRPO baseline across mathematical reasoning benchmarks via dynamic adaptation. (Center) Prompt-GDRO: The adversary learns a non-uniform curriculum. Instead of uniform sampling (dashed line), probability mass (purple bars) shifts to the "reasoning frontier" (bins 6--8), targeting the specific difficulty level where learning is most efficient. (Right) Rollout-GDRO: The adversary optimizes compute utility. Under a fixed global budget (dashed line), it reallocates rollouts (orange bars) from solved tasks (bin 0) to high-variance tasks, scaling exploration with difficulty. Note: Bars represent rollout count per prompt (policy intensity).
• 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 framework employs an Online Difficulty Classifier to dynamically partition prompts based on real-time pass@k. It introduces two independent adversarial feedback loops (not coupled): (1) Prompt-GDRO (Data Distributor) uses an EMA-debiased scorer to shift the sampling distribution toward hard bins, creating a "traveling wave" of difficulty; (2) Rollout-GDRO (Resource Allocator) uses a shadow price $\mu$ to solve a constrained optimization problem, allocating discrete rollout arms ($n_{\min} \dots n_{\max}$) to maximize gradient variance reduction on high-uncertainty tasks.
• Figure 3: The Causal Chain of Curriculum. A triptych comparing the Prompt Distribution (Left), Adversarial Weights (Middle), and Realized Reward (Right) for 1.7B, 4B, and 8B models. This visualizes the mechanism: even when hard bins are rare in the data (dark regions in Left), the adversary applies disproportionate pressure (bright bands in Middle), forcing the model to eventually crack these problems and yield positive rewards (emergence of red/blue bands in Right). This effectively proves that Prompt-GDRO decouples the learning signal from dataset frequency.
• Figure 4: Training Dynamics Quantified (Prompt-GDRO). (Top) The Mean Accuracy Bin Index tracks the rising difficulty of targeted prompts. (Middle) The fraction of prompts in difficulty bands $\ge 3$ (solid) and $\ge 8$ (dashed) highlights scaling laws: 1.7B struggles to clear the $\ge 3$ bar, while 8B rapidly saturates even the $\ge 8$ band. (Bottom) Entropy metrics confirm that the adversary maintains a diverse portfolio of difficulty, preventing mode collapse to a single bin.
• Figure 5: Snapshots of the Learning Distribution. The probability mass of the training set across difficulty bins at four distinct training stages (Start, Early-Mid, Late-Mid, End). These publication-friendly checkpoints reveal the exact shape of the curriculum: note how the 4B model (Middle Row) transitions from a uniform start to a heavy emphasis on accbin_5--accbin_7 by mid-training, whereas the 8B model (Bottom Row) shifts almost its entire mass to the hardest bins by the final step.

Theorems & Definitions (40)

• Remark 1: Two notions of "group"
• Remark 2: Notation: $n$ vs. $k$
• Remark 3: Bins as dynamic groups
• Lemma 1: Entropic GDRO surrogate and softmax best response
• Theorem 4.1: Square-root law for variance-optimal allocation
• Lemma 2: Entropy-regularized maximum and log-sum-exp
• proof
• Corollary 1: Smooth approximation quality
• proof
• Remark 4: KL-robust view and "implicit" regularization