用AI自动学习数据混合权重,避免大模型遗忘旧知识。
Data Mixing Agent: Learning to Re-weight Domains for Continual Pre-training
- 提出可学习的领域重加权框架,通过强化学习自动优化数据混合比例。
- 在数学推理任务中超越基线,在源域和目标域间实现平衡性能。
- 无需重新训练即可跨领域、跨模型泛化,适合快速适配新任务。
在小规模特定任务数据上持续预训练是提升大语言模型在新领域表现的有效方法,但易导致原始能力的灾难性遗忘。常见解决方案是在领域空间中对源域和目标域数据混合进行重加权以实现性能平衡。以往策略依赖人工设计的启发式规则,基于人类直觉或经验结果。本文提出首个模型驱动的端到端框架——数据混合代理(Data Mixing Agent),通过在大量数据混合轨迹上进行强化学习,自动学习可泛化的重加权策略。该代理在数学推理任务的持续预训练中表现优异,优于多个强基线,在源域与目标域基准测试间实现良好平衡。此外,其在未见过的源域、目标模型和领域空间中均表现出良好泛化能力,无需重新训练。直接应用于代码生成领域也验证了其跨目标领域的适应性。进一步分析显示,该代理学到的策略与人类直觉高度一致,且仅需较少源域数据即可实现更优模型性能。
原文摘要 · Abstract (English)
Continual pre-training on small-scale task-specific data is an effective method for improving large language models in new target fields, yet it risks catastrophic forgetting of their original capabilities. A common solution is to re-weight training data mixtures from source and target fields on a domain space to achieve balanced performance. Previous domain reweighting strategies rely on manual designation with certain heuristics based on human intuition or empirical results. In this work, we prove that more general heuristics can be parameterized by proposing Data Mixing Agent, the first model-based, end-to-end framework that learns to re-weight domains. The agent learns generalizable heuristics through reinforcement learning on large quantities of data mixing trajectories with corresponding feedback from an evaluation environment. Experiments in continual pre-training on math reasoning show that Data Mixing Agent outperforms strong baselines in achieving balanced performance across source and target field benchmarks. Furthermore, it generalizes well across unseen source fields, target models, and domain spaces without retraining. Direct application to the code generation field also indicates its adaptability across target domains. Further analysis showcases the agents' well-aligned heuristics with human intuitions and their efficiency in achieving superior model performance with less source-field data.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。