Data Mixing Agent: Het leren herwegen van domeinen voor voortgezette voorafgaande training

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

July 21, 2025

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

cs.AI

Samenvatting

Voortdurende voorafgaande training op kleinschalige, taakspecifieke gegevens is een effectieve methode om grote taalmodellen te verbeteren in nieuwe doelgebieden, maar het brengt het risico van catastrofaal verlies van hun oorspronkelijke capaciteiten met zich mee. Een veelgebruikte oplossing is het herverdelen van trainingsdatamengsels uit bron- en doelgebieden in een domeinruimte om een gebalanceerde prestaties te bereiken. Eerdere strategieën voor domeinherverdeling zijn gebaseerd op handmatige aanwijzingen met bepaalde heuristieken die voortkomen uit menselijke intuïtie of empirische resultaten. In dit werk tonen we aan dat meer algemene heuristieken geparametriseerd kunnen worden door de Data Mixing Agent voor te stellen, het eerste modelgebaseerde, end-to-end framework dat leert om domeinen te herverdelen. De agent leert generaliseerbare heuristieken door middel van reinforcement learning op grote hoeveelheden datamengtrajecten met bijbehorende feedback van een evaluatieomgeving. Experimenten met voortdurende voorafgaande training op wiskundig redeneren laten zien dat de Data Mixing Agent sterke baseline-methoden overtreft in het bereiken van gebalanceerde prestaties over bron- en doelgebiedbenchmarks. Bovendien generaliseert het goed over onbekende brongebieden, doelmodellen en domeinruimten zonder hertraining. Directe toepassing in het veld van codegeneratie geeft ook aan dat het aanpasbaar is over verschillende doeldomeinen. Verdere analyse toont aan dat de heuristieken van de agenten goed aansluiten bij menselijke intuïties en dat ze efficiënt zijn in het bereiken van superieure modelprestaties met minder gegevens uit het brongebied.

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.