Revised: T OWARDS

Revised: THE

Revised: T

Original: O UT - OF -D ISTRIBUTION G ENERALIZATION OF S ELF -S UPERVISED C ONTRASTIVE L EARNING
Revised: RANSFERABILITY OF S ELF - S UPERVISED C ONTRASTIVE L EARNING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Self-supervised learning  attracts much attention recently, since it does not require labeled data for training contrasted to supervised learning.
Revised: Self-supervised learning (SSL) attracts much attention recently, since it does not require labeled data for training contrasted to supervised learning.

Original: Empirical studies observe that it also has better transfer ability than supervised learning.
Revised: Empirical studies have shown that SSL also has promising performance in distribution shift problems, where the downstream and training distributions are different.

Original: However, the theoretical study of the out-of-distribution (OOD) generalization ability of self-supervised learning is still limited.
Revised: However, a theoretical understanding of its transferability under distribution shifts is still limited.

Original: In this paper, we establish a theoretical framework to analyze the OOD performance of contrastive-based self-supervised learning, by connecting the OOD generalization ability with the  data augmentation used in SSL.
Revised: In this paper, we establish a theoretical framework to analyze the transferability of contrastive-based self-supervised learning, by studying the role of data augmentation used in it.

Original: Although some recent work claims that  contrastive learning learns more robust representations than supervised learning, our results suggest that this superiority mainly comes from the data augmentation used, i.e., more data are fed to the model.
Revised: Our results show that downstream distributions in which contrastive learning performs well depend largely on data augmentation.

Original: In the face of more challenging OOD scenarios, the standard contrastive learning still suffers from the same generalization problem as empirical risk minimization (ERM).
Revised: Besides, we also find that contrastive learning can not learn domain-invariant features, which limits its transferability under distribution shifts.

Original: Based on our theoretical results, we propose an augmentation-robust contrastive learning approach, named as ArCL, which significantly improves the OOD performance of  contrastive learning in several datasets.
Revised: Based on these theoretical results, we propose the augmentation-robust contrastive learning, named as ArCL, which provably learns domain-invariant features. Our experiments on several datasets show that ArCL improves the distribution shift transferability of original contrastive learning significantly.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Classical machine learning algorithms commonly assume that training and test samples come from the same distribution, i.e., the IID setting.
Revised: Classical machine learning algorithms commonly assume that training and test samples come from the same distribution, i.e., the IID setting.

Original: However, in real-world applications, this assumption can be easily violated, and classical algorithms can perform poorly due to  thedifference between training and test distribution.
Revised: However, in real-world applications, this assumption can be easily violated, and classical algorithms can perform poorly due to distribution shifts, where the training and test distributions are different.

Original: This motivates the study of out-of-distribution (OOD) generalizationproblem, which aims to develop algorithms that can generalize to unseen test distributions (Blanchard et al., 2011; Muandet et al., 2013; Wang et al., 2021a;
Revised: This issue has motivated a plethora of research in different settings, including transfer learning, domain adaptation and domain generalization (Blanchard et al., 2011; Muandet et al., 2013; Wang et al., 2021a;

Original: Shen et al., 2021).
Revised: Shen et al., 2021).

Original: Obviously, generalization to arbitrary distributions is impossible.
Revised: Obviously, generalization to arbitrary distributions is impossible.

Original: Different characterization of the connection between test and training distributions leads to distinct algorithms.
Revised: Different characterization of the connection between test and training distributions leads to distinct algorithms.

Original: In the supervised learning regime, most of the proposed approaches aim to find the features that extract some invariance across different distributions, and assume that such invariance also holds inunseen test distributions (Peters et al., 2016; Rojas-Carulla et al., 2018;
Revised: In the supervised learning regime, most of the proposed approaches aim to find the features that extract some invariance across different distributions, and assume that such invariance also holds in test distributions (Peters et al., 2016; Rojas-Carulla et al., 2018;

Original: Arjovsky et al., 2019; Mahajan et al., 2021; Jin et al., 2020;
Revised: Arjovsky et al., 2019; Mahajan et al., 2021; Jin et al., 2020;

Original: Ye et al., 2021).
Revised: Ye et al., 2021).

Original: For the self-supervised learning (SSL) setting, related works are still limited.
Revised: For the self-supervised learning (SSL) setting, related works are still limited.

Original: Self-supervised learning has attracted great attention in many fields (He et al., 2020; Chen et al., 2020; Grill et al., 2020;
Revised: Self-supervised learning has attracted great attention in many fields (He et al., 2020; Chen et al., 2020; Grill et al., 2020;

Original: Chen & He, 2021; Zbontar et al., 2021), which firstly learns a representation through a large number of unlabeled training data, and then uses the learned encoder for downstream tasks.
Revised: Chen & He, 2021; Zbontar et al., 2021). It firstly learns a representation through a large number of unlabeled training data, and then fine-tunes the learned encoder to obtain a final model on the downstream task.

Original: Since its twostep nature, self-supervised learning is more likely to meet the OOD issue,  andits generalization ability on downstreamtasks with different distributions  has drawn increasing attention.
Revised: Since its two-step nature, self-supervised learning is more likely to meet the distribution shift issue, i.e., training and downstream distributions could be different. Investigating its transferability under distribution shifts has drawn increasing attention.

Original: Some recent works have studied this issue empirically (Liu et al., 2021;
Revised: Some recent works have studied this issue empirically (Liu et al., 2021;

Original: Goyal et al., 2021; von K¨ugelgen et al., 2021; Wang et al., 2021b).
Revised: Goyal et al., 2021; von K¨ugelgen et al., 2021; Wang et al., 2021b; Shi et al., 2022).

Original: However, athorough theoretical understandingof the OOD generalization ability of contrastive learning is still limited, which  hinders the development oftailor-made OOD algorithms.
Revised: However, a theoretical understanding is still limited, which also hinders the development of algorithms.

Original: SSL, especially contrastive-based SSL (contrastive learning), employs the idea that finding representations that are invariant under some data augmentation, which is similar to  the invariance-based
Revised: In this paper, we try to theoretically study the distribution shift transferability of contrastive based self-supervised learning (contrastive learning).

Original: OOD methods.

Original: From this point of view, we wonder whether contrastive learning learns a domain invariant representation.
Revised: Specifically, we wonder on which downstream distributions the representation obtained by contrastive learning has good performance.

Original: In this paper, we try to theoretically answer the following questions: (a) on which downstream distributions, the representation obtained by contrastive learning has good generalization performance; (b) whether contrastive learning learns a domain invariant representation. We investigate thefirst question by relating contrastive loss to downstream risks, and finds the crucial role of data  augmentation. 
Revised: We study this problem by establishing relationship between the contrastive loss and the downstream risk, and finds the crucial role of data augmentation: on downstream tasks whose distributions are close to the augmented training distribution, contrastive learning provably performs well. Besides, the idea behind contrastive learning is to find representations that are invariant under data augmentation. This is similar to the domain-invariance based supervised learning methods, since applying each augmentation to the training data can be viewed as inducing a specific domain.

Original: For the second question, unfortunately, we discover that contrastive learning fails to learn a domain-invariant representation, even across domains induced by data augmentation.
Revised: Unfortunately, from this point of view, we find that contrastive learning fails to learn a domain-invariant representation, which limits its transferability.

Original: To address this issue, we proposethe augmentation-robust contrastive learning (ArCL) approach, a module that can be applied to various popular contrastive learning algorithms, such as SimCLR and MoCo.
Revised: To address this issue, we propose augmentation-robust contrastive learning (ArCL), a module that can be applied to various popular contrastive learning algorithms, including SimCLR and MoCo.

Original: It forces thelearned representation to align positive samples uniformly, and thus can improve its OOD generalization performance with theoretical guarantee.
Revised: Compared with original contrastive learning, ArCL forces the representation to align positive samples uniformly , and hence provably learns domain-invariant representations.

Original: We conduct several experiments on CIFAR10 and ImageNet  to show the performance of our method.
Revised: Empirically, we conduct ArCL on CIFAR10 and ImageNet respectively to learn encoders, and then test their transferability on various downstream tasks.

Revised: Our experiments show that ArCL improves the original contrastive learning algorithms significantly.

Original: R ELATED W ORK
Revised: R ELATED W ORK

Original: OOD generalization for supervised learning.
Revised: Distribution shift in supervised learning.

Original: OOD generalization, has drawn much attention recently (Blanchard et al., 2011; Muandet et al., 2013; Wang et al., 2021a;
Revised: Distribution shift has drawn much attention recently (Blanchard et al., 2011; Muandet et al., 2013; Wang et al., 2021a;

Original: Shen et al., 2021).
Revised: Shen et al., 2021).

Original: Most of the works focus on learning a representation that performs well on different source domains simultaneously (Rojas-Carulla et al., 2018; Mahajan et al., 2021; Jin et al., 2020), which follows the idea of causal invariance (Peters et al., 2016; Arjovsky et al., 2019).
Revised: Most of the works address this issue by learning a representation that performs well on different source domains simultaneously (Rojas-Carulla et al., 2018; Mahajan et al., 2021; Jin et al., 2020), which follows the idea of causal invariance (Peters et al., 2016; Arjovsky et al., 2019).

