Otimização de Grupo com Robustez Distribucional Orientada ao Aprendizado por Reforço para Raciocínio em LLM

Group Distributionally Robust Optimization-Driven Reinforcement Learning for LLM Reasoning

January 27, 2026

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

cs.AI

Resumo

Os recentes avanços no raciocínio de Modelos de Linguagem de Grande Porte (LLMs) são cada vez mais impulsionados pelo refinamento das funções de perda pós-treinamento e das estratégias de alinhamento. No entanto, os paradigmas padrão de Aprendizado por Reforço (RL), como a Otimização de Política Relativa de Grupo (GRPO), permanecem limitados pela uniformidade estática: amostragem uniforme de prompts e um número fixo de rollouts por prompt. Para dados de raciocínio heterogêneos e de cauda pesada, isso cria ineficiências estruturais que desperdiçam recursos computacionais em padrões já resolvidos, enquanto subtreinam a longa cauda de problemas difíceis. Para resolver isso, propomos a Otimização Robusta de Distribuição de Grupo Multiadversário (GDRO), uma estrutura orientada à otimização que vai além dos modelos de raciocínio uniformes, adaptando dinamicamente a distribuição de treinamento. Introduzimos um Classificador de Dificuldade Online que particiona os prompts em grupos dinâmicos de dificuldade pass@k. Em seguida, propomos dois jogos GDRO independentes para o pós-treinamento: (1) o Prompt-GDRO, que emprega um amostrador bandit com pesos multiplicativos e correção de viés por Média Móvel Exponencial (EMA) para direcionar a margem de dificuldade intensa e aumentar o peso de grupos persistentemente difíceis sem viés de frequência; e (2) o Rollout-GDRO, que usa um controlador de preço-sombra para realocar rollouts entre os grupos, maximizando a redução da variância do gradiente em tarefas difíceis sob um orçamento médio fixo (computacionalmente neutro). Fornecemos garantias de não-arrependimento para ambos os controladores e adicionalmente uma análise proxy de variância que motiva uma alocação ótima de rollouts com raiz quadrada para o Rollout-GDRO. Validamos nossa estrutura no conjunto de dados DAPO 14.1k usando modelos Qwen3-Base. O Prompt-GDRO e o Rollout-GDRO alcançam ganhos relativos médios de +10,6% e +10,1%, respectivamente, na precisão pass@8 nas escalas de 1,7B, 4B e 8B, em comparação com a linha de base GRPO. A análise qualitativa mostra um currículo emergente: os adversários deslocam recursos para a fronteira de raciocínio em evolução, melhorando o desempenho do modelo de raciocínio.

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.