Optimisation Robuste de Distribution de Groupe pour l'Apprentissage par Renforcement dans le Raisonnement des LLM

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

Les progrès récents en raisonnement des grands modèles de langage (LLM) sont de plus en plus pilotés par l'affinement des fonctions de perte post-entraînement et des stratégies d'alignement. Cependant, les paradigmes standards d'apprentissage par renforcement (RL) tels que l'Optimisation de Politique Relative par Groupe (GRPO) restent limités par une uniformité statique : un échantillonnage uniforme des prompts et un nombre fixe de rollouts par prompt. Pour des données de raisonnement hétérogènes et à queue lourde, cela crée des inefficacités structurelles qui gaspillent des ressources de calcul sur des patterns déjà résolus tout en sous-entraînant la longue queue des problèmes difficiles. Pour résoudre ce problème, nous proposons l'Optimisation Robuste de Distribution par Groupe Multi-Adversaire (GDRO), un framework axé sur l'optimisation qui dépasse les modèles de raisonnement uniformes en adaptant dynamiquement la distribution d'entraînement. Nous introduisons un Classificateur de Difficulté en Ligne qui partitionne les prompts en groupes de difficulté dynamiques pass@k. Nous proposons ensuite deux jeux GDRO indépendants pour le post-entraînement : (1) Prompt-GDRO, qui utilise un échantillonneur bandit à poids multiplicatifs débiaisé par moyenne mobile exponentielle (EMA) pour cibler la marge de difficulté intensive et augmenter le poids des groupes persistamment difficiles sans biais de fréquence ; et (2) Rollout-GDRO, qui utilise un contrôleur à prix shadow pour réallouer les rollouts entre les groupes, maximisant la réduction de la variance du gradient sur les tâches difficiles sous un budget moyen fixe (neutre en calcul). Nous fournissons des garanties de non-regret pour les deux contrôleurs ainsi qu'une analyse par proxy de la variance justifiant une allocation optimale des rollouts en racine carrée pour Rollout-GDRO. Nous validons notre framework sur le jeu de données DAPO 14.1k en utilisant les modèles Qwen3-Base. Prompt-GDRO et Rollout-GDRO obtiennent des gains relatifs moyens de +10,6 % et +10,1 %, respectivement, en précision pass@8 aux échelles 1,7B, 4B et 8B par rapport à la baseline GRPO. Une analyse qualitative révèle un curriculum émergent : les adversaires redirigent les ressources vers la frontière de raisonnement évolutive, améliorant les performances du modèle de raisonnement.

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.