Original: To obtain theoretical guarantees for these methods, structural equation models are often assumed (von K ¨ ugelgen et al., 2021;
Revised: To obtain theoretical guarantees for these methods, structural equation models are often assumed (von K ¨ ugelgen et al., 2021;

Original: Liu et al., 2020; Mahajan et al., 2021).
Revised: Liu et al., 2020; Mahajan et al., 2021).

Original: There are also many other approaches.

Original: Distributionally robust optimization optimize a model’s worst-case performance over some uncertainty set directly (Krueger et al., 2021; Sagawa et al., 2019; Duchi & Namkoong, 2021; Duchi et al., 2021).
Revised: Distributionally robust optimization optimize a model’s worstcase performance over some uncertainty set directly (Krueger et al., 2021; Sagawa et al., 2019; Duchi & Namkoong, 2021; Duchi et al., 2021).

Original: Stable learning (Shen et al., 2020; Kuang et al., 2020) learns a set of global sample weights that could remove the confounding bias for all the potential treatments from data distribution.
Revised: Stable learning (Shen et al., 2020; Kuang et al., 2020) learns a set of global sample weights that could remove the confounding bias for all the potential treatments from data distribution.

Original: Disentangled representation learning (Bengio et al., 2013; 
Revised: Disentangled representation learning (Bengio et al., 2013; Tr¨auble et al., 2021;

Original: Tr ¨ auble et al., 2021; Kim & Mnih, 2018) aims to learn representations where distinct and informative factors of variations in data are separated.
Revised: Kim & Mnih, 2018) aims to learn representations where distinct and informative factors of variations in data are separated.

Original: Contrastive self-supervised learning.
Revised: Theoretical understanding of contrastive learning.

Original: Early works such as MoCo (He et al., 2020) and SimCLR (Chen et al., 2020) use the InfoNCE loss to pull the positive samples close while enforcing them away from the negative samples in the embedding space.

Original: These methods require large batch sizes or carefully designed negative sampling strategies.

Original: To obviate these issues, some recent works get rid of negative samples and prevent representation collapse by cross-correlation loss (Zbontar et al., 2021) or optimization tricks (Grill et al., 2020;

Original: Chen & He, 2021).

Original: A number of recent works also aim to theoretically explain the success of contrastive learning in IID settings.
Revised: A number of recent works also aim to theoretically explain the success of contrastive learning in IID settings.

Original: One way to explain it is through the mutual information between positive samples (Tian et al., 2020; Hjelm et al., 2018; Tschannen et al., 2019).
Revised: One way to explain it is through the mutual information between positive samples (Tian et al., 2020; Hjelm et al., 2018; Tschannen et al., 2019).

Original: Arora et al.
Revised: Arora et al.

Original: (2019) directly analyze the generalization of InfoNCE loss based on the assumption that positive samples are drawn from the same latent classes.
Revised: (2019) directly analyze the generalization of InfoNCE loss based on the assumption that positive samples are drawn from the same latent classes.

Revised: In the same setting, Bao et al.

Revised: (2022) establishes equivalence between InfoNCE and supervised loss by giving sharper upper and lower bounds, which also quantifies the effect of the number of negative samples.

Original: Wang & Isola (2020) studies the behavior of InfoNCE loss from the perspective of alignment and uniformity.
Revised: Wang & Isola (2020) studies the behavior of InfoNCE loss from the perspective of alignment and uniformity.

Original: Huang et al.
Revised: Huang et al.

Original: (2022) take data augmentation into account, and provide generalization bounds for k -nearest neighbor error.
Revised: (2022) takes data augmentation into account, and provide generalization bounds for k -nearest neighbor error.

Original: The study of contrastive learning in domain shift setting is still limited.
Revised: The study of contrastive learning in domain shift setting is still limited.

Original: (Shen et al., 2022) and (HaoChen et al., 2022) study contrastive learning in unsupervised domain adaptation setting, where unlabeled target data are obtained. 
Revised: (Shen et al., 2022) and (HaoChen et al., 2022) study contrastive learning in unsupervised domain adaptation setting, where unlabeled target data are obtained. Shi et al.

Revised: (2022) evaluates the performance of SSL, auto-encoders and supervised learning on distribution shift datasets, and conclude that SSL is consistently most robust.

Original: Other robust contrastive learning methods (Kim et al., 2020; Jiang et al., 2020)propose mainly focus on the adversarial robustness , not the distributional robustness studied in this paper.
Revised: Other robust contrastive learning methods (Kim et al., 2020; Jiang et al., 2020) mainly focus on the adversarial robustness, not the distributional robustness studied in this paper.

Original:  P ROBLEM F ORMULATION
Revised:  P ROBLEM F ORMULATION

Original: Given a set of unlabeled data where each sample X is i.i.d.
Revised: Given a set of unlabeled data where each sample X is i.i.d.

Original: sampled from training data distribution D on X ⊆ R d , the goal of Self-Supervised Learning (SSL) is to learn a general encoder f : X → R m for different downstream tasks.
Revised: sampled from training data distribution D on X ⊆ R d , the goal of Self-Supervised Learning (SSL) is to learn an encoder f : X → R m for different downstream tasks.

Original: Contrastive learning is a popular approach of SSL, which augments each sample X twice to obtain a positive pair ( X 1 , X 2 ) , and then learns the encoder f by pulling them close and pushing random samples (also called negative samples) away in the embedding space.
Revised: Contrastive learning is a popular approach of SSL, which augments each sample X twice to obtain a positive pair ( X 1 , X 2 ) , and then learns the encoder f by pulling them close and pushing random samples (also called negative samples) away in the embedding space.

Original: The data augmentation is done by applying a transformation A to the original data, where A is randomly selected from a transformation set A according to the distribution π .
Revised: The data augmentation is done by applying a transformation A to the original data, where A is randomly selected from a transformation set A according to the distribution π .

Original: We use D A to denote the distribution of augmented data under a specific augmentation A , and use D π to denote the distribution of augmented data, i.e., D π = (cid:82) A D A d π ( A ) .
Revised: We use D A to denote the distribution of augmented data under a specific transformation A , and use D π to denote the distribution of augmented data, i.e., D π = (cid:82) A D A d π ( A ) .

Original: The loss of contrastive learning typically consists of two terms (Huang et al., 2022):
Revised: The loss of contrastive learning typically consists of two terms:

Original: [Equation]
Revised: [Equation]

Original: where L align ( f ; D , π ) :
Revised: where L align ( f ; D , π ) :

Original: =
Revised: =

Original: E X
Revised: E X

Original: ∼D E
Revised: ∼D E

Original: ( A 1 ,A 2 ) ∼ π 2 ∥
Revised: ( A 1 ,A 2 ) ∼ π 2 ∥

Original: f
Revised: f

Original: ( A 1 ( X ))
Revised: ( A 1 ( X ))

Original: − f
Revised: − f

Original: ( A 2 ( X ))
Revised: ( A 2 ( X ))

Original: ∥ 2 measures the alignment of positive samples, and L reg ( f ; D , π ) is the regularization term to avoid feature collapse.
Revised: ∥ 2 measures the alignment of positive samples, and L reg ( f ; D , π ) is the regularization term to avoid feature collapse.

Original: For example, the regularization term in InfoNCE loss is
Revised: For example, the regularization term in InfoNCE loss is

Original: [Equation]
Revised: [Equation]

Original: When subscript is omitted, notation ∥ · ∥ stands for L 2 -norm or Frobenius norm for vectors and matrices in this paper.
Revised: When subscript is omitted, notation ∥ · ∥ stands for L 2 -norm or Frobenius norm for vectors and matrices in this paper.

Original: The encoder f is obtained by minimizing L con ( f ; D , π ) .
Revised: The encoder f is obtained by minimizing L con ( f ; D , π ) .

Original: For simplicity, we assume that ∥ f ( x ) ∥ = 1 for every x ∈ X and f is L -Lipschitz continuous throughout this paper.
Revised: For simplicity, we assume that ∥ f ( x ) ∥ = 1 for every x ∈ X and f is L -Lipschitz continuous throughout this paper.

Original: For a downstream task with target distribution D tar , where the labels are obtained, the learned encoder f will be  used.
Revised: For a downstream task with target distribution D tar , where the labels are obtained, the learned encoder f will be fine-tuned and used.

Original: In this paper, we consider a multi-class classification problem X ⊆ R d → Y = { 1 , . ..
Revised: In this paper, we consider a multi-class classification problem X ⊆ R d → Y = { 1 , . .

Original:  , K } .
Revised: . , K } .

Original: We focus on the covariate shift setting, i.e., we allow D tar to be different from D but P ( Y | X ) is fixed.
Revised: We focus on the covariate shift setting, i.e., we allow D tar to be different from D but P ( Y | X ) is fixed.

Original: We use the best linear classifier on top of f for classification, and use its risk to evaluate the performance of f over D tar .
Revised: We use the best linear classifier on top of f for classification, and use its risk to evaluate the performance of f on D tar .

Revised: This evaluation method is also used empirically in Shi et al.

Original: Namely, we define
Revised: Specifically, we define

Original: R

Original: [Equation]
Revised: [Equation]

Original: as the downstream risk of representation f on distribution D tar , where ℓ is some loss function.
Revised: as the downstream risk of representation f on distribution D tar , where ℓ is some loss function.

