Групповая дистрибутивно-робастная оптимизация в обучении с подкреплением для логического вывода больших языковых моделей

Group Distributionally Robust Optimization-Driven Reinforcement Learning for LLM Reasoning

January 27, 2026

Авторы: Kishan Panaganti, Zhenwen Liang, Wenhao Yu, Haitao Mi, Dong Yu

cs.AI

Аннотация

Последние достижения в области логического вывода больших языковых моделей (LLM) все в большей степени определяются совершенствованием функций потерь после обучения и стратегий согласования. Однако стандартные парадигмы обучения с подкреплением (RL), такие как Group Relative Policy Optimization (GRPO), остаются ограниченными статической унификацией: единообразным сэмплированием промптов и фиксированным количеством прогонов (rollouts) на промпт. Для гетерогенных данных логического вывода с тяжелыми хвостами это создает структурные неэффективности, приводящие к растрате вычислительных ресурсов на уже решенные шаблоны при недостаточном обучении на длинном хвосте сложных задач. Для решения этой проблемы мы предлагаем Multi-Adversary Group Distributionally Robust Optimization (GDRO) — оптимизационную среду, которая выходит за рамки моделей с равномерным выводом за счет динамической адаптации распределения обучения. Мы представляем Онлайн-классификатор сложности, который разделяет промпты на динамические группы сложности pass@k. Затем мы предлагаем две независимые GDRO-игры для пост-обучения: (1) Prompt-GDRO, которая использует EMA-дебиазированный мультипликативно-взвешенный бандитский сэмплер для таргетирования интенсивной границы сложности и повышения веса устойчиво сложных групп без частотного смещения; и (2) Rollout-GDRO, которая использует контроллер теневых цен для перераспределения прогонов между группами, максимизируя снижение дисперсии градиента на сложных задачах при фиксированном среднем бюджете (вычислительно-нейтральном). Мы предоставляем гарантии отсутствия сожаления для обоих контроллеров, а также анализ прокси-дисперсии, обосновывающий квадратично-оптимальное распределение прогонов для Rollout-GDRO. Мы проверяем нашу среду на наборе данных DAPO 14.1k с использованием моделей Qwen3-Base. Prompt-GDRO и Rollout-GDRO демонстрируют средний относительный прирост +10,6% и +10,1% соответственно в точности pass@8 для масштабов 1,7B, 4B и 8B по сравнению с базовым GRPO. Качественный анализ показывает emergence учебного плана: адверсарии переносят ресурсы на развивающийся фронтир логического вывода, повышая производительность модели.

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.