用强化学习动态调整训练数据组合,显著提升大模型预训练效率。
AC-ODM: Actor--Critic Online Data Mixing for Sample-Efficient LLM Pretraining
- 通过强化学习构建可参数化的数据混合策略,自动优化训练过程。
- 在Pythia-1B上减少66%训练步数达最优效果,MMLU准确率提升27.5%。
- 支持多种训练场景,适合追求高效预训练的大模型研究者使用。
优化预训练数据构成对大语言模型泛化能力至关重要。尽管动态混合方法能捕捉训练过程的演化特性,但现有方法难以兼顾计算效率、样本效率与结构灵活性。本文提出演员-评论家在线数据混合(AC-ODM),从强化学习视角出发,设计可参数化的策略,理论证明其作为动态线性代理可最大化梯度的建设性干涉。为提升实用性,AC-ODM支持两种模式:(i) 代理模式,将小模型学习的策略迁移至大模型;(ii) 非代理模式,无需先验知识直接端到端训练。实验表明,AC-ODM在多种架构下显著优于基线,在Pythia-1B上实现最优验证困惑度仅需比竞争方法少66%的训练步数,MMLU准确率相对提升27.5%,HumanEval pass@1提升2.23倍,且每步开销仅增加0.4%,内存额外开销仅2%。代码已开源。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。