Shift is Good: Mismatched Data Mixing Improves Test Performance

We thank the reviewer for their feedback.

Q: robustness to other distributions:

A: See answer to question #3.

Q: What if the test distributions were uknown:

A: The main way we view the paper is providing an answer to the question "why might it be better to train on a different distribution then the test distribution". This is a phenomenon encountered in practice especially with training large models with many proficiencies, with much of model training having to do with selecting good training distributions (for a fixed test distribution). We feel this is a current and important question that is not sufficiently acknowledged or studied, or even made explicit, especially from a theoretical perspective. We do also view our work as providing guidance in this "training distribution selection" problem, in the sense of understand what to look for and what direction or kind of training distribution we want, even in setting where the exact proportions are unknown. Nevertheless, we can also imagine some settings in which the proportions are known, e.g. based on quick analysis of a sample from the test distribution.

Q: explainability:

A: We view our results as understanding one purely statistical aspect of positive distribution shift, where the improvement in test error comes from upsampling the components that have a larger impact on the test error and downsampling the components that have a smaller impact on test performance. For example, in the introductory example, because of the power law error scaling, the first few additional training examples in both the Chinese and European histories decrease the error significantly more than a new training data after we have already seen a number of training data in that component. It turns out that 10% of the training data is small enough that an additional example from the Chinese component would decrease the total error more than an additional example from the English component (and of course in expectation). This is the source of positive distribution shift.

Q: It is unclear whether our results would hold for large models:

A: We believe that our results and their implications apply to even large models such as LLaMA (for pretraining and finetuning) given that empirically power law has been shown to capture the error decay of these models well [1-4]. We agree with the reviewer that further experiments in this direction are welcome and leave that for future work.

Here we provide answers to reviewer’s questions:

Q1:

How would this work translate to frontier models/tasks?

A1: Frontier models are usually used for a variety of tasks - e.g. text editing, code generation, chat, reasoning etc., and furthermore are trained using a variety of different datasets. The loss here would be the pre-training loss, which is empirically demonstrated to follow power law scaling, one of the cases we cover [1]. Furthermore, the phenomena on having improvement from training the model on a different train distribution than a test distribution is encountered in practice with training large models with many proficiencies, with much of model training having to do with selecting good training distributions (for a fixed test distribution). We do also view our work as providing guidance in this "training distribution selection" problem, in the sense of understand what to look for and what direction or kind of training distribution we want.

Q2:

Why is there a spike in mismatched learning for the middle and right of Figure 3?

A2: Loss spikes are common in training LLMs, and most of them are recoverable, i.e., after a few steps, the loss gradually goes back to the value before the spike. Figure 3 shows the average curves of 3 training runs, and we observe a spike in one of them. We believe that adding more runs can eventually average out the spikes and produce a smoother curve.

Q3:

How would other test distribution ratios affect the ideal mismatched ratio? For example, instead of 90:10, it becomes 50:50?

A3: Proposition 3.2 (and Corollary 3.3.) shows the effect of test distributions on the ideal mixing ratio. Taking $K=2$, $A_1=A_2=A$, $**p**=(p,1-p)$ and $\alpha=1$, we can see that $**q**^{\star}= \frac{1}{p^{\frac{1}{1+\alpha}}+(1-p)^{\frac{1}{1+\alpha}}}(p^{\frac{1}{1+\alpha}},(1-p)^{\frac{1}{1+\alpha}})$. From this, we see what the improvement is for $p:1-p$ test mixing ratio and what the optimal training distribution $**p**^{\star}$ is. In this special case, taking $p=1-p=½$, we see that the optimal training distribution is the same as the test distribution $**p**^{\star} = (½, ½)$ and that for this specific test distribution there is no improvement, i.e. $L^{\star}((½,½))=L^{same}((½,½))$. We note that this is in line with the rest of our theory. Corollary 7.4 says that in any orthogonal task setup with K tasks, as long as one error function is non-constant, then for all $**p**$ except possibly a zero measure set, there is a positive distribution shift. In particular, for the orthogonal power law setting, if the tasks are equally hard ($A_i$ and $\alpha_i$ are all the same), then the uniform test distribution $**p** = (\frac{1}{K},\dots, \frac{1}{K})$ will not have a positive distribution shift, i.e. the uniform $**p**$ is in the measure zero set excluded by Corollary 7.4.

