Agente de Mezcla de Datos: Aprendiendo a Re-ponderar Dominios para el Pre-entrenamiento Continuo

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

Resumen

El preentrenamiento continuo en datos específicos de tareas a pequeña escala es un método efectivo para mejorar los modelos de lenguaje grandes en nuevos campos objetivo, aunque conlleva el riesgo de olvido catastrófico de sus capacidades originales. Una solución común es re-ponderar las mezclas de datos de entrenamiento provenientes de los campos fuente y objetivo en un espacio de dominio para lograr un rendimiento equilibrado. Las estrategias previas de re-ponderación de dominios dependen de designaciones manuales basadas en ciertas heurísticas derivadas de la intuición humana o resultados empíricos. En este trabajo, demostramos que heurísticas más generales pueden ser parametrizadas al proponer el Agente de Mezcla de Datos, el primer marco basado en modelos y de extremo a extremo que aprende a re-ponderar dominios. El agente aprende heurísticas generalizables mediante aprendizaje por refuerzo en grandes cantidades de trayectorias de mezcla de datos con retroalimentación correspondiente de un entorno de evaluación. Los experimentos en preentrenamiento continuo para razonamiento matemático muestran que el Agente de Mezcla de Datos supera a líneas base sólidas al lograr un rendimiento equilibrado en los puntos de referencia de los campos fuente y objetivo. Además, generaliza bien en campos fuente no vistos, modelos objetivo y espacios de dominio sin necesidad de reentrenamiento. La aplicación directa en el campo de generación de código también indica su adaptabilidad a través de dominios objetivo. Un análisis adicional muestra que las heurísticas del agente están bien alineadas con las intuiciones humanas y su eficiencia para alcanzar un rendimiento superior del modelo con menos datos del campo fuente.

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.