Original: We focus on the case where ℓ is the square loss.
Revised: We focus on the case where ℓ is the square loss.

Original: With a slight abuse of notation, for the classifier h ◦ f : R m → R k
Revised: With a slight abuse of notation, for the classifier h ◦ f : R m → R k

Original: we also use R ( h ◦ f ; D tar ) to denote its risk.
Revised: we also use R ( h ◦ f ; D tar ) to denote its risk.

Original: In this paper, we study the generalization ability of representation learned by contrastive learning on downstream tasks with different distributions.
Revised: In this paper, we study the generalization ability of representation learned by contrastive learning on downstream tasks with different distributions.

Original: At the end of this section, we remark the difference between transformations and augmentations.
Revised: At the end of this section, we remark the difference between transformations and augmentations.

Original: Transformations are deterministic mapping from X to X , and augmentations are random application of them.
Revised: Transformations are deterministic mapping from X to X , and augmentations are random application of them.

Original: The purpose of this hierarchical model is as follows: (a) even A is fixed, the choice of augmentation π is crucial to both the IID and OOD performance of contrastive learning; (b) most augmentation methods used in practical indeed satisfy this model.The former we will see in the following theoretical analysis, and here we give examples to demonstrate the latter.
Revised: We give examples to demonstrate that many augmentation methods used in practice indeed satisfy this model.

Original: Example 2.1.
Revised: Example 2.1.

Original: The augmentations used in SimCLR (Chen et al., 2020) are compositions of random transformations, such as RandomCrop, HorizonalFlip and Color distortion.
Revised: The augmentations used in SimCLR (Chen et al., 2020) are compositions of random transformations, such as RandomCrop, HorizonalFlip and Color distortion.

Original: Each A denotes a specific composition here.
Revised: Each A denotes a specific composition here.

Original: Users will determine probabilities that each transformation will be applied.
Revised: Users will determine probabilities that each transformation will be applied.

Original: If used, the parameters of the random transformation such as crop range are also selected randomly.
Revised: If used, the parameters of the random transformation such as crop range are also selected randomly.

Original: Until the parameters of each used transformations are set, the composited transformation is deterministic.
Revised: Until the parameters of each used transformations are set, the composited transformation is deterministic.

Original: Therefore, we can view these augmentation processes as randomly selecting some deterministic transformations, following the distribution π determined by users.
Revised: Therefore, we can view these augmentation processes as randomly selecting some deterministic transformations, following the distribution π determined by users.

Original: Example 2.2.
Revised: Example 2.2.

Original: Some recent works (Jahanian et al., 2021) propose to augment data by transformations in the latent space.
Revised: Some recent works (Jahanian et al., 2021) propose to augment data by transformations in the latent space.

Original: Suppose we have a pre-trained generative model with an encoder
Revised: Suppose we have a pre-trained generative model with an encoder

Original: T : X → Z and a generator G : Z → X , where Z is a latent space.
Revised: T : X → Z and a generator G : Z → X , where Z is a latent space.

Original: Since Z is usually simple, one can augment data in X by randomly shifting T ( X ) .
Revised: Since Z is usually simple, one can augment data in X by randomly shifting T ( X ) .

Original: In this setting, the augmentation can be parameterized by A θ ( X ) =
Revised: In this setting, the augmentation can be parameterized by A θ ( X ) =

Original: G ( T ( X ) + θ ) for θ ∈ Z .
Revised: G ( T ( X ) + θ ) for θ ∈ Z .

Original: The distribution π of θ is usually chosen to be a normal distribution.
Revised: The distribution π of θ is usually chosen to be a normal distribution.

Original:  O UT - OF - DISTRIBUTION A NALYSIS : THE C RUCIAL R OLE OF D ATA A UGMENTATION
Revised:  T RANSFERABILITY A NALYSIS : THE C RUCIAL R OLE OF D ATA A UGMENTATION

Original: In this section, we theoretically demonstrate that data augmentation plays a crucial role in the superior OOD performance of contrastive learning.
Revised: In this section, we theoretically demonstrate that data augmentation plays a crucial role in the transferability of contrastive learning.

Original: The first step is to establish the connection between contrastive loss L CL ( f  ) and downstream risk R ( f  ) .
Revised: The first step is to establish the connection between contrastive loss L con ( f ; D , π ) and downstream risk R ( f ; D tar ) .

Original: For this purpose, we introduce the ( σ, δ ) -augmentation notion proposed in (Huang et al., 2022), which was used there to study the KNN error of contrastive learning.
Revised: For this purpose, we introduce the ( σ, δ ) -augmentation notion proposed in (Huang et al., 2022), which was used there to study the KNN error of contrastive learning.

Original: Definition 3.1 ( ( σ, δ ) -transformation) .
Revised: Definition 3.1 ( ( σ, δ ) -transformation) .

Original: Let C k
Revised: Let C k

Original: ⊆
Revised: ⊆

Original: X be the set of all points in class k .
Revised: X be the set of all points in class k .

Original: A data transformation set A is called a ( σ, δ ) -transformation on D for some σ ∈ (0 , 1] and δ > 0 , if for every k ∈ [ K ] , there exists C 0 k ⊆ C k such that
Revised: A data transformation set A is called a ( σ, δ ) -transformation on D for some σ ∈ (0 , 1] and δ > 0 , if for every k ∈ [ K ] , there exists C 0 k ⊆ C k such that

Original: [Equation]
Revised: [Equation]

Original: where d A ( x 1 , x 2 ) : = inf A
Revised: and sup

Original: [Equation]
Revised: [Equation]

Original: ,A
Revised: where d A ( x 1 , x 2 ) :

Original: [Equation]
Revised: = inf

Original:  ∈A d ( A 1 ( x 1 ) , A 2 ( x 2 )) for some distance d ( · , · ) .
Revised: A 1 ,A 2 ∈A d ( A 1 ( x 1 ) , A 2 ( x 2 )) for some distance d ( · , · ) .

Original: This definition measures the diversity of data augmentation.
Revised: This definition measures the concentration of data augmentation.

Original: A transformation set with smaller δ and larger σ clusters original data more, i.e., samples from the same class are closer after augmentation.
Revised: A transformation set with smaller δ and larger σ clusters original data more, i.e., samples from the same class are closer after augmentation.

Original: Therefore, one can expect the learned representation f to have better cluster performance, which was measured by KNN error in (Huang et al., 2022).
Revised: Therefore, one can expect the learned representation f to have better cluster performance, which was measured by KNN error in (Huang et al., 2022).

Revised: We remark that most of contrastive learning theories (Arora et al., 2019; Ash et al., 2021; Bao et al., 2022) assume that given class k , positive pairs are obtained by i.i.d.

Revised: sampling from P X ∼D ( · | X ∈ C k ) .

Revised: Compared with this data generation model, our ( σ, δ ) notion is more practical.

Original: We further extend their KNN result to the downstream risk by taking into account the linear classifier on top of f ,  and obtain the following lemma.
Revised: By taking into account the linear classifier on top of f , we extend the KNN result in (Huang et al., 2022) to the downstream risk, and obtain the following lemma.

Original: Lemma 3.1.
Revised: Lemma 3.1.

Original: Define µ k ( f ; D ) :
Revised: For any distribution D and encoder f , define µ k ( f ; D ) :

Original: = E X ∼D f ( X ) for k ∈ [ K ] .
Revised: = E X ∼D f ( X ) for k ∈ [ K ] .

Original: Suppose that A is a ( σ, δ ) - augmentation.
Revised: Suppose that A is a ( σ, δ ) -transformation.

Original: For any representation f : X → R d 1 and linear layer h ∈ R K × d 1 , we haveR ( h ◦ f ; D π ) ≤ c ∥ h ∥ √ Kσ [Equation]
Revised: For any representation f : X → R d 1 and linear layer h ∈ R K × d 1 , we have

Original: L

Original: [Equation]

Original: f ; D , π

Original: [Equation]
Revised: [Equation]

Original: where c is an absolute constant and τ ( σ, δ ) is a constant which only depends on ( σ, δ ) .
Revised: where c is an absolute constant and τ ( σ, δ ) is a constant which only depends on ( σ, δ ) .

Original: The detail form of τ is shown in the appendix.
Revised: The detail form of τ is shown in the appendix.

Original: The first term in equation 3 is the alignment of f , which is optimized during the contrastive learning pre-training.
Revised: The first term in equation 3 is the alignment of f , which is optimized during the contrastive learning pre-training on D .

Original: The second term is determined solely by the ( σ, δ ) quantity of the data augmentation.
Revised: The second term is determined solely by the ( σ, δ ) quantity of the data augmentation.

Original: Larger σ and smaller δ induce smaller τ ( σ, δ ) .
Revised: Larger σ and smaller δ induce smaller τ ( σ, δ ) .

Original: These two terms together quantify the cluster level of the representation f learned by contrastive learning.
Revised: These two terms together quantify the cluster level of the representation f learned by contrastive learning.

Original: The third term is related to the linear layer h on top of f , and will be minimized in downstream training, where each class center is converted to the corresponding ground-truth label.
Revised: The third term is related to the linear layer h on top of f , and will be minimized in downstream training, where each class center is converted to the corresponding ground-truth label.

