TANDEM: Bi-Level Data Mixture Optimization with Twin Networks

Thank you for your constructive comments. We address your concerns below:

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)$.

Yes, 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.

Q2: Performance on downstream tasks.

Thanks 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:

We 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. This 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.

Q3: Computational time with respect to $K$ and $E$. FLOPs during training.

We conduct a computational complexity analysis contrasting TANDEM and vanilla model training on the same dataset. The complexity is as follows:

$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.

Q4: Why not establish an online method and treat $\boldsymbol{u}$ or $\boldsymbol{w}$ as the final model?

Thank 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.

From 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.

Q5: One recent work, Chameleon [3], also considers data mixing in finetuning. Could you compare TANDEM with Chameleon?

There 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. Similar 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.

Finally, 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.

We 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.

[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.

[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.

[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.

[4] Leo Gao, Jonathan Tow, and Baber Abbasi et.al. A framework for few-shot language model evaluation, 2024.

[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,

[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.

[7] Mengzhou Xia, Mikel Artetxe and Chunting Zhou et.al.. Training trajectories of language models across scales. In Association for Computational Linguistics, 2023.

Dear reviewer,

We are delighted to hear that we were able to address your concerns.

We compare the computational complexity of different baselines:

$N$ is the number of parameters. The overall computational cost of these methods is of the order: Aioli > Skill-it > DoReMi $\approx$ DoGE $\approx$ TANDEM. (Skill-it requires training $M$ models with data from each domain for $H$ steps in advance, we use $H = \frac{1}{4}T(K+E)$ in our analysis). Note that though the complexity of DoGE is not particularly large, it requires per-domain gradient computation, which is unfriendly to the computing kernel and takes much more time. For TANDEM, our current implementation includes $T$ times model copying and loading (setting $\boldsymbol{w}=\boldsymbol{u}$ , line 102 ), incuring non-negligible time-cost. A more optimized implementation can greatly reduce the time consumed by TANDEM. We will add detailed computational complexity analysis in the revision.

As for CHAMELEON, it is the most efficient, as it requires only vanilla training ($<C T (E+K)$) to get the "domain correlation". Nevertheless, as we mentioned in the rebuttal, CHAMELEON is not able to capture the data efficacy change in the data-restricted scenario. We compare CHAMELEON and TANDEM. For the 160M/410M/1B model, CHAMELEON gets perplexity 29.85, 27.65, and 27.55, higher than TANDEM's 28.07, 25.0, and 24.35.

Once again, thank you for your valuable feedback and kind consideration.

Best regards,

The authors

Thank you for your constructive comments. We address your concerns below:

Q1: Although the method is presented as novel, it closely builds on DoReMi and DoGE.

While TANDEM does draw inspiration from prior work like DoReMi and DoGE, it represents significant conceptual and methodological advancements in several key aspects:

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.

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.

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.

The synthesis of these elements - theoretical grounding, architectural innovation, and valuable field insights - constitutes innovation beyond simply unifying existing approaches.

Q2: How TANDEM could be used in industry-scale LLM training pipelines and what challenges may arise there.

TANDEM 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.

As 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.

Q3: Limited Task Diversity in SFT with small datasets.

Thanks 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:

We use the same hyperparameters as in section 4.1. The comparison on this larger SFT dataset is as follows:

From the table we see that TANDEM still outperforms all the baselines on this much larger SFT dataset.

Q4: Experiments on larger models.

Thanks for the suggestion, we’ve added 3B scale experiments. Please kindly refer to our response to question 2 for reviewer 3KfP.

Performance on downstream tasks.

Thanks 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

Sensitivity of the final model performance on $K$. The feasibility of adaptive or dynamic $K$

Thanks 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:

From 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.

For adaptive or dynamic $K$, we agree that this can be an interesting and fruitful direction. We left it for future work.

We 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.

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:

Q1: The derivation of Equation (5) could be explained more clearly.

In 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.

Q2: The final paragraph of Section 2.4 reads more like speculation.

There might be a misunderstanding. This section is not speculation. Let us elucidate the rationale underlying the description. In 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.

For 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.

Q3: The rationale behind the choice of hyperparameters.

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.

$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.

$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.

The 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:

TANDEM is not that sensitive to $E$ as long as it makes $T$ large enough so that $\boldsymbol{\alpha}$ is sufficiently updated.

Q4: Experiments on larger models.

Thanks for the suggestion, we’ve added 3B scale experiments. Due to limited space, please refer to our response to question 2 for reviewer 3KfP.

Q5: The improvement of TANDEM in Table 4 (SFT) is not significant.

In 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.

In 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.

Q5: Why are DoReMi and DoGE not as good as Uniform?

As 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.

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.

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.

Q6: In the analysis of the probing step, why do you fix $\boldsymbol{\alpha}$?

We 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.

Q7: Typos, grammatical, typographical issues and other suggestions helps improving the clarity.

Thanks very much for your careful checking, we've fixed these typos and thoroughly polished the revision.

Since 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.

[1] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning, MIT Press, 2016

[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.

[3] Jeremy Howard and Jonathan Whitaker. Can LLMs learn from a single example? 2023

Dear reviewer:

Thank you again for your time reviewing our paper. We would appreciate it if you could confirm that our responses address your concerns. We would also be happy to engage in further discussions to address any other questions that you might have.

Best regards

Q5: Why are DoReMi and DoGE not as good as Uniform?

We fully understand your concern. After careful inspection, we confirmed that the performance of DoReMi is restricted by the a priori trained reference model. Following previous works, in all our experiments, we train the reference model on training data obtained using the Uniform strategy (as mentioned in the paper). 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. Not that this is the default setting of all previous works.

For a more comprehensive comparison, we train another reference model with the natural data mixture ratio (See Figure 4, the data statistics). This setting ensures no severe overfitting happens. The result is as follows:

We will make this clear and report this result to avoid misunderstanding in the revision.

Actually, the instability of DoReMi and DoGE has also been observed in many previous works[1,2,3].

• [1] presents that DoReMi underperforms Uniform (Table 1, the "Average" line, DoReMi-10k) .
• [2] shows that the perplexity of DoReMi can be significantly higher than Uniform in some cases (Table 2, column 2 (GH/C4), column 5 (CC/GH/W). In Table 19, GitHub data constitutes 46.7% of the training set with DoReMi, which is highly skewed. Usually, upweighting Common Crawl and C4 instead of Github leads to lower perplexity.
• [3] figures out that DoReMi is unable to produce competitive sampling weights for any domain, as it often gives the majority of the weight to a single source(Section 5.2, the Findings box).
• From Table 2 in Aioli [2], DoGE can also underperform Uniform in many cases.

So, it is normal that DoRemi and DoGE underperform Uniform in some cases. We believe the implementations were faithfully executed.

Here again, we outline the weakness of these two baselines:

• 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 [4]. 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.

• DoGE: For SFT, DoGE suffers because it relies on the inner product of per-domain gradients, which exhibits large variance (SFT) 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.

Once again, thank you for your valuable feedback and constructive suggestions. Sincerely hope that our clarifications help alleviate your concerns and contribute to a more favorable reassessment of our work.

---

[1] Simin Fan and Matteo Pagliardini, and Martin Jaggi. DOGE: Domain Reweighting with Generalization Estimation. ICML 2024.

[2] Mayee F. Chen, Michael Y. Hu, Nicholas Lourie, Kyunghyun Cho, and Christopher Ré. Aioli: A unified optimization framework for language model data mixing. ICLR, 2025.

[3] J. Parmar et.al. Data, Data Everywhere: A Guide for Pretraining Dataset Construction. EMNLP 2024 oral.

[4] Jeremy Howard and Jonathan Whitaker. Can LLMs learn from a single example? 2023

For Data-Restricted Scenario:

For SFT:

We see that in the data-abundant scenario, the generalization gaps of each domain do not change much and are kept close to 0, $\left|\mathcal{L} _ {\text {val}}^m-\mathcal{L} _ {\text {train }}^m\right| \rightarrow 0$ holds so that Uniform is a valid solution. In the data-restricted scenario, we see an increase in the generalization gap, particularly on small domains (Arxiv, Book, Wikipedia), because the sample repetition kicks in. In the large domains (C4, CommonCrawl) where there is no data repetition, the generalization remains small. SFT is of the same case. Generalization gap $\left|\mathcal{L} _ {\text {val}}^m-\mathcal{L} _ {\text {train }}^m\right|$ increase in small tasks (Word Sem., Inf Ext.), while kept small for large tasks (Text Cat., Text Mat, Text Ent.). The results on All three scenarios validate our analysis.

We believe this empirical result, plus the analysis in our rebuttal, well supports our argument and addresses your concern. Our claim is consistent with the results found in [1] that no existing DMO methods consistently outperform Uniform (in data-abundant scenario) in terms of average test perplexity per group (In abstract).

W3: While you added new experiments to address the concerns with Table 4...

There might be a misunderstanding. We add the new SFT experiments (larger and more diverse) in response to question 3 for reviewer Vq6v, who is concerned that the original SFT tasks are small and lack diversity. We add it here for better comprehensiveness. We also explain why the result in Table 4 is not that significant. The tasks are rather small (6000 samples per task), final metric for the 6 tasks is inconsistent with the loss. While TANDEM significantly decreases the test loss, the performance doesn't improve proportionally.

Thanks for recognizing our work as elegant, interesting, and comprehensive in experiments. We appreciate your insightful comments and would like to take this valuable opportunity to provide a more detailed discussion.

Q1: Polishing Section 2.2 to make it more detailed and accessible.

Thank you for your suggestion. We will thoroughly polish this section. Here we briefly explain the equation (3) $\sim$ (5). This section is about solving the penalized single-level objective by alternatively optimizing $\boldsymbol{u}$, $\boldsymbol{w}$, and $\boldsymbol{\alpha}$. Recall the penalized single-level objective:

$\min _ {\boldsymbol{\alpha}\in\mathcal{A}, \boldsymbol{w}} \mathcal{H} _ {\gamma}(\boldsymbol{\alpha}, \boldsymbol{w}):= \mathcal{L} _ {\text{val}}( \boldsymbol{w})+ \gamma \left(\mathcal{L} _ {\text{train}}( \boldsymbol{\alpha}, \boldsymbol{w}) - \min _ {\boldsymbol{u}} \mathcal{L} _ {\text{train}}(\boldsymbol{\alpha}, \boldsymbol{u})\right)$

• Update of $\boldsymbol{u}$ in Equation~(3) directly comes from $\min_{\boldsymbol{u}} \mathcal{L}_{\text{train}}(\boldsymbol{\alpha}, \boldsymbol{u})$.

• Terms depend on $\boldsymbol{w}$ is $\mathcal{L} _ {\text{val}}( \boldsymbol{w})+ \gamma \mathcal{L} _ {\text{train}}( \boldsymbol{\alpha}, \boldsymbol{w})$. Applying SGD on which results in the update rule in Equation~(4).

• For $\boldsymbol{\alpha}$, in the objective (2), only the $\gamma \left(\mathcal{L} _ {\text{train}}( \boldsymbol{\alpha}, \boldsymbol{w})- \min _ {\boldsymbol{u}} \mathcal{L} _ {\text{train}}(\boldsymbol{\alpha}, \boldsymbol{u})\right)$ 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}$, 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.

Q2: Some of the current claims rely too heavily on speculative reasoning. Additional empirical evidence would strengthen your argument significantly.

Our "broad conclusion" is that Uniform effectively constitutes a valid solution for the bilevel optimization in data-abundant one-epoch training scenario. We are delighted to hear that our current analytical description is sufficient. Let us provide more empirical evidence now. Proposition 1 shows that as long as the per-domain generalization gap $\left|\mathcal{L} _ {\text {val}}^m-\mathcal{L} _ {\text {train }}^m\right| \rightarrow 0$, the Uniform weighting emerges as a valid solution (proof given in Appendix B). The following thing is to show that $\left|\mathcal{L} _ {\text {val}}^m-\mathcal{L} _ {\text {train}}^m\right| \rightarrow 0$ holds in the data-abundant one-epoch training, while not in the data-restricted/multi-epoch training.

We conduct experiments in all the three scenarios and show the change of $\left|\mathcal{L} _ {\text {val}}^m-\mathcal{L} _ {\text {train }}^m\right|$ along training (As figure is not allowed during the dicussion, we temporally outline the results in tables, figures will be given in the revision.)

For Data-Abundant Scenario:

We sincerely appreciate your recognition of our technical contributions. We address your concerns below:

Q1. Does careful mixing still matter when there is more sizable token count?

Thanks 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.

To validate this claim, we train the 160M models on the 6B version of sampled SlimPajama, which constitutes an overtrained (past Chinchilla) case:

We 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.

Q2. Experiments on larger models.

Currently, 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.

Due 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:

Data-restricted scenario:

We 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.

SFT:

As suggested by the reviewer Vq6v, we use a larger SFT dataset. (Please refer to our response to question 3 for reviewer Vq6v)

We see that for a larger 3B model, TANDEM still outperforms the Uniform baseline, showing its effectiveness.

Q3: What happens if you run the same algorithm but with one copy (u = w)?

We 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.

Q3: Typos.

Thanks very much for your careful checking, we've fixed these typos and thoroughly polished the revision.

We 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.

Dear reviewer:

Thank you again for your time reviewing our paper. We would appreciate it if you could confirm that our responses address your concerns. We would also be happy to engage in further discussions to address any other questions that you might have.

Best regards