Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models
Authors: Yuan Li
Zhengzhong Liu
Eric P. Xing
Venue: ICML 2025 poster (2025)
Keywords: Large language models
supervised fine-tuning
data mixing
Primary area: deep_learning->large_language_models

Abstract: Optimizing data mixtures for supervised fine-tuning (SFT) of large language models (LLMs) is critical for developing general-purpose models, yet this area remains underexplored. In this paper, we frame data mixing as an optimization problem and introduce a novel method designed to minimize validation loss. Our approach parametrizes the loss by modeling effective data transferred and leveraging scaling laws for fine-tuning. By experimenting with various small-scale data mixtures, we fit these parameters and derive the optimal weights. We provide both mathematical proofs and empirical results demonstrating that our algorithm achieves excellent overall and individual performance across all domains. Through controlled experiments, we show that models trained with our optimized weights perform on par with those using optimal weights determined via grid search, with per-domain loss only 0.66\% higher than the best domain loss from grid search on average. Additionally, we show that reweighting popular SFT datasets using our method improves both validation loss and downstream performance. Finally, we discuss how our method can generalize to guide data selection for domain-specific models and provide insights into SFT.

Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models
Yuan Li Zhengzhong Liu Eric Xing 12
Abstract
healthcare) procedure termed data mixing. Data mixing has significant impacts on models' performance (Albalak et al., 2023; Xie et al., 2024). However, it is impractical to identify the best data mixture by exhaustively trying every possible combination domain weight.
Optimizing data mixtures for supervised fine- tuning (SFT) of large language models (LLMs) is critical for developing general-purpose models, yet this area remains underexplored. In this paper, we frame data mixing as an optimization problem and introduce novel method designed to mini- mize validation loss. Our approach parametrizes the loss by modeling effective data transferred and leveraging scaling laws for fine-tuning, By exper- imenting with various small-scale data mixtures, we fit these parameters and derive the optimal weights. We provide both mathematical proofs and empirical results demonstrating that our al- gorithm achieves excellent overall and individual performance across all domains. Through con- trolled experiments, we show that models trained with our optimized weights perform par with those using optimal weights determined via grid search, with per-domain loss only 0.66% higher than the best domain loss from grid search on average. Additionally, we show that reweighting popular SFT datasets using our method improves both validation loss and downstream performance. Finally, we discuss how our method can gener- alize to guide data selection for domain-specific models and provide insights into SFT.
Prior research has concentrated on optimizing data mixtures during the pre-training of LLMs. prevalent strategy for determining optimal weights of pre-training data invol defining or parametrizing the loss as a function of domain weights and subsequently minimizing this function. This overarching strategy encompasses various approaches. One such approach employs small proxy models to optimize do- main weights, which are then applied train larger models (Xie et al., 2024; Fan et al., 2024). (Xie et al., 2024; Fan et al., 2024). Alternatively, methods like Data Mixing Laws and RegMix (Ye et al., 2024; Liu et al., 2024) parameterize the relationship between model performance (e.g., loss) and data mixtures. These methods involve employing a function to estimate the performance and then using the function to optimize the data mixture.
Despite significant efforts in pre-training, data mixing for supervised fine-tuning (SFT) remains underexplored. This oversight is unwarranted, as recent research indicates that SFT not only aligns the model stylistically but also infuses knowledge and enhances the model's 's capabilities (Wu et al., 2024; Raghavendra al., 2024). Moreover, existing data mixing methods designed for pre-training are not directly applicable to SFT due to several challenges: 1) Optimizing domain weights using small proxy models may not translate effectively to larger models. Since SFT builds upon models pre-trained with distinct datasets, it is questionable whether the optimal weights derived from proxy models remain valid for the specific large-scale models at SFT stage; 2) Even when the same model used, parameterizing the loss without accounting for the interplay and scaling effects of different domains can lead to poor extrapolation. At the SFT stage, the goal is to train general-purpor model with robust performance across multiple domains. Existing loss minimization does not ensure adequate representation for each domain. fact, these loss parametrizations lead to unbalanced domain weight distributions when scaling up data sizes.
1. Introduction
Large Language Models (LLMs) are general-purpose sys- tems capable of performing wide range of tasks, such following instructions and solving mathematical problems. To train these models, critical factor is to construct train- ing datasets that contain diverse domains (e.g., math, code,
'Carnegie Mellon University *Mohamed bin Zayed University of Artificial Intelligence. Correspondence to: Yuan Li synantil@anddew.cmu..du. Erie Xing @ww.xinmmbmuaaa...a.....
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Data Mixing Optimization
These challenges call for a systematic data mixing method for SFT. In this paper, we introduce Data Mixing Optimiza- tion, method that frames data mixing an optimization problem to determine optimal domain weights for given data budgets. Our approach involves a novel parameterization of validation loss tailored for SFT. This parameterization accounts for both the scaling effects of fine-tuning data and the interactions between different domains. We analyze our optimization method and demonstrate that our method offers two key advantages. First, it shares similar intuition and principles with existing methods (Xie et al., 2024; Fan et al., 2024; Kang et al., 2024) in up-weighting domains that effectively reduce the loss. Second, it is theoretically robust because we rigorously proved that our algorithm ensures balanced performance across all domains.
well-studied at the pre-training stage. Methods like DoReMi (Xie al., 2024) and DoGE (Fan al., 2024) identify opti- mal weights using smaller models and subsequenty apply these weights to train larger models. Other approaches, such as Data Mixing Laws (Ye et al., 2024), RegMix (Liu et al., 2024), and Aioli (Chen et al., 2024a), estimate loss as a function of domain weights and interpolate performance to determine optimal data compositions. However, these methods often struggle to generalize across different model scales or data distributions, limiting their applicability to SFT. To the best of our knowledge, SMART (Renduchin- tala et al., 2024) is the only method explicitly designed to optimize domain weights for SFT; however, it relies on prompt embeddings and remains model and scale-invariant.
In contrast, our proposed approach determines optimal do- main weights that depend on both model size and scale, thus avoiding the limitations of previous methods.
We conducted extensive experiments across various sce- narios, including controlled studies using three distinct do- mains: instruction following, math, and code; and experi- ments of re-weighting popular SFT collections. Our results show that our method effectively derives domain weights that minimize validation loss. Furthermore, defining do- mains as specific tasks (e.g., math, code) can make our method useful for improving downstream performance on specific tasks. We also show that our data mixing optimiza- tion can be extended to guide the training of Nomain-specinf LLMs and provide insights into SFT.
Supervised Fine- Tine-Tunigg SFT, or instruction fine-tuning, involves training pre-trained LLMs with instruction-output pairs (Wei et al., 2022; Chung et al., 2024; Peng et al., 2023). Notable SFT collections include FLAN (Wei et al., 2022), the Tulu series (Wang et al., 2023; Ivison al., 2023: Lam- bert et al., 2024), OpenHermes (Teknium, 2023), Infinity Instruct (Zhao et al., 2024), and Orca (Mukherjee al., 2023), which aggregate data from diverse domains such as math, code, and dialogue to enhance model capabilities. According to Superficial Alignment Hypothesis (Zhou et al., 2024), the primary role of SFT is to align the model's for- matting and style. However, recent studies (Raghavendra et al., 2024) argue that improvements from SFT follow a power law, suggesting that SFT enhances more than just style alignment, including knowledge infusion. Addition- ally, factors such as the number of tasks, model size, in- struction settings, and output length significantly influence LLMs' performance (Puri et al., 2023; Chung et al., 2024; Longpre et al., 2023; Zhao et al.). Unlike existing meth- ods that focus on dataset quality, our work optimizes data mixtures within specified budgets across to enhance SFT.
To summarize, our contributions are three-fold:
Data Mixing Optimization for SFT: We introduce a method that optimizes domain weights for SFT data. Central to this optimization is a novel parameterization of validation loss that models the scaling effects of SIT data and domain interactions. This method is theoreti- cally solid and ensures no domain will underperform. Extensive Empirical Results on Data Mixing: We present extensive SFT results about data mixing across various LLM classes and model sizes. The findings not only validate the effectiveness of our method but also reveal interesting patterns, such as the similarity in optimal weights among models within the same LLM classes, that pave the way for future research.
3. Data Mixing Optimization
In this section, we formulate data mixing as an optimization problem. The objective is to determine the optimal data composition that minimizes the validation loss given a fixed computational budget (i.e., training dataset size in tokens). We introduce a novel algorithm to solve this optimization problem and derive the optimal weights for data domains.
Implications and Insights for SFT: We demonstrate that data mixing optimization can guide data selection for specialized tasks and serve as a framework to un- derstand key aspects of SFT, including data scaling and domain interactions.
2. Related Work
3.1. Problem Formulation
Due to space limitations, we briefly discuss related work here and leave detailed analysis in Appendix A.
Consider fine-tuning an LLM using training dataset D UK D., where D, corresponds to domain i, and K denotes the total number of domains. We similarly denote validation dataset as D = -K= Drail Let N = |D] represent the total number of tokens D and N, = |Da] be the number
Data Mixing, The composition of training data for LLMs has drawn significant interest due to its impact on perfor- mance. Data mixing, also known as domain reweighting, is
2

---