Original: If the regularization term L reg is appropriately chosen, the class centers can be distinguished from each other, and the third term can be reduced to 0 by h .
Revised: If the regularization term L reg is appropriately chosen, the class centers can be distinguished from each other, and the third term can be reduced to 0 by h .

Original: Taking the linear classifier into account, we obtain the following main result.
Revised: Taking the linear classifier into account, we obtain the following main result.

Original: Theorem 3.2.
Revised: Theorem 3.2.

Original: Suppose that A is a ( σ, δ ) -transformation on D , and π is the augmentation distribution on A .
Revised: Suppose that A is a ( σ, δ ) -transformation on D , and π is the augmentation distribution on A .

Original: For an encoder f , define
Revised: For an encoder f , define

Original: [Equation]
Revised: [Equation]

Original: for some positive constants c 1 and c 2 .
Revised: for some positive constants c 1 and c 2 .

Original: Then for any f such that { µ k ( f ; D π ) } Kk =1 are linearly independent and γ reg ( f ; D , π ) > 0 , we have
Revised: Then for any f such that { µ k ( f ; D π ) } Kk =1 are linearly independent and γ reg ( f ; D , π ) > 0 , we have

Original: [Equation]
Revised: [Equation]

Original: where c is an absolute constant.
Revised: where c is an absolute constant.

Original: Remark1.
Revised: Remark

Original: (a) The linear independence of { µ k ( f ; D π ) } Kk =1 condition is weak, since its complement  is a null set in Euclidean space under Lebesgue measure, and have 0 probability under any distribution that is absolutely continuous w.r.t.
Revised: (a) The linear independence of { µ k ( f ; D π ) } Kk =1 condition is weak, since its complement event is a null set in Euclidean space under Lebesgue measure, and have 0 probability under any distribution that is absolutely continuous w.r

Original:  Lebesgue measure.
Revised: t. Lebesgue measure.

Original: (b) The quantity γ − 1 reg is essentially the bound of l 2 norm of the linear classifier h which reduces the third term (cid:80)
Revised: (b) The quantity γ − 1 reg is essentially the bound of L 2 norm of the linear classifier h which reduces the third term (cid:80)

Original: Kk = D π ( C k 
Revised: Kk =1 D π ( C k ) ∥

Original: ∥ e k − h ◦ µ k ( f ; D π ) ∥ in equation 3 to 0 , i.e., mapping the class centers to their corresponding ground-truth labels.
Revised: e k − h ◦ µ k ( f ; D π ) ∥ in equation 3 to 0 , i.e ., mapping the class centers to their corresponding ground-truth labels.

Original: It can be viewed as a measure of “diversity” of the encoder, which encourages the class centers to be orthogonal.

Original: The more orthogonal the class centers, the larger the γ reg and the better the bounds.
Revised: The more orthogonal the class centers, the larger the γ reg and the better the bounds.

Original: The L align term involved in γ reg indicates the phenomenon that better alignment can loose the requirement of diversity.

Original: The theory developed in (Huang et al., 2022)
Revised: The theory developed in (Huang et al., 2022)

Original: shows that popular algorithms such as SimCLR and Barlow Twins can indeed push away the class centers, and the constant γ reg can be ensured to have a constant gap away from 0 .
Revised: shows that popular algorithms such as SimCLR and Barlow Twins can indeed push away the class centers, and the constant γ reg can be ensured to have a constant gap away from 0 .

Original: Therefore, for simplicity we directly impose this requirement on f .
Revised: Therefore, for simplicity we directly impose this requirement on f .

Original: This theorem shows that contrastive learning on distribution D with augmentation π is essentially optimizing the supervised loss (ERM) on  distribution D π  .
Revised: This theorem shows that contrastive learning on distribution D with augmentation π is essentially optimizing the supervised risk on the augmented distribution D π , instead of the original distribution D .

Original: Therefore, the diversity of π is crucial to the OOD generalization performance of contrastive.
Revised: Therefore, the augmentation π is crucial to the transferability of contrastive learning.

Revised: If the downstream distribution D tar is close to D π , the encoder obtained by contrastive learning performs well on it.

Original: The more diverse π is, we can expect D π to cover more downstream distributions, and contrastive learning has better OOD performance.
Revised: The more diverse π is, we can expect D π to approximate downstream distributions more accurately, and contrastive learning has better transferability.

Revised: Besides, if we have some prior knowledge of downstream distribution, we can specify an appropriate augmentation such that D π and D tar are close in order to improve downstream performance.

Original: For example, consider the case where original data D are gray digits, and the downstream data D red are obtained by dying them red.
Revised: For example, consider the case where original data D are gray digits, and the downstream distribution D tar is obtained by dying them red.

Original: Then if we use random coloring as data augmentation and conduct contrastive learning on D , the learned feature will generalize to D red well, even though D red is far away from D .
Revised: Then if we use random coloring as data augmentation and conduct contrastive learning on D , the learned feature will generalize to D tar well, even though D tar is far away from D .

Original:  O UT - OF - DISTRIBUTION L IMITATIONS OF C ONTRASTIVE LEARNING : A D OMAIN I NVARIANCE V IEW
Revised:  T RANSFERABILITY L IMITATION : A D OMAIN -I NVARIANCE V IEW

Original: Our theoretical results in the last section suggests that data augmentation is crucial to the OOD performance of contrastive learning.
Revised: Our theoretical results in the last section suggests that data augmentation is crucial to the transferability of contrastive learning.

Original: However, there exists a fundamental limitation of contrastive learning: the learned representation is not domain invariant, which is important for a model to generalize to unseen test datasets (Arjovsky et al., 2019; Peters et al., 2016).
Revised: However, there exists a fundamental limitation of contrastive learning: the learned representation is not domain-invariant, which is important for a model to generalize to more downstream datasets (Peters et al., 2016; Arjovsky et al., 2019).

Original: Roughly speaking, domain invariance means that the learned model can extract some intrinsic features, which are invariant across different domains.
Revised: Roughly speaking, domain-invariance means that the learned model can extract some intrinsic features, which are invariant across different domains.

Original: Therefore, these features can also be applied to new domains.

Original: In supervised domain generalization, a common setting is to have a training domain set consisting of multiple source domains.
Revised: In supervised domain generalization, a common setting is to have a training domain set consisting of multiple source domains.

Original: In contrastive learning, naturally, we can also consider D A as a transformation-induced domain for each data transformation A , and consider {D A } A ∈A as the training domain set.
Revised: In contrastive learning, correspondingly, we can also consider D A as a transformation-induced domain for each data transformation A , and consider {D A } A ∈A as the training domain set.

Original: Since the goal of contrastive learning is also to align different D A , we naturally expect that the features it learn are domain invariant.
Revised: Since the goal of contrastive learning is also to align different D A , we naturally expect that the features it learn are domaininvariant.

Original: However, since the alignment loss L align in contrastive learning is obtained by averaging  (i.e., taking expectation), the learned features are not invariant even on {D A } A ∈A .
Revised: However, since the alignment loss L align in contrastive learning is obtained by averaging over different transformations (i.e., taking expectation), the learned features are not invariant even across {D A } A ∈A .

Original: In other words, encoders learned through contrastive learning could behave extremely differently in different D A .
Revised: In other words, encoders learned through contrastive learning could behave extremely differently in different D A .

Original: The following toy model gives an illustrative example.
Revised: The following toy model gives an illustrative example.

Original: Proposition 4.1.
Revised: Proposition 4.1.

Original: Consider a two-dimensional classification problem with data ( X 1 , X 2 ) ∼
Revised: Consider a two-dimensional classification problem with data ( X 1 , X 2 ) ∼

Original: N (0 , I 2 ) .
Revised: N (0 , I 2 ) .

Original: The label Y satisfies Y = 

Revised: The label Y satisfies Y =

Original: ( X 1 ≥ 0) , and the data augmentation is to multiply Xby standard normal noise, i.e.,
Revised: ( X 1 ≥ 0) , and the data augmentation is to multiply Xby standard normal noise, i.e.,

Original: [Equation]
Revised: [Equation]

Original: The corresponding transformation-induced domain set is P =
Revised: The corresponding transformation-induced domain set is P =

Original: {D c : D c =
Revised: {D c : D c =

Original: ( X 1 , c · X 2 ) for c ∈ R } .
Revised: ( X 1 , c · X 2 ) for c ∈ R } .

Original: We consider the 0-1 loss in equation 2.
Revised: We consider the 0-1 loss in equation 2.

Original: Then for every ε > 0 , there exists representation f and two domains D c and D c ′ such that
Revised: Then for every ε > 0 , there exists representation f and two domains D c and D c ′ such that

Original: [Equation]
Revised: [Equation]

Original:  but
Revised: ) < ε, but

Original: [Equation]
Revised: [Equation]

Original: This example shows that there exists a representation with arbitrary small contrastive loss, but with very different performance over different transformation-induced domains.
Revised: This example shows that there exists a representation with arbitrary small contrastive loss, but with very different performance over different transformation-induced domains.

Original: The intuition behind this example is that in order to make L align small,  different domains (induced by different augmentations) need to be aligned in an average rather than uniform sense.
Revised: The intuition behind this example is that in order to make L align small, it is sufficient to align different domains (induced by different transformations) in an average rather than uniform sense.

