Holistic Data Scheduler for LLM Pre-training via Multi-Objective Reinforcement Learning

2026-06-23

AI Summary

Existing online data mixing methods are limited to a single optimization objective, making them ill-suited for the multidimensional and dynamic data composition requirements of large language model pretraining. This work formulates data scheduling as a reinforcement learning problem in a continuous control space and introduces the first multi-objective reward function that integrates data-driven, loss-driven, and model-driven perspectives. Leveraging the Soft Actor-Critic algorithm, the approach enables efficient exploration and policy optimization. Experimental results demonstrate that the proposed method achieves the validation perplexity of the best baseline on The Pile using only 56% of the training steps, while also delivering a 7.2% improvement on zero-shot MMLU performance and consistent gains across multiple benchmark evaluations.

Abstract

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.

Problem

• Online Data Mixing
• Large Language Model
• Multi-Objective Optimization
• Data Composition
• Pre-training

Innovation

• Holistic Data Scheduler
• Multi-Objective Reinforcement Learning
• Online Data Mixing
• Soft Actor-Critic
• LLM Pre-training