Agente di Miscelazione dei Dati: Apprendimento per Ricalibrare i Domini nel Pre-Addestramento Continuo

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

July 21, 2025

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

cs.AI

Abstract

Il pre-addestramento continuo su dati specifici per il compito di piccola scala è un metodo efficace per migliorare i grandi modelli linguistici in nuovi campi target, ma comporta il rischio di dimenticanza catastrofica delle loro capacità originali. Una soluzione comune consiste nel ri-pesare le miscele di dati di addestramento provenienti dai campi sorgente e target in uno spazio di dominio per ottenere prestazioni bilanciate. Le precedenti strategie di ri-pesatura del dominio si basano su designazioni manuali con alcune euristiche basate sull'intuizione umana o su risultati empirici. In questo lavoro, dimostriamo che euristiche più generali possono essere parametrizzate proponendo Data Mixing Agent, il primo framework basato su modelli e end-to-end che impara a ri-pesare i domini. L'agente apprende euristiche generalizzabili attraverso l'apprendimento per rinforzo su grandi quantità di traiettorie di miscelazione dei dati con feedback corrispondenti da un ambiente di valutazione. Esperimenti di pre-addestramento continuo sul ragionamento matematico mostrano che Data Mixing Agent supera i forti baseline nel raggiungere prestazioni bilanciate attraverso i benchmark dei campi sorgente e target. Inoltre, generalizza bene su campi sorgente non visti, modelli target e spazi di dominio senza bisogno di ri-addestramento. L'applicazione diretta al campo della generazione di codice indica anche la sua adattabilità attraverso diversi domini target. Un'ulteriore analisi mostra che le euristiche dell'agente sono ben allineate con le intuizioni umane e la loro efficienza nel raggiungere prestazioni superiori del modello con meno dati del campo sorgente.

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.