Title: Complementary Benefits of Contrastive Learning and Self-Training Under Distribution Shift
Venue: NeurIPS 2023
Accepted: yes
Keywords: contrastive learning; self training; distribution shift; semi supervised learning; unsupervised domain adaptation
Abstract: 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.

Complementary Benefits of Contrastive Learning and Self-Training Under Distribution Shift

Abstract

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.

1 Introduction

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

Zachary C. Lipton

Aditi Raghunathan

ERM

ST

CL STOC

65

70

75

80

ID Accuracy

67.8

72.8

79.2

80.1

Aggregate SSL results

(no dist.

shift)

ERM

ST

CL STOC

45

50

55

60

65

70

48.9

57.0

62.4

67.6

Aggregate UDA results

(w dist.

shift)

(a)

STOC amplifies CL

gains under

distribution shift

3

Spurious

Feature

Invariant

Feature

CL

ERM

STOC

ST worsens

dependency on spurious

feature after ERM

ST

Source

Labeled

Target

Unlabeled

Augmentations

CL offers a good

initialization for STOC

2

1

(b)

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 FixMatch [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 domaindependent 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 P X Ď R d to outputs y P Y . We parameterize predictors f ' h ˝ Φ : R d ÞÑ Y , where Φ : R d ÞÑ R k is a feature map and h P R k is a classifier that maps the representation to the final scores or logits. Let P S , P T be the source and target joint probability measures over X ˆ Y with p S and p 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 P U ' p 1 { 2 q ¨ P S p x q ` p 1 { 2 q ¨ P T p x q .

We study two particular scenarios: (i) Unsupervised Domain Adaptation (UDA); and (ii) SemiSupervised Learning (SSL). In UDA, we assume that the source and target distributions have the same label marginals P S p y q ' P T p y q ( i.e ., no label proportion shift) and the same Bayes optimal predictor, i.e ., arg max y p S p y | x q ' arg max y p T p y | x q . 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 ., P S ' P T ' P U . 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 P 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 P via its accuracy, i.e ., A p f, P q ' E P p arg max f p x q ' y q .

Methods. We now introduce the algorithms used for learning from labeled and unlabeled data.

Source-only ERM (ERM) : A standard approach is to simply perform supervised learning on the labeled data by minimizing the empirical risk ř n i ' 1 ℓ p h ˝ Φ p x q , y q , for some classification loss ℓ : R ˆ Y ÞÑ R ( e.g ., softmax cross-entropy) and labeled points tp x i , y i qu n i ' 1 .

Contrastive Learning (CL) : We first use the unlabeled data to learn a feature extractor. In particular, the objective is to learn a feature extractor Φ 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 Φ cl , we learn a linear classifier h on top to minimize a classification loss on the labeled source data. We could either keep Φ 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.

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 Ñ 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%.

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 offthe-shelf pretrained models ( e.g ., on Imagenet [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 ' r x in , x sp s , where x in P R d in is the invariant feature that is predictive of the label across both source P S and target P T and x sp P R d sp is the spurious feature that is correlated with the label y only on the source domain P S but uncorrelated with label y in P T . Formally, we sample y ' Unif t´ 1 , 1 u and generate inputs x conditioned on y as follows:

where γ is the margin afforded by the invariant feature 2 . We set the covariance of the invariant features Σ in ' σ 2 in ¨ p I d in ´ w ‹ w ‹J q . This makes the variance along the latent predictive direction w ‹ to be zero. Note that the spurious feature is also completely predictive of the label in the source data. In fact, when d sp is sufficiently large, x sp is more predictive (than x in ) of y in the source. In the target, x sp is distributed as a Gaussian with Σ sp ' σ 2 sp I d sp . We use w in 'r w ‹ , 0 , ..., 0 s J to refer to the invariant direction/feature, and w sp ' r 0 , ..., 0 , 1 d sp { ? d sp s J 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 P T , where spurious features are absent (to exclude easy-to-generalize features). For UDA, we assume infinite labeled data from P S and infinite unlabeled from P T . Importantly, note that due to distribution shift, population access of 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 Φ . For our toy setup, we consider linear feature extractors i.e. Φ is a matrix in R d ˆ k and the prediction is given by sgn p h J Φ x q . We use the exponential loss ℓ p f p x q , y q ' exp p´ yf p x qq .

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:

For ERM and ST, we train both h and Φ (equivalent to Φ being identity and training a linear head).

Contrastive pretraining. We obtain Φ cl : ' arg min Φ L cl p Φ q by minimizing the Barlow Twins objective [93], which prior works have shown is also equivalent to spectral contrastive and noncontrastive objectives [33, 11]. Given probability distribution P A p a | x q for input x , and marginal 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 p ρ ' 0 q for the theory arguments in the rest of the paper, and in App. D.2 we prove that all our claims hold for small ρ as well. For augmentations, we scale the magnitude of each co-ordinate uniformly by an independent amount, i.e., a ' P A p¨ | x q ' c d x , where c ' Unif r 0 , 1 s 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.

ERM

ST

CL

STOC

60

80

100

Target accuracy

53.7

51.2

79.4

98.3

UDA setup

(a)

300

600

900

CL training epochs

60

80

100

STOC vs. CL in UDA

CL

STOC

ERM

50

100

150

CL training epochs

60

80

100

ID Accuracy

STOC vs. CL in SSL

CL

STOC

ERM

(b)

0.4

0.6

0.8

1.0

1.2

1.4

γ/σ

sp

40

60

80

100

Effect of problem parameters

ERM

ST

CL

STOC

(c)

Keeping the Φ cl fixed, we then learn a linear classifier h cl over Φ cl to minimize the exponential loss on labeled source data (refer to as linear probing ). For STOC, keeping the Φ cl fixed and initializing the linear head with the CL linear probe (instead of source only ERM), we perform ST with (2).

Example 1. For the setup in (1) , we choose γ ' 0 . 5 , σ 2 sp ' 1 . , and σ 2 in ' 0 . 05 with d in ' 5 and d sp ' 20 for our running example. γ { ? d sp controls signal to noise ratio in the source such that spurious feature is easy-to-learn and the invariant feature is harder-to-learn. σ 2 controls the noise in target which we show later is critical in unlearning the spurious feature with CL.

4.1 Simulations and Intuitive Story: A Comparative Study Between SSL and DA

Our setup captures real-world trends in UDA setting. Our toy setup (in Example 1) accentuates the behaviors observed on real-world datasets (Fig. 2(a)): (i) both ERM and ST yield close to random performance (though ST performs slightly worse than ERM); (ii) CL improves over ERM but still yields sub-optimal target performance; (iii) STOC then further improves over CL, achieving nearoptimal target performance. Note that, a linear predictor can improve target performance only by reducing its dependence on spurious feature x sp , and increasing it on invariant feature x in (along w ‹ ). Given this, we can explain our trends if we understand the following: (i) how ST reduces dependence on spurious feature when done after CL; (ii) why CL helps reduce but not completely eliminate the reliance of linear head on spurious features. Before we present intuitions, we ablate over a key problem parameter that affects both the target performance and conditions for ST to work.

Effect of γ { σ sp on success of ST. By increasing the ratio of margin γ and variance of spurious feature on target σ sp (keeping others constant), the problem becomes easier because γ directly affects the signal on x in and reducing σ sp helps ST to unlearn x sp (see App. D.3). In Fig. 2(c), we see that a phase transition occurs for ST, i.e ., after a certain threshold of γ { σ sp , ST successfully recovers the optimal target predictor. This hints that ST has a binary effect, where beyond a certain magnitude of γ { σ sp , ST can amplify the signal on domain invariant feature to obtain optimal target predictor. On the other hand, the performance of CL and ERM improve gradually where CL achieves high performance even at small ratios of γ { σ sp . One way of viewing this trend with CL is that it magnifies the effective γ { σ sp in its representation space, because of which a linear head trained these representations have a good performance at low values of the ratio. Consequently, the phase transition of STOC occurs much sooner then that of ST. Finally, we note that for CL the rate of performance increase diminishes at high values of γ { σ sp because CL fails to reduce dependency along x sp beyond a certain point.

An intuitive story. We return to the question of why self-training improves over contrastive learning under distribution shift in our Example 1. When the classifier at initialization of ST relies more on spurious features, ST aggravates this dependency. However, as the problem becomes easier (with increasing γ { σ sp ), the source-only ERM classifier will start relying more on invariant rather than spurious feature. Once this ERM classifier is sufficiently accurate on the target, ST unlearns any dependency on spurious features achieving optimal target performance. In contrast, we observe that CLperforms better than ERM but is still sub-optimal. This implies that CL ends up decreasing reliance on spurious features (as compared to ERM) but doesn't completely eliminate them. Combining ST and CL, a natural hypothesis explaining our trends is that CL provides a 'favorable' initialization

for ST by sufficiently increasing signal on invariant features. Our empirical findings emphasize an intriguing contrast suggesting that ST and CL improve target performance in complementary ways.

Why disparate behaviors for out-of-distribution vs. in distribution? In the SSL setup, recall, there is no distribution shift. In Example 1, we sample 50 k unlabeled data and 100 labeled data from the same (target) distribution to simulate SSL setup. Substantiating our findings on realworld data, we observe that STOC provides a small to negligible improvement over CL (refer to App. D). To understand why such disparate behaviors emerge, recall that in the UDA setting, the main benefit of STOC lies in picking up reliance on 'good' features for OOD data, facilitated by CL initialization. While contrastive pretraining uncovers features that are 'good' for OOD data, it also learns more predictive source-only features (which are not predictive at all on target). As a result, linear probing with source-labeled data picks up these source-only features, leaving considerable room for improvement on OOD data with further self-training. On the other hand, in the SSL setting, the limited ID labeled data might provide enough signal to pick up features predictive on ID data, leaving little to no room for improvement for further self-training. Corroborating our intuitions, throughout the CL training in the toy setup, when CL doesn't achieve near-perfect generalization, the improvements provided by STOC for each checkpoint remain minimal. On the other hand, for UDA setup, after reaching a certain training checkpoint in CL, STOC yields significant improvement (Fig. 2(b)).

In the next sections, we formalize our intuitions and analyze why ST and CL offer complementary benefits when dealing with distribution shifts. Formal statements and proofs are in App. E.

4.2 Conditions for Success and Failure of Self-training over ERM from Scratch

In our results on Example 1, we observe that performing ST after ERM yields a classifier with nearrandom target accuracy. In Theorem 2, we characterize conditions under which ST fails and succeeds.

Theorem 2 (Informal; Conditions for success and failure of ST over ERM) . The target accuracy of ERM classifier, is given by 0 . 5 ¨ erfc ` ´ γ 2 {p ? 2 d sp ¨ σ sp q ˘ . Then ST performed in the second stage yields: (i) a classifier with « 0 . 5 target accuracy when γ ă 1 { 2 σ sp and σ sp ě 1 ; and (ii) a classifier with near-perfect target accuracy when γ ě σ sp .

The informal theorem above abstracts the exact dependency of γ, σ sp , and d sp for the success and failure of ST over ERM. Our analysis highlights that while ERM learns a perfect predictor along w in (with norm γ ), it also learns to depend on w sp (with norm a d sp ) because of the perfect correlation of x sp with labels on the source. Our conditions depict that when the γ { σ sp is sufficiently small, then ST continues to erroneously enhance its reliance on the x sp feature for target prediction, resulting in near-random target performance. Conversely, when γ { σ sp is larger than 1, the signal in x in is correctly used for predictor on the majority of target points, and ST eliminates the x sp dependency, converging to an optimal target classifier.

Our proof analysis shows that if the ratio of the norm of the classifier along in the direction of w ‹ is smaller than w sp by a certain ratio then the generated pseudolabels (incorrectly) use x sp for its prediction further increasing the component along w sp . Moreover, normalization further diminishes the reliance along w ‹ , culminating in a near-random performance. The opposite occurs when the ERM classifier achieves a signal along w ‹ that is sufficiently stronger than along w sp . Upon substituting the parameters used in Example 1, the ERM and ST performances as determined by Theorem 2 align with our empirical results, notably, ST performance on target being near-random.

4.3 CL Captures Both Features But Amplifies Invariant Over Spurious Features

Here we show that minimizing the contrastive loss (3) on unlabeled data from both P S and P T gives us a feature extractor Φ cl that has a higher inner product with the invariant feature over the spurious feature. First, we derive a closed form expression for Φ cl that holds for any linear backbone and augmentation distribution. Then, we introduce assumptions on the augmentation distribution (or equivalently on w ‹ ) and other problem parameters, that are sufficient to prove amplification.

Proposition 3 (Barlow Twins solution) . The solution for (3) is U J k Σ ´ 1 { 2 A where U k are the top k eigenvectors of Σ ´ 1 { 2 A r ΣΣ ´ 1 { 2 A . Here, Σ A : ' E a ' P A r aa J s is the covariance over augmentations, and r Σ : ' E x ' P U r r a p x q r a p x q J s is the covariance matrix of mean augmentations r a p x q : ' E P A p a | x q r a s .

The above result captures the effect of augmentations through the matrix U k . If there were no augmentations, then Σ A ' r Σ , implying that U k could then be any random orthonormal matrix. On

the other hand if augmentation distributions change prevalent covariances in the data, i.e ., Σ A is very different from r Σ , the matrix U k would bias the CL solution towards directions that capture significant variance in marginal distribution on augmented data, but have low conditional variance, when conditioned on original point x -precisely the directions with low invariance loss. Hence, we can expect that CL would learn components along both invariant w in and spurious w sp because: (i) these directions explain a large fraction of variance in the raw data; (ii) augmentations that randomly scale down dimensions would add little variance along w sp and w in compared to noise directions in their null space. On the other hand it is unclear which of these directions is amplified more in Φ cl . The following assumption and amplification result conveys that when the noise in target p σ sp q is suficiently large, the CL solution amplifies the invariant feature over the spurious feature.

Assumption 4 (Informal; Alignment of w ‹ with augmentations) . We assume that w ‹ aligns with P A p¨ | x q , i.e ., @ x , E a | x r a J w ‹ s ' 1 { 2 ¨ x J diag p d q w ‹ is high. Hence, we assume w ‹ ' d in { ? d in .

✶ ✶ One implication of Assumption 4 is that when w ‹ ' ✶ d in { ? d in , only the top two eigenvectors lie in the space spanned by w in and w sp . To analyze our amplification with fewer eigenvectors from Proposition 3 while retaining all relevant phenomena, we assume w ‹ ' ✶ d in { ? d in for mathematical convenience. While Assumption 4 permits a tighter theoretical analysis, our empirical results in Sec. 4.1 hold more generally for w ‹ ' N p 0 , I d in q .

Theorem 5 (Informal; CL recovers both w in and w sp but amplifies w in ) . Under Assumption 4, the CL solution Φ cl 'r ϕ 1 , ϕ 2 , ..., ϕ k s satisfies ϕ J j w in ' ϕ J j w sp ' 0 @ j ě 3 , ϕ 1 ' c 1 w in ` c 3 w sp and ϕ 2 ' c 2 w in ` c 4 w sp . For constants K 1 , K 2 ą 0 , γ ' K 1 K 2 { σ sp , d sp ' σ 2 sp { K 2 2 , @ ϵ ą 0 , D σ sp 0 , such that for σ sp ě σ sp 0 , ∣ ∣ c 1 { c 3 ´ K 1 K 2 2 d in { 2 Lσ 2 in p d in ´ 1 q ∣ ∣ ď ϵ , and ∣ ∣ | c 2 { c 4 | ´ L ? d sp { γ ∣ ∣ ď ϵ , where L ' 1 ` K 2 2 .

We analyze the amplification of w in { w sp with contrastive learning in the regime where σ sp is large enough. In other words, if the target distribution has sufficient noise along the spurious feature, the augmentations prevent the CL solution from extracting components along w sp . Thus, in our analysis, we first analyze the amplification factors asymptotically p σ sp Ñ8q , and then use the asymptotic behavior to draw conclusions for the regime where σ sp is large but finite.

Theorem 5 conveys two results: (i) CL recovers components along both w in and w sp through ϕ 1 , ϕ 2 ; and (ii) it increases the norm along w in more than w sp . The latter is evident because the margin separating labeled points along w in is now amplified by a factor of | c 2 { c 4 | ' Ω p L ? d sp { γ q in ϕ 2 . Naturally, this will improve the target performance of a linear predictor trained over CL representations. At the same time, we also see that in ϕ 1 , the component along w sp is still significant ( c 1 { c 3 ' O p 1 { Lσ 2 in q ). Intuitively, CL prefers the invariant feature since augmentations amplify the noise along w sp in the target domain. At the same time, the variance induced by augmentations along w sp in source is still very small due to which the dependence on w sp is not completely alleviated. Due to the remaining components along w sp , the target performance for CL can remain less than ideal. Both the above arguments on target performance are captured in Corollary 6.

Corollary 6 (Informal; CL improves OOD error over ERM but is still imperfect) . For γ, σ sp , d sp defined as in Theorem 5, D σ sp 1 such that for all σ sp ě σ sp 1 , the target accuracy of CL (linear predictor on Φ cl ) is ě 0 . 5 erfc p´ L 1 ¨ γ { ? 2 σ sp q and ď 0 . 5 erfc p´ 4 L 1 ¨ γ { ? 2 σ sp q , where L 1 ' K 2 2 K 1 { σ 2 in p 1 ´ 1 { d in q . When σ sp 1 ą σ in a 1 ´ 1 { d in , the lower bound on accuracy is strictly better than ERM from scratch.

While Φ cl is still not ideal for linear probing, in the next part we will see how Φ cl can instead be sufficient for subsequent self-training to unlearn the remaining components along spurious features.

4.4 Improvements with Self-training Over Contrastive Learning

The result in the previous section highlights that while CL may improve over ERM, the linear probe continues to depend on the spurious feature. Next, we characterize the behavior STOC. Recall, in the ST stage, we iteratively update the linear head with (2) starting with the CL backbone and head.

Theorem 7 (Informal; ST improves over CL) . Under the conditions of Theorem 5 and d sp ď K 2 1 ¨ K 2 { 3 2 , the target accuracy of ST over CL is lower bounded by 0 . 5 ¨ erfc p´ | c 2 { c 4 | ¨ γ { p ? 2 σ 2 q q « 0 . 5 ¨ erfc ` ´ L ? d sp { p ? 2 σ sp q ˘ where c 2 and c 4 are the coefficients of feature ϕ 2 along w in and w sp learned by BT.

The above theorem states that when ? d sp { σ sp " 1 the target accuracy of ST over CL is close to 1. In Example 1, the lower bound of the accuracy of ST over CL is erfc ` ´ ? 10 ˘ « 2 showing nearperfect target generalization. Recall that Theorem 6 shows that CL yields a linear head that mainly depends on both the invariant direction w in and the spurious direction w sp . At initialization, the linear head trained on the CL backbone has negligible dependence on ϕ 2 (under conditions in Theorem 6). Building on that, the analysis in Theorem 7 captures that ST gradually reduces the dependence on w sp by learning a linear head that has a larger reliance on ϕ 2 , which has a higher 'effective' margin on the target, thus increasing overall dependency on w in .

Theoretical comparison with SSL. Our analysis until now shows that linear probing with source labeled data during CL picks up features that are more predictive of source label under distribution shift, leaving a significant room for improvement on OOD data when self-trained further. In UDA, the primary benefit of ST lies in picking up the features with a high 'effective' margin on target data that are not picked up by linear head trained during CL. In contrast, in the SSL setting, the limited ID labeled data may provide enough signal in picking up high-margin features that are predictive on ID data, leaving little to no room for improvement for further ST. We formalize this intuition in App. E.

4.5 Reconciling Practice: Implications for Deep Non-Linear Networks

In this section, we experiment with deep non-linear backbone ( i.e ., Φ cl ). When we continue to fix Φ cl during CL and STOC, the trends we observed with linear networks in Sec. 4.1 continue to hold. We then perform full fine-tuning with CL and STOC, i.e., propagate gradients even to Φ cl , as commonly done in practice. We present key takeaways here but detailed experiments are in App. D.4.

Benefits of augmentation for self-training. ST while updating Φ cl can hurt due to overfitting issues when training with the finite sample of labeled and unlabeled data (drop by >10% over CL). This is due to the ability of deep networks to overfit on confident but incorrect pseudolabels on target data [94]. This exacerbates components along w sp and we find that augmentations (and other heuristics) typically used in practice ( e.g . in FixMatch [79]) help avoid overfitting on incorrect pseudolabels.

Can ERM and ST over contrastive pretraining improve features? We find that self-training can also slightly improve features when we update the backbone with the second stage of STOC and when the CL backbone is early stopped sub-optimally ( i.e . at an earlier checkpoint in Fig. 2(b)). This feature finetuning can now widen the gap between STOC and CL in SSL settings, as compared to the linear probing gap (as in 2). This is because STOC can now improve performance beyond just recovering the generalization gap for the linear head (which is typically small). However, STOC benefits are negligible when CL is not early stopped sub-optimally, i.e ., trained till convergence. Thus, it remains unclear if STOC and CL have complementary benefits for feature learning in UDA or SSL settings. Investigating this is an interesting avenue for future work.

5 Connecting Experimental Gains with Theoretical Insights

Our theory emphasizes that under distribution shift contrastive pretraining improves the representations for target data, while self-training primarily improves linear classifiers learned on top. To investigate different methods in our UDA setup, we study the representations learned by each of them. We fix the representations and train linear heads over them to answer two questions: (i) How good are the representations in terms of their ceiling of target accuracy (performance of the optimal linear probe)?-we evaluate this by training the classifier head on target labeled data ( i.e ., target linear probe); and (ii) How well do heads trained on source generalize to target?-we assess this by training a head on source labeled data (source linear probe) and evaluate its difference with target linear probe. For both, we plot target accuracy. We make two intriguing observations Fig. 3):

Does CL improve representations over ERM features? Yes. We observe a substantial difference in accuracy ( « 14% gap) of target linear probes on backbones trained

Source probe

Target probe

50

60

70

80

Target Accuracy

49.0

63.3

57.0

67.1

59.7

79.4

65.3

80.7

Transfer with linear probe

ERM

ST

CL

STOC

with contrastive pretraining ( i.e . CL, STOC) and without it ( i.e ., ERM, ST) highlighting that CL significantly pushes the performance ceiling over non-contrastive features. As a side, our findings also stand in contrast to recent studies suggesting that ERM features might be 'good enough' for OOD generalization [67, 46]. Instead, the observed gains with contrastively pretrained backbones ( i.e . CL, STOC) demonstrate that target unlabeled data can be leveraged to further improve over ERM features.

Do CL features yield perfect linear transferability from source to target? Recent works [40, 76] conjecture that under certain conditions CL representations, linear probes learned with source labeled data may transfer perfectly from source to target. However, we observe that this doesn't hold strictly in practice, and in fact, the linear transferability can be further improved with ST. We first note a significant gap between the performance of source linear probes and target linear probes illustrating that linear transferability is not perfect in practice. Moreover, while the accuracy of target linear probes doesn't change substantially between CL and STOC, the accuracy of the source linear probe improves significantly. Similar observations hold for ERM and ST, methods trained without contrastive pretraining. This highlights that ST performs 'feature refinement' to improve source to target linear transfer (with relatively small improvements in their respective target probe performance). The findings highlight the complementary nature of benefits on real-world data: ST improves linear transferability while CL improves representations.

6 Connections to Prior Work

Our empirical results and our analyses offer a perspective that contrasts with the prior literature that argues for the individual optimality of contrastive pretraining and self-training. We outline the key differences from existing studies here, and delve into other related works in App. A.

Limitations of prior work analyzing contrastive learning Prior works [40, 76] analyzing CL first make assumptions on the consistency of augmentations with labels [39, 11, 73, 44], and specifically for UDA make stronger ones on the augmentation graph connecting examples from same domain or class more than cross-class/cross-domain ones. While this is sufficient to prove linear transferability, it is unclear if this holds in practice when augmentations are imperfect, i.e . if they fail to mask the spurious features completely-as corroborated by our findings in Sec. 5. We show why this also fails in our simplified setup in App. F.1.

Limitations of prior work analyzing self-training Prior research views self-training as consistency regularization, ensuring pseudolabels for original samples align with their augmentations [12, 87, 79]. This approach abstracts the role played by the optimization algorithm and instead evaluates the global minimizer of a population objective promoting pseudolabel consistency. It also relies on specific assumptions about class-conditional distributions to guarantee pseudolabel accuracy across domains. However, this framework doesn't address issues in iterative label propagation. For example, when augmentation distribution has long tails, the consistency of pseudolabels depends on the sampling frequency of 'favorable' augmentations (for more discussion see App. F.2). Our analysis thus follows the iterative examination of self-training [17].

7 Conclusion

In this study, we highlight the synergistic behavior of self-training and contrastive pretraining under distribution shift. Shifts in distribution are commonplace in real-world applications of machine learning, and even under natural, non-adversarial distribution shifts, the performance of machine learning models often drops. By simply combining existing techniques in self-training and constrastive learning, we find that we can improve accuracy by 3-8% rather than using either approach independently. Despite these significant improvements, we note that one limitation of this combined approach is that performing self-training sequentially after contrastive pretraining increases the computation cost for UDA. The potential for integrating these benefits into one unified training paradigm is yet unclear, presenting an interesting direction for future exploration.

Beyond this, we note that our theoretical framework primarily confines the analysis to training the backbone and linear network independently during the pretraining and fine-tuning/self-training phases. Although our empirical observations apply to deep networks with full fine-tuning, we leave a more rigorous theoretical study of full fine-tuning for future work. Our theory also relies on a covariate shift assumption (where we assume that label distribution also doesn't shift). Investigating the complementary nature of self-training and contrastive pretraining beyond the covariate shift assumption would be another interesting direction for future work.

References

Alexandari, A., Kundaje, A., and Shrikumar, A. (2021). Adapting to label shift with biascorrected calibration. In International Conference on Machine Learning (ICML) .

Arora, S., Khandeparkar, H., Khodak, M., Plevrakis, O., and Saunshi, N. (2019). A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint arXiv:1902.09229 .

Azizzadenesheli, K., Liu, A., Yang, F., and Anandkumar, A. (2019). Regularized learning for domain adaptation under label shifts. In International Conference on Learning Representations (ICLR) .

Bardes, A., Ponce, J., and LeCun, Y. (2021). Vicreg: Variance-invariance-covariance regularization for self-supervised learning. arXiv preprint arXiv:2105.04906 .

Baricz, Á. (2008). Mills' ratio: Monotonicity patterns and functional inequalities. Journal of Mathematical Analysis and Applications , 340(2):1362-1370.

Bekker, J. and Davis, J. (2020). Learning from positive and unlabeled data: a survey. Machine Learning .

Ben-David, S., Lu, T., Luu, T., and Pál, D. (2010). Impossibility Theorems for Domain Adaptation. In International Conference on Artificial Intelligence and Statistics (AISTATS) .

Berthelot, D., Carlini, N., Cubuk, E. D., Kurakin, A., Sohn, K., Zhang, H., and Raffel, C. (2019). Remixmatch: Semi-supervised learning with distribution alignment and augmentation anchoring. arXiv preprint arXiv:1911.09785 .

Bishop, C. M. (2006). Pattern Recognition and Machine Learning . Springer.

Blanchard, G., Lee, G., and Scott, C. (2011). Generalizing from several related classification tasks to a new unlabeled sample. Advances in neural information processing systems , 24.

Cabannes, V., Kiani, B. T., Balestriero, R., LeCun, Y., and Bietti, A. (2023). The ssl interplay: Augmentations, inductive bias, and generalization. arXiv preprint arXiv:2302.02774 .

Cai, T., Gao, R., Lee, J., and Lei, Q. (2021). A theory of label propagation for subpopulation shift. In International Conference on Machine Learning , pages 1170-1182. PMLR.

Caron, M., Misra, I., Mairal, J., Goyal, P., Bojanowski, P., and Joulin, A. (2020). Unsupervised learning of visual features by contrasting cluster assignments. Advances in Neural Information Processing Systems , 33:9912-9924.

Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., and Joulin, A. (2021). Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision , pages 9650-9660.

Chapelle, O., Scholkopf, B., and Zien, A. (2006). Semi-supervised learning. 2006. Cambridge, Massachusettes: The MIT Press View Article , 2.

Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. (2020a). A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR.

Chen, X., Chen, W., Chen, T., Yuan, Y., Gong, C., Chen, K., and Wang, Z. (2020b). Self-pu: Self boosted and calibrated positive-unlabeled training. In International Conference on Machine Learning , pages 1510-1519. PMLR.

Christie, G., Fendley, N., Wilson, J., and Mukherjee, R. (2018). Functional map of the world. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition .

Cortes, C., Mansour, Y., and Mohri, M. (2010). Learning Bounds for Importance Weighting. In Advances in Neural Information Processing Systems (NIPS) .

Cortes, C. and Mohri, M. (2014). Domain adaptation and sample bias correction theory and algorithm for regression. Theoretical Computer Science , 519.

Cubuk, E. D., Zoph, B., Shlens, J., and Le, Q. V. (2020). Randaugment: Practical automated data augmentation with a reduced search space. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops , pages 702-703.

Darlow, L. N., Crowley, E. J., Antoniou, A., and Storkey, A. J. (2018). Cinic-10 is not imagenet or cifar-10. arXiv preprint arXiv:1810.03505 .

Deledalle, C.-A., Denis, L., Tabti, S., and Tupin, F. (2017). Closed-form expressions of the eigen decomposition of 2 x 2 and 3 x 3 Hermitian matrices . PhD thesis, Université de Lyon.

DeVries, T. and Taylor, G. W. (2017). Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv:1708.04552 .

Elkan, C. and Noto, K. (2008). Learning classifiers from only positive and unlabeled data. In International Conference Knowledge Discovery and Data Mining (KDD) , pages 213-220.

Ganin, Y., Ustinova, E., Ajakan, H., Germain, P., Larochelle, H., Laviolette, F., Marchand, M., and Lempitsky, V. (2016). Domain-adversarial training of neural networks. The journal of machine learning research .

Gardner, J., Pleiss, G., Weinberger, K. Q., Bindel, D., and Wilson, A. G. (2018). Gpytorch: Blackbox matrix-matrix gaussian process inference with gpu acceleration. In Advances in Neural Information Processing Systems (NeurIPS) .

Garg, S., Balakrishnan, S., and Lipton, Z. (2022a). Domain adaptation under open set label shift. In Advances in Neural Information Processing Systems (NeurIPS) .

Garg, S., Balakrishnan, S., Lipton, Z., Neyshabur, B., and Sedghi, H. (2022b). Leveraging unlabeled data to predict out-of-distribution performance. In International Conference on Learning Representations (ICLR) .

Garg, S., Erickson, N., Sharpnack, J., Smola, A., Balakrishnan, S., and Lipton, Z. (2023). Rlsbench: A large-scale empirical study of domain adaptation under relaxed label shift. In International Conference on Machine Learning (ICML) .

Garg, S., Wu, Y., Balakrishnan, S., and Lipton, Z. (2020). A unified view of label shift estimation. In Advances in Neural Information Processing Systems (NeurIPS) .

Garg, S., Wu, Y., Smola, A., Balakrishnan, S., and Lipton, Z. (2021). Mixture proportion estimation and PU learning: A modern approach. In Advances in Neural Information Processing Systems (NeurIPS) .

Garrido, Q., Chen, Y., Bardes, A., Najman, L., and Lecun, Y. (2022). On the duality between contrastive and non-contrastive self-supervised learning. arXiv preprint arXiv:2206.02574 .

Grandvalet, Y. and Bengio, Y. (2006). Entropy regularization.

Gretton, A., Smola, A. J., Huang, J., Schmittfull, M., Borgwardt, K. M., and Schölkopf, B. (2009). Covariate Shift by Kernel Mean Matching. Journal of Machine Learning Research (JMLR) .

Grill, J.-B., Strub, F., Altché, F., Tallec, C., Richemond, P., Buchatskaya, E., Doersch, C., Avila Pires, B., Guo, Z., Gheshlaghi Azar, M., et al. (2020). Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems , 33:21271-21284.

Gulrajani, I. and Lopez-Paz, D. (2020). In search of lost domain generalization. arXiv preprint arXiv:2007.01434 .

HaoChen, J. Z. and Ma, T. (2022). A theoretical study of inductive biases in contrastive learning. arXiv preprint arXiv:2211.14699 .

HaoChen, J. Z., Wei, C., Gaidon, A., and Ma, T. (2021). Provable guarantees for self-supervised deep learning with spectral contrastive loss. Advances in Neural Information Processing Systems , 34:5000-5011.

HaoChen, J. Z., Wei, C., Kumar, A., and Ma, T. (2022). Beyond separability: Analyzing the linear transferability of contrastive representations to related subpopulations. arXiv preprint arXiv:2204.02683 .

He, K., Fan, H., Wu, Y., Xie, S., and Girshick, R. (2020). Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 9729-9738.

He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep Residual Learning for Image Recognition. In Computer Vision and Pattern Recognition (CVPR) .

Joachims, T. et al. (1999). Transductive inference for text classification using support vector machines. In Icml , volume 99, pages 200-209.

Johnson, D. D., Hanchi, A. E., and Maddison, C. J. (2022). Contrastive learning can find an optimal basis for approximately view-invariant functions. arXiv preprint arXiv:2210.01883 .

Kakade, S. M., Sridharan, K., and Tewari, A. (2008). On the complexity of linear prediction: Risk bounds, margin bounds, and regularization. Advances in neural information processing systems , 21.

Kirichenko, P., Izmailov, P., and Wilson, A. G. (2022). Last layer re-training is sufficient for robustness to spurious correlations. arXiv preprint arXiv:2204.02937 .

Koh, P. W., Sagawa, S., Marklund, H., Xie, S. M., Zhang, M., Balsubramani, A., Hu, W., Yasunaga, M., Phillips, R. L., Gao, I., Lee, T., David, E., Stavness, I., Guo, W., Earnshaw, B. A., Haque, I. S., Beery, S., Leskovec, J., Kundaje, A., Pierson, E., Levine, S., Finn, C., and Liang, P. (2021). WILDS: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning (ICML) .

Krizhevsky, A. and Hinton, G. (2009). Learning Multiple Layers of Features from Tiny Images. Technical report, Citeseer.

Kschischang, F. R. (2017). The complementary error function. Online, April .

Kumar, A., Ma, T., and Liang, P. (2020). Understanding self-training for gradual domain adaptation. In International Conference on Machine Learning , pages 5468-5479. PMLR.

Kumar, A., Raghunathan, A., Jones, R. M., Ma, T., and Liang, P. (2022). Fine-tuning can distort pretrained features and underperform out-of-distribution. In International Conference on Learning Representations .

Lee, D.-H. et al. (2013). Pseudo-label: The simple and efficient semi-supervised learning method for deep neural networks. In Workshop on challenges in representation learning, ICML , volume 3, page 896.

Lipton, Z. C., Wang, Y.-X., and Smola, A. (2018). Detecting and Correcting for Label Shift with Black Box Predictors. In International Conference on Machine Learning (ICML) .

Long, M., Cao, Y., Wang, J., and Jordan, M. (2015). Learning transferable features with deep adaptation networks. In International conference on machine learning , pages 97-105. PMLR.

Long, M., Zhu, H., Wang, J., and Jordan, M. I. (2017). Deep transfer learning with joint adaptation networks. In International conference on machine learning . PMLR.

Loshchilov, I. and Hutter, F. (2016). Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983 .

Ma, M. Q., Tsai, Y.-H. H., Liang, P. P., Zhao, H., Zhang, K., Salakhutdinov, R., and Morency, L.-P. (2021). Conditional contrastive learning for improving fairness in self-supervised learning. arXiv preprint arXiv:2106.02866 .

Mishra, S., Saenko, K., and Saligrama, V. (2021). Surprisingly simple semi-supervised domain adaptation with pretraining and consistency. arXiv preprint arXiv:2101.12727 .

Muandet, K., Balduzzi, D., and Schölkopf, B. (2013). Domain generalization via invariant feature representation. In International Conference on Machine Learning , pages 10-18. PMLR.

Nagarajan, V., Andreassen, A., and Neyshabur, B. (2020). Understanding the failure modes of out-of-distribution generalization. arXiv preprint arXiv:2010.15775 .

Oord, A. v. d., Li, Y ., and Vinyals, O. (2018). Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748 .

Peng, X., Bai, Q., Xia, X., Huang, Z., Saenko, K., and Wang, B. (2019). Moment matching for multi-source domain adaptation. In Proceedings of the IEEE/CVF international conference on computer vision , pages 1406-1415.

Peng, X., Usman, B., Kaushik, N., Hoffman, J., Wang, D., and Saenko, K. (2017). Visda: The visual domain adaptation challenge.

Peng, X., Usman, B., Saito, K., Kaushik, N., Hoffman, J., and Saenko, K. (2018). Syn2real: A new benchmark forsynthetic-to-real visual domain adaptation. arXiv preprint arXiv:1806.09755 .

Quinonero-Candela, J., Sugiyama, M., Schwaighofer, A., and Lawrence, N. D. (2008). Dataset shift in machine learning . Mit Press.

Roberts, M., Mani, P., Garg, S., and Lipton, Z. (2022). Unsupervised learning under latent label shift. In Advances in Neural Information Processing Systems (NeurIPS) .

Rosenfeld, E., Ravikumar, P., and Risteski, A. (2022). Domain-adjusted regression or: Erm may already learn features sufficient for out-of-distribution generalization. arXiv preprint arXiv:2202.06856 .

Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., et al. (2015). Imagenet large scale visual recognition challenge. International journal of computer vision , 115(3):211-252.

Saerens, M., Latinne, P., and Decaestecker, C. (2002). Adjusting the Outputs of a Classifier to New a Priori Probabilities: A Simple Procedure. Neural Computation .

Sagawa, S., Koh, P. W., Lee, T., Gao, I., Xie, S. M., Shen, K., Kumar, A., Hu, W., Yasunaga, M., Marklund, H., Beery, S., David, E., Stavness, I., Guo, W., Leskovec, J., Saenko, K., Hashimoto, T., Levine, S., Finn, C., and Liang, P. (2021). Extending the wilds benchmark for unsupervised adaptation. In NeurIPS Workshop on Distribution Shifts .

Sagawa, S., Raghunathan, A., Koh, P. W., and Liang, P. (2020). An investigation of why overparameterization exacerbates spurious correlations. In International Conference on Machine Learning , pages 8346-8356. PMLR.

Santurkar, S., Tsipras, D., and Madry, A. (2021). Breeds: Benchmarks for subpopulation shift. In International Conference on Learning Representations (ICLR) .

Saunshi, N., Ash, J., Goel, S., Misra, D., Zhang, C., Arora, S., Kakade, S., and Krishnamurthy, A. (2022). Understanding contrastive learning requires incorporating inductive biases. In International Conference on Machine Learning , pages 19250-19286. PMLR.

Schölkopf, B., Janzing, D., Peters, J., Sgouritsa, E., Zhang, K., and Mooij, J. (2012). On Causal and Anticausal Learning. In International Conference on Machine Learning (ICML) .

Scudder, H. (1965). Probability of error of some adaptive pattern-recognition machines. IEEE Transactions on Information Theory , 11(3):363-371.

Shen, K., Jones, R. M., Kumar, A., Xie, S. M., HaoChen, J. Z., Ma, T., and Liang, P. (2022). Connect, not collapse: Explaining contrastive learning for unsupervised domain adaptation. In International Conference on Machine Learning , pages 19847-19878. PMLR.

Shimodaira, H. (2000). Improving Predictive Inference Under Covariate Shift by Weighting the Log-Likelihood Function. Journal of Statistical Planning and Inference .

Shu, R., Bui, H. H., Narui, H., and Ermon, S. (2018). A dirt-t approach to unsupervised domain adaptation. arXiv preprint arXiv:1802.08735 .

Sohn, K., Berthelot, D., Carlini, N., Zhang, Z., Zhang, H., Raffel, C. A., Cubuk, E. D., Kurakin, A., and Li, C.-L. (2020). Fixmatch: Simplifying semi-supervised learning with consistency and confidence. Advances in Neural Information Processing Systems , 33.

Stewart, G. W. (1993). On the early history of the singular value decomposition. SIAM review , 35(4):551-566.

Sun, B., Feng, J., and Saenko, K. (2017). Correlation alignment for unsupervised domain adaptation. In Domain Adaptation in Computer Vision Applications . Springer.

Sun, B. and Saenko, K. (2016). Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision . Springer.

Torralba, A. and Efros, A. A. (2011). Unbiased look at dataset bias. In CVPR 2011 , pages 1521-1528. IEEE.

Van Engelen, J. E. and Hoos, H. H. (2020). A survey on semi-supervised learning. Machine learning , 109(2):373-440.

Venkateswara, H., Eusebio, J., Chakraborty, S., and Panchanathan, S. (2017). Deep hashing network for unsupervised domain adaptation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages 5018-5027.

Wang, D., Shelhamer, E., Liu, S., Olshausen, B., and Darrell, T. (2021). Tent: Fully test-time adaptation by entropy minimization. In International Conference on Learning Representations .

Wei, C., Shen, K., Chen, Y., and Ma, T. (2020). Theoretical analysis of self-training with deep networks on unlabeled data. arXiv preprint arXiv:2010.03622 .

Wu, Z., Xiong, Y., Yu, S. X., and Lin, D. (2018). Unsupervised feature learning via nonparametric instance discrimination. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 3733-3742.

Xie, Q., Luong, M.-T., Hovy, E., and Le, Q. V. (2020a). Self-training with noisy student improves imagenet classification. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10687-10698.

Xie, X., Chen, J., Li, Y., Shen, L., Ma, K., and Zheng, Y. (2020b). Self-supervised cyclegan for object-preserving image-to-image domain adaptation. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XX 16 , pages 498513. Springer.

Yang, X., Song, Z., King, I., and Xu, Z. (2022). A survey on deep semi-supervised learning. IEEE Transactions on Knowledge and Data Engineering .

Zadrozny, B. (2004). Learning and Evaluating Classifiers Under Sample Selection Bias. In International Conference on Machine Learning (ICML) .

Zbontar, J., Jing, L., Misra, I., LeCun, Y ., and Deny, S. (2021). Barlow twins: Self-supervised learning via redundancy reduction. In International Conference on Machine Learning , pages 12310-12320. PMLR.

Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O. (2017). Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations (ICLR) .

Zhang, J., Menon, A., Veit, A., Bhojanapalli, S., Kumar, S., and Sra, S. (2021). Coping with label shift via distributionally robust optimisation. In International Conference on Learning Representations (ICLR) .

Zhang, K., Schölkopf, B., Muandet, K., and Wang, Z. (2013). Domain Adaptation Under Target and Conditional Shift. In International Conference on Machine Learning (ICML) .

Zhang, R. (2019). Making convolutional networks shift-invariant again. In ICML .

Zhang, W., Ouyang, W., Li, W., and Xu, D. (2018). Collaborative and adversarial network for unsupervised domain adaptation. In Proceedings of the IEEE conference on computer vision and pattern recognition .

Zhang, Y., Liu, T., Long, M., and Jordan, M. (2019). Bridging theory and algorithm for domain adaptation. In International Conference on Machine Learning . PMLR.

Zhu, X. and Ghahramani, Z. (2003). Learning from labeled and unlabeled data with label propagation. CMU CALD tech report CMU-CALD-02-107, 2002 .

Appendix

Appendix Table of Contents

G Additional Lemmas

A Other Related Works

Unsupervised domain adaption. Without assumption on the nature of shift, UDA is underspecified [7]. This challenge has been addressed in various ways by researchers. One approach is to investigate additional structural assumptions under which UDA problems are well posed [77, 74]. Popular settings for which DA is well-posed include (i) covariate shift [96, 92, 19, 20, 35] where p p x q can change from source to target but p p y | x q remains invariant; and (ii) label shift [69, 53, 3, 1, 31, 95, 66, 30] where the label marginal p p y q can change but p p x | y q is shared across source and target. Principled methods with strong theoretical guarantees exists for adaptation under these settings when target

49

distribution's support is a subset of the source support. Other works [25, 6, 32, 28] extend the label shift setting to scenarios where previously unseen classes may appear in the target and p p x | y q remains invariant among seen classes. A complementary line of research focuses on constructing benchmarks to develop heuristics for incorporating the unlabeled target data, relying on benchmark datasets ostensibly representative of 'real-world shifts' to adjudicate progress [72, 85, 70, 62, 63]. As a result, various benchmark-driven heuristics have been proposed [54, 55, 82, 81, 99, 98, 26, 79]. Our work engages with the latter, focusing on two popular methods: self-training and contrastive pretraining.

Domain generalization. In domain generalization, the model is given access to data from multiple different domains and the goal is to generalize to a previously unseen domain at test time [10, 59]. For a survey of different algorithms for domain generalization, we refer the reader to Gulrajani and Lopez-Paz [37]. A crucial distinction here is that unlike the domain generalization setting, in DA problems, we have access to unlabeled examples from the test domain.

Semi-supervised learning. To learn from a small amount of labeled supervision, semi-supervised learning methods leverage unlabeled data alongside to improve learning models. One of the seminal works in SSL is the pseudolabeling method [75], where a classifier is trained on the labeled data and then used to classify the unlabeled data, which are then added to the training set. The work of Zhu and Ghahramani [100] built on this by introducing graph-based methods, and the transductive SVMs [43] presented an SVM-based approach. More recent works have focused on deep learning techniques, and similar to UDA, self-training and contrastive pretraining have emerged as two prominent choices. We delve into these methods in greater detail in the following paragraphs. For a discussion on other SSL methods, we refer interested readers to [15, 84, 91].

Self-training. Two popular forms of self-training are pseudolabeling [52] and conditional entropy minimization [34], which have been observed to be closely connected [8, 52, 79, 78]. Motivated by its strong performance in SSL and UDA settings [79, 89, 30, 78], several theoretical works have made attempts to understand its behavior [50, 87, 17]. [87, 12] aims to understand the behavior of the global minimizer of self-training objective by studying input consistency regularization, which enforces stability of the prediction for different augmentations of the unlabeled data. Our analysis of self-training is motivated by the work of Chen et al. [17] which explores the iterative behavior of self-training to unlearn spurious features. The setting of spurious features is of particular interest, since prior works have specifically analyzed the failures of out-of-distribution generalization in the presence of spurious features [60, 71].

Contrastive learning. An alternate line of work that uses unlabeled data for learning representations in the pretraining stage is contrastive learning [36, 61, 13, 16, 88]. Given an augmentation distribution, the main goal of contrastive objectives is to map augmentations drawn from the same input (positive pairs) to similar features, and force apart features corresponding to augmentations of different inputs (negative pairs) [13, 14, 41]. Prior works [11, 44, 38] have also shown a close relationship between contrastive [16, 39] and non-contrastive objectives [4, 93]. Consequently, in our analysis pertaining to the toy setup we focus on the mathematically non-contrastive objective Barlow Twins [93]. Using this pretrained backbone (either as an initialization or as a fixed feature extractor) a downstream predictor is learned using labeled examples. Several works [39, 73, 38, 2, 44] have analyzed the in-distribution generalization of the downstream predictor via label consistency arguments on the graph of positive pairs (augmentation graph). In contrast, we study the impact of contrastive learning under distribution shifts in the UDA setup. Other works [76, 40] that examine contrastive learning for UDA also conform to the augmentation graph view point, making additional assumptions that guarantee linear transferability. In our simplified setup involving spurious correlations, these abstract assumptions break easily when the augmentations are of a generic nature, akin to practice. Finally, some empirical works [58, 57] have found self-supervised objectives like contrastive pretraining to reduce dependence on spurious correlations. Corroborating their findings, we extensively evaluate the complementary benefits of contrastive learning and self-training on real-world datasets. Finding differing results in SSL and UDA settings, we further examine their behavior theoretically in our toy setup.

B More Details on Problem Setup

In this section, we elaborate on our setup and methods studied in our work.

Unsupervised Domain Adaptation (UDA). We assume that we are given labeled data from the source distribution and unlabeled data from a shifted, target distribution, with the goal of performing well on target data. We assume that the source and target distributions have the same label marginals P S p y q ' P T p y q ( i.e ., no label proportion shift) and the same Bayes optimal predictor, i.e ., arg max y p S p y | x q ' arg max y p T p y | x q . Here, even with infinite labeled source data, the challenge lies in generalizing out-of-distribution. In experiments, we assume access to finite data but in theory, we assume population access to labeled source and unlabeled target.

Semi-Supervised Learning (SSL). Here, there is no distribution shift, i.e ., P S ' P T ' P U . We are given a small number of labeled examples and a comparatively large amount of unlabeled examples, both drawn from the same distribution. Without loss of generality, we denote this distribution with P T . The goal in SSL is to generalize in-distribution. The challenge is primarily due to limited access to labeled data. Here, in experiments, we assume limited access to labeled data but a comparatively larger amount of unlabeled in-distribution data. In theory, we assume population access to unlabeled data but limited labeled examples.

Methods. As discussed in the main paper, we compare four methods for learning with labeled and unlabeled data. Table 8 summarizes the main methods and key differences between those methods in UDA and SSL setup. For exact implementation in our experiments, we refer reader to App. C.3.

C Additional Experiments and Details

C.1 Additional setup and notation

Recall, our goal is to learn a predictor that maps inputs x P X Ď R d to outputs y P Y . We parameterize predictors f ' h ˝ Φ : R d ÞÑ Y , where Φ : R d ÞÑ R k is a feature map and h P R k is a classifier that maps the representation to the final scores or logits. With A : X Ñ A , we denote the augmentation function that takes in an input x and outputs an augmented view of the input A p x q . Unless specified otherwise, we perform full-finetuning in all of our experiments on real-world data. That is, we backpropagate gradients in both the linear head h and the backbone ϕ . For UDA, we denote source labeled points as tp x i , y i qu n i ' 1 and target unlabeled points as tp x 1 i qu i m ' 1 . For SSL, we use the same notation for labeled and unlabeled in-distribution data.

C.2 Dataset details

For both UDA and SSL, we conduct experiments across eight benchmark datasets. Each of these datasets consists of domains, enabling us to construct source-target pairs for UDA. The adopted source and target domains are standard to previous studies [76, 30, 70]. 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%. Below provide exact details about the datasets used in our benchmark study.

CIFAR10 We use the original CIFAR10 dataset [48] as the source dataset. For target domains, we consider CINIC10 [22] which is a subset of Imagenet restricted to CIFAR10 classes and downsampled to 32 ˆ 32.

FMoW In order to consider distribution shifts faced in the wild, we consider FMoW-WILDs [47, 18] from WILDS benchmark, which contains satellite images taken in different geographical regions and at different times. We use the original train as source and OOD val and OOD test splits as target domains as they are collected over different time-period. Overall, we obtain 3 different domains (1 source and 2 targets).

BREEDs We also consider BREEDs benchmark [72] in our setup to assess robustness to subpopulation shifts. BREEDs leverage class hierarchy in ImageNet [68] to re-purpose original classes to be the subpopulations and defines a classification task on superclasses. We consider distribution shift due to subpopulation shift which is induced by directly making the subpopulations present in the training and test distributions disjoint. BREEDs benchmark contains 4 datasets Entity-13 , Entity-30 , Living-17 , and Non-living-26 , each focusing on different subtrees and levels in the hierarchy. Overall, for each of the 4 BREEDs datasets (i.e., Entity-13, Entity-30, Living-

Dataset

CIFAR10

FMoW

Visda

Entity13

Entity30

Living17

Nonliving26

Officehome

CIFAR10v1

CINIC10

Year 2002-'13

Year '13-'16

Year '16-'18

Synthetic

Real-1

Real-2

Sub-pop. 1

Sub-pop. 2

Sub-pop. 1

Sub-pop. 2

Sub-pop. 1

Sub-pop. 2

Sub-pop. 2

Product

Real

Sub-pop. 1

Art

17, and Non-living-26), we obtain one different domain which we consider as target. We refer to source and target as follows: BREEDs sub-population 1, BREEDs sub-population 2.

OfficeHome We use four domains (art, clipart, product and real) from OfficeHome dataset [85]. We use the product domain as source and the other domains as target.

Visda We use three domains (train, val and test) from the Visda dataset [64, 63]. While 'train' domain contains synthetic renditions of the objects, 'val' and 'test' domains contain real world images. To avoid confusing, the domain names with their roles as splits, we rename them as 'synthetic', 'Real-1' and 'Real-2'. We use the synthetic (original train set) as the source domain and use the other domains as target.

We summarize the information about source and target domains in Table 3.

Train-test splits We partition each source and target dataset into 80% and 20% i.i.d. splits. We use 80% splits for training and 20% splits for evaluation (or validation). We throw away labels for the 80% target split and only use labels in the 20% target split for final evaluation. The rationale behind splitting the target data is to use a completely unseen batch of data for evaluation. This avoids evaluating on examples where a model potentially could have overfit. over-fitting to unlabeled examples for evaluation. In practice, if the aim is to make predictions on all the target data (i.e., transduction), we can simply use the (full) target set for training and evaluation.

Simulating SSL settings and limited supervision. For SSL settings, we choose the in-distribution domain as the source domain. To simulate limited supervision in SSL, we sub-sample the original labeled training set to 10% and use all the original dataset as unlabeled data. For evaluation, we further split the original holdout set into two partitions (one for validation and the other to report final accuracy numbers).

C.3 Method details

For implementation, we build on top of WILDs [70] and RLSbench [30] open source libraries.

ERM(Source only) training. We consider Empirical Risk Minimization (ERM) on the labeled source data as a baseline. Since this simply ignores the unlabeled target data, we call this as source only training. As mentioned in the main paper, we perform source only training with data augmentations. Formally, we minimize the following ERM loss:

where A is the stochastic data augmentation operation and ℓ is a loss function. For SSL, the ERM baseline only uses the small of labeled data available.

Contrastive Learning (CL). We perform contrastive pretraining on the unlabeled dataset to obtain the backbone ϕ cl . And then we perform full fine-tuning with source labeled data by initializing the backbone with ϕ cl . We use SwAV [13] for contrastive pretraining. The main idea behind SwAV is to train a model to identify different views of the same image as similar, while also ensuring that it finds different images to be distinct. This is accomplished through a swapped prediction mechanism, where the goal is to compute a code from an augmented version of the image and predict this code from other augmented versions of the same image. In particular, given two image features ϕ p x 1 a 1 q and ϕ p x 1 a 2 q from two different augmentations of the same image x 1 , i.e., x 1 a 1 , x 1 a 2 ' A p x 1 q , SwAV computes their codes z a 1 and z a 2 by matching the features to a set of K prototypes t c 1 , ¨ ¨ ¨ , c K u . Then SwAV minimizes the following loss such that ϕ p x 1 a 1 q can compute codes z a 2 and ϕ p x 1 a 2 q can compute codes z a 1 :

where ℓ 1 computes KL-divergence between codes computed with features (e.g. ϕ p x a 1 q ) and the code computed by another view (e.g. z a 2 ). For more details about the algorithm, we refer the reader to Caron et al. [13]. In all UDA settings, unless otherwise specified, we pool all the (unlabeled) data from the source and target to perform SwAV. For SSL, we leverage in-distribution unlabeled data.

We employ SimCLR [16] for the CIFAR10 dataset, aligning with previous studies that have utilized contrastive pretraining on the same dataset [51, 76]. The reason for this choice is that SwA V relies on augmentations that involve cropping images to a smaller resolution, making it more suitable for datasets with larger resolutions beyond 32 ˆ 32 .

Self-Training (ST). For self-training, we apply FixMatch [79], where the loss on labeled data and on pseudolabeled unlabeled data are minimized simultaneously. Sohn et al. [79] proposed FixMatch as a variant of the simpler Pseudo-label method [52]. This algorithm dynamically generates psuedolabels and overfits on them in each batch. FixMatch employs consistency regularization on the unlabeled data. In particular, while pseudolabels are generated on a weakly augmented view of

the unlabeled examples, the loss is computed with respect to predictions on a strongly augmented view. The intuition behind such an update is to encourage a model to make predictions on weakly augmented data consistent with the strongly augmented example. Moreover, FixMatch only overfits to the assigned labeled with weak augmentation if the confidence of the prediction with strong augmentation is greater than some threshold τ . Refer to A weak as the weak-augmentation and A strong as the strong-augmentation function. Then, FixMatch uses the following loss function:

where r y i ' arg max y f y p T weak p x i qq . For UDA, our unlabeled data is the union of source and target unlabeled data. For SSL, we only leverage in-distribution unlabeled data.

We adapted our implementation from Sagawa et al. [70] which matches the implementation of Sohn et al. [79] except for one detail. While Sohn et al. [79] augments labeled examples with weak augmentation, Sagawa et al. [70] proposed to strongly augment the labeled source examples.

Self-Training Over Contrastive learning (STOC). Finally, rather than performing FixMatch from a randomly initialized backbone, we initialize FixMatch with a contrastive pretrained backbone.

C.4 Additional UDA experimemts

C.5 Additional SSL experimemts

C.6 Other experimental details

Augmentations. For weak augmentation, we leverage random horizontal flips and random crops of pre-defined size. For SwAV, we also perform multicrop augmentation as proposed in Caron et al. [13]. For strong augmentation, we apply the following transformations sequentially: random horizontal flips, random crops of pre-defined size, augmentation with Cutout [24], and RandAugment [21]. For the exact implementation of RandAugment, we directly use the implementation of Sohn et al. [79]. Unless specified otherwise, for all methods, we default to using strong augmentation techniques.

Architectures. In our work, we experiment with Resnet18, Resnet50 [42] trained from scratch ( i.e . random initialization). We do not consider off-the-shelf pretrained models ( e.g ., on Imagenet [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 BREEDs datasets, we employ Resnet18 architecture. For other datasets, we train a Resnet50 architecture.

Except for Resnets on CIFAR dataset, we used the standard pytorch implementation [27]. For Resnet on Cifar, we refer to the implementation here: https://github.com/kuangliu/pytorch-cifar . For all the architectures, whenever applicable, we add antialiasing [97]. We use the official library released with the paper.

Hyperparameters. For all the methods, we fix the algorithm-specific hyperparameters to the original recommendations. 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.

We tune the learning rate and ℓ 2 regularization parameter by fixing the batch size for each dataset that corresponds to the maximum we can fit to 15GB GPU memory. We default to using cosine learning rate schedule [56]. We set the number of epochs for training as per the suggestions of the authors of respective benchmarks. For SSL, we run both ERM and FixMatch for approximately 2000 epochs. Note that we define the number of epochs as a full pass over the labeled training source data. We summarize the learning rate, batch size, number of epochs, and ℓ 2 regularization parameter used in our study in Table 7.

Compute infrastructure. Our experiments were performed across a combination of Nvidia T4, A6000, and V100 GPUs.

D Additional Results in Toy Setup

In this section we will first give more details on our simplified setup that captures both contrastive pretraining and self-training in the same framework. Then, we provide some additional empirical results that are not captured theoretically but mimic behaviors observed in real world settings, highlighting the richness of our setup.

D.1 Detailed description of our simplified setup

In this subsection, we will first re-iterate the problem setup in Sec. 4 and provide some comparisons between our setup and those in closely related works. We will then describe the four methods: ERM,

ST, CL, and STOC, providing details on the exact estimates returned by these algorithms in the SSL and UDA settings.

Data distribution. We consider binary classification and model the inputs as consisting of two kinds of features: x ' r x in , x sp s where x in P R d in is the invariant feature that is predictive of the label across both source P S and target P T and x sp P R d sp is the spurious feature that is correlated with the label y only on the source domain P S but uncorrelated with label y in P T . Here, x in P R d in determines the label using the ground truth classifier w ‹ ' Unif p S d in ´ 1 q , and x sp P R d sp is strongly correlated with the label on source but random noise on target. Formally, we sample y ' Unif t´ 1 , 1 u and generate inputs x conditioned on y as follows

where γ is the margin afforded by the invariant feature. We set covariance of the invariant features Σ in ' σ 2 in ¨ p I d in ´ w ‹ w ‹J q to capture structure in the invariant feature that the variance is less along the latent predictive direction w ‹ . Note that the spurious feature is completely predictive of the label in the source data, and is distributed as spherical Gaussian in the target data with Σ sp ' σ 2 sp I d sp .

Why is our simplified setup interesting? In our setup, x in is the hard to learn feature that generalizes from source to target. The hardness of learning this feature is determined by the value of the margin γ and how it compares with size of the spurious feature ( a d sp ). Since, γ { a d sp is small in our setup, x in is much harder to learn on source data (even with population access) compared to the spurious feature x sp which generalizes poorly from source to target. These two types of features have been captured in similar analysis on spurious correlations [71, 60] since it imitates pitfalls emanating from the presence of spurious features in real world datasets ( e.g ., the easy to learn background feature in image classification problems). While this setup is simple, it is also expressive enough to elucidate both self-training and contrastive learning behaviors we observe in real world settings. Specifically, it captures the separation results we observe in Sec. 3.

Differences of our setup with prior works. While our distribution shift settings bears the above similarities it also has important differences with works analyzing self-training and contrastive pretraining individually. Chen et al. [17] analyze the iterative nature of self-training algorithm, where the premise is that we are given a classifier that not only has good performance on source data but in addition does not rely too much on the spurious feature. Under the strong condition of small norms along the spurious feature, they show that self-training can provably unlearn this small dependence when the target data along the spurious feature is random noise. This assumption is clearly violated in setups where the spurious correlation is strong (as in our toy setup), i.e ., the dependence on the spurious feature is rather large (much larger than that on the invariant feature) for any classifier that is trained directly on source data. Consequently, we show the need for 'good' pretrained representations from contrastive pretraining over which if we train a linear predictor (using source labeled data), it will provably have a reduced 'effective' dependence on the spurious feature.

Using an augmentation distribution similar to ours, Saunshi et al. [73] carried out contrastive pretraining analysis with the backbone belonging to a capacity constrained function class (similar analysis also in [40]). Our setup differs from this in two key ways: (i) we specifically consider a distribution

shift from source to target. Unlike their setting, it is not sufficient to make augmentations consistent with ground truth labels, since the predictor that uses just the spurious feature also assigns labels consistent with both ground truth predictions and augmentations on the source data; and (ii) our augmentation distribution assumes no knowledge of the invariant feature, which is why we augment all dimensions uniformly, as opposed to selectively augmenting a set of dimensions. In other words, we assume no knowledge of the structure of the optimal target predictor. For e.g ., if we had knowledge of the spurious dimensions we could have just selectively augmented those. Assuming knowledge of these perfect augmentations is not ideal for two reasons: (a) it makes the problem so easy that just training an ERM model on source data with these augmentations would already yield a good target predictor (which rarely happens in practice); and (b) in real-world datasets perfect augmentations for the downstream task are not known. Hence, we stick to generic augmentations in our setup.

D.2 Discussion on self-training and contrastive learning objectives

In text we will describe our objectives and methods for the UDA setup. In Table 8 we constrast the differences in the methods and objectives for SSL and UDA setups. Recall from Section 2 that we learn linear classifiers h over features extractors Φ . We consider linear feature extractor i.e. Φ is a matrix in R k ˆ d . For mathematical convenience, we assume access to infinite unlabeled data and hence replace the empirical quantities over unlabeled data with their population counterpart. In the UDA setting, we further assume access to infinite labeled data from the source. Note that due to distribution shift between source and target, 'ERM' on infinite labeled data from the source does not necessarily achieve optimal performance on the target. For binary classification, we assume that the linear layer h maps features to a scalar in R such that the prediction is sgn p h J Φ x q . We use the exponential loss ℓ p f p x q , y q ' exp p´ yf p x qq as the classification loss.

Contrastive pretraining. We obtain Φ cl : ' arg min Φ L cl p Φ q by minimizing the Barlow Twins objective [93], which prior works have shown is also equivalent to spectral contrastive and noncontrastive objectives [33, 11]. In Sec. 4, we consider a constrained form of Barlow Twins in (3) which enforces representations of different augmentations a 1 , a 2 of the same input x to be close in representation space, while ensuring feature diversity by staying in the constraint set. We assume a strict constraint on regularization p ρ ' 0 q for the theoretical arguments in the rest of the main paper. In App. E.1.2 we prove that all our claims hold for small ρ as well. In (7), we redefine the pretraining objective with a regularization term (instead of a constraint set) where κ controls the strength of the regularization term, with higher values of κ corresponding to stronger constraints on feature diversity. We then learn a linear classifier h cl over Φ cl to minimize the exponential loss on labeled source data.

1

4

7

10 13 16 19 22 25

k

:

CL feature dimension

20

40

60

80

100

CL vs. STOC across

CL dimension

k

CL

STOC

ERM

0

.

05

0

.

15

0

.

25

0

.

35

0

.

45

κ

:

CL regularization strength

40

50

60

70

80

90

100

CL vs. STOC across

CL regularization strength

κ

CL

STOC

ERM

Augmentations. Data augmentations play a key role in contrastive pre-training (and also as we see later, state-of-the-art self-training variants like FixMatch). Given input x P X , let P A p a | x q denote the distribution over its augmentations, and P A denote the marginal distribution over all possible augmentations. We use the following simple augmentations where we scale the magnitude of each co-ordinate by a uniformly independent amount, i.e .,

The performance of different methods heavily depends on the assumptions we make on augmentations. We try to mirror practical settings where the augmentations are fairly 'generic', not encoding any information about which features are invariant or spurious, and hence perturb all features symmetrically.

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:

For convenience, we keep the feature backbone Φ fixed across the self-training iterations and only update the linear head on the pseudolabels.

STOC(Self-training after contrastive learning). Finally, we can combine the two unsupervised objectives where we do the self-training updates( 2) with h 0 ' h cl and Φ 0 ' Φ cl starting with the contrastive learning model rather than just source-only ERM. Here, we only update h and fix Φ cl .

D.3 Additional empirical results in our simplified setup

We conduct two ablations on the hyperparameters for contrastive pretraining. First, we vary the dimensionality k of the linear feature extractor Φ P R k ˆ d . Second, we vary the regularization strength κ that enforces feature diversity in the Barlow Twins objective (7). In Figure 5 we plot these ablations in the UDA setup.

Varying feature dimension. We find that CL recovers the full set of predictive features ( i.e . both spurious and invariant) only when k is large enough (Figure 5 (left) ). Since the dimensionality of the true feature is 5 in our Example 1, reducing k below the true feature dimension hurts CL. Once k crosses a certain threshold, CL features completely capture the projection of the invariant feature w in . After this point, it amplifies the component along w in . It retains the amplification over the spurious feature w sp even as we increase k . This is confirmed by our finding that further increasing k does not hurt CL performance. This is also inline with our theoretical observations, where we find that for

ERM

ST

CL

STOC

60

80

100

OOD accuracy

53.7

51.2

79.4

98.3

Results in UDA setup

ERM

ST

CL

STOC

0

25

50

75

100

ID accuracy

87.7

51.1

98.1

99.1

Results in SSL setup

suitable w ‹ , the subspace spanned by w in and w sp are contained in a low rank space (as low as rank 2 ) of the contrastive representations (Theorem 5). Once CL has amplified the dependence along w in STOC improves over CL by unlearning any remaining dependence on the spurious w sp . The above arguments for the CL trend also explain why the performance of STOC continues to remain « 100% as we vary k .

Varying regularization strength. In our main theoretical arguments we consider the constrained form of the Barlow Twins objective (3) with a strict constraint of ρ ' 0 (we relax this theoretically as well, see E.1.2). For our experiments, we optimize the regularized version of this objective (7), where the constraint term now appears as a regularizer which enforces feature diversity, i.e . the features learned through contrastive pretraining span orthogonal parts of the input space (as governed under the metric defined by augmentation covariance matrix Σ A ). If κ is very low, then trivial solutions exist for the Barlow Twins objective. For e.g ., ϕ « 0 (zero vector) achieves very low invariance loss. When κ ă 0 . 05 , we find that CL recovers these trivial solutions (Figure 5 (right) ). Hence, both CL and STOC perform poorly. As we increase κ the performance of both CL and STOC improve, mainly because the features returned by Φ cl now comprise of the predictive directions w in and w sp , as predictive by our theoretical arguments for ρ ' 0 (which corresponds to large κ ). On the other hand, when κ is too high optimization becomes hard since κ directly effects the Lipschitz constant of the loss function. Hence, the performance of CL drops by some value. Note that this does not effect the performance of STOC since CL continues to amplify w in over w sp even if it is returning suboptimal solutions with respect to the optimization loss of the pretraining objective.

D.4 Reconciling Practice: Experiments with deep networks in toy setup

In this section we delve into the details of Sec. 4.5, i.e ., we analyze performance of different methods when we make some design choices that imitate practice. First, we look at experiments involving a deep non-linear backbone Φ . Here, the non-linear Φ is learned during contrastive pretraining and

ERM

ST

CL

STOC

60

80

100

OOD accuracy

53.7

51.2

72.6

96.3

Results in UDA setup

ERM

ST

CL

STOC

0

25

50

75

100

ID accuracy

87.7

51.1

99.7

99.8

Results in SSL setup

OOD accuracy

100

80

60

UDA setup: Full-finetuning STOC

(non-linear backbone)

99.2

77.6

CL

70.9

STOC w/o aug.

30

60

90

120 150 180

CL training epochs

60

80

100

ID Accuracy

SSL setup: Full-finetuning STOC

(linear backbone)

CL

STOC

ERM

STOC with aug.

fixed for CL and STOC. Then, we investigate trends when we continue to propagate gradients onto Φ during STOC (we call this full-finetuning). Unlike previous cases, this allows features to be updated.

Results with non-linear feature extractor Φ . In Fig. 7 we plot the performance of the four methods when we use a non-linear feature extractor during contrastive pretraining. This feature extractor is a one-hidden layer neural network (hidden dimension is 500) with ReLU activations. We find that the trends observed with linear backbones in Fig. 6 are also replicated with the nonlinear one. Specifically, we note that STOC improves over CL under distribution shifts, whereas CL is already close to optimal when there are no distribution shifts. We also see that CL and ST individually are subpar. In SSL, we see a huge drop in the performance of ST (over ERM) mainly because we only fit on pseudolabels during ST. This is different from practice where we continue to optimize loss on labeled data points while fitting the pseudolabels. Consequently, when we continue to optimize performance on source labeled data the performance of ST in SSL setup is improves from 51 . 1% Ñ 72 . 6% .

Results with full fine-tuning. Up till this point, we have only considered the case (for both SSL and UDA) where we fix the contrastive learned features when running CL and STOC, i.e ., we only optimized the linear head h . Now, we shall consider the setting where gradients are propagated to Φ during STOC. Note that we still fix the representations for training the linear head during CL. Results for this setting are in Figure 8. We show two interesting trends that imitate real world behaviors.

STOC benefits from augmentations during full-finetuning: In the UDA setup we find that ST while updating Φ cl can hurt due to overfitting issues when training with the finite sample of labeled and unlabeled data (drop by ą 7% over CL). This is due to overfitting on confident but incorrect pseudolabels on target data. This can exacerbate components along spurious feature w sp from source. One reasoning behind this is that deep neural networks can perfectly memorize them on finite unlabeled target data [94]. Heuristics typically used in practice ( e.g . in FixMatch [79]) help avoid overfitting on incorrect pseudolabels: (i) confidence thresholding; to pick confident pseudolabel examples; (ii) pseudolabel a different augmented input than the one on which the self-training loss is optimized; and (iii) optimize source loss with labeled data simultaneously when fitting pseudolabels. Intuitively, thresholding introduces a curriculum where we only learn confident examples in the beginning whose pseudolabels are mainly determined by component along the invariant feature w in . Augmentations prevent the neural network from memorizing incorrect pseudolabels and optimizing source loss prevents forgetting of features learned during CL. When we implement these during fullfinetuning in STOC we see that STOC now improves over CL (by ą 20% ).

Can we improve contrastive pretraining features during STOC? We find that self-training can also improve features learned during contrastive pretraining when we update the full backbone during STOC (see Figure 8 (right) ). Specifically, in the SSL setup we find that STOC can now improve substantially over CL. Recall, that when we fixed Φ cl this was not possible (see E.3 and Fig. 2(b)). This is mainly because STOC can now improve performance beyond just recovering the generalization gap for the linear head (which is typically small). This feature improvement is observed even when we fully finetune a linear feature extractor. Similar trends are also observed with the non-linear backbone. But, it becomes harder to identify a good stopping criterion for CL training. Thus, it remains unclear if STOC and CL have complementary benefits for feature learning in UDA or SSL settings. Investigating this is an interesting avenue for future work.

E Formal Statements from Sec. 4

Recall from Section 2 that we learn linear classifiers h over features extractors Φ . We consider linear feature extractor i.e. Φ is a matrix in R d ˆ k and the linear layer h : R k Ñ R with a prediction as sgn p h J Φ x q . We use the exponential loss ℓ p f p x q , y q ' exp p´ yf p x qq .

E.1 Analysis of ERM and ST: Formal Statement of Theorem 2

For ERM and ST, we train both h and Φ . This is equivalent to Φ ' I d ˆ d being identity and training a linear head h . Recall that the ERM classifier is obtained by minimizing the population loss on labeled source data:

We split Theorem 2 into Theorem 8 and Theorem 9. Before we characterize the ERM solution, we recall some additional notation. Define w in 'r w ‹ , 0 , ..., 0 s J , and w sp ' r 0 , ..., 0 , 1 d sp { ? d sp s J . The following proposition characterizes h ERM and 0-1 error of the classifier on target:

Theorem 8 (ERM classifier and its error on target) . ERM classifier obtained as in (10) is given by

The target accuracy of h ERM is given by 0 . 5 ¨ erfc ` ´ γ 2 {p ? 2 d sp ¨ σ sp q ˘ .

Proof. To prove this theorem, we first derive a closed-form expression for the ERM classifier and then use Lemma 29 to derive its 0-1 error on target. For Gaussian data with the same covariance matrices for class conditional P S p x | y ' 1 q and P S p x | y ' 0 q , Bayes decision rule is given by the Fisher's linear discriminant direction (Chapter 4; Bishop [9]):

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:

Starting with h 0 ST ' h ERM { | | h ERM | | 2 (the classifier obtained with ERM) we perform the following iterative procedure for self-training:

Next, we characterize ST solution:

Theorem 9 (ST classifier and its error on target) . Starting with ERM solution, ST will lead to:

(i) (Necessary condition) h t ST ' w sp as t Ñ8 , such that the target accuracy is 50% for all σ sp ě 1 and γ ď 1 2 ? σ sp .

(ii) (Sufficient condition) h t ST ' w in as t Ñ8 , such that the target accuracy is 100% when the problem parameters γ, σ sp satisfy: γ ě σ sp .

Proof. The proof can be divided into two parts: (i) deriving closed-form expressions for updates on h t ST in terms of h t ´ 1 ST and (ii) obtaining conditions under which the component along w in monotonically increases or decreases with t after re-normalizing the norm of updated h . For notation convenience, we denote h ST with h in the rest of the proof.

Part-1. First, the loss of self-training with classifier h : ' r h in , h sp s where h in P R d in and h sp P R d sp is given by:

ˇ

ˇ

'

E

y

'

U

t´

1

,

1

u

,z

'

N

p

0

,

1

q

exp

´

γ

¨

y

¨

h

J

in

w

where (16) to (17) is implied by simply replacing the definition of target distribution and (17) to (18) is implied by the symmetry of the function with respect to y and ´ y due to the symmetry of the absolute function and Gaussian distribution. For a classifier h t , we denote µ t ' γ ¨ h t in J w ‹ and σ t ' ' σ in p| | h t in | | 2 2 ´p h t in T w ‹ q 2 q ` σ sp ¨ ˇ ˇ ˇ ˇ h t sp ˇ ˇ ˇ ˇ 2 ı . With this notation, we can re-write the loss in (18) as L st p h t q ' E z ' N p 0 ,σ 2 t q r exp p´ | µ t ` z | qs .

Now we derive a closed-form expression of L st p h t q in Lemma 30:

Define the Mill's ratio as r p x q ' exp ` x 2 { 2 ˘ ¨ erfc ` x { ? 2 ˘ ¨ a π { 2 as in Baricz [5]. We will frequently use standard properties of the Mill's ratio. We list them in Lemma 21 for completeness. Define:

Let r h t ` 1 denote the un-normalized gradient descent update at iterate t ` 1 . We have:

'

`

‹

Now we will individually argue about the update of r h t ` 1 along the first d in dimensions and the last d sp dimensions. First, we have:

Notice that the update of h t ` 1 in is split into two components, one along w ‹ and the other along the orthogonal component 2 h t in ´ 2 p h t in J w ‹ q w ‹ . We will now argue that since at initialization, the component along p I ´ w ‹ w ‹J q is zero then it will remain zero. In particular, we have:

With (23), we can argue that if p I ´ w ‹ w ‹J q h t in ' 0 , then p I ´ w ‹ w ‹J q r h t ` 1 inv ' 0 implying that p I ´ w ‹ w ‹J q r h t in ' 0 for all t ą 0 . Hence, we have:

Second, we have the update r h t ` 1 sp given by:

Re-writing the expressions (25) and (26) for the update of r h t ` 1 , we have:

Here, we replace h t sp ' µ t ¨ w ‹ { γ in (25) to get (27). Updates in (27) and (28) show that r h t ` 1 inv remains in the direction of h t in and r h t ` 1 sp remains in the direction of h t sp .

Part-2. Now we will derive conditions under which h t in and h t sp will show monotonic behavior for necessary and sufficient conditions. We will first argue the condition under which ST will provably fail and converge to a classifier with a random target performance. For this, at every t , if we have:

ˇ

ˇ

ˇ

ˇ

then we can argue that as t Ñ8 , we have ˇ ˇ h t sp ˇ ˇ 2 ' 1 and hence, the ST classifier will have random target performance. Thus, we will focus on conditions, under which the norm on ˇ ˇ ˇ ˇ h t sp ˇ ˇ ˇ ˇ 2 increases with t . Re-writing (29), we have:

Plugging in (27) and (28) into (33), we get:

For small enough η , we have the necessary condition for the failure of ST as:

Now we show in Lemma 11 and Lemma 10 that if the conditions assumed in the theorem continue to hold, then we can success and failure respectively.

Lemma 10 (Necessary conditions for ST) . Define α 1 and α 2 as in (20) and (21) respectively. If σ sp ě 1 and γ ď 1 2 ? σ sp , then we have for all t :

Proof. We upper bound and lower bound α 1 and α 2 by using the properties of r p¨q . Recall:

and

We now use Taylor's expansion on r p¨q and we get:

and similarly, we get:

where R 2 ' r 2 p σ 0 q . This is because r 2 p¨q takes positive values and is a decreasing function in σ t (refer to Lemma 21). We now lower bound α 1 p µ t , σ t q and upper bound α 2 p µ t , σ t q :

Substituting the lower bound and upper bound in (36) gives us the following as stricter a necessary condition (i.e., (43) implies (36)):

˜

¸

Now, we will argue the monotonicity of LHS and RHS in (47). Observe that LHS is increasing in µ t and decreasing in σ t and RHS is decreasing in σ t as p σ t r p σ t q ´ 1 q is increasing (and the multiplier is negative). Moreover, if (47) holds true for maximum value of RHS and minimum of LHS, then we would have (36). Thus substituting µ t ' γ and σ t ' σ 0 in LHS and σ t ' σ sp in RHS, we get:

˜

¸

˜

¸

(50)

Taking γ ď 1 2 ? σ sp and substituting R 2 ' r 2 p σ 0 q :

Analytically solving the above expression, we get that (51) is satisfied for all values of σ sp ě 1 when d sp ě 1 . For example, the expression in (51) is also satisfied for the problem parameter used in the running example of the main paper.

As a remark, we note that in the proof of Lemma 10, the conditions derived are loose because of the relaxations made to simply the proof. In principle, the proof (and hence the conditions) can be tightened by carefully propagating second-order terms (which depend on σ t ) in (40).

Lemma 11 (Sufficiency conditions for ST) . Define α 1 and α 2 as in (20) and (21) respectively. If σ sp ď γ , then we have for all t :

Proof. We upper bound and lower bound α 1 and α 2 by using the properties of r p¨q . Recall:

and

We now use Taylor's expansion on r p¨q and we get:

and similarly, we get:

where R 2 ' r 2 p σ 0 q . This is because r 2 p¨q takes positive values and is a decreasing function in σ t (refer to Lemma 21). We now lower bound α 1 p µ t , σ t q and upper bound α 2 p µ t , σ t q :

Substituting the lower bound and upper bound in (52) gives us the following as stricter a sufficient condition (i.e., (59) implies (52)):

Hence, when ' 1 ´ γ 2 σ 2 sp ı ď 0 , we have condition in (64) hold true as r 1 p σ t q is always negative. Hence, the condition γ ě σ sp gives us the necessary condition.

E.1.1 Proof of Proposition 3

For convenience, we first restate the Proposition 3 which gives us a closed form solution for (3) when ρ ' 0 . Then, we provide the proof, focusing first on the case of k ' 1 , and then showing that extension to k ą 1 is straightforward and renders the final form in the proposition that follows.

Proposition 12 (Barlow Twins solution) . The solution for (3) is U J k Σ ´ 1 { 2 A where U k are the top k eigenvectors of Σ ´ 1 { 2 A r ΣΣ ´ 1 { 2 A . Here, Σ A : ' E a ' P A r aa J s is the covariance over augmentations, and r Σ : ' E x ' P U r r a p x q r a p x q J s is the covariance matrix of mean augmentations r a p x q : ' E P A p a | x q r a s .

Proof. We will use ϕ p x q to denote ϕ J x where ϕ P R d . Throughout the proof, we use a to denote augmentation and x to denote the input. We will use P A p a | x q as the probability measure over the space of augmentations A , given some input x P X (with corresponding density) p A p¨ | x q . Next, we use p A p¨q to denote the density associate with the marginal probability measure over augmentations: P A ' ş X P A p a | x q dP U . Finally, the joint distribution over positive pairs A ` p a 1 , a 2 q ' ş X P A p a 1 | x q P A p a 2 | x q dP U , gives us the positive pair graph over augmentations.

Before we solve the optimization problem in (3) for Φ P R k ˆ d for any general k , let us first consider the case where k ' 1 , i.e . we only want to find a single linear projection ϕ . The constraint ρ ' 0 , transfers onto ϕ in the following way:

Under the above constraint we want to minimize the invariance loss, which according to Lemma 22 is given by 2 ¨ ş A ϕ p a q L p ϕ qp a q dP A , where L p ϕ qp¨q is the following linear operator.

Based on the definition of the operator, we can reformulate the constrained optimization for contrastive pretraining as:

where r ϕ p x q ' E a ' P A p¨| x q ϕ p x q ' E c ' Unif r 0 , 1 s d r ϕ J p c d x qs . Note that,

Further, since E a ' P A r ϕ p a q 2 s ' ϕ J Σ ϕ we can now rewrite our main optimization problem for k ' 1 as:

Recall that in our setup both r Σ and Σ A are positive definite and invertible matrices. To solve the above problem, let's consider a re-parameterization: ϕ 1 ' Σ 1 { 2 A ϕ , thus ϕ J Σ A ϕ ' 1 , is equivalent to the constraint } ϕ 1 } 2 2 ' 1 . Based on this re-parameterization we are now solving:

Thus,

which is nothing but the top eigenvector for Σ ´ 1 { 2 A ¨ r Σ ¨ Σ ´ 1 { 2 A .

Now, to extend the above argument from k ' 1 to k ą 1 , we need to care of one additional form of constraint in the form of feature diversity: ϕ J i Σ A ϕ j ' 0 when i ‰ j . But, we can easily redo the reformulations above and arrive at the following optimization problem:

where ϕ 1 i ' Σ 1 { 2 A ϕ i . The above is nothing but the top k eigenvectors for the matrix Σ ´ 1 { 2 A ¨ r Σ ¨ Σ ´ 1 { 2 A . This completes the proof of Proposition 12.

E.1.2 Analysis with ρ ą 0 in Contrastive Pretraining Objective (3)

In (3) we considered the strict version of the optimization problem where ρ ' 0 . Here, we will consider the following optimization problem that we optimize for our experiments in the simplified setup:

where κ ą 0 is some finite constant (note that every ρ corresponds to some κ and particularly ρ ' 0 , corresponds to κ ' 8 ). Let Φ ‹ be the solution for (3) with ρ ' 0 , i.e . the solution described in Proposition 3. Now, we will show that in practice we can provably recover something close to Φ ‹ when κ is large enough.

Theorem 13 (Solution for (78) is approximately equal to Φ ‹ ) . If p Φ is some solution that achieves low values of the objective L cl p Φ , κ q in (78) , i.e ., L cl p p Φ , κ q ď ϵ , then there exists matrix W P R k ˆ k such that:

where γ k ` 1 is the the k ` 1 th eigenvalue for I d ´ Σ ´ 1 { 2 A r Σ Σ ´ 1 { 2 A . Here, λ 1 ď λ 2 ď . . . ď λ d .

Proof. Since we know that L cl p p Φ , κ q ď ϵ , we can individually bound the invariance loss and the regularization term:

Let ϕ ‹ 1 , ϕ ‹ 2 , ϕ ‹ 3 , . . . , ϕ ‹ d be the solution returned by the analytical solution for ρ ' 0 , i.e . the solution in Proposition 3. Now, since Φ ‹ would span R d when Σ A is full rank, we can denote:

Now from Lemma 22, the invariance loss for p ϕ i can be written using the operator L p ϕ qp a q ' ϕ p a q ´ ş A A ` p a,a 1 q p A p a q ϕ p a 1 q d a 1 :

Since, ϕ ‹ i p¨q are eigenfunctions of the operator L [38], we can conclude that:

and if γ 1 ď γ 2 ď γ 3 . . . ď γ d are the eigenvalues for ϕ ‹ 1 , ϕ ‹ 2 , ϕ ‹ 3 , . . . , ϕ ‹ d under the decomposition of L p ϕ qp¨q then:

Recall, we are also aware of a condition on the regularization term: 1 ´ a ϵ κ ď p ϕ J i Σ A p ϕ i ď 1 ` a ϵ κ .

In order to show that the projection of p ϕ i on Φ ˚ is significant, we need to argue that the term ř d j ' k ` 1 ´ η p j q i ¯ 2 is small. The argument for this begins with the condition on invariance loss, and the fact that γ 1 ď γ 2 ď . . . ď γ k ď γ k ` 1 ď . . . ď γ d :

Extending the above result @ i by simply adding the bounds completes the claim of our first result in Theorem 13. Next, we will lower bound the eigenvalue γ k ` 1 . Recall that, ř k j ' 1 ´ η p j q i ¯ 2 ě 1 ´ a ϵ κ ´ ϵ 2 γ k ` 1 . Thus,

We assume that all eigenvalues are strictly positive, which is true under our augmentation distribution. Given, γ k ` 1 ě γ 1 , we can rearrange the above to get:

This completes the claim of our second result in Theorem 13.

E.1.3 Proof of Theorem 5

In this section, we prove our main theorem about the recovery of both spurious w sp , invariant w in features by the contrastive learning feature backbone, and also the amplification of the invariant over the spurious feature (where amplification is defined relatively with respect to what is observed in the data distribution alone). We begin by defining some quantities needed for analysis, that are fully determined by the choice of problem parameters for the model in (6).

From Section 4, we recall the definitions of w in : ' r w ‹ , 0 , . . . , 0 s and w sp : ' r 0 , . . . 0 , w 1 s where w 1 ' 1 d sp { a d sp . Let us now define u 1 , u 2 as the top two eigenvectors of Σ A with eigenvalues λ 1 , λ 2 ą 0 , (note that in our problem setup both Σ A and r Σ are full rank positive definite matrices), and τ : ' a λ 1 { λ 2 . Next we define α as the angle between u 1 and w in , i.e ., cos p α q ' u J 1 w in . Based on the definitions of α and τ , both of which are fully determined by the eigen decomposition of the post-augmentation feature covariance matrix Σ A , we now restate Theorem 5:

Theorem 14 (Formal; CL recovers both invariant w in and spurious w sp but amplifies w in ) . Under Assumption 4 p w ‹ ' 1 d in { ? d in q , the CL solution Φ cl 'r ϕ 1 , ϕ 2 , ..., ϕ k s satisfies ϕ J j w in ' ϕ J j w sp ' 0 @ j ě 3 . For τ, α as defined above, the solution for ϕ 1 , ϕ 2 is:

where 0 ď α, θ ď π { 2 . Let us redefine ϕ 1 ' c 1 w in ` c 3 w sp and ϕ 2 ' c 2 w in ` c 4 w sp .

For constants K 1 , K 2 ą 0 , γ ' K 1 K 2 { σ sp , d sp ' σ 2 sp { K 2 2 , @ ϵ ą 0 , D σ sp 0 , such that for σ sp ě σ sp 0 :

where L ' 1 ` K 2 2 .

Proof. We will first show that the only components of interest are ϕ 1 , ϕ 2 . Then, we will prove conditions on the amplification of w in over w sp in ϕ 1 , ϕ 2 . Following is the proof overview:

When w ‹ ' 1 d in { ? d in , from the closed form expressions for Σ A and r Σ , show that the solution returned by solving the Barlow Twins objective depends on w in and w sp only through the first two components ϕ 1 , ϕ 2 .

For the components ϕ 1 , ϕ 2 , we will show that the dependence along w in is amplified compared to w sp when the target data sufficiently denoises the spurious feature ( i.e ., σ sp is sufficiently large).

Part-I:

We can divide the space R d into two subspaces that are perpendicular to each other. The first subspace is W ' t b 1 ¨ w in ` b 2 ¨ w sp : b 1 , b 2 P R u , i.e . the rank 2 subspace spanned by w in and w sp . The second subspace is W K where W K ' t u P R d : u J w in ' 0 , u J w sp ' 0 u . Then, from Lemma 23 we can conclude that the matrix Σ A can be written as:

where Σ A W K : ' E a ' P A ' Π W K p a qp Π W K p a qq J ‰ is the covariance matrix in the null space of W , and Π W K p a q is the projection of augmentation a into the null space of W , i.e . the covariance matrix in the space of non-predictive (noise) features. Similarly we can define:

Here again r Σ W K : ' E x ' P U ' Π W K p E c ' Unif r 0 , 1 s d p c d x qqp Π W K p E c ' Unif r 0 , 1 s d p c d x qqq J ‰ is the covariance matrix of mean augmentations after they are projected onto the null space of predictive features. The above decomposition also follows from result in Lemma 23.

From Proposition 3, the closed form expression for the solution returned by optimizing the Barlow Twins objective in (3) is U J Σ ´ 1 { 2 A where U are the top-k eigenvectors of:

When w ‹ ' 1 d in { ? d in , then Σ A W K ' r Σ W K ` B where B is a diagonal matrix with diagonal given by 1 3 ¨ diag p r Σ W K q . Further, since diag p r Σ W K q ' p ¨ ✶ d for some constant p ą 0 , the eigenvectors of r Σ W K and Σ A W K are exactly the same. Hence, when we consider the SVD of the expression Σ ´ 1 { 2 A r ΣΣ ´ 1 { 2 A , the matrices Σ A W K and r Σ W K have no effect on the SVD components that lie along the span of the predictive features. In fact, we only need to consider two rank 2 matrices (first terms in (98), (97)) and only do the SVD of Σ ´ 1 { 2 A W ¨ r Σ W ¨ Σ ´ 1 { 2 A W .

There are only two eigenvectors of Σ ´ 1 { 2 A W ¨ r Σ W ¨ Σ ´ 1 { 2 A W . We use λ 1 , λ 2 to denote the eigenvalues of Σ A W , and r cos p α q w ‹ , sin p α q w 1 s J , r sin p α q w ‹ , ´ cos p α q w 1 s J for the corresponding eigenvectors. Similarly, we use r λ 1 , r λ 2 to denote the eigenvalues of r Σ W , and r cos p β q w ‹ , sin p β q w 1 s J , r sin p β q w ‹ , ´ cos p β q w 1 s J for the corresponding eigenvectors. Let SVD U p¨q denote the operation of obtaining the singular vectors of a matrix. Then, to compute the components of the final expression: SVD U p Σ ´ 1 { 2 A r ΣΣ ´ 1 { 2 A q J Σ ´ 1 { 2 A that lies along the span of predictive features (in W ), we need only look at the decomposition of the following matrix:

Based on the above definitions of θ, α, λ 1 , λ 2 , we can then formulate ϕ 1 and ϕ 2 in the following way:

To summarize, using arguments in Lemma 23 and the fact that w ‹ ' 1 d in { ? d in , we can afford to focus on just two rank two matrices Σ A W , r Σ W in the operation: SVD U p Σ ´ 1 { 2 A q r ΣΣ ´ 1 { 2 A . The other singular vectors from the SVD only impact directions that span W K , and the singular vectors obtained by considering only the rank 2 matrices lie only in the space of W .

Part-II:

From the previous part we obtained forms of ϕ 1 , ϕ 2 in terms of: λ 1 , λ 2 , α, θ , all of which are fully specified by the SVD of Σ A W and r Σ W . If we define τ : ' ? λ 1 ? λ 2 , we can evaluate c 1 , c 2 , c 3 , c 4 as:

Now, we are ready to begin proofs for our claims on the amplification factors, i.e . on the ratios c 1 { c 3 , | c 2 { c 4 | .

We will first prove some limiting conditions for c 1 { c 3 , followed by those on | c 2 { c 4 | . For each of these conditions we will rely on the forms for c 1 , c 2 , c 3 , c 4 derived in the previous part, in terms of α, θ, τ (where 0 ď α, θ ď π { 2 ). We will also rely on some lemmas that characterize the asymptotic behavior of α, θ and τ as we increase σ sp . We defer the full proof of these helper lemmas to later sections.

Asymptotic behavior of c 1 { c 3 .

From Lemma 25 and Lemma 26, when γ ' K 1 { ? z and σ sp ' K 2 ? z , then:

where we apply Moore-Osgood when applying limits on intermediate forms. We can do this since τ tan θ approaches a constant, and each of cot α, τ and tan θ are continuous and smooth functions of z (see Lemma 24).

Asymptotic behavior of | c 2 { c 4 | .

When we consider the limiting behavior of c 2 { c 4 z , as we increase z or equivalently σ sp when γ ' K 1 { ? z and σ sp ' K 2 ? z , then we get:

From Lemma 26, cot α tan θ Ñ 0 . Next, if we consider lim z Ñ8 z tan θ { τ ' lim z Ñ8 τ tan θ ¨ z { τ 2 . For z { τ 2 , we invoke Lemma 28, which states that when γ ' K 1 { ? z and σ sp ' K 2 ? z , then:

Further, in our bound on c 1 { c 3 , we derived that τ tan θ Ñ K 1 K 2 2 { p 1 ` K 2 2 q 2 σ 2 in p 1 ´ 1 { d in q . Once again using Moore-Osgood we can plug this along with (108) to get:

Since z ' K 1 ? d sp { γ ,

Since both c 1 { c 3 and | c 2 { c 4 | are continuous functions of z , with lim inf z Ñ8 and lim sup z Ñ8 converging to the limits in 106 and 107 for both quantities respectively, we conclude that @ ϵ ą 0 there exists σ sp 0 such that for all σ sp ě σ sp 0 , the following is true:

This completes both Part-I and Part-II of the proof for Theorem 5.

E.1.4 Proof of Corollary 6

Corollary 15 (CL improves OOD error over ERM but is still imperfect) . For γ, σ sp , d sp defined as in Theorem 5, D σ sp 1 such that @ σ sp ě σ sp 1 , the target accuracy of CL (linear predictor on Φ cl ) is ě 0 . 5 erfc p´ L 1 ¨ γ { ? 2 σ sp q and ď 0 . 5 erfc p´ 4 L 1 ¨ γ { ? 2 σ sp q , where L 1 ' K 2 2 K 1 { σ 2 in p 1 ´ 1 { d in q . When σ sp 1 ą σ in a 1 ´ 1 { d in , the lower bound on accuracy is strictly better than ERM from scratch.

Proof. Recall from Theorem 14, all ϕ j , for j ě 3 , lie in the null space of w in and w sp . Since, the predictive features are strictly contained in the rank two space spanned by w in and w sp , without loss of generality we can restrict ourselves to the case where k ' 2 , and when doing training a head h ' r h 1 , h 2 s J P R 2 over contrastive pretrained representations using source labeled data, we get the following max margin solution:

Without loss of generality we can divide both h 1 and h 2 by h 1 and get the final classifier to be ϕ 1 ` h 2 h 1 ¨ ϕ 2 :

From Lemma 29, we can derive the target accuracy of the classifier h on top of CL representations to be the following:

Finally, from Lemma 27, when γ ' K 1 { ? z and σ sp ' K 2 ? z , then:

Plugging, 109 and 110 into 107 we get the following limit:

where β ' p c 2 γ ` c 4 ? d sp q { p c 1 γ ` c 3 ? d sp q .

Substituting β into the expression c 1 ` βc 2 c 3 ` βc 4 we get:

We first substitute expressions for c 1 , c 2 , c 3 , c 4 from (102), (103), (104) and (105) in the above expression. Then for γ ' K 1 { ? z, σ sp ' K 2 ? z , we substitute the expressions for cot α , tan θ , and τ ' λ 1 { λ 2 with their corresponding closed form expressions (as functions of z ) from Lemma 24. On the resulting expression we apply do repeated applications of L'Hôpital's rule to get the following result:

Based on γ, d sp , σ sp defined in Theorem 5, and (119) we can conclude that D σ sp 1 such that for all σ sp ě σ sp 1 :

Finally, applying (120) to Lemma 29, we conclude the following: When γ ' K 1 K 2 { σ sp , d sp ' σ 2 sp { K 2 2 , there exists σ sp 1 , such that for any σ sp ě σ sp 1 , target accuracy of CL is at least 0 . 5 erfc ´ ´ L 1 ¨ γ ? 2 σ sp ¯ and at most 0 . 5 erfc ´ ´ 4 L 1 ¨ γ ? 2 σ sp ¯ , where L 1 ' K 2 2 K 1 σ 2 in p 1 ´ 1 { d in q .

Comparison with ERM. Recall from Theorem 8 the performance of ERM classifier (trained from scratch) is 0 . 5 erfc ` ´ γ 2 { ? 2 d sp σ sp ˘ . The lower bound on the performance of classifier over CL representations is strictly better than ERM when:

This completes our proof of Corollary 6.

E.2 Analysis of STOC: Formal Statement of Theorem 7

Recall ERM solution over contrastive pretraining. We showed that without loss of generality when k (the output dimensionality of Φ ) is greater than 2, we can restrict k to 2 and the Φ can be denoted as r ϕ 1 , ϕ 2 s J where ϕ 1 ' c 1 w ‹ ` c 3 w sp and ϕ 2 ' c 2 w ‹ ` c 4 w sp . The ERM solution of the linear head is then given by h 1 , h 2 P R :

STOC performs self-training of the linear head over the CL solution. Before introducing the result, we need some additional notation. Let h t denote the solution of the linear head at iterate t . Without loss of generality, assume that the coefficients in ϕ 1 ' c 1 w in ` c 3 w sp and ϕ 2 ' c 2 w in ` c 4 w sp are such that c 2 is positive and c 1 , c 3 , and c 4 are negative. Moreover, for simplicity of exposition, assume that | c 4 | ą | c 3 | .

Theorem 16. Under the conditions of Corollary 15 and when γ 2 σ sp ě ' ´ c 3 ´ c 4 p c 2 ` c 1 q¨ | c 1 | ı _ ' c 4 c 1 ¨ c 2 ı , the target accuracy of ST over CL is lower bounded by 0 . 5 ¨ erfc p´ | c 2 { c 4 | ¨ γ { p ? 2 σ 2 q q ě 0 . 5 ¨ erfc ` ´ L ¨ ? d sp { p ? 2 σ sp q ˘ with L ě 1 .

Before proving Theorem 16, we first connect the condition γ 2 σ sp ě ' ´ c 3 ´ c 4 p c 2 ` c 1 q¨ | c 1 | ı _ ' c 4 c 1 ¨ c 2 ı with the result obtained with contrastive learning.

Remark 1. We first argue that ' ´ c 3 ´ c 4 p c 2 ` c 1 q¨ | c 1 | ı term dominates and hence, if we have γ 2 σ sp ě ' ´ c 3 ´ c 4 p c 2 ` c 1 q¨ | c 1 | ı , then we get the result in Theorem 16. First, recall that as σ sp increases, we have ∣ ∣ ∣ c 3 c 1 ∣ ∣ ∣ converge to 2 Lσ 2 in p d in ´ 1 q K 1 K 2 2 d in , c 2 Ñ 1 and c 1 c 2 Ñ 0 . Using these limits, we get:

which reduces the following condition: d sp ď K 2 1 K 2 { 3 2 ¨ ´ d in 2 Lσ 2 in p d in ´ 1 q ¯ 2 { 3 .

Proof. First, we create an outline of the proof. We argue about the updates of h t showing that both h t 1 and h t 2 increase with | h t 2 | becoming greater than | h t 1 | for some large t . Then we show that | h t 2 | ě | h t 1 | is sufficient to obtain near-perfect target generalization.

Part 1. Recall the loss of used for self-training of h :

Define µ t ' c 1 γh t 1 ` c 2 γh t 2 and σ t ' c 3 σ sp h t 1 ` c 4 σ sp h t 2 . With this notation, we can re-write the loss in (125) as L st p h t q ' E z ' N p 0 ,σ 2 t q r exp p´ | µ t ` z | qs .

Similar to the the treatment in Theorem 9, we now derive a closed-form expression of L st p h t q in Lemma 30:

Define:

Let r h t ` 1 denote the un-normalized gradient descent update at iterate t ` 1 . We have:

Now we will individually argue about the update of r h t ` 1 . First, we have:

and second, we have:

We will now argue the conditions under which h t ` 1 2 increases till its value reaches 1 { ? 2 . In particular, we will argue that when h t 2 is negative, the norm | h t 2 | decreases and when h t 2 becomes positive, then its norm increases. We show that the following three conditions are sufficient to argue the increasing value of h t 2 : for all t , we have (i) µ t ě µ c and | σ t | ă σ c for constant µ c ' | c 1 ¨ γ | { 2 and σ c ' | c 4 σ sp | ; (ii) δ 2 ă 0 ; (iii) | δ 2 | ě δ 1 . In Lemma 18, we argue that our assumption on the initialization of the backbone learned with BT implies the previous three conditions.

Case-1. When h t 2 is negative (and after the update, it remains negative). Then we want to argue the following:

ñ

ñ

ñ

ñ

(137)

Since δ 2 ă 0 , | δ 2 | ě | δ 1 | and h t 2 ă h t 1 ă 0 , we have r ηδ 2 p h t 1 q ´ ηδ 1 p h t 2 qs as positive. This implies inequality (140) to (141) and for small enough η , (141) will continue to hold true.

Case-2. When h t 2 is positive but less than 1 { ? 2 . Then we want to argue the following:

ñ

ñ

ñ

ñ

ñ

ñ

Since δ 2 ă 0 , | δ 2 | ě | δ 1 | , h t 1 ď ´ 1 { ? 2 and 0 ă h t 2 ă 1 { ? 2 , we have r ηδ 2 p h t 1 q ´ ηδ 1 p h t 2 qs as positive. This implies inequality (149) to (150). Focusing on (150), we note that h t 1 ¨ δ 2 is positive and greater in magnitude than h t 2 ¨ δ 1 . Moreover, since h t 2 h t 1 is negative, (150) will continue to hold true.

Now, when h t 2 is positive and greater than 1 { ? 2 , then h t 2 will stay in that region. Convergence of STOC together with conditions of convergence as in Lemma 17 will imply that the at convergence h t 2 will remain greater than 1 { ? 2 , such that h t c 1 h t c 2 ' δ 1 δ 2 . Now we bound the target error of STOC.

Part 2. To bound the accuracy at any iterate t when h t 2 ě 1 { ? 2 , we have from Lemma 29:

We now upper bound and lower bound the fraction c 1 γh t 1 ` c 2 γh t 2 | c 3 σ sp h t 1 ` c 4 σ sp h t 2 | in RHS in (151): (i) c 1 γh t 1 ` c 2 γh t 2 ě c 2 γh t 2 since both c 1 γh t 1 and c 2 γh t 2 have same sign; (ii) | c 3 σ sp h t 1 ` c 4 σ sp h t 2 | ď | c 4 σ sp h t 2 | because | c 4 σ sp h t 2 | ě | c 3 σ sp h t 1 | and they have opposite signs. Hence, from (151), we have:

Substituting the definition of erfc , the expression (152) gives us the required lower bound on the target accuracy.

Lemma 17 (Convergence of STOC) . Assume the gradient updates as in (131) and (132) . Then STOC converges at t ' t c when h t c 1 h t c 2 ' δ 1 δ 2 . For t ą t c , (131) and (132) make no updates to the linear h .

Proof. When the gradient updates δ 1 and δ 2 are such that h t ` 1 1 matches h t 1 , we have convergence of STOC.

ñ

ñ

ñ

ñ

ñ

ñ

(157)

Thus either r ηδ 2 p h t 1 q ´ ηδ 1 p h t 2 qs ' 0 or r ηδ 2 p h t 1 q ` ηδ 1 p h t 2 qs ' 2 h t 2 h t 1 . Since η is such that h 1 ´ ηδ 1 ă 0 , r ηδ 2 p h t 1 q ` ηδ 1 p h t 2 qs ‰ 2 h t 2 h t 1 implying that r ηδ 2 p h t 1 q ´ ηδ 1 p h t 2 qs ' 0 giving us the required condition.

Lemma 18. Under the initialization conditions assumed in Theorem 16, for all t , we have: (i) µ t ě µ c and | σ t | ď σ c for constant µ c ' | c 1 ¨ γ | { 2 and σ c ' | c 4 σ sp | ; (ii) δ 2 ă 0 ; (iii) | δ 2 | ě δ 1 , where δ 1 ' A 1 ¨ p σ t c 3 σ sp ´ c 1 γ q ` A 2 ¨ p σ t c 3 σ sp ` c 1 γ q ´ A 3 c 3 σ sp and δ 2 ' A 1 ¨ p σ t c 4 σ sp ´ c 2 γ q ` A 2 ¨ p σ t c 4 σ sp ` c 2 γ q ´ A 3 c 4 σ sp for A 1 , A 2 and A 3 defined in (127) , (128) , and (129) .

Proof. Recall, µ t ' c 1 γh t 1 ` c 2 γh t 2 and σ t ' c 3 σ sp h t 1 ` c 4 σ sp h t 2 . First, we argue that µ t increases from the initialization value. Notice that µ 0 ' c 1 γh 0 1 ` c 2 γh 0 2 . Due to Corollary 15, we have h 0 2 ě 0 . And since | c 2 | ą | c 1 | , we get µ 0 ě | c 1 γ | as both c 1 and h 0 1 are of same sign. Moreover, as training progresses with h t 1 remaining negative and h t 2 remaining positive, we have µ t stays greater than µ 0 .

Recall the definition of A 1 , A 2 , and A 3 in (127), (128), and (129). Moreover, recall the definition of α 1 p µ t , σ t q and α 2 p µ t , σ t q :

and

Thus, we have α 1 p µ t , σ t q ¨ A 3 ' A 1 ¨ σ t and α 2 p µ t , σ t q ¨ A 3 ' σ t ¨ ´ A 2 ¨ ´ 2 σ t A 3 ¯ . Replacing the definition of A 1 , A 2 , and A 3 in δ 1 and δ 2 , we get:

We now upper bound and lower bound α 1 and α 2 by using the properties of r p¨q . We use Taylor's expansion on r p¨q and we get:

and similarly, we get:

where R 2 ' r 2 p σ 0 q . This is because r 2 p¨q takes positive values and is a decreasing function in σ t (refer to Lemma 21). We now lower bound α 1 p µ t , σ t q and upper bound α 2 p µ t , σ t q :

Part-1. We first prove that δ 2 ď 0 . Substituting the lower bound and upper bound in (163) gives us the following as stricter a sufficient condition (i.e., (168) implies δ 2 ď 0 ):

Thus, if we have µ t ě σ sp ¨p´ c 4 q γ ¨ c 2 , then (168) holds true.

Part-2. Next, we prove that | δ 2 | ě δ 1 . Substituting the lower bound and upper bound in (163) gives us the following as stricter a sufficient condition (i.e., (174) implies | δ 2 | ě δ 1 ):

Thus, if we have µ t ě σ sp ¨p´ c 4 ´ c 3 q γ ¨p c 2 ` c 1 q , then (174) holds true which in-turn implies | δ 2 | ě δ 1 . Plugging in µ t ě µ 0 , we get the required condition.

E.3 Analysis for SSL

For SSL analysis, we argue that the projection learned by contrastive pretraining can significantly improve the generalization of the linear head learned on top, leaving little to no room for improvement for self-training. Our analysis leverages the margin-based bound for linear models from Kakade et al. [45]. Before introducing the result, we present some additional notation. Let Err D p w q denote 0-1 error of a classifier on a distribution D . Define 0-1 error with margin γ as y Err γ p w q ' ř n i ' 1 I r y i w J x i ď γ s n .

Theorem 19 (Corollary 6 in Kakade et al. [45]) . For all classifiers w and margin γ , we have with probability at least 1 ´ δ :

where B is an upper bound on the ℓ 2 norm of the input points x .

When Ą Err γ p w q is close to zero, the denominating term in RHS of (180) is 4 B γ b 1 n . With Proposition 3, CL solution ϕ obtained on the target domain alone (for SSL setup) is w in when k ' 1 . For larger k 's the CL solution is dominated by the ϕ 1 ' w in . Thus, SSL mainly reduces the B on the projected data by reducing the dependency from order ? d to 1 where 1 is the dimensionality of the output of ϕ without altering the margin. Thus, we get a tighter upper bound for linear probing performed on top CL features when compared with linear probing done on inputs directly.

Intuitively, since the target data has only one predictive feature (along w in ), CL directly recovers this predictive feature as it is the predominant direction that minimizes invariance loss.

F Limitations of Prior Work

F.1 Contrastive learning analysis

Prior works that analyze contrastive learning show that minimizers of the CL objective recover clusters in the augmentation graph, which weights pairs of augmentations with their probability of being sampled as a positive pair [39, 11, 73, 44]. When there is no distribution shift in the downstream task, assumptions made on the graph in the form of consistency of augmentations with downstream labels, is sufficient to ensure that a linear probed head has good ID generalization. Under distribution shift, these assumptions are not sufficient and stronger ones are needed. E.g. , some works assume that same-domain/class examples are weighted higher that cross-class cross-domain pairs [40, 76].

Using notation defined in [76], the assumption on the augmentation graph requires cross-class and same-domain weights ( β ) to be higher than cross-class and cross-domain weights ( γ ). It is unclear if examples from different classes in the same domain will be 'connected' if strong spurious features exist in the source domain and augmentations fail to mask them completely ( e.g ., image background may not be completely masked by augmentations but it maybe perfectly predictive of the label on source domain). In such cases, the linear predictor learnt over CL would fail to generalize OOD. In our toy setup as well, the connectivity assumption fails since on source x sp is perfectly predictive of the label and the augmentations are imperfect, i.e ., augmentations do not mask x sp and examples of different classes do not overlap in source ( i.e ., β ' 0 ). On the other hand, since x sp is now random on target, augmentations of different classes may overlap, i.e ., γ ą 0 , thus breaking the connectivity assumption. This is also highlighted in our empirical findings of CL furnishing representations that do not fully enable linear transferability from source to target (see Sec. 5). These empirical findings also call into question existing assumptions on data augmentations, highlighting that perfect linear transferability may not typically hold in practice. It is in this setting that we believe selftraining can improve over contrastive learning by unlearning source-only features and improving linear transferability.

F.2 Self-training analysis

Some prior works on self-training view it as consistency regularization that constrain pseudolabels of original samples to be consistent with all their augmentations [12, 87, 79]. This framework

abstracts the role played by the optimization algorithm and instead evaluates the global minimizer of a population objective that enforces consistency of pseudolabels. In addition, certain expansion assumptions on class-conditional distributions are needed to ensure that pseudolabels have good accuracy on source and target domains. This framework does not account for challenges involved in propagating labels iteratively. For e.g ., when augmentation distribution has long tails, the consistency of pseudolabels depends on the sampling frequency of 'favorable' augmentations. As an illustration, consider our augmentation distribution in the toy setup in Sec. 4. If it were not uniform over dimensions, but instead something that was highly skewed, then a large number of augmentations need to be sampled for every data point to propagate pseudolabels successfully from source labeled samples to target unlabeled samples during self-training. This might hurt the performance of ST when we are optimizing for only finitely many iterations and over finitely many datapoints. This is why in our analysis we instead adopt the iterative analysis of self-training [17].

G Additional Lemmas

In this section we define some additional lemmas that we use in our theoretical analysis in E. Lemma 20 (Upper bound and lower bounds on erfc ; Kschischang [49]) . Define erfc p x q ' 2 ? π ¨ ş 8 x exp p´ z 2 q ¨ dz . Then we have:

`

Lemma 21 (Properties of Mill's ratio [5]) . Define the Mill's ratio as r p x q ' exp x 2 { 2 ¨ erfc ` x { ? 2 ˘ ¨ a π { 2 . Then following assertions are true: (i) r p x q is a strictly decreasing logconvex function; (ii) r 1 p x q ' x ¨ r p x q ´ 1 is an increasing function with r 1 p x q ă 0 for all x; (iii) r 2 p x q ' r p x q ` x 2 ¨ r p x q ´ x is a decreasing function with r 2 p x q ą 0 for all x ; (iv) x 2 ¨ r 1 p x q is a decreasing function of x .

˘

Lemma 22 (invariance loss as product with operator L ) . The invariance loss for some ϕ P R d is given as: 2 ¨ ş A ϕ p a q ¨ L p ϕ qp a q dP A where the operator L is defined as:

Proof. The invariance loss for ϕ is given by:

Lemma 23. If W is the space spanned by w in and w sp , and W K is the null space for W , then for any u P W and any v P W K , the covariance along these directions E a ' P A r a J uv J a s ' 0 .

Proof: We can write the covariance over augmentations after we break down the augmentation a into two projections: a ' Π W p a q ` Π W K p a q

where the last inequality follows from the fact that E a ' P A ' Π W p a q Π W K p a q J ‰ ' E a ' P A r Π W p a qs E a ' P A r Π W K p a qs J , since the noise in the null space of W is drawn independent of the component along W , and furthermore the individual expectations evaluate to zero.

Lemma 24 (closed-form expressions for eigenvalues and eigenvectors of Σ A , r Σ ) . For a 2 ˆ 2 real symmetric matrix ' a, b c, d ȷ the eigenvalues λ 1 , λ 2 are given by the following expressions:

where δ ' a 4 c 2 `p a ´ b q 2 . Further, the eigenvectors are given by U ' ' cos p θ q , sin p θ q sin p θ q , ´ cos p θ q ȷ , where:

For full proof of these statements see [23]. Here, we will use these statements to arrive at closed form expressions for the eigenvalues and eigenvectors of Σ A , r Σ .

Proof. We can now substitute the above formulae with a, b, c, d taken from the expressions of Σ A and r Σ , to get the following values: λ 1 , λ 2 are the eigenvalues of Σ A , with α determining the corresponding eigenvectors r cos p α q , sin p α qs , r sin p α q , ´ cos p α qs ; and r λ 1 , r λ 2 are the eigenvalues of r Σ , with β determining the corresponding eigenvectors: r cos p β q , sin p β qs , r sin p β q , ´ cos p β qs .

Consider the subclass of problem parameters, d sp ' z, γ ' K 1 { ? z and σ sp ' K 2 ? z for fixed constants K 1 , K 2 ą 0 and some variable z ą 0 , which we can vary to give us different problem instances for our toy model in (6).

˜

¨

˛

From Stewart [80], we can use the closed form expression for the singular vectors of a 2 ˆ 2 full rank asymmetric matrix ' a, b c, d ȷ . The singular vectors are given by

Now, substituting the values in the expression from (100), we get singular vectors of the above form where θ P r 0 , π { 2 s satisfies:

Lemma 25 (asymptotic behavior of τ tan θ ) . For γ ' K 1 { ? z , σ sp ' K 2 ? z ,

Proof. In order to determine the asymptotic nature of tan p θ q as z Ñ8 , we take the limit of a slightly different term first, since we have the closed form expression of tan p 2 θ q .

where, tan p 2 θ q is given by:

since it is easy to see that lim z Ñ8 tan 2 p α ´ β q ¨ ´ λ 1 Ă λ 1 λ 2 Ă λ 2 ´ 1 ¯ ' 0 .

If we use tan p α ´ β q ' tan α ´ tan β 1 ` tan α tan β , and substitute the functions of z , for all the quantities in the above expression using Lemma 24, we derive: lim z Ñ8 τ tan 2 θ ' 2 K 1 K 2 2 { p 1 ` K 2 2 q 2 σ 2 in p 1 ´ 1 { d in q .

Since τ Ñ 8 , tan p 2 θ q Ñ 0 , and further from Taylor approximation of tan p 2 θ q , tan p 2 θ q Ñ 2 θ . We can use this to derive the limit for τ tan θ , which would just be 1 { 2 ¨ 2 K 1 K 2 2 { p 1 ` K 2 2 q 2 σ 2 in p 1 ´ 1 { d in q ' K 1 K 2 2 { p 1 ` K 2 2 q 2 σ 2 in p 1 ´ 1 { d in q .

Lemma 26 (asymptotic behaviors of cot α, tan θ ) . For γ ' K 1 { ? z , σ sp ' K 2 ? z following the expressions in Lemma 24,

Proof. For tan θ , since τ Ñ8 , and τ tan θ approaches a constant (from Lemma 25), we conclude lim z Ñ8 tan θ ' 0 . For cot α ,

and,

d

Thus, cot α Ñ 0 .

Lemma 27 (asymptotic behavior of z cot α ) . For γ ' K 1 { ? z , σ sp ' K 2 ? z following the expressions in Lemma 24,

Proof. The expression for z cot α or z { tan α follows from Lemma 24:

where p ' K 2 1 z ´ 1 ` 1 3 d in ¯ ` σ 2 in 3 ´ 1 ´ 1 d in ¯ ´ ´ z 2 ` 2 K 2 2 z 3 ` 1 6 ¯ . Applying L'Hôpital's (relevant expressions are continuous in z ) rule we get: lim z Ñ8 z cot α ' K 1 1 ` 4 { 3 K 2 2 .

Lemma 28 (asymptotic behavior of z { τ 2 ) . For γ ' K 1 { ? z , σ sp ' K 2 ? z following the expressions in Lemma 24,

Proof. For τ ' λ 1 { λ 2 , substituting the relevant expressions from Lemma 24, we get:

where p ' z { 2 ` 2 K 2 2 z { 3 ` 1 { 6 . Applying L'Hôpital's (relevant expressions are continuous in z ) rule we get: lim z Ñ8 z { τ 2 ' 2 σ 2 in { 3 p 1 ´ 1 { d in q 1 ` 4 { 3 K 2 2 .

Lemma 29 (0-1 error of a classifier on target) . Assume a classifier of the form w ' l 1 ¨ w in ` l 2 ¨ w sp where l 1 , l 2 P R and w in 'r w ‹ , 0 , ..., 0 s J , and w sp ' r 0 , ..., 0 , 1 d sp { ? d sp s J . Then the target accuracy of this classifier is given by 0 . 5 ¨ erfc ´ ´ l 1 ¨ γ ? 2 ¨ l 2 ¨ σ sp ¯ .

Proof. Assume p x, y q ' P T . Accuracy of w is given by E P T ' p sign ` w J x ˘ ' y q ‰ .

Using the definition of erfc function, we get the aforementioned accuracy expression.

Lemma 30. For σ ą 0 and µ P R , we have

Proof. The proof uses simple algebra and the definition of erfc function.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to combine contrastive learning and self-training for unsupervised domain adaptation. Experimental results on UDA benchmarks demonstrate the empirical effectiveness of the approach and a thorough study demonstrates the theoretical benefits.\n\nSTRENGTHS:\n- The results on UDA tasks are consistently better than simple contrastive learning and self-training approaches.\n\n- The theoretical analysis is very thorough and demonstrates the fundamental benefits of contrastive learning and self-training individually and of combining them in the case of UDA.\n\nWEAKNESSES:\n- The proposed method STOC is vaguely described in L104-108. It would be clearer to describe it formally in more details.\n\n- The contributions are not clearly stated. Among the methods presented in L90-108, please clarify that STOC is your contribution.\n\n- Some popular UDA benchmarks are missing, many concurrent work are using Office-31 and Office-Home. Evaluation on these datasets would be appreciated.\n\n- Concurrent work are missing, such as [1, 2, 3]. Please position your work against these paper, and compare your approach both conceptually and empirically to the approaches presented in these papers. More generally, it seems that the presented method is only compared to simple baselines and not to methods that are specifically designed to tackle UDA tasks.\n\n[1] DeepJDOT: Deep Joint Distribution Optimal Transport for Unsupervised Domain Adaptation, Damodaran et al., ECCV 2018\n\n[2] Implicit Class-Conditioned Domain Alignment for Unsupervised Domain Adaptation, Jiang et al., ICML 2020\n\n[3] Semantic-aware Message Broadcasting for Efficient Unsupervised Domain Adaptation, Li et al., ArXiv: 2212.02739\n\nQUESTIONS:\n- In Table 1, are concurrent models your reimplementation, or pretrained models from the SwAV and FixMatch papers ?\n\nLIMITATIONS:\nSome limitations are discussed in the analysis, for exemple regarding the limited performance improvement on semi-supervised tasks.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper explores the synergy of self-training and contrastive learning in semi-supervised learning (SSL) and unsupervised domain adaptation (UDA). It discovers their complementary effect in UDA. Furthermore, it proposes Self-Training Over Contrastive learning (STOC) to combine the benefits of the two approach. The STOC algorithm pretrains the self-training model by contrastive learning on source and target data.Finally, it theoretically analyzes the outcome of initializing with contrastive learning under a specific distribution shift.\n\nSTRENGTHS:\n- The paper provides good theoretical support and empirical evidence to defend their hypothesis. It provides valuable insights into the factors that contribute to the success of contrastive learning and self-training.\n- The paper proposes a novel method that can improve unsupervised domain adaptation.\n- The findings of this paper have important implications for incorporating unlabeled data.\n- The paper is well-written.\n\nWEAKNESSES:\n- The combination of semi-supervised learning with self-supervised learning has been proposed in prior work [1], which is not included in the related work. An updated variety of this method can constitute as a relevant benchmark.\n- The assumptions about infinite unlabeled data, linear classifier $h$, and scaling the magnitude of the coordinates for augmentation, are restrictive.\n- The benchmarks are not competitive. There are several self-supervised methods such as DINO that can supplement the benchmarks.\n- There are typos such as in line 40: 'the strong strong results'.\n\n[1] Zhai, Xiaohua, Avital Oliver, Alexander Kolesnikov, and Lucas Beyer. \"S4l: Self-supervised semi-supervised learning.\" In Proceedings of the IEEE/CVF international conference on computer vision, pp. 1476-1485. 2019.\n\nQUESTIONS:\n- The role of the classification head architecture and deep fine-tuning is missing from the experiments.\n\nLIMITATIONS:\nThe paper does not discuss the limitations and societal impacts of this work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper empirically explores and theoretically studies the complementary benefits of using self-training (ST) and contrastive learning (CL) for unsupervised domain adaptation (UDA), where unlabeled data is available from source and target domains, whereas labels are only available from the source domain. Firstly, experiments on many distribution shift datasets show that combining self-training over contrastive learning (STOC) yields significant benefits for target domain accuracy (~5% on average) over just doing CL or ST. However the benefit in-distribution is shown to be much lower (<1%).\n\nMotivated by these findings, the paper delves into the theoretical study of this phenomenon through a toy Gaussian data model, where each example is composed of domain-invariant (useful for source and target) and spurious features (only useful for source). Experiments on the toy model are shown to follow similar trends to real data. Through analysis on this toy model, the paper provides the following intuitions: \n\n(a) ST can learn a good target classifier if it starts from a good enough initial classifier\n\n(b) ERM on source domain cannot provide a good initializer because it focusses too much on spurious features\n\n(c) CL (with some \"generic\" augmentations) learns representations that focus more on domain-invariant features, thus providing better features for ST to improve. Although CL classifier itself is not sufficient for good target performance\n\nThis theoretical model provides many insights into the role of CL and ST. Finally the paper empirically verifies some of these claims through further probing experiments on real data.\n\nSTRENGTHS:\n- The paper studies an unexplored (to my knowledge) space of combining two popular ideas for domain adaptation. The idea is natural and seemingly effective\n\n- Experiments on standard datasets are convincing enough on the efficacy of the proposed approach. Theoretical analysis on the toy example seems solid, and provides useful intuitions & insights into why the proposed approach could work well. Overall the quality of the technical work seems good. The insights about \"CL provides better features while ST improves the head\" is a clean contribution\n\n- The paper is clearly written and easy to follow. Sufficient intuitions are provided for many of the theoretical claims. \n\nI didn't read through all the proofs, but the intuitions made sense and the contrastive learning results seemed believable. On the whole, I think this is a solid contribution and vote for accepting.\n\nWEAKNESSES:\n- One minor criticism is about the presentation of Section 5. That one required multiple passes to follow, and it might help to include a clearer description of what is being tested. For instance in L377, it would help to explain how the 14% number was calculated. Some more questions about this in the next section\n\n- A bit more discussion about prior work, in the main paper, would have been useful to see. Sections 5 and 6 do a bit of this in some parts, but I only partially understand how this analysis different from previous ones, especially for self-training\n\n- The paper discusses contrastive learning and self-training in general, but only tests on one pair of methods (SwaV and FixMatch). Including more methods (even in 1 setting) would be useful.\n\nQUESTIONS:\n(Q1) How many labeled examples were used for experiments in Section 5? I'm wondering if the target probe results in Fig. 3 are due to worse sample complexity rather than expressivity of the features themselves\n\n(Q2) What are the limitations of the proposed theoretical framework? Does the toy model fail in crucial ways to capture realistic distributions? Does the theory for this model predict something that does not hold in practice. A discussion about this would be useful for future work on this topic\n\n(Q3) Does the choice of contrastive learning/self-training method affect the findings?\n\n(Q4) Does linear probe (instead of fine-tuning everything) also work well in practice?\n\nLIMITATIONS:\nSee Q2 above.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the complementary benefits of combining self-training and contrastive pretraining for domain adaptation under distribution shifts. Through an empirical study on 8 benchmarks, the authors demonstrate that applying self-training (FixMatch) after contrastive pretraining (SwAV) yields substantial accuracy gains over either approach alone. To understand this synergy, the authors analyze a simplified theoretical setup indicating that while contrastive pretraining can amplify signal along invariant features, it may retain dependence on spurious source-only correlations. On the other hand, self-training can effectively unlearn these spurious dependencies and improve linear transfer of representations, achieving optimal target performance. The authors further apply linear probing on representations to verify their theoretical findings empirically. Overall, this work highlights the potential of combining self-training and contrastive learning to address challenges posed by distribution shifts, grounded by both empirical evidence and theoretical analysis.\n\nSTRENGTHS:\n**Originality**: The paper explores the promising yet underexplored direction of combining self-training and contrastive learning to address distribution shifts. The extensive empirical study systematically investigates their complementary benefits across diverse benchmarks. The theoretical analysis offers unique insights into their synergistic effects.\n\n**Quality**: The empirical study is relatively extensive, spanning 8 datasets with careful controls. Despite simplifying assumptions, the theory makes falsifiable predictions that are precisely analyzed and align well with observations. Additional probing experiments further validate the theoretical intuitions. \n\n**Clarity**: The paper is clearly presented overall. The problem setup and motivation are lucid, and the methods and experiments are thoroughly explained. The authors effectively convey the core conceptual messages and insights from their theoretical analysis. \n\n**Significance**: Distribution shifts ubiquitously challenge ML model generalization in practice. This paper highlights the significant potential of combining self-training and contrastive learning to address this problem, demonstrating substantial empirical gains and proposing explanations grounded in both theory and experiments. Overall, I think this paper will be of interest to the ML research community of distribution shifts.\n\nWEAKNESSES:\n**Small Pre-training Set**: The authors have chosen to train models using contrastive learning (CL) from scratch on a selection of small datasets from BREEDS and WILDS. The authors have justified this approach by stating, \"We have opted not to use off-the-shelf pretrained models (e.g., on Imagenet [67]) to avoid confounding our conclusions about contrastive pretraining.\" However, in practice, pretraining CL models on small datasets is not a common practice, given the availability of powerful off-the-shelf CL models that have been pretrained on larger datasets (e.g., from SimCLR, iBot to DINOv2). If the paper's results only apply to CL models pretrained on small datasets and do not generalize to more powerful contemporary CL models, the utility of this work could be limited. I would recommend that the authors conduct studies of modern off-the-shelf CL models to strengthen this work. For instance, off-the-shelf ImageNet-pretrained CL models could be used as initialization, followed by self-training, and then evaluation on various shifted datasets (numerous ImageNet shifted datasets are available). Please note, it is not necessary to strive for results during the rebuttal period. This paper is already commendable in my view, I merely want to offer suggestions that could further enhance its impact.\n\n**Limited Combination of CL+ST Algorithms** In this study, the authors have exclusively utilized SwAV+FixMatch, which are indeed effective algorithms. However, the presentation would be more compelling if a wider range of combinations were explored. Given that the authors employ Barlow Twins in their theoretical analysis, would it not be more fitting to also include Barlow Twins as a CL algorithm in the empirical studies?\n\nQUESTIONS:\nSee Weaknesses.\n\nLIMITATIONS:\nSee Weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-29"}