Q4:

How many data points are used for these evaluations? From the Appendix it seems like 400 points which can be too small.

A4: To measure the accuracy of skill composition, we used 400 data points for each evaluation. We argue that 400 data points can already lead to meaningful results because the final test error gap between training with and without mismatch distributions can be as large as 10% in Figure 3.

Q5:

Could the quality of the sampled training points affect these results?

A5: We do online training and we also repeated the experiment 3 times. It seems that the quality of sampled points does not affect the results.

[1] Hoffman et al. An empirical analysis of compute-optimal large language model training. NeurIPS 2022. [2] Ye et al. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. ICLR 2025. [3] Kaplan et al. Scaling Laws for Neural Language Models. 2020. [4] Cherti et al. Reproducible scaling laws for contrastive language-image learning. CVPR 2023.

We thank the reviewer for their feedback.

Q: Framing the problem as distribution shift:

A: First, multi-source domain adaptation is an instance of distribution shift. Second, regarding the relationship to MTL, we emphasize that the main contribution of our paper is in the case of orthogonal tasks, which is clearly different from the MTL literature. We agree with the reviewer that the non-orthogonal cases are less crisp and the boundary with MTL is less clear. We view this as an additional case of where our framework is applicable and will de-emphasize it in the updated version of our paper.

Q1:

How do you justify the assumption that the task mixture at test time is known?

A1: Regarding the assumption that the test mixture distribution is known at test time, the assumption is reasonable if we frame the paper as follows. The main way we view the paper is providing an answer to the question "why might it be better to train on a different distribution then the test distribution". This is a phenomenon encountered in practice especially with training large models with many proficiencies, with much of model training having to do with selecting good training distributions (for a fixed test distribution). We feel this is a current and important question that is not sufficiently acknowledged or studied, or even made explicit, especially from a theoretical perspective. We do also view our work as providing guidance in this "training distribution selection" problem, in the sense of understanding what to look for and what direction or kind of training distribution we want, even in setting where the exact proportions are unknown. Nevertheless, we can also imagine some settings in which the proportions are known, eg based on quick analysis of a sample from the test distribution.

Q2:

Could your framework and results be extended to cases where the test-time mixture is unknown? How would it affect your results?

A2: See the answer above regarding the knowledge of test mixture at test time. We agree with the reviewer that extending the result to the case where the test mixture is unknown is an interesting future direction.

Q3:

Could you provide a recap table or summary of theoretical results to help clarify the contributions?

A3: We systematically study “positive distribution shift” when training with mismatched mixing proportions relative to the test distribution. First, we consider different per-component learning curves and in cases of orthogonal power law, fact memorization, and transfer cases we show that mismatched training and test distributions can be beneficial and we characterize the optimal training mixture and the improvement in test performance and training complexity. Second, we show that only in rare situations the optimal training distribution is equal to the test distribution, while in most cases shift is good. We will add a clearer summary of the main contributions in the updated version of our paper.

Q4:

Could you discuss how your findings relate to or differ from the MSDA literature (e.g., [B]–[E])?

A4: Our work centers on understanding and characterizing how and why might be better to train a different training distribution than the test distribution in a number of relevant settings, and providing guidance in the selection of training distribution problem. On the other hand, based on the references [B-D], in the MDSA problem, the learner receives k distributions $D_i$ corresponding to $k$ sources and a predictor $h_i$ for each of the sources (where the predictor could regression, classification, or probability model). The task of the learner is to find a predictor h that performs well on a distribution $D_T$ made up as a mixture of $D_i$ (where the mixture might be either known or unknown to the learner). So, MDSA is an instance of distribution shift.

