Class Prototype-based Cleaner for Label Noise Learning

Our method follows the two-stage label noise learning framework DivideMix (Li et al., 2020a) and improves the framework with the proposed CPC. CPC comprises class prototypes $C=\{c_{k}\in\mathbb{R}^{1\times d}|k=1,2,...,K\}$, where $c_{k}$ indicates the prototype of $k$-th class and $d$ is the dimension of prototype embedding. Our DNN model consists of a CNN backbone, a classifier head and a projection layer. The backbone maps an image input $x_{i}$ to a feature vector $v_{i}\in\mathbb{R}^{1\times D}$. The classifier takes $v_{i}$ as input and outputs class prediction $p(\tilde{y_{i}}|x_{i})$. The projection layer serves to project the high dimension feature $v_{i}$ to a low-dimensional embedding $v^{\prime}_{i}\in\mathbb{R}^{1\times d}$, where $d<D$.

As shown in Fig. 2, we update the DNN as well as the CPC by iterating a two-stage training pipeline in every epoch. In the first stage, we update CPC as well as the projector in DNN, and utilize the updated CPC to partition label noise. We first calculate the cross-entropy loss of every training sample and fits a GMM to the losses. We utilize the GMM as a label noise cleaner to get a labeled clean set $\mathcal{X}^{GMM}$ and a unlabeled noise set $\mathcal{U}^{GMM}$. The data partition $\mathcal{X}^{GMM}$ and $\mathcal{U}^{GMM}$ are utilized to update the prototypes in CPC and parameters in the projector. Note that we cut off the gradient back-propagation from the projector to the CNN backbone. Then, the updated CPC is employed to re-divide the training data into another two set $\mathcal{X}$ and $\mathcal{U}$. In the second stage, we train DNN model to minimise the EVR in Eq. (1) with data partitioned by the cleaner. In the first $e$ epochs, we wait CPC to warm up, and minimise the EVR of DNNs based on training data partitioned by the GMM cleaner. After the $e$-th epoch, the label noise estimation results of CPC, i.e., $\mathcal{X}$ and $\mathcal{U}$ are employed to train DNNs, while the estimation results of GMM cleaner are only used to update prototypes in CPC. In inference, we utilize DNN classifier for image recognition, directly. In A.5, we further delineate the full framework.

In order to apply class-aware modulation to the label noise partitioning, we propose to learn an embedding space where samples from the same class are aligned with their class prototypes, and leverage the prototypes to recognize noise labels. The prototypes are typically learnt with intra-class consistency regularization, which urges samples in the same class to align with the corresponding class prototype while keeping samples not belonging to the class away. Previous methods (Wang et al., 2022; Li et al., 2020b) apply the intra-class consistency regularization to prototype learning via unsupervised contrastive objectives, e.g., prototypical contrastive objective (Li et al., 2020c), where the unsupervised training labels are typically determined by the similarity between samples and prototypes. The accuracy of the training labels are highly depends on the quality of representation learnt by the CNN encoder, which would be too low to effectively update the prototypes, especially in the early stage of training. In contrast, we empirically find that the GMM cleaner, which is operated under the well evaluated “small-loss prior”, are not as sensitive as the prototypes to the representation quality, and can provide more robust and accurate training labels.

Therefore, we propose to take samples in clean set $\mathcal{X}^{\mathrm{GMM}}$ as positive samples and those in noise set $\mathcal{U}^{\mathrm{GMM}}$ as negative samples to update prototypes. Specifically, given the feature embedding $v^{\prime}_{i}$ of a sample $x_{i}$ from $\mathcal{X}^{\mathrm{GMM}}$, we update prototypes $C$ as well as the parameters of the projector to maximize the score $q(z_{i}=0)$ between $c_{k=y_{i}}$ and $v^{\prime}_{i}$, and minimize the score between $c_{k\neq y_{i}}$ and $v^{\prime}_{i}$ via minimize $L_{\mathcal{X}^{\mathrm{GMM}}}$:

