Mo' Data Mo' Problems: How Data Composition Compromises Scaling Properties

We thank the reviewer for raising their concerns; however, our contributions may have been fundamentally misunderstood by the reviewer.

Main Contribution

So, even if there are no publications on the topic, e.g., the fact of having shift of distributions (and thus, lower performance) in the case you sample not from the target distribution is common knowledge and directly follows from ML grounds.

The reviewer is correct that the fact that “distribution shift between train and test datasets results in reduced performance” has been studied extensively in past work (See our references Sagawa et. al. in the main text and Appendix C.5 detailed related work on distribution shift).

However, what is not known is (1) why distribution shift may arise as dataset size increases and (2) do the harms induced by distribution shift from scaling eclipse the benefits of scaling. This is an incredibly common scenario in the real world. For example, government resource allocation for social programs depends on projections from census data [1]. In smaller counties and states, the decision has to be made on whether adding more data is beneficial or harmful to achieve more accurate and equitable predictions. Moreover, formalizing why adding data does not yield better outcomes is important and it is a central contribution to our work.

What is the problem to be solved? If the practitioner face a situation, knowing D, then they will sample from it. If they have no access to D, then they try to sample some distribution such that it is most close to D.

Without access to D, how is it possible to know what distribution is closest to D? Our central question (Figure 1) is: “How does data composition impact model performance under scaling?”. We examine the situation where there is limited access to D (e.g. a very small sample). A problem arises when we have to decide whether adding data from some other distribution $D_{s_2}$, is beneficial or harmful with respect to overall accuracy, worst group accuracy, and group accuracy disparity.

Seems that it means that the practitioner believes that D_s_1 IS the distribution D. So, when receiving D_s_2 they should take D_s_1 united with D_s_2 as D, or not considering D_s_2 at all.

Suppose we are policymakers in South Dakota, $D_{s_1}$ is the dataset from our state, but $D_{s_2}$ is a very different state like California. If $D_{s_2}$ is not considered at all, $D_{s_1}$ might be too small and uniform to give adequate accuracy/fairness (e,g, only 2.6% of the population is Black [2] ). If we sample from the union of $D_{s_2}$ and $D_{s_1}$, the result would depend on the relationship between $D_{s_2}$ and $D_{s_1}$ (this is captured in our Mixture model Sec. 3).

Crucially, we know that California is a different state but we do not know how it is different. In fact, the Mixture model (see Section 3), would sample from the union of South Dakota and California and we see that it yields worse results than just $D_{s_1}$ in the sequential model with a smaller n (Figure 3a vs Figure 4a). We also know that Louisiana could be as different from South Dakota as California is. However, adding data from Louisiana improves outcomes (Figure 9 vs Figure 3a). We identify Louisiana as a useful state to add using the heuristic we propose (Equation 1).

Answers to the reviewer’s questions:

I see two closing parentheses while having only one open one.

This is a typo where there is an extra closed parenthesis symbol that we will update the paper to exclude the second “)”.

What do you mean saying “that is similar to the test distribution” it is unclear

In our problem setup, the first source $D_{s_1}$ is very small and limited. For example, if you are a very small state (e.g. South Dakota) or a local hospital with a limited number of ICU records. We assume that to evaluate the trade-offs between training on more data vs induced distribution difference, the deployment dataset may be different from the first source but is likely to be most similar to the first source (e.g. $D_{s_1}$).

I do not understand: at the beginning of the work it was stated that “test set” is the union. Here, it is stated that it is just D_d_1.

In section 3, we state that in the data accumulation setting, the training set is a union of sources. Could the reviewer tell us specifically where we said the test set is the union?

We hope our explanations have clarified the novelty of our contributions. As far as we are aware, we agree with the reviewer that there is no prior work that models the mechanisms that induce distribution shift when dataset size is scaled up. Moreover, our theoretical contributions are validated by our empirical results in different settings, and our practical metric bridges our theoretical model with empirically observed divergences. We are happy to continue answering any questions to help the reviewer better understand the contributions of our paper.

We thank the reviewer for their helpful suggestions on our work and we value their positive comments on the quality and presentation of our work.

Please see our answers to reviewer questions below:

1. This is a good point, we will indeed rephrase this sentence to better acknowledge prior work in the distribution shift (in addition to our discussion in C.5 of our supplementary materials).

2. We thank the reviewer for pointing out related work in active learning. We agree that it is important and we will include this in our work. We do highlight that our KL metric is distributional rather than for specific instances as described in equation 15 of Fu et. al. While Zhang et. al. takes a, notably different, model change approach to stop data addition, we believe it is still very wonderful prior work we will include to frame our contribution.

3. In our experimental setup, we specifically wanted to focus on data interventions rather than model interventions suggested by domain adaptation strategies. In our supplementary materials C.5 we discuss in detail our work in relation to domain adaptation in particular.

4. Excellent question. While in our results we show that XGBoost does not suffer from an accuracy drop, the model still performs much worse than it would if more data from a more similar distribution were to be added. In Figures 11b and 11c (supplementary materials), we show that XGB continues to improve substantially when other sources are added like Louisiana.

We value these helpful comments for the reviewer. We believe that taking this data composition framework to understand scaling under possible distribution shifts is an important direction that will inspire future work in richer domains. We are happy to answer any further questions or add experiments and hope that the reviewer will vouch for our work!

We thank the reviewer for taking the time to read and understand our work. We agree with the reviewer that our work is useful in the datasets track – to help formalize practical questions raised by the distribution shift that arises from data scaling.

Answering Reviewer Questions

