On the Out-of-Distribution Generalization of Self-Supervised Learning
Authors: Wenwen Qiang
Jingyao Wang
Zeen Song
Jiangmeng Li
Changwen Zheng
Venue: ICML 2025 poster (2025)
Keywords: Self-Supervised Learning
Representation Learning
Out-of-Distribution
Primary area: deep_learning->selfsupervised_learning

Abstract: In this paper, we focus on the out-of-distribution (OOD) generalization of self-supervised learning (SSL). By analyzing the mini-batch construction during the SSL training phase, we first give one plausible explanation for SSL having OOD generalization. Then, from the perspective of data generation and causal inference, we analyze and conclude that SSL learns spurious correlations during the training process, which leads to a reduction in OOD generalization. To address this issue, we propose a post-intervention distribution (PID) grounded in the Structural Causal Model. PID offers a scenario where the spurious variable and label variable is mutually independent. Besides, we demonstrate that if each mini-batch during SSL training satisfies PID, the resulting SSL model can achieve optimal worst-case OOD performance. This motivates us to develop a batch sampling strategy that enforces PID constraints through the learning of a latent variable model. Through theoretical analysis, we demonstrate the identifiability of the latent variable model and validate the effectiveness of the proposed sampling strategy. Experiments conducted on various downstream OOD tasks demonstrate the effectiveness of the proposed sampling strategy.

On the Out-of-Distribution Generalization of Self-Supervised Learning
Wenwen Qiang" 12*
Jingyao Wang 12
Zeen Song"
Jiangmeng Li Changwen Zheng'
Abstract
SSL models face significant challenges in generalizing out-of-distribution (OOD) data. Understanding and improv- ing the OOD generalization capabilities of SSL is crucial for deploying these models in real-world scenarios where the data distribution can shift over time.
In this paper, we focus on the cut-of-distribution (OOD) generalization of self-supervised learning (SSL). By analyzing the mini-batch construction during the SSL training phase, we first give one plausible explanation for SSL having OOD gen- eralization. Then, from the perspective of data generation and causal inference, we analyze and conclude that SSL learns spurious correlations during the training process, which leads to a re- duction in OOD generalization. To address this issue, we propose a post-intervention distribution (PID) grounded in the Structural Causal Model. PID offers a scenario where the spurious vari- able and label variable is mutually independent. Besides, we demonstrate that if each mini-batch during SSL training satisfies PID, the resulting SSL model can achieve optimal worst-case OOD performance. This motivates us to develop batch sampling strategy that enforces PID constraints through the learning of a latent variable model. Through theoretical analysis, we demonstrate the identifiability of the latent variable model and val- idate the effectiveness of the proposed sampling strategy. Experiments conducted on various down- stream OOD tasks demonstrate the effecctiveness of the proposed sampling strategy.
To investigate the OOD generalization properties of SSL, we propose examining the batch construction process during SSL training. SSL methods are generally categorized into two main types: discrimination-based SSL (D-SSL) (Chen et al., 2020; Grill et al., 2020a) and generation-based SSL (G-SSL) (He et al.. 2022: Tong et al., 2022). The core principle of D-SSL is augmentation invariance, ensuring that the feature representations of two different augmentations of the same sample are similar. In contrast, G-SSL focuses on the mask and reconstruction principle, where portion of sample is masked and then reconstructed using an encoder- decoder structure. Leveraging these principles, augmented samples derived from the same original sample, as well as samples before and after masking, can be considered anchor-related pairs. During SSL training, each pair is treated as distinct class, effectively framing each mini- batch as a multi-class learning task. Consequently, the SSL be perceived as learning distribution over tasks based on discrete training tasks, enabling the trained SSL model to generalize to new, unseen tasks, thus demonstrating its OOD generalization capability. However, machine learning is prone to learning spurious correlations that vary between classes or environments (Wang et al., 2023a; 2022). Therefore, although SSL is highly effective in OOD generalization, it may still face the challenge of mitigating spurious correlations.
training process
1. Introduction
Self-supervised learning (SSL) has emerged as a powerful paradigm for training models without relying on labeled data. Meanwhile, SSL models have achieved competitive superior performance various downstream tasks com- pared to supervised learning approaches (Chen et al., 2020; Grill et al., 2020a; Zbontar et al., 2021; He et al., 2022; Tong et al., 2022). However, despite their superior performance,
Building upon the analysis in Section 3, we examine the aforementioned challenge from the perspectives of data gen- cration and causal inference. First, we conclude that the similarity or reconstruction between samples within a pair is affected by several unobservable factors, such as back- ground or texture information independent of the foreground. We also find that the spurious correlation between the anchor and the unobservable variable can vary with the tasks, mak- ing it difficult to eliminate it using the unified causal crite- rion proposed by (Pearl et al.; Pearl, 2009). Furthermore, we demonstrate that, under these circumstances, the SSL model learns to measure similarity or reconstruct using spurious causal factors. This reliance leads to lack of discriminabil-
'qqual contribution 'Institute of Software Chinese Academy of Sciences, Beijing, China University of the Chinese Academy of Sciences, Beijing, China. Correspondence to: Jiangmeng Li <jjangmeng2019e @ws.l1999ssaaa.c......
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

On the Out-of-Distribuiie Generalization of Self-Supervised Learning
ity within each mini-batch task, preventing the SSL model from effectively learning the true task distribution and con- sequently resulting in diminished OOD generalization. To address this issue, we defin new distribution called the post-intervention distribution (PID), characterized by mu- tual independence between the unobservable variable and the anchor. We demonstrate that when the task distribution adheres to PID. the SSL model trained under this condition achieves the lowest worst-case risk, thereby attaining opti- mal worst-case OOD performance. This insight motivates us to design a new mini-batch sampling strategy that ensures the resulting mini-batche satisfy PID constraints, thereby enhancing the OOD generalization capability of SSL..
E
s
E
x+
llabel
x*
x
xlabel
Figure 1: Equation (1).
The SCM for
Figure 2: The SCM for p² (xx x .label s).
nal sample is then referred to as x/. Thus, the augmented dataset in SSL (whether D-SSL or G-SSL) represented as tr {xt x ) " The pair (x,, forms the i-th pair, and SSL aims learn feature extractor f from these pairs.
Based on the above analysis and discussion, we propose a novel mini-batch sampling strategy consisting of two stages. In the first stage, we aim to learn latent variable model to capture the correlations between different variables, i.e., conditional distributions. We prove the identifiability and uniqueness of the resulting latent variable model under a given equivalence relation. In the second stage, we propose a sufficient condition to obtain the balancing score. Using this, we obtain the mini-batch samples through balancing score matching. We also provide theoretical guarantee that the mini-batches obtained by the proposed sampling strategy approximately satisfy the PID.
The objective of D-SSL methods typically consists of two components: alignment and regularization (Wang & Isola, 2020; Chen et al., 2021a). The alignment part is to maxi- mize the similarity between samples that share the same pair in the embedding space, and the regularization part aims to constrain the learning behavior via inductive bias, e.g., Sim- CLR (Chen et al., 2020) constrains the feature distribution to satisfy a uniform distribution. Meanwhile, G-SSI. methods (He et al., 2022) can be regarded as implementing alignment of samples within a pair based on an aucoding-decoding structure, by inputting sample zt into this structure to gen- erate a sample, and making as consistent as possible with sample z²² It is noteworthy that "alignment" in D-SSL is often implemented based anchor points, that is, viewing one sample in a pair as an anchor, the training process of such SSL methods can be seen as gradually pulling the other sample in this pair towards the anchor. The concept of an- chor is also applicable to G-SSL, where x² is viewed as the anchor, and thus the training process of such SSL methods can be viewed as gradually constraining x! to approach xi.
In summary, we make the following contributions: 1) Anal- ysis of SSL Batch Construction: We provide a detailed anal- ysis of how mini-batch construction in SSL influences OOD generalization; 2) Causal Framework for SSL: We introduce causal framework to understand and mitigate the impact of spurious correlations on SSL models; 3) PID-Based Sam- pling Strategy: We propose a theoretically grounded mini- batch sampling strategy that ensures the generated batches conform to PID, improving OOD performance; 4) Empirical Validation: We validate our approach through extensive ex- periments, demonstrating significant improvements in OOD generalization across multiple tasks. The whole logical structure of this paper is further clarified in Appendix
Regardless of whether it is G-SSL or D-SSL, the anchor can be regarded as learning target. Specifically, SSL can be interpreted as follows: In a data augmentation pair, one sam- ple (the anchor) is designated as the target. By constraining the other augmented sample in the feature space to move toward this anchor, consistency in feature representations achieved. This dynamic adjustment causes samples within the same pair to become tightly clustered, thereby forming an effect similar to local cluster center. In traditional clas- sification problems, the common approach is to first project samples into a label space and then constrain them to move toward their corresponding one-hot labels to achieve super- vision. In contrast, SSL directly applies constraints in the feature space, which means that the anchor effectively takes the role of "label." In this unsupervised setting, the an- chor provides a supervisory signal similar to that of label. Therefore, it can be argued that labels manifest differently across various scenarios, and in SSL scenarios, the anchor represents this "implicit label".
2. Revisiting SSL from Pairwise Perspective
In the training phase, the training data is structured into mini- batches, with each where x; represents the i-th sample and N is the total num- ber of samples. In D-SSL methods such as SimCLR (Chen et al., 2020) and Barlow Twins (Zbontar et al., 2021), each sample in Xer undergoes stochastic data augmentation to generate two augmented views, e.g., for x4 E Xtr. the augmented samples can be represented as x and x². For G-SSL methods, like MAE (He et al., 2022) and VideoMAE (Tong et al., 2022), x, is first divided into multiple small blocks, with some blocks masked, and the remaining blocks reassembled into new sample, denoted as z². The origi-
denoted as X,, = (a; N
Based on the above discussion, when we consider the anchor
2

---

On the Out-of-Distribune Generalization of Self-Supervised Learning
as the label the center of clustering, each mini-batch in the training phase can be viewed as multi-class classification task. Specifically, X = {.," .uquchor NN from N categories, where 2, is the positive sample of the i- th category whose clustering center is arancborr Furthermore, the variability of data across mini-batches implies that each mini-batch corresponds to a distinct training task or domain.
models have been effective in OOD generalization, it may still face the challenge of spurious correlations.
Ji=1 consists of data
We further analyze the above challenge from the perspective of data generation and causal inference. Without loss of generality, for each pair in the SSL training process, we denote the anchor as xaaeel and the other sample as xt. Based on (Zimmermann et al., 2021; Von Kügelgenn et al., 2021), x+ can be regarded as caused by anchor zlabel, an unobserved latent variable s E RR" and an independent noise variable with the following formulation:
3. Motivation and Causal Analysis
In this section, we first offer a plausible explanation for the OOD generalization capability of SSL models from a task distribution perspective (the training tasks and the test task is different, this is the key concept of the OOD general- ization). Next, based on causal inference, we demonstrate that 1) during the training phase, SSL exploits spurious correlations to learn certain tasks, thereby degrading its OOD generalization performance; 2) the tasks addressed in SSL training are difficult to model with a single Structural Causal Model, thus, it is challenging to eliminate the con- founding factors involved in learning different tasks through one unified method. Finally, through theoretical analysis, we present that even in the case of spurious associations, we can further improve the OOD generalization of SSL by constraining the data distribution.
x F plbell + c,
(1)
where F is a reversible injective function. From a causal per- spective, Equation (1) can be reformulated as the Structural Causal Model (SCM) shown in Figure 1. The solid arrow indicates that there is a direct causal relationship between the two variables, e.g., laaeell + x+ states that llaael is the direct cause of obtaining at The dotted line indicates that the relationship between the variables is not clear and varies with different environments. Notably, this paper focuses exclusively on scenarios where the semantic information within x+ is related only to z.label, that is, does not contain any causal semantics related to the task. Next, we exam- ine two examples illustrated in Figure 3. In Figure 3 (a), s represents the assigned color, for example, the color of num- bers varies by category, as in the ColoredMNIST dataset (Arjovsky et al., 2019). Here, Gid denotes the class index. Consequently, within mini-batch during training, samples from different classes may have a different texture color. In Figure 3 (b), indicates assigned stylistic attributes, e.g., sketches, cartoon styles, or photographs, and eid denotes the batch index. This scenario commonly occurs in multi- view or domain generalization contexts, like the tasks in the PACS dataset (Li et al., 2017). Therefore, during training, different batches may exhibit different styles, with samples under each style possessing unique appearance attributes. In both figures, s does not capture the foreground semantics between llabel and x+, and the correlation between x laaell and at may vary depending on the settings.
3.1. Formation of the Problem: Causal Perspective
According to Section 2, can infer that different mini- batches correspond to distinct classification tasks. There- fore, the training process of SSL can be described as follows: given a distribution over tasks and the corresponding data distribution of each task (refer to Appendix E for more details about the definitions of task distribution and data distribution), the SSL model can be regarded as be learned based on various training tasks and their corresponding data. The performance of the SSL model is then evaluated on test tasks that are disjoint from the training tasks. This learn- ing paradigm can be regarded as estimating the true task distribution from discrete training tasks, enabling the SSL model to generalize to new, unseen tasks (i.e., test tasks). This also explains well why the SSL model exhibits good performance in transfer tasks (Chen et al., 2020; Grill et al., 2020a; Zbontar et al., 2021), i.e., it has good OOD gener- alization (refer to Appendix E for more details about the relation between OOD generalization and generalization on task distributions). However, machine learning models are prone to learning spurious correlations during the training phase (Wang et al., 2023a; 2022). For example, compared to the foreground features of input data, researchers have found that machine learning models tend to rely the su- perficial texture information or background information of the data for decision- decision-making (Geirhos et al., 2018; Qiang et al., 2022; Xu et al., 2020). Therefore, although the SSL
Based on Figure 3 (a) and (b), we obtain that the causal relationship between y.label and s changes with unknown en- vironmental variations, making it difficult to eliminate based on a unified causal criterion proposed in (Pearl et al.). From Figure (a), due to the existence of path "z.label..... .... + x², the following proposition states that the correlation between glabel and x+ is inffuenced by 8.
Proposition 3.1. Revisiting SSL.from pairwise perspective and assuming that the two samples in each pair sutisfy Equation (1), we can obtain that the learned SSL model will use non-causal factor, i.e., the unobserved latent variable 8, to measure the similarity or reconstruct in a pair.
Detailed proof of Proposition 3.1 is provided in Appendix
3

---

