Агент смешивания данных: обучение перевзвешиванию доменов для непрерывного предварительного обучения

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

July 21, 2025

Авторы: Kailai Yang, Xiao Liu, Lei Ji, Hao Li, Yeyun Gong, Peng Cheng, Mao Yang

cs.AI

Аннотация

Непрерывное предварительное обучение на небольших наборах данных, специфичных для конкретной задачи, является эффективным методом улучшения крупных языковых моделей в новых целевых областях, однако оно сопряжено с риском катастрофического забывания их исходных возможностей. Распространённое решение заключается в перевзвешивании смесей обучающих данных из исходной и целевой областей в пространстве доменов для достижения сбалансированной производительности. Предыдущие стратегии перевзвешивания доменов основывались на ручном назначении с использованием определённых эвристик, основанных на человеческой интуиции или эмпирических результатах. В данной работе мы доказываем, что более общие эвристики могут быть параметризованы, предлагая Data Mixing Agent — первую модельную, сквозную структуру, которая обучается перевзвешиванию доменов. Агент изучает обобщаемые эвристики с помощью обучения с подкреплением на большом количестве траекторий смешивания данных с соответствующими откликами из оценочной среды. Эксперименты по непрерывному предварительному обучению в области математических рассуждений показывают, что Data Mixing Agent превосходит сильные базовые подходы в достижении сбалансированной производительности на тестах как исходной, так и целевой областей. Более того, он хорошо обобщается на неизвестные исходные области, целевые модели и пространства доменов без необходимости переобучения. Прямое применение в области генерации кода также свидетельствует о его адаптивности к различным целевым доменам. Дополнительный анализ демонстрирует, что эвристики агента хорошо согласуются с человеческой интуицией, а также их эффективность в достижении превосходной производительности модели с использованием меньшего количества данных из исходной области.

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.