SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Saurabh Garg*
Carnegie Mellon University  
sgarg2@andrew.cmu.edu
Amrith Setlur*
Carnegie Mellon University assetlur@andrew.cmu.edu
Zachary C. Lipton
Carnegie Mellon University  
zlipton@andrew.cmu.edu
Sivaraman Balakrishnan
Carnegie Mellon University sbalakri@andrew.cmu.edu
Virginia Smith
Carnegie Mellon University smithv@andrew.cmu.edu
Aditi Raghunathan
Carnegie Mellon University aditirag@andrew.cmu.edu
Self-training and contrastive learning have emerged as leading techniques for incorporating unlabeled data, both under distribution shift (unsupervised domain adaptation) and when it is absent (semi-supervised learning). However, despite the popularity and compatibility of these techniques, their efficacy in combination remains surprisingly unexplored. In this paper, we first undertake a systematic empirical investigation of this combination, finding (i) that in domain adaptation settings, self-training and contrastive learning offer significant complementary gains; and (ii) that in semi-supervised learning settings, surprisingly, the benefits are not synergistic. Across eight distribution shift datasets (e.g., BREEDs, WILDS), we demonstrate that the combined method obtains  $3 - 8\%$  higher accuracy than either approach independently. Finally, we theoretically analyze these techniques in a simplified model of distribution shift demonstrating scenarios under which the features produced by contrastive learning can yield a good initialization for self-training to further amplify gains and achieve optimal performance, even when either method alone would fail.
Even under natural, non-adversarial distribution shifts, the performance of machine learning models often drops [65, 83, 47, 29]. While we might hope to retrain our models on labeled samples from the new distribution, this option is often unavailable due to the expense or impracticality of collecting new labels. Consequently, researchers have investigated the Unsupervised Domain Adaptation (UDA) setting. Here, given labeled source data and unlabeled out-of-distribution (OOD) target data, the goal is to produce a classifier that performs well on the target. Because UDA is generally underspecified [7], researchers have focused on two main paths: (i) domain adaptation papers that explore heuristics for incorporating the unlabeled target data, relying on benchmark datasets ostensibly representative of "real-world shifts" to adjudicate progress [72, 85, 47]; and (ii) theoretically motivated papers that explore structural assumptions under which UDA problems are well posed [77, 74]. This work engages with the former focusing on two popular methods: self-training and contrastive pretraining.
Self-training [75, 52, 79, 90, 86] and contrastive pretraining [13, 16, 93] were both proposed, initially, for traditional Semi-Supervised Learning (SSL) problems, where the labeled and unlabeled data are drawn from the same distribution. Here the central challenge is statistical: to exploit the unlabeled data to learn a better predictor than one would get by training on the (small) labeled data alone. More
Figure 1: Self-training over Contrastive learning (STOC) improves over Contrastive Learning (CL) under distribution shift. (a) We observe that in SSL settings, where labeled and unlabeled data are drawn from the same distribution, STOC offers negligible improvements over CL. In contrast, in UDA settings where there is distribution shift between labeled and unlabeled data, STOC offers gains over CL. Results aggregated across 8 benchmarks. Results on individual data in Table 1 and 2. (b) 2-D illustration of our simplified distribution setup, depicting decision boundaries learned by ERM and CL and how Self-Training (ST) updates those. ①, ②, and ③ summarize our theoretical results in Sec. 4.
recently, these methods have emerged as favored empirical approaches for UDA, demonstrating efficacy on many popular benchmarks [70, 30, 12, 76]. In self-training, one first learns a predictor using source labeled data. The predictor then produces pseudolabels for the unlabeled target data, and a new predictor is trained on the pseudolabeled data. Contrastive pretraining learns representations from unlabeled data by enforcing invariance to specified augmentations. These representations are subsequently used to learn a classifier. In UDA, the representations are trained on the union of the source and target data. In attempts to explain their strong empirical performance, several researchers have attempted analyses under various assumptions on the data, task, and inductive biases of the function class[87, 39, 73, 76, 12, 40, 38, 11]. Despite the strong strong results, there has been surprisingly little work (both empirically and theoretically) exploring when either might be expected to perform best and whether the benefits might be complementary.
In this paper, we investigate the complementary benefits of self-training and contrastive pretraining. Interestingly, we find that the combination yields significant gains in UDA despite producing negligible gains in SSL. In experiments across eight distribution shift benchmarks (e.g. BREEDs [72], FMoW [47], Visda [63]), we observe that re-using unlabeled data for self-training (with Fix-Match [79]) after learning contrastive representations (with SwAV [13]), yields  $>5\%$  average improvement on OOD accuracy in UDA as compared to  $< 0.8\%$  average improvement in SSL (Fig. 1).
Next, we address the question why the combination of self-training and contrastive learning proves synergistic in distribution shift scenarios. To facilitate our analysis, we consider a simplified distribution shift setting that includes two types of features: (i) invariant features that perfectly predict the label; and (ii) domain-dependent features that are predictive of the label in just source. Our theoretical analysis reveals that self-training can achieve optimal target performance but requires a "good" enough classifier to start with. We observe that source-only ERM fails to provide a "good" initialization. On the other hand, contrastive pretraining on unlabeled data performs better than ERM but is still sub-optimal. This implies that contrastive pretraining ends up decreasing reliance on domain-dependent features (as compared to ERM) but doesn't completely eliminate them. Nevertheless, contrastive pretraining does provide a "good" initialization for self-training, i.e., "good" initial pseudolabels on the target unlabeled data. As a result, self-training on top of contrastive learned features effectively unlearns the reliance on domain-dependent features and generalizes perfectly OOD. In contrast, for SSL settings (i.e., in distribution), our analysis highlights that contrastive pretraining already acquires sufficient predictive features such that linear probing with (a small amount of) labeled data picks up those features and attains near-optimal ID generalization.
Finally, we connect our theoretical understanding of "good" representations from contrastive learning and improved linear transferability from self-training back to observed empirical gains. We linearly probe representations (fix representations and train only the linear head) learned by contrastive pretraining vs. no pretraining and find: (i) contrastive pretraining substantially improves the ceiling on
the target accuracy (performance of optimal linear probe) compared to ERM; (ii) self-training mainly improves linear transfer, i.e. OOD performance for the linear probe trained with source labeled data.
2 Setup and Preliminaries
Task. Our goal is to learn a predictor that maps inputs  $x \in \mathcal{X} \subseteq \mathbb{R}^d$  to outputs  $y \in \mathcal{Y}$ . We parameterize predictors  $f = h \circ \Phi : \mathbb{R}^d \mapsto \mathcal{Y}$ , where  $\Phi : \mathbb{R}^d \mapsto \mathbb{R}^k$  is a feature map and  $h \in \mathbb{R}^k$  is a classifier that maps the representation to the final scores or logits. Let  $\mathrm{P_S}, \mathrm{P_T}$  be the source and target joint probability measures over  $\mathcal{X} \times \mathcal{Y}$  with  $p_{\mathrm{S}}$  and  $p_{\mathrm{T}}$  as the corresponding probability density (or mass) functions. The distribution over unlabeled samples from both the union of source and target is denoted as  $\mathrm{P_U} = (1/2) \cdot \mathrm{P_S}(x) + (1/2) \cdot \mathrm{P_T}(x)$ .
We study two particular scenarios: (i) Unsupervised Domain Adaptation (UDA); and (ii) Semi-Supervised Learning (SSL). In UDA, we assume that the source and target distributions have the same label marginals  $\mathrm{P}_{\mathrm{S}}(y) = \mathrm{P}_{\mathrm{T}}(y)$  (i.e., no label proportion shift) and the same Bayes optimal predictor, i.e.,  $\arg \max_y p_{\mathrm{S}}(y \mid x) = \arg \max_y p_{\mathrm{T}}(y \mid x)$ . We are given labeled samples from the source, and unlabeled pool from the target. In contrast in SSL, there is no distribution shift, i.e.,  $\mathrm{P}_{\mathrm{S}} = \mathrm{P}_{\mathrm{T}} = \mathrm{P}_{\cup}$ . Here, we are given a small number of labeled examples and a comparatively large amount of unlabeled examples, both drawn from the same distribution, which we denote as  $\mathrm{P}_{\mathrm{T}}$ .
Unlabeled data is typically much cheaper to obtain, and our goal in both these settings is to leverage this along with labeled data to achieve good performance on the target distribution. In the UDA scenario, the challenge lies in generalizing out-of-distribution, while in SSL, the challenge is to generalize in-distribution despite the paucity of labeled examples. A predictor  $f$  is evaluated on distribution  $\mathrm{P}$  via its accuracy, i.e.,  $A(f,\mathrm{P}) = \mathbb{E}_{\mathrm{P}}(\arg \max f(x) = y)$ .
Methods. We now introduce the algorithms used for learning from labeled and unlabeled data.
1. Source-only  $ERM$  (ERM): A standard approach is to simply perform supervised learning on the labeled data by minimizing the empirical risk  $\sum_{i=1}^{n} \ell(h \circ \Phi(x), y)$ , for some classification loss  $\ell: \mathbb{R} \times \mathcal{Y} \mapsto \mathbb{R}$  (e.g., softmax cross-entropy) and labeled points  $\{(x_i, y_i)\}_{i=1}^n$ .
2. Contrastive Learning (CL): We first use the unlabeled data to learn a feature extractor. In particular, the objective is to learn a feature extractor  $\Phi_{\mathrm{cl}}$  that maps augmentations (for e.g. crops or rotations) of the same input close to each other and far from augmentations of random other inputs [13, 16, 93]. Once we have  $\Phi_{\mathrm{cl}}$ , we learn a linear classifier  $h$  on top to minimize a classification loss on the labeled source data. We could either keep  $\Phi_{\mathrm{cl}}$  fixed or propagate gradients through.
When clear from context, we also use CL to refer to just the contrastively pretrained backbone without training for downstream classification.
3. Self-training (ST): This is a two-stage procedure, where the first stage performs source-only ERM by just looking at source-labeled data. In the second stage, we iteratively apply the current classifier on the unlabeled data to generate "pseudo-labels" and then update the classifier by minimizing a classification loss on the pseudolabeled data [52].
3 Self-Training Improves Contrastive Pretraining Under Distribution Shift
Self-Training Over Contrastive Learning (STOC). Finally, rather than starting with a source-only ERM classifier, we propose to initialize self-training with a CL classifier, that was pretrained on unlabeled source and target data. ST uses that same unlabeled data again for pseudolabeling. As we demonstrate experimentally and theoretically, this combination of methods improves substantially over each independently.
Datasets. For both UDA and SSL, we conduct experiments across eight benchmark datasets: four BREEDs datasets [72]—Entity13, Entity30, Nonliving26, Living17; FMoW [47, 18] from WILDS benchmark; Officehome [85]; Visda [64, 63]; and CIFAR-10 [48]. Each of these datasets consists of domains, enabling us to construct source-target pairs (e.g., CIFAR10, we consider CIFAR10  $\rightarrow$  CINIC shift [22]). In the UDA setup, we adopt the source and target domains standard to previous studies (details in App. C.2). Because the SSL setting lacks distribution shift, we do not need to worry about domain designations and default to using source alone. To simulate limited supervision in SSL, we sub-sample the original labeled training set to  $10\%$ .
Table 1: Results in the UDA setup. We report accuracy on target (OOD) data from which we only observe unlabeled examples during training. For benchmarks with multiple target distributions (e.g., OH, Visda), we report avg accuracy on those targets. Results with source performance, individual target performance, and standard deviation numbers are in App. C.4.
Method	Living17	Nonliv26	Entity13	Entity30	FMoW (2 tgts)	Visda (2 tgts)	OH (3 tgts)	CIFAR→ CINIC	Avg
ERM	60.31	45.54	68.32	55.75	56.50	20.91	9.51	74.33	48.90
ST	71.29	56.79	77.93	66.37	56.79	38.03	10.47	78.19	56.98
CL	74.14	57.02	76.58	66.01	61.78	63.49	22.63	77.51	62.39
STOC (ours)	82.22	62.23	81.84	72.00	65.25	70.08	27.12	79.94	67.59
Table 2: Results in the SSL setup. We report accuracy on hold-out ID data. Recall that SSL uses labeled and unlabeled data from the same distribution during training. Refer to App. C.5 for ERM and ST.
Method	Living17	Nonliv26	Entity13	Entity30	FMoW	Visda	OH	CIFAR	Avg
CL	91.15	84.58	90.73	85.47	43.05	97.67	49.73	91.78	79.27
STOC (ours)	92.00	85.95	91.27	86.14	44.43	97.70	49.95	93.06	80.06
Experimental Setup and Protocols. SwAV [13] is the specific algorithm that we use for contrastive pretraining. In all UDA settings, unless otherwise specified, we pool all the (unlabeled) data from the source and target to perform SwAV. For self-training, we apply FixMatch [79], where the loss on source labeled data and on pseudolabeled target data are minimized simultaneously. For both methods, we fix the algorithm-specific hyperparameters to the original recommendations. For SSL settings, we perform SwAV and FixMatch on in-distribution unlabeled data. We experiment with Resnet18, Resnet50 [42] trained from scratch (i.e. random initialization). We do not consider off-the-shelf pretrained models (e.g., onImagenet [68]) to avoid confounding our conclusions about contrastive pretraining. However, we note that our results on most datasets tend to be comparable to and sometimes exceed those obtained with ImageNet-pretrained models. For source-only ERM, as with other methods (FixMatch, SwAV), we default to using strong augmentation techniques: random horizontal flips, random crops, augmentation with Cutout [24], and RandAugment [21]. Moreover, unless otherwise specified, we default to full finetuning with source-only ERM, both from scratch and after contrastive pretraining, and for ST with FixMatch. For UDA, given that the setup precludes access to labeled data from the target distribution, we use source hold-out performance to pick the best hyperparameters. During pretraining, early stopping is done according to lower values of pretraining loss. For more details on datasets, model architectures, and experimental protocols, see App. C.
Results on UDA setup. Both ST and CL individually improve over ERM across all datasets, with CL significantly performing better than ST on 5 out of 8 benchmarks (see Table 1). Even on datasets where ST is better than CL, their performance remains close. Combining ST and CL with STOC shows an  $3 - 8\%$  improvement over the best alternative, yielding an absolute improvement in average accuracy of  $5.2\%$ .
Note that by default, we train with CL on the combined unlabeled data from source and target. However, to better understand the significance of unlabeled target data in contrastive pretraining, we perform an ablation where the CL model was trained solely on unlabeled source data (refer to this as CL (source only); see App. C.4). We observe that ST on top of CL (source only) improves over ST (from scratch). However, the average performance of ST over CL (source only) is similar to that of standalone CL, maintaining an approximate  $6\%$  performance gap observed between CL and ST. This brings two key insights to the fore: (i) the observed benefit is not merely a result of the contrastive pretraining objective alone, but specifically CL with unlabeled target data helps; and (ii) both CL and ST leverage using target unlabeled data in a complementary nature.
Results on SSL setup. While CL improves over ST (as in UDA), unlike UDA, STOC doesn't offer any significant improvements over CL (see Table 2; ERM and ST results (refer to App. C.5). We conduct ablation studies with varying proportions of labeled data used for SSL, illustrating that there's considerable potential for improvement (see App. C.5). These findings highlight that the complementary nature of STOC over CL and ST individually is an artifact of distribution shift.
4 Theoretical Analysis and Intuitions
Our results on real-world datasets suggest that although self-training may offer little to no improvement over contrastive pretraining for in-distribution (i.e., SSL) settings, it leads to substantial improvements when facing distribution shifts in UDA (Sec. 3). Why do these methods offer complementary gains, but only under distribution shifts? In this section, we seek to answer this question by first replicating all the empirical trends of interest in a simple data distribution with an intuitive story (Sec. 4.1). In this toy model, we formally characterize the gains afforded by contrastive pretraining and self-training both individually (Secs. 4.2, 4.3) and when used together (Sec. 4.4).
Data distribution We consider binary classification and model the inputs as consisting of two kinds of features:  $x = [x_{\mathrm{in}}, x_{\mathrm{sp}}]$ , where  $x_{\mathrm{in}} \in \mathbb{R}^{d_{\mathrm{in}}}$  is the invariant feature that is predictive of the label across both source  $\mathrm{P}_{\mathrm{S}}$  and target  $\mathrm{P}_{\mathrm{T}}$  and  $x_{\mathrm{sp}} \in \mathbb{R}^{d_{\mathrm{sp}}}$  is the spurious feature that is correlated with the label  $y$  only on the source domain  $\mathrm{P}_{\mathrm{S}}$  but uncorrelated with label  $y$  in  $\mathrm{P}_{\mathrm{T}}$ . Formally, we sample  $y \sim \mathrm{Unif}\{-1, 1\}$  and generate inputs  $x$  conditioned on  $y$  as follows:
$$
\mathrm {P} _ {\mathrm {S}}: x _ {\mathrm {i n}} \sim \mathcal {N} (\gamma \cdot \mathrm {y} w ^ {\star}, \Sigma_ {\mathrm {i n}}) x _ {\mathrm {s p}} = \mathrm {y} \mathbf {1} _ {d _ {\mathrm {s p}}}
$$
$$
\mathrm {P} _ {\top}: x _ {\text {i n}} \sim \mathcal {N} (\gamma \cdot \mathrm {y} w ^ {\star}, \Sigma_ {\text {i n}}) x _ {\text {s p}} \sim \mathcal {N} (\mathbf {0}, \Sigma_ {\text {s p}}), \tag {1}
$$
where  $\gamma$  is the margin afforded by the invariant feature $^2$ . We set the covariance of the invariant features  $\Sigma_{\mathrm{in}} = \sigma_{\mathrm{in}}^2 \cdot (\mathbf{I}_{d_{\mathrm{in}}} - w^\star w^{\top})$ . This makes the variance along the latent predictive direction  $w^\star$  to be zero. Note that the spurious feature is also completely predictive of the label in the source data. In fact, when  $d_{\mathrm{sp}}$  is sufficiently large,  $x_{\mathrm{sp}}$  is more predictive (than  $x_{\mathrm{in}}$ ) of y in the source. In the target,  $x_{\mathrm{sp}}$  is distributed as a Gaussian with  $\Sigma_{\mathrm{sp}} = \sigma_{\mathrm{sp}}^2 \mathbf{I}_{d_{\mathrm{sp}}}$ . We use  $w_{\mathrm{in}} = [w^\star, 0, \dots, 0]^\top$  to refer to the invariant direction/feature, and  $w_{\mathrm{sp}} = [0, \dots, 0, \mathbf{1}_{d_{\mathrm{sp}}} / \sqrt{d_{\mathrm{sp}}}]^\top$  for the spurious direction.
Data for UDA vs. SSL For convenience, we assume access to infinite unlabeled data and replace their empirical quantities with population counterparts. For SSL, we sample both finite labeled and infinite unlabeled data from the same distribution  $\mathrm{P_T}$ , where spurious features are absent (to exclude easy-to-generalize features). For UDA, we assume infinite labeled data from  $\mathrm{P_S}$  and infinite unlabeled from  $\mathrm{P_T}$ . Importantly, note that due to distribution shift, population access of  $\mathrm{P_S}$  doesn't trivialize the problem as "ERM" on infinite labeled source data does not achieve optimal performance on target.
Methods and objectives Recall from Section 2 that we learn linear classifiers  $h$  over feature extractor  $\Phi$ . For our toy setup, we consider linear feature extractors i.e.  $\Phi$  is a matrix in  $\mathbb{R}^{d \times k}$  and the prediction is given by  $\operatorname{sgn}(h^\top \Phi x)$ . We use the exponential loss  $\ell(f(x), y) = \exp(-y f(x))$ .
Self-training. ST performs ERM in the first stage using labeled data from the source, and then subsequently updates the head  $h$  by iteratively generating pseudolabels on the unlabeled target:
$$
\mathcal {L} _ {\mathrm {s t}} (h; \Phi) := \mathbb {E} _ {\mathrm {P} _ {\mathrm {T}} (x)} \ell \left(h ^ {\top} \Phi x, \operatorname {s g n} \left(h ^ {\top} \Phi (x)\right)\right)
$$
$$
\text {U p d a t e :} h ^ {t + 1} = \frac {h ^ {t} - \eta \nabla_ {h} \mathcal {L} _ {\mathrm {s t}} \left(h ^ {t} ; \Phi\right)}{\| h ^ {t} - \eta \nabla_ {h} \mathcal {L} _ {\mathrm {s t}} \left(h ^ {t} ; \Phi\right) \| _ {2}} \tag {2}
$$
For ERM and ST, we train both  $h$  and  $\Phi$  (equivalent to  $\Phi$  being identity and training a linear head).
Contrastive pretraining. We obtain  $\Phi_{\mathrm{cl}} \coloneqq \arg \min_{\Phi} \mathcal{L}_{\mathrm{cl}}(\Phi)$  by minimizing the Barlow Twins objective [93], which prior works have shown is also equivalent to spectral contrastive and non-contrastive objectives [33, 11]. Given probability distribution  $\mathrm{P_A}(a \mid x)$  for input  $x$ , and marginal  $\mathrm{P_A}$ , we consider a constrained form of Barlow Twins in (3) which enforces features of "positive pairs"  $a_1, a_2$  to be close while ensuring feature diversity. We assume a strict regularization ( $\rho = 0$ ) for the theory arguments in the rest of the paper, and in App. D.2 we prove that all our claims hold for small  $\rho$  as well. For augmentations, we scale the magnitude of each co-ordinate uniformly by an independent amount, i.e.,  $a \sim \mathrm{P_A}(\cdot \mid x) = \mathbf{c} \odot x$ , where  $\mathbf{c} \sim \mathrm{Unif}[0,1]^d$ . We try to mirror practical settings where the augmentations are fairly "generic", not encoding information about which features are invariant or spurious, and hence perturb all features symmetrically.
$$
\mathcal {L} _ {\mathrm {c l}} (\Phi) := \mathbb {E} _ {x \sim \mathrm {P} _ {\mathrm {U}}} \mathbb {E} _ {a _ {1}, a _ {2} \sim \mathrm {P} _ {\mathrm {A}} (\cdot | x)} \| \Phi (a _ {1}) - \Phi (a _ {2}) \| _ {2} ^ {2}
$$
$$
\text {s . t .} \left\| \mathbb {E} _ {a \sim \mathrm {P} _ {\mathrm {A}}} \left[ \Phi (a) \Phi (a) ^ {\top} \right] - \mathbf {I} _ {k} \right\| _ {F} ^ {2} \leqslant \rho \tag {3}
$$
See App. D.1 for similarities and differences of our setup with prior works.
Figure 2: Our simplified model of shift captures real-world trends and theoretical behaviors: (a) Target (OOD) accuracy separation in the UDA setup (for problem parameters in Example 1). (b) Comparison of the benefits of STOC (ST over CL) over just CL in UDA and SSL settings, done across training iterations for contrastive pretraining. (c) Comparison between different methods in UDA setting, as we vary problem parameters  $\gamma$  and  $\sigma_{\mathrm{sp}}$ , connecting our theory results in Sec. 4.
Keeping the  $\Phi_{\mathrm{cl}}$  fixed, we then learn a linear classifier  $h_{\mathrm{cl}}$  over  $\Phi_{\mathrm{cl}}$  to minimize the exponential loss on labeled source data (refer to as linear probing). For STOC, keeping the  $\Phi_{\mathrm{cl}}$  fixed and initializing the linear head with the CL linear probe (instead of source only ERM), we perform ST with (2).