Year: 2021

Abstract:
We present a method to compute the derivative of a learning task with respect to a dataset. A learning task is a function from a training set to the validation error, which can be represented by a trained deep neural network (DNN). The “dataset derivative” is a linear operator, computed around the trained model, that informs how perturbations of the weight of each training sample affect the validation error, usually computed on a separate validation dataset. Our method, DIVA (Differentiable Validation) hinges on a closed-form differentiable expression of the leave-one-out cross-validation error around a pre-trained DNN. Such expression constitutes the dataset derivative. DIVA could be used for dataset auto-curation, for example removing samples with faulty annotations, augmenting a dataset with additional relevant samples, or rebalancing. More generally, DIVA can be used to optimize the dataset, along with the parameters of the model, as part of the training process without the need for a separate validation dataset, unlike bi-level optimization methods customary in AutoML. To illustrate the flexibility of DIVA, we report experiments on sample auto-curation tasks such as outlier rejection, dataset extension, and automatic aggregation of multi-modal data.

heading: 1 INTRODUCTION
text: 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).

heading: 2 RELATED WORK
text: 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.

heading: 3 METHOD
text: In supervised learning, we use a parametrized model fw(x) to predict a target output y given an input x coming from a joint distribution (x, y) ∼ T . Usually, we are given a training setD = {(xi, yi)}Ni=1 with samples (x, y) assumed to be independent and identically distributed (i.i.d.) according to T . The training set D is then used to assemble the empirical risk for some per-sample loss `,
Ltrain(w;D) = N∑ i=1 `(fw(xi), yi),
which is minimized to find the optimal model parameters wD:
wD = argmin w
Ltrain(w;D).
The end goal of empirical risk minimization is that weights will also minimize the test loss, computed using a separate test set. Nonetheless D is often biased and differs from the distribution T . In addition, from the perspective of optimization, different weighting of the training loss samples can enable or inhibit good learning outcomes of the task T (Lin et al., 2017). Dataset Optimization. In particular, it may not be the case that sampling the training setD i.i.d. from T is the best option to guarantee generalization, nor it is realistic to assume that D is a fair sample. Including in-distribution samples that are too difficult may negatively impact the optimization, while including certain out-of-distribution examples may aid the generalization on T . It is not uncommon, for example, to improve generalization by training on a larger dataset containing out-of-distribution samples coming from other sources, or generating out-of-distribution samples with data augmentation. We call Dataset Optimization the problem of finding the optimal subset of samples, real or synthetic, to include or exclude from a training set D in order to guarantee that the weights wD trained on D will generalize as much as possible.
Differentiable Dataset Optimization. Unfortunately, a naïve brute-force search over the 2N possible subsets of D is unfeasible. The starting idea of DIVA is to instead solve a more general continuous optimization problem that can however be optimized end-to-end. Specifically, we parameterize the choice of samples in the augmented dataset through a set of non-negative continuous sample weights αi which can be optimized by gradient descent along with the weights of the model. Let α = (α1, . . . , αN ) be the vector of the sample weights and denote the corresponding weighted dataset by Dα. The training loss on Dα is then defined as:
LD(w;Dα) = N∑ i=1 αi `(fw(xi), yi). (1)
Note that if all αi’s are either 0 or 1, we are effectively selecting only a subset of D for training. As we will show, this continuous generalization allows us to optimize the sample selection in a differentiable way. In principle, we would like to find the sample weights α∗ = argminα LDtest(wα) that lead to the best generalization. Since we do not have access to the test data, in practice this translates to optimizing α with respect to an (unweighted) validation loss Lval:
α∗ = argminα Lval(wα).
We can, of course, compute a validation loss using a separate validation set. However, as we will see in Section 3.3, we can also use a leave-one-out cross-validation loss directly on the training set, without any requirement of a separate validation set.
In order to efficiently optimize α by gradient-descent, we need to compute the dataset derivative ∇αLval(wα). By the chain rule, this can be done by computing ∇αwα. However, the training function α→ wα that finds the optimal weights wα of the model given the sample weights α may be non-trivial to differentiate or may not be differentiable at all (for example, it may consist of thousands of steps of SGD). This would prevent us from minimizing α end-to-end.
In the next section, we show that if, instead of linearizing the wα end-to-end in order to compute the derivative, we linearize the model before the optimization step, the derivative can both be written in closed-form and computed efficiently, thus giving us a tractable way to optimize α.

heading: 3.1 LINEARIZATION
text: In real-world applications, the parametric model fw(x) is usually a deep neural network. Recent work (Achille et al., 2020; Mu et al., 2020) have shown that in many cases, a deep neural network can
be transformed to an equivalent linear model that can be trained on a simple quadratic loss and still reach a performance similar to the original model. Given a model fw(x), let w0 denote an initial set of weights, For example, w0 could be obtained by pre-training on a large dataset such as ImageNet (if the task is image classification). Following Achille et al. (2020); Mu et al. (2020), we consider a linearization f linw (x) of the network fw(x) given by the first-order Taylor expansion of fw(x) around w0: f linw (x) = fw0(x) +∇wfw0(x) · (w −w0). (2) Intuitively, if fine-tuning does not move the weights much from the initial pre-trained weights w0, then f linw (x) will remain a good approximation of the network while becoming linear in w (but still remaining highly non-linear with respect to the input x). Effectively, this is equivalent to training a linear classifier using the gradients zi := ∇wfw0(xi) as features (Mu et al., 2020). Although f linw (x) is a linear model, the optimal weights wα may still be a complex function of the training data, depending on the loss function used. Achille et al. (2020) showed that equivalent performance can be obtained by replacing the empirical cross-entropy with the regularized leastsquares loss:
LD(w) = N∑ i=1 ‖f linw (x)− yi‖2 + λ‖w‖2 (3)
where y denotes the one-hot encoding vector of the label yi. In Achille et al. (2020), it is shown that linearized models are equivalent from the standpoint of performance on most standard tasks and classification benchmarks, and better in the low-data regime, which is where the problem of “dataset augmentation” is most relevant. The advantage of using this loss is that the optimal weights w∗ can now be written in closed-form as
w∗ = (Z>Z+ λI)−1Z>(Y − fw0(X)), (4) where Z = [z1, . . . , zN ] is the matrix of the Jacobians zi = ∇wfw0(xi). While our method can be applied with no changes to linearization of the full network, for simplicity in our experiments we restrict to linearizing only the last layer of the network. This is equivalent to using the network as a fixed feature extractor and training a linear classifier on top the last-layer features, that is, zi = f L−1 w0 (xi) are the features at the penultimate layer.

heading: 3.2 COMPUTATION OF THE DATASET DERIVATIVE
text: We now show that for linearized models we can compute the derivative∇αwα in closed-form. For the α-weighted dataset, the objective in eq. (3) with L2 loss for the linearized model is written as,
wα = argmin w LD(w;Dα) = argmin w N∑ i=1 α‖w>zi − yi‖2 + λ‖w‖2. (5)
where zi = ∇wfw0(xi) as in the previous section. Note that α‖w>zi − yi‖2 = ‖w>zαi − yαi ‖, where zαi := √ αzi and yαi := √ αyi. Using this, we can reuse eq. (4) to obtain the following closed-form solution for wα:
wα = (Z >DαZ+ λI) −1Z>DαY, (6)
where we have taken Dα = diag(α). In particular, note that wα is now a differentiable function of α. The following proposition gives a closed-form expression for the derivative.
Proposition 1 (Model-Dataset Derivative ∇αwα). For the ridge regression problem equation 5 and wα defined as in equation 6, define
Cα = (Z >DαZ+ λI) −1. (7)
Then the Jacobian of wα with respect to α is given by ∇αwα = ZCα ◦ ( (I− ZCαZ>Dα)Y ) , (8)
Where we write A ◦B ∈ Rn×m×k for the batch-wise outer product of A ∈ Rn×m and B ∈ Rn×k along the common dimension k, i.e., (A ◦B)ijk = aijbik
The Jacobian∇αwα would be rather large to compute explicitly. Fortunately, the end-to-end gradient Lval(wα) if the final validation loss can still be computed efficiently, as we now show. Given a validation dataset Dval, the validation loss is:
Lval(wα) = ∑
(xi,yi)∈Dval `(fwα(xi), yi). (9)
The following gives the expression to optimize α end-to-end with respect to the validation loss. Proposition 2 (Validation Loss Dataset Derivative). Define L as the loss function derivative with respect to the network outputs as,
L =
[ ∂`
∂f (f(x1), y1), · · ·
∂` ∂f (f(xN ), yN ) ] Then the dataset derivative importance weights with respect to final validation is given by
∇αLval(wα) = ZCαZ> × ( L>Y>(I−DαZCαZ>) ) . (10)

