SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jian Chen $^{1}$  Ruiyi Zhang $^{2}$  Tong Yu $^{2}$  Rohan Sharma $^{1}$
Zhiqiang Xu $^{3}$  Tong Sun $^{2}$  Changyou Chen $^{1}$
<sup>1</sup>University at Buffalo  <sup>2</sup>Adobe Research  <sup>3</sup>MBZUAI
{jchen378,rohanjag,changyou}@buffalo.edu
{ruizhang,tyu,tsun}@adobe.com zhiqiang.xu@mbzuai.ac.ae
Learning from noisy labels is a long-standing problem in machine learning for real applications. One of the main research lines focuses on learning a label corrector to purify potential noisy labels. However, these methods typically rely on strict assumptions and are limited to certain types of label noise. In this paper, we reformulate the label-noise problem from a generative-model perspective, i.e., labels are generated by gradually refining an initial random guess. This new perspective immediately enables existing powerful diffusion models to seamlessly learn the stochastic generative process. Once the generative uncertainty is modeled, we can perform classification inference using maximum likelihood estimation of labels. To mitigate the impact of noisy labels, we propose Label-Retrieval-Augmented (LRA) diffusion model<sup>1</sup>, which leverages neighbor consistency to effectively construct pseudo-clean labels for diffusion training. Our model is flexible and general, allowing easy incorporation of different types of conditional information, e.g., use of pre-trained models, to further boost model performance. Extensive experiments are conducted for evaluation. Our model achieves new state-of-the-art (SOTA) results on all standard real-world benchmark datasets. Remarkably, by incorporating conditional information from the powerful CLIP model, our method can boost the current SOTA accuracy by 10-20 absolute points in many cases.
Deep neural networks have achieved extraordinary accuracy on various classification tasks. These models are typically trained through supervised learning using large amounts of labeled data. However, large-scale data labeling could cost huge amount of time and effort, and is prone to errors caused by human mistakes or automatic labeling algorithms [1]. In addition, research has shown that the ability of deep neural network models to fit random labels can result in reduced generalization ability when learning with corrupted labels [2]. Therefore, robust learning methods using noisy labels are essential for applying deep learning models to cheap yet imperfect datasets for supervised learning tasks.
There are multiple types of label noise investigated by previous research. More recent research has focused on studying the more realistic feature-dependent label noise, where the probability of mislabeling a given instance depends on its characteristics. This type of noise is more consistent with label noise in real-world datasets [3, 4, 1, 5, 6, 7]. To address this type of noise, a model is expected to be able to estimate the uncertainty of each training label. Many state-of-the-art methods have primarily relied on the observation that deep neural networks tend to learn simple patterns before
$^{1}$ Code is available at https://github.com/puar-playground/LRA-diffusion
Figure 1: Label denoising as a reverse noising process.
memorizing random noise [8, 9]. This means a temporary phase exists in the learning process, where the model has learned useful features but has yet not started overfitting corrupt labels. At this stage, the model predictions can be used to modify the training labels, so that they are more consistent with model predictions [10, 11, 12, 13, 14, 15]. By correctly modifying the labels, the number of clean training sample increases, which can further benefit the training. These type of approaches, however, is inherently risky because the point at which the model starts to overfit varies with the network structure and dataset. Starting too early can corrupt the training data, while starting too late may not prevent overfitting [16]. Therefore, it is vital to carefully tune the hyper-parameters of the training strategy, such as the number of epochs for warm-up training, learning rate, and uncertainty threshold, to achieve successful training results.
Another class of methods adopts the assumption of label propagation in semi-supervised learning [17, 18], where nearby data points in a feature space tend to have the same label. Therefore, they use neighbor consistency $^2$  regularization to prevent overfitting of the model [19, 20]. The performance highly depends on the quality of the encoder that maps the data to the feature space, as retrieving a neighbor that belongs to a different class could further mislead the training process. Encoders are therefore required to first learn high-level features of the data that can be used for classification, which could be trained simultaneously with the classifier using noisy labels. However, the training can also lead to overfitting or underfitting.
In this paper, by contrast, we formulate the label noise problem from a generative-model perspective, which naturally provides new insights into approaching the problem. Our intuition is to view the noisy labeling process as a stochastic label generation process. Thus, we propose to adopt the powerful diffusion model as the generative building block. Figure 1 illustrates our intuition. In the generative process, we start with a noisy estimation of the label, then gradually refine it to recover the clean label, which is equivalent to the reverse denoising process of the diffusion model.
Specifically, the diffusion model takes a noisy label and some useful conditional information (to be specified) as inputs, and learns to recover/generate the ground-truth labels as outputs. One challenge in this setting is that only noisy labels are available in practice for training. To overcome this issue, we adopt the principle of neighbor consistency, and propose label-retrieval augmentation to construct pseudo clean labels for diffusion model training, where a pre-trained image encoder is used to define the neighborhood of a sample. It is worth noting that the pre-trained image encoder would not be affected by the label noise, because they can be trained in a self-supervised manner [21, 22] or on an additional clean dataset [23, 24]. In fact, pre-training can tremendously improve the model's adversarial robustness [25] and has been used to improve model robustness to label corruption [26]. Another merit of our design is that it is general enough to allow natural incorporation of powerful large pre-trained model such as the CLIP model to further boost the performance.
In addition, the probability nature of diffusion models can also be better equipped to handle uncertainty in the data and label, thus providing more robust and accurate predictions. We call our model LRA-diffusion (label-retrieval-augmented diffusion).
- We formulate learning from noisy labels as modeling a stochastic process of conditional label generation, and propose to adopt the powerful diffusion model to learn the conditional label distribution.
${}^{2}$  Nearby data points tend to have the same label.
- We incorporate the neighbor consistency principle into the modeling, and design a novel label-retrieval-augmented diffusion model to learn effectively from noisy label data.
- We further improve our model by incorporating auxiliary conditional information from large pre-trained models such as CLIP.
- Our model achieves the new state-of-the-art (SOTA) in various real-world noisy label benchmarks, e.g.,  $20\%$  accuracy improvement on noisy CIFAR-100 benchmark.
Diffusion models were initially designed for generative modeling. Recently, it has been extended for classification and regression problems. In this section, we introduce the Classification and Regression Diffusion Models (CARD) [27], which our model is based on.
The CARD model transforms deterministic classification into a conditional label generation process, allowing for more flexible uncertainty modeling in the labeling process [27]. Similar to the standard diffusion model, CARD consists of a forward process and a reverse process. In the forward process, an  $n$ -dimensional one-hot label  $\mathbf{y}_0$  is gradually corrupted to a series of intermediate random vectors  $\mathbf{y}_{1:T}$  which converges to a random variable with a multi-variant Gaussian distribution  $\mathcal{N}(f_q(\mathbf{x}),\mathbf{I})$  (latent distribution) after  $T$  steps, where the mean is defined by a pre-trained  $n$ -dimensional image encoder  $f_{q}$ . The transition steps between adjacent intermediate predictions is modeled as Gaussian distributions,  $q(\mathbf{y}_t|\mathbf{y}_{t - 1},f_q) = \mathcal{N}(\mathbf{y}_t;\pmb {\mu}_t,\beta_t\mathbf{I})$ , with mean values  $\pmb {\mu}_1,\dots ,\pmb {\mu}_T$  and a variance schedule  $\beta_{1},\dots ,\beta_{T}$  where  $\pmb {\mu}_t = \sqrt{1 - \beta_t}\mathbf{y}_{t - 1} + (1 - \sqrt{1 - \beta_t})f_q(\mathbf{x})$ . This admits a closed-form sampling distribution,  $q(\mathbf{y}_t|\mathbf{y}_0,f_q) = \mathcal{N}(\mathbf{y}_t;\pmb {\mu}_t,(1 - \bar{\alpha}_t)\mathbf{I}$ , with an arbitrary timestep  $t$  and  $\pmb {\mu}_t = \sqrt{\bar{\alpha}_t}\mathbf{y}_0 + (1 - \sqrt{\bar{\alpha}_t})f_q(\mathbf{x})$ . The mean term can be viewed as an interpolation between true data  $\mathbf{y}_0$  and the mean of the latent distribution  $f_{q}(\mathbf{x})$  with a weighting term  $\bar{\alpha}_{t} = \prod_{t}(1 - \beta_{t})$ .
In the reverse (generative) process, CARD reconstructs a label vector  $\mathbf{y}_0$  from an  $n$ -dimensional Gaussian noise  $y_{T} \sim \mathcal{N}(f_{q}(\mathbf{x}), \mathbf{I})$  by approximating the denoising transition steps conditioned on the data point  $\mathbf{x}$  and another pre-trained image encoder  $f_{p}$  in an arbitrary dimension. The transition step is also Gaussian for an infinitesimal variance  $\beta_{t}$  [28] (define  $\tilde{\beta}_{t} = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_{t}} \beta_{t}$ ):
$$
p _ {\theta} \left(\mathbf {y} _ {t - 1} \mid \mathbf {y} _ {t}, \mathbf {x}, f _ {p}\right) = \mathcal {N} \left(\mathbf {y} _ {t - 1}; \boldsymbol {\mu} _ {\theta} \left(\mathbf {y} _ {t}, \mathbf {x}, f _ {p}, t\right), \tilde {\beta} _ {t} \mathbf {I}\right). \tag {1}
$$
The diffusion model is learned by optimizing the evidence lower bound with stochastic gradient descent:
$$
\mathcal {L} _ {\mathrm {E L B O}} = \mathbb {E} _ {q} \left[ \mathcal {L} _ {T} + \sum_ {t > 1} ^ {T} \mathcal {L} _ {t - 1} + \mathcal {L} _ {0} \right], \text {w h e r e} \tag {2}
$$
$$
\mathcal {L} _ {0} = - \log p _ {\theta} (\mathbf {y} _ {0} | \mathbf {y} _ {1}, \mathbf {x}, f _ {p}), \mathcal {L} _ {t - 1} = D _ {\mathrm {K L}} (q (\mathbf {y} _ {t - 1} | \mathbf {y} _ {t}, \mathbf {y} _ {0}, \mathbf {x}, f _ {q}) | | p _ {\theta} (\mathbf {y} _ {t - 1} | \mathbf {y} _ {t}, \mathbf {x}, f _ {p}))
$$
$$
\mathcal {L} _ {T} = D _ {\mathrm {K L}} \left(q (\mathbf {y} _ {T} | \mathbf {y} _ {0}, \mathbf {x}, f _ {q}) | | p (\mathbf {y} _ {T} | \mathbf {x}, f _ {p})\right).
$$
Following [29], the mean term is written as  $\pmb{\mu}_{\theta}(\mathbf{y}_t,\mathbf{x},f_p,t) = \frac{1}{\sqrt{\alpha_t}} (\mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\pmb {\epsilon}_\theta (\mathbf{y}_t,\mathbf{x},f_p,t))$  and the objective can be simplified to  $\mathcal{L}_{\mathrm{simple}} = \| \pmb {\epsilon} - \pmb {\epsilon}_{\theta}(\mathbf{y}_t,\mathbf{x},f_p,t)\| ^2$
3 Label-Retrieval-Augmented Diffusion Model
Inspired by CARD, Label-Retrieval-Augmented (LRA) diffusion models reframe learning from noisy labels as a stochastic process of conditional label generation (i.e., label diffusion) process. In this section, we first provide an overview of the our model in Section 3.1 and then introduce the proposed label-retrieval-augmented component in Section 3.2, which can leverage label consistency in the training data. Next, we introduce an accelerated label diffusion process to significantly reduce classification model inference time in Section 3.3. Finally, a new conditional mechanism is proposed to enable the usage of pre-trained models in Section 3.4.
Our overall label-retrieval-augmented diffusion model is illustrated in Figure 2, where a diffusion model is adopted for progressively label denoising, by leveraging both the retrieved labels and auxiliary information from pre-trained models. Our model employs two pre-trained networks, denoted as  $f_{q}$  and  $f_{p}$  encoders, to encode conditional information that facilitates the generation process. The  $f_{q}$  encoder serves as a mean estimator for  $\mathbf{y}_T$ , providing an initial label guess for a given image. This encoder could be a standard classifier trained on noisy labels. On the other hand, the  $f_{p}$  encoder operates as a high-dimensional feature extractor, assisting in guiding the reverse procedure.  $\mathbf{y}_t$  and  $f_{p}(\mathbf{x})$  are concatenated together before being processed. Details about our neural-network architecture design are provided in Supplementary C.
During training, we use labels retrieved from the neighborhood as the generation target  $\mathbf{y}_0$ . Then, in the forward process, the distribution of neighboring labels is progressively corrupted towards a standard Gaussian distribution centered at the estimated mean  $f_{q}(\mathbf{x})$ . During testing, we employ a generalized DDIM method to efficiently compute the maximum likelihood estimation of  $\mathbf{y}_0$ .
Figure 2: Overview of the proposed framework for improving learning performance from noisy labels. The figure depicts the three main components, including (1) using diffusion models to imitate and inverse the label noisig process; (2) using pre-trained encoders (i.e.,  $f_{q}$  and  $f_{p}$ ) within the diffusion model, and (3) the label-retrieval-augmentation approach using the  $f_{p}$  encoder to encourage neighbor consistency of image labels.
3.2 Label-retrieval Augmentation for Training
The noisy nature of labels excludes the availability of clean labels for training. To mitigate the issue, we propose a training strategy based on the concept of retrieval augmented learning [30, 31], such that it is more resistant to label noise. Our main assumption is that in a latent space, data points from different classes form distinctive clusters. Therefore, the majority of a data point's neighbors are expected to have the same label as the point itself. To this end, we used a pre-trained encoder, illustrated as  $f_{p}$  in Figure 2, to map the data into an embedding space and retrieve a label  $y'$  from labels of the  $k$  nearest neighbors  $\{y^{(1)},\dots ,y^{(k)}\}$  in the training set. The diffusion model was then trained to learn the conditional distribution  $p(y'|x)$  of labels within the neighborhood, rather than the distribution of labels  $p(y|x)$  for the data point itself. We empirically select the value of  $k$  based on the KNN accuracy obtained on the validation data, as detailed in Supplementary Section C.
Label-retrieval augmentation enables the model to make use of the information from multiple and potentially more accurate labels to improve its prediction performance. Algorithm 1 describes the training procedure. Additionally, diffusion models are known to be effective at modeling multimodal distributions. By training the model to generate different labels from neighbors based on the same data point, the model can produce stochastic predictions based on the distribution to capture the uncertainty inherent in the data labeling process. As a result, the trained model can be used not only as a classifier, but also as a sampler that simulates the actual labeling process.
3.3 Efficient Inference with generalized DDIM
The iterative generation nature of the classification diffusion model makes its inference efficiency not comparable to traditional classifiers. To enhance the inference efficiency, we propose to incorporate the efficient sampling methods, Denoising Diffusion Implicit Model (DDIM) [32], to accelerate the
Input: training set  $\{\mathbf{X},\mathbf{Y}\}$ , image encoder  $f_{p}, f_{q}$ .
1: while not converged do
2: Sample data  $(\mathbf{x},y)\sim \{\mathbf{X},\mathbf{Y}\}$ ; time slice  $t\sim \{1,\dots ,T\}$ ; and noise  $\epsilon \sim \mathcal{N}(0,\mathbf{I})$
3: Retrieve labels  $\{y^{(1)},\dots ,y^{(k)}\}$  of neighbors in the feature space defined by the encoder  $f_{p}$
4: Sample  $y^\prime \sim \{y,y^{(1)},\dots ,y^{(k)}\}$  , and convert it to a one-hot vector  $\mathbf{y}_0$
5: Take gradient descent step on the loss:
$$
\left\| \boldsymbol {\epsilon} - \boldsymbol {\epsilon} _ {\theta} (\sqrt {\bar {\alpha} _ {t}} \mathbf {y} _ {0} + (1 - \sqrt {\bar {\alpha} _ {t}}) f _ {q} (\mathbf {x}) + \sqrt {1 - \bar {\alpha} _ {t}} \boldsymbol {\epsilon}, \mathbf {x}, f _ {p} (\mathbf {x}), t) \right\| ^ {2}
$$
6: end while
label diffusion process. However, the utilization of the mean estimator  $f_{q}$  makes DDIM incompatible with our setting, as our generation process begins with a non-zero mean Gaussian distribution  $\mathcal{N}(f_q(\mathbf{x}),\mathbf{I})$ . Therefore, we adjust the DDIM method into a more general form that fits our framework. Analogous to DDIM, our sampling process maintains the same marginal distribution as the original  $q(\mathbf{y}_t|\mathbf{y}_0,f_q)$  closed-form sampling process. Detailed derivations can be found in Supplementary A.
With DDIM, the trained model can generate a label vector in much less steps following a pre-defined sampling trajectory of  $\{T = \tau_S >, \dots, > \tau_1 = 1\}$ , where  $S < T$ . Consequently,  $\mathbf{y}_t$  can be computed as:
$$
\mathbf {y} _ {\tau_ {s}} = \sqrt {\bar {\alpha} _ {\tau_ {s}}} \mathbf {y} _ {0} + (1 - \sqrt {\bar {\alpha} _ {\tau_ {s}}}) f _ {q} (\mathbf {x}) + \sqrt {1 - \bar {\alpha} _ {\tau_ {s}}} \boldsymbol {\epsilon}, \tag {3}
$$
where  $\epsilon \sim \mathcal{N}(\mathbf{0},\mathbf{I})$ . Similar to CARD [27], we predict the denoised label  $\tilde{\mathbf{y}}_0$ , a prediction of  $\mathbf{y}_0$  given  $\mathbf{y}_{\tau_s}$ , as:
$$
\tilde {\mathbf {y}} _ {0} = \frac {1}{\sqrt {\bar {\alpha} _ {\tau_ {s}}}} \left[ \mathbf {y} _ {\tau_ {s}} - \left(1 - \sqrt {\bar {\alpha} _ {\tau_ {s}}}\right) f _ {q} (\mathbf {x}) - \sqrt {1 - \bar {\alpha} _ {\tau_ {s}}} \epsilon_ {\theta} \left(\mathbf {y} _ {\tau_ {s}}, \mathbf {x}, f _ {p} (\mathbf {x}), \tau_ {s}\right) \right]. \tag {4}
$$
When  $\tau_{s-1} > 0$ , we can compute  $\mathbf{y}_{\tau_{s-1}}$  given  $\mathbf{y}_{\tau_s}$  from the non-Markovian forward process defined as:
$$
\mathbf {y} _ {\tau_ {s - 1}} = \sqrt {\bar {\alpha} _ {\tau_ {s - 1}}} \tilde {\mathbf {y}} _ {0} + (1 - \sqrt {\bar {\alpha} _ {\tau_ {s - 1}}}) f _ {q} (\mathbf {x}) + \sqrt {1 - \bar {\alpha} _ {\tau_ {s - 1}}} \cdot \boldsymbol {\epsilon} _ {\theta} (\mathbf {y} _ {\tau_ {s}}, \mathbf {x}, f _ {p} (\mathbf {x}), \tau_ {s}). \qquad (5)
$$
As the dimension of label vectors is usually much lower than that of an image, the model can employ fewer steps in the reverse process without compromising generative quality. In our experiments, we use  $S = 10$  and  $T = 1000$ , substantially reducing the time cost of the classification process. Supplementary Figure B gives an example of the label generation (classification) process on the CIFAR-10 dataset.
To further enhance the inference efficiency, we propose a simple and effective trick for computing the maximum likelihood estimation of labels. As the generative process is deterministic given  $\mathbf{y}_T$ , which is sampled from a uni-modal Gaussian distribution, we approximate the maximum likelihood estimation of labels by initiating from the mean, i.e.,  $\mathbf{y}_0 = \mathrm{DDIM}(\mathbf{y}_T = f_q(\mathbf{x}), \mathbf{x})$ . This trick circumvents the time-consuming majority voting approximation that demands repeated generation.
3.4 Flexible conditioning with pre-trained encoders
The original CARD model employs a single model for both the  $f_{p}$  and  $f_{q}$  encoders. However, this limits their representation capacity [33] as the dimension of  $f_{q}(\mathbf{x})$  is typical relatively small, i.e., equalling the number of classes. To mitigate this and improve model performance, we abandon the assumption that  $f_{p} = f_{q}$ , enabling the use of a more powerful pre-trained encoder (e.g., the CLIP image encoder [24]) with arbitrary dimensions for  $f_{p}$ .
Empirically, we find that the model can still achieve satisfactory performance when the magnitude of  $f_{q}(\mathbf{x})$  is small, which means the latent representation  $\mathbf{y}_T = f_q(\mathbf{x}) + \epsilon$  is dominated by the noise term  $\epsilon \sim \mathcal{N}(\mathbf{0},\mathbf{I})$ . In this case, the information provided by  $f_{q}(\mathbf{x})$  to the diffusion process is limited. As a result, we simply set  $f_{q}(\mathbf{x}) = \mathbf{0}$  to avoid handling an additional n-dimensional  $f_{q}$  encoder. For the  $f_{p}$  encoder, one can employ flexible pre-trained models as presented in Section 5. In this paper, we use the SimCLR model trained on the training images (without supervised information) and the pre-trained CLIP model.
Robust loss function and regularization techniques. Several noise-robust loss functions and regularization techniques have been proposed as alternatives to the commonly used cross-entropy loss (CE), which is not robust to label noise. Mean absolute error (MAE) [34] loss has been shown to be robust against noisy labels. Generalized Cross-Entropy (GCE) [35] combines CE and MAE for faster convergence and better accuracy. Symmetric cross-entropy Learning (SL) [36] couples CE with a noise-robust counterpart and has been found to have higher performance than GCE, particularly for high noise rates. Label Smoothing Regularization [37] alleviates overfitting by linearly combining labels with a uniform distribution. Bootstrapping technique [38] combines the labels with the current model prediction. Dynamic bootstrapping [39, 40] uses the prediction confidence to control the weighting in the combination. Neighbor Consistency Regularization (NCR) [19] encourages consistency of prediction based on learned similarity. Our method is also based on the principle of neighbor consistency. However, instead of encouraging consistent predictions among neighbors, our model directly learns from the labels of neighbors. This allows for estimating instance-level uncertainty by learning the label distribution among neighbors, rather than learning a point estimation.
Data recalibration. Data recalibration techniques progressively remove or correct mislabeled data during training to improve the reliability of training data. Wang et al. [11] used the learned similarity and label consistency to identify and discard data with noisy labels. TopoFilter [41] selects clean data by analyzing the topological structures of the training data in the learned feature space. Cheng et al. [4] defines a Bayes optimal classifier to correct labels. Zheng et al. [14] proposed using a likelihood ratio test (LRT) to correct training labels based on predictions. Zhang et al. [15] used LRT to correct labels progressively and provides a theoretical proof for convergence to the Bayes optimal classifier. Dividemix [42], LongReMix [43], and CC [44] treat the low confident data as unlabeled, and then employ semi-supervised learning algorithms [45] for further analysis. C2D [46] combines Dividemix with self-supervised pre-training to boost its performance by improving the quality of the extracted features. Our approach employs the same assumption as TopoFilter that data belonging to the same class should be clustered together with ideal feature representations. However, our technique isn't confined to learned features potentially distorted by label noises. Instead, similar to C2D, our method can effectively leverage the high-quality feature learned by pre-trained encoders to achieve superior accuracy.