Data Mixing Agent: Learning to Re-weight Domains for Continual
  Pre-training
Authors: name: Kailai Yang, name: Xiao Liu, name: Lei Ji, name: Hao Li, name: Yeyun Gong, name: Peng Cheng, name: Mao Yang

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.