Groepsgewijs Distributie-Robuste Optimalisatie-Gestuurde Reinforcement Learning voor LLM Redeneervaardigheid

Group Distributionally Robust Optimization-Driven Reinforcement Learning for LLM Reasoning

January 27, 2026

Auteurs: Kishan Panaganti, Zhenwen Liang, Wenhao Yu, Haitao Mi, Dong Yu

cs.AI

Samenvatting

Recente vooruitgang in het redeneren van Large Language Models (LLM's) wordt in toenemende mate gedreven door de verfijning van verliesfuncties na de training en alignmentstrategieën. Standaard Reinforcement Learning (RL)-paradigma's zoals Group Relative Policy Optimization (GRPO) blijven echter beperkt door statische uniformiteit: uniforme prompt sampling en een vast aantal rollouts per prompt. Voor heterogene, zwaarstaartige redeneergegevens creëert dit structurele inefficiënties die rekenkracht verspillen aan reeds opgeloste patronen, terwijl de lange staart van moeilijke problemen ondergetraind blijft. Om dit aan te pakken, stellen we Multi-Adversary Group Distributionally Robust Optimization (GDRO) voor, een optimalisatiegericht framework dat verder gaat dan uniforme redeneermodellen door de trainingsdistributie dynamisch aan te passen. We introduceren een Online Difficulty Classifier die prompts partitioneert in dynamische pass@k-moeilijkheidsgroepen. Vervolgens stellen we twee onafhankelijke GDRO-spellen voor na de training: (1) Prompt-GDRO, dat een EMA-gedebiased multiplicative-weights bandit sampler gebruikt om de intensieve moeilijkheidsmarge te targeten en hardnekkig moeilijke groepen op te waarderen zonder frequentiebias; en (2) Rollout-GDRO, dat een shadow-price controller gebruikt om rollouts over groepen te heralloceren, waardoor de variantiereductie van de gradient op moeilijke taken wordt gemaximaliseerd onder een vast gemiddeld budget (compute-neutraal). We geven no-regret garanties voor beide controllers en voegen een variantie-proxy analyse toe die een vierkantswortel-optimale rollout allocatie voor Rollout-GDRO motiveert. We valideren ons framework op de DAPO 14.1k dataset met Qwen3-Base modellen. Prompt-GDRO en Rollout-GDRO behalen respectievelijk gemiddelde relatieve winsten van +10.6% en +10.1% in pass@8 nauwkeurigheid over 1.7B, 4B en 8B schalen vergeleken met de GRPO-baseline. Kwalitatieve analyse toont een emergent curriculum: de adversaries verschuiven middelen naar de evoluerende redeneergrens, wat de prestaties van het redeneermodel verbetert.

English

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.