AC-ODM: Actor--Critic Online Data Mixing for Sample-Efficient LLM Pretraining

2025/05/29 by Jing Ma, Ma, Jing, Dang, Chenhao +2
Computer Science · #Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Natural Language Processing Techniques #Semantic Web and Ontologies #Speech and dialogue systems

paper · pdf · doi:10.48550/arxiv.2505.23878

openalex publication_date 2025/05/29 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Optimizing pretraining data composition is pivotal for LLM generalization. While dynamic mixing outperforms static strategies by capturing evolving training dynamics, current methods fail to reconcile computational efficiency with sample efficiency and structural flexibility for diverse pipelines.We introduce Actor--Critic Online Data Mixing (AC-ODM), which approaches data mixing from a reinforcement learning perspective with a parameterized policy that we theoretically prove to act as a dynamic linear surrogate maximizing the constructive interference of gradients. To enhance practical flexibility, AC-ODM supports two operational modes: (i) a proxy mode for fixed, pre-prepared corpora, where a policy learned on a small model is transferred to a larger target; and (ii) a non-proxy mode for direct end-to-end training from scratch without priors. Empirically, AC-ODM significantly outperforms prior methods in convergence speed and downstream accuracy across various architectures. On Pythia-1B, it reaches optimal validation perplexity using up to 66% fewer training steps than competitive baselines, delivering a 27.5% relative improvement in MMLU accuracy and a 2.23 x higher pass@1 on HumanEval, all while incurring a virtually negligible (0.4%) per-step wall-clock increase and only 2% additional memory overhead. Code is available at https://github.com/DANG-ai/AC-ODM.

Citations

• ADMIRE-BayesOpt: Accelerated Data MIxture RE-weighting for Language Models with Bayesian Optimization
• LLM Data Selection and Utilization via Dynamic Bi-level Optimization
• Scaling Laws for Optimal Data Mixtures
• Chameleon: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
• Qwen3 Technical Report
• QuaDMix: Quality-Diversity Balanced Data Selection for Efficient LLM Pretraining
• Meta-rater: A Multi-dimensional Data Selection Method for Pre-training Language Models
• Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training
• Data Mixture Optimization: A Multi-fidelity Multi-scale Bayesian Framework
• PiKE: Adaptive Data Mixing for Large-Scale Multi-Task Learning Under Low Gradient Conflicts
• The Llama 3 Herd of Models
• RegMix: Data Mixture as Regression for Language Model Pre-training
• Efficient Online Data Mixing For Language Model Pre-Training
• Beyond neural scaling laws: beating power law scaling via data pruning
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• GLaM: Efficient Scaling of Language Models with Mixture-of-Experts
• Deduplicating Training Data Makes Language Models Better
• Evaluating Large Language Models Trained on Code
• RoFormer: Enhanced Transformer with Rotary Position Embedding
• The Pile: An 800GB Dataset of Diverse Text for Language Modeling
• Measuring Massive Multitask Language Understanding
• PIQA: Reasoning about Physical Commonsense in Natural Language
• Crowdsourcing Multiple Choice Science Questions

Cited by

• A Component-Based Survey of Interactions between Large Language Models and Multi-Armed Bandits

Related

• Efficient Online Data Mixing For Language Model Pre-Training
• SampleMix: A Sample-wise Pre-training Data Mixing Strategey by Coordinating Data Quality and Diversity
• Multi-armed bandits for resource efficient, online optimization of language model pre-training: the use case of dynamic masking
• Data Mixing Agent: Learning to Re-weight Domains for Continual Pre-training
• DataMan: Data Manager for Pre-training Large Language Models