Gruppenverteilungsrobuste optimierungsgesteuerte Verstärkungslernverfahren für das logische Schließen großer Sprachmodelle

Group Distributionally Robust Optimization-Driven Reinforcement Learning for LLM Reasoning

January 27, 2026

papers.authors: Kishan Panaganti, Zhenwen Liang, Wenhao Yu, Haitao Mi, Dong Yu

cs.AI

papers.abstract

Jüngste Fortschritte im Reasoning großer Sprachmodelle (LLMs) werden zunehmend durch die Verfeinerung von Loss-Funktionen nach dem Training und Alignment-Strategien vorangetrieben. Standard-Paradigmen des Reinforcement Learning (RL) wie Group Relative Policy Optimization (GRPO) bleiben jedoch durch statische Uniformität eingeschränkt: uniformes Prompt-Sampling und eine feste Anzahl von Rollouts pro Prompt. Bei heterogenen, schwer tail-verteilten Reasoning-Daten führt dies zu strukturellen Ineffizienzen, da Rechenleistung für bereits gelöste Muster verschwendet wird, während der lange Schwanz schwieriger Probleme untertrainiert bleibt. Um dies zu adressieren, schlagen wir Multi-Adversary Group Distributionally Robust Optimization (GDRO) vor, ein optimierungsorientiertes Framework, das über uniforme Reasoning-Modelle hinausgeht, indem es die Trainingsverteilung dynamisch anpasst. Wir führen einen Online Difficulty Classifier ein, der Prompts in dynamische Pass@k-Schwierigkeitsgruppen unterteilt. Anschließend schlagen wir zwei unabhängige GDRO-Spiele für die Phase nach dem Training vor: (1) Prompt-GDRO, das einen EMA-entzerrten Multiplicative-Weights-Bandit-Sampler einsetzt, um den intensiven Schwierigkeitsrand zu adressieren und persistent schwierige Gruppen ohne Frequenzverzerrung höher gewichtet; und (2) Rollout-GDRO, das einen Shadow-Price-Controller verwendet, um Rollouts gruppenübergreifend neu zuzuteilen und so die Reduktion der Gradientenvarianz bei schwierigen Aufgaben unter einem festen Mittelwert-Budget (rechenneutral) maximiert. Wir liefern No-Regret-Garantien für beide Controller sowie zusätzlich eine Varianz-Proxy-Analyse, die eine quadratwurzeloptimale Rollout-Zuteilung für Rollout-GDRO motiviert. Wir validieren unser Framework auf dem DAPO-14.1k-Datensatz mit Qwen3-Base-Modellen. Prompt-GDRO und Rollout-GDRO erzielen durchschnittliche relative Verbesserungen von +10,6 % bzw. +10,1 % in der Pass@8-Genauigkeit über 1,7B-, 4B- und 8B-Modelle im Vergleich zur GRPO-Baseline. Eine qualitative Analyse zeigt einen emergenten Lehrplan: Die Adversaries verlagern Ressourcen zur sich entwickelnden Reasoning-Front und verbessern so die Leistung des Reasoning-Modells.

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.