Optimización de Aprendizaje por Refuerzo Basada en Robustez Distribucional Grupal para el Razonamiento de Modelos de Lenguaje Grandes

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

Resumen

Los recientes avances en el razonamiento de los Modelos de Lenguaje a Gran Escala (LLM) están cada vez más impulsados por el refinamiento de las funciones de pérdida posteriores al entrenamiento y las estrategias de alineación. Sin embargo, los paradigmas estándar de Aprendizaje por Refuerzo (RL), como la Optimización de Políticas Relativas por Grupos (GRPO), siguen limitados por una uniformidad estática: muestreo uniforme de prompts y un número fijo de rollouts por prompt. Para datos de razonamiento heterogéneos y de cola pesada, esto crea ineficiencias estructurales que desperdician capacidad de cómputo en patrones ya resueltos, mientras que se entrena insuficientemente la larga cola de problemas difíciles. Para abordar esto, proponemos la Optimización Robusta de Distribución por Grupos Multi-Adversario (GDRO), un marco de optimización-primero que va más allá de los modelos de razonamiento uniformes al adaptar dinámicamente la distribución de entrenamiento. Introducimos un Clasificador de Dificultad Online que divide los prompts en grupos dinámicos de dificultad según pass@k. Luego proponemos dos juegos GDRO independientes para el post-entrenamiento: (1) Prompt-GDRO, que emplea un muestreador bandit de pesos multiplicativos con corrección de sesgo por Media Móvil Exponencial (EMA) para apuntar al margen de dificultad intensivo y aumentar el peso de los grupos persistentemente difíciles sin sesgo de frecuencia; y (2) Rollout-GDRO, que utiliza un controlador de precio sombra para reasignar rollouts entre grupos, maximizando la reducción de la varianza del gradiente en tareas difíciles bajo un presupuesto medio fijo (neutral en cómputo). Proporcionamos garantías de no-lamento para ambos controladores y, además, un análisis proxy de varianza que motiva una asignación óptima de rollouts en raíz cuadrada para Rollout-GDRO. Validamos nuestro marco en el conjunto de datos DAPO 14.1k utilizando modelos Qwen3-Base. Prompt-GDRO y Rollout-GDRO logran ganancias relativas promedio de +10.6% y +10.1%, respectivamente, en la precisión pass@8 a través de escalas de 1.7B, 4B y 8B en comparación con la línea base GRPO. El análisis cualitativo muestra un currículo emergente: los adversarios desplazan recursos hacia la frontera de razonamiento en evolución, mejorando el rendimiento del modelo de razonamiento.

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.