SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jianhong Bai $^{1}$ , Zuozhu Liu $^{1\text{☑}}$ , Hualiang Wang $^{2}$ , Ruizhe Chen $^{1}$ , Lianrui Mu $^{1}$ , Xiaomeng Li $^{2}$ , Joey Tianyi Zhou $^{3}$ , Yang Feng $^{4}$ , Jian Wu $^{1}$ , Haoji Hu $^{1\text{☑}}$
<sup>1</sup>Zhejiang University, <sup>2</sup>The Hong Kong University of Science and Technology,
<sup>3</sup>Centre for Frontier AI Research, A*STAR, <sup>4</sup>Angelalign Technology Inc.
jianhongbai@zju.edu.cn
Data imbalance and open-ended distribution are two intrinsic characteristics of the real visual world. Though encouraging progress has been made in tackling each challenge separately, few works dedicated to combining them towards real-world scenarios. While several previous works have focused on classifying close-set samples and detecting open-set samples during testing, it's still essential to be able to classify unknown subjects as human beings. In this paper, we formally define a more realistic task as distribution-agnostic generalized category discovery (DA-GCD): generating fine-grained predictions for both close- and open-set classes in a long-tailed open-world setting. To tackle the challenging problem, we propose a Self-Balanced Co-Advice contrastive framework (BaCon), which consists of a contrastive-learning branch and a pseudo-labeling branch, working collaboratively to provide interactive supervision to resolve the DA-GCD task. In particular, the contrastive-learning branch provides reliable distribution estimation to regularize the predictions of the pseudo-labeling branch, which in turn guides contrastive learning through self-balanced knowledge transfer and a proposed novel contrastive loss. We compare BaCon with state-of-the-art methods from two closely related fields: imbalanced semi-supervised learning and generalized category discovery. The effectiveness of BaCon is demonstrated with superior performance over all baselines and comprehensive analysis across various datasets. Our code is available at: https://github.com/JianhongBai/BaCon.
Data imbalance [82, 79, 45] and open-ended distribution [14, 72, 17] are inherently intertwined with each other in the real visual world [40]. On the one hand, the frequency of different objects exhibits long-tailed distribution [48, 54] where a small portion of classes dominate the distribution and many classes are associated with only a few examples, as observed in healthcare [42, 75], autonomous driving [41, 36] and species classification [59]. On the other hand, humans or intelligent agents constantly encounter novel visual instances. Such complicated intertwined context poses great challenges to develop effective visual algorithms [28, 55] that apply to real-world scenarios [1, 12].
Tremendous efforts have been devoted to addressing the challenges of long-tail and open-set learning. In long-tail learning (LTL), previous works aim to highlight minority categories by re-sampling [53, 83, 16] or re-weighting [5, 26] approaches, debiasing model prediction from a statistical perspective [49, 43], decoupling feature learning and classifier training [29], or generalizing to semi-supervised [33, 31, 69] and self-supervised [38, 27, 81] scenarios. For open-world learning, existing work can be divided into three sub-domains based on how to handle open-set samples. Robust semi-supervised
Corresponding authors.
Figure 1: Illustration of distribution-agnostic generalized category discovery (DA-GCD).
learning (SSL) [46] assumes that out-of-distribution (OOD) samples may undesirably appear in the unlabeled set, and several techniques [20, 77, 9] are designed to reject these potential open-set samples to improve model robustness. Another line of methods [3, 23, 37, 34, 61, 44, 64] attempt to detect them with close-set knowledge during testing and formulate the task as OOD detection. Furthermore, generalized category discovery (GCD) [60] suggests to classify open-set samples into each class [4, 56, 70]. An earlier work OLTR [40] manage to integrate LTL and open-world learning towards more realistic visual recognition. OLTR handles close-set data imbalance and is aware of open-set instances similar to OOD detection. However, this approach does not meet the requirements for real-world recognition scenarios as: a) it works in a fully-supervised manner and does not effectively utilize a large amount of unlabeled data in the real world. b) it can only detect unseen (open-set) samples rather than classifying them with novel visual concepts.
In this work, we consider a real-world yet more challenging setting: we assume that only a portion of classes are collected and manually annotated (known classes) while massive unlabeled data from both known and open-set categories (novel classes) are available in the real world (Figure 1c). These two types of data, both long-tailed, form the semi-supervised training set, and the model is expected to classify all known and novel classes in a balanced test set, as illustrated in Figure 1. We name this setting Distribution-Agnostic Generalized Category Discovery (DA-GCD) and contrast it with other related settings in Table 1. DA-GCD differs from previous open-world tasks because it considers realistic long-tail settings. It also differs from the pioneering work OLTR [40] by further considering the semi-supervised setting and classifying novel classes, as in up-to-date open-world settings [60, 4].
The DA-GCD setting presents significant challenges that cannot be resolved with existing methods. Firstly, as confirmed by numerous studies in the field of long-tail learning, imbalanced datasets lead to severe performance degradation and model bias [79, 29, 27]. This issue is even more challenging in DA-GCD, where the prior dataset distribution is unavailable in open-world scenarios. Therefore, most LTL methods [43, 50] could not be directly extended to resolve the DA-GCD task. Meanwhile, since only unlabeled samples are available for novel classes (as shown in Figure 1c), the performance on these open-set classes is severely limited. Existing contrastive-based methods in GCD [60, 56] struggle to optimize the feature space of unlabeled samples, resulting in poor alignment and uniformity [66]. Lastly, GCD methods lack tailored designs for imbalanced datasets, leading to significant performance degradation on long-tailed datasets (as discussed in Section 5.2).
To conquer the aforementioned challenges, we propose BaCon, a novel Self-Balanced Co-Advice contrastive framework. BaCon comprises a contrastive-learning branch and a pseudo-labeling branch, which work collaboratively to provide interactive supervision to tackle the imbalanced recognition task under open-world scenarios. Specifically, the contrastive-learning branch provides distribution estimation during training, which is used to regularize the prediction results of the linear classifier for better pseudo-labeling. On the other hand, the generated pseudo-labels are further sampled and debiased to re-balance and provide additional supervision for the contrastive-learning branch. To fuse the knowledge of the two branches and learn a better feature space, we design a novel pseudo-label-based contrastive loss that clusters samples based on their positiveness scores.
We make the following major contributions. 1) We formulate a more realistic DA-GCD task that unifies long-tail and open-world learning and identify unique challenges that cannot be adequately addressed by existing techniques in both fields. 2) We design a novel self-Balanced Co-advice contrastive (BaCon) framework for DA-GCD, which comprises a contrastive-learning branch and pseudo-labeling branch that work collaboratively to tackle the data imbalance and classify novel (open-set) categories simultaneously. 3) We conduct extensive experiments on various datasets and perform fine-grained comparisons with SOTA baselines to validate the effectiveness of BaCon.
Table 1: Comparison of DA-GCD with other similar settings.
Setting	Scenario	Known classes	Novel classes	Involved in training? (novel classes)	Data Distribution
Robust SSL	Semi-supervised	Classify	Reject	✓	Balanced
OOD Det	Depend on method	Classify	Detect	Depend on method	Balanced
GCD	Semi-supervised	Classify	Classify	✓	Balanced
LT-SSL	Semi-supervised	Classify	N.A.	×	Long-tailed
OLTR	Fully-supervised	Classify	Detect	×	Long-tailed
DA-GCD	Semi-supervised	Classify	Classify	✓	Long-tailed
Open-World Learning tackling the scenario where some unknown (open-set) visual concepts or categories are involved during training or real-world applications. OOD detection [37, 24, 39], as a well-known sub-field, considering a distributional shift (either semantic shift [23] or covariate shift [35]) between the train and test stage, with the goal of detecting these OOD samples during testing. Robust semi-supervised learning [46, 20, 77, 9] apply different techniques to reject undesirable OOD samples in the unlabeled set to improve model robustness. More recently, a line of work [60, 4, 56, 70, 52, 51] propose to classify the open-set samples according to their visual concepts. Nevertheless, these methods are built on the class balance assumption, and we observe a consistent performance degradation when generalizing to DA-GCD scenarios. Distribution-agnostic recognition is also investigated by [73], which focuses on the task of novel category discovery.
Long-Tail Learning refers to a scenario where the majority of samples belong to a few classes, while the minority classes have only a small number of samples [63, 15]. Early research on long-tail learning mainly focused on fully-supervised scenarios, with representative methods that highlight the minority samples via re-sampling [53, 83, 16], re-weighting [5, 26], or knowledge transferring [68, 11, 76]. Long-tail learning under semi-supervised [33, 31, 69] and self-supervised [2, 74, 38, 27, 81, 78] settings are also widely investigated and have notable achievements. However, the requirement of the dataset distribution obstructs the application of fully/semi-supervised methods in DA-GCD, while self-supervised methods would disregard the available label information for known samples.
Contrastive Learning [8, 22, 18] provide distinctive and transferable representations by controlling the instance similarity in feature space, achieving significant progress in recent years. The following literature generalizes the idea to different scenarios [30, 62, 27], implements it for solving various downstream tasks [7, 71], and provides theoretical support [66, 67]. However, few attempts have been made at contrastive-based open-world recognition. GCD [60] proposes performing supervised CL on labeled samples and unsupervised CL on all data. OpenCon [56] generating pseudo-positive pairs for closely aligned representations, but this paradigm could lead to another dilemma: the representations and pseudo-labels are interdependent, which means an inferior feature space (e.g., lack of intra-class consistency) could lead to false positive pairs, and in turn deteriorate the learning of feature space. In short, those methods are incompatible with the proposed DA-GCD setting and have inferior representations for samples belonging to novel classes.
As illustrated in Figure 1a, our training set is composed of a labeled subset  $\mathcal{D}^l = \left\{\left(\boldsymbol{x}_i^l,y_i^l\right)\right\} \in \mathcal{X}\times \mathcal{Y}_l$  and an unlabeled subset  $\mathcal{D}^u = \left\{\left(\boldsymbol{x}_i^u,y_i^u\right)\right\} \in \mathcal{X}\times \mathcal{Y}_u$ , where we don't have the label information for any images in novel classes, i.e.,  $\mathcal{V}_l = \{\mathcal{V}_k\}$ ,  $\mathcal{V}_u = \mathcal{V} = \{\mathcal{V}_k\cup \mathcal{V}_n\}$ . Both  $\mathcal{D}^l$  and  $\mathcal{D}^u$  present a long-tail distribution with imbalance ratio  $\rho_l,\rho_u\gg 1$ . The objective of DA-GCD is to assign labels on a disjoint balanced test set containing all classes in  $\mathcal{V}_u$ . The total number of classes in the whole dataset  $\mathcal{D}$  is commonly regarded as a known prior [60, 4, 70], as we can effectively estimate it through off-the-shelf methods [21].
3.2 Revisiting Contrastive Learning
Contrastive-based representation learning methods [8, 22, 18, 30] aim to find an embedding function  $f$  to acquire optimal feature representation  $\pmb{z} \in \mathbb{R}^d$  of the input image  $\pmb{x} \in \mathbb{R}^{CHW}$  with  $\pmb{z} = f(\pmb{x})$ , such that  $\pmb{z}$  retains the discriminative semantic information of the input image. The general contrastive learning loss can be defined as:
$$
\mathcal {L} _ {\mathrm {C L}} (\mathcal {D}) = \frac {1}{| \mathcal {D} |} \sum_ {i \in \mathcal {D}} \left[ \frac {1}{| P (i) |} \sum_ {p \in P (i)} - \log \frac {\exp \left(\boldsymbol {z} _ {i} \cdot \boldsymbol {z} _ {p} / \tau\right)}{\sum_ {a \in A (i)} \exp \left(\boldsymbol {z} _ {i} \cdot \boldsymbol {z} _ {a} / \tau\right)} \right], \tag {1}
$$
where  $\mathcal{D}$  is the training set.  $P(i)$  is the set of indices of positive pairs in  $A(i)$ ,  $|\cdot|$  denotes the operation to compute the cardinality, and  $\tau$  is a hyper-parameter. In unsupervised learning, the positive set is formulated as the two views of the same image [8, 22]. [30] further generalizes the contrastive loss into supervised scenarios where the positive set is the images that share the same ground truth label.
4 Self-Balanced Co-Advice Contrastive Framework
As illustrated in Figure 2, we design a two-branch structure framework for DA-GCD. In Section 4.1, we propose the contrastive-learning branch to estimate the training distribution and further regularize the outputs of the pseudo-labeling branch. We then describe techniques for self-balanced knowledge transfer from the pseudo-labeling branch to the contrastive-learning branch in Section 4.2. Finally, in Section 4.3, a novel contrastive loss is proposed for better novel class learning.
4.1 Regularize Predictions with Dynamic Distribution Estimation
To enhance and re-balance the contrastive representation learning, we propose to learn an auxiliary classifier for pseudo-labeling. Our pseudo-labeling branch includes a cross-entropy loss  $\mathcal{L}_s$  on labeled data, a semi-supervised objective  $\mathcal{L}_u$  (e.g., cross pseudo supervision in [70]) on all samples, and a regularization term  $\mathcal{L}_{reg}$  to align the predictions with the data distribution to avoid non-activated classifiers. However, the distribution of the training set is agnostic in DA-GCD, and simply using the distribution of labeled set  $\pi_{\mathcal{D}^l}$  or a balance prior results in inferior performance (Table 6a). Furthermore, we observed that estimating the distribution from the pseudo-labeling branch itself could accumulate the estimation error and deteriorate model performance.
To this end, we perform estimation on the contrastive-learning branch as an alternative to avoid the accumulation of estimation bias. Concretely, we perform k-means clustering on all samples in the training set  $\mathcal{D}$  to obtain  $C$  clusters with size  $n$ , where  $C = |\mathcal{V}|$ , and normalized the sample number  $n$  to frequency  $\pi_e$ , i.e.,  $\pi_e = n / \sum_{c=1}^{C} (n[c])$ . Moreover, we also need to determine the corresponding relationship between clusters and categories. Therefore, we use the Hungarian optimal assignment algorithm [32] to map  $|\mathcal{V}^l|$  clusters to each known class. For the remaining  $|\mathcal{V}| - |\mathcal{V}^l|$  clusters, we sort them by the cluster size and assign them sequentially to the novel classes. Finally, we regularize the mean prediction of the pseudo-labeling branch with the aligned estimated distribution  $\pi_e$ :
$$
\mathcal {L} _ {r e g} = \mathrm {K L} \left[ \frac {1}{| B |} \sum_ {i \in B} \operatorname {s o f t m a x} \left(f _ {c l s} \left(\boldsymbol {x} _ {i}\right)\right) \| (\operatorname {a l i g n} \left(\boldsymbol {\pi} _ {e}\right)) ^ {p} \right], \tag {2}
$$
where  $B$  refers to the batch size during training,  $f_{cls}$  denote the backbone and the linear classifier of the auxiliary branch,  $\mathrm{KL}[\cdot \| \cdot ]$  is the Kullback-Leibler (KL) divergence between the two distributions, and  $p\in [0,1]$  is a hyper-parameter used to smooth the target long-tailed distribution. To balance estimation accuracy and computation overhead, we re-estimate the dataset distribution every  $r$  epoch. The overall training objective of the pseudo-labeling branch is as follows:
$$
\mathcal {L} _ {\text {c l s}} = \mathcal {L} _ {s} + \eta_ {1} \mathcal {L} _ {u} + \eta_ {2} \mathcal {L} _ {\text {r e g}}. \tag {3}
$$
<sup>1</sup>During inference, we use optimal assignment to determine the correspondence between classifiers and categories (detailed in Section 5.1), hence sequentially assigning the clusters to classes will not affect performance.
Figure 2: Overview of the self-balanced co-advice contrastive framework (BaCon).
4.2 Self-Balanced Knowledge Transfer: Debiasing and Sampling
In Section 4.1, we regularize the pseudo-labeling branch with the estimated dataset distribution  $\pi_e$  acquired from the contrastive-learning branch. In this step, we aim to transfer the knowledge from the pseudo-labeling branch to the contrastive-learning branch in turn and further enhance the representation learning, which is also beneficial to the estimation of data distribution  $\pi$ . However, the long-tail distribution of the underlying dataset places additional requirements on knowledge transferring. On one hand, we should ensure that the knowledge to be transmitted is not affected by the long-tailed distribution. Meanwhile, we hope that this process can help the contrastive-learning branch cope with the issue of data imbalance and the lack of supervision for novel classes.
We design a debiasing and sampling step of the pseudo-labels to meet the two aforementioned requirements respectively for knowledge transfer. First, we apply post-hoc logits adjustment [43] based on the estimated distribution to the predicted logits  $f_{cls}(\boldsymbol{x}_i)$  to eliminate the bias caused by long-tailed distribution:
$$
\widetilde {\boldsymbol {p}} _ {i} = \operatorname {s o f t m a x} \left(f _ {c l s} \left(\boldsymbol {x} _ {i}\right) - k \cdot \log \boldsymbol {\pi} _ {e}\right)), \tag {4}
$$
where  $\widetilde{\pmb{p}}_i$  denotes the rectified class probability prediction of sample  $\pmb{x}_i$ , and  $k$  is a hyper-parameter in [43]. Next, in order to re-balance the learning process and filter low-precision pseudo-labels, we propose to sample the unlabeled instances. Specifically, we sample the pseudo-labels of unlabeled instances in a training batch  $\widetilde{\pmb{p}}_B = \{\widetilde{\pmb{p}}_i : i = 1, \dots, B\}$  according to their prediction class  $c$ , where  $c_i = \arg \max_{i} \widetilde{\pmb{p}}_i$ . Formally, the class-wise sampling rate  $\mathrm{SR}^c$  can be defined as:
$$
\mathrm {S R} ^ {c} = \left\{ \begin{array}{l l} \left(\frac {\pi_ {e} ^ {c}}{\min  \left(\pi_ {e}\right)}\right) ^ {- \alpha}, & c \in \mathcal {Y} _ {B}, \\ \left(\frac {\pi_ {e} ^ {c}}{\min  \left(\pi_ {e}\right)}\right) ^ {- \beta}, & \text {o t h e r w i s e}, \end{array} \right. \tag {5}
$$
where  $\pi_{e}$  is the estimated distribution introduced in Section 4.1,  $\mathcal{Y}_B\subseteq \mathcal{Y}_k$  denotes the set of labels which involved in the current batch, and  $\alpha ,\beta \in [0,1]$  are two hyper-parameters. Setting  $\alpha = \beta = 0$  means sampling all pseudo-labels, while  $\alpha = \beta = 1$  means setting the sampling rate to be inversely proportional to the estimated number of samples in its category. Note that we prioritize selecting samples with higher prediction confidence in each class to remove the potentially false pseudo-labels. The sampled instances  $\{M(\mathcal{D}_u)\}$  complement the original long-tailed distribution to serve as a re-balance term for the labeled subset  $\{\bar{\mathcal{D}}_l\}$  and also provide additional supervision for novel classes. With this debiasing and sampling step, the sampled instances  $\{\mathcal{D}_l\cup M(\mathcal{D}_u)\}$  and their corresponding rectified pseudo-labels mitigate the impact of the imbalance issue and compensate the unlabeled open-set samples simultaneously.
4.3 Pseudo-label-based Soft Contrastive Learning
In Section 4.1 and Section 4.2, we generate pseudo-labels via the pseudo-labeling branch with help from the contrastive-learning branch. Those pseudo-labels could provide additional information to improve the contrastive-learning branch as well. In particular, we adopt a soft contrastive strategy to fully leverage the probabilistic information in pseudo-labels. We design a pair-wise positiveness score to adjust the contribution of different samples to the anchor instance. For image pair  $(\pmb{x}_i,\pmb{x}_j)$ , the positiveness score  $w_{ij} = \mathrm{Sim}(\widetilde{\pmb{p}}_i,\widetilde{\pmb{p}}_j)$ .  $w_{ij}$  is obtained by calculating the similarity of the rectified class probability distribution between  $\pmb{x}_i$  and  $\pmb{x}_j$ . In practice, We implement the similarity metric with the dot product operation, which can be mathematically interpreted as the probability of instance  $\pmb{x}_i$  and  $\pmb{x}_j$  belonging to the same class. We compare different definitions of the positiveness score in Section 5.4. Finally, we formulate the soft contrastive loss by incorporating  $\pmb{w}$  into Eq. 1:
$$
\mathcal {L} _ {\mathrm {C L}} ^ {\text {s o f t}} (\mathcal {D}) = \frac {1}{| \mathcal {D} |} \sum_ {i \in \mathcal {D}} \left[ \frac {1}{\sum_ {j \in A (i)} w _ {i j}} \sum_ {j \in A (i)} - w _ {i j} \cdot \log \frac {\exp \left(\boldsymbol {z} _ {i} \cdot \boldsymbol {z} _ {j} / \tau\right)}{\sum_ {a \in A (i)} \exp \left(\boldsymbol {z} _ {i} \cdot \boldsymbol {z} _ {a} / \tau\right)} \right]. \tag {6}
$$
Claim 1. Assume we train a model  $\theta$  with the proposed soft contrastive loss, and  $\pmb{x}_i$  and  $\pmb{x}_j$  are two samples within a training batch  $B$ . The optimal value of the contrastive logit  $p_{ij}^{*}$  is  $\frac{w_{ij}}{\sum_{k=1}^{|B|} w_{ik}}$ , where  $w_{ij}$  is the corresponding positiveness score for the sample pair  $(\pmb{x}_i, \pmb{x}_j)$  in Eq. 6.
The proof can be found in Appendix A. Claim 1 indicates that minimizing Eq. 6 encourages the similarity of features between two samples to be proportional to the corresponding positiveness score. In this way, we effectively transfer knowledge from the pseudo-labeling branch to contrastive learning. The overall training objective of the contrastive-learning branch consists of an unsupervised CL loss on all samples, a supervised CL loss on labeled instances, and the proposed soft CL loss on both labeled and sampled unlabeled subset:
$$
\mathcal {L} _ {c o n} = \mathcal {L} _ {\mathrm {C L}} ^ {u} (\mathcal {D}) + \gamma_ {1} \mathcal {L} _ {\mathrm {C L}} ^ {s} (\mathcal {D} _ {l}) + \gamma_ {2} \mathcal {L} _ {\mathrm {C L}} ^ {s o f t} (\mathcal {D} _ {l} \cup M (\mathcal {D} _ {u})). \tag {7}
$$
The pipeline of our method is illustrated in Figure 2, and the pseudo-code of BaCon can be found in Appendix B. It's worth noting that the backbone of the two branches is shared and frozen during training (except for the last block). Therefore the two-branch structure only slightly increases computational overheads. During inference, we only utilize the backbone of the contrastive-learning branch and obtain predictions via k-means clustering on the backbone features.