Title: Shift is Good: Mismatched Data Mixing Improves Test Performance
Venue: NeurIPS 2025
Accepted: no
Keywords: data mixing

distirbution shift
Abstract: We consider training and testing on mixture distributions with different training and test proportions. We show that in many settings, and in some sense generically, distribution shift can be beneficial, and test performance can improve due to mismatched training proportions. In a variety of scenarios, we identify the optimal training proportions and the extent to which such distribution shift can be beneficial.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\n### Summary\n\nThis paper challenges the conventional assumption that training and test ratios should be matched for optimal performance. Instead, it demonstrates that mismatching training data ratios is desirable and allows for better generalization performance. The authors provide theoretical support for their claims and show that \"positive distribution shift\" can provide learning speedups in one GPT2-like model.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths: \n* The authors provide theoretical support for the claims made in this paper.\n* The insights from this work hold potential for advancing LLM training, particularly for multi-task and multilingual foundation models.\n\n\n### Weaknesses:\n> Major \n\n* While the current experimental setup demonstrates improvements for a training distribution modified from 90:10 to 75:25 when the test distribution is 90:10, its robustness to other test distribution ratios, such as 100:0, remains unclear. It seems intuitive that more training samples in the dominant class (e.g., English) would benefit the model in such scenarios. Furthermore, given that real-world deployment often involves unknown test distribution ratios, it is uncertain how well this approach would generalize.\n\n* This paper lacks explainability and would benefit from deeper analysis on why/where the improvements come from. For example, Fig 1 could show the individual test errors for English and Chinese, as opposed to the current overall error. In a similar vein, the authors could consider metrics that account for imbalances within the data, such as F1 or Recall.\n\n* Given the current limited experiments, I am not convinced that this work would readily translate to larger models like LLaMA (fine-tuning could be considered if pretraining resources are limited). Additionally, since this seems to be a generic paper, I encourage the authors to demonstrate wider applicability to other machine learning tasks, such as computer vision.\n\n> Minor \n\n* Additionally, this paper has too many typos e.g., L77, L83, L328. Figure 1 is also not referenced.\n\n* Quality:\nProofs seem correct, though i did not check rigourously.\n\n* Clarity:\nVery poorly written. This paper has too many typos and grammatical errors, leading to poor readability. \n\n* Significance:\nEmpirical analysis and evaluation on frontier models is lacking.\n\n* Originality:\nThe theoretical framing of this work appears original.\n\nQUESTIONS:\n* How would this work translate to frontier models/tasks?\n* Why is there a spike in mismatched learning for the middle and right of Figure 3?\n* How would other test distribution ratios affect the ideal mismatched ratio? For example, instead of 90:10, it becomes 50:50?\n* How many data points are used for these evaluations? From the Appendix it seems like 400 points which can be too small.\n* Could the quality of the sampled training points affect these results?\n\nLIMITATIONS:\nWould be good if the authors could include a section for the Limitations of this work, as per declared in the paper checklist.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the impact of distribution shift in a multi-task learning setup, where a model is trained and evaluated on data from a mixture of tasks. The authors focus on a specific type of shift: a change in the _proportion_ of tasks between training and testing. They refer to cases where this shift improves performance as “Positive Distribution Shift.” The paper develops theoretical analyses showing that such shifts are often beneficial, and that training on a different task mixture than the test distribution can lead to improved results on this test distribution.\n\nSTRENGTHS AND WEAKNESSES:\n### Originality\n\n- **Strengths**: The idea of analyzing task mixture proportions and the link with \"data mixture laws\" from Ye et al. is interesting and could provide more insights on the optimal combination of dataset in practice.\n- **Weaknesses**: The framing of the problem as “distribution shift” is somewhat misleading. The shift considered here (changes in task proportions) differs from how distribution shift is typically defined in the literature. This setup is more closely related to multi-task learning (MTL) and multiple-source domain adaptation (MSDA), which have been extensively studied in prior work ([A]–[E]). The paper does not sufficiently acknowledge or engage with this body of literature.\n\n### Quality\n\n- **Strengths**: The theoretical results are presented across both orthogonal and non-orthogonal tasks settings. An empirical result show how learning on a combination of tasks can improve performance for a single task.\n- **Weaknesses**: The assumption that the task mixture at test time is known is strong and not discussed. The findings on optimal task combinations appear to also rely heavily on the assumed power-law form of the error function. Both of these limit the practical applicability of the findings. It would be valuable to explore whether the theoretical insights extend to settings where the test-time task mixture is unknown or must be inferred, which would increase the practical relevance of the findings.\n\n### Clarity\n\n- **Strengths**: The writing is generally clear, and the structure makes it easy to follow the progression of ideas.\n- **Weaknesses**: The paper would benefit from a clearer summary of its contributions. A recap table and/or a concluding discussion of the main results would help understand the scope and implications of the work.\n\n### Significance\n\n- **Strengths**: Understanding how task mixtures affect generalization is important for designing robust multi-task learning systems.\n- **Weaknesses**: The novelty and significance of the results are somewhat diminished by the lack of engagement with prior work that has studied similar setups. For example, Section 5 and Section 6 revisit how learning on a mixture of multiple tasks can improve transfer to test tasks, which are known findings from MTL and MSDA literature (cf [A] for instance), without clearly distinguishing what is new. Without a clearer differentiation from existing MTL and MSDA literature, it is difficult to assess the novelty of the theoretical findings, which weakens the overall contribution.\n\n### References\n\n[A] Ben-David & Schuller. \"Exploiting task relatedness for multiple task learning.\" Learning Theory and Kernel Machines, COLT/Kernel 2003  \n[B] Mansour et al. \"Domain adaptation with multiple sources.\" Neurips 2009  \n[C] Mansour et al. \"Multiple source adaptation and the rényi divergence.\" UAI 2009  \n[D] Hoffman et al. \"Algorithms and theory for multiple-source adaptation.\" NeurIPS 2018.  \n[E] Redko et al. \"Optimal transport for multi-source domain adaptation under target shift.\" AISTATS 2019.\n\nQUESTIONS:\n1. How do you justify the assumption that the task mixture at test time is known? \n2. Could your framework and results be extended to cases where the test-time mixture is unknown? How would it affect your results?\n3. Could you provide a recap table or summary of theoretical results to help clarify the contributions?\n4. Could you discuss how your findings relate to or differ from the MSDA literature (e.g., [B]–[E])?\n\nLIMITATIONS:\nThe paper should acknowledge that the type of distribution shift studied, i.e. known differences in task proportions, is a special case and differs from more general forms considered in the literature. The paper should also discuss the assumption that the test-time task mixture is known, and consider its implications for real-world applicability.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 1, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper consider a senario where the training and test data distributions can be seen as mixtures of some sub-distribution components. It shows that in some cases, having mismatched training and test distribution, i.e. having different mixing weights in the training and test distributions, may help improve the generalization performance.\n\nSTRENGTHS AND WEAKNESSES:\nStrenghths:\n1. The main result that mismatching can actually be helpful is thought-provoking.\n2. The chain of logics is clear.\n3. Comprehensive review of related works.\n\nWeaknesses:\n1. Some parts of the writing are too casual.\n2. Some notations are not well used.\n3. Some definitions are not clear.\n4. A few results and assumptions are not fully justified. \n5. Lack of conclusions and justification of limitations.\n\nQUESTIONS:\n1. Some parts of the writing are too casual. For example, in Line 8-14, the description of the exam example is written too conversationally for an academic paper. Personally, I think this is a good example as a motivation of the work and a warmup for the readers, but this conversational style is more appropriate to be used in oral presentations rather than in a paper.\n2. What are the key differences between mixture setting and compositional setting?\n3. Line 63. Test distribution is defined as a mixture of some \"components\". What are the definitions of the components? Are they all distribution but with non-overlapping supports? Or can their supports overlap? Can all distributions be formulated, in some sense, as mixtures of some underlying components (with careful definitions of the components)?\n4. Line 72. In the definition of $L_N({p},q)$, from my understanding the expectation should be taken with respect to $S\\sim\\mathcal{D}^N_q$ not $S\\sim\\mathcal{D}^N_p$, is that correct?\n5. From my understanding, in the setup of this paper, the properties of models, training algorithms, etc. are all hidden in the power law definition of the error function given in Model 3.1, is this correct? If so, is there any justification of choosing the orthogonal power law particularly to formulate the error function? What makes it a better option than (possibly) other choices? I think this should be clearly explained since it seems that the power lar assumption is highly relied on by the main results and their proofs.\n6. Lin 107, at the end of this line, I think it should be $\\bar{e}_N(q)$, is this correct? If not, then what is the definition of $\\bar{e}(q)$?\n7. Is there any proof of the claim in Line 106 that \"$\\bar{e}_N(q)$ will concentrate around $e(qN)$\"? Also, is $\\bar{e}_N(q)$ a random variable? If so, where does the randomness come from?\n8. Proposition 3.2. What are the intuitions of the assumtions? Also, what is the definition of $n_0(A_i,B_i,\\alpha_i,p_i)$?\n9. Line 79. Does $L^\\text{ratio}_N$ represent the improvement or the training budget?\n10. Line 86. What is $g$ function? I don't see it defined anywhere, or maybe I have missed it.\n11. Corollary 3.3. What is the reason that $A$ and $B_k$'s are not in the bound? Do they have no effect on the sample complexity improvment? If so, then how does the assumption that $A_1=\\cdots=A_K=A$ affect the result?\n12. Section 4. I am not familiar with this field of works. Can you provide a more concrete introduction of orthogonal memorization tasks? For example, what could be the data? What is the mathematical formulation of a task? What are the objectives of the tasks, etc.? Thank you.\n13. Line 240. The concentration property needs to be justified, but doesn't have to be mathematically detailed.\n14. Line 275. What's the difference between $\\Delta^{K-1}$ defined here and $\\Delta_K$ mentioned in the previous text sections.\n15. Condition 7.1. What is the intuition of this condition? What does it tell us about the model and/or the test distribution?\n16. The related works can be put into a early section, and the last section should include a brief conclusion of the work and addressing of the limitations.\n\nLIMITATIONS:\nAddressed.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper investigates a very interesting idea: should the training distribution task mixture ratios match the test distribution rations, with an interesting conclusion, that the answer is no. They show that under power law error decay, often observed in practice, models benefit from slight undersampling of majority task. Even if the tasks are orthogonal to each other.\n\nSTRENGTHS AND WEAKNESSES:\nI think this a super interesting and important research topic. In practice data and its precise mixture can have tremendous impact on model performance. This paper provides a quite unintuitive and thus very valuable insight, that when mixing different tasks in training data ideally one should slightly undersample the majority tasks, even if the training data exactly matches test-time ratios.\n\nThe theory provided is sound, covers both the orthogonal task scenario and transfer learning. It is also backed by appropriate in-silico experiment. While one could desire more extensive experimental validation, I think it is sufficient for a theoretical paper. \n\nThe Fig 3 experiments that could be improved if 1) the 30% mixing ratio of uniform and iid weights would be motivated and 2) performance of 100% uniform sampling would be added, to show both extremes (now we only have 100% iid baseline). For example https://arxiv.org/pdf/2410.05980 argue that uniform sampling improves model generalization. So 100% uniform is also an interesting baseline to have.\n\nQUESTIONS:\nIt's not a question per se, but more of a future work suggestion, that it would be neat to see how this can be adapted if we don't know the task (or mixture) distribution at test time. Could taking the test samples without labels still build an ideal training set for them, based on this theory, by say approximating similarity of two points by their foundation model embedding similarity (a bit similar to the citation above). The assumption that we know the test-time mixture is a bit limiting in practice now.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: SUMMARY:
The paper considers cases where one might train and test on distributions with different mixture proportions, arguing that this may be beneficial. This seems like a big claim but actually it’s a bit obvious. Here the mixture components correspond to completely distinct populations and the problem can be such that there is no transfer among them. Thus because error tends to go down with sqrt(n) and the authors are considering the case where the sample budget is fixed and thus we must allocate the finite budget among the mixture components, we will nearly always choose a slightly different weighting of the components than the test example. The result is presented as a counterintuitive finding of “positive distribution shift” but I think this is really just rediscovering simple motivational cases for active learning. Credit to the authors for making these arguments formal and attempting to weave together a nice analysis. However, I think the problems here sit at the level that the conclusions are on their face obvious given the assumptions the authors are willing to make and they recapitulate known facts and existing intuition about active learning.

STRENGTHS AND WEAKNESSES:
Per above: “The result is presented as a counterintuitive finding of “positive distribution shift” but I think this is really just rediscovering simple motivational cases for active learning.”

QUESTIONS:
None

LIMITATIONS:
None.
TITLE: Official Review
CLARITY: 2
IMPACT: 1
ORIGINALITY: 1
DATE: 2025-07-07