대규모 언어 모델 추론을 위한 그룹 분포 강건 최적화 기반 강화 학습

초록

대규모 언어 모델(LLM)의 추론 분야 최근 발전은 사후 훈련 손실 함수 및 정렬 전략의 정교화에 의해 주도되고 있습니다. 그러나 GRPO(Group Relative Policy Optimization)와 같은 표준 강화 학습(RL) 패러다임은 여전히 정적 균일성, 즉 균일한 프롬프트 샘플링과 프롬프트당 고정된 롤아웟 수에 의해 제약을 받습니다. 이질적이고 두꺼운 꼬리 분포를 보이는 추론 데이터의 경우, 이는 구조적 비효율성을 초래하여 이미 해결된 패턴에 대한 계산 자원을 낭비하는 동시에 어려운 문제들의 긴 꼬리 부분을 충분히 훈련하지 못하게 합니다. 이를 해결하기 위해 우리는 균일 추론 모델을 넘어서 훈련 분포를 동적으로 조정하는 최적화 중심 프레임워크인 Multi-Adversary Group Distributionally Robust Optimization(GDRO)을 제안합니다. 우리는 프롬프트를 동적 pass@k 난이도 그룹으로 분할하는 온라인 난이도 분류기(Online Difficulty Classifier)를 도입했습니다. 그런 다음 사후 훈련을 위한 두 가지 독립적인 GDRO 게임을 제안합니다: (1) 집중적인 난이도 마진을 목표로 하고 빈도 편향 없이 지속적으로 어려운 그룹의 가중치를 높이는 EMA-편향 제거 승법 가중치 밴딧 샘플러(EMA-debiased multiplicative-weights bandit sampler)를 사용하는 Prompt-GDRO와, (2) 고정된 평균 예산(계산량 중립) 하에서 어려운 과제에 대한 그래디언트 분산 감소를 극대화하기 위해 그룹 간 롤아웟을 재배분하는 섀도우 가격 컨트롤러(shadow-price controller)를 사용하는 Rollout-GDRO입니다. 우리는 두 컨트롤러에 대해 무감회(no-regret) 보장을 제공하고, 추가적으로 Rollout-GDRO를 위한 제곱근 최적 롤아웟 할당을 유도하는 분산 프록시(variance-proxy) 분석을 제시합니다. 우리는 Qwen3-Base 모델을 사용하여 DAPO 14.1k 데이터셋에서 우리 프레임워크를 검증했습니다. Prompt-GDRO와 Rollout-GDRO는 GRPO 기준선 대비 1.7B, 4B, 8B 규모에서 pass@8 정확도 기준 평균 각각 +10.6%, +10.1%의 상대적 향상을 달성했습니다. 정성적 분석은 적대자(adversaries)가 자원을 진화하는 추론 최전선(frontier)으로 이동시켜 추론 모델의 성능을 향상시키는 발현적 커리큘럼(emergent curriculum)을 보여줍니다.

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.

대규모 언어 모델 추론을 위한 그룹 분포 강건 최적화 기반 강화 학습

Group Distributionally Robust Optimization-Driven Reinforcement Learning for LLM Reasoning

초록

Support