Coresets for Robust Training of Neural Networks against Noisy Labels

The key idea of our method is to carefully select subsets of data points that allow the neural network to effectively learn from the clean training data, but prevent it to overfit noisy labels. Recent result on optimization and generalization of neural networks show that the Jacobian of typical neural networks exhibits an approximately low-rank structure, i.e., a number of singular values are large and the remaining majority of the spectrum consists of small singular values. Consequently, the Jacobian spectrum can be split into information space $\mathcal{I}$, and nuisance space $\mathcal{N}$, associated with the large and small singular values [32]. Formally, for complementary subspaces $\mathcal{S}_{-},\mathcal{S}_{+}\subset\mathbb{R}^{n}$, and for all unit norm vectors ${\bm{{v}}}\in\mathcal{S}_{+}$, ${\bm{{w}}}\in\mathcal{S}_{-}$, and scalars $0\leq\mu\ll\alpha\leq\beta$, we have

There are two key observations [23, 32]: (1) While learning over the low-dimensional information space is fast and generalizes well, learning over the high-dimensional nuisance space is slow and leads to overfitting; and (2) The generalization capability and dynamics of training is dictated by how well the label, $\bm{y}$, and residual vector, $\bm{r}=f(\bm{W},{\bm{X}})-\bm{y}$, are aligned with the information space. If the residual vector is very well aligned with the singular vectors associated with the top singular values of ${\cal{J}}(\bm{W},{\bm{X}})$, the gradient update, $\nabla{\cal{L}}(\bm{W})={\cal{J}}^{T}(\bm{W},{\bm{X}})\bm{r}$, significantly reduces the misfit allowing substantial reduction in the training error. Importantly, the residual and label of most of the clean data points fall on the information space, while the residual and label of noisy data points fall on the nuisance space and impede training and generalization.

To avoid overfitting noisy labels, one can leverage the first observation above, and iteratively selects subsets $S$ of $k$ data points that provide the best rank-$k$ approximation to the Jacobian matrix. In doing so, gradient descent applied to the subsets cannot overfit the noisy labels. Formally:

where ${\cal{J}}(\bm{W},{\bm{X}}_{S})\in\mathbb{R}^{k\times m}$ is the set of $k$ rows of the Jacobian matrix associated to ${\bm{X}}_{S}$, and $P_{S}={\cal{J}}^{T}(\bm{W},{\bm{X}}_{S}){\cal{J}}(\bm{W},{\bm{X}}_{S})$ denotes the projection onto the $k$-dimensional space spanned by the rows of ${\cal{J}}(\bm{W},{\bm{X}}_{S})$. Existing techniques to find the best subset of $k$ rows or columns from an $n\times m$ matrix have a computational complexity of $poly(n,m,k)$ [3, 10, 20], where $n,m$ are the number of data points and parameters in the network. Note that the subset $S^{*}(\bm{W})$ depends on the parameter vector $\bm{W}$ and a new subset should be extracted after every parameter update. Furthermore, calculating the Jacobian matrix requires backpropagation on the entire dataset which could be very expensive for deep networks. Therefore, the computational complexity of the above methods becomes prohibitive for over-parameterized neural networks trained on large datasets. Most importantly, while this approach prohibits overfitting, it does not help identifying the clean data points.

To achieve a good generalization performance, our approach takes advantage of both the above mentioned observations. In particular, our goal is to find representative subsets of $k$ diverse data points with clean labels that span the information space $\mathcal{I}$, and provide an approximately low-rank Jacobian matrix. The important observation is that as nuisance space is very high dimensional, data points with noisy labels spread out in the gradient space. In contrast, information space is low-dimensional and data points with clean labels that have similar gradients cluster closely together. The set of most centrally located clean data points in the gradient space can be found by solving the following $k$-medoids problem:

