基于 Actor-Critic 的语言模型预训练在线数据混合
机器学习
2025-06-02 v1 人工智能
摘要
预训练数据的覆盖范围和组成显著影响大型语言模型（LLMs）的泛化能力。为了减少训练的碳足迹和财务成本，已有一些数据混合方法被提出，它们将小型代理模型优化后的领域权重应用于训练更大的模型。然而，这些方法并未随训练动态而演化。现有的在线数据混合（ODM）方法通过应用多臂老虎机算法作为数据采样策略来解决这一局限性。然而，它未考虑域内交互。在本文中，我们开发了一种基于 Actor-Critic 的在线数据混合（AC-ODM）方法，该方法通过辅助 Actor-Critic 网络捕获变化的领域权重，并利用奖励函数考虑域内交互。在构建用于预训练大型目标 LLM 的数据集时，我们直接将以小型代理 LLM 作为环境训练的 Actor 用作采样策略。采样策略的转移不仅能确保动态数据混合的效率，还能加速目标 LLM 预训练的收敛。数值结果表明，调用由 410M 参数代理 LLM 获得的采样策略的 AC-ODM-410M，达到 ODM 最优验证困惑度的速度快 71%，并在零样本 MMLU 基准上将准确率提升了 27.5%，在 HumanEval 基准的 pass@1 上提升了约 2.23 倍。
引用
@article{arxiv.2505.23878,
title = {Actor-Critic based Online Data Mixing For Language Model Pre-Training},
author = {Jing Ma and Chenhao Dang and Mingjie Liao},
journal= {arXiv preprint arXiv:2505.23878},
year = {2025}
}
备注
11 pages, 4 figures