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

• URL: http://arxiv.org/abs/2606.24133v1
• Date: Tue, 23 Jun 2026 04:32:46 GMT
• Title: Holistic Data Scheduler for LLM Pre-training via Multi-Objective Reinforcement Learning
• Abstract summary: We introduce the Holistic Data Scheduler (HDS), a novel online data mixing framework.<n>HDS formulates the data scheduling challenge as a reinforcement learning problem in a continuous control space.<n>It integrates a data-driven reward for quality, a loss-driven reward capturing inter-domain influence, and a model-driven reward based on weight norms.
• Score: 7.431998098521794
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• 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.

Related papers

• DataFlex: A Unified Framework for Data-Centric Dynamic Training of Large Language Models [51.48564522455171]
  We present DataFlex, a unified data-centric dynamic training framework built upon LLaMA-Factory.<n>DataFlex supports three major paradigms of dynamic data optimization: sample selection, domain mixture adjustment adjustment, and data reweighting.<n>It provides trainer abstractions and modular components, enabling a drop-in replacement for standard LLM training.
  arXiv  Detail & Related papers  (2026-03-27T08:28:02Z)
• Data Mixing for Large Language Models Pretraining: A Survey and Outlook [52.66761353708619]
  Large language models (LLMs) rely on pretraining on massive and heterogeneous corpora.<n>Unlike sample-level data selection, data mixing optimize domain-level sampling weights to allocate limited budgets more effectively.<n>This paper provides a comprehensive review of data mixing for LLM pretraining.
  arXiv  Detail & Related papers  (2026-03-25T13:30:40Z)
• Towards High Data Efficiency in Reinforcement Learning with Verifiable Reward [54.708851958671794]
  We propose a Data-Efficient Policy Optimization pipeline that combines optimized strategies for both offline and online data selection.<n>In offline phase, we curate a high-quality subset of training samples based on diversity, influence, and appropriate difficulty.<n>During online RLVR training, we introduce a sample-level explorability metric to dynamically filter samples with low exploration potential.
  arXiv  Detail & Related papers  (2025-09-01T10:04:20Z)
• Ring-lite: Scalable Reasoning via C3PO-Stabilized Reinforcement Learning for LLMs [51.21041884010009]
  Ring-lite is a Mixture-of-Experts (MoE)-based large language model optimized via reinforcement learning (RL)<n>Our approach matches the performance of state-of-the-art (SOTA) small-scale reasoning models on challenging benchmarks.
  arXiv  Detail & Related papers  (2025-06-17T17:12:34Z)
• IDEAL: Data Equilibrium Adaptation for Multi-Capability Language Model Alignment [29.703775936837012]
  Large Language Models (LLMs) have achieved impressive performance through Supervised Fine-tuning (SFT) on diverse instructional datasets.<n>When training on multiple capabilities simultaneously, the mixture training dataset, governed by volumes of data from different domains, is a critical factor that directly impacts the final model's performance.<n>We introduce an innovative data equilibrium framework designed to effectively optimize volumes of data from different domains within mixture SFT datasets.
  arXiv  Detail & Related papers  (2025-05-19T06:42:44Z)
• ToReMi: Topic-Aware Data Reweighting for Dynamic Pre-Training Data Selection [28.75333303894706]
  ToReMi is a novel framework that adjusts training sample weights according to their topical associations and observed learning patterns.<n>Our experiments reveal that ToReMi variants consistently achieve superior performance over conventional pre-training approaches.
  arXiv  Detail & Related papers  (2025-04-01T12:06:42Z)
• Dynamic Loss-Based Sample Reweighting for Improved Large Language Model Pretraining [55.262510814326035]
  Existing reweighting strategies primarily focus on group-level data importance.<n>We introduce novel algorithms for dynamic, instance-level data reweighting.<n>Our framework allows us to devise reweighting strategies deprioritizing redundant or uninformative data.
  arXiv  Detail & Related papers  (2025-02-10T17:57:15Z)
• Take the Bull by the Horns: Hard Sample-Reweighted Continual Training Improves LLM Generalization [165.98557106089777]
  A key challenge is to enhance the capabilities of large language models (LLMs) amid a looming shortage of high-quality training data. Our study starts from an empirical strategy for the light continual training of LLMs using their original pre-training data sets. We then formalize this strategy into a principled framework of Instance-Reweighted Distributionally Robust Optimization.
  arXiv  Detail & Related papers  (2024-02-22T04:10:57Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.