heading: 3.3 LEAVE-ONE-OUT OPTIMIZATION
text: It is common in AutoML to optimize the hyper-parameters with respect to a separate validation set. However, using a separate validation may not be practical in limited data settings, which are a main focus of dataset optimization. To remedy this, we now show that we can instead optimize α by minimizing a leave-one-out cross-validation loss that only requires a training set:
LLOO(α) = ∑N i=1 `(fw−iα (xi), yi), (11)
where w−iα are the optimal weights obtained by training with the loss eq. (5) on the entire dataset D except for the i-th sample (xi, yi). This may seem counter-intuitive, since we are optimizing the weights of the training samples using a validation loss defined on the training set itself. It is useful to recall that w−iα minimizes the α-weighted L2 loss on the training set (minus the i-th example):
w−iα = argmin w L−iD (w,Dα) = argminw ∑ j 6=i αj‖fw(xj)− yj‖2 + λ‖w‖2. (12)
Meanwhile, α minimizes the unweighted validation loss in eq. (11). This prevents the existence of degenerate solutions for α.
Computing LLOO naively would require training n classifiers, but fortunately, in the case of a linear classifier with the L2 loss, a more efficient closed-form solution exists (Green & Silverman, 1993; Rifkin & Lippert, 2007). Generalizing those results to the case of a weighted loss, we are able to derive the following expression. Proposition 3. Define
Rα = Z > √ Dα(Z >DαZ+ λI) −1 √ DαZ
Then α-weighted LOOV predictions defined in eq. (12) admit a closed-form solution:
fw−iα (zi) =
[ Rα √ DαY − diag(Rα) √ DαY
diag( √ Dα − √ DαRα)
] i , . (13)
where diag(A) = [a11, . . . , ann] denotes the vector containing the diagonal of A, and the division between vectors is element-wise.
Note that the prediction fw−iα (zi) on the i-th sample when training on all the other samples is a differentiable function of α. Composing eq. (13) in eq. (11), we compute the derivative ∇αLLOO(α), which allows us to optimize the cross-validation loss with respect to the sample weights, without the need of a separate validation set. We give the closed-form expression for ∇αLLOO(α) in the Appendix.

heading: 3.4 DATASET OPTIMIZATION WITH DIVA
text: We can now apply the closed-form expressions for ∇αLval(α) and ∇αLLOO(α) for differentiable dataset optimization. We describe the optimization using Lval, but the same applies to LLOO.
DIVA Reweight. The basic task consists in reweighting the samples of an existing dataset in order to improve generalization. This can curate a dataset by reducing the influence of outliers or wrong labels, or by reducing possible imbalances. To optimize the dataset weights, we use gradient descent in the form: α← α− η∇αLval. (14) It is important to notice that Lval is an unbiased estimator of the test loss only at the first step, hence optimizing using eq. (14) for multiple steps can lead to over-fitting (see Appendix). Therefore, we apply only 1-3 gradient optimization steps with a relatively large learning rate η ' 0.1. This early stopping both regularizes the solution and decreases the wall-clock time required by the method. We initialize α so that αi = 1 for all samples.
DIVA Extend. The dataset gradient also allows us to extend an existing dataset. Given a core dataset D = {(xi, yi)}Ni=1 and an external (potentially noisy) data pool E = {(x̂i, ŷi)}N+Mi=N+1, we want to find the best samples from E to add to D. For this we merge D and E in a single dataset and initialize α such that αi = 1 for samples of D and αi = 0 for samples of E (so that initially the weighted dataset matches D). We then compute ∇αLval(α) to find the top k samples of E that have the largest negative value of ∇αLval(α)i, i.e., the samples that would give the largest reduction in validation error if added to the training set and add them to D. This is repeated until the remaining samples in E all have positive value for the derivative (adding them would not further improve the performance).
Detrimental sample detection. The i-th component of ∇αLval specifies the influence of the i-th sample on the validation loss. In particular, (∇αLval)i > 0 implies that the sample increases the validation loss, hence it is detrimental (e.g., it is mislabeled or overly represented in the dataset). We can select the set of detrimental examples by thresholding∇αLval:
Detrimental( ) = { i : (∇αLval)i ≥ } . (15)

heading: 4 RESULTS
text: For our models we use standard residual architectures (ResNet) models pre-trained on ImageNet (Deng et al., 2009) and Places365 (Zhou et al., 2017). For our experiments on dataset optimization we consider datasets that are smaller than the large scale datasets used for pre-training as we believe they reflect more realistic conditions for dataset optimization. For our experiments we use the CUB-200 (Welinder et al., 2010), FGVC-Aircraft, (Maji et al., 2013), Stanford Cars (Krause et al., 2013), Caltech-256 (Griffin et al., 2007), Oxford Flowers 102 (Nilsback & Zisserman, 2008), MIT-67 Indoor (Quattoni & Torralba, 2009), Street View House Number (Netzer et al., 2011), and the Oxford Pets (Parkhi et al., 2012) visual recognition and classification datasets. In all experiments, we use the network as a fixed feature extractor, and train a linear classifier on top of the network features using the weighted L2 loss eq. (5) and optimize the α weights using DIVA.
Dataset AutoCuration. We use DIVA Reweight to optimize the importance weights of samples from several fine-grain classification datasets. While the datasets have already been manually curated by experts to exclude out-of-distribution or mislabeled examples, we still observe that in all cases DIVA can further improve the test error of the model (Table 1). To understand how DIVA achieves this, in Figure 2 we show the most up-weighted (left) and down-weighted (right) examples on the FGVC Aircraft classification task Maji et al. (2013). We observe that DIVA tends to give more weight to clear, canonical examples, while it detects as detrimental (and hence down-weights) examples that contain multiple planes (making the label uncertain), or that do not clearly show the plane, or show non-canonical poses. We compare DIVA Reweight with two other re-weighting approaches:
Ren et al. (2018), that applies re-weighting using information extracted from a separate validation gradient step, and Chang et al. (2017), which reweighs based on the uncertainty of each prediction (threshold-closeness weighting scheme). For Ren et al. (2018), we set aside 20% of the training samples as validation for the reweight step, but use all samples for the final training (in parentheses). We notice that both baselines, which are designed to reweight noisy datasets, underperform with respect to DIVA on noiseless datasets.
Dataset extension. We test the capabilities of DIVA Extend to extend a dataset with additional
samples of the distribution. In Figure 4 and Table 2 (in the Appendix), we observe that DIVA is
able to select the most useful examples and reaches an optimal performance generalization error
using significantly less samples than the baseline uniform selection. Moreover, we notice that DIVA
identifies a smaller subset of samples that provides better test accuracy than adding all the pool samples to the training set.
Detrimental sample detection. To test the ability of DIVA to detect training samples that are detrimental for generalization, we artificially introduce wrong labels in the dataset. In Section 3.4 we claimed that detrimental examples can be detected by looking at the samples for which the derivative ∇αLLOO(α)i is positive. To verify this, in Figure 3 we plot the histogram of the derivatives for correct and mislabeled examples. We observe that most mislabeled examples have positive derivative. In particular, we can directly classify an example as mislabeled if the derivative is positive. In Figure 3 we report the F1 score and AUC obtained in a mislabeled sample detection task using the DIVA gradients.
Multi-modal learning. Recent multi-modal models such as CLIP (Radford et al., 2021) can embed both text and images in the same vector spaces. This allows to boost the performance on few-shot image classification tasks by also adding to the training set textual descriptions of the classes, such as the label name. However, training on label names may also hurt the performance, for example if the label name is not known by the CLIP model. To test this, we create a few-shot task by selecting 20 images per class from the Caltech-256. We then use DIVA Extend to select an increasing number of
Dataset DIVA Extend Rand. Extend Improvement
Aircrafts 58.00 60.01 +2.01 Cub-200 39.42 42.29 +2.87 MIT Indoor-67 32.54 33.73 +1.19 Oxford Flowers 20.56 23.29 +2.73
Stanford Cars 60.37 62.45 +2.09 Caltech-256 21.97 24.55 +2.59
examples, while it detects as harmful (and hence down-weights) examples that contain multiple268 planes (making the label uncertain), or that do not clearly show the plane, or show non-c nonical269 poses.270
Harmful sample detection. To test the ability of DIVA to detect training samples that are harmful271 for generalization, we artificially introduce wrong labels in the dataset. In Section 3.4 claimed that272 harmful examples can be detected by looking at the samples for which the derivative r↵LLOO(↵)i is273 positiv . To verify this, in Figure 3 w plot the histogram of the derivatives for corr ct and mislabeled274 examples. We observe that indeed most mislabeled examples have positive derivative. In particular,275 we can classify an example as mislabeled if the derivative is positive. In Section 4 we report the F1276 score and AUC obtained in a mislabeled sample detection task using the DIVA gradients.277
Dataset extension. We test the capabilities of DIVA Extend to extend a dataset by adding samples278 from an external source. We consider two settings: one in which the external source is noiseless279 and one in which artificial label noise is introduced, so that randomly adding the external samples280 may be harmful. In the first case, while selecting random examples to add is a viable strategy to281 i prove generalization, we observe that DIVA is able to select the most useful example and reaches282 an optimal performance generalization error using significantly less samples than random selection.283 In the second case, DIVA is able to avoid noisy samples, and significantly outperforms random284 selection, which may in fact decrease the accuracy of the model.285
Figure 3: Distribution of LOO DIVA gradients for correctly labelled and mislabelled samples in Cub-200 dataset (20% mislabelled).
Dataset F1-score (✏ = 0) AUC
Cub200 0.87 0.98 Aircrafts 0.68 0.90 MIT Indoor-67 0.86 0.98 Stanford Cars 0.75 0.93 Caltech-256 0.92 0.99 Oxford Flowers 0.83 0.97
286
Data augmentation. In Section 3.4 we suggested DIVA can be used to find the best probability with287
which to apply a data augmentation procedure in order to improve the generalization performance.288
To test this, we select common data augmentation procedures, horizontal flip and vertical flip, and289
we tune their probability on the Street View House Number [46], Oxford Flowers [48] and the290 Oxford Pets [49] classification tasks. We observe that DIVA assigns different probabilities to each291 transformation depending on the task (Figure 5): on the number classification task DIVA penalizes292 both vertical and horizontal flips, which may confuse different classes (such 2 and 5, 6 and 9). On293 an animal classification task (Oxford Pets) DIVA does not penalize horizontal flips, but penalizes294 vertical flips since they are out of distributions. Finally, on Flowers classification DIVA gives equal295 probability to all transformations (most flower pictures are frontal so all rotations and flips are valid).296
Figure 3: (Left) Distribution of LOO DIVA gradients for correctly labelled and mislabelled samples in CUB-200 dataset (20% of the samples are mislabeled by replacing their label uniformly at random). (Right) DIVA for outlier rejection. We use DIVA on a ResNet-34 network linearization and detect mislabelled sa ples (outliers) in a dataset present with 20% label noise. Selection is based on ∇α(Lval(wα))i > .
8
labels to add to the training set. In Figure 5 (right), we show that DIVA can select the beneficial label embeddings to add in order to improve the few-shot test performance. However, when forced to add all labels, including detrimental ones, the test error starts to increase again.
Data augmentation. To further test the versatility of DIVA, we qualitatively evaluate DIVA Reweight on the task of tuning the probabilities with which we apply a given data augmentation procedure. Let t1, . . . , tK be a set of data augmentation transformations. Let Dtk be the result of applying the data augmentation tk to D. We can create an augmented dataset Daug = D ∪ Dt0 ∪ . . . ∪ DtK , by merging all transformed datasets. We then apply DIVA Reweight on Daug to optimize the weight α of the samples. Based on the updated importance weights we estimate the optimal probability with which to apply the transformation tk as pk = ( ∑ i∈Dtk αi)/( ∑ i αi). In particular we select common data augmentation procedures, horizontal flip and vertical flip, and we tune their probability on the Street View House Number, Oxford Flowers and the Oxford Pets classification tasks. We observe that DIVA assigns different probabilities to each transformation depending on the task (Figure 5): on the number classification task DIVA penalizes both vertical and horizontal flips, which may confuse different classes (such 2 and 5, 6 and 9). On an animal classification task (Oxford Pets) DIVA does not penalize horizontal flips, but penalizes vertical flips since they are out of distributions. Finally, on Flowers classification DIVA gives equal probability to all transformations (most flower pictures are frontal so all rotations and flips are valid).

heading: 5 DISCUSSION
text: In this work we present a gradient-based method to optimize a dataset. In particular we focus on sample reweighting, extending datasets, and removing outliers from noisy datasets. We note that by
developing the notion of a dataset derivative we are capable of improving dataset quality in multiple
disparate problems in machine learning. The dataset derivative we present is given in closed-form
and enables general reweighting operations on datasets based on desired differentiable validation
losses. In cases where a set-aside validation loss is not available we show the use of the leave-one-out framework enables computing and optimizing a dataset “for free” and derive the first closed-form dataset derivative based on the LOO framework.
Dataset DIVA Extend Rand. Extend Improvement
Aircrafts 58.00 60.01 +2.01 Cub-200 39.42 42.29 +2.87
MIT Indoor-67 32.54 33.73 +1.19 Oxford Flowers 20.56 23.29 +2.73
Stanford Cars 60.37 62.45 +2.09 Caltech-256 21.97 24.55 +2.59
examples, while it detects as harmful (and hence down-weights) examples that contain multiple268 planes (making the label uncertain), or that do not clearly show the plane, or show non-canonical269 poses.270
Harmful sample detection. To test the ability of DIVA to detect training samples that are harmful271 for generalization, we artificially introduce wrong labels in the dataset. In Section 3.4 claimed that272 harmful examples can be detected by looking at the samples for which the derivative r↵LLOO(↵)i is273 positive. To verify this, in Figure 3 we plot the histogram of the derivatives for correct and mislabeled274 examples. We observe that indeed most mislabeled examples have positive derivative. In particular,275 we can classify an example as mislabeled if the derivative is positive. In Section 4 we report the F1276
score and AUC obtained in a mislabeled sample detection task using the DIVA gradients.277
Dataset extension. We test the capabilities of DIVA Extend to extend a dataset by adding samples278
from an external source. We consider two settings: one in which the external source is noiseless279
and one in which artificial label noise is introduced, so that randomly adding the external samples280
Figure 3: Distribution of LOO DIVA gradients for correctly labelled and mislabelled samples in Cub-200 dataset (20% mislabelled).
Table 2: DIVA for outlier rejection. We use DIVA on a ResNet-34 network linearization and detect mislabelled samples (outliers) in a dataset present with 20% label noise. Selection is based on r↵(Lval(w↵))i > ✏.
Data augmentation. In Section 3.4 we suggested DIVA can be used to find the best probability with287 which to apply a data augmentation procedure in order to improve the generalization performance.288 To test this, we select common data augmentation procedures, horizontal flip and vertical flip, and289 we tune their probability on the Street View House Number [46], Oxford Flowers [48] and the290 Oxford Pets [49] classification tasks. We observe that DIVA assigns different probabilities to each291 transformation depending on the task (Figure 5): on the number classification task DIVA penalizes292 both vertical and horizontal flips, which may confuse different classes (such 2 and 5, 6 and 9). On293 an animal classification task (Oxford Pets) DIVA does not penalize horizontal flips, but penalizes294 vertical flips since they are out of distributions. Finally, on Flowers classification DIVA gives equal295 probability to all transformations (most flower pictures are frontal so all rotations and flips are valid).296

heading: APPENDIX
text: We structure the appendix as follows: We present additional experiments in Section A and we describe the details of the experiments in Section B. In Section C we provide proofs for the propositions in the paper and additional discussion on the methods.

heading: A ADDITIONAL EXPERIMENTS
text: Validation overfitting. When updating the dataset using the dataset derivative there is a risk of overfitting to the validation set after repeated applications of the derivative. Namely the validation loss is initially an unbiased estimate of the test loss yet after using it to update the dataset repeatedly it eventually will start overfitting. In our settings, we notice that when using a small number of gradient updates (< 5) and with step sizes η ∼ 0.1 we are able to avoid overfitting and improve the test error when optimizing the dataset derivative based on the validation loss. In this experiment we present the final test and validation classification errors of optimized datasets. As we optimize with respect to the validation loss, it is indeed clear that the validation loss decreases dramatically yet more importantly are the effect of the test accuracy.
DIVA Extend plots. In Figure 4 we report the results on all the remaining datasets following the set-up of Figure 4. In Table 2 we report the accuracy of DIVA Extend and uniform sampling on the various datasets when adding 50% of samples from the pool.

heading: B EXPERIMENTAL DETAILS
text: Dataset details For our experiments we utilize several fine-grain classification datasets from the computer vision community that are standard for fine-tuning image classification tasks (CUB-200 Welinder et al. (2010), FGVC-Aircraft, Maji et al. (2013), Stanford Cars Krause et al. (2013), Caltech256 Griffin et al. (2007), Oxford Flowers 102 Nilsback & Zisserman (2008), MIT-67 Indoor Quattoni & Torralba (2009), Street View House Number Netzer et al. (2011), and the Oxford Pets Parkhi et al. (2012)). Some of the datasets do not follow a default train-test split and we use the following splits commonly used in the literature for the datasets,
• Oxford Flowers 102: We use the original 1020 images in the training split without incorporating the default validation set.
• Caltech-256: We split the dataset into a training set with 60 images from each class for training, and use remaining data for testing.
Pre-training setup For the pre-trained networks we use for fine-tuning, we use the pre-trained configurations available on PyTorch’s torchvision package. In particular the reference networks are pre-trained using the ImageNet Deng et al. (2009) dataset. The images embedded by the network are pre-processed via standard resizing and center cropping (256 resize, followed by a 224 cropping).
Regularization parameter λ: To get the best unweighted dataset baseline to compare with the optimized dataset, for each of the un-optimized original datasets, we first search for optimal λ values in λ ∈ {2n for n ∈ {−20,−19, . . . 4}} to measure the classifier’s performance. After selecting optimal λ we proceed with dataset optimization with the optimal λ values. Note that DIVA does not require λ to be optimal and improvements are even more significant for un-optimized λ.
Dataset Derivative Computation In Section C we derive the closed-form dataset derivatives used for DIVA. We computed the closed-form solution analytically and we verified our results using automatic differentiation tools on large number of conditions including synthetic and real data. This has been another method to verify the correctness of the derivative formulas.

heading: DIVA experimental details
text: DIVA Extend For Table 2 and Figure 4 we first split the original training set into 50% training subset and 50% pool subset that will be used to selectively extend the training subset with DIVA or other extension approaches. We run DIVA Extend LOO and uniform sampling to add pool samples to the training set. In both settings we incrementally extend the training subset from the pool (For DIVA, by selecting samples with top DIVA score) in each step we extend an equal number of samples ((# pool samples) // (# number of steps)). In the figure we present the test error as a function of training set size and compare DIVA sample selection with selecting the same number of samples at uniform from the pool set. In Table 2, we present the improvement in test accuracy at the 50% extend mark
of the pool set (e.g. extending 25% of the original training set) between DIVA and uniform sample selection of the same number of samples. For both experiments we use the ResNet-34 architecture.
DIVA Reweigh In Table 1 we use DIVA Reweight LOO with the same parameters for all of the datasets we consider. The DIVA parameters we use are K = 4 for the number of steps and η = 0.15 for the step-size. As with DIVA Extend, we use ResNet 34 for the architecture for the representation.
DIVA validation loss We find the cross entropy loss to work better as the loss function applied to the validation predictions (both in LOO and regular validation). Further for LOO we find it crucial to apply the validation loss only on mis-classified LOO predictions to improve the test accuracy of the model, this can be interpreted as the cross entropy loss with a “hard margin hinge” loss.
Outlier rejection For the results presented in the side-by-side Table and Figure 3 we apply 20% random label noise to each class in the dataset and use DIVA LOO to compute the normalized DIVA gradient (DIVA score) of each sample. The F1 score reports classification by thresholding with = 0 and the AUC is computing by thresholds spanning detection of no samples, to detection of all samples.

heading: C PROOFS OF PROPOSITIONS
text: We write the proposition statements for convenience.
Proposition (Model-Dataset Derivative ∇αwα). For the ridge regression problem eq. (5) and wα defined as in eq. (6), define
Cα = (Z >DαZ+ λI) −1. (16)
Then the Jacobian of wα with respect to α is given by ∇αwα = ZCα ◦ ( (I− ZCαZ>Dα)Y ) . (17)

heading: Proof of Proposition 1:
text: We recall the settings of the problem are Z ∈ Rn×m, Y ∈ Rn×k, wα ∈ Rm×k and the importance weights α ∈ Rn. Here n is the number of samples, m is the number of parameters for each output of the model, and k is the number of classes (represented via the one-hot convention).
The derivation of∇αwα we present is computational in nature and without the loss of generality we consider the derivative for a single output class k = 1 (one-vs-all classification naturally extends). For the single class settings, the relevant dimensions are wα ∈ Rm and ∇αwα ∈ Rm×n (for numerator layout convention of the derivative). To further simplify we consider the derivative entrywise for single index αr,
∂wα ∂αr ∈ Rm×1.
The closed-form solution is, wα = CαZ >DαY.
using chain rule, ∂wα ∂αr = ∂Cα ∂αr × Z>DαY +Cα × ∂Z>DαY ∂αr .
We compute the two parts of the derivative in turn: Let Kα = (Z>DαZ+ λI) so that Kα = C−1α . Then
∂Kα ∂αr = ∂ ∂αr (Z>DαZ+ λI) = Z > r,:Zr,:
where Zr,: ∈ R1×m is the rth row of Z. Next note, CαKα = I
differentiating both sides, ∂CαKα ∂αr = 0.
Applying chain rule we get,
∂CαKα ∂αr = ∂Cα ∂αr Kα +Cα ∂Kα ∂αr ,
rearranging, substituting ∂Kα∂αr , and multiplying to the right by Cα
∂Cα ∂αr = −CαZ>r,:Zr,:Cα.
Next, by direct computation ∂Z >DαY ∂αr satisfies
∂Z>DαY
∂αr = yr · Z>r,:.
Combining the original terms we have
∂wα ∂αr
= − ( CαZ > r,:Zr,:Cα ) × Z>DαY + yr ·Cα × Z>r,:
= CαZ > r,: ( yr − Zr,:CαZ>DαY ) .
Now the full derivative is written as ∇αwα = ZCα ◦ ( (I− ZCαZ>Dα)Y ) Proposition (Validation Loss Dataset Derivative). Define L as the matrix of the loss function derivative with respect to network training outputs as,
L =
[ ∂`
∂f (fwα(x1), y1), · · ·
∂` ∂f (f(xN ), yN )
] .
Then the dataset derivative of the importance weights with respect to final validation loss is given by ∇αLval(wα) = ZCαZ> × ( L>Y>(I−DαZCαZ>) ) . (18)

heading: Proof of Proposition 2:
text: This follows from the chain rule combined with simplification of broadcasting terms. We again consider the single output settings with the single coordinate derivative ∂Lval∂αr which is given as,
∂Lval ∂αr = ∇wLval ∂wα ∂αr .
With
= n∑ i=1 L:,i × z>i = L>Z.
Therefore ( ∇αLval(wα) ) i = ∑ j,k ( ∇αwα ) i,j,k (L>Z)j,k. Now L,Z can be separated into the two terms of∇αwα,
∇αLval(wα) = ZCαZ> × ( L>Y>(I−DαZCαZ>) ) (19)
Next we consider the derivations for the Leave One Out (LOO) framework. In the LOO framework one applies cross-validation to a training set {(z1, y1), . . . (zn, yn)} by running n-fold cross validation, where in each fold, the ith sample (zi, yi) is taken out and is used for validation while the optimal classifier is solved for the remaining of the training task,
w−i,= argmin w ∑ j 6=i `(f(zj), yj). (20)
Then the LOO prediction at the ith index is defined as (fLOO)i = fw−i(zi). Below we prove the LOO predictions can be written in closed-form without explicit cross validation calculations.
Proposition 4 (Closed-form LOO prediction vector). Define the LOO vector predictions as,
fLOO = [fw−1(z1), . . . , fw−n(zn)] >
and define R = Z>(Z>Z+ λI)−1Z
then for the learning task eq. (20) LOO predictions are given as
fLOO = Ry − diag(R)y I− diag(R) . (21)

