多目的強化学習によるLLM事前学習のためのホリスティックデータスケジューラ

要旨

訓練データの構成は、そのソースの多様性と混合戦略によって決定され、大規模言語モデル（LLM）の事前学習における基盤である。オンラインデータ混合（ODM）は、訓練中にデータ混合を適応的に調整する手法であり、効率性を向上させる有望な方向性として注目されている。しかし、既存手法は単一の最適化視点に依存するという制約があり、複雑なLLM事前学習ではデータ構成を多次元から動的に考慮する必要があるという本質的な課題を見過ごしている。この限界を克服するために、我々は新たなオンラインデータ混合フレームワークである Holistic Data Scheduler（HDS）を提案する。HDSはデータスケジューリングの問題を連続制御空間における強化学習問題として定式化し、その安定性とサンプル効率の高さからSoft Actor-Critic（SAC）アルゴリズムを高次元の政策空間の探索に活用する。HDSの中核には、三つの重要な観点——品質を評価するデータ駆動型報酬、ドメイン間の影響を捉える損失駆動型報酬、そして重みノルムに基づくモデル駆動型報酬——を統合した、新規な多目的・総合的報酬関数が存在する。本設計の検証と最適構成の決定のため、我々は様々な規模のLLMに対して系統的な実験を実施した。The Pileベンチマークにおいて、HDSは次善手法の最終検証パープレキシティを44%少ない訓練反復回数で達成した。さらに、MMLU 0-shotタスクでは7.2%の改善に加え、他のベンチマークでも一貫した性能向上を示し、訓練効率と最終的なモデル能力の両方を向上させる能力を実証している。

English

The composition of training data, governed by the diversity of sources and their mixing strategy, is a cornerstone of Large Language Model (LLM) pre-training. Online Data Mixing (ODM), the technique of adaptively adjusting data mixtures during training, has emerged as a promising direction to improve efficiency. However, existing methods are constrained by their reliance on a singular optimization perspective, which fundamentally overlooks the need for complex LLM pre-training to consider the dynamic data composition from multiple dimensions. To overcome this limitation, we introduce the Holistic Data Scheduler (HDS), a novel online data mixing framework. HDS formulates the data scheduling challenge as a reinforcement learning problem in a continuous control space and leverages the Soft Actor-Critic (SAC) algorithm for its stability and sample efficiency in exploring the high-dimensional policy space. At the core of HDS lies a novel multi-objective, holistic reward function that integrates three critical perspectives: a data-driven reward for quality, a loss-driven reward capturing inter-domain influence, and a model-driven reward based on weight norms. To validate our design and determine its optimal configuration, we conducted systematic experiments on LLMs of various sizes. On The Pile benchmark, HDS reaches the final validation perplexity of the next best method with 44% fewer training iterations. Furthermore, it achieves a 7.2% improvement on the MMLU 0-shot task along with consistent gains on other benchmarks, showcasing its ability to enhance both training efficiency and final model capability.