Data Mixing Agent: Lernen zur Neugewichtung von Domänen für kontinuierliches Vortraining

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

Kontinuierliches Vortraining auf kleinen, aufgabenspezifischen Datensätzen ist eine effektive Methode, um große Sprachmodelle in neuen Zielbereichen zu verbessern, birgt jedoch das Risiko eines katastrophalen Vergessens ihrer ursprünglichen Fähigkeiten. Eine gängige Lösung besteht darin, die Mischungen von Trainingsdaten aus Quell- und Zielbereichen in einem Domänenraum neu zu gewichten, um eine ausgewogene Leistung zu erzielen. Bisherige Strategien zur Domänen-Neugewichtung stützen sich auf manuelle Festlegungen mit bestimmten Heuristiken, die auf menschlicher Intuition oder empirischen Ergebnissen basieren. In dieser Arbeit zeigen wir, dass allgemeinere Heuristiken parametrisiert werden können, indem wir den Data Mixing Agent vorschlagen, das erste modellbasierte, end-to-end Framework, das lernt, Domänen neu zu gewichten. Der Agent lernt generalisierbare Heuristiken durch Reinforcement Learning auf großen Mengen von Datenmischungspfaden mit entsprechendem Feedback aus einer Evaluationsumgebung. Experimente zum kontinuierlichen Vortraining im Bereich des mathematischen Denkens zeigen, dass der Data Mixing Agent starke Baselines in der Erzielung einer ausgewogenen Leistung über Quell- und Zielbereichs-Benchmarks übertrifft. Darüber hinaus generalisiert er gut über ungesehene Quellbereiche, Zielmodelle und Domänenräume hinweg, ohne erneut trainiert zu werden. Die direkte Anwendung im Bereich der Code-Generierung zeigt ebenfalls seine Anpassungsfähigkeit über verschiedene Zielbereiche hinweg. Weitere Analysen zeigen, dass die Heuristiken des Agents gut mit menschlichen Intuitionen übereinstimmen und dass sie effizient eine überlegene Modellleistung mit weniger Daten aus dem Quellbereich erreichen.

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.