AI Insight
This paper introduces AC-ODM, a reinforcement learning-based method for optimizing the composition of training data during large language model pretraining. The approach uses an actor-critic framework to dynamically adjust data mixing ratios, achieving significantly faster convergence than existing methods. On Pythia-1B models, AC-ODM reached optimal performance using 66% fewer training steps while improving downstream task accuracy by 27.5% on MMLU benchmarks and 2.23x on HumanEval coding tasks, with minimal computational overhead.
Why it matters
This work addresses a critical bottleneck in LLM training efficiency by reducing the computational resources needed to achieve high performance. The method's negligible overhead and substantial sample efficiency gains could significantly lower the environmental and financial costs of training large language models while improving their generalization capabilities.
Understand the Science
arXiv:2505.23878v2 Announce Type: replace
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.
Source: AC-ODM: Actor–Critic Online Data Mixing for Sample-Efficient LLM Pretraining