where $d_{ij}(\bm{W})=\|\nabla{\cal{L}}(\bm{W},\bm{x}_{i})-\nabla{\cal{L}}(\bm{W},\bm{x}_{j})\|_{2}$ is the pairwise dissimilarity between gradients of data points $i$ and $j$. Note that the above formulation does not provide the best rank-$k$ approximation of the Jacobian matrix. However, as the $k$-medoids objective selects a diverse set of clean data points, the minimum singular value of the Jacobian of the selected subset projected over the subspace $\mathcal{S}_{+}$, i.e., $\sigma_{\min}({\cal{J}}(\bm{W},{\bm{X}}_{S}^{*}),\mathcal{S}_{+})$, will be large. Next, we weight the derivative of every medoid $j\in S^{*}$ by the size of its corresponding cluster $r_{j}=\sum_{i\in V}\mathbb{I}[j={\arg\min}_{s\in S^{*}}d_{is}]$ to create the weighted Jacobian matrix ${\cal{J}}_{r}(\bm{W},{\bm{X}}_{S^{*}})=\text{diag}([r_{1},\cdots,r_{k}])J(\bm{W},{\bm{X}}_{S^{*}})\in\mathbb{R}^{k\times m}$. We can establish the following upper and lower bounds on the singular values $\sigma_{i\in[k]}({\cal{J}}_{r}(\bm{W},{\bm{X}}_{S^{*}}),\mathcal{S}_{+})$ of the weighted Jacobian over $\mathcal{S}_{+}$:

where $r_{\min}\!=\!\min_{j\in[k]}r_{j}$ and $r_{\max}\!=\!\max_{j\in[k]}r_{j}$, and we get an error of $\epsilon$ in approximating the largest singular value of the neural network Jacobian, $\epsilon\leq|\sqrt{r_{\max}}\|{\cal{J}}(\bm{W},{{\bm{X}}}_{S^{*}})\|-\|{\cal{J}}(\bm{W},{\bm{X}})\||$. Now, we apply gradient descent updates in Eq. (2) to the weighted Jacobian ${\cal{J}}_{r}(\bm{W},{\bm{X}}_{S^{*}})$ of the $k$ extracted medoids:

Note that we still need backpropagation on the entire dataset to be able to compute pairwise dissimilarities $d_{ij}$. For neural networks, it is shown that the variation of the gradient norms is mostly captured by the gradient of the loss w.r.t. the input to the last layer of the network [19]. This argument can be used to efficiently upper-bound the normed difference between pairwise gradient dissimilarities [30]:

where $\Sigma^{\prime}_{L}({\bm{z}}_{i}^{L})\nabla_{i}^{L}{\cal{L}}$ is gradient of the loss function ${\cal{L}}$ w.r.t. the input to the last layer $L$ for data point $i$, and $c_{2},c_{2}$ are constants. The above upper-bound is marginally more expensive to calculate than the value of the loss since it can be computed in a closed form in terms of $z^{L}$. Hence, $d^{u}_{ij}=\|\Sigma^{\prime}_{L}({\bm{z}}_{i}^{L})\nabla_{i}^{L}{\cal{L}}-\Sigma^{\prime}_{L}({\bm{z}}_{j}^{L})\nabla_{j}^{L}{\cal{L}}\|_{2}$ can be efficiently calculated. We note that although the upper-bounds $d^{u}_{ij}$ have a lower dimensionality than $d_{ij}$, noisy data points still spread out in this lower-dimensional space, and hence are not selected as medoids. This is confirmed by out experiments (Fig. 1 (a)).

Having upper-bounds on the pairwise gradient dissimilarities, we can efficiently find a near-optimal solution for problem (6) by turning it into a submodular maximization problem. A set function $F:2^{V}\rightarrow\mathbb{R}^{+}$ is submodular if $F(S\cup\{e\})-F(S)\geq F(T\cup\{e\})-F(T),$ for any $S\subseteq T\subseteq V$ and $e\in V\setminus T$. $F$ is monotone if $F(e|S)\geq 0$ for any $e\!\in\!V\!\setminus\!S$ and $S\subseteq V$. Minimizing the objective in Problem (6) is equivalent to maximizing the following submodular facility location function:

where $d_{0}$ is a constant satisfying $d_{0}\geq d^{u}_{ij}(\bm{W})$, for all $i,j\in V$. For maximizing the above monotone submodular function, the classical greedy algorithm provides a constant $(1-1/e)$-approximation. The greedy algorithm starts with the empty set $S_{0}=\emptyset$, and at each iteration $t$, it chooses an element $e\in V$ that maximizes the marginal utility $F(e|S_{t})=F(S_{t}\cup\{e\})-F(S_{t})$. Formally, $S_{t}=S_{t-1}\cup\{{\arg\max}_{e\in V}F(e|S_{t-1})\}$. The computational complexity of the greedy algorithm is $\mathcal{O}(nk)$. However, its complexity can be reduced to $\mathcal{O}(|V|)$ using stochastic methods [29], and can be further improved using lazy evaluation [28] and distributed implementations [31]. Note that this complexity does not involve any backpropagation as we use the upper-bounds calculated in Eq. (9). Hence, the subsets can be found very efficiently, in parallel from all classes. Unlike majority of the existing techniques for robust training against noisy labels that has a large computational complexity, robust training with Crust is even faster than training on the entire dataset. We also note that Problem (10) can be addressed in the streaming scenario for very large datasets [6].

Our experiments confirm that Crust can successfully find almost all the clean data points (c.f. Fig. 1 (a)). The following theorem guarantees that for a small fraction $\rho$ of noisy labels in the selected subsets, deep networks trained with gradient descent do not overfit the noisy labels.

The proof can be found in the Appendix. Note that the elements of the selected subsets are mostly clean, and hence the noise ratio $\rho$ is much smaller in the subsets compared to the entire dataset. Very recently, [32] showed that the classification error of neural networks trained on noisy datasets of size $n$ is controlled by the portion of the labels that fall over the nuisance space, i.e., $\|\Pi_{\mathcal{N}}(\bm{y})\|/\sqrt{n}$. Coresets of size $k$ selected by Crust are mostly clean. For such subsets, the label vector is mostly aligned with the information space, and thus $\|\Pi_{\mathcal{N}}(\bm{y}_{S})\|/\sqrt{k}$ is smaller. Our method improves the generalization performance by extracting subsets $S$ of size $k$ for which $\|\Pi_{\mathcal{N}}(\bm{y}_{S})\|/\sqrt{k}\leq\|\Pi_{\mathcal{N}}(\bm{y})\|/\sqrt{n}$. While defer the formal generalization proof to future work, our experiments show that even under severe noise (80% noisy labels), Crust successfully finds the clean data points and achieves a superior generalization performance (c.f. Fig. 1).

Next, we discuss how to reduce the error of backpropagation on the weighted centers.

There are two potential sources of error during weighted gradient descent updates in Eq. (8). First, we have an $\epsilon$ error in estimating the prominent singular value of the Jacobian matrix. And second, although the $k$-medoids formulation selects centers of clustered clean data points in the gradient space, there is still a small chance, in particular early in training process when the gradients are more uniformly distributed, that the coresets contain some noisy labels. Both errors can be alleviated if we slightly relax the constraint of training on exact feature vectors and their labels, and allow training on combinations of every center $j\!\in\!S$ with a few examples in its corresponding cluster $V_{j}$.

This is the idea behind mixup [51]. It extends the training distribution with convex combinations of pairs of examples and their labels. For every cluster $V_{j}$, we select a few data points $R_{j}\!\!\subset\!\!V_{j}\setminus\{j\},|R_{j}|\ll|V_{j}|$ uniformly at random, and for every data point $(\bm{x}_{i},\!y_{i}),i\!\in\!R_{j}$, we mix it up with the corresponding center $(\bm{x}_{j},\!y_{j}),j\!\in\!S^{*}$, to get the set $\hat{D}_{j}$ of mixed up points:

where $\lambda\sim\text{Beta}(\alpha,\alpha)\in[0,1]$ and $\alpha\in\mathbb{R}^{+}$. The mixup hyper-parameter $\alpha$ controls the strength of interpolation between feature-label pairs. Our experiments show that the subsets chosen by Crust contain mostly clean data points, but may contain some noisy labels early during the training (Fig. 1 (a)). Mixing up the centers with as few as one example from the corresponding cluster, i.e. $|R_{j}|=1$, can reduce the effect of potential noisy labels in the selected subsets. Therefore, mixup can further help improving the generalization performance, as is confirmed by our experiments (Table 2).

The subsets found in Problem (10) depend on the parameter vector $\bm{W}$ and need to be updated during the training. Let $\bm{W}^{\tau}$ be the parameter vector at iteration $\tau$. At update time $\tau\in[T]$, we first classify data points based on the updated predictions $\bm{y}^{\tau}=f(\bm{W}^{\tau},{\bm{X}})$. We denote by $U_{c}^{\tau}=\{(\bm{x}_{i},y_{i})\in\mathcal{D}|f(\bm{W}^{\tau},\bm{x}_{i})=\nu_{c}\}$ the set of data points labeled as $\nu_{c}$ in iteration $\tau$. Then, we find $S(\bm{W}^{\tau})$ by greedily extracting $k\!\cdot\!n_{c}$ medoids from each class, where $n_{c}=|U_{c}|/n$ is the fraction of data points in class $c\in[C]$. Finding separate coresets from each class can further help to not cluster together noisy data points spread out in the nuisance space, and improves the accuracy of the extracted coresets. Next, we partition the data points in every class to by assigning every data point to its closest medoid. Formally, for partition $V_{j}^{\tau}$ we have $V_{j}^{\tau}=\{i\in U_{c}^{\tau}|j={\arg\min}_{j\in S^{\tau}}d^{u}_{ij}\}$. Finally, we take a small random sample $R_{j}^{\tau}$ from every partition $V_{j}^{\tau}$, and for every data point $i\in R_{j}^{\tau}$ we mix it up with the corresponding medoid $j\in S(\bm{W}^{\tau})$ according to Eq. (11), and add the generated set $\hat{D}_{j}^{\tau}$ of mixed up data points to the training set. In our experiments we use $|R_{j}^{\tau}|=1$. At update time $\tau$, we train on the union of sets generated by mixup, i.e. $D^{\tau}=\{\hat{D}_{1}^{\tau}\cup\cdots\cup\hat{D}_{k}^{\tau}\}$, where every data point $i\in\hat{D}_{j}^{\tau}$ is weighted by $r_{i}=|V_{j}^{\tau}|/|R_{j}^{\tau}|$. The pseudo code of Crust is given in Algorithm 1.

Note that while Crust updates the coreset during the training, as almost all the clean data points are contained in the coresets in every iteration, gradient descent can successfully contract their residuals in Theorem 4.1 and fit the correct labels. Since Crust finds a new subset at every iteration $\tau$, we need to use $\alpha\!=\!\min_{\tau}\sqrt{r_{\min}^{\tau}}\sigma_{\min}({\cal{J}}(\bm{W}^{\tau},{\bm{X}}_{S^{\tau}}))$, and $\beta\!=\!\max_{\tau}\sqrt{r_{\max}^{\tau}}\|{\cal{J}}(\bm{W}^{\tau},{\bm{X}}_{S^{\tau}})\|$ in Theorem 4.1, where $\alpha$ and $\beta$ are the minimum and maximum singular values of the Jacobian of the subsets weighted by $\bm{r}^{\tau}$ found by Crust, during $T$ steps of gradient descent updates.

We first evaluate our method on CIFAR-10 and CIFAR-100, which contain 50,000 training images and 10,000 test images of size $32\times 32$ with 10 and 100 classes, respectively. We follow testing protocol adopted in [9, 14], by considering both symmetric and asymmetric label noise. Specifically, we test noise ratio of 0.2, 0.5, 0.8 for symmetric noise, and 0.4 for asymmetric noise.

