DomainPilot: Domain-Level Loss-Guided Two-Stage Data Mixture Optimization for Efficient Language Model Fine-Tuning
DomainPilot employs token-level domain loss monitoring and a two-stage optimization to improve language model fine-tuning by +1.8% to +3.8%.
Key Findings
Methodology
DomainPilot introduces token-level domain loss tracking during training, capturing per-domain learning dynamics without disrupting data pipelines. The first stage fits scaling laws to domain loss trajectories, deriving priors for mixture reallocation. The second stage models cross-domain interactions via controlled sweep experiments to refine proportions. The architecture employs a patch-based design, integrating domain-aware loss computation into existing frameworks like MindSpeed and Megatron-LM with minimal code (~30 lines). Experiments on Qwen3-1.7B show that this approach improves multiple benchmarks (+1.8%-3.8%) without increasing data volume or training cost.
Key Results
- Post-optimization, the model improves +2% on MMLU-Redux, +1.8% on AIME24, +3.8% on LiveCodeBench v5, and +3.6% on BFCL v3, verified across multiple tasks.
- Domain loss trajectories identify underperforming domains such as swe and termagent, whose resource allocation is increased, leading to significant performance gains.
- Results demonstrate that low-cost domain-level signals effectively guide data mixture adjustments, outperforming traditional importance sampling or clustering methods.
Significance
This work addresses the high cost and inflexibility of existing data scheduling methods in large-scale training. By leveraging domain-level loss signals, it enables dynamic, low-overhead adjustment of data proportions, significantly improving multi-task performance. The approach bridges the gap between training feedback and data curation, offering a scalable solution for industry-scale models, reducing manual effort, and enhancing model generalization across diverse domains.
Technical Contribution
The core innovation lies in token-level domain loss monitoring integrated via a patch architecture, combined with a two-stage optimization process guided by scaling and mixing laws. This framework allows real-time, low-cost adjustment of data mixture ratios without altering core training code. The method provides theoretical insights into learning dynamics and practical tools for efficient data management, extending the capabilities of existing training pipelines.
Novelty
This is the first work to utilize token-level domain loss signals for dynamic mixture optimization, coupling it with scaling law fitting and interaction modeling. Unlike prior methods relying on importance scores or clustering, it offers a lightweight, framework-agnostic, and highly effective approach, filling a critical gap in industrial large-scale training workflows.
Limitations
- Validation is currently limited to the 1.7B parameter scale and fine-tuning stage; applicability to larger models and pretraining remains untested. The second-stage interaction modeling is still under development, requiring further validation.
- The patch-based integration may face compatibility issues with certain frameworks like DeepSpeed, and further generalization is needed for broader adoption.
- The approach assumes stable domain labels and loss signals; noisy or ambiguous data may affect the accuracy of the loss-based adjustments.
Future Work
Future efforts include scaling experiments to larger models (e.g., 4B, 8B), extending the framework to pretraining, and developing automated data cleaning pipelines based on domain loss signals. Additionally, integrating sample-level quality scores could further enhance resource allocation precision. Cross-framework compatibility and robustness under diverse training environments will also be prioritized.
AI Executive Summary
Large language models (LLMs) have revolutionized NLP, yet their training efficiency is hampered by costly data management. Traditional data scheduling methods—such as importance sampling, clustering, or auxiliary model training—face significant scalability and infrastructure challenges. These approaches often require extensive computation, complex pipeline modifications, or lack real-time feedback mechanisms, limiting their industrial applicability.
In response, this paper introduces DomainPilot, a novel framework that leverages token-level domain loss monitoring to guide data mixture optimization during fine-tuning. The core idea is to track how each data domain learns over time, using a lightweight patch integrated into existing training pipelines. The first stage fits scaling laws to these loss trajectories, estimating each domain’s difficulty, convergence speed, and residual error. Based on these parameters, it computes a coarse reallocation score, adjusting domain proportions accordingly.
The second stage models cross-domain interactions through controlled sweep experiments, refining the mixture to maximize overall performance. This process is guided by a mixing law that captures pairwise domain effects, enabling precise resource reallocation. The architecture’s patch design ensures minimal disruption, requiring only about 30 lines of code for integration.
Empirical results on the Qwen3-1.7B model demonstrate that the optimized mixture improves performance across multiple benchmarks—such as LiveCodeBench v5 (+3.8%) and BFCL v3 (+3.6%)—without increasing data volume or training cost. These findings validate the effectiveness of domain-level loss signals as a lightweight alternative to expensive data selection or auxiliary models.
Overall, DomainPilot offers a scalable, low-cost solution for dynamic data management in large-scale language model training. It bridges the gap between training feedback and data curation, paving the way for more autonomous, efficient, and adaptable NLP systems. Future work will extend this framework to larger models, pretraining stages, and broader frameworks, further enhancing its industrial relevance and impact.
Deep Analysis
Background
近年来，随着Transformer架构的普及，大规模语言模型（如GPT、BERT、T5等）在自然语言处理领域取得了巨大突破。代表性工作包括OpenAI的GPT系列、Google的T5、Meta的Llama等，推动模型参数规模不断扩大，模型能力显著提升。然而，训练这些模型面临数据质量、数据多样性和调度策略的巨大挑战。传统方法多依赖静态数据混合或手工调优，效率低、成本高。近年来，动态数据调度技术逐步兴起，试图通过样本筛选、比例调整等手段优化训练效果，但存在计算成本高、框架兼容性差等问题。工业界迫切需要一种低成本、易集成、效果显著的动态调度方案，以满足大规模模型的实际需求。
Core Problem
在大规模预训练和微调过程中，数据选择和比例调节成为瓶颈。样本级损失信号难以区分噪声、难度和新颖性，导致重加权策略效果有限。现有的动态调度方法（如DoReMi、CLIMB、ScaleBiO）多依赖训练辅助模型或复杂的聚类算法，计算成本高昂且难以在生产环境中部署。此外，缺乏实时反馈机制，导致数据配置难以动态适应模型学习状态，影响训练效率和模型性能。如何在保证低成本的同时，动态调整多域数据比例，成为亟待解决的问题。
Innovation
本研究提出基于令牌级域损失监控的两阶段数据混合优化框架。第一阶段利用缩放定律拟合每个域的损失轨迹，提取初始损失、收敛速度和极限损失参数，从而估算每个域的学习难度和收敛状态。基于此，计算粗略的比例调整得分，指导数据比例的初步调节。第二阶段通过控制扫描实验，建模不同域之间的交互影响，微调比例以最大化整体性能。架构采用补丁式设计，将域损失计算和比例调整封装为独立模块，集成到现有训练框架（如MindSpeed、Megatron-LM），只需30行代码即可实现。这一创新结合了理论和工程优势，显著降低成本，提高效率。
Methodology
- �� 在数据预处理阶段，利用knapsack packing为每个令牌打上域标签，形成_token-level domain标记。
- �� 在训练中，模型前向传播后，将每个域的损失进行聚合，计算出每个域的平均损失和样本数。
- �� 拟合每个域的损失轨迹，采用幂律或指数函数，提取初始损失、收敛速度和极限损失参数。
- �� 根据拟合参数，计算每个域的调整得分，结合学习难度和收敛速度，得到粗略比例调整方案。
- �� 通过控制扫描实验，调节域比例，拟合交互模型，捕捉不同域间的影响关系。
- �� 利用交互模型，优化比例配置，提升整体训练效果。
- �� 采用补丁架构，将域损失监控和比例调整模块封装为独立补丁，集成到训练框架中，保证低成本和高兼容性。
Experiments
在Qwen3-1.7B模型上进行微调，数据涵盖8个域（如chat、math、safety等），原始比例由手工调优确定。训练中启用域损失监控补丁，实时记录每个域的损失轨迹。利用缩放定律拟合每个域的学习曲线，估算参数并计算调整得分。根据得分调整比例，进行二次扫描实验验证交互模型。最后在多个任务（如MMLU-Redux、AIME24、LiveCodeBench、BFCL）上评估性能，比较优化前后模型表现。实验还包括不同调度策略的对比，验证方法的低成本和高效性。
Results
优化比例后，模型在LiveCodeBench v5中提升3.8%，BFCL v3中提升3.6%，MMLU-Redux提升2%，AIME24提升1.8%。特别是软件工程和工具调用域（swe、termagent）表现最优，说明低效域得到了有效资源重配。微调结果显示，未增加数据或计算成本，性能显著提升，验证了域级损失信号的指导作用。与传统样本重加权和自动聚类方法相比，该方法具有低成本、效果显著的优势，适合工业大规模训练。
Applications
该方法适用于多域、多任务大模型微调，尤其在工业环境中，能自动调节数据比例，提升模型性能。可应用于智能客服、代码生成、科学计算等场景，减少人工调优成本。未来结合样本质量评分，建立自动数据清洗闭环，推动大模型训练的智能化和自动化。
Limitations & Outlook
目前验证仅在1.7B参数模型和微调阶段，尚未在大规模预训练中验证效果。第二阶段的跨域交互建模仍在开发中，实际效果待验证。架构依赖补丁式集成，可能存在兼容性问题。未来需扩展到更大模型和预训练场景，提升通用性和自动化水平。
Plain Language Accessible to non-experts
想象你在准备一场大型派对，里面有不同的区域：游戏区、音乐区、食物区。每个区域的受欢迎程度不同，策划者需要根据实时反馈调整资源分配。传统方法可能平均分配时间和资源，但效果不佳。DomainPilot就像一个聪明的助手，观察每个区域的反馈（比如满意度），根据趋势调整资源比例。它会识别哪些区域还需要帮助，哪些可以减少，从而让整个派对更成功。这种方法比盲目平均投入更聪明，也更省钱。它用一种科学的方式，动态优化资源分配，让每个区域都能发挥最佳效果，最终让派对变得更加精彩。
ELI14 Explained like you're 14
想象你在学校组织一个大型活动，有很多不同的项目，比如运动会、文艺表演、科学展览。每个项目的准备时间和难度都不一样。以前，你可能会平均分配时间和资源，但这样不一定效果最好。有了DomainPilot，就像你有个聪明的助手，能观察每个项目的进展，发现哪些项目还需要更多帮助，哪些可以少投入。助手会根据每个项目的表现，建议你调整资源，比如多帮运动会一点，少花时间在文艺表演上。这样，整个活动就能更顺利，效果也更好。它用一种聪明的方法，帮你在有限的资源下，让每个项目都发挥出最好的水平。
Abstract
The training efficacy of large language models (LLMs) is fundamentally constrained by the quality and composition of training data. Existing dynamic data scheduling methods face critical limitations in industrial-scale pretraining and supervised fine-tuning (SFT): data selection incurs prohibitive O(N) costs on terabyte-scale corpora, mixture optimization schemes introduce severe I/O bottlenecks or require training auxiliary reference models, and sample-level reweighting strategies rely on loss signals that conflate noise, difficulty, and novelty. We present DomainPilot, a domain-level loss-guided two-stage data mixture optimization framework. DomainPilot introduces token-level domain loss monitoring to capture per-domain learning dynamics during training without halting the data pipeline. Building on these signals, we propose a Scaling Law guided coarse optimization stage that fits domain-specific convergence curves and derives a principled prior for mixture adjustment. A subsequent Mixing Law guided fine optimization stage refines the mixture by modeling cross-domain interaction effects through controlled sweep experiments. The entire mechanism is realized via a patch-based architecture that injects domain-aware loss computation into existing training frameworks (e.g., MindSpeed/Megatron-LM) with only ~30 lines of framework-specific adapter code. We validate DomainPilot on the Qwen3-1.7B model during SFT. Compared to the original data mixture, our optimized mixture achieves improvements of +2% on MMLU-Redux, +1.8% on AIME24, +3.8% on LiveCodeBench v5, and +3.6% on BFCL v3, without increasing total data volume or training cost. These results demonstrate that domain-level training signals provide an effective, lightweight alternative to expensive data selection or auxiliary model training for mixture optimization.