1. The reference test set is constructed in our experiments from a very small initial training source $D_{S_1}$ (e.g. South Dakota). We make this assumption to simulate a setting where there is some limited access to the initial dataset of interest but it is then unclear if adding more data would help with outcomes on this initial dataset. We assume that the eventual deployment set might be different from the first source but is likely most similar to the first source.
2. We train from scratch at each dataset size presented in our plots. Thus, each training trajectory shows 5 models over 50 seeds (not 5 - typo in the main text) at for each size n.

Implications of our work

The implications of this work are not clear. In real-world settings, if there exists a datasets similar to one that we currently have but has high divergence does it mean it should not be included in the analysis? Doesn’t not doing so restrict the model from better generalising? Eg. Yelp reviews in MN vs SD.

In the real world, if we knew a dataset is vastly different (e.g. loan applications for credit cards vs loan applications for mortgages), it does make sense to not include it. However, the question that arises is, what level of divergence would harm my model, and what level is acceptable? In all of our datasets, we are looking at the differences that might arise within a dataset. We show that adding data from states (near and far) with reasonable divergences (e.g. there is SOME shift) improves outcomes. Especially when the starting dataset is small (and undiverse) excluding data in the same dataset might not be the best choice. This is a very nuanced point and we will better explain it in our paper.

Thirdly, it looks like adding more data reduces performance disparity between groups and in general helps the least performing group. Reducing disparity is perhaps indicating that the model is generalising better and getting more robust and these should be good things.

Reducing performance disparity by adding more data is indeed a desirable goal. We hope to show that by measuring not just overall accuracy but also disparity and worst group performance, understanding the consequences of data scaling is a multi-faceted pursuit and KL divergence as a lone metric is only just a first heuristic. We hope our work will inspire future work to develop further metrics for determining when to add data for properties like disparity and worst group robustness in particular.

We hope our responses have clarified the questions raised. If we can further provide any explanations to improve the reviewers confidence in our work, please do not hesitate to raise further questions.

We thank the reviewer for reading through our paper and raising interesting questions. We would like to reiterate the contribution and scope of our paper in our response and clarify some misconceptions.

Answers to the reviewers’ questions:

Could the authors provide an appropriate justification or real-world scenario as an example for concentrating exclusively on tabular datasets? In this scenario, if models are not deep learning based would they be prone to large data scaling issues?

Tabular datasets remain fundamental in many applications such as surveys (e.g. census, public health), loan decisions, recidivism prediction, portfolio optimization, and recommender systems. For fairness questions, in particular, Adult and COMPAS are, by far, the most studied datasets by far ([1] Figure 1). Given known issues with the Adult dataset, Folktables (Ding et al.) was developed to provide a wider sample of data for developing policy decisions based on census data. How this much larger sample of census data should be used is a question we seek to answer through our work; as we show, using all available data is not necessarily the best solution. We point to Figure 10 in the appendix where we show the benefits of increasing dataset size. We see that accuracy and worst group accuracy both improve while group disparities decrease monotonically with more data in the tabular setting. This motivates our work in questioning whether these benefits of more data can be applied to all states, even the smaller ones like South Dakota.

In this simpler paradigm with models such as Logistic Regression etc, instead of continuously adding more data for generalizability, would it not make sense to just focus on approaches that curb distribution shift and retrain the model cyclically over certain time periods?

We want to emphasize that state-of-the-art models for tabular datasets are not deep learning models, but tree-based methods and simply MLPs (NN, GB, XBG in Section 5) [2, 3]. In the settings we study, extremely large general tabular datasets are available, but the benefits of incorporating massive tabular data for a specific task in a subset of the data are unknown. Our work follows a line of data-centric works that study the effect of data composition on different state-of-the-art models for the task at hand. In fact, our excess KL metric is designed to evaluate distribution shifts induced by new data sources in order to select the best dataset source. Our technique is a version of curbing distribution shift.

Response to Weakness

LLMs and scope of work:

The focus on tabular datasets significantly narrows the paper's scope, especially considering that data scaling is a critical concern in large language models (LLMs) and other domains outside of tabular data (such as NLP and Vision).

Indeed, we agree with the reviewer that running similar experiments on LLMs, and maybe also VLMs would increase the scope of our paper. However, due to resource constraints, pretraining LLMs from scratch with different data sources is currently outside of our capabilities and our scope. However, we emphasize that the folklore of more data being better is ubiquitous in machine learning even outside of the LLM setting. Furthermore, the question of how many EHR, loan approval, and recommendations datasets should be combined is important for practical decision-making. Our paper gives guidance in this common setting in terms of how data composition can be related to scaling matters.

Generalizability

Furthermore, the results obtained (as well as the approach undertaken) are highly dataset dependent (for e.g., what if I sample the reference test set differently for Folktables? Or use a different state altogether?)

We agree that using a different state altogether might result in different results. However, the contribution of our work is not to claim that “data scaling benefits are always eclipsed by distribution shift” (this conclusion cannot be true if $D_{s_2}$ is sufficiently close to $D_{s_1}$). What we are trying to answer is “How does data composition impact model performance under scaling” (Figure 1 of main text). Our contribution is to give a principled theoretical model for data accumulation that is flexible to describe effects regardless of the dataset. Our empirical results show a there exists claim where we want to urge caution when adding data to illustrate that there are splits of datasets that lead to undesirable effects.

In summary, we rigorously study the question of when and why data composition matters when accumulating training data for improved model performance. Our question is practically grounded in many real-life scenarios where the cleanest source of data is very small, and care must be taken on when and how to supplement the available training dataset for improved performance. We hope this clears up concerns raised by all the reviewers, and we thank the reviewers for engaging this scholarly process with us by carefully reading our paper and considering our contributions.

References [1] Fabris, A., Messina, S., Silvello, G., & Susto, G. A. (2022). Tackling documentation debt: a survey on algorithmic fairness datasets. In Equity and Access in Algorithms, Mechanisms, and Optimization (pp. 1-13).