Our work and the line of work on MDSA differ in a number of ways. The main difference the MDSA literature (including [B-E]) focuses on understanding which predictors $h$ perform well for a given target function and any source distribution $D_T$, while we focus on understanding why and how it might be beneficial to train a model with a different training distribution than a test distribution (and we focus on the case when the test distribution and the error are given by a mixture of components). We also view our work as guiding the training distribution selection problem, in the sense of understanding what to look for and what direction or kind of training distribution we want. None of these is covered by the MDSA literature. In the MDSA literature, there is no notion of a training distribution (or a difference between a training and a test distribution). Namely, in MDSA, the learner is given the domain distributions D_i and the predictors h_i. So the setups (and therefore results) are quite different.

Thank you for the additional discussion.

Our setup does not solely focus on a combination of tasks

The general setup described in Section 2 explicitly considers a test distribution composed of a mixture of tasks. Moreover, the tasks used in the mixture are the same during both training and testing, which reinforces the point that the setup is indeed centered around a fixed combination of tasks.

Power law

To clarify, my original comment was not about the use of power laws, but about how it is connected to prior work that uses learned predictors with specific error rate.

One of my main concerns, as stated in my initial review, is the lack of engagement in the paper with the extensive literature on learning under distribution shift, and this remains unaddressed in both the rebuttal and the discussions. For instance, the whole domain adaptation literature (some references were already provided in my initial review) is directly relevant but never mentioned.

Relevance in practice

I agree that studying learning under mismatched training and test distributions is practically important. However, that is not the setup studied in the paper. The results provides guidance on finding the best mixture of training tasks given a mixture of the same tasks for testing. If the true test distribution (e.g., 10-year-old writers) differs from the approximate one used for evaluation (e.g., adult writers), or even differ from the one(s) available for training, then the result becomes less informative.

I am maintaining my original score. I encourage the authors to include a more thorough discussion of previous work and a better presentation of the practical relevance of the setup studied.

We thank the reviewer for their feedback. We provide answers to the reviewer’s questions.

Q1:

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

A1: We appreciate the reviewer’s feedback regarding the writing style. We feel that such a conversational style is useful for the introduction. We note that the conversational style is kept only in the introduction - we make everything formal and precise in the technical sections. If something is not precise enough, we would be very happy to make it more precise.

Q2:

What are the key differences between mixture setting and compositional setting?

A2: Our compositional setting can be viewed as a mixture setting as follows: we show that mixing the original data distribution $D_{test}$ with an auxiliary distribution $D_{uniform}$ can lead to a lower test error. In the notation of Section 2, we have $D_1 = D_{test}, p_1 = 1$ and $D_2 = D_{uniform}, p_2 = 0$. We emphasize the compositional setting in Section 5 because it is conceptually related to the orthogonal memorization setting in Section 4. In particular, the compositional setting can be understood as learning $N$ orthogonal skills and composing them together at inference time. However, unlike the orthogonal memorization, it is not possible to separate the test distribution in the compositional setting directly into a series of independent components.

Q3:

Line 63. 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)?

A3: The supports of the distributions do not have to be disjoint. For example, in the Europan and Chinese history case, a data sample could be relevant to both Chinese and European history, so it should be counted as a sample for both components. Not all distributions can be represented in this way, but this condition still captures many behaviors - see answer to question 5.

Q4:

Line 72. In the definition of L_N(p,q), from my understanding, the expectation should be taken with respect to SD_q^N and not SD_p^N, is that correct?

A4: Yes, this is a typo. The expectation should be taken with respect to $S\sim D^{N}_{q}$.

Q5:

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.