Original: Therefore, the representation may still suffer a large alignment loss on some rarely selected augmented domains.
Revised: Therefore, the representation may still suffer a large alignment loss on some rarely selected augmented domains.

Original:  A R CL: A UGMENTATION -R OBUST C ONTRASTIVE L EARNING
Revised:  A R CL: A UGMENTATION -R OBUST C ONTRASTIVE L EARNING

Original: The above analysis demystifies the OOD generalization ability of contrastive learning on both the pros and cons.

Original: On the one hand, the data augmentation used by contrastive learning extends the training dataset, thus enabling generalization to more downstream tasks.

Original: On the other hand, contrastive learning is still unable to learn domain-invariant features, even on transformation-induced domains.

Original: To further improve the OOD generalization performance of contrastive learning, we consider how to learn a representation that is invariant across {D A } A ∈A .
Revised: To further improve the transferability performance of contrastive learning, we consider how to learn a representation that is domain-invariant across {D A } A ∈A .

Original: First of all, we need a  definition that mathematically characterize domain invariance, so that we can design algorithms based on it.
Revised: First of all, we need a formal definition that mathematically characterize domain-invariance, so that we can design algorithms based on it.

Original: Here, we borrow thefollowing notion proposed by (Arjovsky et al., 2019).
Revised: Here, we borrow the notion proposed by (Arjovsky et al., 2019).

Original: Definition 5.1.
Revised: Definition 5.1.

Original: We say that adata representation f elicits an invariant predictor h 0 ◦ f across a domain set P , if there is a classifier h 0 simultaneously optimal for all domains in P , that is, h 0 ∈ arg min h R
Revised: We say that a representation f elicits an invariant predictor h 0 ◦ f across a domain set P , if there is a classifier h 0 simultaneously optimal for all domains in P , that is, h 0 ∈ arg min h R

Original: ( h ◦ f ; D ) for all D ∈ P .
Revised: ( h ◦ f ; D ) for all D ∈ P .

Original: This definition is equivalent to learning features whose correlations with the target variable are stable, and has been shown to improve OOD generalization both practically and theoretically.
Revised: This definition is equivalent to learning features whose correlations with the target variable are stable, and has been shown to improve distribution shift transferability both practically and theoretically in supervised learning.

Original: Interestingly, by setting P to be {D A } A ∈A , this definition is well suited to contrastive learning.
Revised: Interestingly, by setting P to be {D A } A ∈A , this definition is well suited to contrastive learning.

Original: Indeed, define the augmentation-robust loss as
Revised: Specifically, define the augmentation-robust loss as

Original: [Equation]
Revised: [Equation]

Revised: which is an uniform version of the original alignment loss L align .

Original:  we have the following results.
Revised: Then we have the following results.

Original: Proposition 5.1.
Revised: Theorem 5.1.

Original: For any two transformations A and A ′ , linear predictor h and representation f , we have
Revised: For any two transformations A and A ′ , linear predictor h and representation f , we have

Original: [Equation]
Revised: [Equation]

Original: Moreover, fix f and let h A = arg min h R ( h ◦ f, D A ) .
Revised: Moreover, fix f and let h A ∈ arg min h R ( h ◦ f, D A ) .

Original: Then we have
Revised: Then we have

Original: [Equation]
Revised: [Equation]

Original: L AR ( f ; D ) is an enhanced version of the original alignment loss L align .

Original: Compared to L align , L AR replaces the expectation over A by the supremum, hence L align ( f ; D , π ) ≤ L AR ( f ; D ) for all f and π .
Revised: Compared with L align , L AR replaces the expectation over A by the supremum operator, hence L align ( f ; D , π ) ≤ L AR ( f ; D ) for all f and π .

Original: If L AR ( f ) is small, the above proposition shows that R ( h ◦ f ; D A ) varies slightly with A , so that the optimal h for D A is close to that for D A ′ .
Revised: If L AR ( f ) is small, the above proposition shows that R ( h ◦ f ; D A ) varies slightly with A , so that the optimal h for D A is close to that for D A ′ .

Original: In other words, representation with smaller L AR tends to elicit the same linear optimal predictors across different domains, a property that does not hold for standard alignment loss.
Revised: In other words, representation with smaller L AR tends to elicit the same linear optimal predictors across different domains, a property that does not hold for original alignment loss.

Original: Small L AR not only requires f to push the positive pairs closer, but also forces this alignment to be uniform across all transformations in A .
Revised: Small L AR not only requires f to push the positive pairs closer, but also forces this alignment to be uniform across all transformations in A .

Original: Therefore, we propose to train the encoder using the augmentation-robust loss for better OOD performance, i.e., replace L align by L AR in contrastive loss equation 1.
Revised: Therefore, we propose to train the encoder using the augmentation-robust loss for better distribution shift transfer performance, i.e., replace L align by L AR in contrastive loss equation 1.

Original: In practical implementation, sup A 1 ,A 2 ∈A ∥
Revised: In practical implementation, sup A 1 ,A 2 ∈A ∥

Original: f ( A 1 ( X )) − f
Revised: f ( A 1 ( X )) − f

Original: ( A 2 ( X ))
Revised: ( A 2 ( X ))

Original: ∥ 22 can not be computed exactly, unless |A| is finite and small.
Revised: ∥ 22 can not be computed exactly, unless |A| is finite and small.

Original: Therefore, we use the following approach to approximate it.
Revised: Therefore, we use the following approach to approximate it.

Original: For each sample X , we first randomly select m augmentations to obtain a set of augmented samples denoted by (cid:98) A m ( X ) .
Revised: For each sample X , we first randomly select m augmentations to obtain a set of augmented samples denoted by (cid:98) A m ( X ) .

Original: Then, rather than taking average as in standard contrastive learning, we consider use
Revised: Then, rather than taking average as in standard contrastive learning, we consider use

Original: [Equation]
Revised: [Equation]

Original: as the alignment term, which is an approximation of sup A 1 ,A 2 ∈A ∥ f ( A 1 ( X )) − f ( A 2 ( X ))
Revised: as the alignment term, which is an approximation of sup A 1 ,A 2 ∈A ∥ f ( A 1 ( X )) − f ( A 2 ( X ))

Original: ∥ 22 .
Revised: ∥ 22 .

Original: The integer m is called the number of views , which controls the approximation accuracy.
Revised: The integer m is called the number of views , which controls the approximation accuracy.

Original: Then, given n unlabeled samples ( X 1 , . . .
Revised: Then, given n unlabeled samples ( X 1 , . . .

Original: , X n ) , the empirical augmentation-robust loss is
Revised: , X n ) , the empirical augmentation-robust loss is

Original: [Equation]
Revised: [Equation]

Original: We remark that our proposed method can be applied to any contrastive learning methods whose objective can be formulated by equation 1, such as SimCLR, Moco, Barlow Twins (Zbontar et al., 2021) and so on.
Revised: We remark that our proposed method can be applied to any contrastive learning methods whose objective can be formulated by equation 1, such as SimCLR, Moco, Barlow Twins (Zbontar et al., 2021) and so on.

Original: We show SimCLR + ArCL as an example in Algorithm 1, and provide MoCo + ArCL in the Appendix.
Revised: We give the detailed form of SimCLR + ArCL as an example in Algorithm 1, and provide MoCo + ArCL in the Appendix.

Original: The difference between Algorithm 1 and SimCLR is only the construction of positive pairs.
Revised: The difference between Algorithm 1 and SimCLR is only the construction of positive pairs.

Original: In each epoch t , for every sample X , we first randomly select m transformations denoted by (cid:98) A ( X ) =
Revised: In each epoch t , for every sample X , we first randomly select m transformations denoted by (cid:98) A ( X ) =

Original: { A 1 , . . .
Revised: { A 1 , . . .

Original: A m } .
Revised: A m } .

Original: Then based on current encoder f
Revised: Then based on current encoder f

Original: and projector g , we select two transformations that have worst alignment (minimal inner product) among all pairs in (cid:98) A ( X ) , and use them to construct the finally used positive pairs of X .
Revised: and projector g , we select two transformations that have worst alignment (minimal inner product) among all pairs in (cid:98) A ( X ) , and use them to construct the finally used positive pairs of X .

Original: The construction of negative pairs and update rule are the same as SimCLR.
Revised: The construction of negative pairs and update rule are the same as SimCLR.

Original: 5. T HEORETICAL J USTIFICATION
Revised: 5. T HEORETICAL J USTIFICATION

Original: We now give theoretical guarantees of our proposed approximation, i.e., bound the gap between (cid:98) L AR ( f ) and L AR ( f ) .
Revised: We now give theoretical guarantees of our proposed approximation, i.e., bound the gap between (cid:98) L AR ( f ) and L AR ( f ) .

Original: For each representation f ∈ F , where F is some hypothesis space, define z f ( x ) :
Revised: For each representation f ∈ F , where F is some hypothesis space, define z f ( x ) :

Original: = sup
Revised: = sup

Original: A 1 ,A 2 ∈A ∥
Revised: A 1 ,A 2 ∈A ∥

