AC-ODM：演员–批评家在线数据混合技术，用于高效处理样本的LLM预训练
AC-ODM: Actor--Critic Online Data Mixing for Sample-Efficient LLM Pretraining
摘要
优化预训练数据的构成对于LLM的泛化能力至关重要。虽然动态混合策略能够捕捉到不断变化的训练过程，从而优于静态策略，但当前的方法在计算效率、样本利用率以及不同模型架构下的结构灵活性方面存在不足。我们提出了Actor–Critic在线数据混合方法（AC-ODM），该方法从强化学习的角度来处理数据混合问题，其参数化策略能够实现梯度的有效叠加，从而提升模型的性能。为了增强其实用性，AC-ODM支持两种操作模式：(i) 针对固定、预先准备好的数据集的代理模式，即将在小型模型中学习的策略迁移到更大的目标模型中；(ii) 直接从头开始进行端到端训练的模式，无需任何先验信息。实验表明，AC-ODM在收敛速度和下游准确率方面均优于其他现有方法。在Pythia-1B上，它仅需较少的训练步数即可达到最佳验证困惑度，相比其他对比方法减少了66%的训练步数；同时，其在MMLU测试中的准确率提升了27.5%，在HumanEval测试中的pass@1得分则提高了2.23倍。此外，每步的计算时间增加仅约为0.4%，内存占用也仅增加了2%。代码可访问于https://github.com/DANG-ai/AC-ODM。
English 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.