DomainPilot: Domain-Level Loss-Guided Two-Stage Data Mixture Optimization for Efficient Language Model Fine-Tuning

2026-07-23

AI Summary

This work addresses the inefficiency of existing dynamic data scheduling methods in large model fine-tuning, which suffer from high computational overhead, I/O bottlenecks, and noisy sample-level loss signals that hinder optimal data mixing. To overcome these limitations, we propose DomainPilot, a novel framework featuring a domain-level loss monitoring mechanism that captures learning dynamics in real time without interrupting training. Our approach operates in two stages: first, it derives a mixing prior by fitting convergence curves of individual domains using scaling laws; second, it refines the mixture ratios by modeling cross-domain interaction effects through mixing laws. DomainPilot can be integrated into mainstream training frameworks with only ~30 lines of code. In SFT experiments on Qwen3-1.7B, it achieves consistent gains—2.0%, 1.8%, 3.8%, and 3.6% improvements on MMLU-Redux, AIME24, LiveCodeBench v5, and BFCL v3, respectively—without additional data or computational cost.

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.

Problem

• data mixture optimization
• large language models
• supervised fine-tuning
• domain-level loss
• training data scheduling

Innovation

• Domain-Level Optimization
• Loss-Guided Mixture
• Two-Stage Optimization
• Scaling Law
• Mixing Law