TANDEM: Bi-Level Data Mixture Optimization with Twin Networks
Authors: Jiaxing Wang
Deping Xiang
Jin Xu
Mingyang Yi
Guoqiang Gong
Zicheng Zhang
Haoran Li
pengzhang liu
Zhen Chen
Ke Zhang
Ju Fan
Qixia Jiang
Venue: NeurIPS 2025 poster ()
Keywords: data mixture optimization
language models
Primary area: deep_learning
TL;DR: A versatile data mixture ratio optimization framework for LLM training that enjoy both theoretical and practical advantages.

Abstract: The capabilities of large language models (LLMs) significantly depend on training data drawn from various domains. Optimizing domain-specific mixture ratios can be modeled as a bi-level optimization problem, which we simplify into a single-level penalized form and solve with twin networks: a proxy model trained on primary data and a dynamically updated reference model trained with additional data. Our proposed method, Twin Networks for bi-level DatA mixturE optiMization (TANDEM), measures the data efficacy through the difference between the twin models and up-weights domains that benefit more from the additional data. TANDEM provides theoretical guarantees and wider applicability, compared to prior approaches. Furthermore, our bi-level perspective suggests new settings to study domain reweighting such as data-restricted scenarios and supervised fine-tuning, where optimized mixture ratios significantly improve the performance. Extensive experiments validate TANDEM's effectiveness in all scenarios.

