AC-ODM: Actor–Critic Online Data Mixing for Sample-Efficient LLM 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 \textbf{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 \textbf{proxy mode} for fixed, pre-prepared corpora, where a policy learned on a small model is transferred to a larger target; and (ii) a \textbf{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$\times$ 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 \url{https://github.com/DANG-ai/AC-ODM}.

Large language models learn from enormous mixtures of text, code, scientific writing, books, and other data sources, but choosing how much to train on each source is difficult. A fixed data mixture can waste computation because the model's needs change throughout training. We propose AC-ODM, a method that treats data mixing as a feedback problem: it observes the training state of the model and adjusts the mixture online. The method favors data sources whose learning signals work well together, helping each training step make more useful progress. In experiments on large pretraining corpora, AC-ODM trains models faster and improves performance on language understanding and code generation tasks. This can reduce the compute and energy needed to reach a target model quality, while also helping researchers better understand how different data sources contribute during pretraining.