大規模言語モデル推論のためのグループ分布ロバスト最適化駆動強化学習

要旨

大規模言語モデル（LLM）の推論技術における最近の進展は、学習後損失関数とアライメント戦略の洗練によってますます推進されている。しかし、Group Relative Policy Optimization（GRPO）のような標準的な強化学習（RL）パラダイムは、静的な均一性、すなわち均一なプロンプトサンプリングとプロンプトあたりの固定数のロールアウトによって、未だ制約を受けている。異質的で裾の重い推論データにおいて、これは構造的非効率を生み出し、既に解決されたパターンに計算リソースを浪費する一方で、困難な問題のロングテールを十分に学習できなくしている。この問題に対処するため、我々は訓練分布を動的に適応させることで、均一な推論モデルを超える最適化優先フレームワークであるMulti-Adversary Group Distributionally Robust Optimization（GDRO）を提案する。我々は、プロンプトを動的なpass@k困難度グループに分割するOnline Difficulty Classifierを導入する。さらに、学習後のための二つの独立したGDROゲームを提案する：(1) **Prompt-GDRO** は、EMAでバイアス補正された乗法重み付けバンディットサンプラーを用いて、集中的な困難度マージンを標的とし、頻度バイアスなく持続的に困難なグループの重みを増加させる。(2) **Rollout-GDRO** は、シャドウプライス制御器を用いてグループ間でロールアウトを再配分し、固定平均予算（計算量中立）の下で困難なタスクにおける勾配分散削減を最大化する。我々は両制御器に対してno-regret保証を提供し、加えてRollout-GDROのための平方根最適ロールアウト配分を動機づける分散代理分析を提供する。我々はQwen3-Baseモデルを用いてDAPO 14.1kデータセット上で本フレームワークを検証する。Prompt-GDROとRollout-GDROは、1.7B、4B、8Bスケール全体でのpass@8精度において、GRPOベースラインと比較して、それぞれ平均+10.6%、+10.1%の相対的な向上を達成した。定性的分析は、敵対者がリソースを進化する推論フロンティアに移行することで推論モデルの性能が向上する、創発的なカリキュラムが現れることを示している。

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