SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Example 1. For the setup in (1), we choose  $\gamma = 0.5$ ,  $\sigma_{\mathrm{sp}}^2 = 1$ , and  $\sigma_{\mathrm{in}}^2 = 0.05$  with  $d_{\mathrm{in}} = 5$  and  $d_{\mathrm{sp}} = 20$  for our running example.  $\gamma / \sqrt{d_{\mathrm{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.  $\sigma_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 near-optimal target performance. Note that, a linear predictor can improve target performance only by reducing its dependence on spurious feature  $x_{\mathrm{sp}}$ , and increasing it on invariant feature  $x_{\mathrm{in}}$  (along  $w^{\star}$ ). 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  $\gamma/\sigma_{\mathrm{sp}}$  on success of ST. By increasing the ratio of margin  $\gamma$  and variance of spurious feature on target  $\sigma_{\mathrm{sp}}$  (keeping others constant), the problem becomes easier because  $\gamma$  directly affects the signal on  $x_{\mathrm{in}}$  and reducing  $\sigma_{\mathrm{sp}}$  helps ST to unlearn  $x_{\mathrm{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  $\gamma/\sigma_{\mathrm{sp}}$ , ST successfully recovers the optimal target predictor. This hints that ST has a binary effect, where beyond a certain magnitude of  $\gamma/\sigma_{\mathrm{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  $\gamma/\sigma_{\mathrm{sp}}$ . One way of viewing this trend with CL is that it magnifies the effective  $\gamma/\sigma_{\mathrm{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 than that of ST. Finally, we note that for CL the rate of performance increase diminishes at high values of  $\gamma/\sigma_{\mathrm{sp}}$  because CL fails to reduce dependency along  $x_{\mathrm{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  $\gamma / \sigma_{\mathrm{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 CL performs 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  $50k$  unlabeled data and 100 labeled data from the same (target) distribution to simulate SSL setup. Substantiating our findings on real-world 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 near-random 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 \cdot \operatorname{erfc}\left(-\gamma^{2} / (\sqrt{2d_{\mathrm{sp}}} \cdot \sigma_{\mathrm{sp}})\right)$ . Then ST performed in the second stage yields: (i) a classifier with  $\approx 0.5$  target accuracy when  $\gamma < 1/2\sigma_{\mathrm{sp}}$  and  $\sigma_{\mathrm{sp}} \geqslant 1$ ; and (ii) a classifier with near-perfect target accuracy when  $\gamma \geqslant \sigma_{\mathrm{sp}}$ .
The informal theorem above abstracts the exact dependency of  $\gamma$ ,  $\sigma_{\mathrm{sp}}$ , and  $d_{\mathrm{sp}}$  for the success and failure of ST over ERM. Our analysis highlights that while ERM learns a perfect predictor along  $w_{\mathrm{in}}$  (with norm  $\gamma$ ), it also learns to depend on  $w_{\mathrm{sp}}$  (with norm  $\sqrt{d_{\mathrm{sp}}}$ ) because of the perfect correlation of  $x_{\mathrm{sp}}$  with labels on the source. Our conditions depict that when the  $\gamma/\sigma_{\mathrm{sp}}$  is sufficiently small, then ST continues to erroneously enhance its reliance on the  $x_{\mathrm{sp}}$  feature for target prediction, resulting in near-random target performance. Conversely, when  $\gamma/\sigma_{\mathrm{sp}}$  is larger than 1, the signal in  $x_{\mathrm{in}}$  is correctly used for predictor on the majority of target points, and ST eliminates the  $x_{\mathrm{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^{\star}$  is smaller than  $w_{\mathrm{sp}}$  by a certain ratio then the generated pseudolabels (incorrectly) use  $x_{\mathrm{sp}}$  for its prediction further increasing the component along  $w_{\mathrm{sp}}$ . Moreover, normalization further diminishes the reliance along  $w^{\star}$ , culminating in a near-random performance. The opposite occurs when the ERM classifier achieves a signal along  $w^{\star}$  that is sufficiently stronger than along  $w_{\mathrm{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  $\mathrm{P_S}$  and  $\mathrm{P_T}$  gives us a feature extractor  $\Phi_{\mathrm{cl}}$  that has a higher inner product with the invariant feature over the spurious feature. First, we derive a closed form expression for  $\Phi_{\mathrm{cl}}$  that holds for any linear backbone and augmentation distribution. Then, we introduce assumptions on the augmentation distribution (or equivalently on  $w^{\star}$ ) and other problem parameters, that are sufficient to prove amplification.
Proposition 3 (Barlow Twins solution). The solution for (3) is  $U_k^\top \Sigma_{\mathsf{A}}^{-1/2}$  where  $U_k$  are the top  $k$  eigenvectors of  $\Sigma_{\mathsf{A}}^{-1/2} \widetilde{\Sigma} \Sigma_{\mathsf{A}}^{-1/2}$ . Here,  $\Sigma_{\mathsf{A}} := \mathbb{E}_{a \sim \mathrm{P}_{\mathsf{A}}} [aa^\top]$  is the covariance over augmentations, and  $\widetilde{\Sigma} := \mathbb{E}_{x \sim \mathrm{P}_{\mathsf{U}}} [\widetilde{a}(x) \widetilde{a}(x)^\top]$  is the covariance matrix of mean augmentations  $\widetilde{a}(x) := \mathbb{E}_{\mathrm{P}_{\mathsf{A}}(a|x)}[a]$ .
The above result captures the effect of augmentations through the matrix  $U_{k}$ . If there were no augmentations, then  $\Sigma_{\mathrm{A}} = \tilde{\Sigma}$ , 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.,  $\Sigma_{\mathrm{A}}$  is very different from  $\tilde{\Sigma}$ , 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_{\mathrm{in}}$  and spurious  $w_{\mathrm{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_{\mathrm{sp}}$  and  $w_{\mathrm{in}}$  compared to noise directions in their null space. On the other hand it is unclear which of these directions is amplified more in  $\Phi_{\mathrm{cl}}$ . The following assumption and amplification result conveys that when the noise in target  $(\sigma_{\mathrm{sp}})$  is sufficiently large, the CL solution amplifies the invariant feature over the spurious feature.
Assumption 4 (Informal; Alignment of  $w^{\star}$  with augmentations). We assume that  $w^{\star}$  aligns with  $\mathrm{P}_{\mathbb{A}}(\cdot \mid x)$ , i.e.,  $\forall x, \mathbb{E}_{a|x}[a^\top w^{\star}] = 1/2 \cdot x^\top \mathrm{diag}(\mathbb{1}_d)w^{\star}$  is high. Hence, we assume  $w^{\star} = \mathbb{1}_{d_{\mathrm{in}}} / \sqrt{d_{\mathrm{in}}}$ .
One implication of Assumption 4 is that when  $w^{\star} = \mathbb{1}_{d_{\mathrm{in}} / \sqrt{d_{\mathrm{in}}}}$ , only the top two eigenvectors lie in the space spanned by  $w_{\mathrm{in}}$  and  $w_{\mathrm{sp}}$ . To analyze our amplification with fewer eigenvectors from Proposition 3 while retaining all relevant phenomena, we assume  $w^{\star} = \mathbb{1}_{d_{\mathrm{in}} / \sqrt{d_{\mathrm{in}}}}$  for mathematical convenience. While Assumption 4 permits a tighter theoretical analysis, our empirical results in Sec. 4.1 hold more generally for  $w^{\star} \sim \mathcal{N}(0,\mathbf{I}_{d_{\mathrm{in}}})$ .
Theorem 5 (Informal; CL recovers both  $w_{\mathrm{in}}$  and  $w_{\mathrm{sp}}$  but amplifies  $w_{\mathrm{in}}$ ). Under Assumption 4, the CL solution  $\Phi_{\mathrm{cl}} = [\phi_1, \phi_2, \dots, \phi_k]$  satisfies  $\phi_j^\top w_{\mathrm{in}} = \phi_j^\top w_{\mathrm{sp}} = 0 \forall j \geqslant 3$ ,  $\phi_1 = c_1 w_{\mathrm{in}} + c_3 w_{\mathrm{sp}}$  and  $\phi_2 = c_2 w_{\mathrm{in}} + c_4 w_{\mathrm{sp}}$ . For constants  $K_1, K_2 > 0$ ,  $\gamma = K_1 K_2 / \sigma_{\mathrm{sp}}$ ,  $d_{\mathrm{sp}} = \sigma_{\mathrm{sp}}^2 / K_2^2$ ,  $\forall \epsilon > 0$ ,  $\exists \sigma_{\mathrm{sp0}}$ , such that for  $\sigma_{\mathrm{sp}} \geqslant \sigma_{\mathrm{sp0}}$ ,  $|c_1 / c_3 - K_1 K_2^2 d_{\mathrm{in}} / 2 L \sigma_{\mathrm{in}}^2(d_{\mathrm{in}} - 1)| \leqslant \epsilon$ , and  $|c_2 / c_4| - L \sqrt{d_{\mathrm{sp}}} / \gamma| \leqslant \epsilon$ , where  $L = 1 + K_2^2$ .
We analyze the amplification of  $w_{\mathrm{in}} / w_{\mathrm{sp}}$  with contrastive learning in the regime where  $\sigma_{\mathrm{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_{\mathrm{sp}}$ . Thus, in our analysis, we first analyze the amplification factors asymptotically  $(\sigma_{\mathrm{sp}} \to \infty)$ , and then use the asymptotic behavior to draw conclusions for the regime where  $\sigma_{\mathrm{sp}}$  is large but finite.
Theorem 5 conveys two results: (i) CL recovers components along both  $w_{\mathrm{in}}$  and  $w_{\mathrm{sp}}$  through  $\phi_1, \phi_2$ ; and (ii) it increases the norm along  $w_{\mathrm{in}}$  more than  $w_{\mathrm{sp}}$ . The latter is evident because the margin separating labeled points along  $w_{\mathrm{in}}$  is now amplified by a factor of  $|c_2 / c_4| = \Omega (L\sqrt{d_{\mathrm{sp}}} /\gamma)$  in  $\phi_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  $\phi_1$ , the component along  $w_{\mathrm{sp}}$  is still significant  $(c_{1} / c_{3} = \mathcal{O}(1 / L\sigma_{\mathrm{in}}^{2}))$ . Intuitively, CL prefers the invariant feature since augmentations amplify the noise along  $w_{\mathrm{sp}}$  in the target domain. At the same time, the variance induced by augmentations along  $w_{\mathrm{sp}}$  in source is still very small due to which the dependence on  $w_{\mathrm{sp}}$  is not completely alleviated. Due to the remaining components along  $w_{\mathrm{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  $\gamma$ ,  $\sigma_{\mathrm{sp}}$ ,  $d_{\mathrm{sp}}$  defined as in Theorem 5,  $\exists \sigma_{\mathrm{sp}_1}$  such that for all  $\sigma_{\mathrm{sp}} \geqslant \sigma_{\mathrm{sp}_1}$ , the target accuracy of CL (linear predictor on  $\Phi_{\mathrm{cl}}$ ) is  $\geqslant 0.5$ $\operatorname{erfc}(-L' \cdot \gamma / \sqrt{2} \sigma_{\mathrm{sp}})$  and  $\leqslant 0.5$ $\operatorname{erfc}(-4L' \cdot \gamma / \sqrt{2} \sigma_{\mathrm{sp}})$ , where  $L' = K_2^2 K_1 / \sigma_{\mathrm{in}}^2 (1 - d_{\mathrm{in}})$ . When  $\sigma_{\mathrm{sp}_1} > \sigma_{\mathrm{in}} \sqrt{1 - 1/d_{\mathrm{in}}}$ , the lower bound on accuracy is strictly better than ERM from scratch.
While  $\Phi_{\mathrm{cl}}$  is still not ideal for linear probing, in the next part we will see how  $\Phi_{\mathrm{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_{\mathrm{sp}} \leqslant K_1^2 \cdot K_2^{2/3}$ , the target accuracy of ST over CL is lower bounded by  $0.5 \cdot \operatorname{erfc}\left(-|c_2/c_4| \cdot \gamma / (\sqrt{2}\sigma_2)\right) \approx 0.5 \cdot \operatorname{erfc}\left(-L\sqrt{d_{\mathrm{sp}}} / (\sqrt{2}\sigma_{\mathrm{sp}})\right)$  where  $c_2$  and  $c_4$  are the coefficients of feature  $\phi_2$  along  $w_{\mathrm{in}}$  and  $w_{\mathrm{sp}}$  learned by  $BT$ .
The above theorem states that when  $\sqrt{d_{\mathrm{sp}}} / \sigma_{\mathrm{sp}} \gg 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  $\operatorname{erfc}(-\sqrt{10}) \approx 2$  showing near-perfect target generalization. Recall that Theorem 6 shows that CL yields a linear head that mainly depends on both the invariant direction  $w_{\mathrm{in}}$  and the spurious direction  $w_{\mathrm{sp}}$ . At initialization, the linear head trained on the CL backbone has negligible dependence on  $\phi_2$  (under conditions in Theorem 6). Building on that, the analysis in Theorem 7 captures that ST gradually reduces the dependence on  $w_{\mathrm{sp}}$  by learning a linear head that has a larger reliance on  $\phi_2$ , which has a higher "effective" margin on the target, thus increasing overall dependency on  $w_{\mathrm{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.,  $\Phi_{\mathrm{cl}}$ ). When we continue to fix  $\Phi_{\mathrm{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  $\Phi_{\mathrm{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  $\Phi_{\mathrm{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_{\mathrm{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  $(\approx 14\%)$  of target linear probes on backbones trained
Figure 3: Target accuracy with source and target linear probes, which freezes backbones trained with various objectives and trains only the head in UDA setup. Avg. accuracy across all datasets. We observe that: (i) ST improves the linear transferability of source probes, and (ii) CL improves representations.
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.