PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chenhao Dang, Jing Ma, Mingjie Liao

arXiv:2606.24133v1cs.LGcs.CL

TL;DR

Existing online data-mixing methods do not fully account for the multiple dimensions of LLM pre-training data composition. HDS addresses this with continuous-control reinforcement learning and a multi-objective reward, matching the strongest baseline’s perplexity with 44% fewer training steps while improving MMLU 0-shot performance by 7.2%.

  • Problem

    Existing online data-mixing approaches inadequately account for the multiple dimensions of data composition required during complex LLM pre-training.

  • Method

    HDS formulates online data mixing as continuous-control reinforcement learning with SAC and a holistic reward integrating data quality, inter-domain influence, and model stability.

  • Results

    HDS outperforms existing baselines across training efficiency and downstream capability, matching AC-ODM’s final perplexity with 44% fewer training steps and improving MMLU 0-shot performance by 7.2%.

  • Takeaways & Limitations

    HDS demonstrates that holistic, multi-objective data scheduling can accelerate LLM pre-training while producing stronger downstream benchmark performance.

Abstract

from arXiv · show

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.

1 Introduction

The introduction presents data composition as central to efficient and capable LLM pre-training, motivating HDS as an online data-mixing framework that uses continuous-control reinforcement learning and a holistic multi-objective reward. HDS combines data quality, inter-domain influence, and model stability signals, achieving faster convergence and improved MMLU performance on Pythia-1B pre-training.

  • Motivation: Pre-training data quality, diversity, and mixing strategy strongly influence LLM capabilities, training efficiency, and economic and environmental costs.The introduction frames the field’s shift from compute-centric Scaling Laws toward Data-Centric AI.
  • Problem: Existing data-mixing methods are static offline or dynamic online, but commonly optimize domain weights from a single perspective rather than multiple dimensions.The paper identifies holistic consideration of data composition as crucial for complex LLM pre-training.
  • Method: HDS formulates online data mixing as a continuous-control reinforcement-learning task and applies Soft Actor-Critic for stable, sample-efficient exploration of high-dimensional policy spaces.The LLM pre-training process serves as the environment, while actions specify subsequent batch sampling probabilities.
  • Reward Design: HDS’s holistic reward integrates data-driven quality, gradient-driven inter-domain influence, and model-centric weight-norm signals for model stability.The reward is designed as a multi-objective system that captures distinct aspects of training dynamics.
  • Results: 22% fewer iterations let HDS reach the previous state-of-the-art final validation perplexity on Pythia-1B trained on The Pile, while MMLU 0-shot performance improved 7.2%.HDS also required 57% fewer iterations than the static The Pile Weights baseline.

2 Methodology: The Holistic Data Scheduler

HDS formulates online data scheduling as a continuous-control Markov decision process, using SAC to adapt domain-mixture weights from comprehensive training-state observations. Its holistic reward combines inter-domain influence, scheduled lexical diversity, and model-stability objectives to guide efficient and stable LLM pre-training.

  • SAC optimization: SAC optimizes the high-dimensional data-mixing policy because of its stability and sample efficiency, while online policies adapt mixtures to the LLM’s learning state.Transitions are stored in replay buffer B, and twin critics mitigate overestimation bias.
  • MDP formulation: HDS models dynamic data scheduling as an MDP with state space S, action space A, transition function P, reward R, and discount factor γ.The transition function is implicitly defined by LLM training dynamics.
  • State and action spaces: The state vector includes data statistics, training step, domain losses, loss deltas, selected-layer weight norms, and weight-norm changes.This representation combines historical trends with the LLM’s instantaneous training state.
  • State and action spaces: The continuous action is a probability-simplex domain-weight vector used to sample the next training batch, and HDS generates it with a stochastic actor policy.The policy samples a specific domain-weight vector from a distribution over continuous actions.
  • Holistic reward: HDS combines inter-domain influence, scheduled lexical diversity, and model stability into a weighted holistic reward.The components respectively capture cross-domain utility, a progression from lower- to higher-MTLD text, and smooth parameter convergence.

3 Experiments and Analysis

Experiments show that HDS improves training efficiency and downstream performance while adding minimal computational overhead. Ablations further indicate that inter-domain influence, lexical-diversity curriculum, and stability rewards jointly drive its gains.

  • Training efficiency: 44% fewer training steps let HDS reach AC-ODM’s final validation perplexity, while matching TPW’s perplexity requires 57% fewer steps.At 41,667 steps, HDS achieves validation perplexity 13.6% lower than TPW and 8.9% lower than ODM.
  • Downstream performance: 0.26915 MMLU 0-shot accuracy and 0.31064 5-shot accuracy make HDS outperform all baselines in both settings.HDS surpasses all baselines by 20k steps and maintains its lead through training.
  • Downstream performance: 0.71000 COPA, 0.69840 PIQA, and 0.65120 WinoGrande accuracies are the highest among methods, while HDS remains competitive on LogiQA at 0.30550.HDS outperforms baselines on three of four evaluated zero-shot datasets; ODM has a marginal advantage on LogiQA.
  • Computational cost: 26.5M added parameters increase model size by approximately 2.6%, while per-step time rises from 2.47s to 2.49s.The per-step increase is less than 1%, indicating minimal operational overhead.
  • Computational cost: 17,917 steps reach target perplexity, corresponding to a 2.21x speedup relative to baseline training time.The efficiency gain is attributed to superior sample quality from the HDS policy at each step.
  • Reward ablation: The complete HDS model achieves the lowest validation perplexity, with removing r_align causing the largest degradation and removing r_diversity producing a persistent performance gap.Removing r_stability performs better than the other ablations but still trails full HDS, acting as a regularizer against erratic updates.

4 Hyperparameter optimization

