Fast Data Mixture Optimization via Gradient Descent
Authors: Haoru Tan
Sitong Wu
Yanfeng Chen
Jun Xia
Ruobing Xie
Bin Xia
Xingwu Sun
Xiaojuan Qi
Venue: ICLR 2026 Poster (2026)
Keywords: Large models
Data-centric AI
AutoML
Primary area: foundation or frontier models, including LLMs

Abstract: While large and diverse datasets have driven recent advances in large models, identifying the optimal data mixture for pre-training and post-training remains a significant open problem. We address this challenge with FastMix, a novel framework that automates data mixture discovery while training only a single proxy model. Instead of relying on predefined heuristics or resource-intensive simulations, FastMix jointly optimizes mixture coefficients and model parameters, substantially improving efficiency and scalability over prior approaches. At the core of FastMix is a reformulation of mixture selection as a bilevel optimization problem.
Under this reformulation, we show that optimizing mixture ratios is mathematically equivalent to assigning per-source loss weights under uniform source sampling. This embeds the mixture coefficients directly into the differentiable iterative optimization objective, enabling efficient, gradient-based optimization of both mixture and model. To solve the optimization problem, FastMix implements an approximate iterative optimization procedure, alternating between (i) updating model parameters on data sampled according to current mixture ratios (inner loop) and (ii) updating mixture ratios based on validation feedback (outer loop). Across pre- and post-training, FastMix outperforms baselines while drastically reducing search cost: in pre-training, it attains an average score of 48.2 with 1.3 GPU-hours ($\times 550$ vs. RegMix; $\times 55$ vs. CLIMB), and in post-training (SFT) it leads with 65.4 with a $+5.5$ gain over the next best, completing search in 2.2 GPU-hours compared to the 115 GPU-hours required by CLIMB/RegMix.

Published as a conference paper at ICLR 2026
FAST DATA MIXTURE OPTIMIZATION VIA GRADIENT DESCENT
Haoru Tan 12
Sitong Wu', Yanfeng Chen RR Jun Xia',
Ruobing Xie".
Bin Xia',
Xingwu Sun', Xiaojuan Qi"
'Tee University of Hong Kong HHunyuan Team, Tencent 'The Chinese University of Hong Kong
ABSTRACT
While large and diverse datasets have driven recent advances large models, identifying the optimal data mixture for pre-training and post-training remains a significant open problem. We address this challenge with FASTMIX, a novel framework that automates data mixture discovery while training only a single proxy model. Instead of relying on predefined heuristics or resource-intensive simulations, FASTMIX jointly optimizes mixture coefficients and model parame- ters, substantially improving efficiency and scalability over prior approaches. At the core of FASTMIX is a reformulation of mixture selection as a bilevel optimiza- tion problem. Under this reformulation, we show that optimizing mixture ratios is mathematically equivalent to assigning per-source loss weights under uniform source sampling. This embeds the mixture coefficients directly into the differen- tiable iterative optimization objective, enabling efficient, gradient- based optimiza- tion of both mixture and model. To solve the optimization problem, FASTMIX implements an approximate iterative optimization procedure, alternating between (i) updating model parameters on data sampled according to current mixture ratios (inner loop) and (ii) updating mixture ratios based on validation feedback (outer loop). Across pre- and post-training. MIX outperforms baselines while dras- tically reducing search cost: in pre-training, attains an average score of 48.2 with 1.3 GPU-hours ( (x550 vs. RegMix; x55 vs. CLIMB), and in post-training (SFT) leads with 65.4 with +5.5 gain over the next best, completing search 2.2 GPU-hours compared to the 115 GPU-hours required by CLIMB/RegMix.
INTRODUCTION
The performance large-scale models (Yang et al., 2024b; Dubey et al., 2024; Touvron et al., 2023; Hu et al., 2024) depends critically on the data used for training. While large and diverse datasets have driven recent advances, identifying the optimal data mixture for pre-training (Shukor et al., 2025) and post-training (Dong et al., 2023) remains significant challenge.
Popular methods such as manual trial-and-error (Yang al., 2023; Tong al., 2024) or proxy-based methods (Liu et al., 2024; Diao et al., 2025) often do not scale well as models grow larger. For example, proxy-based search methods such as RegMix (Liu al., 2024) and CLIMB (Diao al., 2025) have demonstrated strong generalization and stability, yet they require training a large number of proxy models during the search. This results in prohibitive computational overhead, making mixture optimization increasingly impractical as both models and datasets continue to expand. The central question is thus: how can efficiently determine effective data mixtures for large-scale training?
We address this challenge with FASTMIX. a novel framework that automates data mixture discovery while training only a single proxy model. Instead of relying on predefined heuristics or resource- intensive simulations, FASTMIX jointly optimizes mixture coefficients and model parameters, sub- stantially improving efficiency and scalability over prior approaches. At the core of FASTMIX is re- formulation of mixture selection as weighted bilevel optimization problem in (2), Specifically, we show that optimizing mixture ratios mathematically equivalent to assigning per-source loss weights under uniform source sampling. This reparameterization embeds the mixture coefficients directly into the differentiable iterative optimization objective, enabling efficient, gradient-based op- timization of both mixture and model. To solve the optimization problem (Maclaurin et al., 2015;
Corresponding author: Yanfeng Chen and Xiaojuan Qi

---

Published as a conference paper at ICLR 2026
65.0
49.0
Fabti (ous)
+07.
PastMix (omrs)
+5.5
Average Performance
3 x55
+1.0
Average Performance
CLIM
60.0
CLIMB
47.0
x550
Righix
52.5
xx52
RigME
DoRobdi
.2.5
DoRaM:
IDEAL
46.3
0
10
60
50
720
0
10
115
Time-cost (GPU Hours)
Time-cost (GPU Hours)
(a) Pre-training
(b) Post-training
Figure 1: Average Performance versus Time-cost (GPU Hours) comparison for various data mixture strategies. (a) Pre-training: Our proposed FASTMIX (ours) method achieves the highest perfor- mance with the lowest time-cost. The annotations highlight that it is up to 55x more time-efficient than CLIMB (Diao al., 2025) and 550 more time-efficient than RegMix (Liu al., 2024), while providing significant performance gain. (b) Post-training: In this setting, FASTMIX (ours) again demonstrates tfate-of-the-art performance and time-efficiency, outperforming RegMix with a 52x reduction in time-cost and gaining an additional 5.5 performance points over CLIMB. This illus- trates the superior trade-off between performance and time cost achieved by our method.
Franceschi et al., 2018), FASTMIX implements an approximate iterative optimization procedure, alternating between (i) updating model parameters on data sampled according to current mixture ratios (inner loop) and (ii) updating mixture ratios based on validation feedback (outer loop) via gradient-based optimizer (Kingma & Ba, 2014).
Extensive evaluations demonstrate that FASTMIX optimizes data mixtures across model scales and tasks in both pre-training and post-training, outperforming baselines at a fraction of the computa- tional cost (See Fig. 1). pre-training, delivers a top average score of 48 2 and rank across 14 benchmarks (best on 9) with just 1.3 GPU-hours, achieving x550 faster than RegMix (Liu et al., 2024) and x55 than CLIMB (Diao al., 2025). In post-training (SFT), a math-tuned mixture gener- alizes to coding and STEM-QA, reaching 65.4 (+5,5 over next best) in 2.2 GPU-hours versus more than 115 GPU-hours for CLIMB/RegMis, Overall, FASTMIX makes mixture optimization practical and scalable for next-generation large models.
RELATED WORK
Data-centric technologies Tan al. (2025a); Xia et al. (2025); Tan et al. (2023); Wu et al. (2024); Tan et al. (2025b); Wu et al. (2025a); Tan et al. (2025c) play crucial role in the recent rapid progress of Touvron al., 2023; Allal et al., 2024; Yang al., 2023; 2024a; Zhang et al., 2025b; 2026; 2025a; Wu et al., 2025b). The training of modern large models relies heavily strategically mixing data from diverse sources, spanning languages (Yang et al., 2023), modalities (Gunasekar et al., 2023; Yang et al., 2024b), and difficulty levels (He et al., 2025). This data mixture problem (Ge et al., 2024) presents fundamental challenges not only in pre-training (Shukor et al., 2025; Dubey et al., 2024; Yang et al., 2024b) but also in post- post-training (Dong et al., 2023; Ming al., 2025; Tong et al., 2024). Early practice largely relied on manual heuristics, which lack standardization and often fail to generalize across settings. More recently, optimization-be approaches (Xie et al., 2024; Fan et al., 2023; Liu al., 2024) have been introduced to automate mixture selection.
large models (Dubey al., 2024;
Proxy-based methods (Xie al., 2024; Liu et al., 2024; Diao et al., 2025) adopt two two-phase design which a proxy model is trained under candidate mixtures and its performance used to infer op- timal sampling ratios. For example. DoReMi (Xie et al., 2024) trains small proxy to adjust domain weights based on relative losses, then reuses the optimized ratios train larger model. RegMix (Liu et al., 2024) scales this idea by training hundreds of proxy models under different ratios, fit- ting regression model the resulting mixture-performance pairs, and extrapolating the optimal mixture. CLIMB (Diao et al., 2025) improves efficiency by iteratively refining the search region, reducing the number of proxy models required. Other works (Ye al., 2024; Shukor et al., 2025;
2

---

Published as a conference paper at ICLR 2026
Kang al., 2024) study cross-scale transfer: Shukor al. (2025) provide theoretical and empirical evidence that mixtures found on small models generalize to larger ones, while Ye et al. (2024); Kang et al. (2024) report functional relationships between mixture proportions and performance.
In contrast, dynamic methods (Chen et al., 2024; Ming et al., 2025; Albalak et al., 2023) remove the separate search phase by adjusting mixtures on the fly. IDEAL (Ming et al., 2025), for instance, leverages influence functions (Koh & Liang, 2017) to estimate domain contributions downstream performance and to dynamically rebalance training data.
Overall, proxy-based methods such RegMix and CLIMB generally achieve stronger and more stable performance than dynamic approaches, but at substantial computational cost. Our method, FASTMIX, preserves the reliability of proxy-based optimization while cutting search time from hundreds of GPU-hours to nearly one, achieving both higher efficiency and stronger generalization.
3 FASTMIX
3.1 PROBLEM REFORMULATION WITH REPARAMETERIZATTON
Data Mixture Bi-level Optimization Problem. Formally, data mixture optimization can be posed as bilevel optimization problem. Let D = (Di...., ) be a collection of data sources clusters), and let A R denote the mixture weights, where the feasible set A is the probability simplex (o: 2 and Σi= a; = 1). Given mixture a and model parameters w, the training objective is Limin(D w a)) Let (a) be the parameters obtained by (approximately) optimizing this training objective under a. The target to find mixture weights a* that minimize the validation loss, i.e., (w)g (V, w) evaluated at w (a)):
min Learper ( (a))
s.t. w'(a)= arg min Cuain (l,l)))
Σa==1, 1=1
ai 120 0.
(I)
where the inner-loop aims to find the optimal model weights w" (ax) by minimizing the training loss on the dataset given mixture weights a. The outer-loop then seeks to optimize these mixture weights minimize the model's 's final loss target tasks.
While the bi-level formulation is conceptually appealing, is difficult to solve in practice. The crux handling the mixture weights Unlike model parameters which admit efficient gradient-based updates, mixture (sampling) ratios are typically non-differentiable, precluding end-to-end backprop- agation. Consequently, practitioners resort to greedy heuristics or policy-gradient (score-function) updates to adjust a. These procedures are sample-inefficient and scale poorly with the number of data sources, turning mixture search into dominant computational bottleneck.
Differentiable Formulation. Through simple reparameterization. we recast the original bilevel problem into a mathematically equivalent, fully differentiable objective. The key idea is to replace stochastic sampling by mixture ratios with por-source, differentiable loss weights applied under uni- form sampling, so that each source's contribution is controlled continuously via its weight, yielding the following formulation:
min Liarget (w" (a)) s.t. w=(a) arg min W
Σ
Cis Linain (Di, w))
Σ
a: 01=1,
ag 2 0, (2)
a
where Lumin DD, w) denotes the model's 's training loss on source D,, computed under uniform source sampling (each source selected with probability 1/k). The inner-loop finds the optimal model weights. w² (a), by minimizing weighted sum of the training losses from k different data do- mains. The data mixture weight a; serves as the weight for each domain's loss. The outer-loop then aims to optimize these proportions to minimize the model's loss on target tasks. This reparame- terization key: rather than treating mixture ratios as non-differentiable sampling probabilities, we reinterpret them as continuous coefficients that scale each source's loss. Consequently, the mixture weights = (ai,..., an) fully differentiable and amenable to gradient-based optimization. Standard optimizers (e.g., SGD or Adam) can then jointly update the model parameters and the data weights, enabling efficient end-to-end training.
3