A5: Yes, in the setup of this paper, the learning model is specified by the subpopulation error functions $e_k(**n**)$. We study a few cases of different error functions capturing different problems - the orthogonal power law, the orthogonal memorization task (where the error functions follow a 0-1 law), and the transfer case. Why orthogonal power law is an excellent question. First, we take the tasks to be orthogonal as that is the base case when there is no transfer. Secondly, we take the power law because it captures many behaviors based on theoretical analysis (e.g. realizable PAC, agnostic PAC, linear regression) and also empirically ([1], [2], [3], [4]).

Q6:

Line 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}_N(q)?

A6: Yes, at the end of Line 107 it should have been $\bar{e}_N(**q**)$, and in the sentence before on Line 106 it was supposed to say $e(**n**), **n**\sim \text{Mult}(**q**,N)$ will concentrate around $e(**q**N)$. The proof for the concentration property in the orthogonal case is in Appendix A1, in particular Propositions A1-A4.

Q7:

Is there any proof of the claim in Line 106 that "\bar{e}_N(q) will concentrate around e(qN)"? Also, is \bar{e}(qN) a random variable? If so, where does the randomness come from?

A7: Line 106 was supposed to say that “$e(**n**)$ for $**n** \sim \text{Mult}(**q**,N)$ will concentrate around $e(**q**N)$ and we will sometimes exploit this in the analysis, or analyze for $\bar{e}_N(**q**) \approx e(**q**N)$”. The statement was not meant as a formal general claim but rather an indication of what the outline of the main argument is. We thank the reviewer for pointing out these typos.

Q8:

Proposition 3.2. What are the intuitions of the assumptions? Also, what is the definition of n_0(A_i,B_i,\alpha_i,p_i)?

A8: Note that rewriting $\alpha$ in this way does not decrease generality. The intuition is that the error decay will be dominated by the component with the slowest error decay (i.e. smallest $\alpha$) so we want to write the main terms of the error using these components. For the ease of presentation, we didn’t write the definition of the n_0(A_i,B_i,\alpha_i,p_i) because the formula is messy and long but it can be found in the appendix in Equations 9-12 and 14-15.

Q9:

Line 79. Does L_N^{ratio} represent the improvement or the training budget?

A9: $L_N^{\text{ratio}}$ represents the error improvement with the same training budget. $N_{\epsilon}^{\text{ratio}}$ represents the improvement in the training budget to achieve an error at most $\epsilon$.

Q10:

Line 86. What is g function?

A10: They are defined on line 89, i.e. they are any scalar functions. We will clarify that sentence in the updated version of the paper.

Q11:

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 improvement? If so, then how does the assumption that A_1=..=A_k=A affect the result?

A11: In this special case when all $A_i$ are equal, the sample complexity improvement is independent of $A$. This can be seen from the fact that the subpopulation error functions and hence the loss will multiplicatively depend on $A$, so the ratio will not depend on $A$. Corollary 3.3 deals with the special case when all tasks are equally hard (e.g. the introductory example). The dependence can be seen from Proposition 3.2. and Equations 2 and 3. If not all $A_i$ are the same it would introduce an extra term of the form $\frac{\left( \sum_{i=1}^{S} (\alpha_i p_i A_i)^{\frac{1}{1+\alpha_i}}\right)^{\alpha_1}\left( \sum_{i=1}^{S} \frac{(p_i A_i)^{\frac{1}{\alpha_1}}}{\alpha_i^{\frac{\alpha_i}{\alpha_i+1}}}\right)}{\sum_{i=1}^{S}p_i^{1-\alpha_i}A_i}$.

Q12:

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

A12: Let $S$ be the set of possible elements to memorize and $s_1,\dots, s_k \in S$. Assume that the learning rule memorizes all elements it has seen so far and let $M$ be the set of elements the model has seen. Let task $i$ be memorizing element $s_i$, i.e. we get error $0$ if $s_i\in M$ and error $1$ if $s_i \notin M$. For example, this could model fact retrieval or question answering using LLMs. The data in this case would be the text used for (pre)training the LLMs.

Q13:

Line 240. The concentration property needs to be justified, but doesn't have to be mathematically detailed.

