Predicting the impact of dataset composition on model performance

1 INTRODUCTION . The success of large scale machine learning systems depends critically on the quantity and quality of data used during training , and we can not expect these systems to succeed if there is not enough training data or if that data does not cover all the phenomena contained in the test distribution ( BenDavid et al. , 2010 ) . Knowing this , the designer of a machine learning system might create multiple sources of data , with each one targeting a different feature or domain that the model ought to do well on ( Crammer et al. , 2007 ; Wang et al. , 2019a ) . This data-driven design strategy provides powerful tools to improve and evaluate model behavior , but also poses an additional challenge : what is the right way to combine these various data sources ? What is the optimal data collection policy for a given budget ? Our goal is to answer these questions by quantifying the relationship between data sources and model performance – how well will our model do if we were to train it on n samples using a data mixture ( q1 . . . qk ) over ourK data sources . A precise model for predicting model performance will allow us to both identify the optimal data collection policy and quantify cost-performance tradeoffs . The starting point of our work is the recent observation across speech , vision and text ( Hestness et al. , 2017 ; Kaplan et al. , 2020 ; Rosenfeld et al. , 2020 ) that the empirical performance of a model is remarkably predictable , and follows the log-linear formula log ( error ) ≈ −α log ( n ) + C. ( 1 ) In this work , we expand this observation to the multi-data-source setting and discover the surprising fact that the slope of the log-linear relationship ( α ) does not vary with data composition and that the data composition only affects the intercept ( C ) . The simple dependence of log-error on data size allows us to reduce the problem of estimating model error into a learning problem . Our approach is straightforward : we hypothesize that model error follows V ( n , q ) : = exp ( −α log ( n ) +log ( C ( q ) ) ) for a simple parametric functional formC ( q ) , and fit this to observed pairs of ( n , q , error ) that we obtain by subsampling the dataset and re-training a model . We show that there is a natural and simple choice of C ( q ) as a rational function that we derive from optimal experimental design for linear regression , M-estimation , and nonparametric smoothing . The simple and parametric dependence of V ( n , q ) on n allows us to use our resulting estimates to predict model performance under substantial extrapolation in data size . Empirically , the resulting predictions are extremely accurate and hold under substantial extrapolation . On the Amazon review prediction dataset ( Mansour et al. , 2009 ) , we can learn to predict model performance nearly perfectly ( r2 = 0.96 ) from a small dataset of 1200 examples across 3 sources and extrapolate to predict the model error on datasets of up to 4000 examples . We show this high accuracy continues to hold on a real-world task oriented dialogue system ( r2 = 0.93 ) , a multi-domain machine translation system ( r2 = 0.83 ) , and boolean question answering with weak supervision ( r2 = 0.86 ) . In each of the cases , our proposed approach substantially outperforms the best baseline , with the baselines performing worse-than-random in both the machine translation and question answering tasks . Related work Quantifying the effect of data composition on model performance is closely related to the classical ideas of optimal experimental design , as well as more recent machine learning methods such as active learning and data valuation . Our work will draw inspiration from the classical V -optimal experimental design ( John & Draper , 1975 ) as a way to understand how model performance will change with the data collection policies . However , our approach differs substantially beyond this . Instead of making strong linearity assumptions and identifying closed form formulas for model performance , we treat identifying the impact of data sources on errors as itself a prediction problem , which allows us to quantify these effects for neural networks and non-separable objectives . Active learning provides methods for incrementally selecting new points to rapidly reduce a loss ( Hanneke , 2007 ) . These approaches only consider the problem of optimal data collection and do not seek to predict model performance under all data collection strategies ( including suboptimal ones ) , which is critical when making cost-performance tradeoffs across data sources . The model performance predictions produced in our work complements existing work on active learning by providing accurate forecasts of model performance under different data collection strategies . Finally , data valuation methods such as the Shapley value attempt to assign estimate the impact of a data source on model performance ( Ghorbani & Zou , 2019 ; Jia et al. , 2019 ; Ghorbani et al. , 2020 ; Yoon et al. , 2019 ) . These approaches are natural when pricing data sources as part of a market mechanism ( Ohrimenko et al. , 2019 ; Agarwal et al. , 2019 ) due to the axiomatic properties of the Shapley value . Our approach differs in that we seek simply to estimate the performance of a model rather than to assign a single price to examples from a data source . This difference means that axioms such as additivity that are critical for the Shapley value are not relevant for our goal . We show that for the purpose of predicting errors , a rational function ( rather than a linear cost ) follows naturally from optimal experimental design . Our experiments also show that our rational function approximation provides better model performance predictions than a linear , additive model . 2 PROBLEM STATEMENT AND EMPIRICAL OBSERVATIONS . Our goal is to predict the performance of a model as a function of the number of training samples n as well as the dataset composition q , where qk represents the fraction of the training data drawn from data source k. We will now define this goal more formally in terms of the training data distribution , model fitting , and test loss . The training data consists of an n-sample training set pn , q that is created by sampling from the mixture p : = ∑ k∈ [ K ] qkpk where pk are data generating distributions for each of theK data sources and qk are mixture weights with qk ≥ 0 and ∑ k∈ [ K ] qk = 1 . Using this dataset , we learn a prediction model θ̂ that incurs loss ` ( θ̂ ; x , y ) for a training example ( x , y ) . The fitted model is the empirical loss minimizer , which we define as θ̂ ( pn , q ) : = arg min θ∈Θ Epn , q [ ` ( θ ; x , y ) ] . The performance of this classifier is evaluated on a test distribution which may differ from the training distribution by a covariate shift ( i.e . p ( y | x ) = ptest ( y | x ) ) . We are interested in model performance as a function of the data size and composition ( and not a fixed empirical distribution pn , q ) and thus our goal is to predict the model ’ s expected excess loss over draws in both the training and test distributions , L ( n , q ) : = E [ ` ( θ̂ ( pn , q ) ; x , y ) ] − inf θ E [ ` ( θ ; x , y ) ] . Estimating L requires that we hypothesize a relationship between ( n , q ) and the expected model loss . Following earlier observations by Hestness et al . ( 2017 ) , we expect a log-linear relationship between L ( n , q ) and log ( n ) for any fixed q , which implies a possible approximation as log ( L ( n , q ) ) ≈ log ( V ( n , q ) ) : = α ( q ) log ( n ) + C ( q ) . ( 2 ) We now examine this hypothesis in a simple toy example . Linear toy data : We will start with the simplest nontrivial example of linear least-squares regression to study L ( n , q ) . In this example , there are two data sources over x ∈ R2 . The first data source has substantial variability on the first coordinate x0 but not x1 and vice versa for the second data source . The overall generative process is y | x ∼ [ 0.5 , 1 ] > x+ z ∼ Bern ( q ) ∼ N ( 0 , 1 ) x | z = 0 ∼ N ( 0 , [ 1 0 0 0.001 ] ) x | z = 1 ∼ N ( 0 , [ 0.001 0 0 1 ] ) . Let L ( n , q ) be the excess squared loss of a linear least squares model trained with n samples from a mixture q and evaluated on a test distribution with q = 0.5 . What will L ( n , q ) look like ? Figure 1a shows a clear linear relationship between log dataset size ( log ( n ) ) and log ( L ( n , q ) ) . The intercept of the linear relationship seems to vary with the data mixture q , but the slope seems constant . Examining Figure 1a more closely , we find that the extremes of using either data source exclusively ( blue / purple lines ) performs worse than a mix suggesting that log ( L ( n , q ) ) is unlikely to be linear in q . Intuitively , we can think of each data distribution as having a different strength ( i.e . more variance in either x0 or x1 ) and combining the two results in a better data distribution than either alone . We can see this more clearly when we estimate the intercept for each of these lines ( Figure 1b ) . The estimated intercepts show a U-shaped curve that rapidly increases as q → 0 or q → 1 and is generally flat from 0.2 to 0.8 . 3 METHOD AND THEORY . We have observed that in the case of a simple linear regression , the log-error not only follows the relationship outlined in equation 2 , but also that the slope α is constant as we vary the data composition ( and we will further validate this claim on more complex tasks and models in subsequent sections ) . This observation shows we may be able to further simplify the log-linear approximation as log ( L ( n , q ) ) ≈ log ( V ( n , q ) ) : = −α log ( n ) + log ( C ( q ) ) . Now note that this functional form decouples the data size n and mixture proportions C ( q ) into two terms . This is the key observation of our work : log ( V ( n , q ) ) has a very simple dependence on n , and the more complex term C ( q ) has no dependence on n. Therefore we can cast this as a learning problem , where we learn α and a parametric function Cλ ( q ) based on the model ’ s error over a range of q and small n , and extrapolate this for large n using the log-linear dependence of log V on n. Concretely , given a dataset with { n1 . . . nk } we can generate a subsampled dataset with n̂k ∼ Unif ( 0 , nk ) samples from each source . This results in a training set with data size n̂ = ∑ k n̂k and composition q̂k = n̂kn̂ . We fit a model to this subsampled data and compute its loss L ( n̂ , q̂ ) . Given the triple ( n̂ , q̂ , L ( n̂ , q̂ ) ) we can now simply fit the hypothesized functional form , min λ , α Eq̂ , n̂ [ ( log ( L ( n̂ , q̂ ) ) − α log ( n̂ ) + log ( Cλ ( q̂ ) ) ) 2 ] . The experimental data does not specify the functional form of Cλ ( q ) except that it should handle convex functions like those seen in Figure 1b . We will now study V ( n , q ) theoretically and argue that a natural choice is the rational function Cλ ( q ) : = M∑ i=1 ( K∑ k=1 λikqk ) −1 . In the subsequent sections , we will study three settings : ordinary linear regression , M-estimation , and nonparametric regression and show that our hypothesized log-linear approximation arises naturally in all three cases .

Summary: This work studies the problem of predicting model performance with more training data when the data are collected from different sources. The predictor is a function of the number training examples, and the ratio of examples from each source. The predictor needs to be built from a small number of training examples the observed model performance, and applied to larger numbers of training example without actually training the model. The predictor can be used to decide a good data collection policy which is expected to have best model performance. The proposed solution is a simple parametric form of the predictor, which is log-linear in the log of # training examples, and log-rational in the source distribution vector. The solution is motivated by recent literature about the same task for single source. The correctness of the solution is proved for several cases: linear regression, M-estimator and nonparametric binning. The performance of the predictor is then evaluated for a number of real-world tasks: linear regression for Amazon book rating, semantic parsing, machine translation and multitask question answering. The performance is measured by r2 score between the actual performance and predicted performance. The proposed predictor has a clear advantage over the baseline of using a linear predictor.