where $\lambda^{neg}=\frac{1}{K}$ weights the losses between positive pair and negative pairs to avoid under-fitting the positive samples. Given $v^{\prime}_{i}$ of a sample $x_{i}$ from $\mathcal{U}^{\mathrm{GMM}}$, we update prototypes $c_{k}$ as well as the parameters of the projector to minimize the score $q(z_{i}=0)$ between $c_{k=y_{i}}$ and $v^{\prime}_{i}$ via minimizing $L_{\mathcal{U}^{\mathrm{GMM}}}$:

At last, for noise samples in $\mathcal{U}^{\mathrm{GMM}}$ with high classification confidence, the samples are more likely to belong to the class predicted by DNNs, which is potentially valuable to the update of prototypes. Therefore, we collect such training samples $\mathcal{X}^{P}$ from $\mathcal{U}^{\mathrm{GMM}}$ taking the averaged classification confidence of samples in $\mathcal{X}^{\mathrm{GMM}}$ as the threshold. Specifically, given a sample in $\mathcal{U}^{\mathrm{GMM}}$ with the label predicted by DNNs $k=\max_{k}(p(\tilde{y_{i}}|x_{i}))$, the sample is collected into $\mathcal{X}^{P}$ if $p(\tilde{y_{i}}|x_{i})_{k}>average(\{p(\tilde{y_{i}}|x_{i})_{k}|(x_{j},y_{j}|y_{j}=k)\in\mathcal{X}^{\mathrm{GMM}}\})$. Then, we update the prototypes and projectors to minimize $L_{(\mathcal{X}^{P})}$:

The overall empirical risk $L_{C}$ for prototypes and the projector is as follows:

where $\alpha$ is the weight scalar.

CPC distinguishes a clean sample $(x_{i},y_{i})$ with the score $q(z_{i}=0)=\mathrm{sigmoid}(v^{\prime}_{i}c_{k=y_{i}}^{\top})$ and the threshold $\tau$. Samples with $q(z_{i}=0)>\tau$ are classified as clean, and otherwise as noise.

We provide theoretical justification on the efficacy of CPC from the perspective of Expectation-Maximization algorithm, which guarantees that though CPC does not follow the classical prototypical contrastive objective, it can still learn meaningful prototypes and act as an effective cleaner.

We consider training data with label noise $D=(X,Y)={(x_{i},y_{i})}_{i=1}^{N}$ as the observable data, and $Z\in\{0,1\}^{N}$ as the latent variable, where $z_{i}=0$ iff $(x_{i},y_{i})$ is clean (i.e., $y_{i}=\hat{y_{i}}$). The prototypes $C$ in the cleaner are taken as parameters expected to be updated. Then, the negative log likelihood for $D$ given $C$ is as follows:

where $q(z_{i})=p(z_{i}|x_{i},y_{i},C)$. According to the Bayes theorem and Jensen’s inequality , we have

where $-\sum_{D}\sum_{z_{i}\in\{0,1\}}q(z_{i})\log p(y_{i}|C,x_{i})$ is the upper bound of $NLL(D|C)$. Typically, we can adopt the EM algorithm to find the prototypes $C$ that minimize the upper bound by iterating:

E-step: Compute a new estimate of $q(z_{i})$ (i.e., clean or noise) according to prototypes $C^{old}$ from the last iteration:

M-step: Find the prototypes $C$ that minimizes the bound:

In our method, in order to introduce the “small-loss prior” to provide stronger and more robust supervision signals to the learning of CPC, in the E-step, we estimate the distribution of clean or noise of samples, denoted as $q(z^{\prime}_{i})$, via the GMM cleaner instead of $q(z_{i})$ in Eq. (8). And consequently, we replace the $q(z_{i})$ in Eq. (9) to $q(z^{\prime}_{i})$ and find the prototype $C$ minimize the bound. Next, we provide the justification that the EM algorithm still work by proving that $q(z^{\prime}_{i})$ can be considered as an approximation to $q(z_{i})$ in our framework.

