SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jiachen Liang $^{1,2}$ , Ruiming Hou $^{1}$ , Hong Chang $^{1,2}$ , Bingpeng Ma $^{2}$ , Shiguang Shan $^{1,2}$ , Xilin Chen $^{1,2}$
$^{1}$  Institute of Computing Technology, Chinese Academy of Sciences
$^{2}$ University of Chinese Academy of Sciences
jiachen.liang@vipl.ict.ac.cn, {houruibing, changhong, sgshan, xlchen} @ict.ac.cn, bpma@ucas.ac.cn
Traditional semi-supervised learning (SSL) assumes that the feature distributions of labeled and unlabeled data are consistent which rarely holds in realistic scenarios. In this paper, we propose a novel SSL setting, where unlabeled samples are drawn from a mixed distribution that deviates from the feature distribution of labeled samples. Under this setting, previous SSL methods tend to predict wrong pseudolabels with the model fitted on labeled data, resulting in noise accumulation. To tackle this issue, we propose Self-Supervised Feature Adaptation (SSFA), a generic framework for improving SSL performance when labeled and unlabeled data come from different distributions. SSFA decouples the prediction of pseudo-labels from the current model to improve the quality of pseudo-labels. Particularly, SSFA incorporates a self-supervised task into the SSL framework and uses it to adapt the feature extractor of the model to the unlabeled data. In this way, the extracted features better fit the distribution of unlabeled data, thereby generating high-quality pseudo-labels. Extensive experiments show that our proposed SSFA is applicable to various pseudo-label-based SSL learners and significantly improves performance in labeled, unlabeled, and even unseen distributions.
Semi-Supervised Learning (SSL) uses a small amount of labeled data and a large amount of unlabeled data to alleviate the pressure of data labeling and improve the generalization ability of the model. Traditional SSL methods [21, 43, 32, 28, 22, 30, 5, 6] usually assume that the feature distribution of unlabeled data is consistent with that of labeled one. However, in many real scenarios, this assumption may not hold due to different data sources. When unlabeled data is sampled from distributions different from labeled data, traditional SSL algorithms will suffer from severe performance degradation, which greatly limits their practical application.
In real-world scenarios, it is quite common to observe feature distribution mismatch between labeled and unlabeled samples. On the one hand, unlabeled samples could contain various corruptions. For example, in automatic driving, the annotated images used for training can hardly cover all driving scenes, and a plethora of images under varying weather and camera conditions are captured during real driving. Similarly, in medical diagnosis, individual differences and shooting conditions among patients could incur various disturbances in unlabeled data. On the other hand, unlabeled samples could contain unseen styles. For example, in many tasks, the labeled samples are typically real-world photos, while unlabeled data collected from the Internet usually contain more styles that are not present in the labeled data, such as cartoons or sketches. Notably, although the distributions of these unlabeled data may differ from the labeled data, there is implicit common knowledge between them that can compensate for the diversity and quantity of training data. Therefore, it is crucial to enlarge the SSL scope to effectively utilize unlabeled data from different distributions.
In this study, we focus on a more realistic scenario of Feature Distribution Mismatch SSL (FDMSSL), i.e., the feature distributions of labeled and unlabeled data could be different and the feature distributions of test data could contain multiple distributions. It is generally observed that the performance of classical SSL algorithms [22, 30, 6, 9] degrades substantially under feature distribution mismatch. Specifically, in the early stages of training, classical SSL algorithms typically utilize the current model fitted on labeled data to produce pseudo-labels for unlabeled data. However, when the distribution of unlabeled data deviates, the current model are not applicable to unlabeled data, resulting in massive incorrect pseudo-labels and aggravating confirmation bias [3]. Recently, some works [7, 19] attempt to address FDM-SSL. These approaches assume that the unlabeled feature distribution comes from a single source and they only focus on the labeled distribution, which do not always hold in real tasks. However, due to the unknown test scenarios, it is desirable to develop a method to perform well on labeled, unlabeled and even unseen distributions simultaneously.
In this work, we propose a generalized Self-Supervised Feature Adaptation (SSFA) framework for FDM-SSL which does not need to know the distribution of unlabeled data ahead of time. The core idea of SSFA is to decouple pseudo-label predictions from the current model to address distribution mismatch. SSFA consists of two modules, including the semi-supervised learning module and the feature adaptation module. Inspired by [31] that the main classification task can be indirectly optimized through the auxiliary task, SSFA incorporates an auxiliary self-supervised task into the SSL module to train with the main task. In the feature adaptation module, given the current model primarily fitted on labeled data, SSFA utilizes the self-supervised task to update the feature extractor before making predictions on the unlabeled data. After the feature extractor adapts to the unlabeled distribution, refined features can be used to generate more accurate pseudo-labels to assist SSL.
Furthermore, the standard evaluation protocol of SSL normally assumes that test samples follow the same feature distribution as labeled training data. However, this is too restricted to reflect the diversity of real-world applications, where different tasks may focus on different test distributions. It is strongly desired that the SSL model can perform well across a wide range of test distributions. Therefore, in this work, we propose new evaluation protocols that involve test data from labeled, unlabeled and unseen distributions, allowing for a more comprehensive assessment of SSL performance. Extensive experiments are conducted on two types of feature distribution mismatch SSL tasks, i.e., corruption and style mismatch. The experimental results demonstrate that various pseudo-label-based SSL methods can be directly incorporated into SSFA, yielding consistent performance gains across a wide range of test distributions.
Semi-Supervised Learning (SSL). Existing SSL methods [44, 7, 30, 41, 50, 9, 6] typically combine various commonly used techniques for semi-supervised tasks, such as consistency regularization[21, 43, 32, 28] and pseudo-label [22], to achieve state-of-the-art performance. For example, [6] expands the dataset by interpolating between labeled and unlabeled data. [5, 30] use the confident weak augmented view prediction results to generate pseudo-labels for the corresponding strong augmented view. [41, 50, 9, 44] improve the performance by applying adaptive thresholds instead of fixed thresholds. However, these works focus on traditional semi-supervised learning, which assumes that the labeled data and unlabeled data are sampled from the same distribution. Recently, some works [2, 53] address the issue of different feature distributions for labeled and unlabeled data under certain prerequisites. [2] assumes the test and unlabeled samples are drawn from the same single distribution. [53] studies Semi-Supervised Domain Generalization (SSDG), which requires multi-source partially-labeled training data. To further expand SSL to a more realistic scenario, we introduce a new Feature Distribution Mismatch SSL (FDM-SSL) setting, where the unlabeled data comes from multiple distributions and may differ from the labeled distribution. To solve the FDM-SSL, we propose Self-Supervised Feature Adaptation (SSFA), a generic framework to improve the performance in labeled, unlabeled, and even unseen test distributions.
Unsupervised Domain Adaptation (UDA). UDA aims to transfer knowledge from a source domain with sufficient labeled data to an unlabeled target domain through adaptive learning. The UDA methods can be roughly categorized into two categories, namely metric-based methods [23, 35, 25, 34, 49] and adversarial methods [1, 16, 8, 24, 46, 29, 37, 10, 42]. Metric-based methods focus on measuring domain discrepancy to align the feature distributions between source and unlabeled target domains. Adversarial methods [16, 47, 24] use a domain classifier as a discriminator to enforce the
Table 1: Comparison between different problem settings.
Task	Labeled	Unlabeled	Train setting	Test distribution
traditional SSL	scarce	abundant	pl(x) = pu(x)	pl(x)
UDA	abundant	abundant	pl(x) ≠ pu(x)	pu(x)
TTA	abundant	-	pl(x)	pu(x)
FDM-SSL	scarce	abundant	pl(x) ≠ pu(x)	pl(x), pu(x), punseen(x)
feature extractor to learn domain-invariant features through adversarial training. Another task similar to FDM-SSL is Unsupervised Domain Expansion (UDE) [39, 51, 33, 45], which aims to maintain the model's performance on the labeled domain after adapting to the unlabeled domain. The main differences between UDA, UDE and FDM-SSL are the number of labeled data during training and the distribution of the unlabeled data. In FDM-SSL, the labeled data is scarce, and the distribution of unlabeled data is not limited to a specific domain but rather a mixture of multiple domains. These challenges make UDA and UDE methods unable to be directly applied to FDM-SSL.
Test-Time Adaptation (TTA). If the distribution of test data differs from that of training data, the model's performance may suffer performance degradation. TTA methods focus on improving the model's performance with a two-stage process: a model should first adapt to the test samples and then make predictions of them. For example, [31] and [15] introduce an auxiliary self-learning task to adapt the model to test distribution. [38] modifies the scaling and bias parameters of BatchNorm layers based on entropy minimization. [52, 13, 12, 26] try to improve the robustness of test-time adaptation. Unlike most TTA methods which suppose the test samples come from a single distribution, [40] assumes a mixed test distribution that changes continuously and stably. However, these assumptions are still difficult to satisfy in real-life scenarios. In addition, similar to UDA, TTA methods generally require a model trained well on the source domain. In contrast, FDM-SSL setting only requires scarce labeled data and has no restrictions on the distribution of unlabeled data.
For Feature Distribution Mismatch SSL (FDM-SSL) problem, a model observes a labeled set  $D_{l} = \{(x_{i},y_{i})\}_{i = 1}^{N_{l}}$  and an unlabeled set  $D_{u} = \{(u_{j})\}_{j = 1}^{N_{u}}$  with  $N_{u}\gg N_{l}$ , where  $x_{i}$  and  $u_{i}$  are input images and  $y_{i}$  is the label of  $x_{i}$ . The labeled image  $x_{i}$  and unlabeled image  $u_{j}$  are drawn from two different feature distributions  $p_l(x)$  and  $p_u(x)$ , respectively. Note that different from previous works, in FDM-SSL setting the unlabeled samples may come from a mixture of multiple distributions rather than just one distribution: i.e.  $p_u(x) = w_0p_l(x) + \sum_{k = 1}^K w_kp_u^k (x)$ , where  $w_{0}$  and  $w_{k}$  represent the weights of the labeled distribution  $p_l(x)$  and the  $k$ -th unlabeled distribution  $p_u^k (x)$  respectively. The goal of FDM-SSL is to train a model that generalizes well over a large range of varying test data distributions, including labeled, unlabeled and even distributions unseen during training ( $p_{\text{unseen}}(x)$ ). The differences between traditional SSL, UDA, TTA and FDM-SSL, are summarized in Table 1.
The core of FDM-SSL is to enhance the utilization of unlabeled samples in the presence of feature distribution mismatch. On one hand, the shared information, such as patterns and structures, in the unlabeled samples can provide effective cues for model adaptation. On the other hand, these mismatched unlabeled samples can facilitate the learning of a more robust model by exposing it to a wide range of data distributions.
In the FDM-SSL setting, due to the feature distribution shift from  $p_l(x)$  to  $p_u(x)$ , directly applying the SSL model fitted on labeled data to unlabeled data may lead to massive inaccurate predictions on the unlabeled data, thereby aggravating confirmation bias and impairing SSL learning.
This motivates us to propose Self-Supervised Feature Adaptation (SSFA), a unified framework for FDM-SSL, which decouples the pseudo-label predictions from the current model to address distribution mismatch. As illustrated in Figure 1, SSFA consists of two modules: a semi-supervised
learning module and a feature adaptation module. The semi-supervised learning module incorporates a pseudo-label-based SSL learner with a self-supervised auxiliary task that shares a portion of feature extractor parameters with the main classification task. [15] has pointed out that the main task can be indirectly optimized by optimizing the auxiliary task. To this end, the feature adaptation module is designed to update the current model through self-supervised learning on unlabeled data, so as to better match the unlabeled distribution and predict higher-quality pseudo-labels for semi-supervised learning. Therefore, by optimizing the auxiliary self-supervised task individually, the updated model can make more accurate classification decisions in the main task.
Figure 1: The pipeline of SSFA. Let  $x$ ,  $u_w$  and  $u_s$  denote a batch of the labeled data, the weak augmentation and the strong augmentation of unlabeled data respectively,  $\{\cdot\}$  represent the data stream.
4.2 Semi-Supervised Learning Module
In the Semi-Supervised Learning Module, we introduce a self-supervised task as an auxiliary task, which is optimized together with the main task. As shown in Figure 1, the network parameter  $\theta$  comprises three parts:  $\theta_{g}$  for the shared encoder,  $\theta_{c}$  for the main task head, and  $\theta_{s}$  for the auxiliary task head. During training, SSL module optimizes a supervised loss  $\mathcal{L}_x$ , an unsupervised loss  $\mathcal{L}_u$  and a self-supervised auxiliary loss  $\mathcal{L}_{aux}$ , simultaneously. Typically, given a batch of labeled data  $\{(x_b,y_b)\}_{b = 1}^B$  with size  $B$  and a batch of unlabeled data  $\{(u_b)\}_{b = 1}^{\mu B}$  with size  $\mu B$ , where  $\mu$  is the ratio of unlabeled data to labeled data,  $\mathcal{L}_x$  applies standard cross-entropy loss on labeled examples:
$$
\mathcal {L} _ {x} = \frac {1}{B} \sum_ {b = 1} ^ {B} \mathrm {H} \left(x _ {b}, y _ {b}; \theta_ {g}, \theta_ {c}\right), \tag {1}
$$
where  $\mathrm{H}(\cdot, \cdot)$  is the cross-entropy function. Different SSL learners [30, 5] may design different  $\mathcal{L}_u$ . One of the widely used is the pseudo-label loss. In particular, given the pseudo-label  $q_b$  for each unlabeled input  $u_b$ ,  $\mathcal{L}_u$  in traditional SSL can be formulated as:
$$
\mathcal {L} _ {u} = \frac {1}{\mu B} \sum_ {b = 1} ^ {\mu B} \Omega \left(u _ {b}, q _ {b}; \theta_ {g}, \theta_ {c}\right), \tag {2}
$$
where  $\Omega$  is the per-sample supervised loss, e.g., mean-square error [6] and cross-entropy loss [30].
Furthermore, we introduce a self-supervised learning task to the SSL module for joint training to optimize the parameter of the auxiliary task head  $\theta_{s}$ . To learn from labeled and unlabeled data distributions simultaneously, it is necessary to use all samples from both distributions for training. Therefore,  $\mathcal{L}_{aux}$  can be formulated as:
$$
\mathcal {L} _ {a u x} = \frac {1}{(2 \mu + 1) B} \left(\sum_ {b = 1} ^ {B} \ell_ {s} \left(x _ {b}; \theta_ {g}, \theta_ {s}\right) + \sum_ {b = 1} ^ {\mu B} \ell_ {s} \left(\mathcal {A} _ {w} \left(u _ {b}\right); \theta_ {g}, \theta_ {s}\right) + \sum_ {b = 1} ^ {\mu B} \ell_ {s} \left(\mathcal {A} _ {s} \left(u _ {b}\right); \theta_ {g}, \theta_ {s}\right)\right), \tag {3}
$$
where  $\ell_s$  denotes the self-supervised loss function,  $\mathcal{A}_w$  and  $\mathcal{A}_s$  denote the weak and strong augmentation functions respectively. In order to maintain consistent optimization between the main task
and auxiliary task, it is necessary to optimize both tasks jointly. So we add  $\mathcal{L}_{aux}$  to semi-supervised learning module and the final object is:
$$
\mathcal {L} _ {\mathrm {S S F A}} = \mathcal {L} _ {x} + \lambda_ {u} \mathcal {L} _ {u} + \lambda_ {a} \mathcal {L} _ {a u x}, \tag {4}
$$
where  $\lambda_{u}$  and  $\lambda_{a}$  are hyper-parameters denoting the relative weights of  $\mathcal{L}_u$  and  $\mathcal{L}_{aux}$  respectively.
4.3 Feature Adaptation Module
Some classic SSL approaches [30, 6, 5] directly use the outputs of the current classifier as pseudo-labels. In particular, [30] applies weak and strong augmentations to unlabeled samples and generates pseudo-labels using the model's predictions on weakly augmented unlabeled samples, i.e.,  $q_{b} = p_{m}(y|\mathcal{A}_{w}(u_{b});\theta_{g},\theta_{c})$ , where  $p_{m}(y|u;\theta)$  denotes the predicted class distribution of the model  $\theta$  given unlabeled image  $u$ . However, as the classification model  $(\theta_{g},\theta_{c})$  is mainly fitted to the labeled distribution, the prediction on  $\mathcal{A}_w(u_b)$  is usually inaccurate under distribution mismatch between the labeled and unlabeled samples.
To alleviate this problem, we design a Feature Adaptation Module to adapt the model to the unlabeled data distribution before making predictions, thereby producing more reliable pseudo-labels for optimizing the unsupervised loss  $\mathcal{L}_u$ . More specifically, before making pseudo-label predictions, we firstly fine-tune the shared feature extractor  $\theta_g$  by minimizing the self-supervised auxiliary task loss on unlabeled samples:
$$
\mathcal {L} _ {a p t} = \frac {1}{\mu B} \sum_ {b = 1} ^ {\mu B} \ell_ {s} \left(\mathcal {A} _ {w} \left(u _ {b}\right); \theta_ {g}, \theta_ {s}\right). \tag {5}
$$
Here  $\theta_{g}$  is updated to  $\theta_g^{\prime} = \arg \min \mathcal{L}_{apt}$ . Notably, since excessive adaptation may lead to deviation in the optimization direction and largely increase calculation costs, we only perform one-step optimization in the adaptation stage.
After self-supervised adaptation, we use  $\left(\theta_g',\theta_c\right)$  to generate the updated prediction, which can be denoted as:
$$
q _ {b} ^ {\prime} = p _ {m} (y \mid \mathcal {A} _ {w} (u _ {b}); \theta_ {g} ^ {\prime}, \theta_ {c}). \tag {6}
$$
We convert  $q_{b}^{\prime}$  to the hard "one-hot" label  $\hat{q}_{b}^{\prime}$  and denote  $\hat{q}_{b}^{\prime}$  as the pseudo-label for the corresponding strongly augmented unlabeled sample  $\mathcal{A}_w(u_b)$ . After that,  $\theta_q^\prime$  will be discarded without influencing other model parts during training. In the end,  $\mathcal{L}_u$  in the SSL module (Equation 4) can then be computed with cross-entropy loss as:
$$
\mathcal {L} _ {u} = \frac {1}{\mu B} \sum_ {b = 1} ^ {\mu B} \mathbb {I} \left(\max  \left(q _ {b} ^ {\prime}\right) > \tau\right) \mathrm {H} \left(\mathcal {A} _ {s} \left(u _ {b}\right), \hat {q} _ {b} ^ {\prime}; \theta_ {g}, \theta_ {c}\right). \tag {7}
$$
4.4 Theoretical Insights
In our framework, we jointly train the model with a self-supervised task with loss function  $\ell_s$ , and a main task with loss function  $\ell_m$ . Let  $h \in \mathcal{H}$  be a feasible hypothesis and  $D_u = \{(u_i, y_i^u)\}_{i=1}^N$  be the unlabeled dataset. Note that the ground-truth label  $y_i^u$  of  $u_i$  is actually unavailable, but just used for analysis.
Lemma 1 ([31]) Assume that for all  $x, y$ ,  $\ell_m(x, y; h)$  is differentiable, convex and  $\beta$ -smooth in  $h$  and both  $\|\nabla \ell_m(x, y; h)\|$ ,  $\|\nabla \ell_s(x; h)\| \leq G$  for all  $h \in \mathcal{H}$ . With a fixed learning rate  $\eta = \frac{\epsilon}{\beta G^2}$ , for every  $x, y$  such that  $\langle \nabla \ell_m(x, y; h), \nabla \ell_s(x; h) \rangle > \epsilon$ , we have
$$
\ell_ {m} (x, y; h ^ {\prime}) <   \ell_ {m} (x, y; h), \tag {8}
$$
where  $h^\prime$  is the updated hypothesis, namely  $h^\prime = h - \eta \nabla \ell_s(x,y;h)$
Let  $\hat{E}_m(h, D_u) = \frac{1}{K} \sum_{i=1}^{K} \ell_m(x_u^i, y_u^i; h)$  denote the empirical risk of the main task on  $D_u$ , and  $\hat{E}_s(h, D_u) = \frac{1}{K} \sum_{i=1}^{K} \ell_s(x_u^i; h)$  denote the empirical risk of the self-supervised task on  $D_u$ . Under the premise of Lemma 1 and the following sufficient condition, with a fixed learning rate  $\eta$ :
$$
\left\langle \nabla \hat {E} _ {m} (h, D _ {u}), \nabla \hat {E} _ {s} (h, D _ {u}) \right\rangle > \epsilon , \tag {9}
$$
we can get that:
$$
\hat {E} _ {m} \left(h ^ {\prime}, D _ {u}\right) <   \hat {E} _ {m} \left(h, D _ {u}\right), \tag {10}
$$
where  $h^\prime$  is the updated hypothesis, namely  $h^\prime = h - \eta \nabla \hat{E}_s(h,D_u)$ .
Therefore, in the smooth and convex case, the empirical risk of the main task  $\hat{E}_m$  can theoretically tend to 0 by optimizing the empirical risk of self-supervised task  $\hat{E}_s$ . Thus, in our feature adaptation module, by optimizing the self-supervised loss for unlabeled samples, we can indirectly optimize the main loss, thereby mitigating confirmation bias and making full use of the unlabeled samples. As above analysis, the gradient correlation plays a determining factor in the success of optimizing  $\hat{E}_m$  through  $\hat{E}_s$  in the smooth and convex case. For non-convex deep loss functions, we provide empirical evidence to show that our theoretical insights also hold. Figure 2 plots the correlation between the gradient inner product (of the main and auxiliary tasks) and the performance improvement of the model on the
Figure 2: Scatter plot of the gradient inner product between the two tasks, and the improvement from SSFA. We transform the x-axis with  $\log (x) + 1$  for clarity.
test set, where each point in the figure represents the average result of a set of test samples. In Figure 2, we observe that there is a positive correlation between the gradient inner product and model performance improvement for non-convex loss functions on the deep learning model. This phenomenon is consistent with the theoretical conclusion, that is, a stronger gradient correlation indicates a higher performance improvement.
In this section, we evaluate our proposed SSFA framework on various datasets, where labeled and unlabeled data are sampled from different distributions. To this end, we consider two different distribution mismatch scenarios: image corruption and style change, in the following experiments. More experimental details and results are provided in the Appendix.
Datasets. For the image corruption experiments, we create the labeled domain by sampling images from CIFAR100 [20], and the unlabeled domain by sampling images from a mixed dataset consisting of CIFAR100 and CIFAR100 with Corruptions (CIFAR100-C). The CIFAR100-C is constructed by applying various corruptions to the original images in CIFAR100, following the methodology used for ImageNet-C [18]. We use ten types of corruption as training corruptions while reserving the other five corruptions as unseen corruptions for testing. The proportion of unlabeled samples from CIFAR100-C was controlled by the hyper-parameter ratio. For the style change experiments, we use OFFICE-31 [27] and OFFICE-HOME [36] benchmarks. Specifically, we designate one domain in each dataset as the labeled domain, while the unlabeled domain comprises either another domain or a mixture of multiple domains.
Implementation Details. In corruption experiments, we use WRN-28-8 [48] as the backbone except for [6] where WRN-28-2 is used to prevent training collapse. In style experiments, we use ResNet-50 [17] pre-trained on ImageNet [14] as the backbone for OFFICE-31 and OFFICE-HOME. To ensure fairness, we use the same hyperparameters for different methods employed in our experiments.