Agent de Mélange de Données : Apprentissage de la Repondération des Domaines pour le Pré-entraînement Continu

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

July 21, 2025

papers.authors: Kailai Yang, Xiao Liu, Lei Ji, Hao Li, Yeyun Gong, Peng Cheng, Mao Yang

cs.AI

papers.abstract

Le pré-entraînement continu sur des données spécifiques à petite échelle est une méthode efficace pour améliorer les grands modèles de langage dans de nouveaux domaines cibles, mais il risque d'entraîner un oubli catastrophique de leurs capacités originales. Une solution courante consiste à rééquilibrer les mélanges de données d'entraînement provenant des domaines source et cible dans un espace de domaine pour obtenir des performances équilibrées. Les stratégies précédentes de rééquilibrage de domaine reposaient sur une désignation manuelle avec certaines heuristiques basées sur l'intuition humaine ou des résultats empiriques. Dans ce travail, nous démontrons que des heuristiques plus générales peuvent être paramétrisées en proposant Data Mixing Agent, le premier cadre de bout en bout basé sur un modèle qui apprend à rééquilibrer les domaines. L'agent apprend des heuristiques généralisables grâce à l'apprentissage par renforcement sur de grandes quantités de trajectoires de mélange de données avec des retours correspondants d'un environnement d'évaluation. Les expériences de pré-entraînement continu sur le raisonnement mathématique montrent que Data Mixing Agent surpasse des bases de référence solides pour atteindre des performances équilibrées sur les benchmarks des domaines source et cible. De plus, il généralise bien à des domaines source inconnus, des modèles cibles et des espaces de domaine sans nécessiter de réentraînement. Une application directe dans le domaine de la génération de code indique également son adaptabilité à travers différents domaines cibles. Une analyse approfondie met en évidence l'alignement des heuristiques de l'agent avec l'intuition humaine et son efficacité à obtenir des performances de modèle supérieures avec moins de données du domaine source.

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.