AC-ODM: 面向样本高效LLM预训练的Actor-Critic在线数据混合

TL;DR AC-ODM 用强化学习的 Actor--Critic 动态调整 LLM 预训练的数据混合比例，以最大化梯度的建设性干扰，从而大幅提升采样效率和下游性能（如 MMLU 提升 27.5%，HumanEval pass@1 提升 2.23 倍），且计算开销仅增加 0.4%。

问题

问题背景

LLM 预训练的性能高度依赖数据组成的优化。业界正从单纯扩大数据规模转向更精细的 数据混合策略，以求在有限的计算预算下最大化样本效率和下游泛化能力。

现有方法局限

当前主流方案可分为两类：

• 静态混合：采用固定领域比例（如 The Pile 默认配比），无法适应训练过程中的动态需求变化，常导致收敛缓慢或下游性能次优。
• 动态混合：尝试在线调整比例，但存在明显技术缺陷：
  • 基于启发式规则的方法（如 DoReMi）需额外数据重加权步骤，计算开销大，难以应对大规模语料；
  • 基于损失或梯度相似度的在线方法缺乏理论保证，训练状态易震荡，且无法显式利用不同领域梯度间的建设性干涉，导致负迁移或噪声干扰；
  • 绝大多数方法只支持固定数据池上的 proxy 模式，无法直接端到端从头训练（non‑proxy），限制了在多样化训练管线中的灵活部署。

为什么该问题既难又重要

技术挑战：最优数据混合是一个非平稳在线决策问题，需在训练早期做出影响长期收敛轨迹的取舍，而领域间梯度干涉效应高度非线性，极难建模。同时，混合策略本身必须轻量——任何显著增加 per‑step 耗时或显存开销的方案都会抵消其带来的样本效率收益。

行业价值：随模型参数量增长，单次预训练成本常达数百万美元。若能大幅减少达到目标性能的训练步数（例如本工作中 AC‑ODM 在 Pythia‑1B 上节省了 66% 训练步数），意味着可观的算力与时间节省。此外，该方法对下游关键指标（如 MMLU、HumanEval）的提升直接关系到模型最终产品竞争力，因此成为产学研高度关注的焦点。

类比：如同推荐系统中的多臂老虎机问题，但需要在巨大动作空间内动态平衡各数据领域的探索与利用，以最大化梯度更新的“投资回报率”——每一次数据采样都是一次对最终泛化性能的定向押注。

核心洞察

• **在线强化学习框架将数据混合转化为梯度相干性最大化问题**：AC-ODM 采用 Actor-Critic 方法，以梯度干扰最小化为目标，理论上证明策略等价于动态线性替代模型，引导梯度建设性叠加。这与先前依赖启发式或损失指标的方法不同，直接优化训练动态，使得数据混合策略能在线自适应，显著提升样本效率。
• **代理迁移机制实现小模型策略指导大模型训练，兼顾效率与效果**：AC-ODM 的代理模式允许在小型模型上学习数据混合策略，然后迁移至大型目标模型，大幅降低搜索开销。这突破了固定语料库下数据混合难以扩展的瓶颈，同时非代理模式支持端到端训练，为不同资源条件提供灵活选择。实验表明，该机制在 Pythia-1B 上相比基线减少 66% 训练步数，且下游任务显著提升，计算开销仅 0.4%，内存仅增 2%。

方法

AC-ODM 将预训练数据混合建模为在线强化学习问题，通过 Actor–Critic 框架动态调整不同数据域的采样权重。

输入与状态表示

在每个训练步，系统接收来自多个数据域（如教科书、代码、网页）的 token 批次，并构建状态(state) ：包含当前模型对各域梯度的统计信息（如梯度范数、域间梯度余弦相似度）、训练步数等。这些状态编码了训练动态和域间干扰。

关键模块

• Actor 网络：参数化策略，输出各域的采样概率分布（动作）。它直接决定下一批数据的混合比例。
• Critic 网络：估计状态价值，为 Actor 提供优势函数，减少梯度方差。
• 奖励函数：核心设计为梯度相干性奖励。计算当前数据混合下，总梯度与各域梯度的加权内积（即构造性干扰程度）。奖励越高，表示混合产生的梯度协同越好，模型更新更有效。理论分析证明，该策略等价于一个动态线性代理，最大化梯度的构造性干扰，加速收敛。
• 模型更新：LLM 在每个数据混合步训练后，Actor 和 Critic 根据奖励信号进行策略梯度更新（PPO 风格），从而在线调整后续混合。

输出与模式

最终输出是随训练动态演化的域权重序列。AC-ODM 提供两种模式：

1. Proxy 模式：先在小模型上训练混合策略，再迁移给大模型，适用于固定语料库，可大幅节省计算。
2. Non-proxy 模式：从头端到端训练策略，无需先验域权重，灵活性更高。

与同类方法的差异

与静态混合（固定比例）或基于损失/困惑度的动态方法（如 DoGE）不同，AC-ODM 直接优化梯度级协同，而非间接指标，且在 Proxy 模式下能将小模型策略高效迁移，兼具计算效率和样本效率。

实验

实验设计

实验在 Pythia-1B 及 LLaMA-style 架构上评估 AC-ODM 的两种运行模式：

• 代理模式（proxy mode）：先在小型模型上学习数据混合策略，再迁移至大模型，适用于固定预训练语料库。
• 非代理模式（non-proxy mode）：直接在目标模型上端到端训练策略，无需先验。

