AC-ODM: 演员-评论家在线数据混合方法助力大语言模型样本高效预训练

摘要

优化预训练数据组合对于大语言模型的泛化能力至关重要。尽管动态混合策略通过捕捉训练过程中的动态变化优于静态策略，但现有方法无法在计算效率、样本效率以及面向多样化管线的结构灵活性之间取得平衡。为此，我们提出了行为者-评论家在线数据混合方法（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

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.