Group Distributionally Robust Optimization-Driven RL for LLM Reasoning

Abstract

Reasoning post-training with GRPO is typically built on *static uniformity*: uniform prompt sampling and a fixed number of rollouts per prompt. For heterogeneous, heavy-tailed reasoning data, this wastes compute on already-solved patterns while under-training the long tail of hard problems. We cast GRPO post-training as two independent GDRO games (not coupled) over *dynamic difficulty groups* defined by an *online train-time pass@$n$* statistic (computed from the rollouts used for GRPO updates): a *data adversary* that reshapes prompt sampling and a *compute adversary* that redistributes rollouts. **Prompt-GDRO** applies multiplicative-weights reweighting over bins (with an EMA-debiased difficulty score) to upweight persistently hard groups without frequency bias. **Rollout-GDRO** allocates rollouts across bins under a fixed mean budget via a shadow-price controller, improving gradient information efficiency on high-uncertainty groups while preserving the mean train-time rollout budget. Our approach is principled and theory-driven: we provide no-regret guarantees for the Prompt-GDRO game (via an entropy-regularized GDRO surrogate) and a variance-proxy analysis that yields a square-root optimal compute allocation for Rollout-GDRO. On DAPO 14.1k with Qwen3-Base (1.7B/4B/8B), each controller improves pass@8 by 9--13% over GRPO; targeted Qwen3-4B reruns reproduce the Prompt-GDRO gap under a second seed, show second-seed Rollout-GDRO gains on four of five comparable benchmarks, and show low sensitivity to several key hyperparameters.