Original: f ( A 1 x ) − f ( A 2 x ) ∥ 2 and Z F := { z f : f ∈ F} .
Revised: f ( A 1 x ) − f ( A 2 x ) ∥ 2 and Z F := { z f : f ∈ F} .

Original: Then we use the
Revised: Then we use the

Original: Rademacher complexity of Z to define the alignment complexity of F as
Revised: Rademacher complexity of Z to define the alignment complexity of F as

Original: [Equation]
Revised: [Equation]

Original: This quantity is a characteristic of the complexity of F in terms of its alignment .
Revised: This quantity is a characteristic of the complexity of F in terms of its alignment .

Original: If most of the functions in F have good alignment under A , then (cid:101) R n ( F ) is small even though R n ( F ) could be large.
Revised: If most of the functions in F have good alignment under A , then (cid:101) R n ( F ) is small even though R n ( F ) could be large.

Original: Using this definition, we have the following result.
Revised: Using this definition, we have the following result.

Original: Theorem 5.2.
Revised: Theorem 5.2.

Original: Suppose that Θ =
Revised: Suppose that Θ =

Original: B
Revised: B

Original: d A (1) , i.e., the d A -dimensional unit ball.
Revised: d A (1) , i.e., the d A -dimensional unit ball.

Original: Suppose that all A in A are L A -Lipschitz continuous.
Revised: Suppose that all A in A are L A -Lipschitz continuous.

Original: Let π be a distribution on Θ with density p π such that inf θ ∈ Θ p π ( θ ) ≥ c π .
Revised: Let π be a distribution on Θ with density p π such that inf θ ∈ Θ p π ( θ ) ≥ c π .

Original: Then with probability at least 1 − 2 ε , for every f ∈ F ,
Revised: Then with probability at least 1 − 2 ε , for every f ∈ F ,

Original: [Equation]
Revised: [Equation]

Revised: where n is the training sample size and m is the number of views.

Original: The bound in Theorem D.1 consists of two terms O ( n − 1 / 2 ) and O ( m − 1 /d A ) .
Revised: The bound in Theorem 5.2 consists of two terms O ( n − 1 / 2 ) and O ( m − 1 /d A ) .

Original: The first one is due to finite samples, and the second is caused by the finite maximization approximation of the supremum.
Revised: The first one is due to finite sample issue, and the second is caused by the finite maximization approximation of the supremum.

Revised: 1 /c π acts like the ”volume” of the transformation set and is of constant order.

Original: As m increases, the second term decreases and the bound becomes tighter.
Revised: As m increases, the second term decreases and the bound becomes tighter.

Original: Note that we need to assume that the augmentation distribution π behaves uniformly, i.e., inf θ ∈ Θ p π ( θ ) ≥ c for some positive constant c , otherwise we can not ensure that the finite maximization is an acceptable approximation.
Revised: Note that we need to assume that the augmentation distribution π behaves uniformly, i.e., inf θ ∈ Θ p π ( θ ) ≥ c for some positive constant c , otherwise we can not ensure that the finite maximization is an acceptable approximation.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We conduct experiments on several OOD settings to show the performance of our proposed methods.
Revised: We conduct experiments on several distribution shift settings to show the performance of our proposed methods.

Original: To show the adaptability of our approach, we apply ArCL on SimCLR (Chen et al., 2020) and MoCo V2 (He et al., 2020).
Revised: To show the adaptability of our approach, we apply ArCL on SimCLR (Chen et al., 2020) and MoCo V2 (He et al., 2020) respectively.

Original: Following the setting of SimCLR and MoCo V2, we add a projection head after the backbone in the pretraining stage and remove it in the downstream task.
Revised: Following the setting of SimCLR and MoCo V2, we add a projection head after the backbone in the pre-training stage and remove it in the downstream task.

Original: For the linear evaluation, we follow the protocol that the encoder is fixed and a linear classifier on the top of the encoder is trained using test set.
Revised: For the linear evaluation, we follow the protocol that the encoder is fixed and a linear classifier on the top of the encoder is trained using test set.

Original: For the finetuning, we also add a linear classifier on the top of the encoder but we do not fix the encoder.
Revised: For the full fine-tuning, we also add a linear classifier on the top of the encoder but we do not fix the encoder.

Original: 6. E XPERIMENTS ON M ODIFIED CIFAR10 AND CIFAR 
Revised: 6. E XPERIMENTS ON M ODIFIED CIFAR10 AND CIFAR

Original: Setup.
Revised: Setup.

Original: Our representations are trained on CIFAR-10 (Krizhevsky, 2009) using SimCLR modified with our proposed ArCL.
Revised: Our representations are trained on CIFAR-10 (Krizhevsky, 2009) using SimCLR modified with our proposed ArCL.

Original: Different batch sizes ( 256 , 512 ) and number of views ( m = 4 , 6 , 8 ) are[Table] [Table] considered.
Revised: Different batch sizes ( 256 , 512 ) and number of views ( m = 4 , 6 , 8 ) are considered.

Original: We use ResNet-18 as the encoder and a 2-layer MLP as the projection head.
Revised: We use ResNet-18 as the encoder and a 2-layer MLP as the projection head.

Original: We train the representation with 500 epochs.
Revised: We train the representation with 500 epochs.

Original: The temperature is 0 .
Revised: The temperature is 0 .

Original: Warm-up is not used.
Revised: Warm-up is not used.

Original: We conduct linear  evaluation with 100 epochs on different downstream datasets to test the OOD performance of the learned representations.
Revised: We conduct linear [Table] [Table] evaluation with 100 epochs on different downstream datasets to test the OOD performance of the learned representations.

Original: Datasets.
Revised: Datasets.

Original: We use CIFAR10 to train the representation.
Revised: We use CIFAR10 to train the representation.

Original: The data augmentation used follows the setting in SimCLR, i.e., a composition of RandomCrop, RandomHorizontalFlip, ColorJitter and
Revised: The data augmentation used follows the setting in SimCLR, i.e., a composition of RandomCrop, RandomHorizontalFlip, ColorJitter and

Original: Random Grayscale.
Revised: Random Grayscale.

Original: To evaluate the OOD generalization ability of the learned representation, we use different downstream datasets to train the linear classifier and test its accuracy.
Revised: To evaluate the transferability of the learned representation, we use different downstream datasets to train the linear classifier and test its accuracy.

Original: Downstream datasets are generated by modifying the original CIFAR10 and CIFAR100 datasets, i.e., applying different data augmentations to them.
Revised: Downstream datasets are generated by modifying the original CIFAR10 and CIFAR100 datasets, i.e., applying different data augmentations to them.

Original: We choose 5 different augmentations as shown in Table 1.
Revised: We choose 5 different augmentations as shown in Table 1.

Original: The results are shown in Table 2.
Revised: The results are shown in Table 2.

Original: Results.
Revised: Results.

Original: In all settings, our proposed approach improves the OOD performance of SimCLR significantly.
Revised: In all settings, our proposed approach improves the transferability of SimCLR significantly.

Original: As the number of views m increases, the accuracy also increases, which matches our theoretical results.
Revised: As the number of views m increases, the accuracy also increases, which matches our theoretical results.

Original: Besides,  performance improvements tend to saturate, which indicates that m needs not  be very large.
Revised: Besides, the performance improvement brought by increasing m tends to saturate, which suggest that m does not need to be very large.

Original: 6. E XPERIMENTS ON I MAGE N ET T RANSFERRING TO S MALL D ATASETS
Revised: 6. E XPERIMENTS ON I MAGE N ET T RANSFERRING TO S MALL D ATASETS

Original: Setup.
Revised: Setup.

Original: In this part, we will first train an encoder on ImageNet(Deng et al., 2009) using MoCo v modified with our proposed ArCL.
Revised: In this part, we will first train an encoder on ImageNet(Deng et al., 2009) using MoCo v modified with our proposed ArCL.

Original: We select the number of ArCL views from { 2 , 3 , 4 } .
Revised: We select the number of ArCL views from { 2 , 3 , 4 } .

Original: Notice that since there exists asymmetric network architecture in MoCo, there is difference between MoCo and
Revised: Notice that since there exists asymmetric network architecture in MoCo, there is difference between MoCo and

Original: MoCo + ArCL(views=2).
Revised: MoCo + ArCL(views=2).

Original: We use ResNet50 as the backbone and follow all the other architecture settings of MoCo.
Revised: We use ResNet50 as the backbone and follow all the other architecture settings of MoCo.

Original: The memory bank size K is set to 65536 and the batch size is set to 256.
Revised: The memory bank size K is set to 65536 and the batch size is set to 256.

Original: For the training epochs, we propose three schemes: 1) 200 epochs training from scratch, 2) 50 epochs training from the model which is pretrained by MoCo for 800 epochs, 3) 100 epochs training from the model which is pretrained by MoCo for 800 epochs.
Revised: For the training epochs, we propose three schemes: 1) 200 epochs training from scratch, 2) 50 epochs training from the model which is pretrained by MoCo for 800 epochs, 3) 100 epochs training from the model which is pretrained by MoCo for 800 epochs.

