Group Distributionally Robust Optimization-Driven RL for LLM Reasoning

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.

When language models are trained to solve math problems, the usual recipe treats every training problem similarly and gives every problem the same number of practice attempts. This can be wasteful: many problems are already easy for the model, while rare or newly learnable problem types may need more attention. Our work studies how to make this training process adapt to the model’s current weaknesses. We introduce a simple online difficulty signal that watches how often the model succeeds on recent training attempts and groups problems by how hard they currently are. Using this signal, one version of our method gives more training weight to persistently difficult groups, while another gives more practice attempts to uncertain groups without increasing the average number of attempts. In this way, the training process automatically follows the model’s “reasoning frontier” instead of relying on a fixed curriculum. On several math-reasoning benchmarks with Qwen3 models, these adaptive choices improve the model’s chance of solving problems compared with standard GRPO training. The broader lesson is that better reasoning training may come not only from more data or more compute, but from using the available training budget more selectively.