---

Published as a conference paper at ICLR 2026
Algorithm FASTMIX Optimization Algorithm
Initialize model parameters w mixture weights a", inner-loop duration nı and outer-loop duration 722- 2: for 0, 1....... .,.
do
3: if (t) mod 1 * 0 then
4:
// Inner loop: update model parameters (e.g., via the SGD optimizer, and we can change
this update rule to other optimizers, like Adam (Kingma & Ba, 2014))
w'+1 + w - T/e 012 a Cmml
5:
Owe
6:
else
7:
// Outer loop: update mixture weights (e-g.. via the SGD optimizer, and we can change this update rule to other optimizers, like Adam (Kingma & Ba, 2014))
a+ + a' - n/a OLarge ( =2)+
8:
dat
9:
end if
10: end for
11:
Output: the optimized mixture weight final after the final outer loop update.
Proof of equivalence. Let D U D, denote the union of k data sourees (or clusters), and let a = (at...., a4) be mixture weights with E, a; = 1, a; 2 To sample a training example æ, first draw source index Cat(a), then sample D;. The training loss under this mixture sampling is
Luain (D, w a) = EE-Cat(() Ex [((x, w)] =
Σ 1=1
ay Luain (D w),
(3)
where f(x, w) is the per-example loss and Lunin (D w) Ez~D, [6(x,, w)] is the expected loss on source D;. Thus, under mixture sampling, the expected training loss is a convex combination the per-source losses, with coefficients given by the mixture ratios.
3.2 How TO OBTAIN BETTER GENERALIZATION PERFORMANCE?
Like most AutoML algorithms, FASTMIX requires search target, typically defined as perfor- mance metric on a held-out validation However, relying on validation performance alone can lead to overfitting to quirks of the validation data and limited transferability to new scenarios. To improve generalization, to encourage diversity among mixture weights, and (ii) incorporating training loss into the search target to balance validation and training signals.
propose two complementary strategies: (i) entropy- based regularization
Entropy- regularization. Entropy regularization prevents the mixture distribution from col- lapsing onto narrow subset of data sources. Given mixture weights (a1...., og) across k sources, Oiş log og. Minimizing this discourages overly peaked
add the penalty Rentropy E distributions, promoting more uniform weight allocation. This reduces sensitivity to spurious vali- dation patterns and improves robustness by leveraging multiple data sources.
Training loss as an auxiliary target. We further integrate the training loss into the search objective to complement the validation signal. While the validation term reflects oft-ll- generalization, the training term measures how effectively the model fits the mixture as a whole. Combining the two reduces over-reliance on the limited validation set and guides the search toward mixture ratios that generalize more reliably across both in-domain and out-of-domain data.
Joint objective. Together, entropy regularization and the auxiliary training loss yield the following search objective:
Langer (wr) + 3 Lirain (w) + Or log Os;
where 3 > 0 and A 2 are trade-off hyperparameters. Empirically, À is set to small value (e.g., 10-5) encourage diversity without dominating the optimization, while 3 most effective moderate values (e.g., 0.1). We provide a detailed sensitivity analysis of these hyperparameters in ablation studies. Overall, these strategies substantially improve the generalization ability of

---

Published as a conference paper at ICLR 2026
FASTMIX, enabling it to discover mixtures that not only perform strongly on validation benchmarks but also transfer robustly to broader real-world applications.
3.3 OPTIMIZATION
Although the reparameterized formulation enables end-to-end differentiation over both model pa- rameters and data mixtures, the resulting bilevel problem still difficult to solve directly. Accord- ingly, we adopt an iterative procedure (Alg. 1) that alternates between updating the model parameters and the mixture weights (Maclaurin et al., 2015; Liu et al., 2018; Pedregosa, 2016; Franceschi al., 2018). The two key steps are outlined below.
(i) Inner loop (network parameter update). Given current mixture weights a', the model param- eters w are updated for nı steps via stochastic gradient descent (SGD) to minimize the weighted training loss Cuain
o (Ei Ca Di, w²) Thu
w w'
(5)
Owt
where Luiin( (D,, w) denotes the model's 's training loss on source D,. computed under uniform source sampling (each source selected with probability 1/k). This is repeated for ii iterations. Other gradient-based optimizers, such as Adam (Kingma & Ba, 2014), are compatible with our framework. After Tt updates, we denote the resulting parameters as w(+21
(ii) Outer loop (mixture weight update). The mixture weights a' are then updated using validation feedback Carget. Specifically, the model is trained for n2 iterations with the previous mixture weights a', and the resulting parameters w are evaluated on the validation loss Lurget- The mixture weights are updated as:
Laarges (w
a t- ta an-r" m/a
(6)
Oa
In effect, a'11 is updated according to how the validation loss responds after na steps of training under a². This naturally assigns larger weights to data sources that contribute more to improving validation performance. A key consideration is how the gradient is estimated, since this directly impacts both the direction of updates and the efficiency of the search.
In the special case 122 1 with SGD updates, the gradient of the validation loss with respect to at yields a closed-form solution:
Dtargor (w'+1) @w Ow Oa
- V (V, (V, V Curainn w'), da,
Ow
where D; denotes the i-th training source. This shows that por-source training losses directly shape the mixture gradients. The following derivation shows why the formula holds. Under the SGD up- date rule, the weights w at time / + are updated based on the gradient of the loss function with re- spect to the mixture coefficients at: w w'-n'
ni. V w 12 i= a² Lamin Di, w²)]. Taking the deriva-
tive of w++1 with respect to af. we get: Ow'+1
8 w² nhe V. (2* at Luain (Dj, w ) da dat
Since wf is independent of af, the derivative of the first term is zero. Due to the linearity aw*
of the derivative and the sum, only the term corresponding to of remains, hence, da! Vn (D, w').
The formulation in Eq.(7) can be intuitively understood as follows: The gradient with respect to
a; is proportional to the alignment between (i) the validation gradient Va (V, w(+1 and (ii) the training gradient from source D., Vw w²). If these gradients are aligned (positive dot product), the derivative 'i' Va wrain (D,, w²) is negative, so a gradient-descent step on a; increases its weight, emphasizing sources whose updates also reduce the validation loss. If they are opposed (negative dot product), the derivative is positive and step decreases a;, down- weighting sources that harm validation performance. the procedure reallocates mass toward data sources whose training signals most effectively improve the validation objective.
orthogonality yields small updates. Thus,
When n22 > 1, deriving a closed-form gradient becomes intractable, requiring finite-difference ap- proximations or similar techniques, which are often unstable and inefficient. In contrast, f2
5

---

Published as a conference paper at ICLR 2026
admits closed-fomm gradient that is both computationally efficient and empirically effective. Ad- ditional discussion of the multi-step case is provided in the Sec.4.3.1.
3.4 HANDLING NON-DIFFERNTIABLE CASES
Our optimization algorithm is designed for settings where both Lurget and Crain are differentiable. However, in practice, non-differentiable situations may arise. We discuss two representative cases below.
Non-differentiable targets. One common challenge arises when the objective function is non- differentiable, such as when validation performance is measured by discrete metries (e.g., accuracy) rather than a smooth loss. A standard remedy is to apply finieeddiferrene methods to approximate gradients. However, our experiments show that these methods are often slow and suffer from nu- merical instability. In such cases, we propose using differentiable proxy objective, for instance, the supervised tne-tuning (SFT) loss for question- question-answeriig answering tasks, which provides smooth surrogate while remaining aligned with the discrete evaluation metric. This approach has proven to be highly effective in practice.
Long outer-loop horizons. Another challenge arises when the outer-loop duration parameter m2 is greater than one. In this case, computing the gradient of the mixture weights becomes intractable. Without constraints on n2. one would either need to rely on built-in mechanisms in PyTorch (Paszke et al., 2019), such as backpropagation- warrooontttthhhhhhh memory-intensive in
(BPTT), which quickly becomes prohibitively settings, or fall back on general gradient-estimation techniques such as finite differences, which again are slow and unstable. To avoid these pitfalls, restrict ng = 1 whenever possible, which not only yields a closed-form gradient but also delivers the most stable and efficient optimization behavior.
EXPERIMENTS
To comprehensively evaluate the effectiveness of our proposed framework, we conduct experiments on data mixture optimization across different stages of large language model (LLM) training, in- cluding both pre-training and post-training. The compared methods cover a wide spectrum of ap- proaches, ranging from human expert tuning to proxy-based search methods such as DoReMi (Xie al., 2024), RegMix (Liu et al., 2024) and CLIMB (Diao al., 2025). The subsequent sections are organized as follows: Section 4. presents results on pre-training mixture optimization. Section 4.2 reports experiments in post-training settings. Section 4.3 provides ablation studies.
4.1 PRE-TRAINING STAGE EXPERIMENTS
Setups. Following prior work (Liu al., 2024), we conduct our experiments on the Pile dataset (Gao et al., 2020), focusing on the 17 uncopyrighted subsets available on HuggingFace. For mixture optimization in the pre-training stage, we employ small proxy models (e.g., 1M parameters) trained on up to IB tokens. To test the method's generalization ability, consistent with Liu et al. (2024), we use the loss on a representative and diverse part of the training data (the Pile-cc sub-set (Gao al., 2020)) as the search target. For FASTMIX, we employ only a single proxy model, whereas RegMix uses 512 by following (Liu al., 2024) proxy models and CLIMB uses 64 (Diao et al., 2025). For the Human Heuristic baseline, we directly adopt the manually tuned mixture configuration reported in (Liu et al., 2024) to ensure faimesss. After the search stage, use the mixture configurations obtained by each method to train a 1B-parameter model on 25B tokens. For evaluation, we focus on the accuracy of the pretrained model on suite of downstream task benchmarks, including Social IQA (Sap et al., 2019), HellaSwag (Zellers et al., 2019), PiQA (Bisk et al., 2020), et.ll In addition, we also examine the time cost incurred by different methods during the search stage.
Results. As shown in Figure 2, our proposed method, FASTMIX, demonstrates significant advan- tages in both downstream task performance and computational efficiency compared to existing data mixture strategies. It achieves the highest average performance score of 48.2 22 and the best average rank of across all 14 downstream benchmarks, outperforming strong baselines including CLIMB (47.5) and RegMix (47.2). This top ranking underscores its consistent and robust generalization capabilities, further evidenced by its leading results on of the 14 individual tasks. Most notably, FASTMIX offers dramatic improvement in search efficiency, requiring only 1.3 GPU-hours to

---

Published as a conference paper at ICLR 2026
Downstream Task Performance Comparivon
Search Cast (Log Scale)
Soctal IQA
DeRAA-
T.4k
Helldwag
HOA
Landeda
Righlo
ARC-E
COPA
CLDIB
00P
Rights
PauMiix (ture-
MaluKC
C
10
Exalention Seare (Higher is Botter)
GMU Hours (1 (Lawer Batter)
Figure 2: Comparative evaluation of different data mixture strategies in the context of large-scale pretraining, examining their impact both downstream task performance and training efficiency.
identify the optimal mixture. This is orders of magnitude faster than other automated methods, such as CLIMB (71.9 GPU-hours) and RegMix (720.5 GPU-hours), validating the efficacy of our sin- gle proxy model and gradient-based optimization approach. Collectively, these results confirm that FASTMIX not only discovers superior data mixture configurations but also drastically reduces the computational overhead of the search process, offering a scalable and practical solution for large- scale model training.
4.2 POST-TRAINING STAGE EXPERIMENTS
Setups. Building on our pre-training success, we next validated FASTMIX in the post-training stage, aiming to optimize data mixtures for specialized tasks on the Qwenn..5-Math-Instnee 7B model (Hui et al., 2024). For this study, we sourced supervised fine-tuning (SFT) data from eight distinet domains, including Math 1Oph 220k (Open-R1 Team, 2024)), Code (the programming- related subset from the OpenThoughts- 1 14K (Guha et al., 2025)), Dialogue (ShareGPT (RyokoAI, 2023)), and STEM (Platypus (Lee et al., 2023)). Our optimization search objective was a 1:1 weighted sum of scores from two mathematical benchmarks, the simpler GSM8K (Cobbe et al., MAARIO-Matt-Reasonne
2021) and the more challenging gaokao2023en model's generalization capabilities, we extended our test suite beyond math (MATH (Hendrycks et al., 2021), AIME-24 (Jia, 2024)) to include tasks in coding (LiveCodeBench-v>) (Jiang et al., 2024)) and STEM guestion-answering (GPQA-Diamond (Rein et al., 2023)). A significant chal- lenge the post-training setting the absence of very small (e.g.. 10M parameter) proxy models. Therefore, we had to conduct our search using proxy models of approximately billion parame- ters (Qwen2.5-115BB Instructt (Qwen et al., 2025)), with evaluation performed on larger models (7B). This constraint exposed a critical limitation of resource- -nnuuree-nttesivee methods (Liu et al., 2024; Diao et al., 2025), which require training hundreds of proxy models. Given the immense computational cost, our cluster was unable to support hundreds of full IB-model training runs, so we had to re- duce the number of proxy models for both RegMix and CLIMB to just 64. In contrast, FASTMIX's reliance on a single proxy model enabled it to operate efficiently within these resource limitations, highlighting its superior scalability for larger-scale tasks.
2023). To evaluate the
Results. the post-training (SFT) stage, the advantages of FASTMIX are further solidified, demon- strating an even more dominant performance as shown in Figure 3. Our method achieved the highest score across all four benchmarks spanning mathematics, coding, and general question-anse -anseein...seeiing.... resulting in a superior average performance of 65.4 and a top rank of 1, by a significant 5.5 point lead over the next best method, CLIMB (59.9) (Diao al., 2025). Crucially, these results highlight the exceptiooal generalization capability of FASTMIX. While all automated methods used perfor- mance mathematics benchmarks (GSM8K and gaokao2023en) as the guidance signal for opti-

---

Published as a conference paper at ICLR 2026
Downstream Performane
Search Efficiency (Log Scale)
9.7h
Reglix
CUMB
117.4
Dctend
Raphis
Fox38x (Qun)
FartMo (4rk)
16TH
AMb:24
(PGN-Derood
GPU Hours (Loweris Better)
Figure 3: Comparative evaluation of different data mixture strategies in the context of large-scale post-training (SFT), examining the efficiency and downstream task performance.
mization, FASTMIX not only excelled in the math domain but also achieved the best performance on LiveCodeBench (coding) and GPQA-Diamond (STEM QA). This strongly indicates that the data mixture identified by FASTMIX avoids overfitting to the optimization signal and instead fosters more fundamental and comprehensive improvement in the model's capabilities, all while maintain- ing remarkable efficiency by completing its search in just 2.2 GPU hours, substantially faster than RegMix (115.9 hours) and CLIMB (117.4 hours).
4.3 ABLATION STUDY
This section presents comprehensive ablation study to validate the key design choices of our pro- posed FASTMIX framework. All experiments were conducted in pre-training setting, and the reported results are the average accuracy across 14 benchmarks.
4.3.1 ABLATION FOR THE LOOP DURATION HYPER-PARN PARAMETERS 11 AND f22
The hyper-parameter T. represents the inner-loop duration, which is the number of steps the model's parameters are updated for each single update of the mixture weights. We ablated the inner-loop duration (n,) to find the optimal update frequency for the model parameters in Fig. Fig.4(a). Performance rose from 47.3 to a peak of 48.2 as nı increased from 20. The score then plateaued and began to decline beyond = 40, suggesting that a moderately long inner loop is needed to effectively learn from the mixture, while an overly long loop lead an insufficient searching process.
The hyper-parameter n2 represents the outer-loop duration, which determines how many inner- loop steps occur before the mixture weights are updated based on the validation loss. Notably, we derive closed-fomm solution for the mixture weight gradient only when n2 = 1. For cases where n2 * 1, we must rely on general-purpose, but often inefficient, methods like finite-difference algorithms, which did not consider for primary approach. Our experiments confirmed this design choice in Fig.4(a). We tested n2 values {1, 10, 20, 40), and the average performance was highest with n2 at 48.1. Performance consistently decreased n2 increased, falling to 44.2 for n2 40. These results demonstrate that the efficiency and numerical stability our closed-form gradient are not only computationally beneficial but also lead to superior empirical performance.
4.3.2 ABLATION FOR THE REGULARIZATION TERMS.
First, we investigated the effect of the entropy regularization coefficient X. This term encourages a diverse data mixture, preventing the model from focusing too heavily on a single data source. We varied À across a wide range, from 0.1 down to 10 As illustrated in Fig.4(b), the algo- rithm's 's performance quite robust when À less than - However, using an excessively large (e.g., 1.0) significantly disrupts the optimization process, causing the search to fail to converge gigh-performing solution. This suggests that while some degree of regularization is beneficial, a strong entropy penalty can prevent the algorithm from finding the truly optimal, and potentially non-uniform, mixture.
8

---

Published as a conference paper at ICLR 2026
Outer Loop Duration (02)
Training loss ferm (9)
Performaace (A)
Perfortainn (5)
Performance
Parformence
Pefformncc
47.
<hablo Statiw
Rodia Roal
1000
1e
noeat
000
Loop Duration (71)
Entrapy-Temm (A)
Experiment Index
(b)
(c)
Figure 4: This figure presents a series of ablation studies to validate the design choices of our pro- posed FASTMIX framework. (a) An analysis of the inner-loop (n1) and outer-loop (n2) durations.
(b) The effect of varying the entropy (A) and auxiliary training loss (3) coefficients. (c) A compari- son of the final search performance across multiple random initializations.
Next, we ablated the coefficient for the auxiliary training loss 3, which helps guide the search towards solutions that generalize well. We tested values from 0.001 to 0.6. The results indicate that a moderate value for 3 is optimal. Performance was highest when 3 was between 0.1 and 0.3, suggesting that balanced signal from both the validation and training losses is crucial for finding an effective data mixture. Values that are too small provide insufficient regularization, while values that are too large might cause the model to over-optimize for the training distribution at the expense of generalization. The trends for both and / are clearly visualized in Fig.4(b).
4.3.3 ABLATION FOR THE INITIALIZATION.
We next investigated the impact of the initial solution the search results for FASTMIX. We conducted several experiments with different random initializations, and the results are presented in Fig.4(c). As shown in the figure, our method demonstrates remarkable stability. Across 11 different random initializations (labeled E0-E10), the final performance consistently remained high, with average score of 48.34 and standard deviation of only 0.48. The results showed tight clustering, with the highest score at 49.25 and the lowest at 48.53. This consistency suggests that FASTMIX's optimization process is robust and largely independent of its starting point. In all cases, the performance significantly surpassed the baseline RegMix method, which achieved a score of 45.44. This robustness confirms that gradient-basee -aaeentt-badddddeee search effectively navigates the optimization landscape to find high-quality solutions, regardless of the initial conditions.
CONCLUSION
We introduced FASTMIX, efficient framework for discovering data mixtures for large-model training. Our key contribution a weighted bilevel reformulation of mixture selection: via a repa- rameterization, optimizing sampling ratios becomes equivalent to learning per-source loss weights, enabling mixture coefficients to be differentiable. This permits joint, gradient- basedd optimization of both the model and the mixture using a single proxy model rather than hundreds. Across pre-training and bost-training. FastMix delivers superior accuracy with orders-of siceer-offmaageee of-magnitude lower search cost, making data mixture optimization practical, scalable, and robust for next next-generation LLMs.
ACKNOWLEDDGEEN
The work has been supported by Hong Kong Research Grant Council-General Research Fund Scheme (Grant No. 17202422, 17212923, 17215025). Theme-based Research (Grant No. T45- 701/22-R), and Strategic Topics Grant (Grant No. STG3/E-605/25-N) Part of the described research work is conducted in the JC STEM Lab of Robotics for Soft Materials funded by The Hong Kong Jockey Club Charities Trust.

---

Published as a conference paper at ICLR 2026
REFERENCES
Alon Albalak, Liangming Pan, Colin Raffel, and William Yang Wang. Efficient online data mixing for language model pre-training. arXiv preprint arXiv:2312.02406, 2023.
Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Leandro von Werra, and Thomas Wolf. Smollm blazingly fast and remarkably powerful, 2024.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical com- monsense natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Mayee F Chen, Michael Y Hu, Nicholas Lourie, Kyunghyun Cho, and Christopher Ré. Aioli: A uni- fied optimization framework for language model data mixing, arXiv preprint arXiv:2411.05736. 2024.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL https::/arxtt org/abs/ 22100. 14.68.
Shizhe Diao, Yu Yang. Yonggan Fu, Xin Dong, Dan Su, Markus Kliegl, Zijia Chen, Peter Bel- cak, Yoshi Suhara, Hongxu Yin, Mostofa Patwary, Celine Lin, Jan Kautz, and Pavlo Molchanov. Climb: Clustering-base based iterative data mixture bootstrapping for language model pre-training. 13161.
arXiv preprint, preprimtt 2025. URL https:/ www.iv.org/aba....
Guanting Dong, Hongyi Yuan, Keming Lu, Chengpeng Li, Mingfeng Xue, Dayiheng Liu, Wei Wang, Zheng Yuan, Chang Zhou, and Jingren Zhou. How abilities in large language models are affected by supervised fine-tuning data composition. arXiv preprint arXiv:2310.05449 2023.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Simin Fan, Matteo Pagliardini, and Martin Jaggi. Doge: Domain reweighting with generalization estimation. arXiv preprint arXiv:2310. 15393, 2023.
Luca Franceschi. Paolo Frasconi, Saverio Salzo, and Massimilano Pontil. Bilevel programming
for hyperparameter optimization and meta-learning. In International Conference on Machine Learming, 2018.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima. Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.0003 2020.
Ce Ge, Zhijian Ma. Daoyuan Chen, Yaliang Li, and Bolin Ding. Data mixing made efficient: A bivariate scaling law for language model pretraining, arXiv preprint arXiv:2405. 2024.
Etash Guha, Ryan Marten, Sedrick Keh, Negin Raoof, Georgios Smyrnis, Hritik Bansal, Marianna Nezhurina, Jean Mercat, Trung Vu, Zayne Sprague, Ashima Suvarna, Benjamin Feuer, Liangyu Chen, Zaid Khan, Eric Frankel, Sachin Grover, Caroline Choi, Niklas Muennighoff, Shiye Su, Wanjia Zhao, John Yang, Shreyas Pimpalgaonkar, Kartik Sharma, Charlie Cheng-Jie Ji, Yichuan Deng, Sarah Praut, Vivek Ramanujan, Jon Saad-Falcon, Jeffrey Li, Achal Dave, Alon Albalak, Kushal Arora, Blake Wulfe, Chinmay Hegde, Greg Durrett, Sewoong Oh, Mohit Bansal, Saadia Gabriel, Aditya Grover, Kai-Wei Chang, Vaishaal Shankar, Aaron Gokaslan, Mike A. Merrill, Tatsunori Hashimoto, Yejin Choi, Jenia Jitsev, Reinhard Heckel, Maheswaran Sathiamoorthy, Alexandros G. Dimakis, and Ludwig Schmidr. Openthoughts: Data recipes for reasoning models, 2025. URL https: // /axxiv org aba/abaa//255 001.5.66.........
Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, et al. Textbooks are all you need. arXiv preprint arXiv:2306. 1644, 2023.
10

---

Published as a conference paper at ICLR 2026
Zhiwei He, Tian Liang, Jiahao Xu, Qiuzhi Liu, Xingyu Chen, Yue Wang, Linfeng Song, Dian Yu, Zhenwen Liang, Wenxuan Wang, et al. Deepmath-103k, A large-scale, challenging, de- contaminated, and verifiable mathematical dataset for advancing reasoning. arXiv preprint arXiv:2504. 1456, 2025.
Dan Hendrycks. Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03877 2021.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et al. Minicpm: Unveiling the potential small language models with scalable training strategies. arXiv preprint arXiv:2404.06333 2024.
Binyuan Hui, Baosong Yang, Zeyu Cui, Chengpeng Li, Jiaxi Yang, Dayiheng Liu, Lei Zhang, Tianyu Liu, Jiajun Zhang, Bowen Yu, et al. Qwen2.5-math technical report: Toward mathematical expert model via self-improvement. arXiv preprint arXiv:2409. 12122, 2024.
Maxwell Jia. 2024.2024.
AIME_200 2024, 2024.
Congzheng Jiang, Samuel Dooley, Colin White, Manley Jin, Yichen Shen, David Shi, Ruohong Zheng, David Chen, Yilun Zhang, Yifei Li, et al. Livecodebench: Holistic and contamination free evaluation of large language models for code. arXiv preprint arXiv:2403.07977. 2024.
Feiyang Kang, Yifan Sun, Bingbing Wen, Si Chen, Dawn Song, Rafid Mahmood, and Ruoxi Jia. Autoscale: Scale-aware data mixing for pre-training llms. arXiv preprint arXiv:2407.00177, 2024.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6998 2014.
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International conference on machine learning, pp. 1885-1894, PMLR, 2017.
Ariel N. Lee, Cole J. Hunter, and Nataniel Ruiz. Platypus: Quick, cheap, and powerful refinement of lims. arXiv preprint NXXN:2308.003//1 2023.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv preprint arXiv: 09055, 2018.
Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, and Min Lin. Regmix: Data mixture regression for language model pre-training. arXiv preprint arXiv:2407.01492 2024.
Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-based hyperparameter optimiza- tion through reversible learning. In International Conference on Machine Learning, pp. 2113- 2122, 2015.
MARIO-Math-Reasomne
English Translation of Chinese Gaokao 2023 Mathematics Problems. https:/ httss///uuggte /twggingface..eecccco/a.aa co /dtaasese MARTO-Mat Mh--MMathRRaooooo haaemiing/Gaoaaeee
Gaokao2023-Math
0 Mat.okaaa222-Math--a 2023. Accessed: 2025-09-24.
Chenlin Ming. Chendi Qu, Mengzhang Cai, Qizhi Pei, Zhuoshi Pan, Yu Li, Xiaoming Duan, Lijun Wu, and Conghui He. Ideal: Data equilibrium adaptation for multi-capability language model alignment. arXiv preprint arXiv:250s.12772. 2025.
Open-R1 Team. OpenR1 MaanR1-Math-22000 A Large-Scale Dataset for Mathematical Reasoning. https: / www.hgiggaaaccc. co datasataassssssee setattssesssppe-aaaaaeee openn wttpo-ee/ppenll-aaaee /Cppen1-Maataaaaeeee 220k, 2024. Ac- cessed: 14.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.

---

Published as a conference paper at ICLR 2026
Fabian Pedregosa. Hyperparameter optimization with approximate gradient. In International Con-
ference on Machine Leurning, 2016.
Qwen, An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng. Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang. Jiaxi Yang. Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xue, Pei Zhang, Qin Zhu, Rui Men, Runji Lin, Tianhao Li, Tianyi Tang, Tingyu Xia, Xingzhang Ren, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yu Wan, Yugiong Liu, Zeyu Cui, Zhenru Zhang, and Zihan Qiu. Qwen2.5 technical report, 2025. URL https::tp//..c. / a xx
wrxii iv. org/ abs/ 2412 22. 15115.
David Rein, Arnav Gudibande, Jackson Petty, Nishant Balepur, Houman Owhadi, Erik Jones, Yi- Fan Li, Sasha Brown, Jeremy Burnside, Kevin Michael, Josh Albrecht, Samuel R. Bowman, Brian Christian, Samuel Hammond, Ashley Pilipiszyn, Josh Seares, Jordan L. Taylor, and William Saunders. Gpqa: A graduate-level google-proof qa benchmark. arXiv preprint arX0022311.1202,, 2023.
RyokoAL ShareGPT52K, 2023. URL haga:::/:///ugggggafacca... http:///ggggggaaaaaccaaaaaatt/tp ShareGPT52K.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Common- sense reasoning about social interactions. arXiv preprint arXiv:1904. 09728, 2019.
Mustafa Shukor, Louis Bethune, Dan Busbridge, David Grangier, Enrico Fini, Alaaeldin El-Nouby, and Pierre Ablin. Scaling laws for optimal data mixtures. arXiv preprint arXiv:2507.0944 09444, 2025.
Haoru Tan, Sitong Wu, Fei Du, Yukang Chen, Zhibin Wang, Fan Wang, and Xiaojuan Qi. Data prun- ing via waannn-ooe-sampl--one 2023. URL https: wttp/v.ivvvvo//a//ee//aaa org abs//bss000 146..44666 164.
Haoru Tan, Sitong Wu, Wei Huang, Shizhen Zhao, and Xiaojuan Qi. Data pruning by information maximization, 2025a. URL
2550g////.////0....../ 01701.
Haoru Tan, Sitong Wu, Xiuzhe Wu, Wang Wang. Bo Zhao, Zeke Xie, Gui-Song Xia, and Xiaojuan Qi. Understanding data influence with differential approximation, 2025b. URL https:// arxiv.org/ abs/ 2508 14648.
Haoru Tan, Xiuzhe Wu, Sitong Wu, Shaofeng Zhang, Yanfeng Chen, Xingwu Sun, Jeanne Shen, and XIAOJUAN QI. Understanding data influence in reinforcement finetuning. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025c.
Shengbang Tong, Ellis Brown, Penghao Wu, Sanghyun Woo, Manoj Middepogu, Sai Charitha Akula, Jihan Yang, Shusheng Yang, Adithya Iyer, Xichen Pan, Austin Wang, Rob Fergus, Yann LeCun, and Saining Xie. Cambrian-1: A Fully Open, Vision-Centric Exploration of Multimodal LLMs. arXiv preprint arXiv:2406. 16860, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
Sitong Wu, Haoru Tan, Zhuotao Tian, Yukang Chen, Xiaojuan Qi, and Jiaya Jia. Saco loss: Sample- wise affinity consistency for lasiun-language pre-training. In 2024 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 27348-27359, 2024.
Sitong Wu, Haoru Tan, Yukang Chen, Shaofeng Zhang, Jingyao Li, Bei Yu, Xiaojuan Qi, and Jiaya Jia. Mixture-of-score:: Robust image-text data valuation via three lines of code. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 24603-24614 October 2025a.
Sitong Wu, Haoru Tan, Jingyao Li, Shaofeng Zhang, XIAOJUAN QI, Bei Yu, and Jiaya Jia. DLoFT: Gradient-decoupled fine-tuning for generalizable long chain-of-thought reasoning. In The Thirty- ninth Annual Conference on Neural Information Processing Systems, 2025b. URL https:/ openreview n net:/forumm1--/220ttte
12

---

Published as a conference paper at ICLR 2026
Bin Xia, Bohao Peng. Yuechen Zhang. Junjia Huang, Jiyang Liu, Jingyao Li, Haoru Tan, Sitong Wu, Chengyao Wang, Yitong Wang. Xinglong Wu, Bei Yu, and Jiaya Jia. Dreamomni2: Multimodal instruction-based editing and generation, 2025. URL https: //apa//////..o/////////0. 06679.
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy S Liang, Quoc V Le. Tengyu Ma. and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. Advances Neural Information Processing Systems, 36, 2024.
Aiyuan Yang. Bin Xiao, Bingning Wang, Borong Zhang, Ce Bian, Chao Yin, Chenxu Lv, Da Pan, Dian Wang, Dong Yan, et al. Baichuan 2: Open large-scale language models. arXiv preprint arXiv:2309.10305. 2023.
Yang, Baosong Yang, Binyuan Hui, Bo Zheng. Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang. Jian Yang, Jianhong Tu, Jianwei Zhang. Jianxin Ma, Jianxin Yang, Jin Xu, Jin- gren Zhou, Jinze Bai, Jinzheng He. Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang, Mei Li, Mingfeng Xue, Na Ni, Pei Zhang. Peng Wang, Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Siman Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wen- bin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Xuejing Liu, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yugiong Liu, Zeyu Cui, Zhenru Zhang, Zhifang Guo, and Zhihao Fan. Qwen2 technical report, 2024a. URL
2407 10671.
An Yang. Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al. Qwen2. 5 technical report. arXiv preprint arXiv:2412.151155 2024b.
Jiasheng Ye, Peiju Liu, Tianxiang Sun, Jun Zhan, Yunhua Zhou, and Xipeng Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. arXiv preprint arXiv:2403. 16952, 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a ma- chine really finish your sentence? arXiv preprint 1rXiv:195.008830 2019.
Xichen Zhang, Sitong Wu, Haoru Tan, Shaozuo Yu, Yinghao Zhu, Ziyi He, and Jiaya Jia. Smartswitch: Advancing Ilm reasoning by overcoming underthinking via promoting deeper thought exploration, 2025a. URL https:t/arxii.ceoom
/2500. 197677
Xichen Zhang, Sitong Wu, Yinghao Zhu, Haoru Tan, Shaozuo Yu, Ziyi He, and Jiaya Jia. Scaf- grpo: Scaffolded group relative policy optimization for enhancing Ilm reasoning, 2025b. URL https://arxis org/ alss 221005111.0 19807.
Xichen Zhang, Ziyi He, Yinghao Zhu, Sitong Wu, Shaozuo Yu, Meng Chu, Wenhu Zhang, Haoru Tan, and Jiaya Jia. Searchgym: Bootstrapping real-world search agents via cost-effective and high-fidelity environment simulation, 2026. URL https: //tpoi:////vvv..///////e wrw.oorg/aaao. org./bb///bb/.///// 14615
13

---

Published as a conference paper at ICLR 2026
LIMITATIONS AND FUTURE WORK.
Our current FASTMIX framework focuses on optimizing a static data mixture that remains fixed throughout the training process. While this approach has proven highly effective, more dynamic mixture strategy could offer further improvements. Future work will explore methods for optimizing data mixture that adapts over time, such as by incorporating curriculum learning principles. This would allow the model to learn from different data sources at various stages of training, potentially accelerating convergence and enhancing performance.
ETHICS STATEMENT
This research adheres to the ICLR Code of Ethics. The primary contribution of our work is to en- hance the efficiency of the large model training process, which directly leads reduction in energy consumption and computational resources. By making Al development more sustainable, we aim to make positive contribution to society. We acknowledge the importance of the responsible applica- tion of this technology. Our research does not involve the collection or use of any new personally conducted publicly available datasets.
identifiable information, and all experiments
LLM-USAGE STATEMENT
The authors used a large language model to assist with language polishing, grammar correction, and typo identification in this paper. The ideas, methodology, experimental design, and results presented are the sole work of the authors.
BROADER IMPACT
The development of FASTMIX has significant positive implications for the large model training community. By automating and accelerating the data mixture optimization process, framework drastically reduces the computational resources and human effort required for this critical task. This increased efficiency translates directly into reduced energy consumption and lower carbon foot- print, promoting more environmentally sustainable Al research and development.
14


Reviews:
[{"review_id": "YnVkWU36Hl", "reviewer": "Reviewer_Tw57", "summary": "The paper tackles the problem of data mixture optimization for LLM pre-training and SFT. Existing strong methods such as RegMix and CLIMB find good mixtures but do so by training many proxy models, which makes them expensive and slow to iterate on.  The paper proposes FASTMIX, which makes mixture search differentiable by (i) reformulating mixture selection as a weighted bilevel optimization problem and (ii) showing that sampling from a mixture is equivalent in expectation to doing uniform source sampling but applying per-source loss weights. This reparameterization allows the authors to update both model weights and mixture weights in the same training loop with a lightweight hyper-gradient.\n\nThe paper claims:\n\t1.\tOn pre-training over Pile-style subsets, FASTMIX achieves the best average score (48.2) at ~1.3 GPU-hours of search, which the authors say is 55× faster than CLIMB and 550× faster than RegMix while still slightly better in accuracy. \n\t2.\tOn post-training/SFT (math-anchored, tested on coding & STEM QA), FASTMIX reaches 65.4, about +5.5 over the next best method, in 2.2 GPU-hours. \n\nSo the core contribution is a practical, differentiable, single-proxy recipe for data-mixture search that aims to match proxy-based methods in quality but at 1–2 GPU-hours instead of 100+.", "questions": "Can the method enforce per-source lower/upper bounds (e.g. “at least 10% safety-filtered”, “at most 5% web-crawl”)? If yes, how does that change the hyper-gradient update?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The key equivalence of “sampling by α” ≡ “uniform source sampling + per-source loss weights α” is written down cleanly and plugged straight into a bilevel objective. This makes the method implementable in an existing LLM training codebase with only per-batch source tags and a vector of α.\n2. Compute saving is compelling. The claim “1.3 GPU-hours vs 55–550× more for baselines” is helpful, because current RegMix-style methods really do train large banks of proxies. The authors also highlight that they train only one proxy model.", "weaknesses": "No uncertainty / significance reporting. Gains like “+0.7 to +1.0” in pre-training are plausible but small. Without multiple seeds or CIs, especially on the final model, it’s hard to tell whether the improvement is robust or just a friendly seed. This matters because RegMix/CLIMB papers do report variability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "6nzTxnhTGr", "reviewer": "Reviewer_PnEr", "summary": "This paper proposes FASTMIX, a gradient-based framework for optimizing data mixture coefficients in language model training. The key idea is to treat mixture optimization as a differentiable problem by viewing mixture coefficients as loss weights. This enables end-to-end joint optimization of mixture and model parameters with a smaller proxy model. The optimization alternates between inner-loop model updates with outer-loop mixture updates based on validation gradients, using entropy and training-loss for regularization. \n\nIn both pre-training and post-training settings, the authors show FASTMIX outperforms prior mixture-optimization baselines (e.g., CLIMB, RegMix) while being 50–550× more compute-efficient in finding optimal mixture compared to some baselines.", "questions": "- Could the authors clarify what exactly constitutes the “search target” throughout the paper? Does this refer to Eq. 7 as a whole, or only the validation-loss component? Additionally, in practice, was the gradient in Eq. 6 computed using the full expression in Eq. 7, or only its first term?\n- For Fig. 2c, what aspects of the model or training are randomized in the “random initialization” runs?\n- How were the initial mixture weights selected?\n- Does the mixture optimization converge to a stable distribution after a certain number of iterations? How sensitive is the process to different initial mixtures? A plot illustrating mixture trajectories over training would be very helpful.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper shows strong results in downstream performance in both pretraining and post-training settings.\n- The derivation of Eq. 6 is novel and the intuition behind it is clearly presented and insightful.\n- The ablation studies on key hyperparameters are informative", "weaknesses": "- The ADO method by Jiang et al. [1] is, in my opinion, particularly relevant to this paper. But the authors have not compared against ADO in the results. It is relevant for two reasons: (1) it is, to my knowledge, the best data mixture optimization method that does not incur time cost in searching for data mixture and (2) ADO shares the same intuition that one should upsample data mixture that results in the greatest decrease in validation loss. \n- Similarly, the paper misses a few other relevant literature [2, 3, 4]. \n- The paper does not report error bars (in Tables 1 and 2). It would be helpful to include multiple runs to verify that the method reliably outperforms other approaches.\n\n[1] Adaptive Data Optimization: Dynamic Sample Selection with Scaling Laws\n\n[2] Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance\n\n[3] Data Mixture Optimization: A Multi-fidelity Multi-scale Bayesian Framework\n\n[4] ADMIRE-BayesOpt: Accelerated Data MIxture RE-weighting for Language Models with Bayesian Optimization", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "3VTXAG5v55", "reviewer": "Reviewer_Kg4B", "summary": "The paper “Fast Data Mixture Optimization via Gradient Descent (FastMix)” proposes a scalable, differentiable framework for optimizing data mixtures for large-scale model training. The authors address how to efficiently determine the optimal proportions of multiple data sources to maximize performance on a target validation objective, both for pre-training and post-training (SFT) stages.\n\nFastMix reformulates data mixture optimization as a differentiable bilevel optimization problem. Instead of non-differentiable data sampling (drawing examples according to mixture ratios), the authors show that mixture selection can be equivalently represented as per-source loss weighting under uniform sampling. This reparameterization embeds mixture coefficients directly into the training loss, making them continuous and differentiable parameters.\n\nThis leads to significant performance gains over the baselines, while comparatively the time taken is less due to the use of a single proxy model.", "questions": "1. The proposed method assumes that validation loss varies smoothly with mixture weights. How robust is FastMix when this assumption fails, e.g., for highly compositional or discontinuous domains? You could try evaluating on benchmarks such as GSM8K [see Zhang et al, 2024 for an analysis], MATH or HumanEval [see Bradbury et al, 2024 for an analysis] (even one would do for the purpose of the rebuttal).\n\n2. How well do mixture weights learned on smaller proxies transfer to larger models? Ther should be an ablation using multiple proxy sizes (e.g., 0.5B, 1B, 3B).\n\n3. A simple extension towards dynamic mixture optimization could be the following: every K steps, compute $\\alpha_{t+1, \\text{EMA}}$ as a moving average of the current $\\alpha_t$ and the newly computed $\\alpha_{t+1}$, and then resume training the model with the new $\\alpha_{t+1, \\text{EMA}}$. The authors could experiment with this on the Pile itself (Books, Arxiv, Wikipedia etc) which has heterogenous domains where early training may benefit from broad text but later training may become more code/math/etc focused.\n\n4. Can the method generalize to other modalities or multi-domain mixtures? For example, you could try applying FastMix to the MUGEN dataset (which has three modalities, which you can count as data mixture sources) and then optimize mixture weights for downstream video prediction or other multimodal generation task. Another possible experiment is using image-text pretraining corpora such as CC3M, CC12M, LAION-400, and subsequently, FastMix could discover mixture ratios for downstream tasks on validation benchmarks, instead of how the mixtures may have to be manually tuned for CLIP [Gadre et al, 2024][Nguyen et al, 2023].\n\n5. I noticed the mixture visualization in Appendix A and was wondering if it is fair to compare FastMix to RegMix and CLIMB if they don’t incorporate similar regularization to prevent mixture collapse. Can you update your analysis with RegMix and CLIMB variants that prevent this collapse, or can you explain why it is not possible to incorporate such regularization in their frameworks?\n\n6. The approach relies on differentiable validation losses wrt mixture weights. How would it adapt to discrete or non-differentiable metrics such as BLEU, EM, or ROUGE?\n\n7. How sensitive is FASTMIX to the choice of regularization hyperparameters (\\alpha, \\beta) in promoting sparse or diverse mixtures? There should be a regularization ablation showing mixture entropy vs. downstream accuracy, and this would help clarify whether performance depends on sparse or more diffuse mixtures.\n\n8. The method optimizes mixtures based on a single validation dataset. How sensitive are the learned mixtures to the choice of validation task? Does the mixture learned on a particular validation task/dataset generalize to other validation tasks/datasets? Roughly speaking, is there task transfer robustness?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "FastMix presents an interesting advance in data mixture optimization by reframing the problem as a differentiable bilevel optimization, enabling end-to-end gradient-based updates of both model and mixture parameters within a single proxy run. This eliminates the need for computationally intensive multi-proxy searches used in prior works such as RegMix and CLIMB, while delivering superior accuracy and generalization across both pre-training and post-training stages. Its simplicity, efficiency, and demonstrated scalability (spanning 1B to 7B models) make it a practical and elegant solution.\n\nThe method, while seemingly simple, is a novel contribution to the data mixture problem. While this kind of continuous and differentiable transformation is seen in other works, it has not been attempted in mixture optimization for data mixtures before. The bilevel optimization reformulation and its implementation with a single proxy model is also novel. This kind of technique is clearly inspired by advances in hyperparameter optimization. Noteworthy is the direct application to large scale models, opening avenues for e.g. in LLM pre-training. Finally, while other papers usually limit themselves to evaluation on one of pre-training or post-training, this paper does both.\n\nTo summarize,\n1. The authors demonstrate significant performance gains and time savings: up to 550× faster mixture search compared to RegMix and 55× faster than CLIMB, and gains in  pre-training (48.2 avg score) and post-training (65.4 avg score), with top ranks on multiple benchmarks.\n2. The practicality and simplicity of their method, barring the fact that it is not a dynamic method, is a plus point.\n3. The use of a single proxy model as against hundreds is appreciable and makes training with this method more accessible.", "weaknesses": "FastMix represents a conceptually clean and computationally efficient approach to data mixture optimization. However, its limitations arise from the simplifying assumptions underlying its differentiable reparameterization.\n\nFor one, the method assumes smooth, stationary relationships between mixture weights and validation loss. This assumption can break down in real, non-linear training regimes. Further,  the paper’s experimental evaluation, though broad, is limited in terms of cross-scale validation, dynamic adaptation, and proxy–target transferability. The use of a static mixture throughout training restricts its ability to handle evolving data distributions (the authors mention this). Additionally, the efficiency comparisons rely on downscaled baselines, and the generality of results across diverse domains (e.g., multimodal or highly imbalanced datasets) remains untested. Overall, the work is methodologically elegant but still somewhat idealized in scope and validation. Further experimentation is necessary.\n\n1. The authors explicitly note that FastMix optimizes a fixed mixture that remains constant throughout training. The mixture distribution may change over time depending on the scenario. For example, it is possible in non-stationary training regimes that early exposure to simpler data is more beneficial and with time, fine grained data becomes more important.\n\n2. One major issue is that the method assumes the validation loss is a smooth, differentiable function of the mixture weights. This means that it is expected that small changes in \\alpha lead to predictable changes in validation performance. In realistic settings, the loss formulation may exhibit non-smooth regions (although this depends on the parameterization of the network). FastMix’s gradient updates may lead to suboptimal mixtures in this case.\n\n3. Yet another issue is that their method leaves open a proxy-target gap: the experiments rely on 1B param models to optimize mixtures for 7B param target models. There is an implicit scale transfer assumption, that the mixtures should be optimal across sizes. While the paper exhibits strong results, it does not quantify the correlation between proxy and target performance. It is then unclear how well the proxy captures data utility for much larger models.\n\n4. While a relatively minor issue, in the post-training experiments, both RegiMix and CLIMB are limited to 64 proxy models. While this is necessary for ease of computation, the baselines’ potential is undervalued. The efficiency gains of FastMix remain valid but somewhat exaggerated in absolute terms.\n\n5. FastMix optimizes a single scalar validation objective, in which case there is a single downstream goal. It lacks demonstration on multi-domain training and instruction tuning, where mixtures may become biased towards domains that correlate strongly with validation data. The evaluation needs to be more robust and demonstrate realistic tradeoffs in such use cases.\n\n6. All experiments are seemingly conducted on text benchmarks. Experiments on multi-modal mixtures (for e.g. text and image) or even long-tail domains are missing. Thus the claim of scalability remains largely within text.\n\n7. The issue of requiring surrogate validation losses in the case of non-smoothness has been handled cursorily. The authors mostly experiment on the cross-entropy validation loss, and there ought to be a longer discussion about other metrics (e.g. BLEU, Exact Match, etc).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)