Search papers, labs, and topics across Lattice.
This paper introduces Actor-Critic Online Data Mixing (AC-ODM), a novel approach that leverages reinforcement learning to optimize the composition of pretraining data for large language models (LLMs). By employing a parameterized policy that maximizes constructive gradient interference, AC-ODM achieves significant improvements in both convergence speed and downstream task accuracy compared to existing methods. Empirical results show that AC-ODM can reduce training steps by up to 66% while enhancing MMLU accuracy by 27.5% and achieving a 2.23x higher pass rate on HumanEval, all with minimal additional computational cost.
AC-ODM can reduce LLM training steps by 66% while boosting accuracy, reshaping the landscape of sample-efficient pretraining.
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.