Original: For each setting, we search the best start learning rate in { 0 . 03 , 0 . 0075 , 0 . 0015 } and the best temperature in { 0 . 20 , 0 . 15 , 0 .
Revised: For each setting, we search the best start learning rate in { 0 . 03 , 0 . 0075 , 0 . 0015 } and the best temperature in { 0 . 20 , 0 . 15 , 0 .

Original: We use the SGD optimizer and apply the cosine learning rate decay.
Revised: We use the SGD optimizer and apply the cosine learning rate decay.

Original: After the model has been pretrained by contrastive learning, we test its OOD generalization ability with linear evaluation and full finetuning on varieties of small datasets, following the settings in (Ericsson et al., 2021).
Revised: After the model has been pre-trained by contrastive learning, we test its transferability with linear evaluation and full fine-tuning on various downstream small datasets, following the settings in (Ericsson et al., 2021).

Original: For the target datasets, we adopt FGVC Aircraft (Maji et al., 2013), Caltech- 
Revised: For the target datasets, we adopt FGVC Aircraft (Maji et al., 2013), Caltech- (Fei-Fei et al., 2004), Stanford Cars (Krause et al., 2013), CIFAR10, CIFAR100, DTD (Cimpoi et al., 2014), Oxford 102 Flowers (Nilsback & Zisserman, 2008), Food-101 (Bossard et al., 2014)

Original: [Table]
Revised: [Table]

Original: 101 (Fei-Fei et al., 2004), Stanford Cars (Krause et al., 2013), CIFAR10, CIFAR100, DTD (Cimpoi et al., 2014), Oxford 102 Flowers (Nilsback & Zisserman, 2008), Food-101 (Bossard et al., 2014) and Oxford-IIIT Pets (Parkhi et al., 2012).
Revised: and Oxford-IIIT Pets (Parkhi et al., 2012).

Original: For linear evaluation, multinomial logistic regression is fit on the extracted features.
Revised: For linear evaluation, multinomial logistic regression is fit on the extracted features.

Original: When finetuning, the  models are trained for 5,000 steps using SGD with Nesterov momentum.
Revised: For full fine-tuning, the full models are trained for 5,000 steps using SGD with Nesterov momentum.

Original: We calculate the average results of the nine downstream datasets.
Revised: We calculate the average results of the nine downstream datasets.

Original: Results.
Revised: Results.

Original: We compare our methods with the MoCo baseline.
Revised: We compare our methods with the MoCo baseline.

Original: There are about an average rise of 2% in the linear evaluation setting and about an average rise of 3% in the finetuning setting.
Revised: There are about an average rise of 2% in the linear evaluation setting and about an average rise of 3% in the finetuning setting.

Original: In particular, we gain a huge improvement on CIFAR100 finetuning, where our proposed methods can outperform the baseline for up to 10%.
Revised: In particular, we gain a huge improvement on CIFAR100 finetuning, where our proposed methods can outperform the baseline for up to 10%.

Original: It is also worthy to mention that our methods can reach the supervised results on nearly all the datasets.
Revised: It is also worthy to mention that our methods can reach the supervised results on nearly all the datasets.

Original: These facts can all indicate the effectiveness of our proposed methods.
Revised: These facts can all indicate the effectiveness of our proposed methods.

Original: Besides, the results also show that, under the same epoch setting, the average performance grows as the number of the views increase, which fits our theoretical analysis.
Revised: Besides, the results also show that, under the same epoch setting, the average performance grows as the number of the views increase, which fits our theoretical analysis.

Original: One can also notice that there is only a small improvement in the 800+100 epoch setting compared to the 800+50 epoch setting.
Revised: One can also notice that there is only a small improvement in the 800+100 epoch setting compared to the 800+50 epoch setting.

Original: This implies that our methods can also converge in a rather fast speed, which can also be effective within a few epochs of training.
Revised: This implies that our methods can also converge in a rather fast speed, which can also be effective within a few epochs of training.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: In this paper, we theoretically study the OOD generalization ability of contrastive learning.
Revised: In this paper, we theoretically study the transferability of contrastive learning under distribution shifts.

Original: Our results suggest that the data augmentation plays an important role in its OOD performance.
Revised: Our results suggest that the data augmentation plays an important role in its transferability.

Original: We also find a limitation of contrastive learning: it cannot learn domain invariant representations.
Revised: We also find a limitation of contrastive learning: it cannot learn domain-invariant representations.

Original: Based on these theoretical results, we propose the Augmentation-robust Contrastive Learning (ArCL), a module that can be applied to various popular contrastive learning methods.
Revised: Based on these theoretical results, we propose the Augmentation-robust Contrastive Learning (ArCL), a module that can be applied to various popular contrastive learning methods.

Original: It forces the representation to align positive samples uniformly, in order to improve its OOD performance.
Revised: It forces the representation to align positive samples uniformly, in order to improve its distribution shift transferability.

Original: Our experiments on CIFAR10 and ImageNet verify our theoretical results and the superiority of our method.
Revised: Our experiments on CIFAR10 and ImageNet verify our theoretical results and the superiority of our method.

Original: B P ROOFS IN S ECTION
Revised: B P ROOFS IN S ECTION

Original: Proof of Lemma 3.1.
Revised: Proof of Lemma 3.1.

Original: For simplicity we omit the notations D and ϑ in this proof, and use x 1 to denote the positive sample, i.e. the augmented data, use p k to denote D
Revised: For simplicity we omit the notations D and ϑ in this proof, and use x 1 to denote the positive sample, i.e. the augmented data, use p k to denote D

Original: ϑ ( C k ) and use µ k to denote µ k ( f ; D ϑ ) .
Revised: ϑ ( C k ) and use µ k to denote µ k ( f ; D ϑ ) .

Original: From Theorem 2 and Lemma B.1 in (Huang et al., 2022), we have
Revised: From Theorem 2 and Lemma B.1 in (Huang et al., 2022), we have

Original: [Equation]
Revised: [Equation]

Original: for some constant c
Revised: for some constant c

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: Note that τ is decreasing with σ and increasing with δ .
Revised: Note that τ is decreasing with σ and increasing with δ .

Original: Then we have
Revised: Then we have

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: for all linear layer h ∈ R K × d 1 .
Revised: for all linear layer h ∈ R K × d 1 .

Original: Therefore, we obtain
Revised: Therefore, we obtain

Original: [Equation]
Revised: [Equation]

Original: Proof of Theorem 3.2.
Revised: Proof of Theorem 3.2.

Original: By the triangle inequality and equation
Revised: By the triangle inequality and equation

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

Original: Define U = ( µ 1 , . . .
Revised: Define U = ( µ 1 , . . .

Original: , µ K ) ∈ R d 1 × K and let h
Revised: , µ K ) ∈ R d 1 × K and let h

Original: =
Revised: =

Original: U + be the Moore–Penrose inverse of U such that h 0 U
Revised: U + be the Moore–Penrose inverse of U such that h 0 U

Original: = I K .
Revised: = I K .

Original: Then we have
Revised: Then we have

Original: [Equation]
Revised: [Equation]

Original: For every x = ( x 1 , . . .
Revised: For every x = ( x 1 , . . .

Original: x d 1 ) ∈ R d 1 with ∥ x ∥ 2 =
Revised: x d 1 ) ∈ R d 1 with ∥ x ∥ 2 =

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: Therefore, we have ∥ h 0 ∥ ≤ η reg ( f ; D , ϑ ) .
Revised: Therefore, we have ∥ h 0 ∥ ≤ η reg ( f ; D , ϑ ) .

Original: Then, by Lemma 3.1, we complete the proof by
Revised: Then, by Lemma 3.1, we complete the proof by

Original: R ( f ) =
Revised: R ( f ) =

Original: min h R ( h ◦ f ) ≤ R ( h 0 ◦ f ) ≤ cγ reg ( f ; D , ϑ ) L 14 pos ( f ; D , ϑ ) + γ reg ( f ; D , ϑ ) τ ( σ, δ ) .
Revised: min h R ( h ◦ f ) ≤ R ( h 0 ◦ f ) ≤ cγ reg ( f ; D , ϑ ) L 14 pos ( f ; D , ϑ ) + γ reg ( f ; D , ϑ ) τ ( σ, δ ) .

Original: C P ROOF IN S ECTION
Revised: C P ROOF IN S ECTION

Original: Proof of Proposition 4.1.
Revised: Proof of Proposition 4.1.

Original: For any ε > 0 , let t = √ ε/ 2 and f ( x 1 , x 2 ) =
Revised: For any ε > 0 , let t = √ ε/ 2 and f ( x 1 , x 2 ) =

Original: x 1 + tx 2 .
Revised: x 1 + tx 2 .

Original: Then, the alignment loss of f satisfies
Revised: Then, the alignment loss of f satisfies

Original: [Equation]
Revised: [Equation]

Original: Let c = 0 and c ′ = 1 /t .
Revised: Let c = 0 and c ′ = 1 /t .

Original: Then obviously
Revised: Then obviously

Original: R ( f ; D c ) = 0 , but
Revised: R ( f ; D c ) = 0 , but

Original: R ( f ; D c ′ ) =
Revised: R ( f ; D c ′ ) =

Original: P ( X 1 < 0 , X 1 + X 2 ≥ 0) + P ( X 1 ≥ 0 , X 1 + X 2 ≤ 0)
Revised: P ( X 1 < 0 , X 1 + X 2 ≥ 0) + P ( X 1 ≥ 0 , X 1 + X 2 ≤ 0)

Original: = 1 4 .
Revised: = 1 4 .

Original: D P ROOFS IN S ECTION
Revised: D P ROOFS IN S ECTION

Original: Proof of Proposition 5.1.
Revised: Proof of Theorem 5.1.

Original: For any f and h with ∥ h ∥ ≤ c h , we have
Revised: For any f and h with ∥ h ∥ ≤ c h , we have

Original: [Equation]
Revised: [Equation]

Original: for some constant c which depends on c h .
Revised: for some constant c which depends on c h .

Original: proof of Theorem D.1.
Revised: proof of Theorem 5.2.

Original: We first fix the sample ( x 1 , . . .
Revised: We first fix the sample ( x 1 , . . .

Original: , x n ) , and consider the randomness of
Revised: , x n ) , and consider the randomness of

Original: A m ( x i ) .
Revised: A m ( x i ) .

Original: Define A ∗ 1 ( x i ) and A ∗ 2 ( x i ) as
Revised: Define A ∗ 1 ( x i ) and A ∗ 2 ( x i ) as

Original: [Equation]
Revised: [Equation]

Original: Note that their order doesn’t matter.
Revised: Note that their order doesn’t matter.

Original: For every ε > 0 , we have
Revised: For every ε > 0 , we have

Original: [Equation]
Revised: [Equation]

Original: ≥ 
Revised: ≥ P (cid:40) max

Original: [Equation]
Revised: [Equation]

Original: max
Revised: (cid:26)

Original: [Equation]
Revised: [Equation]

Original: ≥ 
Revised: ≥ 1 − n

Original: [Equation]
Revised: [Equation]

Original: −
Revised: (cid:18) log(2 n

Original: [Equation]
Revised: [Equation]

Original:  ≥ 1 − 2 n 
Revised: ε ) c ϑ m (cid:19) 1 d (cid:41) ≥ 1 − 2 n (cid:18)

Original: [Equation]
Revised: [Equation]

Original: ≥ 1 − ε.
Revised: ≥ 1 − ε.

Original: Therefore, with probability at least 1 − ε
Revised: Therefore, with probability at least 1 − ε

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: Besides [Equation] for every f we have
Revised: Besides, for every f we have

Original: [Equation]
Revised: [Equation]

Original: with probability at least 1 − ε .
Revised: with probability at least 1 − ε .

Original: Then we finally obtain that with probability at least 1 − 2 ε , for every f ∈ F ,
Revised: Then we finally obtain that with probability at least 1 − 2 ε , for every f ∈ F ,

Original: [Equation]
Revised: [Equation]

Original: n
Revised: f

Original: [Equation]
Revised: [Equation]

Original: We now give theoretical guarantees of our proposed approximation, i.e., bound the gap between (cid:98) L AR ( f ) and L AR ( f ) . For each representation f ∈ F, where F is some hypothesis space, define z f ( x ) :
Revised: F

Original: = sup
Revised: [Equation]

Original: A 1 ,A 2 ∈A ∥ f ( A 1 x ) − f ( A 2 x ) ∥ 2 and Z F := { z f : f ∈ F} . Then we use the Rademacher complexity of Z to define the alignment complexity of F as
Revised: ε

Original: [Equation]
Revised: [Equation]

Original: This quantity is a characteristic of the complexity of F in terms of its alignment . If most of the functions in F have good alignment under A , then (cid:101) R n ( F ) is small even though R n( F ) could be large. Using this definition, we have the following result. Theorem D.1.
Revised: n

Original: Suppose that Θ = B
Revised: [Equation]

Original: d A (1) , i.e., the d A -dimensional unit ball.
Revised: [Equation] ε [Equation] E A DDITIONAL E XPERIMENTS

Revised: In order to further verify the generality of our method, we conduct more experiments on different settings and compare our objective with different loss.

Revised: [Table]

Revised: E. M O C O ON CIFAR

Revised: Setup.

Revised: In this part, we conduct experiments with MoCo v2 pretrained on CIFAR10.

Revised: We follow the setup of MoCo v2.

Revised: We use ResNet-18 as the encoder and train the representation with 400 epochs.

Revised: The temperature is set to 0.2 and the initial learning rate is set to 0.15.

Revised: The memory bank size is set to 4096.

Revised: The SGD optimizer and a cosine schedule for learning rate are used.

Revised: Warm-up is not used.

Original: Suppose that all
Revised: We conduct linear evaluation with 100 epochs on modified CIFAR10 and CIFAR100 which are used in Section 6.1. Results. Results can be seen in Table 4. We can see that the original results still maintain in the

Revised: MoCo method.

Revised: Our proposed approach improves the transferability of MoCo.

Original:  A in A are L A -Lipschitz continuous.
Revised: As the number of views grows, the accuracy also increases. E. C OMPARISON WITH A VERAGE A LIGNMENT L OSS To make the comparison between ArCL and CL fairer, we add new experiments for original CL. For each sample, we also construct m views and use the expectation of the similarity between positive pairs, namely average alignment loss (AAL), as the learning objective.

Revised: The settings and results are shown in the following part.

Revised: Loss Function.

Revised: For image x i , the normalized features of its m views from the online branch are z

Revised: i 1 , z i 2 , . . .

Revised: , z im and the normalized features from the target branch are z

Revised: i 1 , z ′ i 2 , . . .

Original: Let π bea distribution on Θ with density p π such that inf
Revised: , z ′ im . Our ArCL alignment loss is [Equation] And the average alignment loss should be

Original: [Equation]
Revised: [Equation]

Revised: The uniformity loss keeps the same and the total loss is the sum of the alignment loss and the uniformity loss.

Original: Θ p π ( θ ) ≥ c π .
Revised: Hyperparameters and Datasets. For SimCLR, we train the models on CIFAR10 and evaluate them on modified CIFAR10 and CIFAR100 just as our original paper does. We choose the augmentation views to be 4 and the batch size to be 512. For MoCo, We conduct experiments on ImageNet, with the 800-epochs-pretrained model as the initializing model.

Revised: We train the model for 50 epochs and use the same setting as in Section 6.2.

Revised: Results.

Original: Then with probability at least 1 − 2 ε , for every f ∈ F ,
Revised: Results for SimCLR experiments can be seen in Table 5. We compare the linear evaluation results on modified CIFAR10 and CIFAR100. Results for MoCo experiments can be seen in ?? . We compare both the linear evaluation and finetuning results on small datasets.

Revised: As the experiments show, when contrastive learning methods with average alignment loss use the same amount of data as ArCL, although they can have higher accuracy compared to vanilla methods (those which only use two views) on ID set, they still have a similar performance with vanilla methods on OOD set, which is much worse than ArCL.

Revised: This further verifies the superiority of our method.

Original: [Equation]
Revised: [Table] [Table]

Original: The bound in Theorem D.1 consists of two terms O ( n − 1 / 2 ) and O ( m − 1 /d A ) .
Revised: MNIST-CIFAR is a synthetic dataset proposed by (Shah et al.). It consists of 3 × 64 × 32 synthetic images, each of which is a vertical concatenation of a 3 × 32 × 32 MNIST image from class 0 or class 1, and a 3 × 32 × 32 CIFAR10 image from class 0 or class 1. We follow the distribution settings proposed by (Shi et al., 2022):

Revised: • ID train: 1.0 correlation between MNIST and CIFAR10 labels.

Revised: Contains two classes: class 0 with MNIST “0” and CIFAR10 “automobile”, and class 1 with MNIST “1” and CIFAR “plane”.

Revised: • OOD train: no correlation between MNIST and CIFAR10 labels, images from the two classes of MNIST and CIFAR10 are randomly paired.

Original: The first one is due to finite samples, and the second is caused by the finite maximization approximation of  the supremum.
Revised: We choose the label of CIFAR10 to be the label of the concatenated image. • OOD test: generated similarly to the OOD train set using the test set of MNIST and CIFAR10. In this experiment, we train the model on the ID train set, conduct the linear evaluation on the OOD train set, and calculate the accuracy on the OOD test set. We adopt SimCLR as the CL framework and use the 4-layer CNN as the backbone just as Shi et al.

Revised: (2022) does.

Revised: We fix the base feature size of CNN C = 32 and the latent dimension L = 128 .

Revised: Batch size is set to 128, the optimizer is set to

Original: As m increases, the second term decreases and the bound becomes tighter.
Revised: SGD and the lr. scheduler is set to warmup cosine. The learning rate and weight decay are searched in the range given in Table 2 of Shi et al. Since the training process is easy to converge, we set the training epoch and the linear evaluation epoch to 5. We also compare the average alignment loss (AAL for short) mentioned in the Appendix E.2.

Revised: Results.

Revised: We can see that by using ArCL, SimCLR enjoys a rising performance.

Original: Note that we need to assume that the augmentation distribution π behaves uniformly, i.e., inf θ ∈ Θ p π ( θ ) ≥ c for some positive constant c , otherwise we can not ensure that the finite maximization is an acceptable approximation.
Revised: The accuracy also grows as the number of views grows, which fits our theory. We also notice that the accuracy of SimCLR with AAL does not vary too much as the number of views grows, which indicates that the key point is the usage of multi-view ( adopting the minimum instead of the average). [Table]