データミキシングエージェント：継続的プリトレーニングのためのドメイン再重み付け学習

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.