Agente de Mistura de Dados: Aprendendo a Re-ponderar Domínios para Pré-treinamento Contínuo

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

July 21, 2025

Autores: Kailai Yang, Xiao Liu, Lei Ji, Hao Li, Yeyun Gong, Peng Cheng, Mao Yang

cs.AI

Resumo

O pré-treinamento contínuo em pequenos conjuntos de dados específicos para tarefas é um método eficaz para melhorar modelos de linguagem de grande escala em novos campos de interesse, mas corre o risco de causar esquecimento catastrófico de suas capacidades originais. Uma solução comum é reequilibrar as misturas de dados de treinamento dos campos de origem e de destino em um espaço de domínio para alcançar desempenho equilibrado. Estratégias anteriores de reequilíbrio de domínio dependem de designações manuais com certas heurísticas baseadas em intuição humana ou resultados empíricos. Neste trabalho, demonstramos que heurísticas mais gerais podem ser parametrizadas ao propor o Agente de Mistura de Dados, o primeiro framework baseado em modelo e de ponta a ponta que aprende a reequilibrar domínios. O agente aprende heurísticas generalizáveis por meio de aprendizado por reforço em grandes quantidades de trajetórias de mistura de dados com feedback correspondente de um ambiente de avaliação. Experimentos de pré-treinamento contínuo em raciocínio matemático mostram que o Agente de Mistura de Dados supera baselines fortes ao alcançar desempenho equilibrado em benchmarks de campos de origem e de destino. Além disso, ele generaliza bem em campos de origem não vistos, modelos de destino e espaços de domínio sem necessidade de retreinamento. A aplicação direta no campo de geração de código também indica sua adaptabilidade entre domínios de destino. Análises adicionais mostram que as heurísticas do agente estão bem alinhadas com intuições humanas e sua eficiência em alcançar desempenho superior do modelo com menos dados do campo de origem.

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.