In our experiments, we train ResNet-32 [15] for 120 epochs with a minibatch of 128. We use SGD with an initial learning rate of 0.1 and decays at epoch 80, 100 by a factor of 10 to optimize the objective, with a momentum of 0.9 and weight decay of $5\times 10^{-4}$. We only use simple data augmentation following [15]: we first pad 4 pixels on every side of the image, and then randomly crop a $32\times 32$ image from the padded image. We flip the image horizontally with a probability of 0.5. For Crust, we select coresets of size 50% of the size of the dataset unless otherwise stated.

We report top-1 test accuracy of various methods in Table 1. Our proposed method Crust outperforms all the baselines in terms of average test accuracy. While INCV attempts to find a subset of the training set with heuristics, our theoretically-principled method can successfully distinguish data points with correct labels from those with noisy labels, which results in a clear improvement across all different settings. It can be seen that Crust achieves a consistent improvement by an average of 3.15%, under various symmetric and asymmetric noisy scenarios, compared to the strongest baseline INCV. Interestingly, Crust  achieves the largest improvement of 6% over INCV, under sever 80% label noise. This shows the effectiveness of our method in extracting data points with clean labels from a large number of noisy data points, compared to other baselines.

Here we investigate the effect of each component of Crust and its importance, for robust training on CIFAR-10 with 20% and 50% symmetric noise. Table 2 summarizes the results.

Effect of the coreset. Based on the empirical results, the coreset plays an important role in improving the generalization of the trained networks. It is worthwhile noticing that by greedily finding the coreset based on the gradients, we can outperform INCV already. This clearly corroborate our theory and shows the effectiveness of Crust in filtering the noise and extracting data points with correct labels, compared to other heuristics. It also confirms our argument that although upper-bounded gradient dissimilarities in Eq. (9) has a much lower dimensionality compared to the exact gradients, noisy data points still spread out in the gradient space. Therefore, Crust can successfully identify central data points with clean labels in the gradient space.

Effect of mixup and model update. As discussed in Sec. 4.2, mixup can reduce the bias of estimating the full gradient with the coreset. Moreover, finding separate coresets from each class can further help filtering noisy labels, and improves the accuracy of the extracted coresets. At the beginning of every epoch, Crust updates the predictions based on the current model parameters and extract corsets from every class separately. We observed that both components, namely mixup and extracting coresets separately from each class based on the predictions of the model being trained, further improve the generalization and hence the final accuracy.

Size of the coresets. Fig. 1 demonstrates training curve for CIFAR-10 with 50% symmetric noise. Fig. 1(a) shows the accuracy of coresets of size 30%, 50%, and 70% selected by Crust. We observe that for various sizes of coresets, the number of noisy centers decreases over time. Furthermore, the fraction of correct labels in the coresets (label accuracy) decreases when the size of the selected centers increases from 30% to 70%. This demonstrates that Crust identifies clean data points first. Fig.  1 (b), (c) show the train and test accuracy, when training with Crust  on coresets of various sizes. We can see that coresets of size 30% achieve a lower accuracy as they are too small to accurately estimate the spectrum of the information space and achieve a good generalization performance. Coresets of size 70% achieve a lower training and test accuracy compared to coresets of size 50%. As 50% of the labels are noisy, subsets of size 70% contain at least 20% noisy labels and the model eventually overfits the noise.

WebVision is real-world dataset with inherent noisy labels [24]. It contains 2.4 million images crawled from Google and Flickr that share the same 1000 classes from the ImageNet dataset. The noise ratio in classes varies from 0.5% to 88% (Fig. 4 in [24] shows the noise distribution). We follow the setting in [17] and create a mini WebVision that consists of the top 50 classes in the Google subset with 66,000 images. We use both WebVision and ImageNet test sets for testing the performance of the model trained on coresets of size 50% of the data found by Crust. We train InceptionResNet-v2 [39] for 90 epochs with a starting learning rate of 0.1. We anneal the learning rate at epoch 30 and 60, respectively. The results are shown in Table 3. It can be seen that our method consistently outperforms other baselines, and achieves an average of 5% improvement in the test accuracy, compared to INCV.