heading: Proof of Proposition 4:
text: The proof is reproduced from Rifkin & Lippert (2007) for completeness. Without the loss of generality we derive the LOO prediction of zn. Namely given, {(z1, y1), . . . (zn, yn)} we use {(z1, y1), . . . (zn1 , yn−1)} for training and validate using {(zn, yn}. Denote w−n as be the optimal solution to this training task with regularization parameter λ and define the (currently unknown) LOO prediction as
q = fw−n(zn).
We define the modified learning task consisting of {(z1, y1), . . . (zn−1, yn−1), (zn, q)} where we added the data point (zn, q). Note that the optimal solution with λ regularization to the modified learning task is again w−n since q is taken to have a zero residual. therefore the solution to the modified learning problem can be written in closed-form as,
w−n = (Z>Z+ λI)−1Z>[y1:n−1, q] >.
Using w−n we write the equation for the LOOV prediction q,
q = 〈zn,w−i〉 = z>n (Z>Z+ λI)−1Z>[y1:n−1, q]>. Let R = Z(Z>Z+ λI)−1Z> then we have
q = z>nw−i = Rn,:[Y1:n−1, q] >
and by adding and subtracting R:,n multiplied by [0, yn]> we get,
q = z>nw−i = R:n−1,n[Y1:n−1] > +R:,n[0, q] +R:,n[0, yn] > −R:,n[0, yn]>. Re-arranging we have q −Rnnq = R:,ny − ynRn,n Solving for q we get,
q = R:,ny − ynRn,n
1−Rn,n And without the loss of generality the full prediction vector is given as,
fLOO = Ry − diag(R)y
diag(I−R) (22)
Given the closed-form LOOV expression we may use fLOO for the validation loss to compute the dataset derivative on Lval without any additional validation data. While this may seem contradictory as we are optimizing the dataset validated via the weighting duality between sample loss weighting and data scaling, we define the leave one out value predictions in the weighted dataset settings and evaluate on the original (unweighted) data points as,
fw−iα (zi) (23)
with, w−iα = argmin
w ∑ j 6=i αj`(f(zj), yj). (24)
Therefore the α-weighted LOO term is fw−iα (zi) is faithful to the original distribution despite being trained with the weighted loss eq. (24). We in fact are able to show that α-weighted LOO formulation also admits a closed-form solution that satisfies our definition and for DIVA LOO we utilize the derivative of the closed-form to optimize the dataset.
Proposition. Define Rα = Z
> √
Dα(Z >DαZ+ λI) −1 √ DαZ
Then the α-weighted LOOV predictions defined in eq. (12) admit a closed-form solution:
fw−iα (zi) =
[ Rα √ DαY − diag(Rα) √ DαY
diag( √ Dα − √ DαRα)
] i , (25)
Further the LOO dataset derivative is well-defined and satisfies the following gradient condition,
diag(∇αfLOO) = 0. (26)
The gradient condition diag(∇αfLOO) = 0 implies that the LOO prediction at zi does not depend on αi and ensures that the closed-form solution is well defined in the α-weighted settings and is differentiable.

heading: Proof of Proposition 3:
text: The proof builds on Proposition 4 for the weighted settings.
In general since LOO expression describes the weighting problem, it must be shown that the introduction of the weights do not break the argument of Proposition 4. Considering the optimization problem in eq. (12). We also use the duality between the loss weights and data scaling to note that w−α i can be derived by considering the unweighted LOO eq. (20) with the modified dataset,
{(√α1z1, √ α1y1), . . . ( √ αnzn, √ αnyn)}. (27)
Indeed in this settings with the newly defined data the derivation in Proposition 4 of the final prediction vector of the LOO entries holds with the same optimal solution wα due to the duality between data scaling and loss weights. Nonetheless the closed-form LOO predictions fLOO are evaluated at the data-points √ αizi. Since the model is linear the final predictions at the original data-points of the
weighted training settings are written as fLOO/ √ α. Noting the weighted training problem can be expressed as Ỹ = √ DαY, Z̃ = √ DαZ we use Proposition 4 to write analogously
Rα = Z > √ Dα(Z >DαZ+ λI) −1 √ DαZ (28)
and divide by √ α by multiplying the denominator by √ Dα,
fw−iα (zi) =
[ Rα √ DαY − diag(Rα) √ DαY
diag( √ Dα − √ DαRα)
] i , (29)
Since the derivation at the ith index of the weighted LOO prediction, (fLOO)i = (w−iα ) >zi is entirely independent of αi, we have ∂(fw−iα (zi))i
∂αi = 0.

Review summaries:
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).

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.  

The paper introduces an approach to calculate the derivative of data samples' influence on the validation error. For this, the authors use a recent model linearization method and leverage closed-form leave-one-out loss calculation for linear models. Besides the theoretical derivation, they present consistent improvements in dataset extension, re-weighting, outlier rejection and automatic aggregation of multi-modal data.