DIVA: Dataset Derivative of a Learning Task

1 INTRODUCTION . Consider the following seemingly disparate questions . ( i ) Dataset Extension : Given a relatively small training set , but access to a large pool of additional data , how to select from from the latter samples to augment the former ? ( ii ) Dataset Curation : Given a potentially large dataset riddled with annotation errors , how to automatically reject such outlier samples ? ( iii ) Dataset Reweighting : Given a finite training set , how to reweight the training samples to yield better generalization performance ? These three are examples of Dataset Optimization . In order to solve this problem with differentiable programming , one can optimize a loss of the model end-to-end , which requires differentiating the model ’ s loss with respect to the dataset . Our main contribution is an efficient method to compute such a dataset derivative . This allows learning an importance weight αi for each datum in a training dataset D , extending the optimization from the weights w of a parametric model such as a deep neural network ( DNN ) , to also include the weights of the dataset . As illustrated in the following diagram , standard optimization in machine learning works by finding the weights wα that minimize the training loss Ltrain ( w , Dα ) = ∑ i αi ` ( fw ( xi ) , yi ) on a given ( weighted ) dataset Dα ( dark box ) . We solve a more general learning problem ( light box ) by jointly optimizing the dataset Dα in addition to w. To avoid the trivial solution α = 0 , it is customary in AutoML to optimize Dα by minimizing the validation error computed on a disjoint dataset . This makes for inefficient use of the data , which has to be split between training and validation sets . Instead , we leverage a closed-form expression of the leave-one-out cross-validation error to jointly optimize the model and data weights during training , without the need to create a separate validation set . Ltrain ( w , Dα ) wα Lval ( wα ) Standard optimization finds the best weights to minimize the training loss Dα DIVA finds the best dataset to train on in order so that the weights minimize a validation loss Dα Ltrain ( w , Dα ) wα Lval ( wα ) Dα Ltrain ( w , Dα ) wα Lval ( wα ) Dα The intermediate block in the diagram ( which finds the optimal weights wα for the the training loss on Dα ) is usually non-differentiable with respect to the dataset , or the derivative is prohibitively expensive to compute . DIVA leverages recent progress in deep learning linearization Achille et al . ( 2020 ) , to derive a closed-form expression for the derivative of the final loss ( validation error ) with respect to the dataset weights . In particular , Achille et al . ( 2020 ) have shown that , by replacing cross-entropy with least-squares , replacing ReLu with leaky-ReLu , and performing suitable preconditioning , the linearized model performs on par with full non-linear fine-tuning . We also leverage a classical result to compute the leave-one-out loss of a linear model in closed-form ( Rifkin & Lippert , 2007 ; Green & Silverman , 1993 ) . This allows us to optimize the LOO loss without requiring a separate validation set , setting DIVA apart from bi-level optimization customary in AutoML . To illustrate the many possible uses of the dataset derivative , we run experiments with a simplified version of DIVA to cleanup a dataset of noisy annotations , to extend a training set with additional data from an external pool , identify meaningful data augmentation , and to perform multi-modal expansion using a CLIP model ( Radford et al. , 2021 ) . Rather than using the full linearization of the model derived by Achille et al . ( 2020 ) , we restrict the gradient to its last layer , cognizant that we are not exploiting the full power of LQF and thereby obtaining only a lower-bound of performance improvement . Despite that restriction , our results show consistent improvements from dataset optimization , at the modest computational cost of a forward pass over the dataset to optimize the importance weights . To summarize , our main contributions are : 1 . We introduce a method to compute the dataset derivative in closed form , DIVA . 2 . We illustrate the use of DIVA to perform dataset optimization by minimizing directly the leave-one-out error without the need for an explicit validation dataset . 3 . We perform experiments with a simplified model that , despite not using the full power of the linearization , shows consistent improvements in dataset extension , re-weighting , outlier rejection and automatic aggregation of multi-modal data . Our method presents several limitations . The dataset derivative of a learning task is computed around a point represented by a pre-trained model . It only allows local optimization around this point . Moreover , we only compute a restriction of the linearization to the dimensions spanned by the last few layers . In general , this yields suboptimal results compared to full global optimization from scratch , if one could compute that at scale . Nonetheless , the linearized setting is consistent with the practice of fine-tuning pre-trained models in light of the results of Achille et al . ( 2020 ) , see also ( Radford et al. , 2021 ; Rezende et al. , 2017 ; Mormont et al. , 2018 ; Hosny et al. , 2018 ) . 2 RELATED WORK . AutoML . State of the art performance in image classification tasks often relies on large amount of human expertise in selecting models and adjusting the training settings for the task at hand ( Li et al. , 2020a ) . Automatic machine learning ( AutoML ) ( Feurer et al. , 2019 ; He et al. , 2021 ) aims to automate model selection ( Cawley & Talbot , 2010 ) and the training settings by instead using meta-algorithms for the different aspects of the learning settings . Such methods follow a bi-level optimization framework , optimizing the training settings in the outer level , and traditional model optimization in the inner level ( Jenni & Favaro , 2018 ) . AutoML has focused on achieving better results via automatic model selection ( Deshpande et al. , 2021 ; Feurer et al. , 2019 ) including neural architecture search ( NAS ) ( Zoph & Le , 2016 ; Elsken et al. , 2019 ; Liu et al. , 2019 ) . Other important AutoML topics include hyper-parameter selection ( Li et al. , 2017 ; Akiba et al. , 2019 ) and data augmentation ( Cubuk et al. , 2018 ; Lim et al. , 2019 ; Chen et al. , 2020 ; Behl et al. , 2020 ) , which are closer to our settings of optimizing the dataset weights . Since the main signal for a model ’ s performance is the final validation loss , which requires full optimization of the model for each evaluation , AutoML approaches often incur a steep computational costs . Alternatively , other methods follow alternating optimization of the criteria , such as the work of Ren et al . ( 2018 ) that approximates full network optimization with a single SGD step to learn to reweight the training set dynamically . Differentiable AutoML alleviates outer-optimization costs while optimizing the final validation error via differentiable programming , by utilizing proxy losses and continuous relaxation that enable differentiation . Different approaches to differentiable AutoML include differentiable NAS ( Liu et al. , 2018 ; Wu et al. , 2019 ) , data augmentation ( Liu et al. , 2021 ; Li et al. , 2020b ) , and hyper-parameter optimization ( Andrychowicz et al. , 2016 ) . The DIVA dataset derivative follows the differentiable AutoML framework by enabling direct optimization of the dataset with respect to the final validation error of the model . Importance sampling . While our dataset optimization problem may seem superficially similar to importance sampling , the optimization objective is different . Importance sampling aims to reweight the training set to make it more similar to the test distribution or to speed up convergence . On the other hand , DIVA objective is to optimizes a validation loss of the model , even if this requires making the training distribution significantly different from the testing distribution . Importance sampling methods have a long history in the MCMC machine learning literature where the sampling is conditioned on the predicted importance of samples ( Metropolis & Ulam , 1949 ; Liu , 2008 ) . In deep learning , importance sampling methods have been studied theoretically for linearly-separable data ( Byrd & Lipton , 2019 ) and recently in more generality ( Xu et al. , 2021 ) . Furthermore , there exist many importance sampling heuristics in deep learning training including different forms of hard sample mining ( Shrivastava et al. , 2016 ; Xue et al. , 2019 ; Chang et al. , 2017 ) , weighting based on a focal loss ( Lin et al. , 2017 ) , re-weighting for imbalance , ( Cui et al. , 2019 ; Huang et al. , 2019 ; Dong et al. , 2017 ) and gradient based scoring ( Li et al. , 2019 ) . We emphasize that DIVA ’ s optimization of the sample weights is not based on a heuristic but is rather a differentiable AutoML method driven by optimization of a proxy of the test error . Further , DIVA allows optimization of the dataset weights with respect to an arbitrary loss and also allows for dataset extension computation . LOO based optimization . Leave-one-out cross validation is well established in statistical learning ( Stone , 1977 ) . In ridge regression , the LOO model predictions for the validation samples have a closed-form expression that avoids explicit cross validation computation ( Green & Silverman , 1993 ; Rifkin & Lippert , 2007 ) enabling efficient and scalable unbiased estimate of the test error . Efficient LOO has been widely used as a criterion for regularization ( Pedregosa et al. , 2011 ; Quan et al. , 2010 ; Birattari et al. , 1999 ; Thapa et al. , 2020 ) , hyper-parameter selection ( Hwang & Shim , 2017 ) and optimization ( Wen et al. , 2008 ) . Most similar to our dataset derivative are methods that : ( 1 ) optimize a restricted set of parameters , such as kernel bandwidth , in weighted least squares ( Cawley , 2006 ; Hong et al. , 2007 ) ( 2 ) locally weighted regression methods ( memorizing regression ) ( Atkeson et al. , 1997 ; Moore et al. , 1992 ) , or ( 3 ) methods that measure the impact of samples based on LOO predictions ( Brodley & Friedl , 1999 ; Nikolova et al. , 2021 ) . Dataset selection & sample impact measures . Koh & Liang ( 2017 ) measure the effect of changes of a training sample weight on a final validation loss through per-sample weight gradients , albeit without optimizing the dataset and requiring a separate validation set . Their proposed expression for the per-sample gradient , however , does not scale easily to our problem of dataset optimization . In contrast , in proposition 3 we introduce an efficient closed-form expression for the derivative of the whole datasets . Moreover , in proposition 3 , we show how to optimize the weights with respect to a cross-validation loss which does not require a separate set . In Pruthi et al . ( 2020 ) , the authors present a sample-impact measure for interpretability based on a validation set ; for dataset extension , Yan et al . ( 2020 ) presents a coarse dataset extension method based on self-supervised learning . Dataset distillation and core set selection methods aim to decrease the size of the dataset ( Wang et al. , 2018 ) by selecting a representative dataset subset ( Hwang et al. , 2020 ; Jeong et al. , 2020 ; Coleman et al. , 2019 ; Joneidi et al. , 2020 ; Trichet & Bremond , 2018 ; Killamsetty et al. , 2021 ) . While DIVA is capable of removing outliers , in this work we do not approach dataset selection from the perspective of making the dataset more computationally tractable by reducing the number of samples .

Summary: This paper proposes a method for dataset optimization that learns sample weights without a separate validation dataset. By using squared loss and linearization around a pre-trained model, derivative w.r.t. the sample weights can be obtained as closed-form, which realizes efficient end-to-end learning. In addition, this paper uses leave-one-out cross-validation (LOOCV) loss instead of validation loss on a separate dataset, which enables to conduct evaluation without a separate validation dataset. The experiments are conducted with multiple tasks (Dataset AutoCuration, Dataset extension, detrimental sample detection, multi-modal learning, and data augmentation).