HDS hyperparameter experiments identify practical scaling choices for its state representation, alignment reward, and agent networks. Layer sampling broadens state coverage efficiently, while an agent network of approximately 0.5% of the LLM’s parameters balances performance and computational cost.

  • Layer selection: Layer sampling, such as selecting all even- or odd-indexed layers, effectively trades computational cost for broader state coverage.Using more LLM layers generally improves performance but increases computational cost.
  • Agent network size: 0.5% of the LLM’s parameters is the recommended agent-network size for near-optimal performance without unnecessary computational overhead.Performance generally improves with larger Actor and Critic networks, but gains diminish as size increases.
  • Layer selection: Selecting a few layers from the model’s latter half improves r_align by better reflecting higher-level semantic understanding and gradient alignment.This is recommended for the alignment-reward calculation.
  • Scaling validation: 1.07 perplexity points separated HDS from the ODM baseline at step 5,208 during Pythia-12B pretraining on The Pile.HDS consistently achieved lower validation perplexity than the ODM baseline at every checkpoint, with the gap widening during training.

5 Conclusion

HDS is introduced as an online data mixing framework that formulates LLM pre-training as multi-objective reinforcement learning in a continuous control space. It uses a Soft Actor-Critic agent with a holistic reward integrating inter-domain influence, scheduled lexical diversity, and model stability.

  • HDS reconceptualizes LLM pre-training as a multi-objective reinforcement learning problem in a continuous control space.
  • HDS uses a Soft Actor-Critic agent guided by a holistic reward function.
  • Its reward integrates inter-domain influence, scheduled lexical diversity, and model stability to capture synergies, enforce curriculum, and regularize updates.

A LLM Model Configuration

The language model uses a 16-layer Transformer with a 2048-dimensional hidden state, 16 attention heads, and fixed 1024-token sequences. It encodes positions with RoPE and implements self-attention with FlashAttention for computational efficiency.

  • The model is a 16-layer Transformer with hidden dimension 2048, 16 attention heads, and fixed sequence length 1024 tokens.
  • RoPE provides positional encoding, while FlashAttention implements self-attention to optimize memory access patterns and reduce training time.

B Sensitivity Analysis of Reward Weights

Sensitivity analysis at 20,832 steps identifies [1] [10] as the best reward-weight configuration, while removing any component degrades validation perplexity. Increasing alignment or diversity weights slightly hurts performance, whereas increasing stability preserves the optimum, motivating [1] [10] as the default.

  • Configuration selection: At 20,832 steps, the baseline configuration [1] [10] achieves the optimal validation perplexity of 12.73.The weight order is w_align, w_diversity, w_stability.
  • Configuration selection: Removing any single reward component degrades performance, with w_align = 0 producing the largest perplexity increase to 13.45.This degradation is consistent with the main ablation study.
  • Weight sensitivity: Doubling w_align to 2 or w_diversity to 20 causes slight regressions to 12.81 and 12.78, respectively.These results suggest that over-emphasizing alignment or diversity disrupts the learning balance.
  • Weight sensitivity: Increasing w_stability to 20 maintains the optimal perplexity of 12.73, so [1] [10] is retained as the default configuration.The default minimizes hyperparameter complexity while preserving peak performance.

C Hyperparameter Sensitivity Analysis

This section reports hyperparameter optimization results for HDS, focusing on layer-selection strategies for the state representation and the alignment reward (𝑟align).

  • Hyperparameter optimization: The experiments evaluate hyperparameter sensitivity through detailed optimization results.The analysis specifically examines configuration choices affecting the scheduler.
  • State representation: The state representation is analyzed under different layer-selection strategies.Layer selection is one of the section’s two stated experimental focuses.
  • Alignment reward: The alignment reward (𝑟align) is another focus of the hyperparameter analysis.The section reports optimization experiments involving the alignment reward.

C.1 Layer Selection Analysis

Layer selection substantially affects the model-centric state representation: using more layers improves final perplexity, while layer sampling offers a more computationally efficient alternative. All 24 layers achieve the best performance but at the highest computational cost.

  • C.1 Layer Selection Analysis: Using a broader set of layers produces a more comprehensive state representation and better final perplexity.The layer-selection ablation evaluates the weight-norm component of the state vector.
  • C.1 Layer Selection Analysis: All 24 layers deliver the optimal final perplexity, but they incur the highest computational cost.
  • C.1 Layer Selection Analysis: Sampling all even or odd layers provides an efficient way to construct a rich yet manageable state representation.The passage identifies layer sampling as practical for large-scale applications with constrained computational resources.

C.2 Alignment Reward

Layer selection for the alignment reward has a subtle but discernible effect on performance. Deeper Transformer layers are more suitable because they capture abstract, semantic, and transferable knowledge, while earlier layers perform slightly worse.

  • Layer Selection: Layer selection for calculating r_align has a subtle but discernible impact on performance.Table 8 reports final perplexity on The Pile test set for the layer-selection ablation.
  • Layer Selection: Deeper Transformer layers better capture abstract and semantic information relevant to the transferable knowledge promoted by the alignment reward.This motivates using deeper layers for r_align.
  • Layer Selection: Using earlier layers, such as blocks 2, 4, and 6, results in slightly worse performance, although the differences are not substantial.Earlier layers emphasize syntactic and local features rather than more abstract representations.

D Analysis of Model Weight Norm Evolution

HDS weight-norm trajectories across monitored layers show an initially volatile phase followed by gradual convergence to a stable equilibrium, highlighting the role of the model stability reward.

  • Weight norm evolution: Across all monitored layers, selected-weight L2 norms transition from high initial volatility to gradual convergence at a stable equilibrium.The trajectories are presented in Figure 7 and reflect the model stability reward, r_stability.
Loading 2606.24133v1…