Ottimizzazione Robusta alla Distribuzione di Gruppo Guidata dall'Apprendimento per Rinforzo per il Ragionamento degli LLM

Group Distributionally Robust Optimization-Driven Reinforcement Learning for LLM Reasoning

January 27, 2026

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

cs.AI

Abstract

I recenti progressi nel ragionamento dei Large Language Model (LLM) sono sempre più guidati dal perfezionamento delle funzioni di loss post-addestramento e delle strategie di allineamento. Tuttavia, i paradigmi standard di Reinforcement Learning (RL) come il Group Relative Policy Optimization (GRPO) rimangono limitati da una staticità uniforme: campionamento uniforme dei prompt e un numero fisso di rollout per prompt. Per dati di ragionamento eterogenei e dalla distribuzione heavy-tailed, ciò crea inefficienze strutturali che sprecano risorse computazionali su pattern già risolti, mentre sottopongono a un addestramento insufficiente la coda lunga dei problemi difficili. Per affrontare questo problema, proponiamo il Multi-Adversary Group Distributionally Robust Optimization (GDRO), un framework di tipo optimization-first che supera i modelli di ragionamento uniformi adattando dinamicamente la distribuzione di addestramento. Introduciamo un Classificatore Online della Difficoltà che partiziona i prompt in gruppi di difficoltà dinamici pass@k. Proponiamo quindi due giochi GDRO indipendenti per il post-addestramento: (1) Prompt-GDRO, che impiega un campionatore bandit a pesi moltiplicativi con debiasing EMA per mirare al margine di difficoltà intensivo e aumentare il peso dei gruppi persistentemente difficili senza bias di frequenza; e (2) Rollout-GDRO, che utilizza un controllore shadow-price per riallocare i rollout tra i gruppi, massimizzando la riduzione della varianza del gradiente sui task difficili sotto un budget medio fisso (computazionalmente neutro). Forniamo garanzie di tipo no-regret per entrambi i controllori e, inoltre, un'analisi di proxy della varianza che motiva un'allocazione ottimale dei rollout proporzionale alla radice quadrata per Rollout-GDRO. Convalidiamo il nostro framework sul dataset DAPO 14.1k utilizzando modelli Qwen3-Base. Prompt-GDRO e Rollout-GDRO raggiungono guadagni relativi medi rispettivamente del +10,6% e del +10,1% in accuratezza pass@8 attraverso le scale 1.7B, 4B e 8B, rispetto al baseline GRPO. L'analisi qualitativa mostra un curriculum emergente: gli adversary spostano le risorse verso la frontiera di ragionamento in evoluzione, migliorando le prestazioni del modello di ragionamento.

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.