DUMP：RLベースのLLM向け自動分散レベルカリキュラム学習 ポストトレーニング

DUMP: Automated Distribution-Level Curriculum Learning for RL-based LLM Post-training

April 13, 2025

著者: Zhenting Wang, Guofeng Cui, Kun Wan, Wentian Zhao

cs.AI

要旨

強化学習（RL）ベースのポストトレーニングにおける最近の進展により、大規模言語モデル（LLMs）の推論能力が向上し、複雑なタスクを処理する能力が顕著に改善されました。しかし、既存の手法の多くは、トレーニングデータを統一された全体として扱い、現代のLLMトレーニングが多様な分布（ソースと難易度が異なる）からなるデータの混合物を含むという事実を見落としています。この異質性は、学習効率を最適化するために分布間でトレーニングを適応的にスケジュールする方法という重要な課題を引き起こします。本論文では、分布レベルの学習可能性の概念に基づいた原則的なカリキュラム学習フレームワークを提示します。私たちの核心的な洞察は、ポリシーの優位性の大きさが、特定の分布でのさらなるトレーニングからモデルがどれだけ利益を得られるかを反映しているということです。これに基づいて、RLベースのLLMポストトレーニングのための分布レベルのカリキュラム学習フレームワークを提案します。このフレームワークは、Upper Confidence Bound（UCB）原則を活用して、異なる分布に対するサンプリング確率を動的に調整します。このアプローチは、高い平均優位性（活用）または低いサンプル数（探索）を持つ分布を優先し、適応的で理論的に裏付けられたトレーニングスケジュールを実現します。私たちは、GRPOを基盤とするRLアルゴリズムを用いてカリキュラム学習フレームワークを具体化し、複数の難易度とソースを持つ論理推論データセットでその有効性を実証します。実験結果は、私たちのフレームワークが収束速度と最終的なパフォーマンスを大幅に改善し、LLMポストトレーニングにおける分布を意識したカリキュラム戦略の価値を強調しています。コード: https://github.com/ZhentingWang/DUMP.

English

Recent advances in reinforcement learning (RL)-based post-training have led to notable improvements in large language models (LLMs), particularly in enhancing their reasoning capabilities to handle complex tasks. However, most existing methods treat the training data as a unified whole, overlooking the fact that modern LLM training often involves a mixture of data from diverse distributions-varying in both source and difficulty. This heterogeneity introduces a key challenge: how to adaptively schedule training across distributions to optimize learning efficiency. In this paper, we present a principled curriculum learning framework grounded in the notion of distribution-level learnability. Our core insight is that the magnitude of policy advantages reflects how much a model can still benefit from further training on a given distribution. Based on this, we propose a distribution-level curriculum learning framework for RL-based LLM post-training, which leverages the Upper Confidence Bound (UCB) principle to dynamically adjust sampling probabilities for different distrubutions. This approach prioritizes distributions with either high average advantage (exploitation) or low sample count (exploration), yielding an adaptive and theoretically grounded training schedule. We instantiate our curriculum learning framework with GRPO as the underlying RL algorithm and demonstrate its effectiveness on logic reasoning datasets with multiple difficulties and sources. Our experiments show that our framework significantly improves convergence speed and final performance, highlighting the value of distribution-aware curriculum strategies in LLM post-training. Code: https://github.com/ZhentingWang/DUMP.