Deep Learning on a Data Diet: Finding Important Examples Early in Training

1 Introduction . Recently , deep learning has made remarkable progress driven , in part , by training overparameterized models on ever larger datasets . This trend creates new challenges : the large computational resources required pose a roadblock to the democratization of AI . Memory and resource constrained settings , such as on-device computing , require smaller models and datasets . Identifying important training data plays a role in online and active learning . Finally , it is of theoretical interest to understand how individual examples and sub-populations of training examples influence learning . To address these challenges , we propose a scoring method that can be used to identify important and difficult examples early in training , and prune the training dataset without large sacrifices in test accuracy . We also investigate how different sub-populations of the training data identified by our score affect the loss surface and training dynamics of the model . ⇤This work was carried out while the author was at ServiceNow . It was finalized at Google Brain . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Recent work on pruning data [ 1 , 2 ] , can be placed in the broader context of identifying coresets— examples that provably guarantee a small gap in training error on the full dataset [ 3–7 ] . However , due to the nonconvex nature of deep learning , coreset techniques make conservative estimates that lead to weak theoretical guarantees and are less effective in practice . A different approach recently proposed by Toneva et al . [ 8 ] tracks the number of times through training an example transitions from being correctly classified to misclassified , called a `` forgetting event '' , and find that some examples are rarely forgotten , while others are forgotten repeatedly . Empirically , they observed that training accuracy is not affected by the rarely forgotten training examples and a large fraction of the training data can be removed without any impact on test accuracy . However , since this method relies on collecting forgetting statistics throughout training , the forgetting score is typically calculated in the middle of or at the end of training . Toneva et al . [ 8 ] find that , in their example of a ResNet18 trained on CIFAR-10 for 200 epochs , the Spearman rank correlation between early and late scores is good after about 25 epochs and stabilizes after 75 epochs . Broadly speaking , the ability to prune datasets raises a number of questions : What is the nature of examples that can be removed from the training data without hurting accuracy ? How early in training can we recognize such examples ? How many examples do we need and how does this depend on the data distribution ? These questions may have no generic answers and so , in this work , we begin to pursue them empirically in the context of several standard vision benchmarks and standard network architectures . Answers to these questions may both ( 1 ) lead to new methodologies that could dramatically reduce training times and memory requirements , and ( 2 ) offer important insights into the training dynamics of deep neural networks , and the role of data . Our first finding is that very early in training ( just a few epochs ) , partial forgetting scores identify large fractions of data that can be pruned . Analyzing this puzzling result with a one gradient step analysis of training suggests a very simple heuristic : use the loss gradient norm of individual examples to identify important examples . While this approach does not work when the loss gradient norms are computed at the weights early in training of a single trajectory , we find that , surprisingly , averaging these norms over multiple weight initializations does produce a ranking that correlates strongly with forgetting scores and allows us to prune a significant fraction of examples early in training . Indeed , we can prune 50 % of examples from CIFAR-10 without affecting accuracy , while on the more challenging CIFAR-100 dataset , we can prune 25 % of examples with only a 1 % drop in accuracy . Through a series of empirical studies , we have begun to tease apart the properties of important examples and how they can depend on the data distribution . In particular , we find that the examples with the very highest norms become superfluous as the amount of label noise increases . Indeed , even on clean data , we find that in the high pruning regime , the best population excludes the very highest-scoring examples . 1.1 Contributions . • We propose to score the importance of each training example ( xi , yi ) by its expected loss gradient norm ( GraNd score ) , which , up to a constant , bounds the expected change in loss for an arbitrary example ( x , y ) caused by removing ( xi , yi ) . • We show that pruning training samples with small GraNd scores at initialization allows one to train on much smaller subset of the training data without significant loss in accuracy . While the pruning levels are comparable to those provided by other methods [ 1 , 8 ] , our score is the only one that is well-defined at initialization and early in training . • Our experimental findings suggest that , within the first few epochs of training , the GraNd score is well-approximated by the norm of the error vector ( EL2N score ) , where the error vector is the predicted class probabilities minus one-hot label encoding . In fact , we find that the EL2N score provides an even stronger signal for data-pruning—for CIFAR10 we can prune 50 % of the data , and for the harder CIFAR100 , we can prune as much as 25 % of the data without any loss in test accuracy . • We study the role of examples with the highest EL2N scores , and find that excluding a small subset of the very highest scoring examples produces a boost in performance . This boost in performance is enhanced in a corrupted label regime . • We introduce a method , based on linearly connected modes , for studying the empirical risk surface in terms of the modes of subsets of data , allowing us to identify when , in training , the final performance on subpopulations is determined . We demonstrate that the linearly connected mode at-convergence of empirical risk surface computed on low EL2N score examples is determined much earlier in training compared to high score examples . • Finally , we study how an example ’ s EL2N score connects to the network ’ s training dynamics . We do so by tracking the data-dependent NTK submatrices corresponding to the low or high score examples , and measuring the rate at which it evolves in a scale-invariant way . We find that the NTK submatrix for the high score examples evolves faster throughout training , supporting our hypothesis that high-scoring examples are the ones driving the learning and the changes in the NTK feature space [ 9 ] . 2 Which samples are important for learning ? . 2.1 Preliminaries . We consider supervised classification , where S = { ( xi , yi ) } Ni=1 denotes the training set , drawn i.i.d . from an unknown data distribution D , with input vectors x 2 Rd and one-hot vectors y 2 { 0 , 1 } K encoding labels . For a fixed neural network architecture , let fw ( x ) 2 RK be the logit outputs of the neural network with weights w 2 W ✓ RD on input x 2 Rd . Let be the softmax function given by ( z1 , . . . , zK ) k = exp { zk } / PK k0=1 exp { zk0 } . Let p ( w , x ) = ( f ( w , x ) ) denote the neural network output in the form of a probability vector . For any probability vector p̂ , let ` ( p̂ , y ) = PK k=1 y ( k ) log p̂ ( k ) denote cross-entropy loss . Let w0 , w1 , w2 , . . . , wT be the iterates of stochastic gradient descent ( SGD ) , where , for some sequence of minibatches S0 , S1 , . . . , ST 1 ✓ S of size M , we have wt = wt 1 ⌘ P ( x , y ) 2St 1 gt 1 ( x , y ) , ( 1 ) for gt 1 ( x , y ) = rwt 1 ` ( p ( wt 1 , x ) , y ) , and t = 1 , . . . , T . 2.2 Gradient Norm Score and an infinitesimal analysis . Fix a training set S. Due to training with SGD from a random initialization , the weight vector at time t > 0 , wt , is a random variable . The expected magnitude of the loss vector is our primary focus : Definition 2.1 . The GraNd score of a training example ( x , y ) at time t is t ( x , y ) = Ewt kgt ( x , y ) k2 . Here we describe conditions under which the GraNd score controls the contribution of a training example to the change in the training loss . In order to simplify our analysis , we approximate the training dynamics as if they were in continuous time . A key quantity in our analysis is the time derivative of the loss for a generic labeled example ( x , y ) : t ( ( x , y ) , St ) = d ` ( ft ( x ) , y ) dt ( where ft ( · ) = fwt ( · ) ) , i.e. , the instantaneous rate of change in the loss on ( x , y ) at time t , where the gradient is computed on the minibatch St. By the chain rule , t ( ( x , y ) , St ) = gt ( x , y ) dwt dt . ( 2 ) This relates to our discrete time dynamics via dwtdt ⇡ wt+1 wt = ⌘ P ( x0 , y0 ) 2St gt ( x 0 , y0 ) . Our goal is to understand how removing a training point from minibatch St affects t ( ( x⇤ , y⇤ ) , St ) for any ( x⇤ , y⇤ ) . If a training point ( x , y ) is not in the minibatch St , then the effect is trivial . We thus study t ( ( x⇤ , y⇤ ) , S ) in order to be able to rank all the training examples . Lemma 2.2 . Let S¬j = S \ ( xj , yj ) . Then for all ( x⇤ , y⇤ ) , there exists c such that k t ( ( x⇤ , y⇤ ) , S ) t ( ( x⇤ , y⇤ ) , S¬j ) k ckgt ( xj , yj ) k. ( 3 ) Proof . For a given example ( x⇤ , y⇤ ) , the chain rule yields t ( ( x⇤ , y⇤ ) , S ) = d ` ( ft ( x ⇤ ) , y⇤ ) dt = d ` ( ft ( x ⇤ ) , y⇤ ) dwt dwt dt . Since the weights are updated using SGD , we have dwt dt = ⌘ P ( xj , yj ) 2St gt ( xj , yj ) . Letting c = ⌘k d ` ( ft ( x ⇤ ) , y⇤ ) dwt k , the result follows . At any given training step , given the current location wt , the contribution of a training example ( x , y ) 2 to the decrease of loss on any other example , is bounded by Eq . ( 3 ) . Since the constant c does not depend on the training example ( x , y ) 3 , we only consider the gradient norm term , kgt ( x , y ) k. The expected value of this gradient norm is exactly the GraNd score of ( x , y ) . In other words , examples with a small GraNd score in expectation have a bounded influence on learning how to classify the rest of the training data at a given training time4 . We therefore propose to rank training examples by their GraNd scores , larger norm meaning more important for preserving t ( x ) . For an arbitrary input x 2 Rd , let ( k ) t ( x ) = rwtf ( k ) t ( x ) denote the kth logit gradient . Then GraNd can be written as5 t ( x , y ) = E PK k=1 rf ( k ) ` ( ft ( x ) , y ) T ( k ) t ( x ) 2 . ( 4 ) Under the cross entropy loss , rf ( k ) ` ( ft ( x ) , y ) T = p ( wt , x ) ( k ) yk . When { ( k ) t ( x ) } k are roughly orthogonal across logits , and are of a similar size across logits and training examples x , then we can approximate GraNd by just the norm of the error vector . Definition 2.3 . The EL2N score of a training sample ( x , y ) is defined to be Ekp ( wt , x ) yk2 . Our experimental results suggest that this approximation becomes accurate after a few epochs of training ( see Section 3 ) . These approximations are also in agreement with the empirical results reported in [ 9 , 10 ] . Fort and Ganguli [ 10 , Sec . 5.1 ] demonstrate that the mean logit gradients are nearly orthogonal among classes throughout training . The authors demonstrate that per-example gradients cluster around the mean logit gradient . Fort et al . [ 9 , Figs . 12D-14D ] provide evidence that the mean logit gradient directions evolve rapidly early in training and then stabilize ( as measured by the cosine distance between mean logit gradient vectors at different times in training ) .

Summary: The paper proposes two metrics to rank training set samples for reducing the training set size:  1. L2-norm of the error vector (EL2N). 2. L2-norm of the gradient vector (GraNd).  These can be evaluated cheaply and relatively early in training (10-20 epochs) or even at initialization. For instance, by keeping only the top 50% CIFAR10 of training data according to this metric, one can achieve the same test accuracy as when using the whole training set.  The authors evaluate their method on several datasets and architectures, and run several experiments that relate these metrics to change in NTK during training or the onset of linear mode connectivity, and more.