In our method, $q(z^{\prime}_{i})=p(z^{\prime}_{i}|l(p(\tilde{y}_{i}|x_{i}),y_{i}))$, where $\tilde{y}_{i}\sim p(\tilde{y_{i}}|x_{i},\theta)$, which is the label predicted by the DNN parameterized by $\theta$. As introduced in section 4.1, in the first stage of each epoch, the CPC’s estimation results $z_{i}\sim q(z_{i})$ are utilized to divide training samples into a labeled set for clean data $\mathcal{X}=\{(x_{i},y_{i})|z_{i}=0\}$ and an unlabeled set for noise data $\mathcal{U}=\{(x_{i},y_{i})|z_{i}=1\}$. Then the parameters of DNNs, which we denote as $\theta$, are optimized using Eq. (1) in the second stage. There exists an optimal $\theta^{*}$ with respected to $z_{i}$, with which the softmax output $p(\tilde{y}_{i}|x_{i})$ of DNNs satisfies:

where $\ell(p(\tilde{y_{i}}|x_{i}),y_{i})$ is the cross-entropy loss between the network prediction and the annotated label. With these loss values, the subsequent GMM cleaner can easily distinguish samples of $\mathcal{X}$ from samples of $\mathcal{U}$. In other words, under the optimal $\theta^{*}$, the estimation of the GMM cleaner would be consistent with the partition of CPC, i.e., $z^{\prime}_{i}=z_{i}$. In practice, in each epoch, we takes the $\theta$ optimized to minimize Eq. (1) as an approximation to the optimal $\theta^{*}$ with respect to $z_{i}$, and consequently we can get $q(z^{\prime}_{i})$ as an approximation to $q(z_{i})$. Therefore, we can see that with the “small loss prior” introduced into the prototype learning, the EM optimization procedure would still work, which guarantees CPC can learn meaningful prototypes and act as an effective cleaner. In appendix A.4, we further present more details and empirical results to demonstrate the approximation is hold in practice.

Datasets. We evaluate our method on the following popular LNL benchmarks. For CIFAR-10 and CIFAR-100 (Krizhevsky et al., 2009), we experiment with two types of synthetic noise: symmetric and asymmetric, which are injected into the datasets following the standard setup in (Li et al., 2020a). Clothing1M (Xiao et al., 2015b) and WebVision1.0 (Li et al., 2017a) are two large-scale real-world label noise benchmarks. Clothing1M contains 1 million images in 14 categories acquired from online shopping websites, which is heavily imbalanced and most of the noise is asymmetric (Yi & Wu, 2019b). WebVision1.0 contains 2.4 million images crawled from the web using the concepts in ImageNet-ILSVRC12 (ILSVRC12). Following convention, we compare with SOTAs on the first 50 classes of WebVision, as well as the performance after transferring to ILSVRC12.

Real-world noise benchmarks. We evaluate our method on real-world large scale data sets, and compare our method with latest SOTA label noise learning methods, including DivideMix(Li et al., 2020a), LongReMix(Cordeiro et al., 2022), NGC(Wu et al., 2021), GJS(Englesson & Azizpour, 2021), ELR+(Liu et al., 2020), AugDMix(Nishi et al., 2021) and NCR(Huang et al., 2021). For WebVision, we measure the top1 and top5 accuracy on WebVision validation set and ImageNet ILSVRC12 validation set. We take ResNet50-based DivideMix (Zheltonozhskii et al., 2021) as baseline. As shown in Table 1, our CPC improves top1 and top5 accuracy over baseline model on WebVision by 3.33% and 2.81%, respectively. Our method achieves competitive performance on WebVision, and shows stronger transferable capability, outperforming other competitors on the ILSVRC12 validation set significantly. For Clothing1M, we apply the strong augmentation strategy (Nishi et al., 2021) to DivideMix as our baseline, and rerun the method three times. Our method achieves 75.4% accuracy on this challenging benchmark, outperforming all the other SOTAs. We also notice that though NCR achieves SOTA result on WebVision, it shows moderate performance compared to ELR+, DivideMix and AugDMix on Clothing1M containing asymmetric noise with imbalanced data distribution. It reveals that our method could be more robust across different label noise scenarios.