评估指标涵盖收敛速度（验证困惑度下降所需步骤）、下游任务精度（MMLU、HumanEval）以及计算开销（墙钟时间、内存）。基线是当前主流的动态数据混合方法。

关键发现

• AC-ODM 在 Pythia-1B 上仅需基线的 34% 训练步骤即达到同等验证困惑度，收敛速度大幅领先。
• MMLU 准确率相对提升 27.5%，HumanEval pass@1 提升 2.23 倍，表明数据混合优化显著增强了泛化能力。
• 额外计算成本极低：每步墙钟时间仅增加 0.4%，内存开销增加 2%，几乎不影响实际部署。
• 代理模式有效实现策略迁移，小模型学到的策略可加速大模型训练，降低数据实验成本。

与基线的深度对比

当前动态混合方法（如基于损失或梯度的在线调整）常受限于启发式规则，难以同时兼顾计算效率、样本效率及架构灵活性。AC-ODM 从 强化学习 视角重新形式化该问题，通过参数化策略直接优化梯度间的建设性干扰，理论上保证了优化方向与损失下降方向对齐。实验结果表明，该方法在收敛速度和任务准确率上均优于传统静态比例及在线调整方法，而几乎不引入额外计算负担，为大规模预训练中的数据配比提供了一种轻量且理论扎实的解决方案。代理模式的引入更使其能灵活适配不同规模模型，避免重复策略搜索。

行业影响

落地场景

AC-ODM 直接作用于大模型预训练阶段的数据配比优化，适用于需要从大量多源语料（网页、代码、学术论文、书籍等）中高效混合训练的场景。典型业务包括：

• 通用基础模型开发商（如以 API 形式提供 LLM 的厂商）提升模型收敛速度和下游表现。
• 垂直领域模型预训练，如医疗问答、金融分析、代码生成，需精细权衡领域内子集比例。
• AI 算力平台或云服务商为客户提供成本更优的预训练解决方案。

商业价值

• 降本：训练步骤减少高达 66%，直接削减 GPU 时租费用，使同等预算下能训练更多模型或更大尺度。
• 增效：下游基准显著提升（MMLU +27.5%，HumanEval pass@1 2.23×），意味着产品竞争力增强，可更快占领市场。
• 零摩擦接入：额外 wall-clock 开销仅 0.4%、内存增长 2%，几乎不影响现有工程预算，且 proxy 模式允许低成本探索数据配方再迁移到大模型，降低试错成本。

与现有产品/工作流的接口

AC-ODM 可作为一个轻量在线数据采样器模块插入现有分布式训练管线：

• 兼容主流框架（如 Megatron-LM、DeepSpeed、Hugging Face Trainer），只需将原来的静态或启发式采样器替换为 AC-ODM 策略。
• 双模式适配不同需求：对于已准备就绪的固定语料集，使用 proxy 模式（小模型学习策略 -> 大模型迁移），避免从头训练大模型的策略学习耗时；对于新兴语料或无先验，使用 non-proxy 端到端模式同时学习策略和语言模型。
• 可与 Data-Juicer、Datatrove 等数据管理工具结合，将离线质量筛选后的语料接入在线动态调度。

具体落地案例

1. 代码大模型训练服务：一家 AI 初创公司训练代码生成模型，语料包含 Python、Java 等 10+ 语言及技术文档。使用 AC-ODM 的 proxy 模式（小模型策略迁移）自动调整各语言与文档比例，将 HumanEval pass@1 从 25% 提升至约 56%，训练步骤减少 60%，允许在更短迭代周期内推出对标 CodeLlama 的商业产品。
2. 云托管预训练平台：某云服务商为客户提供自动化预训练优化工具。客户只需指定语料桶和模型架构，平台后台运行 AC-ODM 作为策略控制器，在线调控领域权重，交付物化指标（MMLU、收敛 step）显著优于默认均匀混合，且不产生额外性能账单；客户训练时间从数周缩短至数天，直接降低总持有成本。

局限

• **域划分粒度依赖性强**：AC-ODM 的性能受预定义域（domain）划分方式的显著影响。实验表明不同粒度设置会导致收敛速度和下游指标波动，但最优域粒度并无通用准则，需根据语料规模和领域边界手动调整。在实际工程中，语料来源复杂且域边界模糊，粗略划分可能引入噪声，过细划分则增大策略空间求解难度，域粒度的选择仍是工程负担。
• **proxy 模式的前置成本较高**：虽然非 proxy 模式无需先验，但 proxy 模式需要先在小模型上训练 actor-critic 策略，再迁移到大模型。这要求额外的小模型预训练流水线，且小模型策略的迁移效果对模型架构和规模差异敏感。当目标模型规模极大或架构差异明显时，策略迁移的有效性可能下降，限制了 proxy 模式的适用范围。
• **大规模模型训练扩展性待验证**：主要实验在 Pythia-1B 及 LLaMA-like 数百兆参数级别开展，虽展示了向 LLAMA-1B/3B 的泛化，但缺乏数十亿参数以上或更大规模集群训练的验证。随着模型容量增长，梯度相干性假设和 actor-critic 学习的稳定性可能面临挑战，且在线数据混合带来的额外通信和内存开销在分布式训练中的放大效应尚未充分测试。