On the
Generalization of Self-Supervised Learning
backgrounds, textures, ctc. that are unique to pair, e.g., the color of small number is red.
stylistic attributes of samples, c.g., sketches, cartoon images, photographs
eid
eid
x
x
xlaael
index
class
x+
x
llabel
batch index
(a) Example task related to ColoredMNIST datasct
(b) Example task related to PACS datasct
Figure 3: Two specific instances illustrate the variability in the causal relationship between a.label and 8 due to environmental changes. The black squares are variables and the arrows indicate causality.
A.1. Notably, when SSL models measure the similarity or reconstruct between paired elements using non- con-causal fac- tors and we can not use one unified method to eliminate the non-causal factors for all related tasks, the extracted repre- sentations may incorporate semantics irrelevant to the task. From the pairwise perspective, this may result in SSL failing to effectively learn each specific task, thereby hindering the modeling of the task distribution and ultimately reducing the OOD generalization ability of SSL.
be recovered from the latent variable x+ via the inverse
mapping:(z/abul, s) = F F='(x c) = (Fatabal (2+
Egiaba), F, fx - Ea)). where denotes the combined noise introduced in the transformation process.
Assumption 3.3 implies that glabel Lpis/z and the intu- itive explanation of Assumption 3.3 can be found in Ap- pendix F. Based on Section 2 and Section 3.1, both D-SSL and G-SSL share a common objective: aligning the positive sample in a pair with its corresponding anchor. Thus, the learning objectives of D-SSL and G-SSL can be unified as maximizing py (atbbbblll=** they achieve py (a-label | ). For example, the training data is first projected the feature space by f, then SimCLR uses contrastive loss to achieve py (z. (2:t while MAE employs the n--norm achieve py (a.label) (x)). We obtain:
3.2. Motivation: Post- Inst-innervention Distribution
The difference lies in how
As shown in Figure 3, regardless of the correlation between s and zalabel, the generation mechanism of x+ is invariant. Because SCMs can be considered as a joint probability dis- tribution, we use the following distribution set to represent the joint probability distribution related to Figure 1:
Theorem 3.4. From a Bayesian perspective, the align- ment component of the SSL objective, i.e., encouraging paired samples to be similar in the feature space, can be interpreted as maximizing the conditional likelihood
D- (a s) plant lalatdd slata pos po pocctii p ptoaaaid so). (2)
Instead of exploring what the specific structure of "y.label...
Pf (xlaabl | x+). Given a model f. the risk over a batic drawn from environment e E DD treated as a distributional constraint, can be defined as: (f) (gr lllog py denotes the joint distribution in environment e. Under As- sumption 3.3, if f arg min (f) all E PID, then f is minimax optimal across all environments in D, that is, f arg miny maxegD L (p label x x²)).
.s zt", we propose to consider using Post- Interentervnnnonn Distribution (PID) to model "(=",>"" labol s), defined as:
Definition 3.2. If the distribution p(x+ ..lbbll s) p(at [glabb]] (clabe))((s), then ((xtt,labo s) is de- fined as PID. glabel and are independent in PID.
(y (z.abel | x=t)], where pf (x+ (cct,,labbo) ,llabel)
We use pT to denote distributions belonging to the PID family. As we can see, p(x+ (altealbblll, 8) is both component of p (x+ ..labll,. s) and a result of the unchanged causal mechanism "s - x + plabel" in Figure 1. Then, the SCM of pf Pl aalabel 8) is shown as Figure 2. this new distribution, because there are no paths between s and llaell we can obtain that + and glabel are only correlated through the stable causal relation = t- laaell Then, from proba- bilistic perspective, what we argue is that compared to SSL models trained on batches satisfying other distribution con- straints in D, SSL models trained on batches that meet the PID distribution constraint have the lowest worst-case risk. To support this, we build upon (Pearl, 2009) by introducing an assumption regarding the invertibility of functions:
Detailed proof of Theorem 3.4 is provided in Appendix A.2. Theorem 3.4 implies that when D is sufficiently large and diverse, no other f obtained from training on any distri- bution can achieve better worst-case OOD performance than the PID (refer to Appendix H for more detail about why the worst-case OOD performance is effective). Notablely, transferring Figure 1 to Figure 2 is similar to backdoor ad- justment in causal inference (Pearl al.). However, from backdoor adjustment pespective, it is straightforward to explain why PID can improve the OOD performance of D-SSL: during the learning of each task, PID eliminates the influence of background semantic confounding (Qiang et al., 2022). However, for G-SSL, regardless of the relationship between s and gllabll, due to the generation purpose, G-SSL
Assumption 3.3. There exist functions Fytttat, F,, and noise variables Cglabal, Ca such that the pair (ableell 8) can
4

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
inherently requires encoding background semantics. Thus, explaining the improvement of OOD performance of G- SSL from the backdoor adjustment perspective is incorrect. Therefore, Theorem 3.4 is provided to explain why PID can improve the OOD performance of both D-SSI. and G-SSL. simultaneously. Also, an intuitive explanation of Theorem 3.4 shown in Appendix F. Moreover, Theorem 3.4 mo- tivates us to design a new mini- sampling strategy to ensure that the resulting mini-batches satisfy PID, thereby improving the OOD generalization of SSL.
Note that k; Q, and T are determined by the type of chosen exponential family distribution and thus independent of e, this guides us to constrain all batch tasks to share these pa- rameters during the training phase. For ease of calculation, we set Qi(+) = exp(-/ - and K as the feature normaliza- tion operator. For X5, since it varies with e, we implement it as the output of a network. Specifically, we first average all the data of a batch, then feed it into a learnable network g. and output the corresponding X For T, we need to guaran- tee it to be a sufficient statistic, one simple way to implement this is the constant transformation. Considering the identifia- bility of the parameters, we implement it as T.y (+) = Daj X where A = [ayy] is a learnable parameter. Up to this point, we obtain the implementation of pfr (s) (alael)) as Ps.A (s) (s|zzlabe)). Then, we implement the conditional genera- tive model in each e E 'D with parameters 0 = (I, A) as: p& (x+ s|zlabel) = pr (x= |s, x alabel) )p (s)c/
4. The Proposed Method
In this section, we present the proposed method which con- sists of two stages. In the first stage, we use a regularized latent variable model (RLVM), e.g., variational autoencoder (VAE) (Kingma & Welling, 2013b), to learn the under- lying distribution plx+, ,,label,, s) for each batch task. In the second stage, based on commonly used instruments in causal effect estimation (Rosenbaum & Rubin, 1981; King & Nielsen, 2019; Austin, 2011). we first introduce a way to calculate the propensity score. Then, we use the learned dis- tribution and propensity score obtain a sampling strategy that can create PID based on training data.
Motivated by the VAE, we estimate the above conditional generative model with the following regularized evidence lower bound (ELBO) in each batch distribution e:
C5.0 E Bgg anabott (log Pr (xc [ - -KL(yo(s/= , || A - a En, A. (4)
Regularizer
4.1. Learning Regularized Latent Variable Model
where A., is the column vector of A, KL(-) is the KL-divergence, and a is a hyperparameter. As for go x .2 it is implemented by learnable network & that outputs the mean and variance, and we use reparame- terization trick (Kingma & Welling, 2013a) to deal with it during training. The last term of Equation (4) is to constrain the column vector orthogonality of A. The training process of Equation (4) is similar to meta-learning, e.g., Prototype Networks (Snell al., 2017), because we construct series of tasks in the training phase. Thus, from meta-learning perspective, training with Equation (4) also indicates that the learned @ can be adaptable for all available tasks.
As shown in Equation (2), to learn the underlying joint distri- bution p(z+, ..labll s) for each batch task, we need to know p(z (zlabel, s), p(zelabel), ((sl in each batch task. Because that p(x+ [.llabll s) is the unchanged causal mech- anism, so we can use a unified f to model [z.abel, 8) in all tasks. Based on the discussion in Section 2, we obtain that glabel is regarded as the label. So, p(zzlabel) can be regarded as the label distribution, and we can represent it with the same uniform distribution in all tasks. Based on the mean-field approximation (Blei et al., 2017; Sriperumbudur et al., 2013) which can be expressed as a closed form of the true prior, we obtain that when the causal relationship between the latent covariate and the label changes with the tasks, an exponential family distribution has the ability to model the conditional distribution p(s p(s|,z/hhe) lalhbee), thus, we have the following assumption for each batch task:
We further show that we can uniquely recover the model parameter up to an equivalence relation. Specifically, we first give the definition of the equivalence relation based on (Motian et al., 2017):
Assumption 4.1. Denote the mini-batch task index as e, the correlation between glaael and in the data distribution pf (#tt,alabel Llaeel, 8) of task can be characterized by:
Definition 4.2. (f,g. 9, Af'w(f', ,,,,), if and ony if there exists an invertible matrix W E R and vector b E RR A(f A(f- (x)) WA'(F (z)) + +, Vr E X
pr (sppaaall
II
exp Σ (1 (3) KeIat
Then, motivated by (Khemakhem et al., 2020), the identifia- bility condition of 8 can be presented as:
where n is the dimension of the latent variable 8, k is the dimension of each sufficient statistic, is the i-th element of 8. Q = [Q.]: s + R" is the base measure, T = [Ty]: s + RM is the sufficient statistics, K" = [K(]: x label + R" is the normalizing constraint, and X = [AG]] plabbl -> R
Theorem 4.3. Suppose that ((ctaall) pr |s, 15,zlaban )ll ((spaaaabbll and the generation pro- cess of X+ can be represented by the SCM depicted in Figure 1, a sufficient condition for 0 = (f.g.A) to be AA-identifibble is given as: 1) Suppose that
5

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
Pe Maylabll , 8)) = pr (1 pr(x+|plabel |alabel, s). Ue is the char- acteristic function of pe(xxt f(g.label, 8)), and the set fxt bostloss (x+)= 0} has measure zero; 2) The sufficient statis- ties T are differentiable almost everywhere, and [Tillsssk are linearly independent on any subset of X with measure greater than zero; 3) There exist nk + 1 distinct pairs (ub :ea,, .ea),... (alabelk; ki; ens) such that the nk x nk matrix = (X (1 (zlibel)- - -x00 (x/pbe)) Aenu (x (alabe)) x (algbal) is invertible.
Then, given function ba(s), we present sufficient condi- tion that it can be the balancing score:
Corollary 4.6. Let ba(s) be a function of 8, a sufficient condition that ba(s) can be regarded as a balancing score that there exists a function vo such that må(s) (bb((s)).
The proof of Corollary 4.6 can be obtained based on Theo- rem I and Theorem 2 in (Rosenbaum & Rubin, 1981). We use ba (8) to denote the balancing score for a specific batch task of SSL. Then, the corresponding propensity score can be represented as mi = [p= (arlabels) [s)] which can be derived from pf (5) (s)zllabe) as defined in Equation (3):
Detailed proof of Theorem 4.3 is provided in Appendix A.3. In Equation (4), we constrain the column vector or- thogonality of A, this can lead to the linearly independence of elements of T, thus, the second assumption of Theorem 4.3 holds. Meanwhile, according to Section 2, we can ob- tain that each ancestor training sample can be regarded as a class. by combining different classes with each other, we can construct adequate tasks, thus, the third assumption of Theorem 4.3 can easily holds. Therefore, based on The- orem 4.3, we can obtain that & can be uniquely recovered. Moreover, we do not theoretically prove that the latent vari- able model can directly identify the spurious variable s. In this paper, the identification of s is based on Assumption 4.1. Theorem 4.3 is also based on Assumption 4.1. The key result of Theorem 4.3 is that we can uniquely identify & and g,, A), thus, is conditionally identified (the detailed explanation of the identifiability of spurious variable 8 is provided in Appendix G).
P9.A (g..(()a))abb)) (p" (z)phel)
p¹ (1rabells) = 2== Pa. A (s
(5)
Jabel (aaaabl)) )) (auabel)'
where p (ar = 1/nu, because that pe (zraabl) is defined empirically as a uniform distribution.
Based on Corollary 4.6, we set s as identical transforma- tion and propose to use the propensity score computed from Equation (5) directly as our balancing score, e.g.. bu(s) (s). Next, we derive the proposed sampling strategy, When given the training data X = (xit , j.label] mu "mu" pairs, we can obtain X of Equation (5) based on the mean of the entire dataset. Then, for each pair, we firstly take one s based on the learned (5(() x " ..labll) and sec- ondly obtain ba(s) by setting mu mu in Equation (5) (please refer to the Step 3 in Appendix G for the explanation of the spatial structure of and the rationale for sampling from g [s|xt, ..label)) )) Finally, we match ba(s) of the selected pair with S a N - 1 different pairs that have the same/closest balaneing score. The detailed sampling strategy is shown in Algorithm 1. Meanwhile, we set the distance metric d(-) in Algorithm as the measurement of distribution distance, c.g., JS-divergence (Endres & Schin- delin, 2003), because that the propensity score itself can be seen as a distribution. Denote the distribution obtained from Algorithm as p(æ 2 glabel ss) then we have:
Ji=l with
According to Assumption 3.3, both an!hell and s can be obtained through an invertible neural network. Training such an invertible neural network typically requires training data in the form of {(z/ahe), 84, 22." R
) However, we do not adopt this mechanism directly because we do not have access to such training data, In particular, we are unable to provide the corresponding 8; for each pair. Therefore, we opt to use a RLVM approach instead.
4.2. The Proposed Mini-Batch Sampling Strategy
As shown in (Rosenbaum & Rubin, 1981), balancing score matching has become a useful tool in the average treatment effect estimation. One of its purposes is to reveal the true causal relationship from the observational data, defined as:
Theorem 4.7. If d((bass)), ((bb((s))bb(s))))) ba(s;)) = 0 in Algorithm 1, the obtained mini-batch is regarded as sampling from a PID, e.g., p(xt, lnhell
la s) = pp'(x"" (x talbel s).
Detailed proof and high-level explanation of Theorem 4.7 is provided in Appendix A.4 and F. Based on The- orem 4.7, at each step, we achieve perfect matching (i.e., (s()) ba(s:)), and the obtained mini-batch samples can be regarded as sampled from the PID. It is worth noting that concepts such as the balancing score, propensity score, and Equation (5) are not novel to this work, because they have been extensively studied by previous researchers (Rosen- baum & Rubin, 1981; King & Nielsen, 2019; Austin, 201 1). However, what we aim to highlight is that the novelty of our approach lies in the proposed sampling strategy, e.g.,
Definition 4.4. A balancing score ba(s) is function of covariate that satisfies: .... (b)())))
From (Rosenbaum & Rubin, 1981), we can obtain that many functions can be used as a balancing score, among them, propensity score p(zlabells) |s) is the coarsest one. Motivated by this, given the hatch task with "mu" pairs, we define the propensity score under the SSL scenario as:
Definition 4.5. The propensity score for a batch task in SSL scenario is mi(s) = [p(r]abo |))
6

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
Algorithm 1 Proposed Mini-Batch Sampling Strategy
transfer learning, and few-shot learning. We introduce the experimental setups in the corresponding sections. Finally, we perform ablation studies. All results reported are the av- erages of five runs performed on NVIDIA RTX 4090 GPUs. More experiments are shown in Appendix C.
Input: Training dataset X =xx-- _x, ar/abel mm balancing score function ba(+), distance metric d ba(()) x bo(-) - R Output: Mini-batch data D consisting of a + 1 examples
1: DP < 0; 0
2: while i = 0
5.1. Benchmark Datasets
3: Randomly sample (xz², : adabel) from Xagg add to
DPI: Compute (8) from (att s!abell)
For unsupervised learning, we select ImageNet-100 100 (Tian et al., 2020) and ImageNet (Deng et al., 2009). For semi- supervised learning, we select ImageNet (Deng et al., 2009) for evaluation. For transfer learning, we select PASCAL VOC (Everingham et al., 2010) and COCO (Lin al., 2014). For few-shot learning. we evaluate the proposed method on Omniglot (Lake et al., 2019), minilmageNet ( (Vinyals et al., 2016), and CIFAR-FS (Bertinetto et al., 2018).
4:
i+i+1
5: end while
6: for 1 iissis a do
7: j arg arg mingJ extr, Aom Aor d(bb(ss)), ba(s)))
8: Add (xt,
to DPI
9: i -++11
10: end for
5.2. Empirical Analysis
Algorithm 1, which enables direct sampling of a mini-batch that satisfies the post-intervention distribution under perfect matching conditions, i.e., when d(ba(ss), ba(s))) = O.
In this article, we primarily address the OOD generalization of SSL. Our experimental design consists of the follow- ing steps: First, we validate that the proposed sampling strategy enhances the performance of SSL. methods in in- distribution scenarios using unsupervised tasks. Second, we classify OOD tasks by difficulty into semi-supervised tasks, transfer learning tasks, and few-shot learning tasks, and sub- sequently evaluate the proposed sampling strategy these tasks. Meanwhile, we conduct experiments generative SSL, the evaluation is provided in Appendix C.I.
As shown in Section 4.2 and Section 4.1, the proposed method has two main phases with the following complexity analysis per mini-batch (mini-batch size B, dataset size D): 1) Regularized Latent Variable Model Training: (a) (s(stt aallett ): Each sample requires a forward pass with cost O(Co), totaling O(B CG)). (b) p ( [s, xlabel), Each sample incurs cost O(Cy), totaling O(B-Cy (c) for X Computed once per mini-batch with cost O(Cg). (d) KL- Divergence: Involves operations over the latent dimension n and sufficient statistic dimension k. contributing O (BBn-Kk -n-k). (e) Orthogonality Regularization: Requires (O k²), which constant when n and are small. Thus, the training phase complexity is approximately: O(B (C.- + Cf + k) + C, + k²);; ; 2) Algorithm 1: (a) Propensity Score Calculation: For each sample, computing scores across the D candidates costs O(D- n k),, leading to total of O(D². n k) for the mini-batch. (b) Matching Operation: A brute- force matching over D samples yields an additional O(D?).
Experimental setup. Our proposed sampling strategy is compatible with any D-SSL or G-SSL model. the stan- dard training procedure of SSL, a mini-batch is randomly sampled from the training data before each iteration. In con- trast, our method replaces this random sampling step with structured mini-batch construction process defined by Al- gorithm 1. Specifically, our approach integrates seamlessly into existing SSL frameworks by substituting the mini-batch sampling component with Algorithm 1, while leaving all other aspects of the SSL training pipeline unchanged. As result, the overall training procedure and hyperparame- ter settings remain identical to those used in the bascline methods. Therefore, for all our experiments, we retain the original hyperparameter configurations to ensure a fair and consistent comparison.
Therefore, the sampling phase has an overall complexity of approximately: O(D2... n k); 3) Overall Complexity: The combined complexity per mini-batch is: O(B (Ca + +C + n. k) + Cg +n k2 + D2 n k). The symbols Cz.. Cj and C represent the computational cost for a single forward pass (or operation) of each respective network module.
Results on unsupervised learning tasks. Table shows the top-1 and top-5 linear classification accuracies on ImageNet- and ImageNet for unsupervised learning task. We can observe that applying the proposed method achieves stable performance improvement, and significantly outper- forms the state-of-the-art (SOTA) methods on all datasets.
5. Experiments
In this section, we first introduce the datasets used in exper- iments. Next, we evaluate our method! on multiple tasks, including unsupervised learning, semi-supervised learning,
Results on semi-supervised learning tasks. Table 2 shows the results on ImageNet for semi supervipd learning task. No matter 1 1% or 10% of the labels are available in 1000
'Codes of the proposed Sampling Strategy can be found in /tttps:/giitee thub. www./mmm///////e..../..ee/ /M/M/MM.TAAAA TASA/ PPD-SSS.
7

---

On the Out-of-Distriben Generalization of Self-Supervised Learning
Table I: The Top-1 and Top-5 classification accuracies of linear classifier on the ImageNet- 100 dataset and the Top-1 results for ImageNet dataset with ResNet-50.
MiConda's
.40
320
ImngeNes-100
IntgeNes
Method
Top-l
Top-5
400
Fpects
1000 Fignels
SinC1
20200)
70.15 +
89.25 +
60:34
0:21
70:45
0:30
MHCo sl., 2020)
72.50 0.12
91.64
0.14
71.16
0.25
711
SimSiam (Chen Ne. 2021)
2301
021
92.6
0.27
70.86 1.3:4
71.37
0:22
2000
100 Sice
0.50
Barlow Twins (Phontoret et al., 2021)
75.97
0:25
92.91
0.119
00.2
0:15
73:29
0.15
SwAW (Cazon 2020)
15.78 1: 0.15
92.86
70:78
0.34
75.32
0.11
Figure 4: Influence of the hy-
Figure 5: Influence of the hy- perparameter a.
RELIC
DINO (Cama tl. 2021)
2012)
98.88
15.45
0.18
1:0.1
71.98
1 0.26
7394
72.17
0.29
perparameter a.
VICRegL. (Burdes 2022)
75.96
0.19
92.97 1: 0.26
72.14
75.07
0.25
SanCLR Oues
73.32 8: 0.15
91.74
72.24
0.20
73.66 0.25
5.3. Ablation Study
McCo Oues
74.71 1: 0.22
93.89
1: 0.17
72.04
1 0.21
74.06
1:006:000 0.20
SamSiam Ours
75.66 1: 0.18
95.02
1: 0.21
72,96
0.22
73.77
0:07:::00 0.17
Bullow Twins a Ours
17.37 2: 0.18
94.99
2: 0.20
73.08
a 0.21
75.89 :: 0.17
Influence of the batch size hyperparameter As shown in Theorem 4.7, a suitable is important. To explore whether the SSL model is more sensitive to original batch size or a, we conduct experiments using ImageNet and BYOL. Figure 4 shows that the performance of BYOL rapidly deteriorates with batch size, while B remains stable over a wide range of batch sizes. Thus, although the proposed sampling strategy has a high require- ment on the SSL method is less sensitive to a.
SwAV Ouss
76.99
0.11
73.25
77.42
0.21
DINO Qurs
77.47 0.15
75.09
0.17
VICRegL. a Ouzs
78.20 0.14
95.07 +:
74.91
0.34
77.77
0.21
Table 2: The semi-supervised learning accuracies (+ 95% confidence interval) on the ImageNet dataset with the ResNet-50 pre-trained on the ImageNet dataset.
Mathod
Epochs
Top-1
Top
Top-1
Top-5
al. 20000
6199 a 0.1
$4.6 02
Influence of a. In Equation 4, a as hyperparameter, controls the weight of the term that constrains the orthogo- nality of the column vectors in the matrix A. It prevents the model from learning redundant or interdependent features, enhancing its generalization. To evaluate its impact, we assess the model performance with varying a (ranging in [0.001, 0.01, 0.1, 1, 10)) on ImageNet- 100. Figure 5 show that performance peaks at a = 1, which is also our setting.
(YOL et 20000)
54.8
0.3
68.0 02
88.5 022
BYOL- Ours
46.5
63.6 0.3
35.6
MoCh Ours
87.4 0.2
0.3
71.4 0.2
90.2
SiLCLK (Chen et 20009
1000
0.2
15.5 0.0.1
65.6 a 0.1
818 0.2
MiCo (He at al., 20200
1000
$2.3 0.1
77.9 0.2
69.4 0.1
38.0 0.2
BYOL (Gril at... 20206)
1000
56.3 0.2
176-1:022
69.7 ::
893
Barkow Twins (Zhontar et all, 2021)
1080
35.0
70.2 0:0.1
67.3 =
022
KBLIC (lomase 2002)
1000
5:.2.:002
6::0:00
889 02
VICReaL (Burdes 2022)
54.91.0.1
1.26
02
SiaCLR
0.
0.2
02
MoCo
0.2
0.2
71.2
0.1
995
BYOL- Ours
1000
$8.9
72.1
91.2=01 01
Barlow Twins Ones
1000
37.6 0.2
90.6 1 0.1
68.9
02
91.8 02
6. Related work
epochs, the improvement brought by the proposed method reaches more than 3% on T and 2% on Top-5 results. This further demonstrates the effectiveness of our method.
SSL is an effective unsupervised representation learning paradigm, aimed at learning general representations suitable for various downstream tasks. From (Jaiswal et al., 2020; Kang et al., 2023), existing SSL models can be divided into D-SSL and G-SSL. The D-SSL methods, e.g., SimCLR (Chen et al., 2020), BYO (Grill et al., 2020a), and Mocov3 (Chen et al., 2021b), are modeled based on the augmentation invariance principle. The G-SSL methods, e.g., MAE (He et al., 2022), iBOT (Zhou et al.), SMA (Xie al., 2024), are modeled based on the mask and reconstruction principle. In real-world scenarios, the data distribution can shift over time. Thus. improving the OOD generalization of SSL is crucial. While various methods (Ni et al., 2021; Bai et al., 2023; Liu et al., 2022; Qiang et al., 2024; Guo et al., 2024; Li et al., 2022; Song et al., 2024) have been proposed with impressive performance, a remaining challenge is they have to contend with trade-offs between inductive biases or approaches without theoretical guarantees. In this paper, we extend the understanding of SSL. from the perspective of task learning and analyze its OOD generalization through causal inference and batch construction.
Results on transfer learning tasks. Table 3 shows the results on the most commonly used object detection and instance segmentation protocol (Chen et al., 2020; Zbontar et al., 2021) for transfer learning, where introducing our method achieves stable improvements in all tasks.
Results few-shot learning tasks. Table 4 shows the effect of the proposed sampling strategy on standard few- shot transfer learning tasks. Compared to the original base- lines, introducing our proposed method achieves remark- able performance improvement, achieving more than 5% improvement. These results demonstrate the superiority of the proposed method under data-scarce conditions.
In summary, we can observe that when the SSL methods trained based on mini-batches generated by our pro- posed sampling strategy, they all further improve their per- formance and by at least 2%. This shows that our sampling strategy is effective in reducing the spurious correlations in the distribution of the mini- batch task, which leads to better causal learning and improves the OOD generalization.
Causality Analysis in SSL plays a crucial role by helping
8

---

On the
Generalization of Self-Supervised Learning
Table 3: Transfer learning on object detection and instance segmentation with C4-backbone. "AP" is the average precision, "APx" represents the average precision when the IoU (Intersection and Union Ratio) threshold is N%.
VOC 07 detection
VOC 07+12 detection
COCO detection
COCO instance segmentation
Method
APw 50
AP
AP
AP
AP7s
AP
AP
Apmask
AP
Supervised
74.4
42.4
42.7
81.3
53.5
58.8
58.2
38.2
41.2
54.7
33,3
35.2
SimCLR (Chen et al, 2020)
75.9
46.8
50.1
81.8
55.5
61.4
57.7
37.9
40.9
54.6
33,3
35.3
MoCo (He et al., 2020)
77.1
46.8
52.5
82.5
57.4
64.0
58.9
39.3
42.5
55.8
34.4
36.5
BYOL (Grill et al., 20206)
77.1
47.0
49.9
81.4
55.3
61.1
57.8
37,9
40,9
54.3
33.2
35.0
SimSiam (Chen He, 2021)
77.3
48.5
52.5
82.4
57.0
63.7
59.3
39.2
42.1
56.0
34,4
36.7
SwAV (Caron al., 2020)
75.5
46.5
49.6
82.6
56.1
62.7
58.6
38.4
41.3
55.2
33.8
35.9
VICRegL (Bardes et al., 2022)
75.9
47.4
52.3
82.6
56.4
62.9
59.2
39.8
42.1
56.5
35.1
36.8
SimCI.R R + Ours
77.6
50.1
51.7
85.3
58.4
63.99
59.2
40.6
43.9
57.1
35.9
37.1
MoCo + Ours
79.4
50.2
54,9
$6.1
60.2
66.1
614
42.1
44,9
59.2
36.9
38.8
BYOL+ + Ours
79.1
50.4
51.9
83.9
58.7
64.1
60.6
39.9
43.7
56.2
35.1
38.6
SimSiam Ours
80.5
50.8
54.4
85.2
66.1
62.3
42.5
43,9
58.1
37.2
39.8
SwAV Ours
77.9
49.3
51.8
84.9
58.1
65.8
62.1
40.2
43.9
56.9
37.3
37.9
VICRegL. + Ours
77.9
50.4
53.9
85.2
58.8
65.3
63.1
42.2
45.3
59.1
37.8
39.9
Table 4: Few-shot transfer learning accuracies (+ 95% confidence interval) on minilmageNet, Omniglot, and CIFAR-FS.
Omniglot
minilmageNett
CIFAR-FS
Method
(5,1)
(5,5)
(20,1)
(5,1)
(5,5)
(20,1)
(5,1)
(5,5)
(20,1)
SimCLR (Chen al., 2020)
90.83 =: 0.21
97.67 = 0.21
81.67 1 0.23
42.32 a: 0,38
51.10
0.37
36.36
0.36
49.44 a 0.30
60.02
0.29
39.29 = 0.30
McCo (Hi et al., 2020)
87.83 1: 0.20
93.52 0.19
80.03 1 0.21
40.56
0.34
19411
0.37
36.52
0.3
45.35
0.31
58.11
1 0.32
37.89 0.32
SwAV (Caron et al., 2020)
91.28 : 0.19
97.21
0.20
82.02
0.20
44.39
0.36
54.1
49.39
a 0.29
62.20 + 0.30
40.19 0.32
SimCLR
91.15
0.21
02
53.18 0.16
67.91
46.94 0.11
MaCo 4 Ours
93.22 + 0.21
97.93 - 0.19
88.93 + 0.22
46,93 0.21
0122--021
41.12
024
51.76 + 0.22
66.42 + 0.21
44.93 0.23
SwAV Ours
96.24 0.26
98.76 0.22
91.96
49.15 0.21
64.28 0.29
42,23
52.64 0.24
70.18 0.21
48.19 0.14
to identify and understand the underlying relationships be- tween variables. Recent works (Sontakke et al., 2021: Zuo et al., 2021; Qiang et al., 2022; Wang al., 2024a; Qiang et al., 2023) have focused on developing methods that lever- age causal inference to extract more robust feature represen- tations. For instance, (Song et al., 2023) used causal invari- ance to obtaín causal SSL representations; ( (Von Kügelgen et al., 2021) studied the identifiability of latent represen- tations, aiming to improve efficiency. However, most of them build causal analysis in-distribution, but ignore the influence of spurious correlations under OOD generaliza- tion settings. In this paper, we explore the essential reasons for spurious correlations in SSL and propose a method that makes the relationships between variables free from spuri- ous correlations. Meanwhile, duo to the limited space in the following, we have discussed the content of Spurious Correlation in Appendix D.
provide a causal analysis of this issue and theoretically examine the intrinsic reasons for incorporating spurious cor- relations during the learning process. Based our causal analysis, we demonstrate that when mini-batches satisfy specific distribution, e.g., PID, SSL models achieve optimal worst-case OOD performance. This insight guides us to propose a new sampling strategy that ensures the resulting mini-batches satisfy the PID constraints. Extensive theoreti- cal and empirical analyses demonstrate its effectiveness.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning and Self-Supervised Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
A Acknowledgements
7. Conclusion
The authors would like to thank the anonymous reviewers for their valuable comments. This work is supported in part by the China Postdoctoral Science Foundation, Grant No. N0.2024M753356, and in part by the National Natural Sci- ence Foundation of China, No. 62406313.
In this paper, we first establish the connection between mini-batches formed during the SSL training phase and multi-class tasks. Next, we explain the rationale for OOD generalization of SSL from a multi-task learning perspective. We then analyze how existing SSL models, when learning mini-batch tasks, rely on spurious correlations to measure sample similarity, leading to suboptimal performance. This reliance affects the SSL model's approximation of the task distribution, resulting in reduced OOD generalization. We
References
Aberdam, A., Litman, R., Tsiper, S., Anschel, O., Slossberg,
R., Mazor, S., Manmatha, R., and Perona, P. Sequence-
9

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
to-sequence contrastive learning for text recognition. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 15302-15312, 2021.
Chen, X. and He, K. Exploring simple siamese represen- tation learning. In Proceedings of the IEEE/CVE Con- ference Computer Vision and Pattern Recognition, pp. 15750- 2021.
Arjovsky, M., Bottou, L.,, Gulrajani, I,, and Lopez-Paz, D. Invariant risk minimization. arXiv; Machine Learn- ing.arXiv: Machine Learning, Jul 2019.
Chen, X., Xie, S., and He, K. An empirical study of training self-supervised vision transformers. Cornell University arXiv, Cornell University arXiv, Apr 2021b.
Austin, P. C.C. An introduction to propensity score methods for reducing the effects of confounding observational studies. Multivariate behavioral research, 46(3):399 424, 2011.
Deng, J., Dong, W. Socher, R., Li, L.-J,, Li, K., and Fei-Fei, Imagenet: A large-scale hierarchical image database. In 2009 IEFE conference computer vision and pattern recognition, pp. 248 118-2555,00 Icee, 2009.
Bai, J., Liu, Z., Wang, H., Hao, J., Feng, Y., Chu, H., and Hu, data in self-
Dosovitskiy, A., Beyer, I.., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. / 1929, 2020.
H. On the effectiveness of supervised long-tail learning. In The Twelfth International Conference on Learning Representations, 2023.
Bardes, A., Ponce, J., and LeCun, Y. Vicregl: Self- supervised learning of local visual features. Advances in Neural Information Processing Systems, 35:8799-8810, 2022.
Endres, D. M. and Schindelin, E. new metric for prob- ability distributions. IEEE Transactions Information theory, 49(7): (777):1888-88880 2003.
Everingham, M., Van Gool, L., Williams, C. K., Winn, J., and Zisserman, A. The pascal visual object classes (voc) challenge. International journal of computer vision, 88 (2):303-338, 2010.
Bertinetto, L., Henriques, J. F., Torr, P. H., and Vedaldi, A. Meta-leaming with differentiable closed-form solvers. arXiv preprint 1at..08880000330 2018.
Geirhos, R., Rubisch, P., Michaelis, C., Bethge, M., Wich- mann, F., and Brendel, W. Imagenet- trained cnns are biased towards texture; increasing shape bias improves accuracy and robustness. International Conference Learning Representationss, Internatonnnneeeeiinnnno Conference Learning Representations, Sep 2018.
Blei, D. M., Kucukelbir, A., and McAuliiff,, D. Variational inference: A review for statisticians. Journal of the Amer- ican Statistical Association, pp. 859-877, Apr 2017. doi: 10. 111180001621459922171288888 URL http: doi.orgg10 1ttp:/111.10000011111
1285773..
Caron, M., Misra, I. Mairal, J., Goyal, P., Bojanowski, P., and Joulin, A. Unsupervised learning of visual features by contrasting cluster assignments. CoRR, abs/2006.09882,, 2020. URL /attps:/aaxxiv.oxt abs/2006.0980 09882.
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. Bootstrap your own latent-a new approach celf-superviee learning. Advances in neural information processing systems, 33:21271-21284, 2020a.
Caron, M., Touvron, H., Misra, I., Jegou, H., Mairal, J.. Bojanowski, P., and Joulin, A. Emerging proper- ties in self-supervised vision transformers. In 2021 IEEE/CVF International Conference on Computer Vi- sion (ICCV), Oct 2021. doi: 10.1 00.109//cc448922.2021. 00951. URL http: / dx doi. www/o0.../......o 1109/ 1ccv46922.2021 202.222.2222. 2091.00911. 00951.
Grill, J.-B., Strub, F, Altché, F., Tallec, C., Richemond, P. H., Buchatskaya, E., Doersch, C., Pires, B. A., Guo, Z. D., Azar, M. G., et al. Bootstrap your own latent: A new approach to self-supervised learning. arXiv preprint M.Xiv:200.0075:3 2020b.
Guo, H., Ba, Y., Hu, J., Si, L., Qiang, W., and Shi, L. Self- supervised representation learning with meta comprehen- sive regularization. In Proceedings of the AAAI Confer- ence on Artificial Intelligence, volume 38, pp. 1959-1967, 2024.
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual rep- resentations. In International conference on machine learning, pp. 16977 PMLR, 2020.
Hamidich, K., Zhang, H., Sankaranarayanan. S., and Ghas- semi, M. Views can be deceiving: Improved ssl through feature space augmentation. Proceedings of the Interna- tional Conference on Machine Learning, 2024.
Chen, T., Luo, C., and Li, L. Intriguing properties of con- trastive losses. Advances in Neural Information Process- ing Systems, 34:11 14:11884-11 1111-114444 11845, 2021a.
10

---

On the Out-of-Distribuuie Generalization of Self-Supervised Learning
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Lake, B. M., Salakhutdinov, R., and Tenenbaum, B. The omniglot challenge: 3-year progress report. Current Opinion in Behavioral Sciences, 29:97-104, 2019.
Li, D., Yang, Y., Song, Y.-Z., and Hospedales, M. Deeper, broader and artier domain generalization. In 2017 IEEE International Conference on Computer Vision (ICCV), Oct 2017. doi: 10.11099iiec.22117.59 URL http: / / dx. doi wwo.o//g////o/oo 111//199/cccc 2017.
He, K., Gkioxari, G., Dollár, P., and Girshick, R. Mask cnn. In Proceedings of the IEEE international conference on computer vision, pp. 2961-2969, 2017.
He, K., Fan, H., Wu, Y., Xie, S., and Girshick, R. Mo- mentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and puttern recognition, pp. 9729-9738, 2020.
Li, J., Qiang, W., Zheng, C., Su, B., and Xiong. H. Metaug: Contrastive learning via meta feature augmentation. In International conference machine learning, pp. 12964- 12978. PMLR, 2022.
Lin, T.-Y,, Maire, M., Belongie, S., Hays, J., Perona, P., Ra- manan, D., Dollár, P., and Zitnick, C. L. Microsoft coco: Common objects in context. In European conference on computer vision, pp. 740-755. Springer, 2014.
He, K., Chen, X., Xie, S., Li, Y., Dollár, P., and Girshick, R. Masked autoencoders are scalable vision learners. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 16000-16009, 2022.
Lin, T.-Y., Dollár, P., Girshick, R., He, K., Hariharan, B., and Belongie, S. Feature pyramid networks for object detection. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2117-2125, 2017.
Huang. W., Han, A., Chen, Y., Cao, Y., Xu, Z., and Suzuki, T. On the comparison between multi-modie and single- modal contrastive learning. Advances in Neural Informa- tion Processing Systems, 37:81549-88 81605, 2024.
Jaiswal, A., Babu, A. R., Zadeh, M. Z., Banerjee, D., and Makedon, F. A survey on contrastive self-supervised learning. Technologies, 9(1):2, 2020.
Liu, X., Wang, Z., Li, Y.-L., and Wang, S. Self-supervised learning via maximum entropy coding. Advances in Neu- ral Information Processing Systems, 35:34091-34105. 2022.
Kang, M., Song, H., Park, S., Yoo, D., and Pereira, S. Bench- marking self-supervised learning on diverse pathology datasets. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3344- 3354, 2023.
Motian, S., Piccirilli, M., Adjeroh, D. D. A., and Doretto, G. Unified deep supervised domain adaptation and gener- alization. In 2017 IEEE International Conference Computer Vision (ICCV), Oct 2017. doi: 10.1109/ccey 2017.609. URL http: // dx doi. 10lgg//0..0 wtt.:rgg///..........11.11.11.1. 1109/ iccv 2017.609
Kay, W., Carreira, J., Simonyan, K., Zhang, B., Hillier, C., Vijayanarasimhan, S., Viola, F., Green, T., Back, T., Natsev, P., Suleyman, M., and Zisserman, A. The kinetics human action video dataset, 2017. URL https:// mexiv.org/e
Murray, N., Marchesotti, L., and Perronnin, F. Ava: A large- scale database for aesthetic visual analysis. In 2012 IEEE conference on computer vision and pattern recognition, pp. 2408-2415. IEEE, 2012.
Khemakhem, I. Kingma, D.. Monti. R., and Hyvarinen, A. Variational autoencoders and nonlinear ica: A unify- ing framework. In International conference on artificial intelligence and statistics, pp. 2207-22111111 2217. PMLR, 2020.
Ni, R., Shu, M., Souri, H., Goldblum, M., and Goldstein, T. The close relationship between contrastive learning and meta- meta-learning. In International conference on learning representations, 2021.
King, G. and Nielsen, R. Why propensity scores should not be used for matching. Political analysis, 27(4):435-455 2019.
Park, G. Y., Jung, C., Lee, S., Ye, C., and Lee, S. W. Self- supervised debiasing using low rank regularization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12395-12405, 2024.
Kingma, D. and Welling, M. Auto-encoding variational bayes. arXiv: Machine Learning,arXiv: Machine Learn- ing, Dec 2013a.
Pearl, J. Causality. Cambridge university press, 2009.
Pearl, J., Glymour, M., Jewell, N., Balke, A., Chickering, D., Galles, D., Geiger, D.,, Goldszmidt, M., Kim, J., Rebane, G., Shpitser, I., Tian, J., Verma, T., Bareinboim, E., Chen,
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint 131. 2013b.

---

On the Out-of-Distriaation Generalization of Self-Supervised Learning
B., Forney, A., Li, A., and Mohan, K. Causal inference in statistics primer.
Tian, Y., Krishnan, D., and Isola, Contrastive multiview coding. In European conference computer vision, pp. 776-794, Springer, 2020.
Qiang, W., Li, Zheng, C., Su, B., and Xiong, H. Interven- tional contrastive learning with meta semantic regularizer. In International Conference Machine Learning, pp. 18018-18030. PMLR, 2022.
Tomasev, N., Bica, I., McWilliams, B., Buesing, L.,, Pascanu, R., Blundell, C., and Mitrovic, Pushing the limits of self self-suuer resnets: Can we outperform supervised learning without labels imagenet? arXiv preprint arXiv:2201.00 10200100519,, 2022.
Qiang, W. J., Su, B., Fu, J., Xiong, H., and Wen, J.-R. Meta attention -enneratonnn network for cross gross-granularity few-shot learning. International Journal of Computer Vision, 131(5):1211-1233 2023.
Tong, Z., Song, Y., Wang, J., and Wang, L. Videomae: Masked autoencoders are data-efficient learners for self- supervised video pre-training. Advances in neural infor- mation processing systems,
Qiang, W. Song, Z., Gu, Z., Li, J., Zheng, C., Sun, F., and Xiong, H. On the generalization and causal explana- tion self-supervised learning. International Journal of Computer Vision, pp. 1-28, 2024.
2022.
Vinyals, O., Blundell, C., Lillicrap. T., Wierstra, D., et al. Matching networks for one shot learning. Advances in neural information processing systems, 29, 2016.
Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Rad- ford, A., Chen, M., and Sutskever, I. Zero-shot text-to- image generation. In International conference on ma- chine learning, pp. 8821-8831. Pmlr, 2021.
Von Kügelgen, J., Sharma, Y., Gresele, I.., Brendel, W., Schölkopf, B., Besserve, M., and Locatello, F. Self- supervised learning with data augmentations provably isolates content from style. Advances neural informa- tion processing systems, 34:16451-16467, 2021.
Ren, S., He, K., Girshick, R., and Sun, J. Faster r-cnn: Towards real-time object detection with region proposal networks. Advances in neural information processing systems, 28, 2015.
Wang. J., Qiang, W., Ren, Y., Song. 7., Su, X., and Zheng, C. Hacking task confounder in meta-learning. arXiv preprint arXiv:2312.0577 2023a.
Rosenbaum, P. R. and Rubin, D. B. The central role of the propensity score in observational studies for causal effects. Dec 1981. doi: 10.21236/adal 14514. URL http://....////aa1111112521111
Wang, Song. Z., Qiang, W., and Zheng, C. Unleash model potential: Bootstrapped meta self-supervised learning. arXiv preprint arXiv:2308. 14267, 2023b.
http: / /xxx/xxx///////////
Wang. J., Qiang, W., Su, X., Zheng, C., Sun, F,, and Xiong, H. Towards task sampler learning for meta-learning, In- lernational Journal of Computer Vision, pp. 1-31, 2024a.
Snell, J., Swersky, K., and Zemel, R. Prototypical networks for few-shot learning. Advances in neural information processing systems, 30, 2017.
Wang. J., Qiang, W., and Zheng, C. Explicitly modeling generality into self-supervised learning. arXiv preprint arXiv:2405.0105: 2024b.
Song, Z., Su, X., Wang J.. Qiang, W., Zheng, C., and Sun, F. Towards the sparseness of projection head in self- supervised learning. arXiv preprint arXiv:2307.08913 2023.
Wang. T. and Isola, P. Understanding contrastive represen- tation learning through alignment and uniformity on the hypersphere. In International Conference on Machine Learning, pp. 9929-9939 PMLR, 2020.
Song. Z., Qiang. W. Zheng, C., Sun, F., and Xiong, H. On the discriminabilityy of self-supervised representation learning, arXiv preprint wrXiv:2407.13541, 2024.
Wang. X., Saxon, M., Li. J., Zhang, H., Zhang, K., and Wang, W. Y. Causal balancing for domain generalization. 2022.
Sontakke, S. A., Mehrjou, A., Itti, L., and Schölkopf, B. Causal curiosity: Rl agents discovering self-supervised experiments for causal representation learning. In Inter- national conference on machine learning, pp. 9848-9858. PMLR, 2021.
arXiv preprint arXiv:2206.0522
Wu, Y., Kirillov, A., Massa, F., Lo, W.Y and Gir- shick, R. Detectron2. https: //gi http://ggtt//./////// thub. com/ 2019.
aceeookresaarcch/etteett..
Sriperumbudur, B., Fukumizu, K., Gretton, A., Hyvärinen, A., and Kumar, R. Density estimation in infinite di- mensional exponential families. arXiv: Statistics The- ory,rXivv Statistics Theory, Dec 2013.
Xie, J., Lee, Y., Chen, A. S., and Finn, C. Self- guided masked autoencoders for domain aonaii- self- supervised learning. arXiv preprint 4rXiv:2402.14789, 2024.
12

---

On the Out-of-Distriaation Generalization of Self-Supervised Learning
Xu, M., Zhang, J., Ni, B., Li, T., Wang, C., Tian, Q. and Zhang, W. Adversarial domain adaptation with domain mixup. Proceedings of the AAAI Conference Artificial Intelligence, pp. 1502-6509 Jun 2020. doi: 10.1609/ aaai. maai.v34i0461133 URL http: /txpxdx..... www..doi.org//....1.11 .org/10. 1609/aasa 03.4144.6660 6123.
Yasmeen, U.. Shah, J. H., Khan, M. A., Ansari, G. J., Rehman, S. U., Sharif, M., Kadry, S., and Nam, Y. Text detection and classification from low quality natural im- ages. Intell. Autom. Soft Comput,
2020.
Yuan. L., Hou, Q., Jiang, Z., Feng. J., and Yan, S. Volo: Vision outlooker for visual recognition. JEEE transac- tions on pattern analysis and machine intelligence, 45(5); 3575-6586 2022.
Zare, S. and Van Nguyen, H. Evaluating and improving domain invariance in contrastive self-supervised learning by extrapolating the loss function. IEEE Access, 11:
187758-137765 2023.
Zbontar, J., Jing, L., Misra, I., LeCun, Y., and Deny, S. Barlow twins: Self-supervised learning via redundancy reduction. arXiv preprint arXiv:2103.03230, 2021.
Zhou, J., Wei, C., Wang, H., Shen, W., Xie, C., Yuille, A., and Kong, T. ibot: Image bert pre-trane pre-training with online tokenizer.
Zhu, W., Liu, S., Fernandez- Grandnd Granda, C., and Razavian, N. Making supe-supe learning robust spurious corre- lation via learnig-sppeed aware sampling. arXiv preprint arXiv:2311.16361 2023.
Zimmermann, R. S., Sharma, Y., Schneider, S., Bethge, M., and Brendel, W. Contrastive learning inverts the data generating process. In International Conference Machine Learning, pp. 12979-12990. PMLR, 2021.
Zuo, X., Cao, P., Chen, Y., Liu, K., Zhao, J., Peng, W., and Chen, Y. Improving event causality identification via self-supervised representation learning on external causal statement. arXiv preprint arXiv:2106.01654 2021.
13

---

On the Out-of-Distributie Generalization of Self-Supervised Learning
Appendix
The Appendices provide additional details that support the main findings and methods proposed this paper. is organized into the following sections:
Appendix A contains the proofs of the presented theorems.
Appendix B provides details for the experimental settings for each experiment.
Appendix C showcases additional experiments that were omitted in the main text due to page limitations.
Appendix D provides the related works for spurious correlation in SSL.
Appendix E explains the differences and connections between task distribution and data distribution.
Appendix F provides the intuitive explanation of several concepts, assumption, and theorems mentioned in the proposed
methodology.
Appendix G provides explanation of the identifiability of spurious variable.
A. Proofs
This section provides the complete proof of Proposition and Theorem in the main text.
A.1. Proof of Proposition 3.1
Before giving the detailed proofs of Proposition 3.1, we first provide the problem definition. Given multiple pairs of samples in an SSL task, let x label be the anchor specific pair, then the remaining samples involving two classes of being x lahel and not ,.label Let g.aabll and zlabel represent the label variables of being ,.label and not z.aabel, since these are binary classification tasks, llabbll and gllabll belong to the set +1. Note that any multi-classification task can be decomposed into binary tasks.
We assume that the labels are drawn from two different probabilities, with balanced sampling probabilities for label values, i.e. Pazlabel 1) lazlabbl -1) = 0.5. Our conclusions also hold for imbalanced distributions. Next, we consider two d-dimensional factors Fz+ and F. representing the knowledge to tackle the two labels. Both are drawn from the
Gaussian distribution:
Fgt ~ (a [laell Paaeet
F, NV "Ma, I
where Mlabel Ms E N denote the mean vectors, while oabell and o denote the covariance vectors. We examine the spurious correlations in SSL. To simplify our analysis, we define Pse the varying correlations that result from different spurious correlations across batches.
Training a single model will result in the optimal model for the target, incorporating non-causal features from the other sample pairs. To substantiate this, we derive the optimal SSL model as follows:
P(xlabel Fat,Fa) P(F+., F.)
P(xlabel Fat, F.)
lay Pat, F.)
Eglabe((-1.)) P(alabel, Fat,
Using the independence of F, and F.- given label we rewrite the joint probability Platabel, Fat, F2) as:
P(zlabel Fgt, F,) P(alabell FB+) P(F, labeel Fst)
labll Fat) P(F, (2 label)
P( label Faa)
Σ
P(F., glabel la label)
/labdl//1.11
P(allabel)
Σ
((Folalaboo P(ylabbll ( labely
glabnl ( 11.1
14

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
Assuming that Fx+ and Fs are drawn from Gaussian distributions, and P(Y/j. Fat, F.) sigmoid where Uس and b3 are the regression vectors for the optimal Bayesian classifier, we have:
Piahel
Fx- + :gg Fs).
Iabel
Jabel
P(x Fat,FE F.) P(a Inbel -PF.) label Fxt) P((llabel) P((F lalabel)
Σ
P(F ( P P(zlabel label)
glabel [(-11)
glabel, Maabel Pot labal
glabal,
mg Po
- lahanl b P.
O &
Pece
+ (1 -Pecle
glabal yysbal Pa- + m3 13)
glabel
(yabaal Fot -2g mg... (
= Pace
Tabot
+ (1 - Peele
Let:
8+
Mlabel
Fxt
14 0
F.
Label
8 Mlabel
11g Fxt
F.
Rabel
0
Substituting 8+ and 3- back into the original equation, we have:
P( lapllabll| Fat, Cu+,Ful Fe)
Page gt +(1- 1111-pp--- ) slubel
1 1+
Paoe: glabel gt +(1 Psole glabel g
When the samples are easy to distinguish, e.g., the similarity the augmented sample from different pairs is not 1:
P( label Fat Fa) =
1
+ earlabe (B+ +3-)
Combining with the expressions for 8+ and B we get:
P(z.aabl| Fot, F,)=
=
Zwlabel Taaml F.+
1 +
In this case, the optimal SSL model only utilizes its own factor Fr and assigns zero weight to the non-causal factor F, from task Ty. Thus, it is difficult to distinguish between the different pairs, the optimal model has non-zero weights for non-causal factors for each task. When the samples are diffieult to distinguish, e.g., the most extreme case, the similarity of the augmented sample from different pairs equal
have:
laalabell Fs)
+ g-
Combining with the expressions for 8+ and B get:
laalll F Fat, F.)
baball
elpial p.t + EyF. Jabel
1 +
In this case, the optimal classifier incorporates both factors Fx+ and F.. Thus, if Pae a 0.5, the optimal classifier assigns non-zero weights to non-causal factors for each task.
A.2. Proof of Theorem 3.4
Proofs. Here, we provide proof of the minimax optimality of the SSL model trained on PID. The SSL model trained on PID Pl (x+ (att,alabel) .label) has py ( (abel [x)b = (z.labell/ Now, consider the expected cross-entropy loss of this classifier on an
15

---

On the Out-of-Distributie Generalization of Self-Supervised Learning
unseen test distribution p::
L (1 (pP)(( (clbbellxx))) ( -IB (p ,label log logp Pr
(s:label (x:t)
p (a(laabel) p (zlabel|
-Pppe walabaly
plaaaaa log (z + E Epr(st bllbar log
pfl (zclabe) = (pP"((zaabee)) + (X, .)) log (pllbaabell (
ppl (abla p (2llabbel (x
(p (zlbbed)))))) (ulaabaa
Eg (xptcx/slhbal,)) log
Consider that glabel Lprs and labeel -psss we we get:
p (2clabel (s)
(1 (x+))) = L (v (P + (x Bx (alabol),)
log p PP (sclabe) lat ( log pp (x+ [zolabal
= (p + E (arhaba :s)
B (x
= Lp" (x ) KLIp" (x (z label
s)]l (x+1s)]
Thus we have the cross entropy loss of (x laholl in any environment is smaller than that p (x (alaa (random guess):
L (xbabellt L. (x (x.labbc))))) s -E -Pprr 1 kT (x+ (
which means:
(x (2gabbell 1xbbbll(xt)) L E3EE
ce(( Lp (albee)) 50.
max
where the performance of p (x+ ,label) is at least good as random guess in any environment. Since we assume the environment diversity, that is for any pf with æ label Les, there exists an environment such that p² (a |z performs worse than a random guess. So we have:
-L' max (Ld (a ( - Ld (p 50 < max ( (x - L (p (a e'EE
3'E3
This inequality shows that the model trained on PID does not perform worse than the random guess model. The KI. divergence term non-negative, ensuring that the difference between the loss of the PID-trained model and the random guess model is non-positive. Now we want to prove that we can obtain pd (a.labll We E, llabel Les, glabel Loslxt, pf (z
(x+)) pPl (P (abel xx+) for any s E S with
we have:
D² px |zlabel 5)
pf (allabll ) = p² (a (aalabella* (xt,s) = D² (a label)
Epr (alaball))) (s) (x plabel]]
(ar lallabel, 8) = Pp Expt(r [pP1 (xt xclabel]]
pp (albell xx s) p (a (x+).
Thus, we have the following minimax optimality p (crabell)) = arg minpyer max L (ps (polalabel) (xtb))l(()))))))))))))))) Based on the above analyses, we have f is the minimax optimal across all elements in D, egg. f arg, min mf (p label x+)).
A.3. Proof of Theorem 4.3
Proofs. We aim to establish the identifiability of the key parameters in the VAE that govern the spuriously correlated covariate features. The proof proceeds through the following two parts: (i) Introducing Auxiliary Variables: We introduce both and label as auxiliary variables to better capture the dependency structure. (ii) Incorporating Causal Mechanism: We specify the causal mechanism that generates + as r label 8) e plabel (ar) + e, where represents the noise term.
16

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
First, transform the equality of the marginal distributions over the observed data into the equality of noise-free distribution. Let start by considering two sets of model parameters, i.e., 0 = (f, A) and 0' = (f', g' A/). Suppose that the marginal distributions of the observed data are equal across these two parameterizations:
(a po (atabee .) pa (x ( .e) Ve E Emain
This condition implies that for each E &min. the distributions of xt conditioned on rlabel and e must be the same for both sets of parameters, This leads us to the following equation:
/ Po.A(Z/ai e)ps (xt (L=t|Z,,, x labblllldd
/ Po e)p/(<=* (2, zlablldd
Now, the noise term (x flabel (Z)) appears in both expressions, as it governs the relationship between xt and the latent variable Z. This results in:
/ Pa.A Pa.A(Zlalabe ce)p(xx - flabel (Z))dZ (
/ Po (g plaeel -e)p(cx - -falabe) ((label(Z)) (Z))dZ
Next, we define the volume of the matrix A as vol(A) := v (e A A) and introduce a change of variables the integrals. On the lent-hand side, we change variables to x flabel (Z), while on the right-hand side, we ø+ = = flabel (Z). Since is injective, we can write (x+ (glabel Z), which implies that Z can be recovered from x+ and xlaabll Thus, arrive at the following equality:
Jac P.A.F.aaaaaaa (x (pc
- PorAA cl (x (p
dE
x+) d=
The key insight here is that both sides describe the same transformation, albeit with different parameterizations. We now introduce convolution representation for the probability distributions. Specifically, define:
( = -Po.A( plabelt ) / ppo
where J / (x++) denotes the Jacobian of flabel Similarly, we define the corresponding expression on the right-hand side. We then introduce the convolution operator * and apply the Fourier transform to both sides:
(Po.A.faelaaa,e * (x+) (P At , glaball e (x+).
Then, we use for the convolution operator, and use F[-] to designate the Fourier transform. The characteristic function of e is then dh = [pe]. Exploit the properties of the Fourier transform to transform the convolution into multiplication. This means that the Fourier domain, we have F ([(Pg.A..zzlata, * (x) = F F'
F(pelww) Meanwhile, we dropped (w)) from both sides non-zero almost everywhere (by assumption of the Theorem).
fg..ff.rlahal e (xx)) = P'',A', A' glabal pelbaal se (x+))
which shows that the probability distributions for x under the two parameterizations are equal.
Obtaining the above results, simplify the expression by removing terms that depend on xt, label or e, and then take the logarithm of both sides of the equation. This results in:
log vol.; (x+) +
Σ (log Q.(t (x+)) - log W (alaael)
Tislf (x (x*)) (
k +
= log vol.p 1 (x+) +
Σ (log (((((((x+))))) - log W (2 label)
(){ Σ ( (zlabel)).
Next, we evaluate this equation at several points (Co, 26 a.bhel), (e1, ..Jablll ..., (Enk, label) provided by assumption (3) of the theorem. We evaluate the above equations at these points to obtain &: + equations, and subtract the first equation from the
17

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
remaining k equations obtain:
Whoh (urabel) W (z/abel)
(xlabel) 200 (=o(bbbel))
log
log W (arabel) w (s/abel)
=f (x+xxxx)))))) ( (z/abel) (x)
This creates system of equations, which will be solved in the next step.
Let L be the matrix defined in assumption (3) and L' similarly defined for X (C' is not necessarily invertible). Define wi'e (slebelyw"" (xr (x(abel) and = [bx]
by = > log with (agebel) w (afabae (afabet)
Then, we can obtain the matrix form:
+
We multiply both sides of Equation by L-T get:
(
+
Where A - L-'L' and L. To complete the proof, we must demonstrate that A is invertible. By the definition of T, its Jacobian exists and is an nk x n matrix with rank n. Consequently, the Jacobian of T e f' also exists and has rank which implies that A is of rank as well. We mainly consider two cases: (i) If k = L, then A is invertible since A E g*xx; and (ii) If 1, define x = f (x) and (z, = (Z,, (T.))
Suppose for any choice of =;, ...,. at. the family
(T
dT.(2) dzf
is never linearly independent. This implies
dB
that 7: lies within subspace of & with a dimension of at most k 1. Let h be a non-zero vector orthogonal to T:(R) Then for x E IR, we have (dTT,a)) h = By integrating, we find that (Ti(x), h) = const. Since this holds for all R and h + 0, we conclude that the distribution is not strongly exponential. Thus, by contradiction, there must exist are linearly independent.
dr
points z/. x? ..... h such that
(dT.(2))
dT, dI,(2.)
d2!
dat
Next, collect these points into k vectors (x+ ,.... ......)))) xx) and concatenate the l Jacobians Jr(x') evaluated at each of those vectors horizontally into the matrix Q = (J.()))..... (T((") (#*)).
Similarly, define Q' as the concatenation of the Jacobians of (f' f()) evaluated at those points. Then the matrix Q is invertible. By differentiating Equation for each x', we get Q = AQ' The invertibility of Q implies the invertibility A and Q'. This completes the proof.
A.4. Proof of Theorem 4.7
Proofs. In Algorithm 1, we describe the procedure for uniformly sampling a different labels. Specifically, we define this process as sampling llabll (a/abel Jabel .r plabel) } using the following procedure: alt
r Jabel ~ U11,2,... .... mu) \ [claabe]
label UT1,2,.... ..., mu, \ (alabel x1 Jabell
x2
label U(,,,, mu) \ (alabel label label Jabel)
Here, U represents the uniform distribution over the specified set, and glabel is the label selected in the current training instance, from which we exclude during the sampling of other labels.
Let Doalance denote the balanced data distribution, which is drawn from the distribution & (a+ .laeel). On the other hand, let De be the data distribution corresponding to the environmental distribution p(x+ ..label.. We assume that for every balancing score match, the labels in the balanced dataset are perfectly matched. Under this assumption, for all environments E Exrans we have:
(2 (I.label (ba*(s) p( laallaelllbc (ba (s)),
18

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
where (s) denotes the balancing score specific environment e and sample 8. By the definition of a balancing score, p(zlabel|s) s) = p(z//bellba" (s)) and & (xglabl|s) 10 (xlabellba (ba*()))
then we have:
f (glabbllss p(x)abells).
and thus, by substitution, we have (1rlabell b) - 1 (slabel (s)). Consequently, we conclude:
8d (sllabelss) (s) p(z
This implies that the balanced distribution pd ()) is identical to the original data distribution p(zlabel|s). Therefore, the joint distribution of the balanced data is:
pd (xt x label .) =p"(x*
.laabll , , 8).
Finally, this shows that the balanced dataset Dallansdd can be regarded as being sampled from a PID (Permuted and Injected Distribution). The implication is that the balanced dataset behaves in the same manner as the original distribution, but with some permutations introduced by the balancing procedure.
B. Experimental Settings
In this section, we provide the details of the settings and datasets for each experiment.
Unsupervised Learning Following the widely adopted protocol (Chen et al., 2020; Wang et al., 2024b), we freeze the feature extractor and train a supervised linear classifier on top of it. The Adam optimizer is used, with Momentum set to 0.8 and weight decay set to 10-4 The linear classifier is trained for 500 epochs, with a batch size of 128. The learning rate starts at 5 10-2 and decays to 5 x 10-6, For this experiment, we utilize several benchmark datasets to evaluate the model's performance. CIFAR-10 and CIFAR- are small-scale image classification datasets consisting of 60,000 32x32 color images in 10 and 100 classes, respectively. STL-10 another small-scale dataset that contains 100,000 unlabeled images and 5,000 labeled examples from 10 classes, with higher image resolution (96x96). Tiny ImageNet contains more challenging small-scale benchmark. For these datasets,
100,000 64x64 images across 200 classes and serves we use ResNet-18 18 as the feature extractor. For larger datasets, we employ ImageNet-100 (a subset of ImageNet with 100 classes) and the full ImageNet dataset, which consists of over 1.2 million images in 1,000 classes, using ResNet-50 as the feature extractor.
Semi-Supervised Learning In accordance with the standard protocol (Zbontar al., 2021), we create two balanced subsets by sampling 1% and 10% of the training dataset. Specifically, we use the ImageNet dataset, a large-scale benchmark for visual recognition tasks, comprising 1.2 million images in 1,000 categories. The subsets contain 1 1% and 10% of the labeled training data, which are used for fine-tuning the model. The models are fine-tuned for 50 epochs, with learning rates set to 0.05 and 1.0 for the classifier and 0.0001 and 0.01 for the backbone on the 1% and 10% subsets, respectively.
Transfer Learning We conduct three transfer learning experiments, including object detection and instance segmentation, transfer other domains, and vased tasks. For object detection, we evaluate the model on two benchmark datasets: Pascal VOC and COCO. Pascal VOC is widely used for object detection tasks, containing around 20,000 images across 20 categories. We train Faster R-CNN (Ren et al., 2015) model on the combined VOC 2007 and 2012 datasets (VOC 07+12), which contains around 16,000 images, and adjust the learning rate at 18K and 22K iterations. We also conduct experiments on a smaller version of Pascal VOC. the VOC 07 set (5K images), with a reduced number of iterations. For instance segmentation, we use the COCO 2017 dataset, which contains over 1 118,000 images and covers 80 object categories. We train a Mask R-CNN (He et al., 2017) with the standard Ix schedule and C4-backbone (Wu et al., 2019), reporting results on the validation split.
Few-shot Learning The protocol outlined in (Wang et al., 2024b; 2023b) is followed for few-shot learning, where we evaluate the proposed method on three standard few-shot learning benchmarks: minilmageNet, Omniglot, and CIFAR-FS minilmageNet is a widely used few-shot learning benchmark derived from the ImageNet dataset, consisting of 60,000 84x84 images across 100 classes. Omniglot is a dataset designed for character recognition, containing 1,623 different characters from 50 different alphabets, making it suitable for testing few-shot learning algorithms. CIFAR-FS is a few-shot version of the CIFAR- 100 dataset, specifically adapted for few-shot learning tasks, containing 100 classes with 600 images per class. For each task, N samples without class-level overlap are randomly selected, and K-times data augmentation is applied to
19

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
create N-way K-shot task. The model is optimized using stochastic gradient descent (SGD) with momentum and weight decay values set to 0.9 and 10-4, respectively. The trained model's 's performance is then evaluated unseen samples drawn from new classes, testing its ability to generalize in few-shot scenarios.
C. Additional Experiments
C.1. Evaluation on Generative SSL
To examine the model's impact on generating SSL, we conducted a series of experiments using the ImageNet-IK dataset (Deng et al., 2009). We started with self-supervised pre-training on the ImageNet- 1K (INIK) training set. Next, we evaluated the representations through supervised training using two methods: (i) end-to-end fine-tuning and (ii) linear probing. We reported the top-1 validation accuracy for single 224x224 crop. For these experiments, we employed ViT-Large (ViT-L/16) (Dosovitskiy et al., 2020) the backbone. ViT-Large is significantly larger (an order of magnitude bigger) than ResNet-50 (He et al., 2016) and has a tendency to overfit. The following section provides a comparison of the models.
Table 5: Comparison between models.
Method
scratch, original
scratch, our impl.
baseline MAE
MAE + Ours
Top
76.5
82.5
84.9
86.4
Table 6: Comparisons with previous results on ImageNet-IK using the ImageNet- IK training set for pre-training, except for the tokenizer in BETT, which was pre-trained 250M DALLE data (Ramesh al., 2021).
Method
pre-train data
ViT-B
ViT-L
VIT-
ViT-Hus
DINO
INIK
82.8
MoCo
INIK
83.2
84.1
BEiT
INIK+DALLE
83.2
85.2
MAE
INIK
83.6
85.9
86.9
87.8
MAE=Ours
INIK
85.9
87.4
88.6
89.3
Comparisons with self-supervised methods. Table we compare the fine-tuning results of self-supervised ViT models. Our method has shown steady improvement from bigger models. We obtain 88.6% accuracy using ViT-H (224 size). The previous best accuracy, among all methods, using only INIK data, is 87.1 (512 size) (Yuan et al., 2022), based on advanced networks. We improve over the state-of the-rrrrtttt by nontrivial margin in the highly competitive benchmark of INIK (no external data). Our result is based vanilla VIT, and we expect advanced networks will perform better.
Object detection and segmentation. We fine-tune Mask R-CNN (He et al., 2017) end-to-end on COCO (Lin et al., 2014). The VIT backbone is adapted for use with FPN (Lin et al., 2017). We apply this approach to all entries in Table We report box AP for object detection and mask AP for instance segmentation. Compared to supervised pre-training, our MAE performs better under all configurations (Table 7). Meanwhile, we also transfer the learned VideoMAE Ours Kinetics-400 (Kay et al., 2017) the downstream action detection dataset AVA (Murray et al., 2012). Following the standard setting, we evaluate top 60 common classes with mean Average Precision (mAP) as the metric. The results are
Table 7: COCO object detection and segmentation using ViT Mask R-CNN baseline.
Apbox
Apmusk
Method
pre-train data
ViT-B
ViT-L
ViT-B
ViT-L
supervised MoCo v3
INIK w/ labels
47.9
49.3
42,9
43.9
INIK
47.9
49.3
42.7
44.0
BEiT MAE
INIK+DALLE
49.8
53.3
44.4
47.1
INIK
50.3
53.3
44.9
47.2
MAE + Ours
INIK
52.5
55.9
46.4
49.7
20

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
Table 8: Performance on for text recognition.
Methods
IIIT5K
IC03
SimCLR (Chen et al., 2020)
SeqCLR (Aberdam el al., 2021)
35.7
43.6
SimCLR + Ours
18.7
19.0
SeqCLR Ours
38.5
47.4
Table 9: Performance Kinetics-700. "Extra Labels" denotes that if the pre-trained models additionally fine-tuned the pre-training dataset with labels before being transferred to AVA. T x refers to the frame number and sample rate.
Method
Backbone
Pre-train Dataset
Extra Labels
T
GFLOPs
Param
mAP
VideoMAE
VIT-L
Kinetics 700
No
16
597
305
36.1
VideoMAE
VIT-I.
Finetics-700
Yes
16 x 4
597
305
39.3
VideoMAE Ours
ViT-L
Kinetics-700
No
16 x
597
305
38.7
VidcoMAE Ours
ViT-L
Kinetics-700
Yes
16 X
597
305
42.1
shown in Table From the results, we can observe that our method achieves stable improvements.
Evaluation for VAE To illustrate without loss of generality, we take SimCLR as representative SSL method for evaluation. The results are shown in Figure 6. First, our proposed method only modifies the mini-batch construction process during the training phase of SimCLR. Even though train a VAE, it does not affect other components of SimCLR's 'ss training pipeline, including the training objective, network architecture, and optimization algorithm. Second, training VAE independently on ImageNet and using its feature extractor for evaluation yields an accuracy of 35.4%, in contrast to SimCLR's Rss 70.2%. We then use the parameters of the VAE's feature extractor to initialize the feature encoder of SimCLR and retrain SimCLR from this initialization. The resulting accuracy is 68.7%, which is 1.4% lower than that of SimCLR trained from scratch. In comparison, SimCLR combined with our method achieves an accuracy 73.3%. These results demonstrate the fairness of our evaluation and confirm that the performance gain is not due the additional VAE.
C.2. Evaluation on More Modalities
The proposed method can be applied in various fields and domains, e.g., instance segmentation, video tracking, sample generation, etc., mentioned before. Here, we provide the experiments of the proposed method on text modality-basee datasets, i.e., IC03 and IIITK (Yasmeen et al., 2020), which we have conducted before. We follow the same experimental settings as mentioned in (Aberdam et al., 2021). The results shown in Table 8 demonstrate that the proposed method achieves stable effectiveness and robustness in various modalities combined with the above experiments.
80
SimCLR (Scratch)
70
VAE Feature Extractor
SimCLR (VAE Initialization)
Accuracy (%)
SimCLR + Ours
50
40
20
10
Methods
Figure 6: Evaluation for the VAE-based components.
21

---

On the Generalization of Self-Supervised Learning
Table 10: Performance comparison on PACS dataset.
Method
Photo
Sketch
Cartoon
Painting (Unseen)
Average
SimCLR
86.4
85.1
87.2
74.3
80.7
SimCLR+Ours
88.0
87.4
90.1
79.2
85.0
BYOL
83.9
84.6
82.7
64.5
74.2
BYOL+Ours
84.2
86.9
85.0
70.8
78.9
Table 11: Performance comparison OfficeHome dataset
Method
A
A
A R
C- A
C P
C->
P C
P R
R -
R R-+ C
R P
SimCLR
58.2
63.5
69.8
78,9
69.7
66.8
63.4
52.3
58.4
56.1
72.9
71.0
SinsLR+Ours
61.1
65.2
71.9
81.1
72.0
68.2
67.5
59.1
59.9
61.2
74.8
73.5
C.3. Evaluation on OOD Tasks
In addition validating the proposed method on standard and few-shot transfer learning scenarios, we also specifically test on benchmark datasets targeting the -itor--distribution (OOD) problem, including PACS, OfficeHome, Waterbirds, and ColoredMNIST. Specifically, evaluate the performance of SSL baselines before and after introducing the proposed PID on these three datasets. For PACS, we follow the experimental setup in Section 5.2 to evaluate the most commonly used SSL baselines (i.e., SimCLR and BYOL) on three domains (i.c., Photo, Sketch, and Cartoon) training on these domains and testing on all four domains, including Photo, Art, Cartoon, and Sketch, as well as the average performance. The results are shown Table 10. For OfficeHome, we randomly select one domain as the source domain for training and another as the target domain. The labels for the source domain are predefined, whereas the labels for all target domains are unknown. We then evaluate the performance change of SimCLR before and after introducing PID, with results shown in Table 11. For the Waterbirds dataset, we adopt the implementation from (Zare& & Van Nguyen, 2023). During training, waterbirds (landbirds) are predominantly paired with water (land) backgrounds. However, at test time, the distribution of backgrounds is altered, creating domain shift. In our evaluation, we report both the average performance and the performance of the worst-performing group. The results are shown in Table 12. From the results, we can observe that our method consistently yields improvements, with particularly significant gains observed for the worst-performing group, indicating the effectiveness of our approach in addressing domain shifts. Finally, for the Colored-MNIS" dataset, we follow the experimental settings mentioned in (Huang al., 2024). In the benchmark, the task involves classifying 10-digit classes, where 10% of the labels are randomly reassigned evaluate the model's robustness label noise. During training, images belonging to class '0' (or "I') are colored red (or green) with probability of 77.5%, and another random color with probability of 22.5%. At test time, the color scheme is reversed, allowing us to assess the model's reliance on color cues for classification. The results shown in Table 13 demonstrate that our method enhances out-of-distribution (OOD) test accuracy by nearly 10%, highlighting the improvement in the model's generalization ability.
D. Related Works for Spurious Correlation
In the recent work SSL, there has been growing interest in understanding its vulnerability spurious correlations (Hamidieh et al., 2024; Wang al., 2022; 2023a). These correlations arise when models learn associations from data that do not truly reflect the underlying causal structure, but instead are coincidental or context-speciff patterns (Pearl, 2009). This susceptibility can undermine the effectiveness of SSL., particularly when dealing with diverse data environments.
Some works have been proposed to alleviate the effects of spurious correlations SSL. Hamidieh et al. (Hamidieh al., 2024) introduced a method that counteracts these correlations by expanding the feature space, thereby providing more diverse training views to mitigate misleading associations. Park et al. (Park al., 2024) proposed that spuriously correlated attributes make neural networks inductively biased towards encoding lower effective rank representations and used rank regularization to eliminate biased samples. Another notable contribution comes from Chen et al. (Z.u et al., 2023), who explored the use of a data reweighting strategy to reduce the importance of data samples that may contain spurious correlations. These methods attempt to eliminate spurious correlations by filtering or enhancing SSL samples at the sample level. Although this approach has proven effective by excluding samples that may contain spurious correlations, difficult to ensure that the learned features are still reliable due to the partial unobservability of spurious correlations and
22

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
Table 12: Performance comparison on Waterbirds dataset.
Method
Test Accuracy
Worst Group
SimCLR
76.2
19.2
SimCLR + Ours
78.0
24.9
MAE
74.9
17.6
MAE Ours
77.2
22.1
Table 13: Performance comparison on ColoredMNIST dataset.
Method
Accuracy(Sa) (%)
SimCLR
12.7
SimCLR Ours
23.5
MAE
15.1
MAE + Ours
24.9
variable coupling. In contrast, our work directly addresses the impact that spurious correlations might cause, utilizing the independence between unobserved variables and anchors under post-intervention distributions ensure the reliability the learned representations.
E. Task Distribution & Data Distribution
Task Distribution: Task distribution refers to a set tasks and their underlying distribution, where each task has its own specific objectives and associated data distribution. It is often used in meta-learning or multi-task learning scenarios to describe the diversity and variation across tasks.
For example, in meta-learning scenario, the task distribution could include:
A "cat dog" classification task (Task 1).
A "car vs airplane" classification task (Task 2).
A "bird vs fish" classification task (Task 3).
These tasks form the task distribution, and the meta-learning model is trained across this task space.
Data Distribution: Data distribution refers to the statistical distribution of data samples within single task, typically described as the joint distribution P(X,Y Y) of input X and labels Y.
Task distribution deseribes the variability between tasks in learning system, focusing on generalization across tasks. Data distribution focuses the variability within single task, addressing adaptation specific data characteristics. The two concepts are hierarchical: task distribution governs the diversity of tasks, while each task has its own distinct data distribution.
Reformulation of OOD Generalization as Generalization on Task Distributions: We organize the whole process into the following steps:
Step First, we provide the formal definition of task distribution.
Without loss of generality, let
classification task as example. We define XØ [(= vf ) training dataset, where xi represents sample, vf represents the corresponding label, denotes the dataset index. and N denotes the number of samples in the dataset. For classification task, the goal is to learn classifier y" (y/ (x?f), so that for any given sample x," the corresponding label can be predicted.
N + +o0, X can be approximated as containing all the information necessary for the classification task and can thus be regarded as a complete dataset for a classification task. Simply put, the elements of a classification task include: the
23

---

On the Out-of-Distribution Generalization of Self-Supervised Learning
classifier and the dataset. We denote a task as fr (v. lut)). Then, the discrete distribution of tasks can be expressed as {X pa (v/l== [x) where M represents the number of tasks.
Furthermore, when a is different, the label space corresponding to y is also different. For example, when a = 1. the label space is {Cat, D and when a = 2, the label space is {Plane, Train). If M + toc, (X%, pp (,(l=1))) [ can be regarded as a complete task distribution.
Step 2: Next, we reformulate SSL from the perspective of task distribution.
In Section 2 and Section 3.1, we explain why a batch in SSL can be viewed as a task. Simply put, for given mini-batch, it can be expressed as: X {a 2 aanchor,a 12h where N denotes the number of ancestor samples in the mini-batch, represents the index of the mini-batch, and x aanchor,a can be regarded as the label of the augmented sample. Meanwhile, the classifier to be learned for each mini-batch is modeled as p (Fanchor,a (m))
Notably, the classifiers for all tasks in SSL are learned using the same classifier, i.c., the classifiers for all tasks aim to learn p l=)).. For example, SimCLR models the classifier using a contrastive loss, while MAE models it using the L2- L2-norm. Therefore, whether D-SSL or G-SSL is used, as M + +00, {(X tr.a p(= aa )o as a task distribution, where M represents the number of tasks.
can be approximated
Step 3: Finally, we reformulate the OOD generalization of SSL as generalization on task distributions.
In traditional machine learning, given training data, the goal is learn p(y|æ)) This can be understood as modeling the data distribution p(ær, y) as p(z) ((z)p(((a)) (y)))) where p(y(z) is learned from the training data and transferred to the test data distribution p(x). This approach assumes that the training and test data are identically and independently distributed, i.e., p( =
p(Ztest) = p(z). and p(Zerain: Ntrain) = p( ((etest Vtest) = p(x, y). Consequently, p(r)p train (ylæ) = p(z) )p test (g|z)) leading to p prrrii (y|z) = ptsst (v)
By analogy, when each data sample is treated as a task, the corresponding learning objective becomes p(p"(r x This learning goal is similar to that in meta-learning [1-2], where the goal is to learn a function that can output the classifier for a given task dataset. Therefore, when the training data are drawn from task distribution, the learning objective is to model the task distribution, i.c., to learn (p'((x(z)) ([(akkk i)). such that it applies to both training and test tasks. Since training and test tasks are different, from the perspective of the training tasks, the test tasks represent OOD scenarios. However, from the perspective of the task distribution, both training and test tasks belong to the same task distribution.
Thus, from the viewpoint of traditional machine learning, SSL can be considered as training with mini-batches of size where each training sample is a training task. One open problem is how to model p(p' (ylxz) p(task i)). Since we define the classifier p(ar plæanehor, (x()) for each SSL training task as identical, ((p((yla) (|(as) i)) can be directly modeled as p(a) aanhhora ( which applies to any sample from any task.
In conclusion, combining Step 1-3, we reformulate OOD generalization as generalization on task distributions.
F. Intuitive Explanations for Assumption 3.3 and Theorem 4.7
Assumption 3.3 illustrates that regardless of whether e E D er E PID, x+ is generated under the control of two variables, s and laaell Therefore, given x*,, s and glabel are conditionally independent, regardless of the correlation between them.
From Assumption 3.3, the optimal / should be Fginhe However, without additional constraints, it is difficult to obtain this optimal f. Theorem 3.4 provides way obtain another good f. detined as f* in the theorem. Why is f* considered good? This is because Theorem 3.4 implies that when D is sufficiently large and diverse, optimal f trained on one distribution will perform worse than random guessing in some other environment. Under such conditions, no other I obtained from training on any distribution can achieve better worst-case OOD performance than the PID. Why is focusing on the worst-case scenario better than other cases? During training, we minimize the worst-case scenario, which involves minimizing: max (py (z. (aclabel (aabell=)))))=== )).. For any f. the term maxeep L (z [x)) is always greater than or equal to L (py labell ()+) for any specific environment e. If we learn an f that minimizes the worst-case term max L (p [x=))), then we naturally minimize L (p (azlabol )) for all in D. This ensures robustness across all scenarios, making the worst-case optimization strategy effective for improving OOD performance.
The high-level explanation of Theorem 4.7 can be presented as follows: 1) From Definition 4.4, it follows that if ba(s)
24

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
can be identified, then and glaeel are conditionally independent given ba(s); 2) In this paper, ba(s) is implemented as described in Equation (5) in the main text. The key challenge lies in obtaining 8. As shown in Section 4.1, explain the identifiability of 8, well as how each label is modeled using a distribution for s. During implementation, we sample from this distribution to generate a series of discrete vectors that approximate associated with specific label; 3) From Equation (2) in the main text, we have: p(x+, x ,8) p(x+ [s.lael, ()p((zlabe)) (p((s) (..ael ). If we select sample pairs for a mini-batch such that all pairs share the same ba(s), the resulting mini-batch can be considered as constructed under the same ba(s). In other words, the samples in the mini-batch are conditioned on bo(s), Combined with the argument in Point 1). we have: p(x+, glabel, s) = p( [a ()p(z)abe))p(s)) which ensures that the mini-batch satisfies PID. The key to achieving PID is ensuring that all sampled examples in the mini-batch have consistent ba(s), i.e., the background information is the same. This allows SSL training to focus on foreground information while disregarding background information.
G. More Explanation for the Identifiability of Spurious Variable
To better address the identifiability of spurious variables in the context of SSL., we organize the response into the following steps:
Step 1: First, we need to clarify that in Section 2 and Section 3.1, propose new perspective for understanding SSL. Taking classification as an example, under this new perspective, each mini-batch during the training phase of SSL can be treated as an independent multi-clasiffication task. Different mini-batches correspond to different classification tasks. In contrast, the traditional perspective of SSL considers the entire dataset as a single task for unsupervised learning.
Therefore, under the new perspective, the training samples in each mini-batch can be considered labeled. Whether these labels are accurate is not our concern for now, as this falls under the domain of Bayesian error. Consequently, in this sense, spurious variables can be identifiable.
We first explain what we mean by the distribution of tasks, using classification as an example. A learning task can be narrowly defined as assigning a label to each sample in a dataset, where the label types are finite. This dataset can represent the task, and the entirety of the dataset can he regarded as the data distribution of that task. Thus, different tasks correspond to different datasets, with distinct label types (tasks with the same label types are considered the same). In this way, a task distribution is essentially the distribution over these datasets, with each element of the task distribution corresponding to a specific dataset.
Step 2:
Step 3: Next, we point out that in this paper, the spurious variable indeed takes values in an infinite space since it is represented by high- dibeesiienii vector. The values of this vector can be arbitrary. We must define s as taking infinite values because, as discussed in Section 2 and Section 3.1, we reinterpret SSL as learning a task distribution where the label types involved are infinite.
Different labels may correspond to different latent variables. These differences are represented by different distributions, i.e., we model the distribution (s)x .alaee) using latent variable model. This allows us to derive the distribution of the spurious variable s for any given label. The values of the probability density can be understood as the degree of correlation between a specific label and a particular value of the latent variable. Hence, given a label, once its conditional distribution (s|x ...labely is determined, we can estimate the corresponding spurious variable 8 through sampling.
We do not theoretically prove that the latent variable model can directly identify the spurious variable s. In this paper, the identification of is based on a strong assumption Assumption 4.1 in the paper. This assumption is justified as follows:
Step 4:
Based on the literature (Blei et al., 2017: Sriperumbudur et al., 2013), which expresses the true prior in closed form, we deduce that when the causal relationship between the latent covariate and the label changes with the tasks, an exponential family distribution is capable of modeling the conditional distribution p( lable.)
Combining Step 1, Step 2, and Step 3. we satisfy the condition that the causal relationship between the latent covariate and the label changes with the tasks.
Step& 5: Theorem 4.3 is also based on Assumption 4.1. The key result of Theorem 4.3 is that we can uniquely identify o and (T, g, A). However, this strong assumption imposes certain limitations on the accuracy of spurious variable identification, which is topic for future research. Despite this strong assumption, our experimental results demonstrate the effectiveness
25

---

On the Out-of-Distributin Generalization of Self-Supervised Learning
of our method.
H. Explanation on Why Worst-Case is Good for OOD Generalization in SSL
(zzlabll(=))) For any f. the worst-case loss During training, we minimize the loss in the worst-case scenario: max, e (p) is always greater than or equal to the loss in any specific environment et 2 (py((llbellz
maxeeD
(xl(xaabel((zt) (p
If we learn an S that minimizes the worst-case loss max.e (py (zlabel )x+)), then ʃ naturally minimizes (x(aylalballzz for all e in D. This ensures robust performance across all scenarios, making the worst-case optimization strategy particularly effective for enhancing OOD generalization.
Therefore, the worst-case can help us to better learning Pj (adabell)) ). A better (y can achieve less empirical risk in the training distribution. As shown in Step 3 of "Reformulation of OOD Generalization as Generalization on Task Distributions" in Appendix E, from the viewpoint of traditional machine learning. SSL can be considered as training with mini-batches of size I, where each training sample is a training task. Then, from PAC theory, we can obtain that minimizing the empirical risk can lead to minimize expected risk, and the expected risk is calculated based on a series of test tasks. Thus, the worst-case can improve the OOD generalization in SSL.
I. The Logical Structure and Viewpoint
To make it easier for readers to understand, we further clarify the logical structure and viewpoints of this paper in the following steps:
Step 1: Reformulate SSL from the perspective of task distribution. From the perspective of task distribution, SSL can be understood learning a distribution of tasks. Each task during training is classification task: for G-SSL, the classifier modeled using the I2-norm, while for D-SSL, the classifier is modeled using contrastive loss. It is crucial to emphasize that we unify the concepts of alignment, classifier, and loss function, as they are essentially the same in our formulation. For further details, please refer to "Reformulation of OOD Generalization as Generalization on Task Distributions" in Appendix F.
Step 2: Model the learning process each task using a fuzzy SCM. In our new understanding, we model the learning process of each task using a fuzzy SCM. It is considered "fuzzy" because the relationship between s and æ lalael is unclear, as shown in Figure 1. We elaborate on this in Section 3.1, explaining that the relationship between and plabel varies across tasks and cannot be captured using single SCM.
Step 3: Explain how spurious variables impact OOD generalization in SSL. In Section 3.1, we explain that under our new understanding, current SSL methods may learn spurious variables, which negatively impact OOD generalization performance. Specifically, spurious variables affect OOD generalization as follows:
Spurious variables make it challenging to learn each specific task properly.
This, in turn, hinders the modeling of the task distribution.
Consequently, SSL performs poorly on test tasks, where test tasks differ from training tasks.
This ultimately undermines OOD generalization.
Step 4: Core argument - improving SSL's OOD performance despite spurious variables. At this point, the core argument of this paper emerges. Through Theorem 3.4, we demonstrate that even in the presence of spurious variables, possible to propose a method to enhance OOD generalization. From Assumption 3.3, the optimal f* should be Fuina. However, without additional constraints, obtaining the optimal f challenging. Theorem 3.4 implies that when D is sufficiently large and diverse, optimal fr trained on one distribution will perform worse than random guessing in some other environments. Under such conditions, no other f obtained from training on any distribution can achieve better worst-case OOD performance than the PID. In other words, even in the presence of spurious variables, e.g. G-SSL, our proposed approach can improve the OOD performance of SSL.
Step 5: Why is the worst-case scenario critical for improving SSL's OOD performance? This is a key insight into how our approach improves SSL's OOD performance. During training, we minimize the loss in the worst-case
26

---

On the Out-of-Distribation Generalization of Self-Supervised Learning
scenario: max (py (py(aaabel)x For any J. the worst-case loss max,eD (py (alabella is always greater than or equal to the loss in any specific environment L (p (a: }). If we learn an f that minimizes the worst-case loss max.ED (ps (p (xraaabbl (x))), )) then f naturally minimizes L (py(xlabel, (xt) for all in D. This ensures robust performance across all scenarios, making the worst-case optimization strategy particularly effective for enhancing OOD generalization.
Step 6: The proposed approach - achieving PID. Our approach achieves PID based on SCM. Why does Algorithm achieve PID? high level explanation is presented in "Intuitive Explanations for Theorem 4.7" Appendix F.
27


Reviews:
[{"review_id": "ve5rsyocsy", "reviewer": "Reviewer_VtXq", "summary": "This paper explores whether self-supervised learning possesses out-of-distribution (OOD) generalization capabilities and investigates the reasons behind its potential failure. To address this, the authors propose a Post-Intervention Distribution (PID), grounded in the Structural Causal Model. PID enables accurate OOD generalization by disentangling spurious correlations between features and labels. The authors introduce a simple yet seemingly effective mini-batch resampling technique and provide a substantial number of supporting theorems. However, I find the number of theorems somewhat excessive. I suggest that the authors consolidate the most essential ones into key theorems and present them in the manuscripts for better clarity and impact.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Response to Weaknesses 1**:\n\nThank you for pointing this out. Due to space limitations, we reported the experimental results of Colored-MNIST and PACS in **Tables 9** and **Tables 10** in **Appendix C.3** of the original submission. In the final version, we will move these results to the main body of the paper.\n\n----\n\n**Response to Weaknesses 2**: \n\nThank you for pointing this out. We explain this issue through the following steps:\n\n**Step 1: How are augmented data pairs formed in SSL**\n\nIn D-SSL, each sample in a mini-batch undergoes stochastic data augmentation to generate two augmented views, e.g., for ${{x_i}}$, the augmented samples can be represented as $x^1_i$ and $x^2_i$. For G-SSL, $x_i$ is first divided into multiple small blocks, with some blocks masked, and the remaining blocks reassembled into a new sample, denoted as $x^1_i$. The original sample is then referred to as $x^2_i$. Thus, the augmented dataset in SSL (whether D-SSL or G-SSL) is represented as $X_{tr}^{aug} = ( {x_i^1,x_i^2} )_{i=1}^N$. $(x_i^1,x_i^2)$ forms the $i$-th pair. \n\nThe above statement can be found in the first paragraph of **Section 2** in the original submission.\n\n**Step 2: How is the anchor formed in SSL**\n\nThe objective of D-SSL methods typically consists of two components: alignment and regularization. The alignment part is to maximize the similarity between samples that share the same pair in the embedding space, and the regularization part aims to constrain the learning behavior via inductive bias. It is noteworthy that “alignment” in D-SSL is often implemented based on anchor points, that is, viewing one sample in **a pair** as an anchor, the training process of such SSL methods can be seen as gradually pulling **the other sample in this pair** (a pair consists of two augmented samples) towards the anchor. Meanwhile, G-SSL can be regarded as implementing alignment of samples within a pair based on an encoding-decoding structure, by inputting sample $x^1_i$ into this structure to generate a sample, and making it as consistent as possible with sample $x^2_i$. The concept of anchor is also applicable to G-SSL, where $x^2_i$ is viewed as the anchor, and thus the training process of such SSL methods can be viewed as gradually constraining $x^1_i$ to approach $x^2_i$.\n\nThe above statement can be found in the second paragraph of **Section 2** in the original submission.\n\n**Step 3: How does the anchor become a label in SSL**\n\nBased on Step 2, regardless of whether it is G-SSL or D-SSL, the anchor can be regarded as a learning target. Specifically, SSL can be interpreted as follows: In a data augmentation pair, one sample (the anchor) is designated as the target. By constraining the other augmented sample in the feature space to move toward this anchor, consistency in feature representations is achieved. This dynamic adjustment causes samples within the same pair to become tightly clustered, thereby forming an effect similar to a local cluster center.\n\nIn traditional classification problems, the common approach is to first project samples into a label space and then constrain them to move toward their corresponding one-hot labels to achieve supervision. In contrast, SSL directly applies constraints in the feature space, which means that the anchor effectively takes on the role of a “label.” In this unsupervised setting, the anchor provides a supervisory signal similar to that of a label. Therefore, it can be argued that labels manifest differently across various spaces—in the feature space, the anchor represents this “implicit label.”\n\n----\n\n**Response to Weaknesses 3**:\n\nThank you for pointing this out. We clarify this issue through the following steps:\n\n**Step 1: How we obtain $s$ for a given pair**\n\n$s$ is a vector obtained based on the distribution $q_\\phi^e(s \\mid x^+, x^{\\text{label}})$. In other words, when $\\phi$ is given, $s$ is also determined. The learning process of $\\phi$ is described in Section 4.1 of the original submission. For a given pair, e.g., $(x_i^+, x_i^{\\text{label}})$, we sample once from $q_\\phi^e(s \\mid x_i^+, x_i^{\\text{label}})$ to obtain $s_i$.\n\n**Step 2: The computation of $ba(s_i)$ in Algorithm 1**\n\nBased on Equation (5) in the original submission, we compute $ba(s_i)$. Specifically, this computation is performed with respect to the entire dataset. Given the full training dataset, the pair $(x_i^+, x_i^{\\text{label}})$, and $s_i$, the $x_j^{\\text{label}}$ involved in Equation (5) are traversed across the entire dataset. Then, according to Definition 4.5, we obtain $ba(s_i)$, which is a vector.\n\n**Step 3: High-level explanation of Algorithm 1 and the identifiability of the spurious variable $s$**\n\nA high-level explanation of **Algorithm 1** is provided in **Appendix F**, while a high-level explanation regarding the identifiability of the spurious variable $s$ is provided in **Appendix G**.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "yes", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "This paper explores the relationship between self-supervised learning and OOD generalization.", "essential_references_not_discussed": "yes", "other_strengths_and_weaknesses": "Strengths:\n1. Comprehensive experiments and theoretical justifications.\n2. The proposed method is simple yet appears to be effective.\n\nWeaknesses:\n\n1. The authors should include additional OOD benchmark datasets in the main experiments, such as Colored-MNIST and PACS.\n\n2. I find the concept of $x_{i}^{anchor}$  mentioned in line 90 somewhat confusing, particularly regarding why it can be directly transformed into the $x^{label}$ used in Equation (1).\n\n3. How is the balancing score function specifically implemented? Is it learnable? Does it produce a scalar output?", "other_comments_or_suggestions": "see weakness", "questions_for_authors": "see weakness", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Qrrhao6BgH", "reviewer": "Reviewer_1Q1V", "summary": "The paper explores the **out-of-distribution (OOD) generalization** of self-supervised learning (SSL). It analyzes how mini-batch construction in SSL training influences OOD generalization and argues that SSL models often learn **spurious correlations**, which hinder their ability to generalize to unseen distributions. To address this issue, the paper introduces a **post-intervention distribution (PID)** based on **Structural Causal Models (SCMs)**. This ensures that spurious variables and label variables remain independent, improving OOD generalization.\n\nFurthermore, the authors propose a **mini-batch sampling strategy** that enforces PID constraints through a latent variable model. They provide theoretical proof of the identifiability of their method and validate it with empirical results. Experiments on various downstream OOD tasks demonstrate that their approach significantly enhances SSL’s generalization performance.\n\n### **Strengths:**\n1. **Novel Causal Perspective on SSL OOD Generalization**  \n   - The paper offers a compelling causal analysis of why SSL struggles with OOD generalization and how spurious correlations arise.\n  \n2. **Innovative Mini-Batch Sampling Strategy**  \n   - Unlike traditional batch sampling, the method ensures **spurious correlations are minimized**, leading to better OOD generalization.\n\n3. **Strong Empirical Performance**  \n   - The proposed method consistently improves performance across diverse benchmarks, including **unsupervised, semi-supervised, transfer learning, and few-shot learning tasks**.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Response to Claims And Evidence**: \n\nThank you for pointing this out. In **Appendix C.1**, we report the results of MAE. Now, we present the results of VideoMAE.\n\nWe transfer the learned VideoMAE + Ours on Kinetics-400 [1] to downstream action detection dataset AVA [2]. Following the standard setting [3], we evaluate on top 60 common classes with mean Average Precision (mAP) as the metric under IoU threshold of 0.5. The result is shown as follows: \n\n| Method            | Backbone      | Pre-train Dataset | Extra Labels  | $T \\times \\tau $ | GFLOPs | Param | mAP  |\n| ----------------- | ------------- | ----------------- | ------------- | ---------------- | ------ | ----- | ---- |\n| VideoMAE          | ViT-L         | Kinetics-700      | No            | $16 \\times 4 $   | 597    | 305   | 36.1 |\n| VideoMAE          | ViT-L         | Kinetics-700      | Yes           | $16 \\times 4 $   | 597    | 305   | 39.3 |\n| VideoMAE + Ours   | ViT-L         | Kinetics-700      | No            | $16 \\times 4 $   | 597    | 305   | 38.7 |\n| VideoMAE + Ours   | ViT-L         | Kinetics-700      | Yes           | $16 \\times 4 $   | 597    | 305   | 42.1 |\n\nIn the above table, \"Extra Labels\" denotes that if the pre-trained models are additionally fine-tuned on the pre-training dataset with labels before transferred to AVA. $T \\times \\tau $ refers to frame number and corresponding sample rate. In the final version, we will add these results to the main body of our submission.\n\n[1] The kinetics human action video dataset. arXiv preprint, 2017.\n\n[2] Ava: A video dataset of spatio-temporally localized atomic visual actions. CVPR, 2018.\n\n----\n\n**Response to Weaknesses 1**: \n\nThank you for pointing this out. We provide some additional ImageNet-inspired examples and explanations to address the concern:\n\n**Object vs. Background:**\n  Consider an ImageNet class like \"Golden Retriever.\" The causal semantics for recognizing a Golden Retriever primarily reside in the dog’s shape, fur texture, and facial features. Although many images might have different backgrounds—such as parks, beaches, or urban settings—these background elements (which could be captured by $s$) are not causally responsible for the image being classified as a Golden Retriever. In this case, $x^{\\rm label}$ would capture the object-specific features, while $s$ would account for non-causal variations like the background.\n\n**Intra-Class Variability:**\n  Take another class, such as \"Volcano.\" A volcano can be pictured under different weather conditions, from different angles, and with various surrounding landscapes. While these environmental or stylistic factors vary widely, the key causal semantics—such as the volcano’s structure, cone shape, and crater—remain consistent. Again, $s$ may vary (capturing changes in lighting, weather, or background) without affecting the causal information needed to identify a volcano.\n\n----\n\n**Response to Weaknesses 2**: \n\nThank you for pointing this out. We provide the model efficiency and memory footprint results of the proposed method trained on 8 NVIDIA Tesla V100 GPUs. \n\n| Method      | Training Time (Hours) |          |Memory Footprint (GB) |          |\n| ----------- | --------------------- | -------- | --------------------------------- | -------- |\n|             | CIFAR-10              | ImageNet | CIFAR-10                          | ImageNet |\n| SimCLR      | 10.4                  | 101.9    | 23.3                              | 221.6    |\n| SimCLR+Ours | 12.7                  | 106.2    | 29.7                              | 230.7    |\n| MAE         | 13.8                  | 115.5    | 26.9                              | 244.9    |\n| MAE+Ours    | 16.4                  | 122.2    | 31.2                              | 252.2    |\n\nFor specific computation complexity, please refer to  **Response to Weaknesses 1 & Questions 3** of **Rebuttal for Reviewer fDza**\n\n----\n\n**Response to Weaknesses 3**: \n\nThank you for pointing this out. To illustrate without loss of generality, we take SimCLR as a representative SSL method. First, our proposed **Algorithm 1** only modifies the mini-batch construction process during the training phase of SimCLR. Even though we train a VAE, it does not affect other components of SimCLR’s training pipeline, including the training objective, network architecture, and optimization algorithm. Second, training a VAE independently on ImageNet and using its feature extractor for evaluation yields an accuracy of **35.44%**, in contrast to SimCLR's **70.15%**. We then use the parameters of the VAE's feature extractor to initialize the feature encoder of SimCLR and retrain SimCLR from this initialization. The resulting accuracy is **68.71%**, which is **1.45%** lower than that of SimCLR trained from scratch. In comparison, SimCLR combined with our method achieves an accuracy of **73.32%**. These results demonstrate the fairness of our evaluation and confirm that the performance gain is not due to the additional VAE.", "claims_and_evidence": "The claims in the paper are partially supported by the evidence in the experiment part. However, the reviewer is concerned about the lack of evaluation on mask-autoencoder based pre-training methods like [1,2] in the main paper. Although some results are provided in the supplementary, it can be worthwhile to add the comparison in the experiments in the main paper together with contrastive-based methods, as in the analysis and proof part, the authors formulate discriminated-based and generative-based methods with a unified framework.\n\n[1] He, K., Chen, X., Xie, S., Li, Y., Doll´ ar, P., and Girshick, R. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 16000–16009, 2022.\n[2] Tong, Z., Song, Y., Wang, J., and Wang, L. Videomae: Masked autoencoders are data-efficient learners for self-supervised video pre-training. Advances in neural information processing systems, 35:10078–10093, 2022.", "methods_and_evaluation_criteria": "The reviewer thinks the evaluation setting and metrics make sense for the claims.", "theoretical_claims": "The reviewer checked partially of the proofs. Specifically, theoretical claims in Section 3 and 4.1 are checked.", "experimental_designs_or_analyses": "The reviewer thinks that experimental designs are valid.", "relation_to_broader_scientific_literature": "The paper makes a theoretical and practical contribution by introducing a causal approach to improving OOD generalization in SSL. While the method is empirically validated, its assumptions, computational cost, and feasibility in large-scale applications could be explored further.", "essential_references_not_discussed": "The reviewer does not come up with essential related literature that is not discussed.", "other_strengths_and_weaknesses": "1. Starting from Line 162, left column, the authors propose an assumption that \"the semantic information within x+ is related only to xlabel, that is, s does not contain any causal semantics related to the task.\". They provide two examples about this assumption. However, as natural images are not restricted to numbers/styles discussed in the two example, the reviewer is concerned about the assumption. More examples in ImageNet should be provided against it.\n2. As additional training is required, the reviewer is concerned about the training efficiency of the proposed method. The authors should provide evaluation on it.\n3. As more parameters (VAE) are introduced in the proposed method, the reviewer is concerned about the fairness of comparison. Some discussion should be provided.", "other_comments_or_suggestions": "Please refer to Other Strengths And Weaknesses part.", "questions_for_authors": "Please refer to Other Strengths And Weaknesses part.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LZUkoALmm1", "reviewer": "Reviewer_vxgt", "summary": "This work propose a minibatch sampling strategy to select pairs of samples in the mini-batch to enhance the OOD geralization ability of SSL methods. By investigating on a causal perspective from the constructed SCM model, the method propose a Post-Intervention Distribution, which can be realized by balancing score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Response to Weaknesses 1 & Experimental Designs Or Analyses**: \n\nThank you for pointing this out. We clarify this issue through the following steps:\n\n**Step 1: How the original submission constructs the OOD task**\n\nThe transfer learning task and the few-shot learning task can be regarded as OOD (out-of-distribution) tasks, as the training and test datasets in these tasks follow different data distributions. Meanwhile, in **Appendix C.3** of the original submission, we also provide evaluation results on two OOD datasets, namely the Colored-MNIST dataset and the PACS dataset.\n\n**Step 2: Results on Waterbirds dataset and CMNIST dataset**\n\nFor Waterbirds, we follow the implementation in ZARE et al. (2023) \"Evaluating and Improving Domain Invariance in Contrastive Self-Supervised Learning by Extrapolating the Loss Function\". During training, waterbirds (landbirds) predominantly appear on water (land) backgrounds; however, the distribution is altered at test time. We report both the average and worst-group performance. The results in the table below demonstrate that our method yields consistent improvements, particularly enhancing performance for the worst-performing group.\n\n| Method      | Test Accuracy | Worst Group |\n| ----------- | ------------- | ----------- |\n| SimCLR      | 76.2          | 19.2        |\n| SimCLR+Ours | 78.0          | 24.9        |\n| MAE         | 74.9          | 17.6        |\n| MAE+Ours    | 77.2          | 22.1        |\n\nFor Colored-MNIST, we follow the implementation in Huang et al. (2024) \"On the Comparison between Multi-modal and\nSingle-modal Contrastive Learning\". The task is a 10-class digit classification, with 10% of the labels randomly reassigned. During training, images belonging to class ‘0’ (or ‘1’) are colored red (or green) with a probability of 77.5%, and another random color with a probability of 22.5%. For the test set, the coloring scheme is reversed relative to training, which allows us to evaluate the extent to which the model relies on color cues for classification. The results in the table below show that our method improves OOD test accuracy by nearly 10%.\n\n| Method      | Test Accuracy |\n| ----------- | ------------- |\n| SimCLR      | 12.7          |\n| SimCLR+Ours | 23.5          |\n| MAE         | 15.1          |\n| MAE+Ours    | 24.9          |\n\nThese results further demonstrate that our proposed method effectively improves the OOD generalization performance of SSL.\n\n----\n\n**Response to Question 2**:\n\nThank you for pointing this out. According to Assumption 3.3 in the original submission, both $ x^{\\rm label} $ and $s$ can be obtained through an invertible neural network. Training such an invertible neural network typically requires training data in the form of $(x^{\\rm label} _i, s _i, x _i^+) _{i=1}^N $. However, we did not adopt this mechanism directly because we did not have access to such training data. In particular, we were unable to provide the corresponding $ s_i $ for each pair. Therefore, we opted to use a **Learning Latent Variable Model** approach instead.", "claims_and_evidence": "I do not find any evident errors in the claims.", "methods_and_evaluation_criteria": "The method seems convincing, essentially when the balancing condition holds, there is no changes of distribution in spurious features, therefore the SSL method will focus on the invariant features.", "theoretical_claims": "The results of Theorem 3.4 should be correct that the loss minimizes the worst-case risk, this is a well-defined target in invariant learning literatures. I am not sure if Theorem 4.3 is correct as I am not familiar with the identifiability theory.", "experimental_designs_or_analyses": "This work primarily addresses the OOD generalizability of SSL methods; however, the experiments do not include any OOD datasets, such as _Waterbirds and CMNIST_. Conducting experiments directly on OOD datasets would help evaluate the effectiveness of the proposed sampling strategy.", "relation_to_broader_scientific_literature": "This work relates to the self-supervised learning literatures and domain generalization, as well as causal inference literatures.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "1. The biggest concern for me is the experiment setting, i.e., it does not involve any OOD dataset for evaluation, however, the main goal of this work is to enhance OOD generalization ability of SSL methods.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Why using a exponential family distribution to model $p(s \\mid x^{label})$? Why not using a reversible neural nets which may achieve higher expressibity and simplify the design?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HWha6ZBbzM", "reviewer": "Reviewer_fDza", "summary": "This paper focuses on the out-of-distribution generalization of self-supervised learning. The authors first give one plausible explanation for SSL having OOD generalization, then analyze and conclude that SSL learns spurious correlations during the training process from the perspective of generation and causal inference. To address this issue, they further propose a post-intervention distribution (PID) grounded in the Structural Causal Model. Experiments verify the advantages of their method.\n\n## Update After Rebuttal\nThank you for addressing my concerns. I maintain my positive score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Response to Weaknesses 1 & Questions 3**: \n\nThank you for pointing these out. The proposed method has two main phases with the following complexity analysis per mini-batch (batch size $B$, dataset size $D$):\n\n**Step 1: Latent Variable Model Training:**\n   - **$q_\\phi(s|x^+,x^{\\rm label})$:** Each sample requires a forward pass with cost $O(C_\\phi)$, totaling $O(B \\cdot C_\\phi)$.\n   - **$p_f(x^+|s,x^{\\rm label})$:** Each sample incurs a cost $O(C_f)$, totaling $O(B \\cdot C_f)$.\n   - **$g$ for $\\lambda^e$:** Computed once per mini-batch with cost $O(C_g)$.\n   - **KL-Divergence:** Involves operations over the latent dimension $n$ and sufficient statistic dimension $k$, contributing $O(B\\cdot n\\cdot k)$.\n   - **Orthogonality Regularization:** Requires $O(n\\cdot k^2)$, which is constant when $n$ and $k$ are small.\n\n   Thus, the training phase complexity is approximately:   $\n   O\\Big(B\\cdot (C_\\phi + C_f + n\\cdot k) + C_g + n\\cdot k^2\\Big).\n   $\n\n**Step 2:  Algorithm 1**\n   - **Propensity Score Calculation:** For each sample, computing scores across the $D$ candidates costs $O(D\\cdot n\\cdot k)$, leading to a total of $O(D^2\\cdot n\\cdot k)$ for the mini-batch.\n   - **Matching Operation:** A brute-force matching over $D$ samples yields an additional $O(D^2)$.\n\n   Therefore, the sampling phase has an overall complexity of approximately:  $\n   O(D^2\\cdot n\\cdot k).\n   $\n\n**Step 3: Overall Complexity**\n\nThe combined complexity per mini-batch is: $\nO\\Big(B\\cdot (C_\\phi + C_f + n\\cdot k) + C_g + n\\cdot k^2 + D^2\\cdot n\\cdot k\\Big).\n$\n\nThe symbols $C_\\phi$, $C_f$, and $C_g$ represent the computational cost for a single forward pass (or operation) of each respective network module. For specific computational resources and time, please refer to **Response to Weaknesses 2** of **Rebuttal for Reviewer 1Q1V**\n\n----\n\n**Response to Weaknesses 2**\n\nThank you for pointing this out. We provide a deeper exploration of **Assumption 3.3** and **Assumption 4.1** in our response.\n\nFor **Assumption 3.3**, we implicitly assume that the noise is independent of both $ s $ and $ x^{\\rm label} $. However, in many practical scenarios, noise may be correlated with either the latent variables or the observed features—for example, sensor noise that correlates with lighting conditions in image data—which can interfere with the separation between causal and non-causal factors.\n\nRegarding **Assumption 4.1**, the main concern lies in the potential mismatch between the true conditional distribution in real-world data and the assumed exponential family. If the actual distribution is more complex or exhibits behaviors that go beyond this family—such as multi-label or multi-instance characteristics—then the applicability of our method may be compromised.\n\n----\n\n**Response to Weaknesses 3**: \n\nThank you for pointing this out. In the final version, we will add a table to illustrate all terms and variables related to our method.\n\n----\n\n**Response to Questions 1**: \n\nThank you for pointing this out. We explain this issue through the following steps:\n\n**Step 1: How do we implement PID**\n\nAccording to *Definition 4.4* in the original submission, $s$ and $x^{\\rm label}$ are conditionally independent given $ba(s)$. Based on this, if all pairs in a mini-batch share the same $ba(s)$, then within this mini-batch, $s$ and $ x^{\\rm label} $ can be considered independent. Consequently, such a mini-batch can be viewed as being sampled from a **PID**.\n\n**Step 2: How is this integrated into SSL**\n\nIn the training phase of SSL, a mini-batch is typically sampled from the training data prior to each iteration. In standard SSL, this mini-batch is randomly sampled. In contrast, our method constructs the mini-batch using **Algorithm 1** from the original submission. That is, our approach embeds into SSL by replacing the mini-batch sampling process with **Algorithm 1**, without altering any other part of the SSL training procedure.\n\nAccording to **Algorithm 1**, the core criterion for selecting samples is to ensure that the $ ba(s) $ values of each pair are as similar as possible. This ensures that the resulting mini-batch has consistent $ ba(s) $ across all samples, thereby forming a PID.\n\n----\n\n**Response to Questions 2**: \n\nThank you for pointing this out. Instead of proposing specific metrics or evaluation criteria, we run a toy experiment on the COCO dataset [1] with two different experimental settings: 1) training and testing the SSL model on full images; 2) training and testing the SSL model on foreground images. Setting 2) can be thought of as not being subject to background semantic confounding. For the Top 1  classification accuracy, the results of  SimCLR are 39.66 and 50.19, the results of  SimCLR + Ours are 45.25 and 51.48. We observe that our method gives closer results in both settings and significantly outperforms SimCLR. Thus, it can be concluded that our method learns less spurious correlations. \n\n[1] Microsoft coco: Common objects in context. ECCV, 2014.", "claims_and_evidence": "It seems convincing.", "methods_and_evaluation_criteria": "The method and evaluation criteria make sense.", "theoretical_claims": "I have checked the proofs of theoretical claims, and the entire theoretical derivation and claims seem appropriate.", "experimental_designs_or_analyses": "I have checked the experimental designs and analyses, and the experimental results are impressive.", "relation_to_broader_scientific_literature": "By analyzing the mini-batch construction during the SSL training phase, this paper gives one plausible explanation for self-supervised learning (SSL) having OOD generalization. Moreover, this paper also analyzes and concludes that SSL learns spurious correlations during the training process, which leads to a reduction in OOD generalization.", "essential_references_not_discussed": "The essential references seem sufficient.", "other_strengths_and_weaknesses": "Strengths:\n1. This article provides a wealth of theoretical analysis, making the entire work more solid.\n2. The experimental results are impressive.\n\n\nWeaknesses:\n1. While the theoretical aspects are robust, the practical implementation of these concepts, especially the integration of causal inference in SSL, might be complex and computationally intensive. This could limit its applicability in environments with constrained computational resources.\n2. Some of the causal assumptions made may not hold in all real-world scenarios, which could affect the generalizability of the findings. A deeper exploration of these assumptions, including conditions under which they may not be valid, would provide a more comprehensive view of the method’s applicability.\n3. Some key terms and variables used throughout the paper could be defined more clearly to avoid ambiguity, enhancing the paper’s accessibility to a broader audience.", "other_comments_or_suggestions": "See weaknesses (above) and questions (below).", "questions_for_authors": "1. Could you provide more detailed insights into how the Post-Intervention Distribution (PID) is specifically integrated into the self-supervised learning process? It would be beneficial to understand the operational steps or algorithms used to enforce PID constraints during mini-batch preparation.\n2. You mentioned that the proposed method minimizes spurious correlations. Can you discuss any specific metrics or evaluation criteria used to measure the extent of spurious correlations before and after applying your method?\n3. Could you discuss the scalability of your proposed method, particularly in terms of computational resources and time required as the dataset size increases? Is the method feasible for large-scale real-world applications where computational efficiency is critical?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The reviewers generally agree that this paper should be accepted. I concur, albeit on the margins. The reviewers found the theory provided to be solid but raised some questions on both the range of the empirical validation and the generalizability of the proposed method, given notable additional computational needs.

*Strengths.* The theoretical motivations of the work are robust and comprehensive.

*Weaknesses.* The proposed method requires some unverifiable assumptions that may not fully hold in practice. The empirical results give some cover for this, but could be expanded. The proposed method also requires notably more computation resources, which introduces questions about tradeoffs between compute and added performance.

The authors provided clarifications and additional OOD dataset experiments and computational estimates. The set of OOD datasets could be expanded, e.g., Domainbed/WILDS datasets. However, the strengths of the paper outweigh the weaknesses.
Decision: Accept (poster)