Data Mixing Optimization
of tokens in domain i. We introduce domain weight vector (w1, 102- .... wx), where each weight w; represents the proportion of domain i. The weight vector satisfies the constraints of the K -i-imensional probability simplex:
loss when trained exclusively on data from different distri- butions (Hernandez et al., 2021). For domain i, the total effective data can be calculated as the sum of in-domain data, denoted as [Da], and effective data transferred to domain i, denoted as [D] [D = |2] + |-
K
w WEAR
W WERK
2
0
Vi,
W 1
1
Inspired by Hernandez et al. (2021), we represent [D]'] as a function of the dataset size from other domains |Dyal:
Each domain weight is calculated as w;
N
|
k; 1
(2)
The model parameters 0 are trained by minimizing the em- pirical loss over the training dataset D. yielding the param- eters 0- = arg ming L(0,D), D). We use 0' and 0" (N, w) interchangeably to emphasize its dependence on data com- position (N, w). After training, the model 0" is evaluated on the validation set Doal by computing the aggregated loss across all validation domains!:
where i is constant associated with the model architecture and data distribution differences, and ag is scaling coef- ficient. We apply Equation (2) to estimate the equivalent dataset size transferred from other domains when calculat- ing the domain validation loss.
3.3. Scaling Law for Fine-Tuning
K
I( Dual)
Σ L(O" ((
We estimate the validation loss as a function of the effective dataset size. According to Zhang et al. (2024), the scaling law for fine-tuning data follows a power-law relationship:
The objective is to seek the domain weight vector w such that, after training on the weighted domains, the model's validation loss is minimized. We formalize this into an optimization problem:
L(D) = C- 1 IDI + E,
(3)
where I. is the loss, and C encompasses factors such as model size. In our setting, since these factors are held constant, C acts as multiplicative term. 3 is the scaling coefficient, and E represents the irreducible loss. This scal- ing law for fine-tuning suggests that loss decreases rapidly with an initial increase in data but gradually plateaus as the dataset size grows. These implications align with previous studies that fine-tuning with a small number of examples can achieve decent results (Zhou et al., 2024), while larger fine-tuning datasets continue to enhance downstream perfor- mance (Raghavendra et al., 2024).
min w
(N,, (),, D
subject to
0 wi S
E 21, K}.
K 1,
1=1 where 0" (N,, w) = arg min L(0, D).
There is no closed-form solution for the optimization prob- lem 1. Specifically, finding the optimal w* involves an iterative optimization process. In each iteration, the model is trained using specified weight vector, and the weights are adjusted based on the validation loss to enhance perfor- mance. However, this procedure is computationally expen- sive, as each iteration necessitates retraining the entire LLM. To address this issuc, we propose an estimation of validation loss through the ideas of effective data from transfer (Her- nandez et al., 2021) and scaling law for fine-tuning (Zhang al., 2024).
Leveraging effective data from transfer in Equation (2) and scaling law for fine-tuning in Equation (3), we estimate the validation loss for domain i as:
i( T(O', ~ C. : (N. + N
8 + E; 2 C. (N, + ky (D,,
(4)
+ E,
Here, Dr, represents the aggregated datasets from all do- mains except domain i, and [D, = N - N,. 3, is a scaling coefficient, and C₁ h; and E, are domain-specific con- stants.
3.2. Effective Data from Transfer
With this formula, we fit five parameters for each domain in Equation (4) using only small amount of data. Motivated by Kang et al. (2024), we vary the data size for domain i while keeping the data sizes of other domains |Dy,l con- stant. Specifically, the data size for domain is varied five times, denoted as 7: for t = 0, 2, 4, where t 0 corresponds to the original allocation N. For each t, the corresponding validation loss I.(0) is recorded. To fit the
Effective data from transfer refers to the equivalent amount of training data that model would need to achieve the same
We compute the total validation loss across all domains, as- signing equal weights to each to develop a general purpose LLM that excels in multiple areas. Alternatively, weighting multipli- ers can be applied to individual domains to prioritize model per- formance in specific areas. We will discuss fine-tuning domain- specific L.MSs in Section 5.4.
general-purpose
3

---

Data Mixing Optimization
parameters, we follow the standardized procedure in Hoff- mann al. (2024) and Zhang al. (2024) and minimize the residuals using Huber loss (Huber, 1964) with 6 = 0.001. We further use Trust Region Method (Conn et al., 2000) with a nonlinear constraint (Wächter & Biegler, 2006) to ensure that the effective data from transfer kę (N N,,; remains less than the original quantity N -N,. The objective is formalized as:
Algorithm Data Mixing Optimization
Input: Data budget No: sample size N (N < No): number of domains K; data sources Si...., SK: valida- tion sets Dual
(Dial .....
.....,Daly Drel;; perturbation ratios
Tl..... .44
# Initialization for = 1 to K do
Dj + Sample(S;, N) end for
Hubers ( - )
min C.
Train an initial model on D U D,; let 8* be the trained model.
(5)
1=0
subject to k IN - Ni < IN - -
C. L(6", L(
where {(N(") = C. (N!") + kg IN N.0000) -9:
# Perturbation Experiments for t = 1 to 4 do
+
E,
for j = to K do
D/ Sample R
3.4. Optimization Algorithm
D(t) (D),)) UD! (1)
Using the fitted parameters for each domain, we formulate the optimization problem for given data budget No and domain weight vector as follows:
Train a model on D(1)_ - 010)
c [,((()) Deal)
end for end for
min
L 7N00 w), D.
# Parameter Estimation for i = 1 to K do
K
Solve for {Ci, kis ais 3i, Eij via Trust Region:
R min
Σ
Ci (w;No + hi (No (. + Es
w
1=1
min
Huber (L(N(") - I,!"),
subject to
0
f wi
1
ViE (1,..,,.,,K)
subject to kg - N I IN WV-M No1.
K
Σ
end for
(6)
# Domain Weight Optimization
Solve for w = (Wi,..., wx) via SLSQP:
Note that the objective is convex function (proof of convexity in Appendix B) but has no closed-form solu- tion. To this end, we solve this optimization to find op- timal weight vector w using sequential least squares pro- gramming (SLSQP) algorithm (Nocedal & Wright, 2006), gradient-based method that handles both equality and inequality constraints. At each iteration, SLSQP con- structs quadratic approximation of the objective function ΣK C. (w;N + As (N );N)):)=* -8, dient information to refine the solution.
K
C, (w: No + ki (N w: -B, + E.
min W
2=1
subject to S wi S 1, Σ wi =
Output: w
B. + E, using gra-
from Instrrct (BAAI, 2024)). math (sampled from OpenMathInstruut-2 (Toshniwal et al., 2024)), and code (sampled from OpenCoder (Huang et al., 2024)). Each do- main has unit sample size of n 660,000 tokens. We set perturbation ratios to 1/2, 1/3, 2. and 3, and calculate the parameters by solving Equation (5). We use Llama3.2-3B as an example.
In summary, we integrate the aforementioned components and present the algorithm for determining the optimal data weights in Algorithm 1.
4. Interpretations and Analysis
In this section, we interpret our data mixing algorithm by examining key estimated parameters from perturbation ex- periments and analyzing mathematical implications.
Loss and Parameters. To build intuition, we first ana- lyze the validation loss for each domain. In Figure I(a), we present a matrix comparing the token-level perplexity (PPL) of each domain as domain data decreases from 3n to n tokens. The most notable finding is that increasing do-
Settings. We consider a scenario with three distinct domains: general instruction following (IF) (sampled
4

---

Data Mixing Optimization
1.30
0.44
-6:33
043
Thaningg Donan
0.40
Code
0.001
5.316
013
Weight
0.36
2,20
0.32
4:2
0.28
valoorio Doman
Domain Data Sizt odied tukong)
0.24
0
100
200
300
400
500
(X]
(b)
No (1e6 takens)
Figure 1. (a) Difference in per-domain perplexity between domain data sizes of 3n and n, while keeping each domain's 's data size fixed at n. (b) Overall perplexity in perturbation experiments with domain data sizes set to 1/3n, 1/2n, n, 2n, and 3n, while keeping each domain's data size fixed at n.
Figure 2. Estimated optimal weights across data budgets.
data budget increases, the changes in weights slow down. Domain's weight does not continue to increase or decrease significantly as the data budget increases; instead, it reaches a plateau. In other words, math data, which we discuss as being relatively less effective in reducing the loss, still maintains decent proportion. We next analyze why this is desirable property of our optimization algorithm to ensure that no domain will underperform.
main data significantly reduces its validation loss. Among all domains, IF data exhibits the most substantial impact. Furthermore, increasing IF data also decreases PPL in the math and code domains. This observation aligns with our intuition during SFT: math and code problems are formatted as instructions, enhancing instruction following capabílities enables LLMs to better address these tasks.
We then explore how the loss relates to the parameters in Equation (4) using the following parametrization (repeated here for clarity):
No Domains Left Behind. As shown in Figure 2, esti- mated domain weights do not experience significant changes as the data budget increases. This is a desirable property for developing a general-purpose model that achieves strong overall performance while maintaining robust capabilities across all domains. This behavior is driven by power-law relationship: initially, increasing data significantly reduces validation loss, but as data volume continues to scale, the marginal benefits diminish, necessitating the inclusion of data from other domains. Specifically, for each domain 8L is large when w; is small and
B. + E,
C=C; (N₂ + k; I IN-
All estimated parameters for the three domains are presented in Table 1, and the fitted loss curves are shown in Figure 1(b). The parameters 3, and C, are crucial for governing the trend of the loss curve, with B; serving as a decisive factor as the number of training tokens increases. We observe that 3; for the instruction following domain data has the largest value of 0.051. This indicates that scaling up IF data is the most efficient method for reducing loss, consistent with findings in Figure I(a), whereas math data has the minimal effect.
i, the partial derivative thr. decreases w, increases. This diminishing returns effect ensures that no single domain disproportionately dominates the training process as the data budget grows, maintaining balanced performance across all domains. We provide for- mal proof in Appendix C and present empirical results in Section 5 to support this argument.
Table 1. Estimated parameters Ct. k. a., 8. and E, for instruc- tion following (IF), math, and code domains.
Domain
C.
h:
3
E;
5. Experiments
IF
1.1562
0.1948
0.5288
0.0510
1.0967
Math
0.7512
0.0401
0.4467
0.0430
1.4934
Code
0.9820
0.1235
0.5235
0.0439
1.2679
In this section, we present extensive experiments conducted under two scenarios: (1) mixing three distinct domains IF, math, and code-given fixed data budgets; and (2) re- weighting domains for two popular SFT collections: Tulu3 (Lambert et al., 2024) and Orca (Mukherjee al., 2023). We also evaluate models trained with re-weighted data on downstream tasks demonstrate the practical utility of our method. Following Hernandez et al. (2021), we use token- level PPL as the primary performance metric, measured on held-out validation data for each domain. We provide exper- imental details, including data processing, hyperparameters, and evaluation benchmarks in Appendix D.
Estimated Optimal Domain Weights. We then solve Equation (6) and plot the weights under different data bud- gets in Figure 2. We made two important observations: (1) The optimal weights are scale dependent. This find- ing resonates with the arguments in Kang et al. (2024) and challenges the assumption of scale independence presented in many papers. It further echoes Dong et al. (2023) by showing that different domains scale differently, resulting in relative weight adjustments based on scale; (2) As the
5

---

Data Mixing Optimization
@wen2.5-0.55
Qwen2.5-1.58
Llamm332-33
Llama3.1-88
20M
200M
Figure surface plots PPL data budgets 5M, 20M, and 200M tokens. Each plot illustrates PPL results from models trained with different domain mixture weights derived from grid search (black dots). The weights estimated by our method are highlighted with yellow stars. Only the IF and math weights are shown on the x and y axes, respectively, as code weight dependent.
5.1. Mixing Domains at Fixed Data Budgets
higher than the best PPL. Further analysis of overall PPL and model classes is discussed in Appendix E.
Settings. To assess our approach, we consider three dis- tinct domains:
IF, math, and code, using the same data sources as described in Section 4. For each data budget (5M, 20M, and 200M tokens), compute the optimal weights, mix the three domains accordingly, and train the models. We use four LLMs from two classes: Qwen2.5 (0.5B and 1.5B) and Llama (3.2-3B and 3.B-8B parameters).
Best (Grid Search)
3.0
2.923
2.8
24
2.343
2,343
2.2
2:776
Baselines. We use data mixtures with weights determined by grid search as baselines. Specifically, we perform a grid search over the set (0.125, 0.25, 0.75) retaining only combinations where the proportions sum 1, This results in total of 21 valid mixtures for comparison.
2.0
Quen25-0..5 SB
Qhwrn2 5-1 1B
Liema32-30
Liamedh 1-88
Figure 4. Comparison of model's PPL with the best PPL in grid search using a data budget of 20M tokens.
We conduct a detailed investigation into PPL for each do- main (all statistics in Appendix E). We compare domain PPL with the best domain PPL across all 21 grid search mixtures. Most often, the best domain PPL is achieved when the domain's 's data is dominant. For Llama 3. 1-B and Llama 3.2-3B, the differences domain PPL. are not signif- icantly greater than each other, supporting that our method can prevent any domain from underperforming. However, for the Qwen model, the IF domain shows relatively larger deficit compared to the best IF loss in grid search than the math and code domains. This indicates that our weighting scheme, which upweights math and code relative to the best IF weights that prioritize IF, significantly affects IF domain performance. This likely because IF is more challenging to learn for small models evidenced by higher domain
Results. In Figure 3, present surface plots of PPL for all 21 mixtures identified through grid search. Across different data budgets (5M, 20M, and 200M tokens) for the same model, we observe that the optimal weights are scale-dependent, i.e., the best weight configurations vary according to the size the data budget. In addition, the weights estimated by our method (highlighted with a yel- low star) effectively locate regions that achieve satisfactory, though not always optimal, performance. Training models with our weights further demonstrates that they are near- optimal. Specifically, as shown in Figure 4 (complete results in Appendix E), our model's PPL almost on par with the best PPL by models in grid search. On average, across models and data budgets, our method's PPL. only 0.66%
6

---

Data Mixing Optimization
PPL than other domains), therefore requiring more data.
Table 3. Per-domain validation loss (PPT.) for reweighting Tulu3
collection.
Table Comparison of our model's s domain PPL with the best domain PPI. in grid search.
Weight
General
Knowledge
Math
Code
Precise IF
Safeły
Avg.
Llama3.
Original
3,589
1.762
1.298
1.691
2.280
2.355
2.162
Domain
Model
Scale
Equal_T
4,955
2.310
1.425
1.999
3.545
2.670
2.817
IF
Math
Code
Equal_1
4.638
2,183
1.492
1.949
3.207
3.171
2.773
5M
0,054
0,024
0.005
Ours
3.422
1.629
1.421
1.615
2.581
2.235
2.150
Qwen2.5-0058
20M
0.122
0.036
0.009
Llamad. 1-708
200M
0.212
0.043
0.035
Original
3.527
1.880
1.290
1.681
2.167
2.214
2.127
5M
0.021
0.004
-0.003
Equal T
3.744
2,040
1.371
1.721
2.274
2.312
2.244
Qwen2.5-1.5B
20M
0.061
0.019
0.020
Equal_
3,553
1.943
1.404
1.736
2.254
2,380
2.212
200M
0.061
0.045
0.014
Ours
3.409
1.774
1.374
1,638
2.056
2.231
2,080
Qwen2.5-326
SM
0.065
0.049
0.040
1lamm3.2-38
20M
0.053
0.051
0.035
Original
3.701
1.673
1.262
1.644
2.276
2.433
2.165
Equal T
3,844
1.734
1.294
1.645
2.413
2.546
2.246
200M
0.054
0.053
0.039
Equal_1 3.891
1.760
1.268
1.665
2.390
2.491
2.244
SM
0,045
0,036
0.035
Ours 3.891
1.672
1.292
1.654
2.047
2.239
2.133
L-ama3.1-8B
20M
0.039
0.032
0.054
200M
0.049
0.042
0.066
as discussed in Section 4. In contrast, our method assigns a relatively smaller proportion to math and allocates more weight to domains that experience more significant de- crease in loss when scaling the same amount of data. We defer the discussion of validation loss results for the Orea collection to Appendix E, as it conveys similar message.
5.2. Re- Weighting SFT Collections
Settings. We further examine our method by re- re-weighting two popular SFT collections: Tulu3 and Orca. These collec- tions reflect different conceptions of domains. Tulu3 iden- tifies "skills" for LLMs and generates/collects data based on these skills. We use six of their identified skills as do- mains: general, knowledge recall, math, code, precise IF, and safety. Orca samples from and combines prior SFT collections, including TO, FLAN, NIV, and COT. We fol- low their categorization and treat the prevalence of these datasets as domains. For domains requiring more data than originally available, we use repeated sampling. Details of data processing are shown in Appendix D.1.
5.3. Performance on Downstream Tasks
We have seen that our method achieves remarkable perfor- mance in decreasing validation loss. We examine whether this translate to downstream tasks with specific metrics for downstream tasks.
Settings. For Tulu3 collection, we also evaluation on six general tasks, each correcponds identified skills AGIEval (Zhong et al., 2023), IFEval (Zhou et al., 2023),
Baselines.
We compare our method with three baselines: Original: The selected instances maintain the same token proportions as the original dataset.
MMLU (Hendrycks et al.), GSM8K (Cobbe et al., 2021), HumanEval (Chen et al., 2021), and safety (averaged over ToxiGen (Hartvigsen et al., 2022) and TruthfulQA (Lin et al., 2022)). For Orca collection, we evaluate on AGIEval (Zhong et al., 2023), HellaSwag (Zellers et al., 2019), safety (averaged over ToxiGen (Hartvigsen et al., 2022) and Truth- fulQA (Lin et al., 2022)), which does not necessarily corre- spond to domains categorization of Orca, but it as significant overlap with the original paper (Mukherjee et al., 2023).
Equal_T: All domains are weighted equally in tokens, i.c., instances sampled to ensure an equal token count from each domain.
Equal I: All domains are weighted equally with regard to the number of items.
Results. In Table 3, we present the domain PPL and aver- age PPL of our method and baselines on held-out datasets. We find that our method achieves the best performance in reducing validation loss and maintains relatively low per- domain PPL across all domains, often being the lowest among all baselines. Through examining per-domain PPL, we gain insight into our method: the original Tulu3 col- lection assigns& a dominant proportion of around 50% to math data, resulting in the lowest PPL in the math domain. However, this allocation is not the most efficient due to the diminishing returns of increasing data single domain,
Results. We present the performance of re-weighted Tulu3 collection on downstream tasks in Table 4. We ob- serve that our method exhibits impressive performance: for Llama3.1 1-8B and 1lama3.1-700 training with our weights leads to the best overall results, while for Qwen2.5-32B, it only underperforms compared to the best-performing Equal_ I weight by a small margin. Additionally, we ob- serve that our method does not cause any domain to perform significantly worse. One exception the IFEval results for
7

---

Data Mixing Optimization
Table 4. Downstream performance for different data mixing methods for re-weighting Tulu3 collection.
Model
Weight
AGIEval
IFEval
MMLU
GSM8K
HumanEval
Safety
Average
Original
33.9
54.5
56.4
65.3
47.0
43.9
50.1
Equal_T
32.1
51.9
55.8
65.2
48.2
41.4
49.1
Llama3.1-8D
E
32.8
42.7
54,9
61.9
50.0
43.9
47.7
Ours
36.1
48.2
56.6
62.3
53.7
45.3
50.4
Original
47.7
65.3
75,4
86.1
44.5
47.4
61.1
Equal_T
52.4
46.7
76.8
85.5
71.3
49.0
63.6
Llama3.1-708
Equal_
51.4
45.8
76.9
83.9
51.2
51.8
60.2
Ours
53.8
51,2
76.5
80.4
68.3
52.7
63.8
Original
49.8
57.7
80.2
89.5
55.5
53.5
64.4
EquaL_
55.8
61.2
81.0
89.8
61.6
56.4
67.6
Qwen2,5-32B
Equal
57.4
64.3
80.9
91.4
62.8
54.3
68.5
Ours
58.3
61.3
81.0
90.8
60.5
56.1
68.0
Llama models, which are at least 6% lower than those with the original weights. Comparing weights of our method and original weights (full analysis in Appendix E.3), we find that that we assign more weights to domains like precise IF compared to the original. However, the total data for Tulu3 remains fixed; consequently, we have repeatedly sample precise IF data multiple times, which might lead models to overfit to this domain. This might explains why our model performs less well tasks like FFval. However, we spec- ulate that with more diverse data for the domain through synthetic data generation, our method could achieve better performance.
cialtics. This is achieved by introducing domain-specifer weighting factors, Yi transforming the objective into min- L(0*(NN w). Dyal). To illustrate this,
imizing Σ we present a case study on training a medical chatbot, where validation dataset only consists of medical domain.
Settings. We mix data from two distinct domains: general instruction following (IF) and medical. General IF data sampled from Alpaca- GPT4c and medical data sampled from PubMedQA (Jin et al., 2019) 3. The validation and testing datasets are derived from the validation and testing sets of PubMedQA. We set the data budget to 10M tokens. All other experimental settings, including unit sample size for each domain and perturbation ratios, remain consistent with previous experiments.
We show results for re-weighting the Orea collection in Appendix E.4. provides different story compared the Tulu3 collection: downstream performance does not correlate with validation loss, and all methods perform sim- ilarly in downstream averages. The best-performing data weights on the held-out validation dataset do not neces- sarily transform to superior downstream performance here. This is not a surprise, since the evaluation can be consid- ered out-of-distribution and dissimilar our held-out dataset. Additionally, because the training domains de- fined by their data sources-for example, the FLAN domain internally includes various tasks and data distributions-ie is difficult to determine which portion of the training data will enhance downstream performance. This emphasizes the importance of ensuring that the validation set closely matches the actual test settings. A good way achieve this, in the Tulu3 collection, is to clearly identify domains tasks/skills, making it easier to collect similar data.
Results. Through perturbation experiments and optimiza- tion of Equation (6), we find that the optimal weights are 67.7% for Alpaca-GPT4 and 32.3% for PubMedQA. This result somewhat unexpected, as it suggests that when the validation dataset is purely medical, our method recom- mends a training mixture with more general IF data than medical data to achieve optimal performance medical tasks. This finding aligns with the cocktail effect (Brief et al., 2024): fine-tuning exclusively on the target task is not always the most effective strategy, Instead, inclusion of different domains-where models are trained on a mixture of related tasks-can enhance performance.
Table 5. Comparison of validation loss (PPI.) and accuracy be- tween our mixture and all medical
Metrie
Our Mixture
Medical Data
5.4. Extension: Training Domain-Specifie LLMs.
PPL (1)
5.42
5.48
Our formulation in Equation (6) assumes equal weight- ing across all domains when calculating validation loss, aligning with our objective to train general-purpose mod- els. However, this formulation can be extended to guide data mixing for training models with one or several spe-
Accuracy (1)
76.8%
76.4%
In Table 5, we compare mixture with weights determined
8

---

Data Mixing Optimization
by our method (67.7% IF data and 32.3% medical QA data) against dataset of the same size composed exclusively of medical QA data. Our mixture outperforms training solely on medical data in both validation loss and downstream performance. We explain this by arguing that general IF data enhance LLMs' language understanding and response generation abilities, which are critical for answering med- ical questions effectively. This case study provides strong evidence that our method offers practical guidance for data mixing during the SFT stage, particularly in selecting unin- tuitive yet effective data combinations.
References
Albalak, A., Pan, L., Raffel, C., and Wang, W. Y. Efficient online data mixing for language model pre-training. In RO- FoMo: Robustness of Few-shot and sero-shot Learning in Large Foundation Models, 2023.
BAAL Infinity-instruet, Aug 2024.
URL
https://huggingg www.gggginnnnn.goooooomom co Stlatssess/BAAT Infinity-Inetn Instruct..
Bach, S., Sanh, V., Yong, Z. X., Webson, A., Raffel, C., Nayak, N. V., Sharma, A., Kim, T., Bari, M. Févry, T., al. Promptsource: An integrated development environ- ment and repository for natural language prompts. In Pro- ceedings of the 60th Annual Meeting of the Association for Computational Linguisties: System Demonstrations, pp. 93-104, 2022.
6. Discussion
What is a Domain? We follow Xie et al. (2024) and define domains based on data provenance in our experiments. For the SFT stage, this definition is nuanced: some dataset may consist of single task, while others contain multiple tasks with diverse distributions. We do not attempt to differentiate these variations because our method, from loss perspective, optimizes domain weights irrespective of what each domain L.e it is agnostic to the domain's 's composition. However, as demonstrated in Section 5.3, clearly defining domains by tasks (e.g., math, code) facilitates the selection of appropriate validation datasets for downstream tasks and streamlines follow-up data generationncooleccion to further enhance model capabilities.
Boyd, S. and Vandenberghe, L. Convex optimization. Cam- bridge university press, 2004.
Brief, M., Ovadia, O., Shenderovitz, G., Yoash, N. B., Lem- berg, R., and Sheetrit, E. Mixing it up: The cocktail effect of multi-task fine-tuning on Ilm performance- a case study in finance. arXiv preprint arXiv:2410.0/ 109, 2024.
Chen, M., Tworek, J., Jun, H., Yuan, Q., de Oliveira Pinto, H. P., Kaplan, J.. Edwards, H., Burda, Y., Joseph, N., Brockman, G., Ray, A., Puri, R., Krueger, G., Petrov, M., Khlaaf, H., Sastry, G., Mishkin, P., Chan, B., Gray, S., Ryder, N., Pavlov, M., Power, A., Kaiser, L., Bavar- ian, M., Winter, C.,, Tillet, P., Such, F, P., Cummings, D., Plappert, M., Chantzis, F., Barnes, E.,, Herbert-Voss, A., Guss, W. H., Nichol, A., Paino, A., Tezak, N., Tang, J., Babuschkin, I., Balaji. S.. Jain, S., Saunders, W. Hesse, C., Carr, A. N., Leike, J., Achiam, J., Misra, V., Morikawa, E., Radford, A., Knight, M., Brundage, M., Murati, M., Mayer, K., Welinder, P., McGrew, B., Amodei, D., McCandlish, S., Sutskever, I., and Zaremba, W. Evaluating large language models trained on code. 2021.
Future Work. Our experiments suggest a promising di- rection for synergizing our method with synthetic data gen- eration. We speculate that generating new data for domains, similar to Toshniwal et al. (2024), rather than repeatedly upsampling existing data as in our experiments, can fully reveal the potential of our approach. Another intriguing direction is extend our parameterization of validation loss to directly model downstream performance (Isik et al., 2024; Chen et al., 2024b), thereby optimizing weights for targeted tasks.
Conclusion. In this paper, we introduce a novel method to optimize data mixtures for SFT of LLMs. Our method parametrizes loss as a function of domain weights while simultaneously modeling the scaling effects of fine-tuning and the interplay between domains. This design enables our method to derive optimal weights for any data scale and ensure that no domain underperforms. Extensive ex- perimental results validate the effectiveness of our method. Furthermore, we demonstrate that this method offers valu- able insights and practical guidance for SFT.
Chen, M. F., Hu, M. Y., Lourie, N., Cho, K., and Ré, C. Aioli: A unified optimization framework for language model data mixing. arXiv preprint arX1v:2411.05735, 2024a.
Chen, Y., Huang, B., Gao, Y., Wang, Z., Yang. J., and Ji, H. Scaling laws for predicting downstream performance in 1lms. arXiv preprint arXiv:2410.08527,, 2024b.
Impact Statement
This paper presents work whose goal is to advance the field of training LLMs. There are many potential societal conse- quences of our work, none which we feel must be specifi- cally highlighted here.
Chung, H. W., Hou, L., Longpre, S., Zoph, B., Tay, Y., Fe- dus, W., Li, Y., Wang, X., Dehghani, M., Brahma, S., et al. Scaling instruction eestruttionfineuuueed language models. Journal of Machine Learning Research, 25(70): 2024.
9

---

Data Mixing Optimization
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint 2t..x::2110011440 2021.
Isik, B., Ponomareva, N., Hazimeh, H., Paparas, D., Vas- silvitskii, S., and Koyejo, S. Scaling laws for down- stream task performance of large language models. arXiv preprint arXiv:2402.04177, 2024.
Conn, A. R., Gould, N. and Toint, P. L. Trust region methods. SIAM, 2000.
Ivison, H., Wang, Y., Pyatkin, V., Lambert, N., Peters, M., Dasigi, P., Jang. J., Wadden, D., Smith, N. A., Beltagy, I., et al. Camels in changing climate: Enhancing Im adaptation with tulu 2. arXiv preprint arXiv:2311.1070. 2023.
Dong, G., Yuan, H., Lu, K., Li, C., Xue, M., Liu, D., Wang, W., Yuan, Z., Zhou, C., and Zhou, J. How abilities in large language models are affected by supervised fine-tuning data composition. arXiv preprint arXiv:2310.0523 2023.
Jin, Q., Dhingra, B., Liu, Z.,, Cohen, W., and Lu, X. Pub- MedQA: dataset for biomedical research question an- swering. In Inui, K., Jiang, J., Ng. V., and Wan, X. (eds.), Proceedings of the 2019 Conference on Empir- ical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) pp. 2567-2577, Hong Kong, China, November 2019. Association for Compu- tational Linguistics. doi: 10. 10.18653/v1/D19.11 URL https: ftttss://aclarrtholl..oo wwtphhoooooo..ooo11-//111110
Fan, S., Pagliardini, M., and Jaggi, M. DOGE: domain reweighting with generalization estimation. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL htttps:/ http:/rwwwwieeeeefeorreeree id-7rf26bM2q44
Hartvigsen, T., Gabriel, S., Palangi, H., Sap, M., Ray, D., and Kamar, E. Toxigen: A large-scale machine-generated dataset for adversarial and implicit hate speech detection. In Proceedings of the 60th Annual Meeting of the Asso- ciation for Computational Linguistics (Volume I: Long Papers), pp. 3309-3326, 2022.
Kang. F., Sun, Y., Wen, B., Chen, S., Song, D., Mah- mood, R., and Jia, R. Autoscale: Automatic prediction compute-optimal data composition for training Ilms. arXiv preprint arXiv:2407.20177 2024.
Lambert, N., Morrison, J., Pyatkin, V., Huang, S., Ivison, H., Brahman, F., Miranda, L. J. V., Liu, A., Dziri, N., Lyu, S., et al. T\" ulu 3: Pushing frontiers in open language model post-training. arXiv preprint 2rXiv:2411.15112 2024.
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding, In International Conference on Learning Representations.
Lin, S., Hilton, J., and Evans, O. Truthfulqa: Measuring how models mimic human falsehoods. In Proceedings of the 60th Anmual Meeting of the Association for Compu- tational Linguistics (Volume I: Long Papers), pp. 3214 3252,20 2022.
Hernandez, D., Kaplan, J., Henighan, T., and McCan- dlish, S. Scaling laws for transfer. arXiv preprint arXiv:2102.01293, 2021.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Vinyals, O., Rae, W., and Sifre, L. Training compute-optimal large lan- guage models. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS 22, Red Hook, NY, USA, 2024. Curran Associates Inc. ISBN 9781 1713870088.
Liu, Q., Zheng, X., Muennighoff, N., Zeng, G., Dou, L., Pang. T., Jiang, J., and Lin, M. Regmix: Data mixture as regression for language model pre-training. arXiv preprint arXiv:2407.01499 2024.
Longpre, S., Hou, L., Vu, T., Webson, A., Chung, H. W. Tay, Y., Zhou, D., Le, Q V., Zoph, B., Wei, J., et al. The flan collection: Designing data and methods for effec- tive instruction tuning. In International Conference on Machine Learning, pp. 12664-2 PMLR, 2023.
Huang. S., Cheng, T., Liu, J. K., Hao, J., Song, L., Xu, Y., Yang, J., Liu, J., Zhang, C., Chai, L., al. Openeoder: The open cookbook for top-tier code large language mod- els. arXiv preprint arXiv:2411.000 1104905, 2024.
Mukherjee, S., Mitra, A., Jawahar, G., Agarwal, S., Palangi. H., and Awadallah, A. Orea: Progressive learning from complex explanation traces of gpt-4. arXiv preprint arXiv:2306.02772 2023.
Huber, P. J. Robust estimation of a location parameter. The Annals of Mathematical Statisties, 35(1):73-101, 1964.
Nocedal, J. and Wright, S. J. Quadratic programming. Nu- merical optimization, pp. 448 492, 2006.
10

---

Data Mixing Optimization
Peng, B., Li, C., He, P., Galley, M., and Gao, Instruc- tion tuning with gpt-4, arKiv preprint arXiv:330..03277, 2023.
Wang, Y., Ivison, H., Dasigi, P., Hessel, J., Khot, T., Chandu, K., Wadden, D., MacMillan, K., Smith, N. A., Beltagy, I., et al. How far can camels go? exploring the state of instruction tuning on open resources. Advances in Neural Information Processing Systems, 36:74764-74787 2023.
Puri, R. S., Mishra, S., Parmar, M., and Baral, C. How
many data samples is an additional instruction worth? In Vlachos, A. and Augenstein, I. (eds.), Findings of the As- sociation for Computational Linguistics: EACL 2023, pp- 1042-1057, Dubrovnik, Croatia, May 2023. Association for Computational Linguistics. doi: 10. 10.18633/1/2023. findings-cacl.77. URL https: https///ccaaattooooao.oo / org/2023. findingsseaccl.....
Wei, J., Bosma, M., Zhao, V., Guu, K.. Yu, A. W. Lester, B., Du, N., Dai, A. M., and Le, Q. V. Finetuned language models are zero-shot learners. In International Confer- ence on Learning Representations, 2022. https: / //opereview ew. net/ forum? (ttmmmiiigg22rGroee0e0
Wu, E., Wu, K., and Zou. J. Finetunebench: How well do commercial fine-tuning apis infuse knowledge into Ilms? arXiv preprint arXiv:241 2.0.0059 2024.
Raghavendra, M., Nath, V., and Hendryx, S. Revisiting the superficial alignment hypothesis. arXiv preprint arXiv:2410.037/1 2024.
Xie, S. M., Pham, H., Dong, X., Du, N., Liu, H., Lu, Y., Liang, P. S., Le, Q. V., Ma, T., and Yu, A. W. Doremi: Optimizing data mixtures speeds up language model pre- training. Advances Neural Information Processing Systems, 36, 2024.
Renduchintala, H. S. V. N. S. K., Bhatia, S., and Ramakr- ishnan, G. SMART: Submodular data mixture strat- egy for instruction tuning. In Ku, L.-W., Martins, A., and Srikumar, V. (eds.), Findings of the Association for Computational Linguistics: ACL 2024, pp. 12916- 12934, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10. 1/.28 findings-acl. ac1.666. URL https: /aclaataoonnooo..... org/ 2024.022222
Ye, J., Liu, P., Sun, T., Zhou, Y., Zhan, J., and Qiu, X. Data mixing laws: Optimizing data mixtures by pre- dicting language modeling performance. arXiv preprint arXiv:2403. 16952, 2024.
acldingg-aal..1e
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Asso- ciation for Computational Linguistics, pp. 4791-4800 2019.
Rudin, W. et al. Principles of mathematical analysis, vol- ume 3. McGraw-hil New York, 1964.
Teknium. Openhermes 2.5: An open dataset of synthetic data for generalist I1m assistants, 2023. URL https: https://huggingface co/datasets/ tekniumm/ppenlermess
Zhang, B., Liu. Z., Cherry, C., and Firat, O. When scaling meets LLM finetuning: The effect of data, model and finetuning method. In The Twelfth International Confer- ence Learning Representations, 2024. URL https: // //openreview.ret wwweteeeettuunne forum?icc 5tC.ddeddss.....ss.....
Toshniwal, S., Du, W., Moshkov, I., Kisacanin, B., Ayrapetyan, A., and Gitman, Openmathinstruct-22 Accelerating for math with massive open-source in- struction data. arXiv preprint 4rXiv:2410001560, 2024.
Zhao, H., Andriushchenko. M., Croce, F., and Flammarion, N. Long is more for alignment: A simple but tough-to- beat baseline for instruction fine-tuning. In Forty-first International Conference on Machine Learning.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv:2302. 2023.
Zhao, H., Du, L., Ju, Y., Wu, C., and Pan, T. Beyond iid: Optimizing instruction learning from the perspective of instruction interaction and dependency. 2024. URL https: / /aaxivivivv./////i//.a/////aatvv org/ ha//abb///555 2ttp//499..077 07045.
Wächter, A. and Biegler, L. T. On the implementation of ontrrior-point filter line-search algorithm for large-scale nonlinear programming. Mathematical programming, 106:25-57, 2006.
Zhong. W., Cui, R., Guo, Y., Liang, Y., Lu, S., Wang, Y., Saied, A., Chen, W., and Duan, N. Agieval: A human-centric benchmark for evaluating foundation mod- els. arXiv preprint S.VV:2304.006666 2023.
Wang, Y., Mishra, S., Alipoormolabashi, P., Kordi, Y., Mirzaei, A., Naik, A., Ashok, A., Dhanasekaran, A. S., Arunkumar, A., Stap, D., et al. Sunprrnaturaliistructiosss Generalization via declarative instructions on 1600+ nlp tasks. In Proceedings of the 2022 Conference on Em- pirical Methods in Natural Language Processing, pp. 5085-5109 2022.
Zhou, C.,, Liu, P., Xu, P., Iyer, S., Sun, J., Mao, Y., Ma. X., Efrat, A., Yu, P., Yu, L.,, et al. Lima: Less is more for alignment. Advances in Neural Information Processing Systems, 36, 2024.

---

Data Mixing Optimization
Zhou, J., Lu, T., Míshra, S., Brahma, S., Basu, S., Luan, Y., Zhou, D., and Hou, L.. Instruction-folloiinng evaluation for large language models. arXiv preprint aTXIv:2311.071112222. 2023.
12

---

Data Mixing Optimization
A. Related Work
Data Mixing. The composition of training data in LLMs is an emerging research field because of its significant impact on performance. Most studies on data mixing, or known as domain reweighting. focus on the pre-training stage. Data mixtures were typically determined heuristically or through ablation studies, such as upweighting domains with higher-quality text (Touvron et al., 2023) those distilled from advanced LLMs (Lambert et al., 2024). However, these intuitively determined weights often rely on practitioners' expertise and may not be optimal (Albalak et al., 2023). DoReMi (Xie et al., 2024) pioneered a systematic approach to data mixing by formulating it as group distributionally robust optimization, using a small proxy model to minimize excess loss and applying the resulting weights to larger models. Similarly, DoGE (Fan et al., 2024) enhances generalization by determining domain weights through optimizing gradient alignment between domains. Other methods like Data mixing laws (Ye et al., 2024), RegMix (Liu et al., 2024), and Aioli (Chen et al., 2024a) estimate the loss as function of domain weights and interpolate performance to identify the optimal compositions.
These methods have underlying assumptions that the optimal weights are model or scale-invan invariant. However, weights optimized for smaller models or scales often fail to generalize to larger ones or different data distributions (Albalak et al., 2023). This issue is particularly pronounced in SFT, where models may have various pre-training data distributions. To address this, AutoScale (Kang et al., 2024) extrapolates optimal weights from small to larger scales, ensuring scale-dependent domain weighting. However, AutoScale can amplify minor weight differences, resulting in extreme or unbalanced domain weights. This may lead to the underrepresentation of certain domains, reducing generalization and effectiveness in SFT.
To the best of our knowledge, SMART (Renduchintala et al., 2024) is the only method that shares the same objective with our work to optimize domain weights for supervised fine-tuning. It uses submodular functions based on prompt embedding to determine the domain importance. However, the use of prompt embedding in SMART indicates its assumption that the optimal weights are model and scale-invariant. In contrast, our method suggests to determine optimal weights dependent on the model and scale for SFT.
Supervised Fine-Tunning SFT, also known as instruction fine-tuning, refers to training pretrained LLMs with instruction- output pairs (Wei et al., 2022; Chung et al., 2024; Peng et al., 2023). SFT data is sourced from diverse domains, such as math, code, and dialogue. Wei et al. (2022) first introduced a large-scale SFT collection called FLAN, which aggregates 62 tasks in natural language inference and commonsense reasoning. Chung et al. (2024) expanded FLAN collection to over 1800+ tasks by incorporating datasets like Super-Natral Instructions (Wang et al., 2022) and PromptSource (Bach et al., 2022). Other notable open fine-tuning collections include Tulu series (Wang et al., 2023; Ivison et al., 2023; Lambert al., 2024), OpenHermes (Teknium, 2023), Infinity Instruct (Zhao et al., 2024), and Orca Mukherjee al. (2023), where aggregate data from various domains to enhance model capabilities.
In addition to SFT collections, current research investigates mechanisms and dataset considerations for training general- purpose LLMs. For example, Zhou et al. (2024) proposed Superficial Alignment Hypothesis, arguing that an LLM's capabilities are almost entirely developed during pre- pre-training stage, and SFT is for format/style alignment. They demonstrated that carefully curated 1, 1,000 instruction- resioose pairs can achieve effective instruction following abilities, with implications that the diversity and quality of the data outweigh quantity. However, Raghavendra al. (2024) challenged this hypothesis. They contended that the improvement from SFT follows a power law; thus, SFT enhances more than formatting or style alignment, such as knowledge infusion. The importance of data diversity is widely acknowledged. supported by Puri et al. (2023) by showing an additional variant of instruction worthies many data samples. Moreover, Chung et al. (2024) and Longpre et al. (2023) analyzed factors like the number of tasks, the size of the models, and instruction settings (e.g.. few-shot, thain-of-the -houhht, and they found that these factors can have significant impact on the performance of LLMs. Zhao et al. suggested that lengthy outputs, which contain more learnable information and are harder to overfit, contribute to the success of SFT. Our work diverges by focusing not on dataset quality but on optimizing data mixtures within a given data budget and across various datasets.
B. Proof: Convexity of Objective Function
In this section, we analyze the convexity of the objective function. Define
f(w) = (w N + KN - w N)
13

---

Data Mixing Optimization
where a0 (0, B 0. O, and N O. We further define
x(w) w N ks k w N
and let
h(x) = =
Thus, original function is ((w) h (x((w) [a(w)]] -B We aim show that (w) convex in w.
We first examine the concavity of z(w). Observe that æ (w) can be rewritten as z=s) = - & (N s) by letting = w N, which lies in [0, .J. We compute the first and second derivatives of zss) with respect to
z' (s)
k a (N 5)0-1
<"(s) = k a(a 1) (N - s)
Since (0, we have 1- < 0 and thus a(a 1) < 0. Consequently, x" z²(s)
indicating that x(s) is concave in Because is affine transformation of w, the function x(w) is also concave in E [0. 1].
Next, we check the convexity and monotonicity of h(x) = x for x >> and B > 0. The second derivative of h is
(x) = B(B + x 8-2 0 for > 0, B > 0,
so x is convex. Furthermore,
1'(x) = -B x -8-1
B-1 < 0
which shows h is decreasing on (0, oo).
We now apply a composition rule from convex analysis (Boyd Vandenberghe, 2004):
Lemma B.1 (Composition Rule for Convexity). If g is concave and convex and con-increasing on the relevant domain,
then h g is convex.
In our case, aw) is concave in w, and h(x)) is convex and non-increasing in a. Therefore, the composition
(ww h<((w)
is convex in w. To generalize, consider summing over all domains
= L....,K: K:
K
K Σ 1=1
F(w)= = EL(w)
(w; N (N w; N)a)) -8 + E.l.
under the constraints w; 2 0 and Σ my = For each a; E (0, 1) and 3 > 0, then each summand L:(w) is convex in w,, Since a sum of convex functions is also convex, the overall objective remains convex w. Hence one can find a global minimum of F(w) using standard convex optimization methods. We adopt sequential least squares programming (SLSQP).
C. Proof: No Domains Left Behind
We show that our objective function produces a balanced allocation of domain data without explicitly imposing domain- specific constraints. For convenience, recall the following optimization problem:
K
C (w:N + k(NN -w,N)a) 3 + E,
min
subject to
0 5 wi 5 =i=1 w1=,.... -....,K. K,
W: 1.
14

---

Data Mixing Optimization
We rewrite this problem in standard constrained form follows:
K
min wERN
F(w) Σ
Li(w)
Σ
CCwwNN+hiC wiN)aa) - + Eb
K
subject to
g(w) = Σ
Wi: 1 1=0,
1=1
h;(w) = w S 0, == 1, I,....AA K.
Since F'w) is convex (proof in Appendix B) and the set Âκ = {w E R E wy = 1} is compact, the Weierstrass extreme value theorem (Rudin al., 1964) guarantees a unique global minimizer w E AA.
The Lagrangian is
K
C(w, X, (m+3)= F(w) + A
Σ wi - 1)
Σ
My Wi:
where À E R and (4g 2 O. By the
(KKT) conditions, any optimal w must satisfy
K
VF(w") + +
Σ 1g
Vhu(w") = O,
with complementary slackness /4, (-w;) O, and i 0, M 2 E, t 1.
Suppose that for all i, we have 0 < wf 1. Then - is strictly negative, so μ = 0 by complementary slackness. Since Vg(w*)) (I...., 1), the KKT condition becomes
+ +
0,
which implies
OF Ow
OF Ow2
OF Owk (w")
(w*)
(w*)
Because F(w) ((w)=Ek L((w), each partial derivative satisfies
OF
K
θLg (w").
(w*) =
Ow;
Owi
In our objective function, (w)) could have partial dependence w; indirectly if (N wy WyyN N) changes. However, each whenever * j. Consequently, if some domain has
domain is governed mostly by w, That is, large negative gradient for small w, the optimizer increases W/ until partial derivatives balance across domains.
aL,
OL
Thet
For boundary points where some w 0 or w 1, complementary slackness dictates how the multipliers μ4 adjust the stationarity condition. If w 0. then from stationarity.
OF
(w") u4 + X 0, ow;
and -w is active, so n 2 Domains j with 0 < w < have n1 = O. and thus
OF Owy
+ X" = 0, (w")
implying
0F
OF
Ow; (w") = X² + 1 2
(w") -X"
ow;
15

---

Data Mixing Optimization
Hence, if : O, domain i cannot possess strictly more negative gradient than the domains with positive allocation; otherwise, it would have been optimal to increase w An analogous argument covers the W 1, showing that domain saturates only if it truly yields the best marginal gain up that boundary. Because each L; typically has large negative derivative when w; is small (and milder derivative when W; is large), the solution cannot assign zero weight to a domain that still has significantly larger marginal return than others. Similarly, in the interior case, equality of partial derivatives forces additional resources toward domains that can quickly reduce their loss, so domains with high marginal benefit do not remain under-allocated. In summary, we present the following theorem.
Theorem C.1 (No Domains Left Behind). The optimization problem
K
mMiM
L (.(W)
WEAK
admits unique global minimizer w AK, and:
(i) wf is in the interior of Ax, no domain is under-allocated if can still yield larger marginal improvement than those domains currently receiving weight.
(ii) (f 0 for some domain i, then the marginal gain of allocating data to is no larger than that of any other domain with w O. Consequently, no domain with strictly higher marginal return is left at zero.
D. Experimental Details
D.1. Data Processing
We provide data processing details for all datasets used in the paper below:
Alpaca-GPT4::
We merged Instruction and Input as the input for the model. Infinity Instruct": We used 625 version. Given that the dataset encompasses both mathematical and code data, we undertook filtering process eliminate items featuring math-related keywords such as "algebra" and "geometry," as well code keywords like "Python" and "Bash." This ensures the dataset remains exclusively tailored for general instruction following.
OpenMaathnnstuuct22: We used train_1M version.
OpenCoder': We used fi filtered infreidiifiiity... iinst version.
Tulu38: We filtered items larger than 4096 tokens to fit our context window. We also filter non-English data and Tulu3 hardcoded data.
Orea²: We merged System.prompt and Questiqo i on as the input for the model. We used 300M tokens of the original dataset.
D.2. Hyperparamters
We use a cosine learning rate scheduler, set the batch size to 256 for training models, and use a sequence length of 4096 tokens. For perturbation experiments, train the model for 3 epochs and select the model with the lowest validation loss. The maximum training steps determined by the data budget follows: 200 steps for 5M budget, 400 steps for 20M, and 2,500 steps for 200M. For the Tulu3 and Orca experiments, set the maximum training steps to 6,000. All experiments are conducted on NVIDIA H100 GPUs. The learning rate for each model is show in Table 6:
D.3. Evaluation Benchmarks
We use LM-Evluation harness 10 (with default settings if not specified) to assess models' performance on downstream tasks.
Details of our downstream benchmarks are discussed below:
"naps://uuggingfaaecc/ultat colaaaseeeiiiiilllllaaaeeeggg "httss//huggingfacecco
hhupp://github. wwt.:lu...ooommuuuaeeee http://hwwa..toomuuooeeee//////n
16

---

Data Mixing Optimization
Table 6. Learming rates for models in the experiments
Model Class
Qwen2.5
Qwen2.5
Qwen2.5
Llama3.2
Llama3.1
Llama3.
Model Size
0.5B
1.5B
32B
3B
8B
70B
Learning Rate
2 x 10-5
2 x 10-6
10-6
10-5
1 10-5
1 10
AGIEval (Zhong al., 2023): AGIEval is a general benchmark. We specific subset of AGIEval tasks that are
multiple-choice and English-only.
IFEVal (Zhou et al., 2023): IFEVal measures precise instruction-followwe ability, We report prompt-level loose accuracy.
MMLU (Hendrycks et al.): MMLU assesses knowledge across a wide range of academic subjects, including mathe- matics, philosophy, and law. We use a zero-shot setting.
GSM8K (Cobbe al., 2021): GSM8K evaluates mathematical problem-solving. We use an 8-shot setting with thain-of-thought prompts.
HumanEval (Chen et al., 2021): HumanEval evaluates code generation capabilities. We use zero-shot setting and
report the pass@1 score.
ToxiGen (Hartvigsen et al., 2022): ToxiGen assesses safety by measuring toxicity levels. We use zero-shot setting
with unnormalized accuracy.
TruthfulQA (Lin et al., 2022): TruthfulQA evaluates hallucination, which we include as a measure of general safety.
We use the test version truthfulqa.mc in zero-shot setting.
HellaSwag (Zellers et al., 2019): HellaSwag tests language understanding and commonsense reasoning. We report unnormalized accuracy using a zero-shot setting.
E. Additional Results
E.1. Further Analysis on Figure
In Figure 3, an intriguing observation that models within the same class exhibit similar performance patterns concerning domain weights. We hypothesize that these models share data distributions despite differences size. Future work could explore the relationship between pre-training and SFT data in this direction. If our hypothesis is validated, this could allow us streamline method by estimating loss parameters using smaller models with similar data distributions and then applying the optimized weights to larger models.
E.2. Comparing Models Trained with Our Weights versus Grid-Search Weights
Table 7 presents the performance of models trained using our weighting method against the best results obtained through grid search. The overall PPL scores show that our method performs very close to the grid search results across all models and data budgets. This suggests that our approach is highly effective in approximating the optimal weights without the extensive computational cost grid search. Notably, the gap between our method and grid search does not significantly expand as the data budget increases, indicating that our weighting strategy scales well with larger datasets.
Per-domain PPL results shown Table 8 further validate the robustness of our method. Across domains including IF, Math, and Code, our weights consistently achieve performance comparable the best grid search results. This demonstrates that method generalizes well across diverse tasks. The consistent performance across scales and domains highlights the efficiency and reliability of our weighting approach, making it a practical guide for data mixing.
E.3. Analysis of Weights
We present weights for each domain determined by baseline methods and our method for Tulu3 experiments and Llama3.1 1- 70B as study case. A key observation is the significant difference between our method and the original weights. The original weights heavily prioritize the math domain (57.77%), while drastically underweighting other domains like precise IF (2.05%) and safety (4.65%). In contrast, our method demonstrates a more balanced weighting. While still emphasizes certain domains like safety (25.25%) and general (19.26%) it significantly reduces the weight for math (12.81% 1)) compared the original weights. This redistribution indicates that our method prioritizes broader and more equitable distribution of
17

---

Data Mixing Optimization
Table Overall PPI. of models trained with weights and with weights in grid search.
Data Budget
Best
Model
Ours
(Grid Search)
5M
2.9847
2.9692
Qwen2.5-0.5B
20M
2.9293
2.8952
200M
2.8210
2.7621
5M
2.3954
2.3980
Qwen2.5-1.5B
20M
2.3631
2.3428
200M
2.3154
2.2931
2.5154
2.5085
1lama3.2-3B
20M
2.3986
2.3993
200M
2.2536
2.2505
5M
2.2343
2.2259
Llama3.1 1-8B
20M
2.1756
2.1633
200M
2.0762
2.0542
Table 8. Per-domain PPL of models trained with weights and with weights in grid search.
IF
Math
Code
Model
Scale
Ours
Best
Ours
Best
Ours
Best
5M
5.228
5.174
1.532
1.508
2.195
2.190
Qwen2.5-0.5B
5.120
4.998
1.506
1.470
2.162
2.153
200M
4.850
4.638
1.467
1.424
2.148
2.113
5M
3.865
3.844
1,409
1.405
1.912
1.915
Qwen2.5-115B
20M
3.796
3.735
1.395
1.376
1.899
1.878
200M
3.661
3.600
1.393
1.348
1.892
1.878
5M
3.734
3.670
1.630
1.581
2.182
2.142
Llama3.2-23
20M
3.543
3.490
1.567
1.516
2.086
2.051
200M
3.309
3.255
1.464
1.410
1.988
1.949
5M
3.186
3.141
1.505
1.469
2.011
1.977
Llama3.1-88
20M
3.110
3.071
1,455
1.423
1.961
1.907
200M
2.919
2.870
1.402
1.360
1.907
1.841
weights across domains, in which we show in Table lead to better generalization and performance across diverse tasks.
Table 9, Weights for each domain determined by baseline methods and our method for Tulu3 experiments using Llama3. 1-70B (%)
Method
General
Knowledge Recall
Math
Code
Precise IF
Safety
Original
16.42
5.77
57.77
13.33
2.05
4.65
Equal T
16.67
16.67
16.67
16.67
16.67
16.67
E
32.99
12.57
24.30
13.65
10.43
6.04
Ours
19.26
16.92
12.81
10.10
15.63
25.25
18

---

Data Mixing Optimization
E.4. Validation Loss and Downstream Performance of Re- Re-Weighting Orca Collection
Table 10 and Table highlight important discrepancy between validation loss PPL and downstream task performance. In Table 10, our method achieves competitive even superior PPL scores compared other weighting strategies, particularly for larger models like Llama3.1-70B and Qwen2.5-32B. However, Table 11 reveals that lower PPL does not consistently translate to better performance on downstream tasks such as AGIEval, HellaSwag, and Safety. For instance, while our method achieves the lowest average PPL for Llama3. 1-70B and Qwen2.5-32B, its downstream performance is not always the best, as seen with Llama3. 1-8B, where the Equal T weighting outperforms ours on average.
This inconsistency can be attributed the distribution difference between the validation dataset and the testing dataset for downstream tasks. The validation dataset may not fully capture the complexity diversity of real-world tasks, leading selection of the model that does not generalize well to downstream applications. Therefore, it indicates that method, while effective in reducing validation loss, requires careful curation of validation dataset that aligns with downstream tasks to translate its powerfulness to downstream applications.
Table 10. Validation loss for reweighting Orca collection.
Table 11. Downstream performance for reweighting Orca collec- tion.
Weight
TO
CoT
Flan
Niv
Avg.
Weight
AGIEval
HellaSwag
Safety
Avg.
Llama3.1-8B
Llama3.1-88
Equal_T
1.782
1.808
1.948
1.645
1.796
Equal_T
38.3
59.5
39.7
45,9
E
1.728
1.813
1,930
1.580
1.763
Equal
39.6
59.2
38.9
45.9
Original
1.726
1.845
1.924
1.665
1.790
Original
39.7
59.5
40.3
46.5
Ours
1.781
1.826
1.949
1.626
1.795
Ours
38.2
59.9
39.7
45.9
Llama3.
1-70B
Llama3. 1-70B
Equal T 1.629
1.679
1.785
1.571
1.666
Equal T
52.6
67.6
49.4
56.5
Equal I 1.608
1.682
1.770
1.515
1.644
Equal_I
53.7
67.8
41.9
54.5
Original
1.600
1.720
1.772
1.581
1.668
Original
53.3
67.2
51.1
57.2
Ours
1.620
1.651
1.767
1.534
1.643
Ours
52.6
67.6
44.3
54.8
Qwen2.5-32B
Qwen2.5-32B
Equal_T
1.643
1.670
1.798
1.553
1,666
Equal_T
59.8
64.4
58.5
60.9
Equal_I
1.626
1.673
1.787
1.511
1,649
Equal I
61.2
64.3
59.9
61.8
Original
1.617
1.699
1.789
1.560
1,666
Original
60.6
64.4
55.9
60.3
Ours
1.613
1.662
1.787
1.533
1.649
Ours
61.7
64.4
54.6
60.3
19


Reviews:
[{"review_id": "HhugGTeR8O", "reviewer": "Reviewer_rbjG", "summary": "This paper proposes an algorithm to optimize proportions of data sources for the supervised fine-tuning (SFT) stage in LLM training by minimizing validation loss. Given that jointly optimizing the mixing weights and model parameters is expensive, the paper makes an approximation. It makes uses of data scaling laws (Zhang et al. 2024) to parameterize the validation loss for each domain and borrows the ideas of 'data transfer' (Hernandez et al 2021) to come up with a 5-parameter regression model to estimate validation loss from each domain. Given this model, the mixture optimization becomes a convex function, which is solved using gradient methods.  The paper shows that the model takes into account interactions between domains and thus assigns an equitable proportion to each domain. Results using two LLM families (Qwen 2.5 and Llama 3.1) show that the method obtains slightly worse performance compared to an exhaustive grid search. The paper also shows that the method can be used to rebalance dataset mixtures in existing SFT datasets.\n\n## update after rebuttal\nIn my initial review, I asked the authors to justify why they needed to perform a 2 stage regression for a 5 parameter model and added pointers to several papers (from statistics literature) which argued why that should not be done. The authors did not answer this question adequately. They mentioned that the single stage regression is necessary only when there are a large number of parameters - this is not the case: A single stage model should always yield the best performance while a 2 stage model will likely yield biased results (e.g. overinflated regression coefficients, overfitting to the training set etc). If the authors want to show that the two stage model is equivalent to a one stage model, they should report results using both models. Without such an empirical justification, I find it hard to trust the results from the paper. Hence, I am not changing my scores after the rebuttal.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you! We will address your comments as follows.\n\n---\nQ:  Validation loss approximation in MSE\n\nA: The MSE of the validation loss estimation in the experiments from Section 5.1 is presented in Table 0.1.1 (https://anonymous.4open.science/r/Data_Mixture-5D78/src/Table0.1.1.png). Additionally, the MSE of our estimation for the re-weighting experiments with Tulu3 and Orca are shown in Table 0.1.2 in the link). While not perfect, our method provides a reasonable estimate of the validation loss.\n\n---\nQ: Low validation losses do not always translate to downstream evaluations.\n\nA: In the re-weighting Orca experiments, we concluded that a low validation loss does not necessarily translate to improved downstream evaluations if the validation datasets are not properly sampled. This is intuitive and applies to many machine learning scenarios, highlighting the importance of careful validation dataset selection. Our key claim is that our method, which optimizes for validation loss, will generally lead to improved downstream performance when validation dataset is properly chosen.\n\nSo now I think the primary concern lies in why the improvement in downstream performance for re-weighting Tulu3 appears marginal (in Table 4). We have noticed this issue and, over the past month, we conducted additional experiments to investigate it. Our responses are as follows: In our paper, we mentioned using repetitive sampling for domains with higher weights than the original dataset, which we speculated to be the cause of the observed downperformance. Taking the domain weights of Llama 3 70B in Tulu3 as an example (shown in Table 0.1.3 in the link)), the \"Precise IF\" domain in the original Tulu3 dataset constitutes about 2%, but in our calculation, we aimed to increase it to approximately 15%. This means that, for the same amount of data, the same item in the \"Precise IF\" domain appears around 7x more frequently than in the original dataset when using repetitive sampling if we keep the training dataset size fixed. We hypothesized that this repetitive sampling strategy might cause unintended model behaviors, potentially explaining the suboptimal performance. To address this, we replaced repetitive sampling with a re-writing method, inspired by [1] and [2]. Re-writing can be intuitively understood as paraphrasing repetitive items: it does not add new information but reduces repetition in the sampled dataset. We then trained the model on the resulting datasets and attached the downstream performance in \"downstream.png\" in the link. From the results, we observed that our method led to more significant improvements in downstream benchmarks, validating the effectiveness of our approach.\n\n---\nQ: why is the 5-parameter regression done in 2 stages\n\nA: Our decision to use a two-stage estimation was driven by empirical and intuitive considerations rather than strict theoretical requirements. We have conducted experiments initially estimating all five parameters in a single stage using the Huber loss and found that the resulting domain weights are very similar to those obtained via a two-stage approach. So in our context, the impact of using one versus two stages might not affect the domain weights significantly if both estimations form similar curves of validation loss. \nWe went through the materials you provided, and we noted that the issues associated with two-stage estimation are most pronounced in models with a large number of variables and when dependencies among explanatory variables are poorly captured. In our method, however, we only have 5 parameters (this might explains why multi-step estimation is not that problematic) and parameters are not treated equally. In Eq. (6), \\beta exerts a much stronger influence on the validation loss when the dataset is large. Since the validation loss (from the trained model) can be noisy, we want to use additional information—specifically, the pairwise slopes across different training runs—to determine \\beta more reliably.\n\n---\nQ：'Data budget' N0 and 'sample size' N?\n\nA: The data budget N0 refers to the total number of tokens allocated for the final training dataset, which we use to optimize the distribution of domains. The sample size N is the number of tokens from each domain used to fit the parameters in Equation 4.\n\n---\nQ: In Figure 3, the PPL of the proposed method (yellow star) is always better than the grid search though this is not the case in the tables presented in Appendix F.\n\nA: The yellow star in Figure 3 represents the points which our method identifies, and are near-optimal rather than always outperforming the grid search. The yellow stars are located in regions that correspond to good PPL values, and they are not necessarily the best PPL achieved by the grid search. Many of the yellow stars are positioned at the boundaries, slightly off the darkest areas, indicating that while our method performs well, it may not always match the best results from the grid search.", "claims_and_evidence": "Some of the claims can be improved:\n* The paper should report how well the validation loss estimates works by reporting mean squared error.", "methods_and_evaluation_criteria": "Some of the methods in the paper need to be improved:\n* Appendix B (P656): For estimating the domain specific validation loss in Equation (4), the paper employs a two-step linear regression in whereby the first stage estimates \\beta and the second step estimates the remaining variables with \\beta fixed. Such a multi stage approach is referred to in statistics as 'stepwise selection' and it is known to suffer from a number of issues e.g. biases in regression coefficients, problems in the presence of multi-collinearity. See references below:\n\nReferences:\n* https://link.springer.com/book/10.1007/978-3-319-19425-7\n* https://hbiostat.org/rmsc/multivar\n* https://journalofbigdata.springeropen.com/articles/10.1186/s40537-018-0143-6", "theoretical_claims": "The proofs seem reasonable.", "experimental_designs_or_analyses": "Yes, they seem reasonable.", "relation_to_broader_scientific_literature": "The paper introduces a new technique to optimize weighting of data sources in SFT. Past work has mostly focussed on mixture optimization in the context of pre-training. The only prior work which has studied this problem for SFT is SMART (Renduchintala et al., 2024) - their approach differs from this work in using prompt embedding and the use of submodular functions.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths:\n* Proposes a method to optimize mixture ratios for SFT\n* Method shows competitive but slightly weaker ratios compared to grid search\n* Method can be used to rebalance mixing ratios in existing SFT datasets\n* The proposed approach yields equitable mixture ratios such that no single dataset dominates the overall mixture\n\nWeaknesses:\n* It would be good to see some results on how well the validation loss approximation using the 5-parameter models holds by reporting MSE.\n* Though the method yields the mixture set with a low validation set loss, this does not always translate to performance improvements in downstream evaluations.\n* Some of the aspects of the method are not explained or justified clearly e.g. why is the 5-parameter regression done in 2 stages i.e. finding \\beta in the 1st stage and the rest of the parameters in the 2nd stage? There is evidence from the multiple regression literature showing that a stepwise selection suffers from multiple issues (See above).", "other_comments_or_suggestions": "* What is the difference between 'data budget' N0 and 'sample size' N?\n* In Figure 3, the PPL of the proposed method (yellow star) is always better than the grid search though this is not the case in the tables presented in Appendix F.", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "fVulcxkqR6", "reviewer": "Reviewer_AAeY", "summary": "This paper proposes a data mixing optimization problem for supervised fine-tuning of LLMs. The problem is based on a scaling law determined by neural scaling laws for SFT. The paper proves convexity properties about the optimization problem and applies it to optimize mixtures for several finetuning settings.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you! We will address your comments as follows.\n\n---\nQ1 : Improvements appear marginal\n\nA1 : We have noticed this issue and, over the past month, we conducted additional experiments to understand it. Our responses are as follows: In our paper, we mentioned using repetitive sampling for domains with higher weights than the original dataset, which we speculated to be the cause of the observed downperformance, not the problem with our method. Taking the domain weights of Llama 3 70B in Tulu3 as an example (shown in Table 0.1.3 https://anonymous.4open.science/r/Data_Mixture-5D78/src/Table0.1.3.png), the \"Precise IF\" domain in the original Tulu3 dataset constitutes about 2%, but in our calculation, we aimed to increase it to approximately 15%. This means that, for the same amount of data, the same item in the \"Precise IF\" domain appears around 7x more frequently than in the original dataset when using repetitive sampling. We hypothesized that this repetitive sampling strategy might cause unintended model behaviors, potentially explaining the suboptimal performance. To address this, we replaced repetitive sampling with a re-writing method, inspired by [1] and [2]. Re-writing can be intuitively understood as paraphrasing repetitive items: it does not add new information but reduces repetition in the sampled dataset. We then trained the model on the resulting datasets and attached the downstream performance in https://anonymous.4open.science/r/Data_Mixture-5D78/src/downstream.png. From the results, we observed that our method led to more significant improvements in downstream benchmarks, validating the effectiveness of our approach.\n\nTo further validate our method, we present the PPL-step plot in https://anonymous.4open.science/r/Data_Mixture-5D78/src/PPL_step.png as you suggested (we use Llama 3.1-8B model as an example, with the original dataset’s weight for comparison). Our method achieves comparable performance with faster computation. Ultimately, our method converged at a lower PPL.\n\nFinally, we want to emphasize that the baselines we are comparing against are already highly optimized (e.g., the best configuration from grid search in Section 5.1, and the original weights in Tulu3, which required extensive manual labor and multiple iterations to fine-tune [3]). Therefore, the advantage of our approach lies not only in performance improvement but also in its full autonomy and relatively lower computational cost, as we will explain further below.\n\n---\nQ2: Computation cost\n\nA2: The computation cost of our approach primarily lies in running on domain samples over N tokens to determine the parameters in Eq6. We compare the computation cost of our method, and cost of running the entire training dataset in (we use deepspeed.profiling.flops_profiler to estimate overall number of floating-point operations (FLOP) as a metric).\n\nFor experiments in Section 5.1 (see https://anonymous.4open.science/r/Data_Mixture-5D78/src/cost_5.1.png), our method costs computation of ~7x compared to training on 5M tokens, ~3x compared to training on 20M tokens, and 0.7x compared to training on 200M tokens. While for grid search, the current granularity runs each experiment for 21 times, which is substantially more.\n \nFor re-weighting Tulu3 experiments (see https://anonymous.4open.science/r/Data_Mixture-5D78/src/cost_tulu3.png), our method costs ~1.5x compared to one training round, which is impressive.\n\nSimilar stories ring true in re-weighting Orca experiments (see https://anonymous.4open.science/r/Data_Mixture-5D78/src/cost_orca.png). \n\nTherefore, our method demonstrates impressive advantages in terms of computation costs compared to traditional manual methods seen in [3], which takes extensive runs to determine the data mixture.\n\n\n[1] Wizardlm: Empowering large language models to follow complex instructions.\n\n[2] Datagen: Unified synthetic dataset generation via large language models.\n\n[3] Tulu 3: Pushing frontiers in open language model post-training.", "claims_and_evidence": "The paper claims to develop a data mixing optimization framework and empirically validate it. It is well-supported.", "methods_and_evaluation_criteria": "The methods and evaluation criteria make sense, and are well tested on a variety of benchmarks.", "theoretical_claims": "I checked the proofs in the Appendix. They follow standard procedures and appear correct.", "experimental_designs_or_analyses": "The experiments appear standard.", "relation_to_broader_scientific_literature": "The paper contributes to the data mixing literature by focusing on data mixing for SFT. It introduces a new method in this space.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n1. The paper is generally easy-to-read with key results flowing intuitively.\n\nWeaknesses:\n1. Improvements in Table 3-4, appear marginal. It is hard to discern the advantage of the proposed method in this case and whether the improvement is potentially random. Typically, we may see a plot of the validation loss over steps to show that the method allows for achieving the same performance at faster computation. Similarly, improvements in Figure 4 show that grid searching is marginally better than the proposed method, at roughly the same level of improvement as the comparison in Table 3-4. Ultimately, it is hard to discern what constitutes a meaningful performance improvement for this SFT data mixing task.\n\n2. Similarly, it would be useful to measure overall compute cost of this approach, from estimating the scaling laws all the way to optimization. Is there a compute advantage to the proposed method? For example, since grid searching over a handful of mixing ratios seems to improve over the proposed method, but surely the proposed method is computationally better than grid searching?", "other_comments_or_suggestions": "Some grammatical errors in Section 3.1 (e.g., Consider \"fine-tuning\")", "questions_for_authors": "See Strengths + Weaknesses", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vHRzKXeP7V", "reviewer": "Reviewer_jCLn", "summary": "This paper addresses the problem of data mixing for supervised fine-tuning (SFT) of large language models (LLMs). The authors frame data mixing as an optimization problem and introduce a method to minimize validation loss. Their approach involves parameterizing the loss by modeling effective data transferred and leveraging scaling laws for fine-tuning. The authors provide mathematical proofs and empirical results to demonstrate the effectiveness of their algorithm.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you! We will address your comments as follows.\n\n--- \nQ: The generalization of the method to different data domains (excluding IF, Math and Code) could be explored further.\n\nA: In Section 5.2, we conducted experiments beyond the three domains (IF, Math, and Code), exploring additional domains such as Safety, Knowledge, and General. We also experimented domains that are not well-defined in the Orca dataset. Furthermore, in Section 5.4 and Appendix F5, we included experiments in the Healthcare domain to further validate the generalizability of our method.\n\n---\nQ: The paper could benefit from a more detailed discussion of the limitations of the proposed method.\n\nA: Thank you for your suggestion! One limitation of our work is that our method parameterizes validation loss. Since we are in the post-training stage and targeting downstream applications, modeling validation loss is less straightforward than directly modeling downstream performance. However, modeling downstream performance is challenging due to the lack of a clear mapping between performance metrics and data size, like what’s in scaling law that connects data size with validation loss. Another limitation is that our method relies on parameter estimation, which can be sensitive at times. A potential solution is to train the model multiple times and compute the mean of validation loss when fitting parameters, though this comes with the tradeoff of increased computational cost.\n\n---\nQ: Computation cost & scalability\n\nSee Q2 and A2 under Reviewer AAeY about computation cost.\n\nOur method demonstrates strong scalability across various data sizes. In Section 5.1, we conducted experiments with 5M, 20M, and 200M tokens. In Section 5.2, the data size was approximately 500M tokens, equivalent to around 1M items. For context, industrial-level models like Qwen 2.5 and DeepSeek V3 typically use 1M–3M items for SFT. Thus, the data sizes in our experiments are quite substantial.\n\n---\nQ: It would be helpful to include a more detailed comparison with other data mixing techniques, highlighting the advantages and disadvantages.\n\nA: While most existing data mixing techniques are applied during the pre-training stage, to our knowledge, SMART is the only method designed for the post-training stage, and we have already incorporated it in our paper. However, SMART has its limitations, as it relies on embeddings of data items and is agnostic to the models used. This is problematic because model performance is significantly influenced by the choice of base models. Other post-training data mixing efforts are either not publicly disclosed or require extensive manual effort.\n\n\n---\nQ: How the proposed method can be extended to guide data selection for domain-specific LLMs? How would the parameterization of validation loss be adapted in such cases?  \n\nA: The procedure for using our method to guide data selection for domain-specific LLMs follows the same approach as in general cases. The only difference is that the validation loss calculation is focused solely on the validation dataset for the specific domain, excluding other domain datasets. As it turns out, domains that contribute little to the validation loss for a specific domain will have very low domain weights.\n\n---\nQ: How sensitive is the performance of the proposed method to the choice of validation dataset? Could you discuss any potential strategies for selecting an appropriate validation dataset to ensure good generalization performance?\n\nA: Good questions! In general, when there is a significant distribution difference between the validation and downstream test sets, the best performance on the validation set (typically used to select the model with lowest validation loss) is unlikely to translate to best downstream performance. This applies to our method as well. Developing effective strategies for selecting an appropriate validation dataset to ensure strong generalization is similar to those for training datasets. Key principles include broad domain coverage, diversity in prompts/templates/questions, and a balanced distribution of difficulty. For a comprehensive list of guidelines, we refer to the Kimi K1.5 technical report [4], which provides extensive criteria for data collection and filtering. Additionally, to enhance the validation dataset, synthetic data generation using LLMs can be employed to create in-distribution data. A good example is OpenMathInstruct-v2 [5], which generates high-quality, in-distribution math instruction data based on existing datasets. Similar approach can be applied to other domains as well.\n\n[4] Kimi k1. 5: Scaling reinforcement learning with llms\n\n[5] Openmathinstruct-2: Accelerating ai for math with massive open-source instruction data", "claims_and_evidence": "The claims made in this submission are generally well-supported by the evidence provided. The authors present both theoretical analyses and experimental results to validate their method.", "methods_and_evaluation_criteria": "The proposed method, Data Mixing Optimization, and the evaluation criteria, including experiments on controlled data mixtures and popular SFT collections, are appropriate for the problem.", "theoretical_claims": "I have checked the correctness of the proofs and the theoretical claims appear to be sound.", "experimental_designs_or_analyses": "I have checked the soundness/validity of the experimental designs and analyses, and they seem to be well-designed and the analyses are appropriate.", "relation_to_broader_scientific_literature": "The key contributions of this paper are related to the broader scientific literature through its focus on optimizing data mixtures for LLM training. It builds upon existing work in data mixing for pre-training and addresses the underexplored area of data mixing for SFT.", "essential_references_not_discussed": "The paper adequately discusses essential references. However, it could benefit from a more in-depth discussion of recent work on scaling laws and their implications for LLM training.", "other_strengths_and_weaknesses": "Strengths:\n\n* The paper introduces a novel and theoretically sound method for data mixing in SFT.   \n\n* It provides a comprehensive set of experiments and results across various scenarios.   \n\n* The paper is well-written and easy to follow.   \n\nWeaknesses:\n\n* The paper could benefit from a more detailed discussion of the limitations of the proposed method.   \n\n* The generalization of the method to different data domains (excluding IF， Math and Code) could be explored further.", "other_comments_or_suggestions": "* Consider adding a section discussing the computational cost and scalability of the proposed method.   \n\n* It would be helpful to include a more detailed comparison with other data mixing techniques, highlighting the advantages and disadvantages.", "questions_for_authors": "* Could you elaborate on how the proposed method can be extended to guide data selection for domain-specific LLMs, as mentioned in the paper? How would the parameterization of validation loss be adapted in such cases?   \n\n* In the experiments, how sensitive is the performance of the proposed method to the choice of validation dataset? Could you discuss any potential strategies for selecting an appropriate validation dataset to ensure good generalization performance?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pq137TnQhZ", "reviewer": "Reviewer_wneA", "summary": "This paper introduces a way to decide the optimal proportion of different domain datasets for LLM fine-tuning. The Scaling law is utilized to learn the optimal proportions, improving fine-tuning performance in various settings.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your feedback! We will discuss your comments as follows.\n\n---\nQ:  Relatively marginal improvement of \"Ours\" compared to \"Original\"\n\nA: We have noticed this issue and, over the past month, we conducted additional experiments to understand it. Our responses are as follows: In our paper, we mentioned using repetitive sampling for domains with higher weights than the original dataset, which we speculated to be the cause of the observed downperformance, not the problem with our method. Taking the domain weights of Llama 3 70B in Tulu3 as an example (shown in Table 0.1.3 https://anonymous.4open.science/r/Data_Mixture-5D78/src/Table0.1.3.png), the \"Precise IF\" domain in the original Tulu3 dataset constitutes about 2%, but in our calculation, we aimed to increase it to approximately 15%. This means that, for the same amount of data, the same item in the \"Precise IF\" domain appears around 7x more frequently than in the original dataset when using repetitive sampling. We hypothesized that this repetitive sampling strategy might cause unintended model behaviors, potentially explaining the suboptimal performance. To address this, we replaced repetitive sampling with a re-writing method, inspired by [1] and [2]. Re-writing can be intuitively understood as paraphrasing repetitive items: it does not add new information but reduces repetition in the sampled dataset. We then trained the model on the resulting datasets and attached the downstream performance in https://anonymous.4open.science/r/Data_Mixture-5D78/src/downstream.png. From the results, we observed that our method led to more significant improvements in downstream benchmarks, validating the effectiveness of our approach.\n\nTo further validate our method, we present the PPL-step plot in https://anonymous.4open.science/r/Data_Mixture-5D78/src/PPL_step.png (we use Llama 3.1-8B model as an example, with the original dataset’s weight for comparison). Our method achieves comparable performance with faster computation. Ultimately, our method converged at a lower PPL.\n\n---\nQ: Will this method consider synergy effects of different domains?\n\nA: This is a very intriguing and insightful question. To give a brief answer: our method partially accounts for synergy effects. Specifically, it does not consider pairwise influence between domains but instead looks at the aggregated effect of all other domains on the current domain, as detailed in Section 3.2. That said, we’d be happy to share some of our experiences. Early in the project, we attempted to model pairwise synergy effects between domains, initially focusing on just two domains. Appendix F5 illustrates this with instruction following and healthcare domains. We found that instruction following data indeed enhanced performance on healthcare tasks more than healthcare data alone. This finding is interesting, and seems intuitive: effectively responding to healthcare questions requires a model to understand and follow prompts—skills that are improved by instruction following data. However, when we introduced a third domain, the intercorrelations between the two original domains changed, almost becoming unrelated. This led us to conclude that the behaviors of SFT are complex and case-specific. Nevertheless, understanding the underlying mechanisms of SFT remains both an intriguing and challenging task.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "This work relates to data mixing literature in general. It proposes a new method by using the scaling laws to predict the optimal proportions.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "## Strengths\n- The paper is very well organized and easy to read.\n- The idea of using Scaling Laws to predict validation loss and therefore the optimal dataset proportions is very interesting.\n- Theoretical proof is provided.\n- Experiments are done thoroughly on various settings with different models and datasets.\n\n## Weaknesses and Questions\n- I did not spot any critical weakness in this manuscript. One minor concern is the relatively marginal improvement of \"Ours\" compared to \"Original\". But considering that it is hard to expect groundbreaking improvements just by changing the dataset proportion, the results are acceptable.\n- Out of curiosity, will this method consider synergy effects of different domains? For instance, some arbitrary pair of domains can have synergic effects, which may impact the optimal proportion that should be assigned. Relevant analyses or discussions might be very interesting.", "other_comments_or_suggestions": "N/A - see Other Strengths And Weaknesses", "questions_for_authors": "N/A - see Other Strengths And Weaknesses", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents an approach to data mixing optimization for the supervised fine-tuning of LLMs. The authors frame data mixing as an optimization problem aimed at minimizing validation loss, introducing a method that parametrizes this loss by modeling effective data transferred and leveraging scaling laws for fine-tuning. The proposed algorithm achieves performance often on par with computationally expensive grid search methods. Furthermore, the paper showcases the practical utility of their approach by demonstrating improvements in both validation loss and downstream performance when reweighting popular SFT datasets. 

During the review process, several constructive critiques were raised that the authors should address in the camera-ready version to further strengthen the manuscript. A primary concern revolved around the two-stage parameter estimation approach detailed in Appendix B. Reviewers requested a more thorough justification for this method compared to a single-stage approach and suggested providing empirical evidence of their equivalence or the benefits of the two-stage method, potentially by reporting results from both. Additionally, it was suggested to report the mean squared error (MSE) of the validation loss estimation to quantify the accuracy of the approximation. Finally, the authors should consider elaborating on the limitations of their method, including its reliance on validation loss as a proxy for downstream performance and the potential sensitivity of parameter estimation.
Decision: Accept (poster)