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

# DomainPilot: Domain-Level Loss-Guided Two-Stage Data Mixture Optimization for Efficient Language Model Fine-Tuning
Source: [https://arxiv.org/html/2607.22769](https://arxiv.org/html/2607.22769)
###### 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 prohibitiveO​\(N\)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 presentDomainPilot, a domain\-level loss\-guided two\-stage data mixture optimization framework\. DomainPilot introducestoken\-level domain loss monitoringto capture per\-domain learning dynamics during training without halting the data pipeline\. Building on these signals, we propose aScaling Law guided coarse optimizationstage that fits domain\-specific convergence curves and derives a principled prior for mixture adjustment\. A subsequentMixing Law guided fine optimizationstage refines the mixture by modeling cross\-domain interaction effects through controlled sweep experiments\. The entire mechanism is realized via apatch\-based architecturethat injects domain\-aware loss computation into existing training frameworks \(e\.g\., MindSpeed/Megatron\-LM\) with only∼\\sim30 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%\+2\\%on MMLU\-Redux,\+1\.8%\+1\.8\\%on AIME24,\+3\.8%\+3\.8\\%on LiveCodeBench v5, and\+3\.6%\+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\.

## 1Introduction

The performance of large language models \(LLMs\) is increasingly determined not by architectural innovations alone, but by the curation, composition, and weighting of training data\[[3](https://arxiv.org/html/2607.22769#bib.bib1),[2](https://arxiv.org/html/2607.22769#bib.bib2)\]\. In the data\-centric AI paradigm, even modest adjustments to data mixture ratios can yield improvements comparable to doubling model parameters\[[10](https://arxiv.org/html/2607.22769#bib.bib3)\]\. Despite this, industrial\-scale training pipelines still rely heavily on manual heuristics for multi\-domain data blending, leaving substantial gains unrealized\.

Existing approaches to dynamic data scheduling fall into three broad categories, each with fundamental limitations in the pretraining or large\-scale SFT regime:

Data selectionmethods such as LESS\[[8](https://arxiv.org/html/2607.22769#bib.bib4)\], DSIR\[[9](https://arxiv.org/html/2607.22769#bib.bib5)\], and Quad\[[7](https://arxiv.org/html/2607.22769#bib.bib6)\]identify high\-value subsets by computing per\-sample gradients or importance scores\. While effective for fine\-tuning, these techniques require anO​\(N\)O\(N\)forward pass over the entire corpus—a cost approaching one full training epoch—rendering them prohibitive for terabyte\-scale pretraining\.

Data mixture optimizationmethods such as DoReMi\[[10](https://arxiv.org/html/2607.22769#bib.bib3)\], CLIMB\[[1](https://arxiv.org/html/2607.22769#bib.bib7)\], and ScaleBiO\[[5](https://arxiv.org/html/2607.22769#bib.bib8)\]adjust domain proportions based on model feedback\. DoReMi trains an auxiliary reference model \(3×\\timescompute overhead\) and has only been validated at 30B tokens, two orders of magnitude below production pretraining\. CLIMB relies on clustering to discover domains automatically, an unnecessary step when industrial data teams already maintain explicit domain taxonomies\. ScaleBiO’s intra\-batch dynamic rebalancing introduces severe I/O bottlenecks by converting sequential reads into random accesses, invalidating prefetch caches and degrading GPU utilization\.

Sample reweightingmethods, exemplified by DataFlex\[[11](https://arxiv.org/html/2607.22769#bib.bib9)\]and RHO\-1\[[4](https://arxiv.org/html/2607.22769#bib.bib10)\], modify per\-sample loss contributions during training\. This paradigm avoids extra data traversal and is therefore the only lightweight option viable at scale\. However, existing weighting strategies treatsample\-level lossas a proxy for data quality, overlooking the fact that high loss can indicate noise, difficult domain content, novel valuable information, or short\-text statistical instability—phenomena that are indistinguishable from a single scalar\.

Compounding these algorithmic limitations is aframework portabilitybarrier\. DataFlex is tightly coupled to the HuggingFace Trainer ecosystem \(modifying\_inner\_training\_loopandcompute\_loss\), whereas industrial pretraining predominantly runs on Megatron\-LM\-derived frameworks such as MindSpeed, which possess entirely independent training loops, data loaders, and loss computation paths\. Select and Mix modes—which require rebuilding the DataLoader—are infeasible under Megatron’s memory\-mapped binary data pipeline\. Only Weight\-like interventions, which merely alter loss computation, are practically migrable\.

### 1\.1Core Motivations

These observations motivate five core research questions that guide our work:

1. 1\.Absence of data\-driven mixture optimization\.Multi\-domain training mixtures are typically set by manual tuning and remain static throughout training, ignoring the fact that different domains exhibit disparate learning dynamics\.
2. 2\.Inadequacy of sample\-level loss as a quality signal\.A scalar loss cannot disentangle noise, difficulty, novelty, and instability; weighting decisions based solely on this signal risk amplifying corrupted data\.
3. 3\.Incompatibility of existing schedulers with industrial\-scale training\.Online selection, reference\-model training, and dynamic batch rebalancing all incur costs or infrastructure changes that are unacceptable in production pretraining\.
4. 4\.Unexplained epoch\-boundary loss drops in SFT\.The “staircase” loss reduction observed at epoch boundaries in repeated SFT training lacks systematic explanation, blurring the distinction between generalization and memorization\.
5. 5\.Missing bridge between training feedback and data cleaning\.Offline rule\-based filtering lacks a mechanism to leverage live training signals for identifying and removing low\-quality source data\.

### 1\.2Contributions

We introduceDomainPilot, a domain\-level loss\-guided framework that addresses the above limitations through the following contributions:

1. 1\.Token\-level domain loss monitoring\.We implement a non\-intrusive patch that tags each token with its domain identifier during preprocessing and aggregates per\-domain losses during forward propagation, synchronized across data\-parallel ranks with negligible overhead\.
2. 2\.Two\-stage mixture optimization pipeline\.Stage 1 fits domain\-specificScaling Lawsto extract learning\-dynamic priors \(convergence loss, convergence speed, initial amplitude\) and computes a coarse reallocation\. Stage 2 employsMixing Lawsweep experiments centered on the Stage 1 output to model cross\-domain interactions and refine proportions within a±20%\\pm 20\\%local neighborhood\.
3. 3\.Patch\-based architecture\.The framework adopts a two\-layer design—a framework\-agnostic algorithm layer \(pure PyTorch\) and a thin framework\-specific adapter layer \(∼\\sim30 lines per framework\)\. This enables deployment on MindSpeed/Megatron\-LM without modifying the underlying training codebase\.
4. 4\.Empirical validation on Qwen3\-1\.7B\.Our optimized mixture improves over the original by up to\+3\.8%\+3\.8\\%on LiveCodeBench v5 and\+3\.6%\+3\.6\\%on BFCL v3, with consistent gains on MMLU\-Redux \(\+2%\+2\\%\) and AIME24 \(\+1\.8%\+1\.8\\%\), at no additional data or compute cost\.

## 2Related Work

We organize related work along the three axes introduced in Section[1](https://arxiv.org/html/2607.22769#S1): data selection, mixture optimization, and sample reweighting\. Table[1](https://arxiv.org/html/2607.22769#S2.T1)provides a high\-level comparison\.

Table 1:Comparison of dynamic data scheduling paradigms\. “Pretrain viable” indicates suitability for TB\-scale pretraining or large\-scale SFT\.### 2\.1Data Selection

Data selection aims to identify the most valuable training subset without altering the mixture ratios of the retained data\. LESS\[[8](https://arxiv.org/html/2607.22769#bib.bib4)\]uses low\-rank gradient similarity to select samples most influential for target tasks\. DSIR\[[9](https://arxiv.org/html/2607.22769#bib.bib5)\]formulates selection as importance resampling toward a target distribution\. Quad\[[7](https://arxiv.org/html/2607.22769#bib.bib6)\]jointly optimizes for quality and diversity\. While theoretically appealing, all three require evaluating every candidate sample—a cost that scales linearly with corpus size and becomes impractical whenNNexceeds10910^\{9\}tokens\.

### 2\.2Data Mixture Optimization

Mixture optimization adjusts theproportionsacross domain\-level data sources rather than selecting individual samples\. DoReMi\[[10](https://arxiv.org/html/2607.22769#bib.bib3)\]trains a small proxy model alongside a reference model and up\-weights domains where the proxy exhibits higher excess loss\. The approach is principled but requires training the reference model from scratch \(3×\\timestotal compute\) and has only been demonstrated at 30B tokens, far below production pretraining scales\. CLIMB\[[1](https://arxiv.org/html/2607.22769#bib.bib7)\]discovers domains via clustering and iteratively searches for optimal proportions; however, when domain labels are already explicitly managed by data engineering teams, automatic clustering provides limited additional value\. ScaleBiO\[[5](https://arxiv.org/html/2607.22769#bib.bib8)\]proposes intra\-batch dynamic rebalancing, but changing per\-batch sampling distributions breaks sequential\-read assumptions in memory\-mapped data pipelines, causing prefetch cache misses and GPU idle time that can negate model\-quality gains\.

### 2\.3Sample Reweighting

Sample reweighting modifies the contribution of each training example to the global loss without changing which samples are seen or in what proportion\. Dynamic Loss Reweighting\[[11](https://arxiv.org/html/2607.22769#bib.bib9)\]systematically studies linear, quadratic, and extreme\-value weighting strategies as a function of per\-sample loss\. RHO\-1\[[4](https://arxiv.org/html/2607.22769#bib.bib10)\]extends this to token\-granular selective loss\. Because reweighting only touches the loss computation, it incurs virtually no I/O or synchronization overhead, making it the sole paradigm feasible for trillion\-token pretraining\.

The critical limitation of existing reweighting methods is their reliance onsample\-level lossas the sole signal\. As noted in Section[1](https://arxiv.org/html/2607.22769#S1), sample loss conflates at least four distinct phenomena \(novelty, noise, domain difficulty, and short\-text instability\)\. DataFlex partially addresses this through its Warmup strategy—disabling weighting until the model has partially adapted—but does not resolve the fundamental ambiguity of the signal\.

### 2\.4Framework Portability

DataFlex is implemented within LLaMA\-Factory\[[12](https://arxiv.org/html/2607.22769#bib.bib12)\], heavily overriding HuggingFace Trainer internals \(compute\_loss,\_inner\_training\_loop\)\. Industrial pretraining frameworks such as Megatron\-LM\[[6](https://arxiv.org/html/2607.22769#bib.bib11)\]and MindSpeed build their own data loaders, training loops, and distributed communication patterns\. Rebuilding the DataLoader for Select or Mix modes is infeasible under Megatron’s binary mmap pipeline\. Consequently, only Weight\-like interventions—which require merely intercepting per\-sample loss and domain labels—are practically portable\. DomainPilot’s patch\-based architecture is designed explicitly around this constraint\.

## 3Methodology

DomainPilot comprises four interconnected components: \(1\) token\-level domain loss monitoring that captures per\-domain learning dynamics at training time; \(2\) Scaling Law guided coarse optimization that derives a principled prior for mixture reallocation; \(3\) Mixing Law guided fine optimization that refines proportions by modeling cross\-domain interactions; and \(4\) a patch\-based architecture that realizes these mechanisms without modifying the underlying training framework\. We describe each in turn\.

### 3\.1Domain Loss Monitoring

The foundation of DomainPilot is the ability to measure, at every training step, how much loss each domain contributes\. Unlike sample\-level loss—which conflates noise, difficulty, and novelty—domain\-level loss aggregates signals across thousands of tokens within a domain, averaging out sample\-specific outliers and revealing genuine learning trends\.

#### Token\-level domain identification\.

During preprocessing, theMultiDomainPackedHandlerperforms knapsack packing over sub\-samples from multiple domains\. While packing, it tags each token with its sourcedomain\_id, producing paired binary files \(\_packed\_domain\.bin/\.idx\) aligned with the standard token indices\. Because packing can combine sub\-samples from different domains within a single sequence, domain labels are maintained attoken granularity, yielding a domain\-id vector of length equal to the sequence length\.

#### Data loading and forward propagation\.

TheDecoderPackedMTFDatasetloads token\-level domain IDs alongside input IDs and injects them into the sample dictionary\. The SFT trainer broadcasts domain IDs through the standard batch\-communication path, passing them into the model’sforwardmethod as keyword arguments\.

#### Domain\-wise loss aggregation\.

Inside the model forward pass, after the standard logits computation, a lightweight post\-processing step groups tokens bydomain\_idand accumulates\(loss\_sum,token\_count\)\(\\text\{loss\\\_sum\},\\text\{token\\\_count\}\)pairs per domain\. These statistics are stored inself\.\_last\_domain\_losses\. At each logging interval, a training\-log patch performs DP All\-Reduce across data\-parallel ranks to obtain globally consistent domain losses, which are then written to TensorBoard and console logs\.

The entire pipeline adds<<1% throughput overhead because it reuses existing loss tensors and only introduces lightweight indexing and reduction operations\.

### 3\.2Scaling Law Guided Coarse Optimization

With domain loss trajectories in hand, Stage 1 fits adomain\-specific Scaling Lawto characterize how each domain’s loss evolves with training progress\. The functional form is:

Li​\(D\)=ai​D−αi\+biL\_\{i\}\(D\)=a\_\{i\}D^\{\-\\alpha\_\{i\}\}\+b\_\{i\}\(1\)
whereDDis the number of training steps \(proxy for data volume\),aia\_\{i\}is the initial loss amplitude,bib\_\{i\}is the irreducible \(converged\) loss, andαi\\alpha\_\{i\}is the convergence speed\. Fitting this curve to the monitored trajectories yields three interpretable parameters per domain:

- •aia\_\{i\}: Reflects initial domain difficulty and pretraining familiarity\.ai≈0a\_\{i\}\\approx 0indicates the model already possesses strong prior knowledge of the domain\.
- •bib\_\{i\}: The asymptotic loss; higher values indicate inherently harder domains\.
- •αi\\alpha\_\{i\}: Learning efficiency; smaller values mean slower convergence and typically signal insufficient data volume\.

#### Four\-factor reward score\.

We derive a coarse reallocation score by combining these parameters:

scorei=biminj⁡bj×maxj⁡αjαi×Licurr−bimaxj⁡\(Ljcurr−bj\)\.\\text\{score\}\_\{i\}=\\frac\{b\_\{i\}\}\{\\min\_\{j\}b\_\{j\}\}\\times\\frac\{\\max\_\{j\}\\alpha\_\{j\}\}\{\\alpha\_\{i\}\}\\times\\frac\{L\_\{i\}^\{\\text\{curr\}\}\-b\_\{i\}\}\{\\max\_\{j\}\(L\_\{j\}^\{\\text\{curr\}\}\-b\_\{j\}\)\}\.\(2\)
The three factors respectively capture: \(i\) convergence difficulty \(higherbib\_\{i\}needs more data\), \(ii\) learning speed \(lowerαi\\alpha\_\{i\}needs more data\), and \(iii\) remaining improvement headroom\. The new mixture proportion is then:

pinew=piold⋅scorei∑jpjold⋅scorej\.p\_\{i\}^\{\\text\{new\}\}=\\frac\{p\_\{i\}^\{\\text\{old\}\}\\cdot\\text\{score\}\_\{i\}\}\{\\sum\_\{j\}p\_\{j\}^\{\\text\{old\}\}\\cdot\\text\{score\}\_\{j\}\}\.\(3\)
Table[2](https://arxiv.org/html/2607.22769#S3.T2)reports the fitted parameters for Qwen3\-1\.7B SFT\.

Table 2:Fitted Scaling Law parameters for Qwen3\-1\.7B SFT domains\.
#### Interpretation for Qwen3\-1\.7B\.

termagentexhibits the slowest convergence \(α=0\.085\\alpha=0\.085\) and receives a\+25%\+25\\%boost\.swe\(α=0\.098\\alpha=0\.098\) receives\+15%\+15\\%\.chat, conversely, converges fastest \(α=0\.548\\alpha=0\.548\) and is reduced by−8%\-8\\%to free budget\.mathmaintains its proportion because it is already learning efficiently despite high difficulty\.

### 3\.3Mixing Law Guided Fine Optimization \(Planned\)

While Scaling Law characterizes intra\-domain learning in isolation, real\-world mixture optimization must account forcross\-domain interactions—e\.g\., code data may indirectly improve mathematical reasoning through shared logical structure\. We outline a second\-stageMixing Lawrefinement procedure, which remains future work\.

#### Modeling cross\-domain interactions\.

Letpnewp^\{\\text\{new\}\}denote the mixture produced by Stage 1\. We model the validation loss of domainiiunder mixtureppas a second\-order expansion aroundpnewp^\{\\text\{new\}\}:

Li​\(p\)≈Li​\(pnew\)\+∑jAi​j​\(pj−pjnew\)\+∑j,kBi​j​k​\(pj−pjnew\)​\(pk−pknew\),L\_\{i\}\(p\)\\approx L\_\{i\}\(p^\{\\text\{new\}\}\)\+\\sum\_\{j\}A\_\{ij\}\(p\_\{j\}\-p\_\{j\}^\{\\text\{new\}\}\)\+\\sum\_\{j,k\}B\_\{ijk\}\(p\_\{j\}\-p\_\{j\}^\{\\text\{new\}\}\)\(p\_\{k\}\-p\_\{k\}^\{\\text\{new\}\}\),\(4\)
where𝐀\\mathbf\{A\}captures first\-order sensitivity andℬ\\mathcal\{B\}captures pairwise interaction effects\. The interaction coefficients are estimated from a small grid of controlled sweep experiments in which each domain proportion is perturbed within a±20%\\pm 20\\%neighborhood ofpnewp^\{\\text\{new\}\}while holding the total token budget fixed\.

#### Stage 2 optimization\.

Given the fitted interaction model, the final mixturepfinalp^\{\\text\{final\}\}is obtained by solving a constrained optimization problem that minimizes the aggregate predicted validation loss subject to∑ipi=1\\sum\_\{i\}p\_\{i\}=1andpi≥0p\_\{i\}\\geq 0\. Because Stage 2 requires additional sweep experiments, the empirical results reported in Section[5](https://arxiv.org/html/2607.22769#S5)reflect only Stage 1 \(Scaling Law\) optimization; Stage 2 validation is left as immediate future work\.

### 3\.4Patch\-Based Architecture

DomainPilot is implemented as a set of minimal, non\-intrusive patches rather than a monolithic framework fork\. This design respects the operational reality that production training stacks \(MindSpeed, Megatron\-LM, DeepSpeed\) are under active development and cannot be frozen for algorithmic modifications\.

#### Two\-layer design\.

The architecture separatesframework\-agnostic algorithm logic\(pure PyTorch tensor operations\) fromframework\-specific thin adapters\(interceptors for batch construction, forward arguments, and loss computation\)\. Adding support for a new training framework requires only implementing the adapter interface \(∼\\sim30 lines of Python\), leaving the underlying framework untouched\.

#### MindSpeed adapter overview\.

For MindSpeed, the patch set comprises:

- •arguments\_patch: Exposes\-\-enable\-domain\-lossCLI flags\.
- •data\_handler\_patch: Persists token\-level domain IDs during preprocessing\.
- •gpt\_dataset\_patch: Injects domain IDs into the sample dictionary\.
- •gpt\_model\_patch: Appends domain\-wise loss aggregation toforward\.
- •training\_utils\_patch: Broadcasts domain IDs and accumulates global statistics\.

This patch\-based approach guarantees that upstream framework updates can be merged without conflict, and the algorithm layer can be validated independently via unit tests against synthetic tensors\.

## 4Experiment

### 4\.1Model and Training Setup

All experiments are conducted on theQwen3\-1\.7Bbase model, a densely activated transformer with 1\.7 billion parameters\. The supervised fine\-tuning \(SFT\) stage is executed on the MindSpeed training framework, which is built atop Megatron\-LM and represents the industrial standard for large\-scale LLM pretraining and fine\-tuning in our deployment environment\.

Training data comprises eight domains:chat,convagent,if\(instruction following\),math,safety,science,swe\(software engineering\), andtermagent\. The original mixture ratios are determined by manual heuristic tuning\. In the optimized mixture, ratios are adjusted according to the two\-stage pipeline described in Section[3](https://arxiv.org/html/2607.22769#S3)\.

### 4\.2Domain Loss Monitoring

During training, we activate the Domain Loss monitoring patch to record per\-domain loss trajectories\. Domain identifiers are injected at the token level during data preprocessing: theMultiDomainPackedHandlertags each token with its source domain while performing knapsack packing, producing paired\_packed\_domain\.bin/\.idxfiles\. At each forward step, losses are aggregated by domain ID via DP All\-Reduce synchronization across data\-parallel ranks\. Monitoring overhead is negligible \(<<1% throughput degradation\) because the patch only intercepts existing loss tensors without altering the data\-loading pipeline\.

### 4\.3Evaluation Benchmarks

We evaluate on four representative benchmarks covering diverse capabilities:

- •MMLU\-Redux: Comprehensive knowledge reasoning across 57 subjects\.
- •AIME24: Mathematical competition\-level reasoning, reported as Pass@1 / Cons@64\.
- •LiveCodeBench v5: Code generation capability, reported as Pass@1 / Pass@5\.
- •BFCL v3: Function calling and tool\-use ability\.

### 4\.4Baselines

We compare four model states:

1. 1\.Base: The pretrained Qwen3\-1\.7B checkpoint without any SFT\.
2. 2\.SFT \(report\): Scores reported in the official Qwen3 technical report for reference\.
3. 3\.SFT \(original\): Our reproduction using the original, manually tuned data mixture\.
4. 4\.SFT \(optimized\): Our reproduction using the mixture adjusted by DomainPilot’s two\-stage pipeline\.

## 5Results

### 5\.1Main Results

Table[3](https://arxiv.org/html/2607.22769#S5.T3)summarizes the benchmark scores across the four model states\.

Table 3:Benchmark comparison across model states\. Optimization column shows relative improvement of SFT \(optimized\) over SFT \(original\)\.The optimized mixture consistently outperforms the original mixture across all benchmarks\. Notably,LiveCodeBench v5andBFCL v3—corresponding to thesweandtermagentdomains that Scaling Law analysis identified as under\-represented—exhibit the largest relative gains \(\+3\.8%\+3\.8\\%and\+3\.6%\+3\.6\\%, respectively\)\. MMLU\-Redux improves by\+2%\+2\\%, while AIME24 sees a\+1\.8%\+1\.8\\%gain\. These improvements are achievedwithout increasing total training data or compute budget; only the inter\-domain proportions are reallocated\.

### 5\.2Analysis by Capability Domain

#### Code Generation \(swe\)\.

LiveCodeBench v5 Pass@1 rises from 23\.2 to 27\.0\. The Scaling Law analysis had flaggedsweas having low convergence speed \(α=0\.098\\alpha=0\.098\) and high convergence loss \(b=1\.990b=1\.990\), signaling insufficient data\. The optimized mixture increases thesweproportion by\+15%\+15\\%, which directly translates into the largest benchmark improvement \(\+3\.8%\+3\.8\\%\)\. This validates that domain\-level loss signals can accurately pinpoint bottleneck domains\.

#### Agent / Tool Use \(termagent\)\.

BFCL v3 improves from 54\.5 to 58\.1\.termagentwas identified as the slowest\-learning domain \(α=0\.085\\alpha=0\.085\), and its mixture share is boosted by\+25%\+25\\%\. The resulting\+3\.6%\+3\.6\\%gain confirms that even domains with small initial representation can yield substantial downstream improvements when adequately resourced\.

#### Mathematical Reasoning \(math\)\.

AIME24 improves modestly from 41\.1/10\.6 to 42\.9/11\.7\. The Scaling Law parameters formath\(b=2\.117b=2\.117,α=0\.112\\alpha=0\.112\) indicated that the domain is inherently difficult but already learning efficiently under the original mixture\. Consequently, DomainPilot recommends maintaining themathproportion, and the observed micro\-gain is consistent with expectations\.

#### General Knowledge \(chat\+ others\)\.

MMLU\-Redux improves by\+2%\+2\\%\. Thechatdomain was flagged as potentially over\-represented \(α=0\.548\\alpha=0\.548, fastest convergence\), and its proportion is reduced by−8%\-8\\%, freeing budget forsweandtermagent\. The fact that general\-knowledge performance still improves suggests that the original mixture had misallocated resources away from higher\-impact domains\.

### 5\.3Comparison with Stronger Baselines

To contextualize the gains, we compare against two stronger baselines in Table[4](https://arxiv.org/html/2607.22769#S5.T4)\.

Table 4:Comparison against stronger baselines: Qwen2\.5\-3B \(larger model\) and DeepSeek\-R1\-Distill\-Qwen\-1\.5B \(distilled model\)\.#### vs\. Qwen2\.5\-3B\.

Despite having only 57% of the parameters, the optimized Qwen3\-1\.7B surpasses Qwen2\.5\-3B on BFCL v3 \(\+15\.6%\+15\.6\\%relative\) and MATH\-500 \(\+10\.1%\+10\.1\\%relative\)\. This suggests that data mixture optimization can unlock effective capacity equivalent to a 2×\\timesparameter scaling at a fraction of the training cost\.

#### vs\. DeepSeek\-R1\-Distill\-Qwen\-1\.5B\.

The optimized model widens its lead over this knowledge\-distilled counterpart on all benchmarks, notably LiveCodeBench \(\+27\+27absolute points\)\. This underscores that native training with principled mixture optimization outperforms distillation\-based approaches when data composition is carefully tuned\.

### 5\.4Cost\-Benefit Analysis

Table[5](https://arxiv.org/html/2607.22769#S5.T5)compares the cost\-effectiveness of different improvement strategies\.

Table 5:Cost\-benefit comparison of improvement strategies\.Mixture optimization requires only lightweight sweep experiments \(∼\\sim10% of a full SFT run\) yet delivers gains comparable to 50% of the benefit from doubling model size, with zero inference\-cost penalty\.

## 6Conclusion

We presented DomainPilot, a domain\-level loss\-guided framework for two\-stage data mixture optimization\. By monitoring per\-domain loss trajectories at token granularity through a non\-intrusive patch architecture, DomainPilot extracts scalable signals that characterize how each domain learns during training\. These signals feed a Scaling Law coarse\-optimization stage that identifies under\- or over\-represented domains, followed by a Mixing Law fine\-optimization stage that models cross\-domain interactions through controlled sweep experiments\.

Empirical validation on Qwen3\-1\.7B SFT demonstrates consistent improvements across diverse benchmarks:\+3\.8%\+3\.8\\%on LiveCodeBench v5,\+3\.6%\+3\.6\\%on BFCL v3,\+2%\+2\\%on MMLU\-Redux, and\+1\.8%\+1\.8\\%on AIME24\. Crucially, these gains are achieved solely by reallocating existing data proportions, with no increase in total data volume, training compute, or inference cost\. Comparison against Qwen2\.5\-3B further suggests that principled mixture optimization can unlock effective capacity equivalent to a2×2\\timesparameter scaling\.

#### Limitations\.

Our current validation is limited to the 1\.7B parameter scale and the SFT stage; pretraining\-scale validation remains future work\. The Mixing Law sweep experiments \(Stage 2\) are not yet complete, so the reported results reflect only Scaling Law coarse optimization\. Additionally, the Weight mechanism migration \(Section[3](https://arxiv.org/html/2607.22769#S3)\) remains at the architectural\-design stage and has not been experimentally validated\.

#### Future work\.

We identify four immediate directions:

1. 1\.Larger\-scale validation\.Apply DomainPilot to Qwen3\-4B and 8B models, and to the pretraining stage, to verify that domain\-level scaling laws hold across model sizes and training regimes\.
2. 2\.Cross\-architecture transfer\.Implement the thin adapter layer for HuggingFace Trainer and native Megatron\-LM to demonstrate framework portability\.
3. 3\.Weight mechanism integration\.Combine domain\-level reweighting with sample\-level quality scores \(e\.g\., offline LLM\-based ratings\) to disentangle noise from genuine difficulty\.
4. 4\.Data cleaning feedback loop\.Use domain loss andΔ\\DeltaLoss signals to identify low\-quality source documents and establish a closed\-loop pipeline from training feedback to data curation\.

## References

- \[1\]S\. Diaoet al\.\(2025\)Nemotron\-climb: clustering\-based iterative data mixture optimization\.Advances in Neural Information Processing Systems38\.Cited by:[§1](https://arxiv.org/html/2607.22769#S1.p4.1),[§2\.2](https://arxiv.org/html/2607.22769#S2.SS2.p1.1)\.
- \[2\]J\. Hoffmann, S\. Borgeaud, A\. Mensch, E\. Buchatskaya, T\. Cai, E\. Rutherford, D\. d\. L\. Casas, L\. A\. Hendricks, J\. Welbl, A\. Clark,et al\.\(2022\)Training compute\-optimal large language models\.arXiv preprint arXiv:2203\.15556\.Cited by:[§1](https://arxiv.org/html/2607.22769#S1.p1.1)\.
- \[3\]J\. Kaplan, S\. McCandlish, T\. Henighan, T\. B\. Brown, B\. Chess, R\. Child, S\. Gray, A\. Radford, J\. Wu, and D\. Amodei\(2020\)Scaling laws for neural language models\.arXiv preprint arXiv:2001\.08361\.Cited by:[§1](https://arxiv.org/html/2607.22769#S1.p1.1)\.
- \[4\]Z\. Lin, Z\. Gou, Y\. Gong, X\. Liu, Y\. Shen, R\. Xu, C\. Lin, Y\. Yang, J\. Jiao, N\. Duan, and W\. Chen\(2024\)Rho\-1: not all tokens are what you need\.arXiv preprint arXiv:2404\.07965\.Cited by:[§1](https://arxiv.org/html/2607.22769#S1.p5.1),[§2\.3](https://arxiv.org/html/2607.22769#S2.SS3.p1.1)\.
- \[5\]R\. Pan, D\. Zhang, H\. Zhang, X\. Pan, M\. Xu, J\. Zhang, R\. Pi, X\. Wang, and T\. Zhang\(2025\)ScaleBiO: scalable bilevel optimization for llm data reweighting\.Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics\.Cited by:[§1](https://arxiv.org/html/2607.22769#S1.p4.1),[§2\.2](https://arxiv.org/html/2607.22769#S2.SS2.p1.1)\.
- \[6\]M\. Shoeybi, M\. Patwary, R\. Puri, P\. LeGresley, J\. Casper, and B\. Catanzaro\(2019\)Megatron\-lm: training multi\-billion parameter language models using model parallelism\.arXiv preprint arXiv:1909\.08053\.Cited by:[§2\.4](https://arxiv.org/html/2607.22769#S2.SS4.p1.1)\.
- \[7\]A\. Wettig, T\. Li, M\. Kim, Z\. Yao, and D\. Zhang\(2024\)QuRating: selecting high\-quality data for training language models\.International Conference on Machine Learning\.Cited by:[§1](https://arxiv.org/html/2607.22769#S1.p3.1),[§2\.1](https://arxiv.org/html/2607.22769#S2.SS1.p1.2)\.
- \[8\]M\. Xia, S\. Malladi, S\. Gururangan, S\. Arora, and D\. Chen\(2024\)Less: selecting influential data for targeted instruction tuning\.arXiv preprint arXiv:2402\.04333\.Cited by:[§1](https://arxiv.org/html/2607.22769#S1.p3.1),[§2\.1](https://arxiv.org/html/2607.22769#S2.SS1.p1.2)\.
- \[9\]S\. M\. Xie, T\. Liang, and T\. Ma\(2023\)Data selection for language models via importance resampled mcmc\.Advances in Neural Information Processing Systems36\.Cited by:[§1](https://arxiv.org/html/2607.22769#S1.p3.1),[§2\.1](https://arxiv.org/html/2607.22769#S2.SS1.p1.2)\.
- \[10\]S\. M\. Xie, H\. Pham, X\. Dong, N\. Du, H\. Liu, Y\. Lu, P\. Liang, T\. Ma, and A\. W\. Yu\(2023\)DoReMi: optimizing data mixtures speeds up language model pretraining\.Advances in Neural Information Processing Systems36\.Cited by:[§1](https://arxiv.org/html/2607.22769#S1.p1.1),[§1](https://arxiv.org/html/2607.22769#S1.p4.1),[§2\.2](https://arxiv.org/html/2607.22769#S2.SS2.p1.1)\.
- \[11\]Z\. Zhao, M\. Qiang, M\. Chen, L\. Ma, R\. Yu, H\. Feng, S\. Sun, Z\. Meng, X\. Ma, X\. Yang, Q\. Cai, R\. An, B\. Zeng, Z\. H\. Wong, C\. Shen, R\. He, Z\. Han, Y\. Zheng, F\. Fu, C\. He, B\. Cui, Z\. Li, W\. E, and W\. Zhang\(2026\)DataFlex: a unified framework for data\-centric dynamic training of large language models\.arXiv preprint arXiv:2603\.26164\.Cited by:[§1](https://arxiv.org/html/2607.22769#S1.p5.1),[§2\.3](https://arxiv.org/html/2607.22769#S2.SS3.p1.1)\.
- \[12\]Y\. Zheng, R\. Zhang, J\. Zhang, Y\. Ye, Z\. Luo, and Y\. Zhang\(2024\)LlamaFactory: unified efficient fine\-tuning of 100\+ language models\.arXiv preprint arXiv:2403\.13372\.Cited by:[§2\.4](https://arxiv.org/html/2607.22769#S2.SS4.p1.1)\.