SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Wanxing Chang<sup>1</sup> Ye Shi<sup>1,2</sup> Jingya Wang<sup>1,2*</sup>
$^{1}$ ShanghaiTech University  $^{2}$ Shanghai Engineering Research Center of Intelligent Vision and Imaging
{changwx,shiye,wangjingya}@shanghaiitech.edu.cn
Learning with noisy labels (LNL) poses a significant challenge in training a well-generalized model while avoiding overfitting to corrupted labels. Recent advances have achieved impressive performance by identifying clean labels and correcting corrupted labels for training. However, the current approaches rely heavily on the models predictions and evaluate each sample independently without considering either the global or local structure of the sample distribution. These limitations typically result in a suboptimal solution for the identification and correction processes, which eventually leads to models overfitting to incorrect labels. In this paper, we propose a novel optimal transport (OT) formulation, called Curriculum and Structure-aware Optimal Transport (CSOT). CSOT concurrently considers the inter- and intra-distribution structure of the samples to construct a robust denoising and relabeling allocator. During the training process, the allocator incrementally assigns reliable labels to a fraction of the samples with the highest confidence. These labels have both global discriminability and local coherence. Notably, CSOT is a new OT formulation with a nonconvex objective function and curriculum constraints, so it is not directly compatible with classical OT solvers. Here, we develop a lightspeed computational method that involves a scaling iteration within a generalized conditional gradient framework to solve CSOT efficiently. Extensive experiments demonstrate the superiority of our method over the current state-of-the-arts in LNL. Code is available at https://github.com/changwxx/CSOT-for-LNL.
Deep neural networks (DNNs) have significantly boosted performance in various computer vision tasks, including image classification [33], object detection [61], and semantic segmentation [32]. However, the remarkable performance of deep learning algorithms heavily relies on large-scale high-quality human annotations, which are extremely expensive and time-consuming to obtain. Alternatively, mining large-scale labeled data based on a web search and user tags [49, 37] can provide a cost-effective way to collect labels, but this approach inevitably introduces noisy labels. Since DNNs can so easily overfit to noisy labels [4, 79], such label noise can significantly degrade performance, giving rise to a challenging task: learning with noisy labels (LNL) [50, 52, 46].
Numerous strategies have been proposed to mitigate the negative impact of noisy labels, including loss correction based on transition matrix estimation [35], re-weighting [60], label correction [76]
and sample selection [52]. Recent advances have achieved impressive performance by identifying clean labels and correcting corrupted labels for training. However, current approaches rely heavily on the models predictions to identify or correct labels even if the model is not yet sufficiently trained. Moreover, these approaches often evaluate each sample independently, disregarding the global or local structure of the sample distribution. Hence, the identification and correction process results in a suboptimal solution which eventually leads to a model overfitting to incorrect labels.
In light of the limitations of distribution modeling, optimal transport (OT) offers a promising solution by optimizing the global distribution matching problem that searches for an efficient transport plan from one distribution to another. To date, OT has been applied in various machine learning tasks [11, 83, 28]. In particular, OT-based pseudo-labeling [11, 73] attempts to map samples to class centroids, while considering the inter-distribution matching of samples and classes. However, such an approach could also produce assignments that overlook the inherent coherence structure of the sample distribution, i.e. intra-distribution coherence. More specifically, the cost matrix in OT relies on pairwise metrics, so two nearby samples could be mapped to two far-away class centroids (Fig. 1).
In this paper, to enhance intra-distribution coherence, we propose a new OT formulation for denoising and relabeling, called Structure-aware Optimal Transport (SOT). This formulation fully considers the intra-distribution structure of the samples and produces robust assignments with both global discriminability and local coherence. Technically speaking, we introduce local coherent regularized terms to encourage both prediction- and label-level local consistency in the assignments. Furthermore, to avoid generating incorrect labels in the early stages of training or cases with high noise ratios, we devise Curriculum and Structure-aware Optimal Transport (CSOT) based on SOT. CSOT constructs a robust denoising and relabeling allocator by relaxing one of the equality constraints to allow only a fraction of the samples with the highest confidence to be selected. These samples are then assigned with reliable pseudo labels. The allocator progressively selects and relabels batches of high-confidence samples based on an increasing budget factor that controls the number of selected samples. Notably, CSOT is a new OT formulation with a nonconvex objective function and curriculum constraints, so it is significantly different from the classical OT formulations. Hence, to solve CSOT efficiently, we developed a lightspeed computational method that involves a scaling iteration within a generalized conditional gradient framework [59].
Our contribution can be summarized as follows: 1) We tackle the denoising and relabeling problem in LNL from a new perspective, i.e. simultaneously considering the inter- and intra-distribution structure for generating superior pseudo labels using optimal transport. 2) To fully consider the intrinsic coherence structure of sample distribution, we propose a novel optimal transport formulation, namely Curriculum and Structure-aware Optimal Transport (CSOT), which constructs a robust denoising and relabeling allocator that mitigates error accumulation. This allocator selects a fraction of high-confidence samples, which are then assigned reliable labels with both global discriminability and local coherence. 3) We further develop a lightspeed computational method that involves a scaling iteration within a generalized conditional gradient framework to efficiently solve CSOT. 4) Extensive experiments demonstrate the superiority of our method over state-of-the-art methods in LNL.
Learning with noisy labels. LNL is a well-studied field with numerous strategies having been proposed to solve this challenging problem, such as robust loss design [82, 70], loss correction [35, 56], loss re-weighting [60, 80] and sample selection [52, 31, 41]. Currently, the methods that are delivering superior performance mainly involve learning from both selected clean labels and relabeled corrupted labels [46, 45]. The mainstream approaches for identifying clean labels typically rely on the small-loss criterion [31, 77, 71, 14]. These methods often model per-sample loss distributions using a Beta Mixture Model [51] or a Gaussian Mixture Model [57], treating samples with smaller loss as clean ones [3, 71, 46]. The label correction methods, such as PENCIL [76], Selfie [63], ELR [50], and DivideMix [46], typically adopt a pseudo-labeling strategy that leverages the DNNs predictions to correct the labels. However, these approaches evaluate each sample independently without considering the correlations among samples, which leads to a suboptimal identification and correction solution. To this end, some work [55, 45] attempt to leverage  $k$ -nearest neighbor predictions [6] for clean identification and label correction. Besides, to further select and correct noisy labels robustly,
OT Cleaner [73], as well as concurrent OT-Filter [23], designed to consider the global sample distribution by formulating pseudo-labeling as an optimal transport problem. In this paper, we propose CSOT to construct a robust denoising and relabeling allocator that simultaneously considers both the global and local structure of sample distribution so as to generate better pseudo labels.
Optimal transport-based pseudo-labeling. OT is a constrained optimization problem that aims to find the optimal coupling matrix to map one probability distribution to another while minimizing the total cost [40]. OT has been formulated as a pseudo-labeling (PL) technique for a range of machine learning tasks, including class-imbalanced learning [44, 28, 68], semi-supervised learning [65, 54, 44], clustering [5, 11, 25], domain adaptation [83, 12], label refinery [83, 68, 73, 23], and others. Unlike prediction-based PL [62], OT-based PL optimizes the mapping samples to class centroids, while considering the global structure of the sample distribution in terms of marginal constraints instead of per-sample predictions. For example, Self-labelling [5] and SwAV [11], which are designed for self-supervised learning, both seek an optimal equal-partition clustering to avoid the models collapse. In addition, because OT-based PL considers marginal constraints, it can also consider class distribution to solve class-imbalance problems [44, 28, 68]. However, these approaches only consider the inter-distribution matching of samples and classes but do not consider the intra-distribution coherence structure of samples. By contrast, our proposed CSOT considers both the inter- and intra-distribution structure and generates superior pseudo labels for noise-robust learning.
Curriculum learning. Curriculum learning (CL) attempts to gradually increase the difficulty of the training samples, allowing the model to learn progressively from easier concepts to more complex ones [42]. CL has been applied to various machine learning tasks, including image classification [38, 84], and reinforcement learning [53, 2]. Recently, the combination of curriculum learning and pseudo-labeling has become popular in semi-supervised learning. These methods mainly focus on dynamic confident thresholding [69, 29, 75] instead of adopting a fixed threshold [62]. Flex-match [78] designs class-wise thresholds and lowers the thresholds for classes that are more difficult to learn. Different from dynamic thresholding approaches, SLA [65] only assigns pseudo labels to easy samples gradually based on an OT-like problem. In the context of LNL, CurriculumNet [30] designs a curriculum by ranking the complexity of the data using its distribution density in a feature space. Alternatively, RoCL [85] selects easier samples considering both the dynamics of the per-sample loss and the output consistency. Our proposed CSOT constructs a robust denoising and relabeling allocator that gradually assigns high-quality labels to a fraction of the samples with the highest confidence. This encourages both global discriminability and local coherence in assignments.
Optimal transport. Here we briefly recap the well-known formulation of OT. Given two probability simplex vectors  $\alpha$  and  $\beta$  indicating two distributions, as well as a cost matrix  $\mathbf{C} \in \mathbb{R}^{|\alpha| \times |\beta|}$ , where  $|\alpha|$  denotes the dimension of  $\alpha$ , OT aims to seek the optimal coupling matrix  $\mathbf{Q}$  by minimizing the following objective
$$
\min  _ {\mathbf {Q} \in \Pi (\boldsymbol {\alpha}, \boldsymbol {\beta})} \left\langle \mathbf {C}, \mathbf {Q} \right\rangle , \tag {1}
$$
where  $\langle \cdot, \cdot \rangle$  denote Frobenius dot-product. The coupling matrix  $\mathbf{Q}$  satisfies the polytope  $\Pi(\alpha, \beta) = \left\{\mathbf{Q} \in \mathbb{R}_{+}^{|\alpha| \times |\beta|} | \mathbf{Q} \mathbb{1}_{|\beta|} = \alpha, \mathbf{Q}^{\top} \mathbb{1}_{|\alpha|} = \beta \right\}$ , where  $\alpha$  and  $\beta$  are essentially marginal probability vectors. Intuitively speaking, these two marginal probability vectors can be interpreted as coupling budgets, which control the mapping intensity of each row and column in  $\mathbf{Q}$ .
Pseudo-labeling based on optimal transport. Let  $\mathbf{P} \in \mathbb{R}_+^{B \times C}$  denote classifier softmax predictions, where  $B$  is the batch size of samples, and  $C$  is the number of classes. The OT-based PL considers mapping samples to class centroids and the cost matrix  $\mathbf{C}$  can be formulated as  $-\log \mathbf{P}$  [65, 68]. We can rewrite the objective for OT-based PL based on Problem (1) as follows
$$
\min  _ {\mathbf {Q} \in \Pi \left(\frac {1}{B} \mathbb {1} _ {B}, \frac {1}{C} \mathbb {1} _ {C}\right)} \left\langle - \log \mathbf {P}, \mathbf {Q} \right\rangle , \tag {2}
$$
where  $\mathbb{1}_d$  indicates a  $d$ -dimensional vector of ones. The pseudo-labeling matrix can be obtained by normalization:  $B\mathbf{Q}$ . Unlike prediction-based PL [62] which evaluates each sample independently,
Figure 1: (Top) Comparison between classical OT and our proposed Structure-aware OT. Classical OT tends to mismatch two nearby samples to two far-away class centroids when the decision boundary is not accurate enough. To mitigate this, our SOT generates local consensus assignments for each sample by preserving prediction-level and label-level consistency. Notably, for vague samples located near the ambiguous decision boundary, SOT rectifies their assignments based on the neighborhood majority consistency. (Bottom) The illustration of our curriculum denoising and relabeling based on proposed CSOT. The decision boundary refers to the surface that separates two classes by the classifier. The  $m$  represents the curriculum budget that controls the number of selected samples and progressively increases during the training process.
OT-based PL considers inter-distribution matching of samples and classes, as well as the global structure of sample distribution, thanks to the equality constraints.
Sinkhorn algorithm for classical optimal transport problem. Directly optimizing the exact OT problem would be time-consuming, and an entropic regularization term is introduced [19]:  $\min_{\mathbf{Q} \in \Pi(\alpha, \beta)} \langle \mathbf{C}, \mathbf{Q} \rangle + \varepsilon \langle \mathbf{Q}, \log \mathbf{Q} \rangle$ , where  $\varepsilon > 0$ . The entropic regularization term enables OT to be approximated efficiently by the Sinkhorn algorithm [19], which involves matrix scaling iterations executed efficiently by matrix multiplication on GPU.
Problem setup. Let  $\mathcal{D}_{train} = \{(\mathbf{x}_i, y_i)\}_{i=1}^N$  denote the noisy training set, where  $\mathbf{x}_i$  is an image with its associated label  $y_i$  over  $C$  classes, but whether the given label is accurate or not is unknown. We call the correctly-labeled ones as clean, and the mislabeled ones as corrupted. LNL aims to train a network that is robust to corrupted labels and achieves high accuracy on a clean test set.
4.1 Structure-Aware Optimal Transport for Denoising and Relabeling
Even though existing OT-based PL considers the global structure of sample distribution, the intrinsic coherence structure of the samples is ignored. Specifically, the cost matrix in OT relies on pairwise metrics and thus two nearby samples could be mapped to two far-away class centroids. To further consider the intrinsic coherence structure, we propose a Structure-aware Optimal Transport (SOT) for denoising and relabeling, which promotes local consensus assignment by encouraging prediction-level and label-level consistency, as shown in Fig. 1.
Our proposed SOT for denoising and relabeling is formulated by adding two local coherent regularized terms based on Problem (2). Given a cosine similarity  $\mathbf{S} \in \mathbb{R}^{B \times B}$  among samples in feature space, a one-hot label matrix  $\mathbf{L} \in \mathbb{R}^{B \times C}$  transformed from given noisy labels, and a softmax prediction matrix  $\mathbf{P} \in \mathbb{R}_{+}^{B \times C}$ , SOT is formulated as follows
$$
\min  _ {\mathbf {Q} \in \Pi \left(\frac {1}{B} \mathbb {1} _ {B}, \frac {1}{C} \mathbb {1} _ {C}\right)} \left\langle - \log \mathbf {P}, \mathbf {Q} \right\rangle + \kappa \left(\Omega^ {\mathbf {P}} (\mathbf {Q}) + \Omega^ {\mathbf {L}} (\mathbf {Q})\right), \tag {3}
$$
where the local coherent regularized terms  $\Omega^{\mathbf{P}}$  and  $\Omega^{\mathbf{L}}$  encourages prediction-level and label-level local consistency respectively, and are defined as follows
$$
\Omega^ {\mathbf {P}} (\mathbf {Q}) = - \sum_ {i, j} \mathbf {S} _ {i j} \sum_ {k} \mathbf {P} _ {i k} \mathbf {P} _ {j k} \mathbf {Q} _ {i k} \mathbf {Q} _ {j k} = - \left\langle \mathbf {S}, (\mathbf {P} \odot \mathbf {Q}) (\mathbf {P} \odot \mathbf {Q}) ^ {\top} \right\rangle , \tag {4}
$$
$$
\Omega^ {\mathbf {L}} (\mathbf {Q}) = - \sum_ {i, j} \mathbf {S} _ {i j} \sum_ {k} \mathbf {L} _ {i k} \mathbf {L} _ {j k} \mathbf {Q} _ {i k} \mathbf {Q} _ {j k} = - \left\langle \mathbf {S}, (\mathbf {L} \odot \mathbf {Q}) (\mathbf {L} \odot \mathbf {Q}) ^ {\top} \right\rangle , \tag {5}
$$
where  $\odot$  indicates element-wise multiplication. To be more specific,  $\Omega^{\mathbf{P}}$  encourages assigning larger weight to  $\mathbf{Q}_{ik}$  and  $\mathbf{Q}_{jk}$  if the  $i$ -th sample is very close to the  $j$ -th sample, and their predictions  $\mathbf{P}_{ik}$  and  $\mathbf{P}_{jk}$  from the  $k$ -th class centroid are simultaneously high. Analogously,  $\Omega^{\mathbf{L}}$  encourages assigning larger weight to those samples whose neighborhood label consistency is rather high. Unlike the formulation proposed in [1, 16], which focuses on sample-to-sample mapping, our method introduces a sample-to-class mapping that leverages the intrinsic coherence structure within the samples.
4.2 Curriculum and Structure-Aware Optimal Transport for Denoising and Relabeling
In the early stages of training or in scenarios with a high noise ratio, the predictions and feature representation would be vague and thus lead to the wrong assignments for SOT. For the purpose of robust clean label identification and corrupted label correction, we further propose a Curriculum and Structure-aware Optimal Transport (CSOT), which constructs a robust curriculum allocator. This curriculum allocator gradually selects a fraction of the samples with high confidence from the noisy training set, controlled by a budget factor, then assigns reliable pseudo labels for them.
Our proposed CSOT for denoising and relabeling is formulated by introducing new curriculum constraints based on SOT in Problem (3). Given curriculum budget factor  $m \in [0,1]$ , our CSOT seeks optimal coupling matrix  $\mathbf{Q}$  by minimizing following objective
$$
\min  _ {\mathbf {Q}} \langle - \log \mathbf {P}, \mathbf {Q} \rangle + \kappa (\Omega^ {\mathbf {P}} (\mathbf {Q}) + \Omega^ {\mathbf {L}} (\mathbf {Q}))
$$
$$
\left. \mathbf {Q} \in \left\{\mathbf {Q} \in \mathbb {R} _ {+} ^ {B \times C} \mid \mathbf {Q} \mathbb {1} _ {C} \leq \frac {1}{B} \mathbb {1} _ {B}, \mathbf {Q} ^ {\top} \mathbb {1} _ {B} = \frac {m}{C} \mathbb {1} _ {C} \right\}. \right. \tag {6}
$$
Unlike SOT, which enforces an equality constraint on the samples, CSOT relaxes this constraint and defines the total coupling budget as  $m \in [0,1]$ , where  $m$  represents the expected total sum of Q. Intuitively speaking,  $m = 0.5$  indicates that top  $50\%$  confident samples are selected from all the classes, avoiding only selecting the same class for all the samples within a mini-batch. And the budget  $m$  progressively increases during the training process, as shown in Fig. 1.
Based on the optimal coupling matrix  $\mathbf{Q}$  solved from Problem (6), we can obtain pseudo label by argmax operation, i.e.  $\hat{y}_i = \mathrm{argmax}_j\mathbf{Q}_{ij}$ . In addition, we define the general confident scores of samples as  $\mathcal{W} = \{w_0,w_1,\dots ,w_{B - 1}\}$ , where  $w_{i} = \mathbf{Q}_{i\hat{y}_{i}} / (m / C)$ . Since our curriculum accelerator assigns weight to only a fraction of samples controlled by  $m$ , we use  $\mathrm{topK}(S,k)$  operation (return top- $k$  indices of input set  $S$ ) to identify selected samples denoted as  $\delta_{i}$
$$
\delta_ {i} = \left\{ \begin{array}{l l} 1, & i \in \operatorname {t o p K} (\mathcal {W}, \lfloor m B \rfloor) \\ 0, & \text {o t h e r w i s e} \end{array} , \right. \tag {7}
$$
where  $\lfloor \cdot \rfloor$  indicates the round down operator. Then the noisy dataset  $\mathcal{D}_{train}$  can be split into  $\mathcal{D}_{clean}$  and  $\mathcal{D}_{corrupted}$  as follows
$$
\mathcal {D} _ {c l e a n} \gets \left\{(\mathbf {x} _ {i}, y _ {i}, w _ {i}) | \hat {y} _ {i} = y _ {i}, \delta_ {i} = 1, (\mathbf {x} _ {i}, y _ {i}) \in \mathcal {D} _ {t r a i n} \right\},
$$
$$
\mathcal {D} _ {\text {c o r r u p t e d}} \leftarrow \left\{\left(\mathbf {x} _ {i}, \hat {y} _ {i}, w _ {i}\right) | \hat {y} _ {i} \neq y _ {i}, \left(\mathbf {x} _ {i}, y _ {i}\right) \in \mathcal {D} _ {\text {t r a i n}} \right\}.
$$
4.3 Training Objectives
To avoid error accumulation in the early stage of training, we adopt a two-stage training scheme. In the first stage, the model is supervised by progressively selected clean labels and self-supervised by unselected samples. In the second stage, the model is semi-supervised by all denoised labels. Notably, we construct our training objective mainly based on Mixup loss  $\mathcal{L}^{mix}$  and Label consistency
loss  $\mathcal{L}^{lab}$  same as NCE [45], and a self-supervised loss  $\mathcal{L}^{\text{simsiam}}$  proposed in SimSiam [15]. The detailed formulations of mentioned loss and training process are given in Appendix. Our two-stage training objective can be constructed as follows
$$
\mathcal {L} ^ {\text {s u p}} = \mathcal {L} _ {\mathcal {D} _ {\text {c l e a n}}} ^ {\text {m i x}} + \mathcal {L} _ {\mathcal {D} _ {\text {c l e a n}}} ^ {\text {l a b}} + \lambda_ {1} \mathcal {L} _ {\mathcal {D} _ {\text {c o r r u p t e d}}} ^ {\text {s i m s i a m}}, \tag {9}
$$
$$
\mathcal {L} ^ {\text {s e m i}} = \mathcal {L} _ {\mathcal {D} _ {\text {c l e a n}}} ^ {\text {m i x}} + \mathcal {L} _ {\mathcal {D} _ {\text {c l e a n}}} ^ {\text {l a b}} + \lambda_ {2} \mathcal {L} _ {\mathcal {D} _ {\text {c o r r u p t e d}}} ^ {\text {l a b}}. \tag {10}
$$
5 Lightspeed Computation for CSOT
The proposed CSOT is a new OT formulation with nonconvex objective function and curriculum constraints, which cannot be solved directly by classical OT solvers. To this end, we develop a lightspeed computational method that involves a scaling iteration within a generalized conditional gradient framework to solve CSOT efficiently. Specifically, we first introduce an efficient scaling iteration for solving the OT problem with curriculum constraints without considering the local coherent regularized terms, i.e. Curriculum OT (COT). Then, we extend our approach to solve the proposed CSOT problem, which involves a nonconvex objective function and curriculum constraints.
5.1 Solving Curriculum Optimal Transport
For convenience, we formulate curriculum constraints in Probelm (6) in a more general form. Given two vectors  $\alpha$  and  $\beta$  that satisfy  $\| \alpha \| _1\geq \| \beta \| _1 = m$ , a general polytope of curriculum constraints  $\Pi^{\mathrm{c}}(\alpha ,\beta)$  is formulated as
$$
\Pi^ {c} (\alpha , \beta) = \left\{\mathbf {Q} \in \mathbb {R} _ {+} ^ {| \boldsymbol {\alpha} | \times | \beta |} \mid \mathbf {Q} \mathbb {1} _ {| \beta |} \leq \boldsymbol {\alpha}, \mathbf {Q} ^ {\top} \mathbb {1} _ {| \boldsymbol {\alpha} |} = \beta \right\}. \tag {11}
$$
For the efficient computation purpose, we consider an entropic regularized version of COT
$$
\min  _ {\mathbf {Q} \in \Pi^ {c} (\boldsymbol {\alpha}, \boldsymbol {\beta})} \left\langle \mathbf {C}, \mathbf {Q} \right\rangle + \varepsilon \left\langle \mathbf {Q}, \log \mathbf {Q} \right\rangle , \tag {12}
$$