SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dongmin Park<sup>1</sup>, Seola Choi<sup>1</sup>, Doyoung Kim<sup>1</sup>, Hwanjun Song<sup>2</sup>, Jae-Gil Lee<sup>1*</sup>  
<sup>1</sup> KAIST, <sup>2</sup> AWS AI Labs
{dongminpark, seola.choi, dodokim, jaegil}@kaist.ac.kr, hwanjun.song@amazon.com
Data pruning, which aims to downsize a large training set into a small informative subset, is crucial for reducing the enormous computational costs of modern deep learning. Though large-scale data collections invariably contain annotation noise and numerous robust learning methods have been developed, data pruning for the noise-robust learning scenario has received little attention. With state-of-the-art Re-labeling methods that self-correct erroneous labels while training, it is challenging to identify which subset induces the most accurate re-labeling of erroneous labels in the entire training set. In this paper, we formalize the problem of data pruning with re-labeling. We first show that the likelihood of a training example being correctly re-labeled is proportional to the prediction confidence of its neighborhood in the subset. Therefore, we propose a novel data pruning algorithm, Prune4ReL, that finds a subset maximizing the total neighborhood confidence of all training examples, thereby maximizing the re-labeling accuracy and generalization performance. Extensive experiments on four real and one synthetic noisy datasets show that Prune4ReL outperforms the baselines with Re-labeling models by up to  $9.1\%$  as well as those with a standard model by up to  $21.6\%$ .
By virtue of ever-growing datasets and the neural scaling law [1, 2], where the model accuracy often increases as a power of the training set size, modern deep learning has achieved unprecedented success in many domains, e.g., GPT [3], CLIP [4], and ViT [5]. With such massive datasets, however, practitioners often suffer from enormous computational costs for training models, tuning their hyper-parameters, and searching for the best architectures, which become the main bottleneck of development cycles. One popular framework to reduce these costs is data pruning, which reduces a huge training set into a small subset while preserving model accuracy. Notably, Sorscher et al. [6] have shown that popular data pruning approaches can break down the neural scaling law from power-law to exponential scaling, meaning that one can reach a desired model accuracy with much fewer data. Despite their great success, the impact of label noise on data pruning has received little attention, which is unavoidable in real-world data collection [7, 8, 9].
Noisy labels are widely known to severely degrade the generalization capability of deep learning, and thus numerous robust learning strategies have been developed to overcome their negative effect in deep learning [10]. Among them, Re-labeling [11], a family of methods that identify wrongly labeled examples and correct their labels during training by a self-correction module such as self-consistency regularization [12], has shown state-of-the-art performance. For example, the performance of DivideMix [13] trained on the CIFAR-10N [7] dataset containing real human annotation noise is nearly identical to that of a standard model trained on the clean CIFAR-10 dataset. Consequently, it is evident that this excellent performance of re-labeling must be carefully considered when designing a framework for data pruning under label noise.
(a) Pruning Performance.
Figure 1: Key idea of Prune4ReLU: (a) shows data pruning performance of Prune4ReLU and existing sample selection methods on CIFAR-10N with DivideMix; (b) shows how the neighborhood confidence affects the re-labeling correctness; (c) shows the goal of Prune4ReLU that maximize the neighbor confidence coverage to the entire training set, thereby maximizing the re-labeling accuracy.
(b) Re-labeling by Neighborhood.
(c) Goal of Prune4ReLU.
In this paper, we formulate a new problem of data pruning with re-labeling for a training set with noisy labels, which aims to maximize the generalization power of the selected subset with expecting that a large proportion of erroneous labels are self-corrected (i.e., re-labeled). Unfortunately, prior data pruning and sample selection algorithms are not suitable for our problem because the re-labeling capability is not taken into account, and have much room for improvement as shown in Figure 1(a). Popular data pruning approaches (denoted as Forgetting [14] and GraNd [15] in blue and yellow, respectively) favor hard (i.e., uncertain) examples because they are considered more beneficial for generalization [16]; however, because it is very difficult to distinguish between hard examples and incorrectly-labeled examples [17], many of the incorrectly-labeled examples can be included in the subset causing unreliable re-labeling. In addition, the small-loss trick [18] (denoted as SmallLoss in green) for sample selection favors easy examples because they are likely to be correctly labeled; however, they are not beneficial for generalization at a later stage of training. Therefore, this new problem necessitates the development of a new data pruning approach.
Accordingly, we suggest a completely novel approach of finding a subset of the training set such that the re-labeling accuracy of all training examples is preserved as much as possible with the model trained on the subset. The first challenge in this direction is how to estimate whether each example can be re-labeled correctly even before fully training models on the candidate subset. The second challenge is how to find the subset that maximizes the overall re-labeling accuracy of the entire training set in an efficient manner.
Addressing these two challenges, we develop a novel framework, called Prune4ReL. For the first challenge, we define the concept of the neighborhood confidence which is the sum of the prediction confidence of each neighbor example in the selected subset. We show that, as in Figure 1(b), an example with high neighborhood confidence is likely to be corrected by Re-labeling methods. We further provide theoretical and empirical evidence of this argument. For the second challenge, we show that the overall re-labeling accuracy is maximized by selecting a subset that maximizes the sum of its reachable neighborhood confidence for all training examples, as shown in Figure 1(c). Furthermore, because enumerating all possible subsets is a combinatorial optimization problem which is NP-hard [19], we provide an efficient greedy selection algorithm that expands the subset one by one by choosing the example that most increases the overall neighborhood confidence.
Extensive experiments on four real noisy datasets, CIFAR-10N, CIFAR-100N, WebVision, and Clothing-1M, and one synthetic noisy dataset on ImageNet-1K show that Prune4ReLU consistently outperforms the eight data pruning baselines by up to  $9.1\%$ . Moreover, Prune4ReLU with Re-labeling models significantly outperforms the data pruning baselines with a standard model by up to  $21.6\%$ , which reaffirms the necessity of data pruning with re-labeling.
2.1 Robust Learning under Noisy Labels
A long line of literature has been proposed to improve the robustness of DNNs against label noise, and refer to [10] for a detailed survey for deep learning with noisy labels. While some studies have
focused on modifying the architectures [20, 21, 22] and the loss functions [23, 24, 25], others have opted for sample selection approaches [26, 18, 27, 28] that select as many clean examples as possible while discarding noisy examples based on a certain cleanness criterion, e.g., small-loss [18]. Note that, these works do not consider the compactness or efficiency of the selected subset. Meanwhile, to further exploit even noisy examples for training, Re-labeling [11, 29] approaches try to correct noisy labels and reuse them for training with a re-labeling module, e.g., a heuristic rule [11]. Notably, according to recent benchmark studies on real-world noisy datasets [7], a family of Re-labeling methods with self-consistency regularization [12] has shown state-of-the-art performance.
In general, Re-labeling methods with self-consistency regularization are based on
$$
\mathcal {L} _ {R e - l a b e l i n g} (\tilde {\mathcal {D}}; \theta , \mathcal {A}) = \sum_ {(x, \tilde {y}) \in \tilde {\mathcal {D}}} \mathbb {1} _ {[ C _ {\theta} (x) \geq \delta ]} \mathcal {L} _ {c e} (x, \tilde {y}; \theta) + \lambda \sum_ {x \in \tilde {\mathcal {D}}} \mathcal {L} _ {r e g} (x; \theta , \mathcal {A}), \tag {1}
$$
where  $\tilde{\mathcal{D}} = \{(x_i,\tilde{y}_i)\}_{i = 1}^m$  is a given noisy training set obtained from a noisy joint distribution  $\mathcal{X}\times \hat{\mathcal{Y}}$ ,  $\theta$  is a classifier,  $\mathcal{A}$  is a strong data augmentation,  $\mathcal{C}_{\theta}(\cdot)$  is a prediction confidence score, and  $\delta$  is a threshold to identify confident (clean) examples for the supervised loss  $\mathcal{L}_{ce}(x,\tilde{y};\theta)$ , i.e., cross-entropy. The noisy labels are implicitly corrected by the self-consistency loss  $\mathcal{L}_{reg}(x;\theta)$  leveraging the power of strong augmentations [30]. DivideMix [13], ELR+ [31], CORES [32], and SOP+ [33] are popular approaches belonging to this family. DivideMix uses a co-training framework to further improve the re-labeling accuracy, and SOP+ introduces additional learnable variables combined with a self-consistency loss. For simplicity, we call this Re-labeling family with self-consistency regularization as "Re-labeling" throughout the paper. Despite their effectiveness, Re-labeling models tend to require more computation time due to additional data augmentations, multiple backbones, and longer training epochs, which raises a need to enhance its efficiency.
2.2 Data Pruning
In order to achieve high generalization performance with a selected subset, most data pruning approaches often prioritize the selection of hard or uncertain examples. Specifically, uncertainty-based methods [34, 35, 36] favor selecting less confident examples over more confident ones, as the former is assumed to be more informative than the latter. Similarly, geometry-based methods [37, 38] focus on removing redundant examples that are close to each other in the feature space, and loss-based methods [14, 15, 39, 40] favor selecting the examples with a high loss or gradient measured during training. However, these existing methods may not be effective in realistic scenarios under label noise because noisy examples also exhibit high uncertainty and could be wrongly considered informative for training [14]. Meanwhile, some recent works reported that existing data pruning methods do not work well at high pruning ratios [41, 42]. To alleviate this drawback, AL4DP [41] shows that mixing various levels of uncertain examples is better for data scarcity, Moderate [43] aims to select examples with the distances close to the median, and CCS [42] proposes a coverage-based method that jointly considers data coverage with sample importance. While a few works attempted to improve the robustness of sample selection against label noise by filtering the noise [19], no work yet considers the effect of data pruning on noise-robust learners such as Re-labeling.
We formalize a problem of data pruning with re-labeling such that it finds the most informative subset  $S$ , where a model  $\theta_S$  trained on  $S$  maximizes the re-labeling accuracy of the entire noisy training set  $\tilde{\mathcal{D}} = \{(x_i, \tilde{y}_i)\}_{i=1}^m$ . Formally, we aim to find an optimal subset  $S^*$  such that
$$
\mathcal {S} ^ {*} = \underset {\mathcal {S}: | \mathcal {S} | \leq s} {\operatorname {a r g m a x}} \sum_ {(x, \tilde {y}) \in \tilde {\mathcal {D}}} \mathbb {1} _ {[ f (x; \theta_ {\mathcal {S}}) = y ^ {*} ]}: \theta_ {\mathcal {S}} = \underset {\theta} {\operatorname {a r g m i n}} \mathcal {L} _ {R e - l a b e l i n g} (\mathcal {S}; \theta , \mathcal {A}), \tag {2}
$$
where  $y^{*}$  is the ground-truth label of a noisy example  $x$ ,  $f(x;\theta_{S}) \in \mathbb{R}^{c}$  is a  $c$ -way class prediction of the example  $x$  from the Re-labeling model  $\theta_{S}$ , and  $s$  is the target subset size.
Finding the optimal subset  $S^{*}$  through direct optimization of Eq. (2) is infeasible because the ground-truth label  $y^{*}$  is unknown in practice. In addition, the subset should be found at the early stage of
${}^{2}$  Maximizing the re-labeling accuracy is equivalent to correctly re-labeling all training examples. This formulation enables the model to exploit all clean labels for training, leading to a satisfactory generalization.
training, i.e., in a warm-up period, to reduce the computational cost [16]. To achieve these goals in an accurate and efficient way, in Section 3.1, we first introduce a new metric, the reduced neighborhood confidence, that enables estimating the re-labeling capacity of a subset even in the warm-up period. Then, in Section 3.2, we propose a new data pruning algorithm Prune4ReLU using this reduced neighborhood confidence to find a subset that maximizes the re-labeling accuracy.
3.1 Reduced Neighborhood Confidence
As a measurement of estimating the re-labeling accuracy, we use the confidence of neighbor examples for each target noisy example  $x$ , because the noisy examples are known to be corrected by their clean neighbor examples with self-consistency regularization [44]. Specifically, once an augmentation of a noisy example has a similar embedding to those of other clean neighbors in the representation space, the self-consistency loss can force the prediction of the noisy example to be similar to those of other clean neighbors as a way of re-labeling. This property is also evidenced by a theory of re-labeling with a generalization bound [45]. Thus, the neighboring relationship among examples can be a clear clue to estimate the re-labeling accuracy even in the early stage of training.
We define a neighborhood and its reduced neighborhood confidence to utilize the relationship of neighboring examples in Definitions 3.1 and 3.2.
Definition 3.1. (NEIGHBORHOOD). Let  $\mathcal{B}(x_i) = \{x\colon ||\mathcal{A}(x_i) - x||\leq \epsilon \}$  be a set of all possible augmentations from the original example  $x_{i}$  using an augmentation function  $\mathcal{A}$ . Then, given a noisy training set  $\tilde{\mathcal{D}}$ , a neighborhood of  $x_{i}$  is defined as  $\mathcal{N}(x_i) = \{x\in \tilde{\mathcal{D}}\colon \mathcal{B}(x_i)\cap \mathcal{B}(x)\neq \emptyset \}$ , which is the set of examples that are reachable by the augmentation  $\mathcal{A}$ .
Definition 3.2. (REDUCED NEIGHBORHOOD CONFIDENCE). The reduced neighborhood confidence  $C_{\mathcal{N}}(x_i; \mathcal{S})$  of an example  $x_i$  is the sum of the prediction confidence  $C(\cdot)$  of its neighbors  $x_j \in \mathcal{N}(x_i)$  in a given reduced (i.e., selected) subset  $\mathcal{S}$ , which is formalized as
$$
C _ {\mathcal {N}} \left(x _ {i}; \mathcal {S}\right) = \sum_ {x _ {j} \in \mathcal {S}} \mathbb {1} _ {\left[ x _ {j} \in \mathcal {N} \left(x _ {i}\right) \right]} \cdot C \left(x _ {j}\right), \tag {3}
$$
and its empirical reduced neighborhood confidence is computed by using the cosine similarity among the augmentations of all possible pairs of example embeddings,
$$
\hat {C} _ {\mathcal {N}} \left(x _ {i}; \mathcal {S}\right) = \sum_ {x _ {j} \in \mathcal {S}} \mathbb {1} _ {[ s i m \left(\mathcal {A} \left(x _ {i}\right), \mathcal {A} \left(x _ {j}\right)\right) \geq \tau ]} \cdot s i m \left(\mathcal {A} \left(x _ {i}\right), \mathcal {A} \left(x _ {j}\right)\right) \cdot C \left(x _ {j}\right), \tag {4}
$$
where  $sim(\cdot)$  is the cosine similarity between the augmentations  $\mathcal{A}(x)$  of two different examples in the embedding space, and  $\tau$  is a threshold to determine whether the two examples belong to the same neighborhood. Unlike Eq. (3), Eq. (4) is calculated as a weighted sum of prediction confidences with cosine similarity to approximate the likelihood of belonging to the neighborhood.
Based on these definitions, we investigate the theoretical evidence of employing the reduced neighborhood confidence as a means to estimate the re-labeling capacity of a subset.
Theoretical Evidence. A subset  $S$  with a high value of the total reduced neighborhood confidence, the sum of the reduced neighborhood confidence of each example in  $S$ , allows a Re-labeling model to maximize its re-labeling accuracy in the entire training set. We formally support this optimization by providing a theoretical analysis that extends the generalization bound in the prior re-labeling theory [45] to data pruning.
Assumption 3.3. (EXPANSION AND SEPARATION). Following the assumption in [45], the  $\alpha$ -expansion and  $\beta$ -separation assumptions hold for the training set  $\mathcal{D}$ . The  $\alpha$ -expansion means that an example is reachable to the  $\alpha$  number of augmentation neighbors on average, i.e.,  $\mathbb{E}_{x\in \mathcal{D}}[|\mathcal{N}(x)|] = \alpha$ . The  $\beta$ -separation means that data distributions with different ground-truth classes are highly separable, such that the average proportion of the neighbors from different classes is as small as  $\beta$ .
Under these assumptions, we can obtain a training accuracy (error) bound of a Re-labeling model trained on a subset  $S$  as in Theorem 3.4.
Theorem 3.4. Assume that a subset  $S \in \tilde{\mathcal{D}}$  follows  $\alpha_{\mathcal{S}}$ -expansion and  $\beta_{\mathcal{S}}$ -separation, where  $\alpha_{\mathcal{S}} \leq \alpha$ . Then, the training error of a Re-labeling model  $\theta_{\mathcal{S}}$  trained on  $S$  is bounded by the inverse of the total reduced neighborhood confidence  $\sum_{x \in \tilde{\mathcal{D}}} C_{\mathcal{N}}(x; S)$  such that,
$$
\operatorname {E r r} \left(\theta_ {\mathcal {S}}\right) \leq \frac {2 \cdot | \mathcal {S} | \cdot \operatorname {E r r} \left(\theta_ {\mathcal {M}}\right)}{\sum_ {x \in \tilde {\mathcal {D}}} C _ {\mathcal {N}} (x ; \mathcal {S})} + \frac {2 \cdot \alpha_ {\mathcal {S}}}{\alpha_ {\mathcal {S}} - 1} \cdot \beta_ {\mathcal {S}}, \tag {5}
$$
where  $\theta_{\mathcal{M}}$  is a model trained with the supervised loss in Eq. (1) on a given clean set  $\mathcal{M} \subset \mathcal{S}$ .
Proof. We extend the label denoising theorem in [45] by incorporating the influence of the subset to the expansion factor  $\alpha_{\mathcal{S}}$ . The complete proof is available in Appendix A.
Since  $\beta_{S}$  is usually very small, its effect on the error bound is negligible. Then, the bound highly depends on the total reduced neighborhood confidence. That is, as the total reduced neighborhood confidence increases, the error bound becomes tighter. This theorem supports that we can utilize the reduced neighborhood confidence for the purpose of maximizing the re-labeling accuracy.
Empirical Evidence. To empirically support Theorem 3.4, we validate the correlation between the empirical reduced neighborhood confidence and the re-labeling accuracy using CIFAR-10N, which is a real-world noisy benchmark dataset.
Specifically, we train DivideMix [13] on the  $20\%$  randomly selected subset  $\mathcal{S}$  for a warm-up training epoch of 10 and calculate the empirical reduced neighborhood confidence  $\hat{C}_{\mathcal{N}}(x; \mathcal{S})$  for the entire training set. Next, we fully train DivideMix [13] on the random subset  $\mathcal{S}$ . Last, we divide the entire training set into 15 bins according to the obtained  $\hat{C}_{\mathcal{N}}(x; \mathcal{S})$  and verify the average re-labeling accuracy for each bin.
Figure 2: Correlation between neighborhood confidence and re-labeling accuracy on a  $20\%$  randomly selected subset of CIFAR-10N.
Figure 2 shows how the re-labeling accuracy changes according to the empirical reduced neighborhood confidence in Eq. (4). (The term "empirical" is simply omitted hereafter.) The re-labeling accuracy shows a strong positive correlation with the reduced neighborhood confidence. Interestingly, as the neighborhood confidence increases, its utility in improving the re-labeling accuracy decreases, eventually reaching a convergence point after surpassing a certain threshold.
3.2 Data Pruning by Maximizing Neighborhood Confidence Coverage
We present a new data pruning algorithm called Prune4ReLU which optimizes the total reduced neighborhood confidence defined in Eq. (4). This objective is equivalent to identifying a subset that maximizes the re-labeling accuracy on the entire training set, as justified in Theorem 3.4. Therefore, the objective of Prune4ReLU is to find the subset  $S^{*}$ , which is formulated as
$$
\mathcal {S} ^ {*} = \underset {\mathcal {S}: | \mathcal {S} | \leq s} {\operatorname {a r g m a x}} \sum_ {x _ {i} \in \tilde {\mathcal {D}}} \sigma \left(\hat {C} _ {\mathcal {N}} \left(x _ {i}; \mathcal {S}\right)\right), \tag {6}
$$
where  $\sigma(z)$  is a utility function of the reduced neighborhood confidence  $\hat{C}_{\mathcal{N}}(x_i; S)$  in improving the re-labeling accuracy. By the observation in Figure 2, we define  $\sigma(z)$  as a non-decreasing and concave function where  $\sigma(0) = 0$ . In our implementation, we use the positive of the tanh function as the utility function, i.e.,  $\sigma(z) = \tanh(z)$ . However, directly solving Eq. (6) is computationally expensive and impractical due to its NP-hard nature as a Set-Cover problem [19]. Accordingly, we employ an approximation solution to efficiently address this combinatorial optimization.
Optimization with Greedy Approximation. We present a practical solution for solving the optimization problem stated in Eq. (6) using a greedy approximation. The objective function satisfies both the monotonicity and submodularity conditions, indicating that the return of the objective function monotonically increases and the marginal benefit of adding an example decreases as the subset grows. Thus, a greedy sample selection can be employed as in Algorithm 1.
$^{3}$ RandAug [30] is used as the augmentation function for the reduced neighborhood confidence.
In detail, we begin with an empty set  $S$  and initialize the reduced neighborhood confidence  $\hat{C}_{\mathcal{N}}$  to 0 (lowest confidence) for all training examples (Line 1). Next, at every step, we select an example  $x$  that maximizes the marginal benefit  $\sigma (\hat{C}_{\mathcal{N}}(x) + C(x)) - \sigma (\hat{C}_{\mathcal{N}}(x))$  of Eq. (6), and update the reduced neighborhood confidence  $\hat{C}_{\mathcal{N}}$  based on the similarity scores (Lines 3-7). To further improve robustness and efficiency, we introduce a class-balanced version, Prune4ReLU $_B$ , of which the detailed process is elaborated in Appendix B.
Algorithm 1 Greedy Neighborhood Confidence
INPUT:  $\tilde{\mathcal{D}}$ : training set,  $s$ : target subset size, and  $C(x)$ : confidence from warm-up classifier
1: Initialize  $\mathcal{S} \gets \emptyset$ ;  $\forall x \in \tilde{\mathcal{D}}$ ,  $\hat{C}_{\mathcal{N}}(x) = 0$
2: repeat
3:  $x = \operatorname{argmax}_{x\in \hat{\mathcal{D}}\setminus \mathcal{S}}\pmb {\sigma}(\hat{C}_{\mathcal{N}}(x) + C(x)) - \pmb {\sigma}(\hat{C}_{\mathcal{N}}(x))$
4:  $\mathcal{S} = \mathcal{S}\cup \{x\}$
5: for all  $v \in \tilde{\mathcal{D}}$  do
6:  $\hat{C}_{\mathcal{N}}(v) + = \mathbb{1}_{[sim(x,v)\geq \tau ]}\cdot sim(x,v)\cdot C(x)$
7: until  $|\mathcal{S}| = s$
OUTPUT: Final selected subset  $S$
In Theorem 3.5, we guarantee the selected subset  $\mathcal{S}$  obtained by our greedy solution achieves a  $(1 - 1/e)$ -approximation of the optimum.
Theorem 3.5. Since Eq. (6), denoted as  $OBJ$ , is a monotone, submodular, and non-negative function on  $x$ , the greedy solution provides a set with a  $(1 - 1 / e)$ -approximation of the optimum. Formally,
$$
O B J (\mathcal {S}) \geq (1 - 1 / e) \cdot O B J \left(\mathcal {S} ^ {*}\right). \tag {7}
$$
Proof. We prove the monotonicity and submodularity of Eq. (6). If the two conditions are satisfied, Eq. (7) naturally holds. See Appendix C for the complete proof.