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: The paper proposes a novel method for computing the derivative of the weights assigned to datapoints in a dataset with respect to the generalization error of a learning task. The core idea being that one can take a pretrained model, and compute its derivative with respect to a fine-tuning task/ dataset and understand which datapoints are important and which datapoints are not important. This can then be used for tasks such as 1) outlier detection, 2) dataset expansion, and 3) dataset reweighting. Comparisons of the proposed work to other previous work reveals that the current work achieves better results on dataset reweighting.