Title: TANDEM: Bi-Level Data Mixture Optimization with Twin Networks
Venue: NeurIPS 2025
Accepted: yes
Keywords: data mixture optimization

language models
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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?\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nSee weakness.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths：\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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n1. 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).\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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.\n\nQUESTIONS:\nGiven 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-10"}