Synthetic noise benchmarks. We evaluate the performance of CPC on CIFAR-10 and CIFAR-100 datasets with symmetric label noise level ranging from 20% to 90% and asymmetric noise of rate 40%. We take AugDMix as the baseline, and compare our method with latest SOTA methods, where DivideMix, LongReMix and Aug-DMix are semi-supervised learning based methods. Following NGC and GJS, we run our method three times with different random seeds and report the mean and standard deviation. For other methods, e.g., ProtoMix (Li et al., 2021), we report the best results reported in their papers. As shown in Table 2, though with a baseline method as strong as AugDMix, our method brings about performance improvement across all noise levels as well as noise types consistently, and establishes new SOTAs on CIFAR-10 and CIFAR-100. Additionally, we notice that, under asymmetric noise set-up, semi-supervised learning based methods consistently outperform other methods that achieve SOTA results on WebVision benchmark,including NGC, GJS and NCR. The results reveal that semi-supervised learning based method could be more robust to asymmetric noise, while our method achieves SOTA performance among them.

Is CPC a better label noise cleaner? We evaluate the performance of label noise cleaner under both symmetric and asymmetric label noise set-ups. For symmetric noise, we use CIFAR-100 with 90% noise as benchmark to reveal the relationship between CPC and the significant performance improvement under this set-up. For asymmetric noise, we employ the most commonly adopted CIFAR-10-asym40% as benchmark. The AUC of clean/noise binary classification results of a cleaner is calculated as the evaluation metric. We take the original class-agnostic GMM cleaner (GMM_agn) proposed in DivideMix as baseline, and compare it to our CPC and the aforementioned naive class-aware GMM cleaner (GMM_awr). Furthermore, we also implement another version of CPC that trained based on the class-aware GMM cleaner. To distinguish these two CPC, we denote the regular one trained based on conventional class-agnostic GMM cleaner as CPC_agn, and the other one as CPC_awr. As shown in Figure 3, in both cases, the regular CPC_agn outperforms the baseline GMM_agn as well as GMM_awr, which demonstrates our class prototype-based method is the better label noise cleaner. As for the comparison between GMM_agn and GMM_awr, we find that in the situation of high symmetric noise, though GMM_agn shows better performance in the early stage of training, GMM_awr outperforms it in the second half stage of training. In the case of asymmetric noise, GMM_awr, which tend to classify hard clean samples in clean categories as noise wrongly, consistently underperforms GMM_agn across the whole training period. The results further prove that our class prototype-based method is the better choice for applying class-aware modulation to label noise cleaning, which is more robust across different noise types. Moreover, we find that in the case of asymmetric noise, CPC_agn achieves higher AUC compared to GMM_agn, which shows our method can partially make up for the shortcomings of GMM_agn. In the case of symmetric noise, we find that GMM_agn can further improve the performance of CPC, where CPC_awr achieves the best performance among the four cleaners.

Is the GMM cleaner beneficial to the learning of prototypes? In our method, we propose to leverage the GMM cleaner to facilitate the learning of prototypes via the “small loss prior”. To validate the effectiveness of our method, we first compare the quality of prototypes learnt in CPC with prototypes learnt in another prototype-based label noise learning method MoPro (Li et al., 2020b). We take WebVision as benchmark and utilize prototypes to classify test samples via measuring the similarity between samples and prototypes. The results show that, on the first 50 classes of WebVision, our prototype achieves a top1 accuracy of 78.44%, while MoPro’s accuracy is 72.23%, which demonstrates that our method is able to learn better prototypes. To further verify the contribution of the GMM cleaner, we remove the GMM cleaner and learn class prototypes in CPC via the typical prototypical contrastive objective as in MoPro. In experiments, we find that without the help of the GMM cleaner, the learnt prototypes generate less accurate data partition that further drawing back the overall training framework for DNNs, which proves the benefits of the GMM cleaner to our method. For more details and discussion, please refer to A.3.