Data Mixing Agent: Learning to Re-weight Domains for Continual Pre-training

这是一篇关于大语言模型（LLM）持续预训练（Continual Pre-training）中数据混合策略的学术论文总结。

论文标题

Data Mixing Agent: Learning to Re-weight Domains for Continual Pre-training
（数据混合智能体：学习为持续预训练重新加权领域）

1. 研究背景与问题 (Problem)

• 持续预训练的困境：为了将通用大语言模型适配到新的特定领域（如数学推理、代码生成），通常需要进行持续预训练。然而，直接仅在目标领域数据上训练会导致灾难性遗忘（Catastrophic Forgetting），即模型丧失在原始领域（如通用知识）的能力。
• 现有解决方案的局限：常见的解决方法是混合源领域（Source）和目标领域（Target）的数据。现有的数据混合策略（Data Mixing）通常依赖人工设计的启发式规则（基于人类直觉或经验结果）来调整不同领域数据的比例（Re-weighting）。
• 核心痛点：
  1. 人工启发式规则缺乏通用性，难以发现更广泛、更深层的混合规律。
  2. 现有的动态调整方法（如基于损失函数的动态批加载）往往计算开销大或需要针对每个任务重新设计。
  3. 缺乏一种能够自动学习“如何混合数据”的端到端模型。

2. 方法论 (Methodology)

作者提出了 Data Mixing Agent，这是首个基于模型、端到端的领域重加权框架。该方法将领域重加权问题建模为马尔可夫决策过程（MDP），并利用强化学习（RL）来学习通用的数据混合启发式规则。

核心组件：

1. MDP 建模：

   • 状态 (State)：包含之前的数据分布历史以及环境反馈（模型在目标领域的表现）。
   • 动作 (Action)：当前步骤中各个领域（Domain）的数据分配比例（概率分布）。
   • 奖励 (Reward)：基于轻量级评估环境（Evaluation Environment）对模型检查点（Checkpoint）在通用能力和目标能力上的评估反馈。

2. 训练流程 (Training Pipeline)：

   • 轨迹采样 (Trajectory Sampling)：随机采样大量数据混合轨迹。为了覆盖高质量和低质量的混合策略，设计了一种基于归纳偏置（Inductive Biases）的评分算法来筛选轨迹（例如：当前分布不应与上一步偏离过大，应逐渐向目标分布靠拢等）。
   • 代理模型训练 (Proxy Training)：在采样到的每条轨迹上训练小型代理模型（Proxy Models），并在每个重加权步骤收集环境反馈。
   • 强化学习优化 (RL Optimization)：
     • 使用 保守 Q 学习 (Conservative Q-Learning, CQL) 算法进行离线强化学习。CQL 能够防止对分布外（Out-of-Distribution）动作的 Q 值高估，这对于利用包含“坏”轨迹的数据集至关重要。
     • 先进行监督微调（SFT）作为热身，再使用 CQL 进行策略优化。
   • 模型架构：Agent 本身是一个轻量级的 Transformer 解码器（仅 210 万参数），输入为历史轨迹和反馈，输出为下一步的领域分布。

3. 在线应用：

   • 在目标模型的持续预训练过程中，Agent 根据当前的训练状态和历史反馈，实时（On-the-fly）预测下一个步骤的最佳领域数据分布。

3. 关键贡献 (Key Contributions)

1. 首个基于模型的领域重加权方法：提出了 Data Mixing Agent，实现了从“人工规则”到“模型自动学习”的范式转变，能够端到端地学习数据混合策略。
2. 发现并参数化通用启发式规则：证明了存在通用的、与模型和数据无关的启发式规则（Heuristics），可以通过强化学习被 Agent 捕捉。这些规则不仅符合人类直觉（如增加科学类数据提升 MMLU 表现），还能发现人类未察觉的规律。
3. 卓越的泛化能力：
   • 跨领域泛化：在数学推理领域训练的 Agent，无需重新训练即可直接应用于代码生成领域。
   • 跨模型泛化：在不同架构和规模的模型（如 LLaMA 变体、Pythia）上均表现优异。
   • 跨空间泛化：在 2 维（基于数据源）和 52 维（基于 Nvidia 领域分类器）的不同领域空间定义下均有效。
4. 高效性：相比基线方法，Agent 能以更少的源领域数据 Token 预算实现更优的平衡性能，且 Agent 本身的推理开销极低。

4. 实验结果 (Results)

实验主要在数学推理和代码生成两个目标领域进行，对比了静态方法（RegMix）、动态方法（DBL）和纯监督微调的 Agent（DataAgentSFT）。

• 性能提升：
  • 在数学推理任务中，DataAgentRL 在 8 个通用基准和 4 个数学基准上的平均表现比强基线 RegMix 高出 3.02%，比 DBL 高出 3.53%。
  • 成功缓解了灾难性遗忘，在提升目标领域能力的同时，显著保留了通用领域能力（方差更低，性能更稳定）。
• 泛化性验证：
  • 在未见过的目标模型（LLaMA-FWE, LLaMA-Nemo）和未见过的目标领域（代码生成）上，无需重新训练即可取得 SOTA 级别的平衡性能。
  • 直接迁移到代码生成领域，平均性能比 RegMix 高出 1.45%。
• 数据效率：
  • Agent 能够更早地触发停止训练（Early Stopping），在达到同等性能时，比 RegMix 节省了约 2.14B 的源领域 Token 预算。
  • 在源领域数据不可用（使用合成数据估计初始状态）的情况下，Agent 依然有效。
• 可解释性：
  • Agent 学到的策略与人类直觉高度一致（例如：在 MMLU 评估中增加“科学”和“健康”类数据，减少“时尚美容”类数据）。
  • Agent 采用了更精细的三阶段策略：早期优先源数据（热身）-> 中期快速增加目标数据 -> 后期重新引入源数据以稳定性能。

5. 意义与结论 (Significance)

• 自动化数据工程：Data Mixing Agent 将数据混合从“人工试错”转变为“自动化学习”，为持续预训练提供了一种可复现、可扩展的解决方案。
• 降低计算成本：通过优化数据混合比例，减少了不必要的源领域数据消耗，降低了训练 Token 预算，同时避免了为每个新任务重新设计混合策略的巨大成本。
• 通用启发式空间：论文揭示了领域重加权中存在丰富的、可学习的启发式空间，这为未来探索更复杂的数据混合策略（如多模态、多任务）奠定了基础。
• 局限性：目前评估主要集中在数学和代码领域，对于分布差异极大的领域（如从数学到医疗），泛化能力可能下降；且主要应用于持续预训练场景，大规模预训练中的适用性有待进一步探索。

总结：该论文提出了一种利用强化学习自动学习数据混合策略的创新框架，有效解决了持续预训练中的灾难性遗忘问题，并在多个领域和模型上证明了其优越的泛化能力和数据效率。