TANDEM: Bi-Level Data Mixture Optimization with Twin Networks
Jiaxing Wang"),
Deping Xiang",, Jin Xu?,
Mingyang Yi
Guoqiang Gong',
Zicheng Zhang', Haoran Li Pengzhang Liu Zhen Chen', Ke Zhang',
Ju Fan", Qixiang Jiang'
JJD.com
University of Oxford Renmin University of China
UUniversity of Chinese Academy of Sciences
(vangjiaxing41 xiangdeping1, gongguoqiangg, Stg.g//....ogoo..ooooooooooooooa (liupengzhang, ( |
jiangizia)@jd.ccom {y fyimingyang, fanj) edu.
jin. wwwstats. ac.uk
1ihaoran21 maaennmmmaall..mmmmmmammmmmmmmmm ucas. ac
Abstract
The capabilities of large language models (LLMs) significantly depend on training data drawn from various domains. Optimizing domain-specific mixture ratios can modeled as bi-level optimization problem, which we simplify into a single- level penalized form and solve with twin networks: proxy model trained on primary data and a dynamically updated reference model trained with additional data. Our proposed method, Twin Networks for bi-level DatA mixturE optiMiza- tion (TANDEM), measures the data efficacy through the difference between the twin models and up- up-weights domains that benefit more from the additional data. TANDEM provides theoretical guarantees and wider applicability, compared to prior approaches. Furthermore, our bi-level perspective suggests new settings study domain reweighting such as data-restricted scenarios and supervised fine- mningg where optimized mixture ratios significantly improve the performance. Extensive experiments validate TANDEM's effectiveness in all scenarios.
1 Introduction
The success of large language models (LLMs) largely relies on extensive training data collected from diverse domains, including chat logs [18. academic writings [32], mathematical problems [45), and code repositories (21). The emergent capabilities observed in LLMs are substantially influenced by
the specific composition of cross-domain corpora [33 38). Therefore, it is important carefully
balance the proportions of domain-specine data in training sets to ensure models develop intended and balanced capabilities for target domains.
Optimizing the mixture ratios data domains can be formulated as bi-level optimization problem, in which the inner loop optimizes model parameters for fixed ratio training data, while the outer loop searches for the best mixture ratio on validation data. Due to the difficulty of exactly solving this bi-level problem, transform it into single-level optimization problem. Specifically, the inner-level optimization is viewed as Lagrangian penalty within the outer-level objective. This perspective naturally motivates the introduction of twin networks: proxy model trained exclusively the primary training data, and a reference model exposed to additional validation data. Interestingly, this twin-network formulation relates closely prior methods such as DoReMi [37] and DoGE
"Equal Contribution.
t Corresponding to Mingyang and Pengzhang Liu
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

77. offering insights into addressing their limitations. Based on this formulation, we propose Twin Networks for bi-level DatA mixturE optiMization (TANDEM), an algorithm that measures the domain data efficacy through the twin models' disparities and ultimately approximates the original bi-level optimization problem. Unlike DoReMi, which employs static reference model, TANDEM dynamically updates both models. Additionally, TANDEM enhances stability and provides theoretical convergence guarantees compared to DoGE, as it aggregates multiple updating steps and avoids relying directly gradient estimation, thereby mitigating issues related to high variance.
The bi-level formulation also emphasizes that future research on data domain weighting could focus more on scenarios with limited foomain-speciin data and supervised fine-tuning (SFT), rather than exclusively on traditional pretraining settings with abundant data. From a bi-level optimization viewpoint, assigning equal mixture ratios becomes a valid solution for single-epoch training when domain data abundant, aligning with recent empirical findings highlighting uniform mixing strategies as competitive baselines [2] Despite the prevalence of big data, limited data scenarios are quite common, particularly within specific domains, since large datasets frequently consist of many heterogeneous smaller datasets [33]. Furthermore, SFT often requires domain data to be visited multiple times, which creates generalization gap that leads to non-trivial solution for the bi-level problem. It is precisely in these cases that optimizing data mixture ratios can yield significant improvements.
While previous methods like DoReMi and those built on data mixing laws 21 23 16 are not directly
applicable these newer scenarios, TANDEM can effectively extended. Our experimental results demonstrate TANDEM's 's effectiveness in such settings.
Our contributions can be summarized as follows:
We introduce TANDEM, an effective and efficient algorithm for data mixture optimization that utilizes twined proxy and reference networks approximate the bi-level objective. TANDEM enjoys theoretical convergence guarantees.
We highlight that data mixture optimization is particularly beneficial in scenarios with limited data availability rather than traditional pretraining setups with abundant domain data.
We empirically show TANDEM's effectiveness across standard and data-limited scenarios, showing its superiority over set of competitive data mixture optimization methods.
2 Methodology
We formulate the problem of finding the optimal data mixture ratio as bi-level optimization. To solve the problem, we propose our penalty-based algorithm TANDEM, as presented in Figure [bb] and Algorithm| in Appendix A TANDEM draws insights from many previous works, while improving upon them. Besides, by inspecting the bi-level optimization formulation, we suggest broadening the scope of research on data mixture optimization under limited-domain data and supervised fine-tuning (SFT), rather than focusing solely on conventional data-rich pretraining settings. The notations used in this paper are summarized in Appendix
2.1 Problem Formulation
Consider training an LLM on data composition from M domains, (Di, D D (e.g. Wikipedia, CommonCrawl). Data mixture optimization (DMO) refers to the problem of finding the optimal proportions of data for each domain
[ai, 02, ., am] over probability simplex A == {a E M m= 4mm. = 1, am 2 0}. For any data mixture ratio a, and its correspond- ing model parameter w(a) obtained by training loss, goal is to minimize the validation loss Lval (w))) over a. The optimization problem is formulated as follows:
M
min (va((w(()) = Σ C (w(a)) s.t. w(ax) E arg min Lirain (a, w) = Σ aEA
On Cm (
(1)
By splitting the data into training and validation sets, we can construct the validation and training loss L (w) and L (w) on domain Dm. Intuitively, in the bi-level problem, the outer level problem seeks the optimal domain weights for validation loss with the model weights obtained by
the inner level reweighted training loss. Similar to [7) given the learned final mixture ratio a'

---

stepa
Steps
Stage One
K sleps
R. steps
stepe
Prusy Model R2
Data mixture miño learning
Striis
(f1)
glesin U gr
wiki
to Oel
Misture Ratio
Mixture Ratio
wild
widh
Stage Iwo
Update Ratio
Update Kanoo Cua
Updane Ratio Eur Cu
fs)
3a)
(16)
training
learned mixture ratio.
lg
ty
(a) The Two-Stage DMO.
(b) Computation Procedure of TANDEM
Figure 1: (a) The two-stage data mixture optimization. Optimal mixtures are first learned and then utilized to train the final model. (b) The computation procedure of TANDEM, twined proxy model (green) and reference model (orange) are used to determine the update of the mixture ratio (pink).
construct the final training set by sampling D model is trained (Figure I(a)
SMM m=
-UNIF (Dm. upon which the final
2.2 Twin Networks for Bi-level Data Mixture Optimization
Solving the bi-level optimization problem (1 is challenging due to its nested structure. By viewing the inner-level problem as a constraint and subsequently incorporating into the outer-level problem
Lagrangian penalty, (I) can be reformulated as single-level problem 30 [9
min (a, w) (w) + (Cuain (c, w) min Curiin u)) (c,,)) A,
(2)
Here, the auxiliary variable is introduced as proxy of va(a) E 5*) == arg min Lerain (ar, w)..
The constrain in (T) is transferred into the penalization Lraain (ax, w) min, Lirain (ot, u.). Clearly, by properly invoking y + oc, the solution of 2 will approximate the original (l. This claim is justified by Proposition
[30]. We refer readers the Appendix A for more information. Next, proceed to illustrate our algorithm of optimizing the penalized Lagrange problem (
Algorithm Procedure Besides the mixture ratio a, optimizing 2 deals with two models: proxy model and reference model w. As shown in Figurq(((b) the optimization processes of a, w. and u are indexed by f, & and k.
Update u: Firstly, given a data mixture ratio (!), we find the arg min Craain (9 u) in the penalization term under certain mixture ratio at Our u is updated for K steps to approximate the optimal one:
(t)
Tu V Vuranan (a u
+++1
(green line with arrow in Figure ((b)). The proxy model is trained on the whole training set and maintained through the DMO process.
Update on w: w serves as a reference model updated on both the training and the validation sets. Similar the proxy model u, w is updated for steps before one data mixture ratio a update, which is used to optimize the inner problem (2). Note that the term related to w is (va(ww) + + (a, w),
the update rule then becomes:
w (t)
Neo Cvr (a("),e(" + + Cirain ( 2 (1)
k+1
(orange line in Figure l(b).. Intuitively, training w for multiple steps provides more subtle guidance for mixture ratio update. This will be elaborated in the update part.
Unlike the proxy model u, which maintained throughout training (green line in Figure 1(b) do not maintain an independent reference model w, but rather synchronize the starting point of w and
by setting (t) (!) initialization of the K updates. By doing so, control the disparity

---

between w and u during the optimization. Intuitively, w and u should not diverge from each other as acts as a proxy of w(a)= E (*)a) and w approximates wa)a) E (2)a) := arg min Ha, w). Clearly, the ideal w (ax) under penalized problem will approximate one ur (x) when y -> oo, since the penalized Lagrange problem approximates the original problem ①
The mixture ratio update in (5) seeks solution of (2) That says, a data mixture ratio yields trained model with good performance on data from the validation set under all domains. The term in (2) related is the a weighted domain-wiss loss. Applying the projected gradient descent gives the following update rule:
Update
7 (Cmmin(ax, (Ct,,))=- minl, Limin u)). Recall that Liaiin (ax by definition,
(++1) = IIA a - n
CliM (a(1)
(f)
Cl:A (a
(6) 2
(5)
reference model
model
(Pink line in Figure l(b) where (-) projects the updated mixture ratio to the probabilistic
simplex. The post-updated ui (1) and w are applied to capture the domain-wiss wise loss difference C (w ) (u ) Since the uk and wxx (t) respectively trained training set and training set plus validation set, the aforementioned gap captures the gain of incorporating the additional validation data. Larger loss differences indicate that the model gets significant improvement by consuming more data, thus the corresponding domains are up-weighted.
(e)
(t)
In each episode i, K steps training on and w are conducted to probe the proper direction of a update. Notably, the parameters of proxy model u and reference model w are synchronized at the beginning of probing, forming Twin Networks for bi-level DatA mixturE optiMization (TANDEM) framework. Since the updating of a requires u, w probing, in practice, we decrease the frequency of updating to reduce the computational cost, and leave the proxy model trained freely for E steps before the next a update (see Figure [(b). Altogether the updates u, w. a, the data mixture optimization (DMO) problem can be solved efficiently. The overall computation graph of TANDEM outlined in Figure I(b) and a detailed workflow is summarized in Algorithm in Appendix A
Convergence Analysis Next, we explore the convergence rate of our proposed method. For non-convex bi-level optimization problem, it is standard to study its first-order stationary convergence result e.g.. [30 [9]. Notably, our problem ( is constrain problem over a E A. Thus, it should be considered first-order stationary condition as in 9 30). The convergence result our method is summarized in the following theorem.
Theorem (Informally). For sufficiently large T, under mild assumptions (detailed in Appendix A converges to first-order stationary point of problem in the rate of (T à) by properly selecting "y, K, E, Na, Two and qu-
for ()) obtained in TANDEM Algorithm 1
The proof is left in Appendix A Theorem 1
indicates that TANDEM theoretically ensures the
optimality of the learned mixture ratio.
2.3 TANDEM Improves Existing DMO Methods
We discuss the relationship between TANDEM and two existing methods, DoReMi 37 and DoGE [7]. By comparing the hyper-gradient A that determines a updates in different methods, we show that our TANDEM draws common insights from previous works and improves upon them.
DoReMi updates the mixture ratio according the excess loss of proxy model relative reference model w trained on uniformly sampled data. DoGE (7] pioneers bi-level optimization to settle data mixtures and tracks the influence (28) of each domain the validation data. Specifically, ADOGE Luain (t) (+)). L (a) u (+))) The inner product (+) is essentially first- order approximation of the domain-wiss loss difference between a proxy model and reference model, where the reference model is obtained by one-step update the validation data.
Outlined in Table we see that A of all three methods takes the form of per-domain loss difference between a proxy model and a reference model. Nevertheless, contrasting DoReMi, which adopts
DooRMM and DoGE utilize exponential gradient descent update where a (t+1) a (t) (x)))()pp))))) exp(-n.A)

---

Table I: Summary of A in DoReMi, DoGE and TANDEM. All three methods update a with two models.
Figure 2: SFT exhibits lower gradi- ent alignment than pretraining
1.0
Method
Hyper-gradient
pretrate
0.8
DoReMi
max Mi
(1)
clibd (or, w.
8 Go
0.6
prony model
reformnce model Cl:M (ar (1) 7A
DoGE
MIM
(0 nA
(ual
0.4
train
reference
proxy
0.2
TANDEM
tmin (a
1:Ad
Cl:M Lanin
(a
0.0
reforence model
proxy model
0
200
400
600
800
1000
Steps
fixed reference model, the reference model to in TANDEM is dynamic, which better captures the current training status. In its original form, DoGE incurs significant memory overhead as maintains the per-domain gradients. Relaying explicitly on the gradient estimation also makes it vulnerable to instability issues in high gradient variance scenarios like supervised fine-tuning (SFT). Our TANDEM better adapts, as increasing the probing step K helps reduce the variance on A, as will be elaborated in Section 4.3] Figure 2 we show that SFT exhibits higher gradient variance than pretraining. The alignment of 91 and 82 cos(g1, 92) is used as proxy of the gradient variance, where and J2 are gradients evaluated on different batches. Lower alignment cos 92) indicates larger gradient variance.
2.4 Domain Reweighting Beyond Traditional Settings
In the above, we discuss our method to determine the data mixture ratio by solving bi-level problem ( The circumstances under which DMO is most effective remain to be explored. Theoretically, the standard training setting sets Orn = 1/M for every m. Let us check the following proposition.
Proposition 1. Assume Liman Lig the uniform data mixture ratio Ü. M for 1, 1,..., M constitutes a valid solution of the bi-level mixture optimization
The proof is left in Appendix As can be seen, when the generalization gap L mum approaches zero, uniform weighting emerges as a valid solution, thus making the reweighting less significant. This holds for the conventional data-abundant scenario. When om = Jे. the train data and validation data are independently identically distributed, so the loss gap is approximately zero for the first epoch training. This aligns with the empirical observations that uniform weighting highly competitive that many DMO methods can not consistently outperform 2
When there is limited data in specific domains or the trained model overfits some domain, since both phenomenons result in a large generalization gap, the reweighting technique over data domains becomes significant. In fact, even though I.LMs common for specific domains to be relatively small, e.g.. specialized scientific literature, low-resource languages, or domain-specific user interactions. In practice, the data-restricted scenario is ubiquitous. Furthermore, in SFT, repeated passes over the same domain data exacerbate overfitting and widen the generalization gap [10], which presents more interesting opportunities for domain reweighting. Please refer to Appendix F. for empirical evidence.
trained massive datasets overall, it is
3 Related Work
Data mixture optimization is drawing increasing attention for designing LLMs with comprehensive and balanced capabilities. Our work follows the recent trend of formulating the DMO as a bi-level optimization. We summarize the related literature from these two strands of research in the following.
Data Mixture Optimization Conventional industry practices determine the optimal cross- domain composition with human expertise G. and-error, various heuristics has been explored. DoReMi 37 settles the mixture ratio by
T1]. To circumvent the exhaustive trial-
In the data-ab abundant training scenario, samples are consumed only once.
5

---

minimizing the worst-domain excess loss relative to well-trained reference model, pursuing good performance in all domains. 123 [6] fit global data mixing laws, predict the per- formance on other mixtures, and search for those with good performances. Other works cap- ture the inter-domain relations through domain embeddings [39] or training several models [3] Still-It (3] then utilize the inter-domain relations to establish a local data mixing law that predicts the per-domain loss. Aioli [2] improves Skill-It by dynamically updating the inter-domain relation matrix according to current model states. Never- theless, these approaches remain theoretically un- grounded due to their inductive nature, and incorpo-
Table 2: Computational complexity of different DMO methods.
Method
Computational Complexity
Vanilla Train
CTE 7 CCTE
rating the fitting of the mixing law incurs additional approximation error. DOGE 7 pioneers bi-level optimization to settle data mixtures and tracks the influence (28) of each domain on the validation set. However, assessing the influence relies directly on the per-domain gradient estimation, which under- mines its efficacy in high gradient variance scenar-
DoReMi
DoGE
2CTE
Skill-it
( + 1JCTE + CC
Aioli
CTE + MCTK + +3 M2CT UTE+2CTT + CT
TANDEM
ios while incurring large memory overhead. Comparison of the computational complexity of these strategies is given in Table 2 where C' is the complexity of one step training of model. T is the update number of ax Typically, the overall computational cost of these methods is of the order: Aioli > Skill-it DoReMi z DoGE. z TANDEM 5 [ Note that, although the complexity of DoGE is not particularly large, it requires per-domain gradient computation, which is unfriendly to the computing kernel and takes significantly more time in practice.
Bi-level Optimization Bi-level optimization has been an important research topic in many scientific disciplines, however, solving the bi-level optimization problem is challenging due to the complicated dependency of the upper-level and lower-level problems. Typical bi-level optimization algorithms [24]
29 12 4 15 requires estimation of the implicit gradient, which requires second-order derivatives on the lower-level variables. Incorporating the Hessian makes it computationally prohibitive for
large-scale problems. Recently 30 19 20 pioneer the penalized methods, where the inner-level problem is reformulated into an penalty. As first-order methods avoid the estimation of the Hessian or the Jacobian. Though effective in theory, their practical applications in large-scale LLM settings remain rarely explored. The most similar to ours is the recent ScaleBio 277, which belongs to this category, and specially tailored for large-scale data mixture optimization problems. Nevertheless, the solving procedure is different, TANDEM enjoys more stable training by synchronizing u and w periodically.
maadient-based approaches, the penalized
4 Experiments
In this section, we compare TANDEM to state-of-the-art algorithms in Section 4.2 Then we analyse the effectiveness of each design ingredient in Section
4.1 Experimental Setup
We consider three application scenarios: conventional data-abundant pretraining, data-restricted training, and supervised fine-tuning. A brief summarization of the experimental setup is introduced below, while complete hyper-parameter settings and implementation details in Appendix C
Data-Abundant bundant Scenario: For the data-aboundant scenario, we train 160M GPT-style LMs M on a 6B sampled version of SlimPajama [31] as in (2]. SlimPajama consists of 7 domains: ArXiv, Books, CommonCrawl. C4, Github, StackExchange, and Wikipedia. The statisties of this sampled We set E = 20, K = 5, train with batch size and context length 2048 for 40000 steps (with respect updates of proxy model u. so the mixture ratio a is updated for 2000 steps.) as 2 Though the SlimPajama-68 corpus exhibits significant domain imbalance, 40K
corpus are given in Figure 4
'We assume backward takes 2x computation as the forward process. Skkill-it trains M models with data from each domain for H steps, H AT' in our analysis.

---

Data Abundant Pretraining.
Data Restricted Pretraining.
Supervised Fine-tuning.
0.22
0.20
Text tea.
Celechang
Text
0.16
0.20
0.14
0.30
0.10
16000 32000 40009
1000
2000
4000
5000
Stups
Staps
Figure 3: Step-wise data mixture ratio evolution under three scenarios. We repeat the DMO times and the 95% confidence intervals are given. (a) data-abundant pretraining (b) data-restricted pretraining and (c) supervised fine-tuning.
steps of training doesn't deplete even the smallest domain, so this setting constitutes a abundab one-epoch scenario. The penalty constant is set to I across all the experiments. Experiments are conducted on eight NVIDIA Hopper H-100s.
For the data-restricted scenario, we construct 300M sampled version of SlimPajama and keep the domain distribution unchanged. GPT-style LMs of size 160M, 410M and IB are trained to examine the scalability of TANDEM. In this scenario, we train with K E = 5. batch size 128, and context length 512 for 5000 steps, which ensures that samples in small domains like Arxiv, Books, StackExchange, and Wikipedia are exposed more than once. After DMO, the learned mixture ratio is utilized to resample the 300M corpus for the final model training.
Data-Restricted Scenario:
Supervised Fine-tuning: For supervised fine-tuning, we use 6 major categories (containing 99 tasks) from Natural Instructions (25 (5), Textual Entailment, Answer Verification, Text Matching, Information Extraction, Word Extraction and Text Categorization. Prioritizing diversity of capabilitties and formats, this corpus is comprised of open- open-ended text generation, multiple choice, and True/False tasks. Due the limited space, the statisties
given in Appendix D In this scenario, train pretrained Qwen2-500M model 40) with K = 10, F = 10, batch size 64, and context length 512 for 5000 steps. Different from pretraining where the majority of samples only trained once, in instruction tuning, each sample is on average exposed 1.15 times.
4.2 Comparisons with State of the Arts
We compare TANDEM against various state-of-the-arr methods. Uniform A simple baseline that uniformly mixes groups and requires zero extra training runs. DoReMi (37] adopts the idea of distributionally robust optimization and searches for the data mixture ratio by minimizing the worst- domain excess loss over reference model trained with the uniform strategy. DOGE (7] solves the DMO problem by tracking the data influence of each domain on the validation set and up-weights the most influential domains. Skill-It trains several models fit inter-domain relation matrix, which is later used to establish an incremental data mixing law that induces a mixture ratio a update rule. Aioli (2 improves Skill-It by dynamically updating the inter-domain relation matrix according to current model states. For the baselines, We use the published implementations. Averaged results from 3 runs reported. Due to limited space, the standard deviation is given in Appendix B
Pretraining In Table (Upper), we see that TANDEM discovers mixture ratios with comparative performance. As discussed in Section 2.4] in this scenario, the uniform strategy is highly competitive. The mixture ratio for the baselines is obtained from Aioli [2]. We show the step-wis wise data mixture ratio evolution of TANDEM during DMO in Figure (Left). Note that the uniform strategy is not the only valid solution to the DMO problem, and TANDEM finds another solution that performs equally well. The detailed learned mixture ratio of each method is given in Appendix G).
'DoRMMi: https:/gghhhb...oommm/mmmmmmclld DOGE; https://github..co//lvviaaaaaa Skill-It and Aioli: https://gthubb.cm//taaasssssssse
7

---

Table 3: Comparison for the 160M GPT-style model on SlimPajama in the data-abundant sce-
nario (Upper) and data-restricted scenario (Lower). Per-domain perplexity is reported and "Average"
represents the exponential of the average loss across all domains as 7 2). denotes the results reported use the mixture ratio given in Aioli [
Methods
Arxiv
Book
C4
CommonCrawl
Github Stackexchange Wikipedia
Average
Uniform
11.46
62.53
66.43
59.29
6.71
13.98
28.31
25.74
DoReMit
12.71
80.09
82.60
71.76
5.75
14.26
29.54
28.32
Data
DoGEf
12.89
51.50
54.32
49.34
8,48
16.77
37.21
26.60
Abundent
Skill-Itt
11.76
62.24 64.58
59.84
6.36
12.36
34.87
25.87
Regime
Aiolit
11.47
61,89
65.52
58.24
6.74
14.08
28.48
25.66
TANDEM
11.53
61.82
65.92
58.86
13.76
27.26
25.43
Uniform
18.05
65.86
71.05
63.76
9.37
17.94
34.27
31.53
DoReMi
18,90
80.29 89.05
79.02
10.24
19.74
43.20
36.91
Data
DoGE
17.76
60.88 65.08
58.81
9.00
17.71
33.94
30.10
Restricted
Skill-It
20.93
52.00 57.11
49.50
8.77
16.74
40.49
29.24
Regime
Aioli
17.68
62.48
69.02
61.44
9,26
17.79
33.06
30.67
TANDEM
16.85 52.75 56.82
51.11
8.99
18.21
32.52
28.07
Figure SlimPajama-6B Statistics.
Table Comparison for models of different sizes.
50
160M
410M
IB
gittad wiklsedie
Avg.
Time
Avg.
Time
Avg.
Time
Uniform
31.53
29.59
29.91
DoReMi
36.91
16min 54.61
31
56.53
41 min
20
DoGE
30.10
65min
27.45
172min
Skill-It
29.24
28
27.70
59
27.15
80min
10
Aioli
31.19
36
28.79
64min
28.07
93min
3.4
2.8 3.1
TANDEM
28.07
26min
25.00
57
24.35
77
Data-Restricted Pretraining For the data-restricted training scenario, TANDEM significantly outperforms baselines shown in Table 3 (Lower). For instance, achieves 28.07 averaged perplexity, surpassing the most competitive Skill-It by 1.17 and the uniform baseline by 3.46. In this scenario, the uniform strategy no longer competitive. Equally assigning 1 weights potentially leads to overfitting in small domains (repeated multiple times), while leaving the large domains underfitting. The mixture ratios learned with TANDEM and other baselines are shown Figure and the step- wise data mixture ratio evolution during the data mixture optimization is given in Figure (Middle). We see that after DMO, CommonCrawl and C4 take the majority, driven by their extensive lexical diversity and complex semantics. Nevertheless, compared to the original data distribution, TANDEM up-weights the small domains Arxiv, Books, Github, StackExchange and Wikipedia from 3.4%. 3.7%, 4.2%, 2.8%, 3.1 to 8.9%, 8,5%, 7.6%, 11.5%, 7.9% respectively, preventing these small domains from being overwhelmed while avoiding potential overfitting. Besides, we inspect the scalability of TANDEM with three differentt scales (160M, 410M, IB) in Table 4 TANDEM consistently outperforms the baselines with a large margin while not incurs much computational overhead. We omit DoGE for the 1B model experiment due to its large memory consumption.
Supervised Fine-tuning The results for supervised fine-tuning are given in Table The overall averaged accuracy as well the test loss reported. From Table TANDEM outperforms or on par with the other data mixture optimization methods out of 6 task clusters. Showing its effectiveness in the supervised fine-tuning scenario. Besides, consider more fine-grained task-level SET (instead of the task cluster) case, please refer to Appendix
4.3 Analyses
To evaluate the effectiveness of each design component, conduct ablation experiments. The model used the 160M GPT--ttyle model unless specified. Due to limited space, we focus on the effect of synchronizing and w. the effect of the probing steps K, and the robustness of TANDEM to model

---

Amdi TANDEM
Mixture Ratto
0.3
backs
0.2
Mixture
0.2
gih
vikipedio
0.1
0.4
0.0
arsiv
books
github
stack.
willpedia
Model Size xB
Figure 5: Mixture ratio learned by different methods. Figure Impact of different model sizes.
Table 5: Comparison for the 500M Qwen-2 model on mixture of 6 major categories (99 tasks in
lotal) of supervised fine-tuning tasks.
Method
Textual Ent.
Answer Ver.
Text Mat.
Inf. Ext.
Word Sem.
Text Cat.
Test Loss
Avg. Metric
Uniform 84.4
75.1
86.2
77.8
88.3
83.2
0.231
82.5
DoReMi
85.3
74.9
83.6
78.1
87.0
79.0
0.249
81.6
DoGE
81.4
76.7
86.5
78.6
88.3
82.4
0.297
82.4
Skill-It
84.1
75.0
86.4
78.3
87.9
83.7
0.232
82,6
Aioli
83.9
75.8
86.2
78.0
88.3
84.0
0.229
82.7
TANDEM
85.3
76.3
86.2
78.6
88.5
84.9
0.208
83.3
scales. For more comprehensive analysis, please refer to Appendix including the effectiveness of DMO in the real-world large-scale data (past chinchilla [[4] scenario, sensitivity of the final result on K and E, the performance of the DMO pretrained model on downstream tasks, as well as the performance of the proxy model u.
The Effect of Synchronizing and One obvious characteristic of our method is that the proxy model u and the reference model w are synchronized by setting w uip (8) We show the effect of the synchronization by inspecting Dist(u, w) ||u wil2 along the DMO training process. From Figure 7(a) we see that with synchronization, the distance between u and w is well controlled under 1.5e gradually contracts during the whole DMO process. This contraction is critical for wr and a in the penalized problem (2) to converge to the original bi-level optimization problem (T) as discussed in Section On the contrary, independently maintaining the proxy model u and the reference model incurs blown-up distance Dist(u. u, w) as shown in Figure 7(b)
The Effeet of the Probing Steps K During DMO, the hyper-gradient A determines the update of a. To validate the effectiveness of K in reducing the variance
0.8
TANOSM
TORDUR
consistency costa,
of A, we trace cos (A, Á through the training. Á the hyper-gradient evaluated using another batch of data other than that of A, so cos Á serves as
0.4
0.2
0.0
proxy of the variance, the better A and Á aligns, the
-0.2
smaller the variance of hyper-gradient is. We inspect under the SFT setting (with the 500M Qwen model) where the gradients exhibited large variance. During the training, is fixed to prevent the interference of inaccurate mixture ratio.
200
400
600
800
Steps
Figure 8: The variance of A decreases with more probing steps.
that DoGE exhibits the largest A variance it explicitly depends the noisy parameter gradient estimation. As increases, the variance A decrease, leading to more reliable updates of the mixture ratio. Nevertheless, large K increases the computational cost. So in practice and we need to deliberately choose the proper K.
From Figure
The Effect the Model Scale investigate how the model size will impact the final mixture ratio. In Figure 6 compare learned with models of size 160M, 410M, and IB. We see that learned

---

x10
x10
1.0
2000
2005
0.5
1000
2000
3000
4000
5000
500
1000
1500
2000
2500
(a) Dist(u u, w) with synchronization.
(b) Dist( w) w/o synchronization.
Figure 7: The (u, w) evolution comparison during DMO with and without u, synchronization.
a with larger models are slightly "sharper" than the smaller ones. More specifically, the 1B model further increases the weights of the already large CommonCrawl and C4 while down-weights the others. For large models, due to the increasing capability of memorizing samples, smaller domains less likely to be overwhelmed, while the risk of potential overfitting increases. The capability of capturing this subtle difference further demonstrates the effectiveness of TANDEM. Nevertheless, the optimal mixture ratios under different model scales share the same trend, smaller model can work as a valid proxy for efficient searching.
Conclusion
In this paper, we propose TANDEM. a principled, efficient and versatile data mixture optimization framework. By solving the DMO bi-level optimization problem, TANDEM ensures the optimality of the learned mixture ratios, along with O(T-)) convergence rate. Besides the algorithmic contribution, from the bi-level optimization perspective, we further demonstrate the limitation conventional data-abundant setting in DMO and advocate new settings like data-restricted scenario as well as supervised fine-tuning. Extensive experiments and analysis are conducted to demonstrate effectiveness of our approach. Our work deepens the understanding of data mixture optimization and expands its application scenarios.
References
[1] Stella Biderman. Hailey Schoelkopf. Quentin Gregory Anthony, Herbie Bradley, Kyle O' Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, and Ed- ward Raff eta.al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, 2023.
[2] Mayee F. Chen, Michael Y. Hu, Nicholas Lourie, Kyunghyun Cho, and Christopher Ré. Aioli: unified optimization framework for language model data mixing. International Conference on Learning Representations, 2025.
[3] Mayee F. Chen, Nicholas Roberts, Kush Bhatia, Jue Wang, Ce Zhang, Frederic Sala, and Christopher Ré. Skill-itt a data-driven skills framework for understanding and training language models. In Neural Information Processing Systems, 2023.
[4] Tianyi Chen, Yuejiao Sun, and Wotao Yin. Closing the gap: Tighter analysis of alternating stochastic gradient methods for bilevel problems. Advances Neural Information Processing Systems, 2021.
[5] Sang Keun Choe, Sanket Vaibhav Mehta, Willie Neiswanger Hwijeen Ahn, Pengtao Xie, Emma Strubell, and Eric Xing. Making scalable meta learning practical. Advances in Neural Information Processing Systems, 2024.
[6] Nan Du Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang, Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson.
10

---

Kathleen Hellstee,,,,,mm, Toju Duke, Lucas Dixon, Kun Zhang, Quoe Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. GLaM: Efficient scaling of language models with mixture-of-expeeeee eexeeetss............ee In Proceedings of the 39th International Conference on Machine Learning. 2022.
[7] Simin Fan, Matteo Pagliardini, and Martin Jaggi. Doge: Domain reweighting with generaliza- tion estimation. In International Conference on Machine Learning. 2024.
[8] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang. Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite. Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024.
[9] Saced Ghadimi, Guanghui Lan, and Hongchao Zhang. Mini-batch stochastic approxima- tion methods for nonconvex stochastic composite optimization. Mathematical Programming, (55()))267-305 2016.
[10] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016. Book in preparation for MIT Press.
[11] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, and Abhishek Kadian et.al. The llama herd of models. arXiv preprint, arXiv:2407.2178. 2024.
[12] Riccardo Grazzi, Massimiliano Pontil, and Saverio Salzo. Bilevel optimization with a lower- level contraction: Optimal sample complexity without warm-start. In Journal of Machine Learning Research, 2023.
[13] Shangmin Guo, Yi Ren, Stefano V Albrecht, and Kenny Smith. Sample relationships through the lens of learning dynamics with label information. In Workshop on Interpolation Regularizers and Beyond at NeurIPS, 2022.
[14] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katherine Millican, George van den Driessche, Bogdan Damoe, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack William Rac, and Laurent Sifre. An empirical analysis of compute-optimal large language model training. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[15] Mingyi Hong, Hoi-To Wai, Zhaoran Wang, and Zhuoran Yang. A two-timescale stochastic algorithm framework for bilevel optimization: Complexity analysis and application to actor- critic. SIAM J. Optim., 2023.
[16] Feiyang Kang, Yifan Sun, Bingbing Wen, Si Chen, Dawn Song. Rafid Mahmood, and Ruoxi Jia. Autoscale: Automatic prediction of compute-optimet data compositions for training LLMs. arXiv preprint, arXiv:2407.20117 2025.
[17] Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal- gradient methods under the bolyak--ojasiewicz condition. In European Conference on Machine Learning, 2016.
[18) Andreas Köpf, Yannic Kilcher, Dimitri von Rütte, Sotiris Anagnostidis, Zhi Rui Tam, Keith Stevens, Abdullah Barhoum, Due Minh Nguyen, Oliver Stanley, Richárd Nagyfi, Shahul ES, Sameer Suri, David Alexandrovich Glushkov, Arnav Varma Dantuluri, Andrew Maguire, Christoph Schuhmann, Huu Nguyen, and Alexander Julian Mattick. Openassistant conversations democratizing large language model alignment. In Neural Information Processing Systems, 2023.
[19] Jeongyeol Kwon, Dohyun Kwon, Stephen Wright, and Robert Nowa. A fully first-order method for stochastic bilevel optimization. In International Conference Machine Learning, 2023.
[20] Jeongyeol Kwon, Dohyun Kwon, Stephen Wright, and Robert D Nowak. On penalty methods for nonconvex bilevel optimization and first-order stochastic approximation. In International Conference Learning Representations, 2024.

---

[21] Raymond Li, Loubna Ben allal, and Yangtian Zi et.al. Starcoder: may the source be with you! Transactions on Machine Learning Research, 2023.
[22] Chin-Yew Lin. ROUGE: A package for automatic evaluation of summaries. In Text Summariza- lion Branches Out, 2004.
[23] Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang. Jing Jiang, and Min Lin. Regmix: Data mixture as regression for language model pre-training. In International Conference on Learning Representations, 2025.
[24] Dagréou Mathieu, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. A framework for bilevel optimization that enables stochastic and global variance reduction algorithms. In Advances in Neural Information Processing Systems, 2022.
[25] Swaroop Mishra, Daniel Khashabi, Chitta Baral, and Hannaneh Hajishirzi. Cross-task gener- alization via natural language crowdsourcing instructions. In Association for Computational Linguistics, 2022.
[26] Maher Nouiched, Maziar Sanjabi. Tianjian Huang, Jason D Lee, and Meisam Razaviyayn. Solving a class of non-conver min-max games using iterative first order methods. Advances in Neural Information Processing Systems, 2019.
[27] Rui Pan, Jipeng Zhang, Xingyuan Pan, Renjie Pi, Xiaoyu Wang, and Tong Zhang. Scalebio: Scalable bilevel optimization for Ilm data reweighting. ar Xiv preprint, arXiv:2406. S141:22006119766. 2024.
[28] Garima Pruthi, Frederick Liu, Satyen Kale, and Mukund Sundararajan. Estimating training data influence by tracing gradient descent. In Advances in Neural Information Processing Systems, 2020.
[29] Amirreza Shaban, Ching-An Cheng, Nathan Hatch, and Byron Boots. Truncated back- propagation for bilevel optimization. In International Conference on Artificial Intelligence and Statistics, 2019.
[30] Han Shen and Tianyi Chen. On penalty-based bilevel gradient descent method. In International Conference on Machine Learning, 2023.
[31] Daria Soboleva, Faisal Al-Khateeb, Joel Hestness, and Jacob Robert Steeves Nolan Dey Open- tensor: Robert Myers. Slimpajama: A 627b token, cleaned and deduplicated version of redpajam.
cartaned-and-aaee and- reddddddullicated
reddpaaa 2023.
[32] Ross Taylor, Marcin Kardas, Guillem Cucurull, Thomas Scialom, Anthony Hartshorn, Elvis Saravia, Andrew Poulton, Viktor Kerkez, and Robert Stojnic. Galactica: A large language model for science, arXiv preprint, arXiv:2211 1.09085, 2022.
[33] Yee Whye Teh. On big data learning for small data problems. KDD '18, page 3, New York, NY, USA, 2018. Association for Computing Machinery.
[34] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. arXiv preprint, 1rXiv:2302.13971, 2023.
[35] Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaci, Anjana Arunkumar, Arjun Ashok, Arut Selvan Dhanasekaran, Atharva Naik, and David Stap et al. Super-natural instructions: generalization via declarative instructions on 1600+ tasks, In Conference on Empirical Methods in Natural Language Processing.. 2022.
[36] Quan Xiao, Songtao Lu, and Tianyi Chen. generalized alternating method for bilevel learning under the polyak-Eooasiewiizz condition. In Neural Information Processing Systems, 2023.
[37] Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoe Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. In Neural Information Processing Systems, 2023.
12

---

[38] Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. models via importance resampling. In Neural Information Processing Systems, 2023.
Data selection for language
[39] Wanyun Xie, Francesco Tonin, and Volkan Cevher. Chameleon: A flexible data-mixing framework for language model pretraining and finetuning. In International Conference Machine Learning 2025.
[40] An Yang, Baosong Yang, and Binyuan Hui elal. Qwen2 technical report. arXiv preprint, arXiv:2407.10671, 2024.
[41] Jiasheng Yc, Peiju Liu, Tianxiang Sun, Jun Zhan, Yunhua Zhou, and Xipeng Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. In International Conference on Learning Representations, 2025.
[42] Mingyang Yi, Lu Hou, Jiacheng Sun, Lifeng Shang. Xin Jiang, and Qun Liu. Improved ood generalization via adversarial training and pretraing. In International Conference on Machine Learning, 2021.
[43] Mingyang Yi, Ruoyu Wang, and Zhi-Ming Ma. Characterization of excess risk for locally strongly convex population risk. Advances in Neural Information Processing Systems, 2022.
[44] Mingyang Yi, Ruoyu Wang, Jiacheng Sun, Zhenguo Li, and Zhi-Ming Ma. Breaking cor- relation shift via conditional invariant regularizer. In International Conference on Learning Representations, 2023.
[45] Longhui Yu, Weisen Jiang, Han Shi, Jincheng YU, Zhengying Liu, Yu Zhang, James Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. In International Conference on Learning Representations, 2024.
13

---

A Convergence of Bi-level Data Mixture Optimization
A.1 The Proposed Algorithm
As mentioned in the main part of this paper, the data mixture optimization problem is bi-level. Here
we present theoretical analysis of the proposed TANDEM. The original bi-level optimization problem
P: min Lval (w (a:) s.t. w (a) E (a) == arg min Lirain (a, w)
(6)
aEA
is difficult to solve owing to the imposed hard constraint. We turn to the Lagrangian problem of P such that
min Lval (w) + (Cerain w) min Cirain ( (a,, w)
aελ,
There is vast variety of existing literature, e.g. [30 19 36 discuss the relationship between the Lagrangian problem Py and the original problem P. In a word, when taking sufficient large, the solution of P., will approximate the solution of P. Thereafter, we develop the algorithm to P, to solve P. as we did in this paper.
Algorithm Twin Networks for bi-level DatA mixturE optiMization (TANDEM)
Input:
Train set Drrain validation set Deal comprised of M domains.
Episode number T, Episode length E, Probing length for each episode K, Learning rate Iju, ?uu iJa for (reference), (proxy) and a (mixture) respectively. Initialize proxy model parameters u and domain weights a.. T'
1:
2: for
Mixture ratio a update.
3: Set ( un
(t) for k 0 to
4:
do
5:
uk+1 u
[+)
(8)
T/u V.. Cutaram (a uxx))
6:
w (t)
wki1 w Tho (o. Lval (w (t)) + Lamin(( (,(
(t)
7:
end for
8:
1:M a (1+1) = ll(( N/
xxx
Cl:M (ae uk
reference mudel
provy model
I/ Free model weights update.
9:
for to E 1
u the 1/u Va (441) , ).).
11: end for
12:
Set u ((+1) ug
13:
end for
Output:
domain weights a (?).
full workflow of Twin Networks for bi-level DatA mixturE optiMizatio (TANDEM) is shown in Algorithm[T] TANDEM alternates between updating the mixture ratio a and the proxy model u, whereas u is used to approximate the minimum Ltrain (a, w). Update of the mixture ratio requires probing the data efficacy of each domain. During probing, we optimize the reference model w, as well as the proxy model u for K steps. u and u are respectively trained on the train set and the train set plus validation set, their incurred loss gap captures the gain of incorporating the additional validation data. TANDEM then up-weights domains that benefit more from additional data. Notably, the proxy model u and reference model w synchronized the beginning of probing. Since the updating of a requires u, w probing, we decrease the frequency of updating a to reduce the computational cost, and leave the u trained freely for E steps before the next update.
A.2 Convergence Rate
Next, we explore the convergence rate of the proposed Algorithm T For the original problem P ( its convergence property under the iterates obtained by solving the Lagrange problem 7 has been
14

---

explored, under different regularity conditions 30 19 In this paper, we will present our results under
mild Polyak-Lojasiewiez (PL) condition 17 study the bi-level optimization problem. Technically, we will impose the following Assumptions to derive the convergence rate. Before illustrating our Assumptions, we need the following definitions to simplify the notations. We denote
26 42), which has been widely imposed to
H.y w) Lma(w) + 7 (Crrain(ae w) min Lerain wi),
(8)
with St (a) (w w E arg min (a, w)), and (0xa)) {w w E arg min Lerain (a, w)).
We further define
H(a)
inf (E)-(a)
(w)(
where H(xx) is well-defined since Lirain (or, w) is continuous to w. Besides that, the minimum of H(aa) is exactly the solution original problem P (6)
Assumption 1 (PI. condition). For any a A, both Ltrain (ax, w) and Hy(ax, w) satisfy the PL
inequality with coefficient and pay, respectively. That says:
Liaain (oe, w) -min Lerain (ox, w) S 2/
V Wuwatrinaa, w)ll
(10)
and
1
Hy (a, w) min (a, 1 (H, 2p
(11)
40
Moreover, the coefficient μ4 satisfies t Nr =
Assumption (Smoothness). For any a A.A.
1. Both Vu w) and (w are Lipschitz continuous to a (hold for V to Luraaail( (oe, w)) and w coefficient L.
2. For any E A, both Ciraia(a, (a, w) and Lval(w) are Lipschitz continuous with coeffi- cient B.
Assumption 3 (Bounded Hessian). For any a E A. w E (a)), there exists positive constants X, p. satisfying Hessian matrices V Vew (a, w) >= X and Va Ptraiii w) 3
Assumption 4 (Lipschitz Hessian). For any E A, Lerain(Or, w) is twice-times continuous differen- riable, and the Hessian matrices Cerain w), V Veo Strain (o, w) are all Lipschitz continuous with coefficient H.
Assumption 5 (Bounded Loss Function), The non-negative loss function Lirain (ox, w), (val(w) are uniformly bounded by positive constant D.
Notably, for the bounded Hessian condition, due the structure of Lussin (or, w), by the lower bounded V wow Cu (w) and an upper bounded V Ctrain (w) for any m. Moreover, it worth noting that for bi-level optimization problem, is standard impose some regularity conditions to the Hessian matrix. For examples, the smooth Hessian in [30 19 and lower bounded Hessian in
be implied
19 Therefore, the imposed bounded Hessian Assumption can considered as a mild condition.
Next, we present a lemma to characterize the gap between the gradient of VaH((x) and V. V.Hlaa, w).
Lemma 1. Under Assumptions
1 for any given and (a) E (a))),,, it holds
(aal(a)= da
0 H, (o, w (a))
1
HB2 122
(2+1)
pLB
(12)
Proof. Without loss of generality, suppose that H((x) = Lvail (w" (ax)), due to the chain rule, we have
H.H(a) VaLval( (a)) Va (ax) T (u (a:)..
(13)
For matrices A. A À means A Ad positively semi-definite matrix.
15

---

To simplify the notation, we denote (a) and w" (a) as wo; and w" in the sequel. For the penalized problem, given any (a), we have
0
Hy (a,, wx,)
da
Va (a, wi,) + Vaw, V, Vaw, w
Va Ha, wi,)
Y (VaLtrain (a, w.) V aitraiu (a; w²) Vaw" Vw w²))
(14)
= VaCiiain(aa, w. V Ltrain (a, w²) V Cawainn w²) (w))(w - w w²))
- Vaw" (V "wwla("") + +V Curaiu (a, w"))
+ Va vwwvall"" V Ctrain w (w, w").
Besides that, we have
w (VwLyi(W"") + +VV Cerain (erriin((, (ce, w"))
= Vaw' (V (VwLralw") VworallW.)
+ +vaww" V (e) V wa w.)) + V uwn (a, wn) w")(w" w")) (15)
+ + aww Ferain (a, w" (w,, - ww):
due the optimal conditions V Fra + V Ctrain (a, w and V Fimin (a, w*)
V Vaw*
V TV wow ftrain w²) = 0. Combining the two above equations and 13), we get
H. (a, (;; - -. -VaH(a)
da
y (VaCrain wi) - Cererain (ac w²) V Ciraiiaaaa w (w, - w"))ll
Naw (VwLtrain(a,, w') - Cw w;,) + V2 Cirain w (w,; w
Vaw-T (VwLval(w") (w
pl: 1 w. S SyH ( +1) - A
(16)
due the bounded Hessian Assumption Smoothness Assumption and Lipchitz Assumption Then, due to the PL condition and Smoothness Assumption we know there exists w (the projection of w" to (a)) satisfies
1/00 -w; <
V Lirain w:)]l
1
(1I w Lval(ww) + + V Ctrain (w; + V/V (w))))
7/1
(17)
V Luaa (w;))
7/4
7/4
Combining this with inequality (16). we obtain the conclusion under such wt. Finally, due to the Lemma A.5 in [26]. we know that V. H., (or, wx,) is invariant over w E (*)a),, we prove our conclusion.
From Lemma 1 we know that the gap between the gradients of original problem H(a) and its Lagrange version H. (a, w (ax) can be extremely small by invoking penalty parameter y - co. Thus, it implies that we can compute the gradient of the Lagrange problem to implement the gradient- based method. Next, we illustrate a useful lemma, which characterizes the Lipschitz continuity of w (a) and w" w((a) to a.
16

---

Lemma 2. Under Assumptions
7 andd for any a1, E A, it holds
1|w" (ax) w" (a2))l - L || a2l1,
(18)
/
for any w (a1) S* (a) and w (012) E (a2) satisfies w" (xx2) arg min | w" (a1))l- On the other hand, holds
||a ||a
|e (ax1) (2)()l v
a2l
(19)
H4y
for any wx (a1) and w., (ax)) arg minnwestaan |/w w (ai))l
Proof. The two conclusions
directly obtained from Lemma A.3 in [26]. prove the second
conclusion as an example. The proof can be similarly generalized to the first conclusion. Due to the formulation of V Hyy (a,, w). we know
V w (01, (a2))l = [ Hy w (a2)) - V (a (22))
= > V Surainn (a1, w (a2)) V Curaian w (a2))
(20)
I + S -2ll:
by Assumption 2 On the other hand, by invoking the Assumption and I17 we get
V Wa w (o2)) = V (w (az)) + v Straiu (01, w (a2))
(21)
2 H4 || w (a2) - w (a1)
where w (a1) ) the projection of w" (o2)) (a1))) By combining the two above inequalities, we obtain our second conclusion.
From this lemma, we can obtain the following Lipschitz smoothness of V H, w (a)) w.r.t. a by the following Lemma. It worth noting that VaHH (a)) is invariant over wi (ax) E (7) as discussed in the proof of Lemma 1 Thus, the Ha (a, (a)) is well-defined.
Lemma 3. Under Assumptions !
lndd 2 VaHl (a, w (a)) has semi-Lipschin gradient such that
I
(i)
1.
aaHtyan (",,,, - VaHt,(ay (a2)) SBB
+
||cn
azll-
(22)
Hy 1
bn
Proof. From
14).
we see
|| V (a wi - V Ha w (a2))
S 7 VaLirinn ( (a:)) VaLtrain (02, w (ar))
+ 7 VaLtrain (a w² (a2)) Va w (a:))
(23)
$ BB (1lw, (a:) w (2z) + ( - w(2)))l
BB b
4) ||on - a211,
My 11
which proves conclusion.
Notably, for the original optimization problem, there exists a constraint a E A. Thus, to prove the first order stationary condition of constrain problem min Ha. we consider the generalized projected gradient stable condition, i.e., 1 ||a - (a(()) na V (() < for some small tue
positive This is standard first-order stationary condition for Non-convex optimization problem
with constrains 30 42 42.9. 9.. Due Lemma
0 we know that
(0) lll(aa (1) nl V.H( ())
| - llla na VaHa(a (a (1)
aa
[I Te V.Ha(( (a no V. H. wa
To
(24)
< (D) - Vath(a (
5 0 ())
17

---

when + co. Here the second inequality is from the Lipschitz continuity of projection operation [9] Thus, the above inequality indicates that to prove the first-order stationary condition of H(a). it is sufficient to prove the first-order stationary condition of (a, w). Next we present our main theorem, i.e., the formal version of Theorem which shows the convergence result of (a)) under Algorithm|] [
Theorem 1. For sufficiently large T, under Assumptions for the a obtained in Algorithm
holds
min -Iaa (t) 1la(a na =o
(25)
ISST 7a
by selecting = T K
log
min
log 2-2-1
j. 1)ax = 1
Thn the I
log
ply
log (-
The L.
Proof. As mentioned in (24). is sufficient to prove the first-order stationary condition for H- w)
Due to the Lipchitz smoothness it, we have
Hy (a (++1), w (a(())) H, w),(a")))
S (a wi )). a (++1) (1)
+ L
lla (t+)) (1
2
= (VaFi(a
[++] a (1)
(vaHh(a!" ) V
Va (a
10 u1 (1)) a (++1) -()1))(())))))) (a)
(26)
L (110)) -a00
2
< L |
Ja - =000
V. (a wi - VaF wx 22
2//a
2
where the last inequality is due to the property of projection operator and Jensen's 's inequality. Let us define ((+1) = (I Na Va H,,,( (a w (c)()))) we proceed to upper bound the gap
between (1-1) ( and (ta (1) ||. By the triangle inequality
||a (+1) -1 laa (4t) - -a 1
$ 11a ((1) - -a 11
II (a To VaJ (a m x x2) 11A (at)
no (a (e
(27)
S M/w VaJ (1) wx x V. H, (a (1) (ce (<'
S B wi (a (2)l + 7 Bl (t) - w (a )l,
where the last inequality is from the smoothness Assumption tively the projection of w K and (0) to (a(()) and ("(a(()) Firstly, due to the Pl-condition and Lipschitz smoothness of Liraiu. we have
2 w (a) and w (a) are respec-
Erain (a (9) (uttt) Lirain (a ((( <
<< (V (anaiiai (u(() u u(())
T. (1 (0
2
2L VerLerainaa (a (1 )12
- L (Cania((a!)) (t) (1 Lirain (a((t) w (a (0)))
14
(28)
which implies
(1) w (a
#
K4 L.
Ctrain(a ((.1)) Luainu( w0 (a (1)
11
K+
<
D
(29)
L
1
T
18

---

due to the selection of K and E and the PL condition of Ltrain (ar, w) (which implies the error bound condition [7].. Similarly, we can prove that
1 T,
K (Hy(a ,0 Hyyae H,,a"",,;(a"))
14
K
1
Hz
(1+7)
D
(30)
lny
o
Combining (26 27 29
and 300g we have
aa(l1)] - al0l
1a722)
H4 (1+1) w (a (1+1)) (a (al))) so that, by combining
+0
(31)
47a
T
we get
T
(0 T
a
(attt)) - a 1001| + - -I
(++))
|| -a ||a (1+1) - ==0
nl
t=1
Ha (0), (a
()))
infa H. (a +
(22
aa T
T
o T
(32)
Due to the definition of (++1), combining this with (24). and the value of y proves our conclusion.
B Proof of the Proposition
Proposition 1. Assume C
Cua the uniform mixture ratio ām M
1 for m 1, I,...,M M
constitutes a valid solution of
Proof. Let Lm
Extata [
= L To establish the above, it suffices to show:
C ((((w(())) (ā))
C (ww (a:)
m=1
m=1
which is equivalent to:
M
M Σ
am cm (w (ax) S
am
(w (a))
(33)
m=1
m=1
Because u² (ax is the minimizer of the inner-level problem under uniform weighting, we have:
am L. Lm (a)( <
Σ am
LT (w)
(34)
m=1
for any w. In particular, by choosing w w" (ax), the right-hand side becomes the expression in 133 , completing the proof.
C Implementations and Hyper-parameters
We use the gradient descent optimizer for the K step u and w up- date. Their learning rates w/th and w are kept the same that the domain-wise loss difference Lm (wx)) L.m (ux) faithfully reflects the gain of incorporating the additional validation data. For the E step update, utilize the Adam optimizer for fast training.
Implementation Details
19

---

Table 6:
of TANDEM for different application scenarios
Data Aundent
Data Restricted
Supervised Fine-tuning
GPT-like 160M
160M
410M
1B
Qwen2-0.5B
Batch Size
128
128
128
32
Learning Rate T/u
5e-5
5e-4
Se-4
5e-4
4e-6
Learning Rate
le-2
le-2
le-2
le-2
1e-2
Learning Rate The
le-2
le-2
le-2
le-2
le-2
Learning Rate Tlxx
2e-3
4e-3
4e-3
4e-3
4e-3
Learning Rate Scheduler
Cosine
Cosine
Cosine
Cosine
Cosine
Penalty Y
1.0
1.0
1.0
1.0
1.0
Probing Steps K
5
5
5
5
10
Free Training Steps E
Total Steps (w.r.t u))
40000
5000
5000
5000
5000
Context Length
2048
512
512
512
512
Weight Decay
le-2
1c-2
lc-2
1e-2
1c-2
Gradient Clipping
1.0
1.0
1.0
1.0
1.0
Table 7: Statistics of the SFT data.
Table 8: Statistics of the task-level SFT data.
Method
Num. Sample
Method
Num. Sample
Textual Entailment
79332
SQuAD1.1
6498
Answer Verification
13195
AMRSum
6500
Text Matching
47297
MuTual
6500
Information Extraction
31053
SemEval
5996
Word Extraction
17294
SST2
6495
Text Categorization
89572
BoolQ
6500
peeeer-paraeete settings The detailed hyper-parameters for the TANDEM algorithms in the con- ventional data-abundant pretraining scenario, the data-restricted training scenario, and the supervised fine-tuning are shown in Table
D Supervised Learning Data Statistics
In the supervised fine-tuning case, we use 6 major categories from Natural Instructions 25 35
Textual Entailment, Answer Verification, Text Matching, Information Extraction, Word Extraction. and Text Categorization, each constitutes a task cluster. This corpus comprises 99 tasks ranging from open-ended text generation, multiple choice, and True/False tasks. The statistics are given Table 7
Besides, we delve into a more fine-grained task-level SFT case, and select tasks SQUADI.I, AMRSum, MuTual, SemEval, SST2, and BoolQ. The statistics are given in Table 8 For the text generation tasks SQuAD1.1 and AMRSum, we report the Rouge-L [22] score. For the multi-choice tasks (MuTual, SemEval) and Yes/No tasks (SST2, BoolQ), reported as well. We experiment with Qwen2-5008 500M with K 20, E = 10, batch size 32, and context length 512 for 2000 steps. The result is shown in Table 9 In this case, the improvement in test loss is more significant than that in the final evaluation metnes, likely due to the imperfect alignment between them.
focus on the accuracy. The test loss
Table 9; Comparison for the 500M Qwen-2 model on a mixture of 6 tesk-level SFT tasks.
Method
SQUADI.1
AMRSum MuTual
SemEval
SST2
BoolQ
Test Loss
Avg. Metrie
Uniform
72.62
45.18
72.72
89.75
87.75
80.29
0.591
74.72
DoReMi
71.40
43.40
70.97
88.50
87.00
77.43
0.686
73.11
DoGE
71.26
44.81
71.67
89.00
88,30
81.04
0.563
74.35
Skill-It
72.14
44.21
73.07
89,40
88.40
80.34
0.539
74.60
Aioli
72.35
45.01
73.57
89.10
88.10
79.64
0.542
74.63
TANDEM
72.73
45.19
73.77
89.70
88.70
80.04
0.508
75.03
20

---

Table 10: Comparison of models of different sizes in the Lata-abundant pretraining scenario and data restricted scenario with standard deviation.
Data Aundant Regime
Data Restricted Regime
160M Avg.
160M Avg.
410M Avg.
1B Avg.
Uniform
25.74.0.13
31.53
29.59
29.91 10.09
DoReMi
28.32
36.91 10.09
54.61 110.60
56.53
DoGE
26.60+0.21
30. 10.005
27 27.45+00.02
Skill-It
25.87
29.24
27.70+00.03
27.15
Aioli
25.6640.14
31.19
28.7940.06
28.07 +0.03
TANDEM
25.43 10.15
28.07 10007
25.00 10000
24.35 1:0003
Data Abundant Pretraining.
Gap
Pretraining.
Gen. Gap Supervised Fintuning.
1.00
0.40
ab
0.30
New
0.40
0.20
0.05
0.10
0.90
0.00
0.00
1000
2000
3000
6000
0
1000
2009
3000
4.00
1000
2000
2009
4000
Stops
evolution under three scenarios. (a) data- abundant pretraining (b) dettrrestd pretraining and (c) supervised fine-tuning.
Figure 9: Step-wise generalization gap 1 train
E
Comparison with Standard Deviation
We test each method in Table 3 3 times. The averaged perplexity and standard deviation are reported in Table
F
Additional Experiments
F.1
Generalization Gap Analyses
To validate analysis in Section 2.4 plot the change of generalization gap 14 Ch during training under the three circumstances Figure 9 In the data-abundant scenario, the generalization gaps of each domain do not change much and are kept close to 0. According to the analysis in Section 2.4. Uniform is valid solution in this scenario. the data-restricted scenario, we see an increase in the generalization gap. particularly repetition kicks in. In the large domains (C4, CommonCrawl) where there is no data repetition, the generalization remains small. SFT is of the (Word Sementics, Information Extraction), while kept small for large tasks (Text Categorization, Text Matching, Textual Entailment). The results all three scenarios validate analysis in Section 2.4
small domains (Arxiv, Book, Wikipedia), because the sample
case. Generalization gap increase in small tasks
F.2 Larger Scale Experiments
To validate the effectiveness of TANDEM in practice, we conduct experiments to test its performance under more realistic conditions. This includes scaling up both the dataset size and the model size.
Table 11: Comparison for larger data scale (6B tokens).
Arxiv
Book
C4
CC
Github
Stackexchange
Wikipedia
Average
Uniform
10.28
32.37
38.53
34.57
5.86
10.76
15.26
17.09
TANDEM
10.26
29.59
32.52
29.52
6.01
10.68
16.02
16.25
21

---

Table 12: Comparison for larger (3B) model in the data restricted regime.
Arxiv
Book
C4
CC
Github
Stackexchange
Wikipedia
Average
Uniform TANDEM
17.20
53,83
60.14
55,45
9.85
22.08
41.76
31.03
15.73
42.85
43,60
39.59
7.49
17.26
29.20
23.81
Table 13: Comparison for the larger (3B) model in the supervised fine-tuning.
Method
Textual Ent. Answer Ver. Text Mat. Inf. Ext.
Word Sem.
Text Cat.
Avg. Metrie
Test Loss
Uniform
91.9
74.8 76.9
88.8
79.2 80.1
88.8
85.5
84.9
0.189
TANDEM
92.4
89.6
88.0
87.9
85.8
0.174
Larger Data Scale We train the 160M models on the full 6B version of sampled SlimPajama, which constitutes a more practical past Chinchilla [14] (Chinchilla optimal requires =3.2B data for the 160M model) case. From Table IT
that TANDEM still outperforms Uniform, though it seems that the improvement (~ 1) in this case not as significant as in the 300M data case (~ 3.5). As the training proceeds and the perplexity goes down, further gains become inherently harder to achieve. This result is still significant, showing that careful mixing still matters in this "over-trained"
Larger Model Scale We conduct DMO with 3B models in the data-restricted scenario (Pythia 2.9B) and SFT (Qwen-2.5-3B)) Further scaling up necessitates engineering upgrades to fit the two models (u and w) within the 80GB memory, while also demanding substantial time and computational resources. We skip the data-abundant scenario where Uniform is already valid solution. From Table 12 we see that as the model goes larger, inappropriate mixture ratios (Uniform) lead to more obvious negative outcomes(e.g. severe overfitting on the overly sampled small domains). While TANDEM consistently generates proper mixture ratios. For SFT (Table 13), TANDEM still outperforms the Uniform baseline, showing its effectiveness.
F.3 Sensitivity of K and
K: K' is the number of probing steps used to estimate the proper update direction (33 We conduct a sensitivity analysis with performance. From Table we see that too small K may results in less fidel a update direction, thus sub-optimal data mixture ratio and higher perplexity. When K large enough for update probing, increasing K will not induce further benefits. In our experiments, for relatively stable pretraining, we use a smaller = 5 in the data-abundant and data-restricted scenarios. The gradient variance in SFT is much higher, so we use larger K = 10.
1,3, 5, 10 to see how will affect the final model
Given the total number of training steps (amounts of data to be used), E' determines the number of updates during DMO: T total steps We test E 1, 5, 10 in the 160M data-restricted scenario. From Table [15 we see that TANDEM is not sensitive to as long as T is large enough so that a is sufficiently updated. In experiments, for the data-abundant scenario, we chose F. 20 that the mixture ratio is updated for T = 2000 steps. As one a update requires additional steps w updates, larger T means higher additional computational cost. To reduce the computational overhead, set E in the data-restricted scenario and SFT to ensure T 1000 and = 500 respectively.
E:
Table 14: The effect of K
Table 15: The effect of E
k=1
k=3
k=5
k=10
E=1 E=5
E=10
Perplexity 29.57 28.31
28.07
28.05
Perplexity 27.99 28.05
28.56
22

---

Table 16: Downstream evaluation after training on SlimPajama in the data-restricted scenario.
Method
ARC-C
ARC-E
BoolQ
HellaSwag
LAMBADA
PIQA
WinoGrande
Average
Uniform
25.76
19.58
60.70
18.62
9.88
45,43
50.04
32.85
DoReMi
21.36
19.58
60.61
11,66
6.44
42.94
49.57
31.06
DoGE
22.03
24.87
51.35
18.50
10.07
44.45
47.28
31.39
Skill-it
18.98
26.10
50.82
24.99
11.47
43.36
50.59
31.40
Aioli
20.34
19.58
61.71
19.59
10.89
44.18
49.41
32.41
TANDEM
20.34
20.81
57.55
25.59
11.20
40.53
51.07
32.92
Table 17: The effectiveness of the proxy model
Method
Data Abundant (Perpl.)
Data Restricted (Perpl.)
SFT (Acc.)
Uniform
25.74
31.53
74.72
Online (Proxy)
25.72
29.63
74.49
Two-Stage (default)
25.43
28.07
75.03
F44
Downstream Tasks
We also evaluate the 160M model pre-trained with SlimPajama data ARC-C, ARC-E, BoolQ, HellaSwag, LAMBADA, PiQA and WinoGrande using the Language Model Evaluation Harness g From Table
16
that TANDEM outperforms all the baselines, validating its effectiveness.
F.5 The Effectiveness of the Proxy Model
might be expected that the online-trained proxy model could also perform well. We evaluate the proxy model u, which has been trained with adaptive domain weights (1) in all the three scenarios, and compare it to the default two-stage (DMO and then train) trained model. Counter-intuitivele the online trained model falls behind the default two-stage trained model as well as the uniform baseline (Table This result coincides with the findings in previous works, e.g., DoReMi and DoGE. We hypothesize that the frequent change of data distribution deteriorates the training process.
the performance
F.6 A priori Trained Reference model Restricts DoReMi
In the experiments, DoReMi 37 doesn't perform well in the data-restricted scenarios. We hypothe- size that the performance of DoReMi is restricted by the a priori trained reference model. Following previous works, we train the reference model on training data obtained using the Uniform strategy. This reference model, however, is not well-suited for the data-restricted scenario, as the many small domains overfit and cannot provide faithful signals for DMO. For more comprehensive comparison, we train another reference model with the natural data mixture ratio (See Figure the data statistics). This setting ensures no severe overfitting happens. The result in Table T8 validates our hypothesis.
G Visualization of the Optimized Mixture Ratios.
We visualize the mixture ratio learned in each application scenario in Figure 10 Figure I1 and Figure 12 For the baselines, the average proportion over the entire training trajectory is taken. While for TANDEM, we the average mixture ratio at the last 10% training trajectory.
Table 18: DoReMi is restricted by the priori trained reference model
Method 160M
410M
1B
Uniform
31.53
29.59
29.91
Natural
30.97
27.82
27.30
DoReMi--attral
30.83
27.26
26.07
TANDEN
28.07
25.00
24.35
23

---

DoRaft
DOGB
0.4
Mixture Ratio
0.3
0.2
0.1
wikipedia
Figure 10: Mixture ratio learned by different methods in the data-abundant pretraining.
DoReft
0.4
DOGE
Skillin
Mixture Ratio
0.3
0.2
wikipedia
(a) 160M Model
DoRaM
0.5
TAND
Mixture Ratio
0.3
0.2
0.0
andi
books
c4
github
stack.
wikipedin
(b) 410M Model
0.6
Suilitt
0.5
TAND
0.4
Micture
0.3
0.1
0.0
arxiv
books
64
github
stack.
wikkpedia
(c) 1B Model
Figure 1I: Mixture ratio learned by different methods the data-restricted pretraining.
24

---

0.35
Dohalt
0.30
poos Skil--tt
Mixture Ratoo
0.10
0.05
0.00
Ent.
Figure 12: Mixture ratio learned by different methods in the supervised fine-tuning.
H
Summary of Notations
Table 19: Summary the notations used throughout this paper. Variables only used theoretical analysis are grayed for better readability.
Topic
Notation
Explanation
Data Sets
M
The number of domains.
Dm
The m-th domain data. Train set.
Dimin
Dall
Validation set.
Models & Parameters u
Parameters of the proxy model.
Parameters of the reference model.
w
Optimal solution for the lower level problem.
w
Optimal solution for the penalized problem. Solution set for the lower level problem.
S*
(a)
S (a)
Solution set for the penalized problem.
Data mixture ratio The probability simplex.
A
Problems & Losses
Lm
Train loss on the m-th domain.
train
Val
Validation loss on the m-th domain.
Luain
Overall train loss weighted by the mixture ratio
Lval
Overall validation loss.
H(e)
The upper-level loss with the lower-level problem optimized.
H. (o,, w)
The loss of the penalized problem.
Function Properties
1
PL coefficient of the lower-level problem. PL coefficient of the penalized problem.
L
Lipschitz constant for V. Crain (a,, w) W. Lipschitz constant for Cuain (a, w) and Lval (w) on w.
À and p
Hessian Curinn (or, w) č À and V va (a, w) Lipschitz constant for V2 22.0 Curain (a,, w) and V Vinin Surin (a, w). Upper bound for the train/validation loss. Lipschitz constant for V., Hy (a, (a)))
Ly
Train
Mixture ratio training step The total number of a update.
T
k
and ur update step.
Free u update step.
The number of u, w probing update for one a update. The number of free update.
E
Ja
The learning rate on a. The learning rate
N
The learning rate on w.
7
The penalty strength.
25

---

I Limitations and Future Works
Despite the advancements introduced this work, several challenges remain open for future research. The limitations of this paper are as follows: (1) Due to limited computational resources, our ex- periments were conducted using models of up to billion parameters. Although our experimental results demonstrate TANDEM's ' effectiveness at this scale, the constraint on model size limits our ability verify whether our findings generalize to significantly larger models, such as those with 405 billion parameters. (2) Our current data mixture optimization (DMO) approach is validated naturally occurring domain splits. For example, the SlimPajama corpus consists of seven domains sourced from different origins. The impact of using more fine- fine-grained and intentionally designed domain splits on DMO performance remains unexplored and presents an interesting direction for future investigation.
coarse,
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: This paper focuses on data mixture ratio optimization in LLM training. The abstract and introduction sufficiently summarize the paper's contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitations of the paper are discussed in the "Limitations and Future Works" section in the Appendix.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
27

---

Answer: [Yes]
Justification: The full set of assumptions and complete (and correct) proof' are provided in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The implementation details and Spper-parameters settings are given in the Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
28

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification:
The code will be released upon the compliance review process of the company.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.com cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ praips.cc/ubblic/guuueee CoodeSumeeieeeeee ssi iessonPolll mollicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The core experimental setting is outlined in the paper, and the details are given in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We have provided the 95% confidence interval of the DMO process and the standard deviation of the test results in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
29

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The type of compute workers and time of execution are provided in experi- ments.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't I make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:////aaarrseee//o/aa///a/aaa
Answer: [Yes]
Justification: The research conducted conforms with the NeurIPS Code of Ethics in every respect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
Answer: [No]
Justification: The goal of this paper is to build better understanding of data mixture optimization and thus improve the LLM training. doesn't perform negative societal impacts as long the LLM is properly used.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
30

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper doesn't 't provide pretrained language models, image generators, or scraped datasets.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging. and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The original papers that produced the dataset are properly cited.
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
31

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
32

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development this research does involve LLMs.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy bettts://neurre /0tps//eec.o0ooeeeeee////2////// for what should or should be described.
33


Reviews:
[{"review_id": "TOssotV5qG", "reviewer": "Reviewer_3KfP", "summary": "This paper tackles the problem of reweighting domains for language model pretraining. The key insight is to cast the problem as a bilevel optimization problem and used penalty-based to solve using gradient updates (as opposed to requiring Hessian information). Theoretically, they prove convergence to a first-order stationary point. Empirically, the method is shown to outperform previous methods such as DoReMi and DoGE for data limited settings (for normal settings, the uniform baseline is hard to beat) and SFT.", "questions": "Given that the u and w are meant to be close to each other, what happens if you run the same algorithm but with one copy (u = w)?\n\nline 127: typo \"constrain problem\" => \"constrained problem\"\n\nline 157: \"by solving bi-level problem\" => \"by solving the bi-level problem\", etc.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The problem of automatically selecting data mixes is of critical importance in language model pretraining. The paper presents a principled formulation based on bilevel optimization and show a practical algorithm to actually solve the problem. The formulation nicely and automatically captures important intuitions about what happens in the data limited regime and as the model size grows.\n\nIt is unclear whether the experimental settings reflect real-world LM training regimes. For example, 300M is a very small number of tokens for a 1B model (Chinchilla optimal would be already 20B tokens). Of course we are sometimes data limited but only on some domains, so is it possible to set up the experiment to have a much more sizable token count but have some rare domains that really matter for downstream evals? Also, given that most small models are overtrained (past Chinchilla), I'd be curious how the different methods perform; does careful mixing still matter? I am also quite curious about how this scales to 7B parameters and beyond, but of course I don't expect the authors to necessarily have the compute.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your recognition of our technical contributions. We address your concerns below: \n\n>Q1. Does careful mixing still matter when there is more sizable token count?\n\nThanks for this insightful question. Generally, DMO is important as long as we are seeking a powerful and balanced model, even there are huge amount of tokens. For the ease of understanding, assume that different model capabilities emerge from different domains and are governed by their scaling laws each (neglecting the inter-domain relationships). DMO is effective as long as the benefits of scaling up certain domains outweigh the losses incurred by scaling down the other domains. Besides improving the overall model performance, it can be especially critical to balance the different capabilities. \n\nTo validate this claim, we train the 160M models on the 6B version of sampled SlimPajama, which constitutes an overtrained (past Chinchilla) case: \n|Method|Arxiv|Book|C4|CC|Github|Stackexchange|Wikipedia| Average|\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n|Uniform|10.28| 32.37|38.53|34.57|5.86|10.76|15.26|17.09| \n|TANDEM|10.26|29.59|32.52|29.52|6.01|10.68|16.02|16.25|\n\nWe see that TANDEM still outperforms Uniform. Though it seems that the improvement ($\\sim$ 1)  is not as significant as that in the 300M data case ($\\sim$ 3.5), we want to clarify that as the training proceeds and the perplexity goes down, further gains become inherently harder to achieve. This result is still significant, showing that careful mixing still matters in this overtrained case. \n\n> Q2. Experiments on larger models.\n\nCurrently, the model scales (160M/410M/1B for data restricted case and 0.5B for SFT) are chosen following previous works, referring DoReMi (280M/510M/760M/1B), DoGE (60M/82M/124M), Skill-it (125M/1.3B) and Aioli (160M). The most recent work, released in May 2025, CHAMELEON, as mentioned by Reviever SkSi, also uses an 86M model for DMO. Nevertheless, we agree that scaling up the experiments would significantly strengthen the persuasiveness of our findings. \n\nDue to the limited computation and time, we conduct DMO with 3B models in the data-restricted scenario (Pythia 2.9B) and SFT (Qwen-2.5-3B). Scaling up to 7B+ models necessitates engineering upgrades to fit the two models ($\\boldsymbol{u}$ and $\\boldsymbol{w}$) within the 80GB memory, while also demanding substantial time and computational resources, so we resort to 3B experiments. We skip the data-abundant scenario as discussed in Section 2.4, Uniform is already a valid solution. The results are as follows:\n\n**Data-restricted scenario:**\n|Method|Arxiv|Book|C4|CC|Github|Stackexchange|Wikipedia| Average|\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n|Uniform|17.20|53.83|60.14|55.45|9.85|22.08|41.76|31.03|\n|TANDEM|15.73|42.85|43.60|39.59|7.49|17.26|29.20|23.81|\n\nWe see that as the model goes larger, inappropriate mixture ratios (Uniform) lead to more obvious negative outcomes(e.g., much more severe overfitting on the overly sampled small domains). While TANDEM consistently generates proper mixture ratios.\n\n**SFT:**\n\nAs suggested by the reviewer Vq6v, we use a larger SFT dataset. (Please refer to our response to question 3 for reviewer Vq6v)\n| Method  | Textual Ent. | Answer Ver. | Text Mat. | Inf. Ext. | Word Sem. | Text Cat. | Avg. Metric $\\uparrow$ | Test Loss $\\downarrow$ |\n|:-------:|:-----------:|:----------:|:--------:|:--------:|:--------:|:--------:|:-------------:|:-----------:|\n| Uniform | 91.9        | 74.8       | 88.8     | 79.2     | 88.8     | 85.5     | 84.9          | 0.189       | \n| TANDEM  | 92.4        | 76.9       | 89.6     | 80.1     | 88.0     | 87.9     | 85.8          | 0.174       |\n\nWe see that for a larger 3B model, TANDEM still outperforms the Uniform baseline, showing its effectiveness.\n\n>Q3: What happens if you run the same algorithm but with one copy (u = w)?\n\nWe can not simply set $\\boldsymbol{u} = \\boldsymbol{w}$. Actually, TANDEM deliberately deals with these two models. They are supposed to be close to each other as imposed by the Lagrange penalty, while they can not be exactly the same, as it is exactly the difference between them that motivates the update of $\\boldsymbol{\\alpha}$ (Equation (5)). So, if we maintain just one copy $\\boldsymbol{u} = \\boldsymbol{w}$, TANDEM will not return the optimized mixture ratio as expected. \n\n\n>Q3: Typos.\n\nThanks very much for your careful checking, we've fixed these typos and thoroughly polished the revision.\n\nWe thank the reviewer again for the review. If you have any more questions about our replies, or any further questions or concerns, please let us know."}, {"review_id": "P7asmStZgM", "reviewer": "Reviewer_jR7t", "summary": "The paper proposes TANDEM, a new algorithm for DMO in LLM training. It reformulates the bi-level optimization of domain mixture ratios into a single-level penalized problem, which is then solved using a twin-network approach: a proxy model trained on primary data and a reference model trained on additional validation data. TANDEM measures the impact of data from different domains by comparing these models, dynamically up-weighting domains that improve performance. TANDEM offers better stability, theoretical convergence guarantees, and strong empirical results—especially in underexplored settings like supervised fine-tuning and data-restricted regimes.", "questions": "1. Please provide a more detailed explanation of how Equation (5) is derived. The current presentation is abrupt and lacks sufficient intermediate steps for readers to follow.\n\n2. One of your central claims is that “data mixture optimization is particularly beneficial in scenarios with limited data availability rather than traditional pretraining setups with abundant domain data.” However, the last paragraph of Section 2.4 is mostly high-level speculation. It would be more convincing if you could include either a theoretical extension (e.g., following from Proposition 1) or empirical evidence that directly supports this claim.\n\n3. Many hyperparameters vary across scenarios (e.g., different values of $K$, $E$, and even context length). Could you elaborate on the rationale behind these choices? Is TANDEM’s performance robust to such variations?\n\n4. You argue that uniform mixing performs well in the first (data-abundant) scenario. However, even in the last two scenarios, DoReMi and DoGE consistently underperform the uniform baseline. I’m concerned whether these baselines were properly implemented and tuned.\n\n5. In the analysis of the probing step $K$, why do you fix $\\alpha$? Doesn’t this deviate from how the actual algorithm works? \n\n6. (Not affecting the score):\n(1) Some notations (e.g., $\\prod_A$) should be better explained.\n(2) In Line 155, it is misleading to refer to cosine similarity directly as “variance.” \n(3) For better clarity, please annotate all figures and tables (e.g., Table 3, Figures 5 and 6) with the corresponding scenario (data-abundant, data-restricted, or SFT).", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n1. The core idea of reformulating DMO as a single-level penalized problem and solving it via twin networks is elegant and conceptually appealing.\n\n2. I especially appreciate Sections 2.3 and 2.4, which clearly connect TANDEM with prior methods (e.g., DoReMi and DoGE), offering a unified view that highlights both similarities and improvements, and go beyond algorithmic design to discuss **when and why** DMO matters.\n\n**Weaknesses**\n1. The methodology section lacks clarity in several places, and many claims are made without sufficient justification or analysis. For example, (1) in Section 2.2, the derivation of Equation (5) could be explained more clearly; (2) the final paragraph of Section 2.4 reads more like speculation, with no analytical or empirical support and no follow-up on Proposition 1.\n\n2. The experimental setup relies on models that are too small to draw strong conclusions. Most experiments are done with models <1B, with only one section using a 1B model. While I understand the computational cost of DMO, these models are not sufficiently representative. It is also unclear why the SFT experiments use only a 0.5B model—does it really impose a high cost?\n\n3. The experimental results raise several concerns: (1) Table 4 does not report any statistical significance, and the observed differences appear too subtle to be convincing; (2) in all settings, several baselines (e.g., DoReMi and DoGE) perform worse than uniform, which is surprising and concerning. While the paper argues that uniform mixing works well in data-abundant settings, the poor performance of baselines in the other two scenarios warrants deeper investigation.\n\n4. The writing quality is unprofessional, with numerous grammatical and typographical errors throughout the paper, including L62, L99, L111, L115, L149, L151, and the caption of Table 2.\n\n----\n\nI would be happy to raise my score if the authors can adequately address the weaknesses and questions I raised.", "was_revised": true, "final_justification": "I believe the presentation of the paper has improved substantially after the rebuttal, and many of the previously qualitative arguments are now sufficiently supported. Therefore, I have raised my score. However, the use of relatively small models (understandable given the computational cost) and the statistical significance of some experimental results (e.g., Tables 2 and 4) remain concerns.", "rebuttal": "Thank you for highlighting the strengths of our method—its elegance, comprehensive scope, and its deeper discussion of when and why DMO matters. We believe this provides meaningful insights into the DMO problem. We address your concerns below:\n\n>Q1: The derivation of Equation (5) could be explained more clearly.\n\nIn the objective (2), only the $\\gamma (\\mathcal{L} _ {\\text{train}}( \\boldsymbol{\\alpha}, \\boldsymbol{w})- \\min_{\\boldsymbol{u}} \\mathcal{L}_{\\text{train}}(\\boldsymbol{\\alpha}, \\boldsymbol{u}))$ is related to $\\boldsymbol{\\alpha}$. To see how this objective induces the update rule Equation (5), we first inspect the training loss. By definition, it is the $\\boldsymbol{\\alpha}$ weighted domain-wise loss: $\\mathcal{L} _ {\\text{train}}(\\boldsymbol{\\alpha}, \\boldsymbol{w}) = \\sum _ {m=1}^M \\boldsymbol{\\alpha} _ m\\mathcal{L} _ {\\text{train}}^{m}(\\boldsymbol{w})$. The gradient w.r.t $\\boldsymbol{\\alpha}$ is exactly the domain-wise loss: $\\nabla _ {\\boldsymbol{\\alpha}}\\mathcal{L} _ {\\text{train}}(\\boldsymbol{\\alpha}, \\boldsymbol{w}) = \\mathcal{L} _ {\\text{train}}^{1:M}(\\boldsymbol{\\alpha}, \\boldsymbol{w})$. After the $K$ steps update of $\\boldsymbol{u}$ and $\\boldsymbol{w}$ (Equation (3) (4)), substitute $\\boldsymbol{u}_k$ and $\\boldsymbol{w}_k$ results in the update rule Equation (5). $\\Pi _ {\\mathcal{A}}(\\cdot)$ projects the updated $\\boldsymbol{\\alpha}$ into a probabilistic simplex. Running equation (3), (4) and (5) gives the solution of the bilevel DMO solution as shown in Theorem 1. We will make this clearer in the revision.\n\n>Q2: The final paragraph of Section 2.4 reads more like speculation.\n\nThere might be a misunderstanding. This section is not speculation. Let us elucidate the rationale underlying the description. \nIn Proposition 1, we see that as long as the per-domain generalization gap $ | \\mathcal{L} _ {\\text{val}}^m - \\mathcal{L} _ {\\text{train}}^m | \\rightarrow 0$, uniform weighting emerges as a valid solution. But why in the conventional data-abundant scenario do the generalization gaps tend to 0? Note that for uniform $\\boldsymbol{\\alpha}$, the training data and validation data are independently identically distributed (See Equation 1, $\\mathcal{L} _ {\\text{val}}$ always takes uniform $\\boldsymbol{\\alpha}$). The validation loss can be used as a proxy for population risk. Recall that in empirical risk minimization, mini-batch stochastic gradient follows the gradient of the true population risk as long as no examples are repeated ([1], Page 277), so the $\\mathcal{L} _ {\\text{val}}^m$ and  $\\mathcal{L} _ {\\text{train}}^m$ decrease at a similar speed during the training when all the training data are consumed only once (the first epoch), that is, our conventional data-abundant scenario. So $|\\mathcal{L} _ {\\text{val}}^m - \\mathcal{L} _ {\\text{train}}^m| \\rightarrow 0$ in this case. In the data-restricted scenario (In a general sense, including cases where the overall data volume is large, but certain domains are relatively small), the situation changes. Samples from smaller domains are inevitably upsampled, mini-batch stochastic gradient no longer follows the gradient of the population risk (the validation loss). The generalization gap $|\\mathcal{L} _ {\\text{val}}^m - \\mathcal{L} _ {\\text{train}}^m|$ increases. A typical example is overfitting on the smaller domains. In this way, the uniform $\\boldsymbol{\\alpha}$ is no longer a solution for the DMO problem. \n\nFor the empirical evidence, from Table 2, we see that in the conventional data-abundant scenario, no one approach can significantly outperform Uniform (supported also by [2]), showing that Uniform might already be a valid solution. While in the data-restricted scenario, a number of approaches including DoGE, Skill-it, Aioli and TANDEM outperform Uniform by a large margin. \n\n>Q3: The rationale behind the choice of hyperparameters.\n\n**Context length:**  In the data-abundant pretraining scenario, we use the context length 2048, the same as in Aioli [1], so that we can compare directly with the results given in Aioli~[1]. For the data-restricted scenario and SFT, we keep a unified context length of 512 as the generated token lengths in the SFT tasks are all smaller than 512. The context length will not affect the results. In the data-restricted scenario, we test context length 2048 and TANDEM reaches a final perplexity 28.04, 25.17, 25.21 for 160M, 410M, and 1B models, which is very similar to the 512 length results. \n\n**$K$:**   $K$ is the number of probing steps used to estimate the proper $\\boldsymbol{\\alpha}$ update direction (Equation (3)$\\sim$(5)). Larger $K$ helps suppress the gradient variance and provides more fidel $\\boldsymbol{\\alpha}$ update direction. Pretraining is relatively stable, so we use a smaller $K=5$ in the data-abundant and data-restricted scenarios. The gradient variance in SFT, however, is much higher (Figure 2), so we use a larger $K=20$. Furthermore, we conduct a sensitive analysis on how $K$ will affect the final model performance. The result shows that $K$ should be sufficiently large to ensure proper $\\boldsymbol{\\alpha}$ update, while further increasing $K$ beyond this point will not induce additional benefits. Please kindly refer to our response to question 6 for reviewer Vq6v. \n\n**$E$:**   Given the total number of training steps (amounts of data to be used), $E$ determines the number of $\\boldsymbol{\\alpha}$ updates during DMO: $T = \\frac{total\\  steps}{E}$. In the data-abundant scenario, we chose $E=20$ so that the mixture ratio is updated for $T = 2000$ steps. As one $\\boldsymbol{\\alpha}$ update requires an addition $K$ steps $\\boldsymbol{w}$ updates, a larger $T$ means higher additional computational cost. To reduce the computational overhead, we set $E$ in the data-restricted scenario and SFT to ensure $T=1000$ and $T=200$ respectively. \n\nThe value of $E$ is not deliberately tuned. Here we provide a sensitive analysis on $E$. We test $E=1,5,10$ in the 160M data-restricted scenario, and the result is as follows:\n|           | E=1  | E=5  | E=10 |\n|:---------:|:----:|:----:|:----:|\n| Perplexity|27.99 |28.05 |28.56 |\n\nTANDEM is not that sensitive to $E$ as long as it makes $T$ large enough so that $\\boldsymbol{\\alpha}$ is sufficiently updated.\n\n>Q4: Experiments on larger models.\n\nThanks for the suggestion, we’ve added 3B scale experiments. Due to limited space, please refer to our response to question 2 for reviewer 3KfP.\n\n>Q5: The improvement of TANDEM in Table 4 (SFT) is not significant.\n \nIn Table 4, TANDEM doesn't yield particularly strong empirical results because the final metric for the 6 tasks is inconsistent with the loss. While TANDEM significantly decreases the test loss, the performance doesn't improve proportionally. \n\nIn response to question 3 for reviewer Vq6v, we conducted a much larger-scale SFT experiment. On this larger dataset, TANDEM demonstrates much significant advantages over the baselines, showing the effectiveness of the proposed method. Please kindly refer to our answer to question 3 for reviewer Vq6v.\n\n>Q5: Why are DoReMi and DoGE not as good as Uniform?\n\nAs discussed in Section 2.4, Uniform is already a valid solution for the bi-level DMO problem in the data-abundant scenario, so it is hard for DoReMi and DoGE to outperform Uniform significantly. We then illustrate why the two methods suffer in the other scenarios.\n\n**DoReMi:**  The problem of DoReMi lies in that it depends on a pre-trained static reference model. Imagine that in the data-restricted scenario and SFT, the reference model overfits on certain small domains (performs extremely low training loss, so the excessive losses are high), DoReMi may incorrectly continue to up-weight the domains to reduce the corresponding excessive losses. Besides, we know that, different from pretraining where the training loss decreases steadily, in SFT, there are always sudden downward jumps in training loss when the data repeats [3]. There are also chances that the SFT loss goes even to 0 for a large model. These pose greater challenges for DoReMi to obtain a good reference model.  \n\n**DoGE:**   In Table 2, we see that DoGE outperforms Uniform in the data-restricted scenario. While for SFT, DoGE suffers because it relies on the inner product of per-domain gradients, which exhibits large variance and is difficult to provide reliable $\\boldsymbol{\\alpha}$ update directions (Figure 8). On the contrary, our TANDEM does not rely on the gradient estimation and utilizes $K$ steps probing to get more reliable $\\boldsymbol{\\alpha}$ updates. \n\n>Q6: In the analysis of the probing step, why do you fix $\\boldsymbol{\\alpha}$?\n\nWe fix $\\boldsymbol{\\alpha}$ to isolate the effect of $K$ in suppressing the hyper-gradient variance. Update of $\\boldsymbol{\\alpha}$ will introduce an extra variable--the value of the mixture ratio $\\boldsymbol{\\alpha}$. Nevertheless, we agree that an analysis on how $K$ directly affects the final model performance is important. We add a sensitive analysis, please kindly refer to our response to question 6 for reviewer Vq6v. \n\n\n>Q7:  Typos, grammatical, typographical issues and other suggestions helps improving the clarity.\n\nThanks very much for your careful checking, we've fixed these typos and thoroughly polished the revision.\n\nSince we have replied to all the questions, if you find our answers satisfactory, we respectfully ask that you please consider increasing the score. If you have any more questions or concerns, please let us know and we will be happy to answer.\n\n\n\n\n[1] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning, MIT Press, 2016\n\n[2] Mayee F. Chen, Michael Y. Hu, Nicholas Lourie, Kyunghyun Cho, and Christopher Re. Aioli: A unified optimization framework for language model data mixing. In International Conference on Learning Representations, 2025.\n\n[3] Jeremy Howard and Jonathan Whitaker. Can LLMs learn from a single example? 2023"}, {"review_id": "f4cOtzEcAI", "reviewer": "Reviewer_Vq6v", "summary": "This paper introduces TANDEM, a novel method for data mixture optimization in the context of large language model (LLM) training. The authors formulate the optimization of domain-specific mixture ratios as a bi-level optimization problem, and transform it into a penalized single-level formulation. They propose a twin-network architecture where a proxy model is trained on primary data, and a reference model is dynamically trained with additional data. The difference in loss between these two models serves as a signal to reweight the domain mixture. TANDEM is theoretically grounded, enjoys convergence guarantees, and is empirically validated across three representative scenarios: data-abundant pretraining, data-restricted pretraining, and supervised fine-tuning.", "questions": "1. Scalability to LLM-scale models (10B+ parameters): While 160M to 1B models are explored, can TANDEM scale to 10B+ scale? Are there practical bottlenecks (e.g., memory, wall-clock time) for training two models in parallel at scale?\n2. Downstream Generalization: The current evaluation is focused on perplexity or loss. Have you tested whether mixture ratios learned by TANDEM improve generalization in downstream tasks, especially zero-shot / few-shot ability in different benchmarks for pretraining models (such as MMLU, GSM8K)?\n3. The authors provide a useful analysis on how increasing K reduces gradient variance (Fig. 8). However, it remains unclear how sensitive the final model performance (e.g., perplexity, accuracy) is to different values of K, especially across different domains or tasks. Ablation studies on these dimensions, as well as the feasibility of adaptive or dynamic K, would further strengthen the claims.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths：\n1. The paper is well-organized with clear motivation, rigorous derivations, and strong empirical validation.\n2. The transformation of the bi-level optimization into a penalized single-level form is theoretically grounded, with convergence analysis included.\n3. The method is efficient and scalable, avoiding Hessian computation, and is applicable to a wide range of scenarios, especially under-explored settings like SFT.\n4. Extensive experiments in multiple scenarios (pretraining, low-resource, SFT) show TANDEM outperforming strong baselines.\n\nWeakness：\n1. Although the method is presented as novel, it closely builds on DoReMi and DoGE. The conceptual contribution lies more in unifying and refining rather than innovating from scratch.\n2. The paper lacks discussion on how TANDEM could be used in industry-scale LLM training pipelines and what challenges may arise there.\n3. Limited Task Diversity in SFT: The SFT experiments focus on only six tasks with small datasets; broader evaluation would strengthen generalizability claims.", "was_revised": true, "final_justification": "I appreciate the authors' thorough and comprehensive rebuttal. I have carefully reviewed the responses and the additional experiments presented, which addressed most of my concerns. I revised my score accordingly.", "rebuttal": "Thank you for your constructive comments. We address your concerns below: \n\n>Q1: Although the method is presented as novel, it closely builds on DoReMi and DoGE.\n\nWhile TANDEM does draw inspiration from prior work like DoReMi and DoGE, it represents significant conceptual and methodological advancements in several key aspects: \n\n**Theoretical Foundation:**  Unlike DoReMi and DoGE's heuristic approach (group DRO \\& alternating model parameters $\\boldsymbol{w}$ and mixture weights $\\boldsymbol{\\alpha}$), TANDEM establishes a rigorous bi-level optimization framework with provable convergence guarantees (Theorem 1), representing a fundamental theoretical contribution.\n\n**Architectural Innovation:**  The twin-network design (Figure 1) introduces novel synchronization mechanisms between proxy and reference models, addressing critical limitations: Dynamic reference model vs DoReMi's static reference model. Light-weight loss comparison vs  DoGE's cumbersome and noisy per-domain gradient estimation. The proposed method is elegant and conceptually appealing, as the reviewer jR7t stated.\n\n**Valuable Field Insights:**   As the reviewer jR7t recognized, this paper goes beyond the algorithmic design to discuss when and why DMO matters. Providing valuable insights to the DMO field.\n\nThe synthesis of these elements - theoretical grounding, architectural innovation, and valuable field insights - constitutes innovation beyond simply unifying existing approaches.\n\n\n>Q2:  How TANDEM could be used in industry-scale LLM training pipelines and what challenges may arise there.\n\nTANDEM can be used to construct the final training set after samples in each domain are properly pre-processed and de-duplicated, aiming at building a well-performing and balanced model.\n\nAs the reviewer pointed out, directly applying TANDEM to industry-scale LLM poses significant challenges. Maintaining the reference model and proxy model simultaneously requires nearly doubled memory. One potential bypass is to use smaller models for data mixture optimization and transfer the learned mixture ratio to LLM training. Previous works like DoReMi and DoGE have demonstrated the effectiveness of this transfer. This also applies to TANDEM. In Figure 6, we show that the optimal mixtures learned with models of different scales are quite consistent. So the $\\boldsymbol{\\alpha}$ learned with a smaller model can be safely applied to a larger model train.\n\n>Q3: Limited Task Diversity in SFT with small datasets.\n\nThanks for your suggestion. We have already tested TANDEM on tasks of different formats, including open-ended text generation, multiple choice, and True/False tasks. Here we further increase the data scale of the SFT experiments. We use 6 major categories (containing 99 tasks) from Natural Instructions, each contains tasks of different formats. The statistics are as follows:\n| Method               | Num. Sample |\n|:--------------------:|:-----------:|\n| Textual Entailment   | 79332       |\n| Answer Verification  | 13195       |\n| Text Matching        | 47297       |\n| Information Extraction | 31053    |\n| Word Extraction      | 17294       |\n| Text Categorization  | 89572       |\n\nWe use the same hyperparameters as in section 4.1. The comparison on this larger SFT dataset is as follows:\n| Method   | Text Ent. | Ans Ver. | Text Mat. | Inf Ext. | Word Sem. | Text Cat. | Avg Metric $\\uparrow$ | Test Loss $\\downarrow$ |\n|:--------:|:---------:|:--------:|:---------:|:--------:|:---------:|:---------:|:-----------:|:----------:|\n| Uniform  | 84.4    | 75.1 | 86.2  | 77.8   | 88.3 | 83.2 | 82.5 |  0.231       | 0.231      |\n| DoReMi   | 85.3      | 74.9     | 83.6      | 78.1     | 87.0      | 79.0      | 81.6        | 0.249      |\n| DoGE     | 81.4      | 76.7     | **86.5**  | **78.6** | 88.3      | 82.4      | 82.4        | 0.297      |\n| Skill-It | 84.1      | 75.0     | 86.4      | 78.3     | 87.9      | 83.7      | 82.6        | 0.232      |\n| Aioli    | 83.9      | 75.8     | 86.2      | 78.0     | 88.3      | 84.0      | 82.7        | 0.229      |\n| TANDEM   | **85.3**  | **76.3** | 86.2      | **78.6** | **88.5**  | **84.9**  | **83.3**    | **0.208**  |\n\nFrom the table we see that TANDEM still outperforms all the baselines on this much larger SFT dataset. \n\n>Q4: Experiments on larger models.\n\nThanks for the suggestion, we’ve added 3B scale experiments. Please kindly refer to our response to question 2 for reviewer 3KfP.\n\n> Performance on downstream tasks.\n\nThanks for the suggestion, we've added downstream few shot experiments and indepth discussion. Please kindly refer to our response to the Q2 for reviewer SkSi. MMLU & GSM8K are too hard for models and data of this size, so we evaluate ARC-C, ARC-E, BoolQ, HellaSwag, LAMBADA, PiQA and WinoGrande as in Aioli\n\n> Sensitivity of the final model performance on $K$. The feasibility of adaptive or dynamic $K$\n\nThanks for your suggestion. Here we add a sensitivity analysis on $K$, aiming at inspecting how vary $K$ will affect the final model performance. We test $K=1,5,10$ in the data-restricted scenario, and the result is as follows:\n|           | K=1  | K=3  | K=5 | K=10 |\n|:---------:|:----:|:----:|:----:|:----:|\n| Perplexity|29.57 |28.31 |28.06 |28.05 |\n\nFrom this table, we see that too small $K$ may result in less fidel $\\boldsymbol{\\alpha}$ update direction, thus sub-optimal data mixture ratio and higher perplexity. When $K$ is large enough for $\\boldsymbol{\\alpha}$ update probing,   increasing $K$ will not induce further benefits.\n\nFor adaptive or dynamic $K$, we agree that this can be an interesting and fruitful direction. We left it for future work. \n\n\nWe thank the reviewer again for the review. If you have any more questions about our replies, or any further questions or concerns, please let us know."}, {"review_id": "4yD0TKnpqJ", "reviewer": "Reviewer_SkSi", "summary": "This paper introduces TANDEM, an algorithm designed to optimize the domain-specific mixture ratios of training data for LLMs. The authors frame data mixture optimization (DMO) as a bi-level optimization problem and reformulate it into a single-level penalized objective.\nTANDEM solves this by employing \"twin networks\": a proxy model trained on training data and a reference model trained on both training and validation data. By the difference in losses between these twin models, the method up-weights domains that show more significant improvement from additional validation data.\nExtensive experiments across data-abundant pretraining, data-restricted pretraining, and supervised fine-tuning scenarios validate TANDEM's effectiveness and superiority over existing methods.", "questions": "1. Why not treat the proposed algorithm as an online method? Since proxy model $u$ or reference model $w$ has been trained with adaptive domain weights $\\alpha^t$, is it possible to train a base model using procedure of TANDEM directly? (Treat $u$ or $w$ as final LM.) I'm curious the performance of them.\n2. One recent work, Chameleon [1], also considers data mixing in finetuning. Could you compare TANDEM with Chameleon? That would be great if you can also show some empirical results.\n\n[1] Xie, Wanyun, Francesco Tonin, and Volkan Cevher. \"Chameleon: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning.\" ICML 2025.", "limitations": "See weakness.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 1, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strength**\n1. The paper's re-formulation of the bi-level DMO problem into a single-level penalized objective is insightful.\n2. TANDEM's dynamic updating and synchronization of both proxy and reference models address limitations of prior work.\n3. It's important to consider Data-Abundant and Data-Restricted scenarios separately.\n4. It's interesting to consider data mixing in supervised fine-tuning. \n\n**Weakness**\n1. The derivation from Eq (2) to Eq (4) is unclear. Does $\\nabla L_{val}(\\alpha, w_k)$ in Eq (4) mean the gradient on validation set is computed with domain weights $\\alpha^t$ rather than uniform weights? However, shouldn't it be uniform weights according to Eq (2)? \n2. The empirical results on SlimPajama only show perplexity, but it doesn't test on downstream tasks. Although average perplexity can represent model's performance, downstream tasks can show models' real capability. Many data mixing works like Doge, RegMix, Chameleon [1] show various downstream tasks' accuracy.\n3. Although time of TANDEM is competitive to other data mixing methods as shown in Table 3, the computational cost is related to $K$ and $E$. Could you discuss more about computational time with respect to $K$ and $E$? In addition, could you report FLOPs during training?", "was_revised": true, "final_justification": "- It's interesting to consider data-restricted scenarios.\n- The paper expands data mixing methods into supervised fine-tuning.\n- I would suggest including the performance of downstream tasks in the paper.\n- Comparing computational cost with other baselines should be added.", "rebuttal": "Thank you for your constructive comments. We address your concerns below:\n\n>Q1: $\\nabla L _ \\text{val}\\left(\\boldsymbol{\\alpha}, \\boldsymbol{w} _ k\\right)$ in Equation 2 should be $\\nabla L _ \\text{val}\\left( \\boldsymbol{w} _ k\\right)$.\n\nYes, there is a typo in Equation (4). The gradient on the validation set is computed with the uniform weights. Thank you for pointing this out, we will make it clearer in the revision.\n\n\n>Q2: Performance on downstream tasks.\n\nThanks for the suggestion, we evaluate the 160M models trained on SlimPajama on ARC-C, ARC-E, BoolQ, HellaSwag, LAMBADA, PiQA and WinoGrande using the Language Model Evaluation Harness [4] as in Aioli [1]. The result is as follows:\n| Method   | ARC-C | ARC-E | BoolQ | HellaSwag | LAMBADA | PiQA | WinoGrande | Average |\n|:--------:|:-----:|:-----:|:-----:|:---------:|:-------:|:----:|:----------:|:-------:|\n| Uniform  | **25.76** | 19.58 | 60.70 | 18.72     | 9.98    | **45.43**| 50.04      | 32.85   |\n| DoReMi   | 21.36 | 19.58 | 60.61 | 11.66     | 6.44    | 42.94| 49.57      | 31.06   |\n| DoGE     | 22.03 | 24.87 | 51.35 | 18.50     | 10.07   | 44.45| 47.28      | 31.39   |\n| Skill-it | 18.98 | **26.10** | 50.82 | 24.99     | **11.47**   | 43.36| 50.59      | 31.40   |\n| Aioli    | 20.34 | 19.58 | **61.71** | 19.59     | 10.89   | 44.18| 49.41      | 32.41   |\n| TANDEM   | 20.34 | 20.81 | 57.55 | **25.59**     | 11.20   | 40.53| **51.07**      | **32.92**   |\n\nWe see that TANDEM outperforms all the baselines, though not as significantly as on the perplexity. Actually, as many recent works[1][5][6][7] pointed out, lower perplexity is not necessarily positively correlated with higher performance on downstream tasks. \nThis disparity might be caused by the distribution shift between the pre-training data and the downstream evaluation task. Recall that the bilevel formulation directly optimizes the model performance on validation data, which is, in our case also the pre-training data instead of the downstream tasks. When an algorithm is completely unaware of the downstream tasks,  we do not expect it to universally outperform the others. We do believe that by establishing a comprehensive and diverse validation system, TANDEM will significantly improve the development of LLM. We will add an in-depth discussion on this in the revision. \n\n>Q3:  Computational time with respect to $K$ and $E$. FLOPs during training.\n\nWe conduct a computational complexity analysis contrasting TANDEM and vanilla model training on the same dataset. The complexity is as follows:\n\n|  Method | Computational Complexity |\n|:---------:|:---:|\n| Vanilla Model Train|$C T (E + K)$ |\n| TANDEM             |$C T (E + K) +  C  T  K + \\frac{2}{3}C T$ |\n\n$C$ is the complexity of one step training of an ordinary model. $T$ is the update number of $\\boldsymbol{\\alpha}$. $K + E$ compose the mixture ratio update interval ($K$ steps of probing, and $E$ steps of free $\\boldsymbol{u}$ training). In this way, training on $\\boldsymbol{u}$ takes $C T (E + K) $ computation, additional update of $\\boldsymbol{w}$ takes $ C  T  K $, and update on $\\boldsymbol{\\alpha}$ takes $\\frac{2}{3}C T$ (forward process of two models $\\boldsymbol{u}$ and $\\boldsymbol{w}$, here we simply assume backward takes 2$\\times$ computation as the forward process). Please refer to Figure 1 for a better understanding. Compared to vanilla model training on the same dataset~(train only $\\boldsymbol{u}$ for $T(K+E)$ steps), which takes $C  T  (K+E)$ computation, TANDEM only introduce $C  T  K + \\frac{2}{3}C T$ additional computation, less than 2$\\times$ vanilla training in most cases. This complexity dominates the FLOPs during training. We will include FLOPs in the experiments in the revision.\n\n>Q4:  Why not establish an online method and treat $\\boldsymbol{u}$ or $\\boldsymbol{w}$ as the final model?\n\nThank you for the suggestion. We evaluate the model $\\boldsymbol{u}$ which has been trained with adaptive domain weights $\\boldsymbol{\\alpha^{t}}$ in all the three scenarios ($\\boldsymbol{w}$ is (re)initialized with $\\boldsymbol{u}$ periodically so we consider just $\\boldsymbol{u}$), and compare it to the default two-stage (DMO and then train) trained model. \n| Method             | Data Abundant (Perpl.) | Data Restricted (Perpl.) | SFT (Acc.) |\n|:------------------:|:---------------------:|:-----------------------:|:---------:|\n| Uniform            |         25.74         |          31.53          |   74.72   |\n| Online             |         25.72         |          29.63          |   74.49   |\n| Two-Stage (default)|         25.43         |          28.07          |   75.03   |\n\nFrom the table above, we see that the performance of the online trained model $\\boldsymbol{u}$ falls behind the default two-stage trained model as well as the uniform baseline. This result coincides with the findings in previous works, e.g., Table 6 in DoReMi(The proxy model underperforms the baselines) and Figure 8 in DoGE. We hypothesize that the frequent change of training data distribution deteriorates the training process. Though TANDEM in its current form is not suitable to be used as an online method, we totally agree that establishing online methods remains an important future work.\n\n>Q5:  One recent work, Chameleon [3], also considers data mixing in finetuning. Could you compare TANDEM with Chameleon?\n\nThere might be a misunderstanding. The \"finetuning\" in CHAMELEON much resembles (continual) pretraining where the data used (Wiki-40 and Stack-decup) are also unlabeled plain texts and consumed only once, while in our SFT setting we focus on the \"prompt \\& answer\" pair data, which usually induces much noisy gradients, allows multi-epoch training and suffer from overfitting problems. \nSimilar to Skill-it, CHAMELEON determines the optimal mixture ratio according to a pre-computed fixed $M \\times M$ domain correlation matrix ($A$ in skill-it and $XX^T$ in CHAMELEON), which encodes the data efficacy of each domain. However, this kind of method fades when data in certain domains quickly depletes and upsampling kicks in (our data-restricted scenario) because the data efficacy changes, the domain correlation matrix is no longer aligned with the current data characteristics.\n\nFinally, note that CHAMELEON is released on May 30, 2025, after the submission deadline of NeurIPS 2025, so by no means at that time can we compare CHAMELEON with TANDEM. Due to the limited rebuttal duration, we defer this comparison to the revision. \n\nWe thank the reviewer again for the review. If you have any more questions about our replies, or any further questions or concerns, please let us know.\n\n[1] Mayee F. Chen, Michael Y. Hu, Nicholas Lourie, Kyunghyun Cho, and Christopher Ré. Aioli: A unified optimization framework for language model data mixing. In International Conference on Learning Representations, 2025.\n\n[2] Mosh Levy, Alon Jacoby, and Yoav Goldberg. Same task, more tokens: the impact of input length on the reasoning performance of large language models. In Annual Meeting of the Association for Computational Linguistics.\n\n[3] Xie, Wanyun, Francesco Tonin, and Volkan Cevher. Chameleon: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning. In International Conference on Machine Learning, 2025.\n\n[4] Leo Gao, Jonathan Tow, and Baber Abbasi et.al. A framework for few-shot language model evaluation, 2024.\n\n[5] Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better downstream: Implicit bias matters for language models. In International Conference on Machine Learning 2023,\n\n[6] Yi Tay, Mostafa Dehghani and Samira Abnar. Scaling laws vs model architectures: How does inductive bias influence scaling? In Conference on Empirical Methods in Natural Language Processing, 2023.\n\n[7] Mengzhou Xia, Mikel Artetxe and Chunting Zhou et.al.. Training trajectories of language models across scales. In Association for Computational Linguistics, 2023."}]

Author rebuttal:

Meta-review: The paper proposes TANDEM, a bi-level data mixture optimization framework solved via twin networks. It reformulates domain reweighting into a single-level penalized problem, where a proxy model is trained on primary data and a dynamically updated reference model is trained with additional data. The difference between the two serves as a signal to adapt mixture ratios, yielding a principled, theoretically grounded, and practically efficient algorithm. TANDEM provides convergence guarantees and shows broad applicability across three scenarios: data-abundant pretraining, data-restricted pretraining, and supervised fine-tuning (SFT). Experiments demonstrate consistent improvements over strong baselines, especially in restricted and SFT regimes.

### Strengths of the paper:

- The reviewers (3KfP, Vq6v, SkSi) highlight the elegance and theoretical grounding of reformulating bi-level optimization with twin networks, noting that it unifies and clarifies prior methods (DoReMi, DoGE) while providing provable convergence.

- The method is versatile and broadly applicable, validated across data-abundant, data-restricted, and SFT settings, making it more comprehensive than most prior work.

- The paper is technically solid, offering both rigorous proofs and extensive empirical evaluation with multiple model scales and domains.

- Reviewers appreciated the insightful analysis of when and why DMO matters, going beyond algorithmic design to field-level insights.

- The rebuttal added larger-scale (3B) experiments, more diverse SFT evaluation (99 tasks), downstream task results, sensitivity analyses, and cost comparisons, strengthening the submission.

### Weaknesses of the paper:

- Reviewer jR7t noted presentation issues in the methodology, particularly the derivation of Eq. (5) and speculative claims in Section 2.4. While improved in rebuttal, clarity could still be strengthened in revision.

- The evaluation primarily uses relatively small models (<1B, some 3B). While this is standard in DMO literature due to compute constraints, it raises questions about scalability to 10B+ industry-scale models.

- Some empirical results, especially in SFT (Table 4), showed subtle gains without statistical significance, and the surprising underperformance of DoReMi/DoGE compared to Uniform prompted concerns about baseline tuning. The rebuttal provided clarifications and additional experiments, but skepticism remains.

- The computational cost, while analyzed, still involves training two networks, which may be heavy for larger-scale application.

- Writing quality was flagged as uneven in places, with grammar/notation issues that should be carefully polished.

### Primary reasons for Accept (Poster)

The primary reasons for recommending Accept (Poster) are that TANDEM provides a principled, elegant, and theoretically grounded framework for data mixture optimization, offering both unification of prior approaches and practical improvements in underexplored settings (restricted data and SFT). Despite limitations in model scale and presentation, the rebuttal convincingly addressed reviewer concerns with new empirical evidence, downstream evaluations, and detailed clarifications. The contribution is timely, impactful, and technically rigorous, making it a valuable addition to the NeurIPS program.

### Summary of the discussion and rebuttal

The authors engaged thoroughly with reviewers. For R-3KfP, they validated DMO importance in overtrained regimes and added 3B-scale experiments, showing TANDEM remains effective. For R-Vq6v, they expanded SFT evaluation to 99 tasks and added downstream results (ARC, BoolQ, HellaSwag, etc.), confirming generalization. For R-SkSi, they clarified derivations, reported FLOPs, and compared against Chameleon, showing TANDEM’s superiority in restricted regimes. For R-jR7t, they polished methodological explanations, provided empirical validation of generalization gap dynamics, and clarified why DoReMi/DoGE may underperform Uniform. While some concerns remain about scale and clarity, the rebuttal upgraded multiple scores and established consensus on the paper’s merit. Overall, TANDEM is recognized as a solid, rigorous, and timely contribution, meriting acceptance as a poster.
Decision: Accept (poster)