A13: The approximation is reasonable because $e_k(**n**)$ has a nice behaving closed form and $**n**$ concentrates around $**q**N$. The proof would be similar to the orthogonal case (Proposition A2). We will add a short justification to the updated version.

Q14:

Line 275. What's the difference between \Delta^{K-1} defined here and \Delta_K mentioned in the previous text sections.

A14: Yes, they are the same thing. We will update section 7 to make the notation consistent.

Q15:

Condition 7.1. What is the intuition of this condition? What does it tell us about the model and/or the test distribution?

A15: Condition 7.1 represents the very special case when there is no positive distribution shift, so the intuition is that it restricts the subpopulation error function to a very specific set of functions that do not represent a typical case. In the orthogonal case, Condition 7.1 implies that all the subpopulation error functions in terms of the mixture are constant. Lemma 7.5 further shows that Condition 7.1 uniquely determines the subpopulation error functions in terms of the mixture up to a constant.

Q16:

The related works can be put into an early section, and the last section should include a brief conclusion of the work and addressing of the limitations.

A16: We will add a conclusion section at the end and also address limitations and future work. Regarding the related works section, we put it at the end of the paper so that a) we get to the main results faster and b) the flow of the paper is more natural this way because our paper is not directly based on any one previous work.

[1] Hoffman et al. An empirical analysis of compute-optimal large language model training. NeurIPS 2022. [2] Ye et al. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. ICLR 2025. [3] Kaplan et al. Scaling Laws for Neural Language Models. 2020. [4] Cherti et al. Reproducible scaling laws for contrastive language-image learning. CVPR 2023.

We would like to thank the reviewer for the thoughtful and constructive review. We are grateful that the reviewer found our research topic “super interesting and important,” and our findings offer “quite unintuitive and thus very valuable insight.” We also appreciate your positive comments on the soundness of our theoretical approach and the backing of appropriate experiments. We agree with the reviewer that an interesting future direction is when the task distribution is not known at test time, and we will leave it for future work.

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

A: Thank the reviewer for the constructive suggestions. The mixing ratio 30% here is a hyperparameter, and in our preliminary experiments, we found that setting this hyperparameter to a number between 10% ~ 70% generally improves the test accuracy. This method is motivated by our calculation for the orthogonal memorization task: properly upsampling low-probability skills can lead to better memorization (Line 197), and mixing the original data distribution with $D_{uniform}$ is one way to effectively upsample low-probability skills. We will expand the discussion of the motivation in our camera-ready version.

Using 100% uniform sampling would lead to very high test errors. It is worth noting that each data point in $D_{uniform}$ only contains one skill (which corresponds to just one step of the CoT for skill composition), and thus the model will fail to learn how to compose skills with CoT if it only trains on $D_{uniform}$. Only by mixing the original data distribution with $D_{uniform}$ can the model learn the skills and how to compose them properly. This is why we mix them together rather than using just one data distribution.

Overall, our experiments here demonstrate that mixing the original data distribution with another distribution can lead to a better test error than just training on the original one. The mixed distribution with 30% mixing ratio we tested is one way to achieve this goal, and we do not claim that the mixed distribution is optimal.

We thank the reviewer for their feedback.

Q: This is just rediscovering motivation for active learning:

A: The active learning setting, as typically studied, is rather different, even if relaxed, from the setup of our paper. For example, pool based active learning, where individual examples can be selected from a sample, or query based are both different, and almost always the benefit comes from being adaptive to the data so far.

In any case, we could not find a discussion of a mixture setting with analysis of the optimal proportions and discussion of being better with a train distribution different from the test. If we are missing something, we would be happy to get the references and see what the reviewer has in mind.

Q: The conclusions are on their face obvious:

A: Furthermore, based on a small informal survey of the experts in the field with whom we discussed our results while writing the paper, the results did not seem so immediately obvious. Yes, the results might be obvious in hindsight because the math is very simple. But we need to ask the reviewer whether the answer to the motivating example in the introduction was immediately obvious to them before reading on.