The Trade-off between Label Efficiency and Universality of Representations from Contrastive Learning
Abstract
The pre-train representation learning paradigm is a recent popular approach to address distribution shift and limitations in training data.
This approach first pre-trains a representation function using large unlabeled datasets from multiple tasks by self-supervised (e.g., contrastive) learning, and then learns a simple classifier on the representation using small labeled datasets from the downstream target tasks.
The representation should have two key properties: label efficiency (i.e., ability to learn an accurate classifier with a small amount of labeled data) and universality (i.e., usefulness across a wide range of downstream tasks).
In this paper, we focus on contrastive learning and systematically study the trade-off between label efficiency and universality both theoretically and empirically.
We empirically show that this trade-off exists in different models and datasets.
Theoretically, we propose a data model with a hidden representation and provide analysis in a simplified linear setting.
Our analysis shows that compared to pre-training on the target task, pretraining on diverse tasks leads to a larger sample complexity for learning the optimal classifier, and thus has worse prediction performance.
1. Introduction
The pre-train representation learning paradigm is a recent successful approach to utilize large-scale unlabeled data to address the challenges of labeled data scarcity and distribution shift.
Different from the traditional supervised learning approach using a large set of labeled data, representation
Correspondence to: Zhenmei Shi < zhmeishi@cs.wisc.edu > .
First Workshop of Pre-training: Perspectives, Pitfalls, and Paths Forward at ICML 2022 , Baltimore, Maryland, USA, PMLR 162, learning first pre-trains a representation function using largescale diverse unlabeled datasets by self-supervised learning (e.g., contrastive learning), and then learns a predictor on the representation using a small labeled dataset for a downstream target task.
The pre-trained model is sometimes referred to as a foundation model (Bommasani et al., 2021), and has achieved good performance in many applications, e.g., BERT (Devlin et al., 2019), GPT-3 (Brown et al., 2020) and CLIP (Radford et al., 2021).
While the foundation models exist for different applications and with different learning methods, the following two properties are key to their success: (1) label efficiency : with the pre-trained representation, only a small amount of labeled data is needed to build accurate predictors for downstream target tasks; (2) universality : the pre-trained representation can be used across various downstream tasks.
This work points out and studies a potential trade-off between label efficiency and universality, though ideally one would like to have these two key properties simultaneously.
We focus on contrastive learning which is one exemplary pre-training method, while we believe similar observations and insights apply to other methods.
Empirically, we observe that such a trade-off indeed exists when the representation is pre-trained via contrastive learning on a large and diverse unlabeled dataset consisting of different tasks (or distributions).
Since pre-training on diverse tasks is widely used in practice, such a trade-off deserves to be fully understood.
More precisely, we perform controlled experiments comparing two cases: (1) specific representation pre-trained on a specific unlabeled dataset similar to that of the target task; (2) universal representation pre-trained on the union of diverse unlabeled datasets.
The diverse data include the specific unlabeled dataset and some other datasets that can be quite different from that of the target task, which mimics the practical scenario that the foundation model is pre-trained on diverse data to be widely applicable for various downstream tasks.
We observe that when the specific dataset is large, universal representation leads to worse prediction on the target, i.e., adding the diverse datasets in pre-training can harm the label efficiency though enhancing the universality.
We also observe that if the specific dataset is small, then adding diverse datasets can help.
These suggest the following explanation: diverse unlabeled datasets have both positive and negative impacts for prediction on a specific target task.
On one hand, they share some useful semantic features with the target task, that can help learn these features in the representation and improve the prediction.
On the other hand, they can have many other features not so useful for the target, and encoding such features into the representation essentially down-weights the useful ones and thus hurts prediction.
When the specific dataset is large, the positive impact is marginal and overwhelmed by the negative impact, leading to the trade-off.
Theoretically, we provide analysis formalizing the intuition that the diverse unlabeled datasets have both positive and negative impacts by helping the representation learn various semantic features.
We propose a hidden representation model for the data, which first generates a hidden representation containing various features, and then uses it to generate the label and the input.
We then provide analysis in a simplified setting with linear data and representation functions.
We show that contrastive learning can learn hidden features invariant to the transformations, and thus allows an accurate predictor if the label depends only on such invariant features.
On the other hand, when pre-trained on diverse unlabeled data, it encodes all invariant features from different tasks and essentially emphasizes those common features but down-weights those specific to the target task.
This then leads to a larger sample complexity for prediction on the target task and thus a worse generalization performance.
Related Work.
We only discuss the most related ones here and include more in Appendix A. Cole et al.
point out the “diversity-difficulty trade-off”: pre-training on pooled datasets leads to worse performance on the in-domain task compared to pre-training on the in-domain dataset only.
Similarly, Bommasani et al.
call for further research on the issue of specialization vs. diversity in foundation model training.
Neither of them provide a systematic investigation or theoretical analysis on the trade-off.
Our work aims to provide empirical and theoretical analysis on how the pre-training data leads to the trade-off.
We also note that existing theoretical analysis (e.g., (Arora et al., 2019; HaoChen et al., 2021)) typically assumes that the pre-training data distribution is the same as the target distribution, while the difference between the two is critical for the trade-off focused in this work.
Thus, our work proposes a new analysis approach.
2. Theoretical Analysis
In this section, we will analyze the sample complexity of learning the predictor on top of the pre-trained representation.
Here, the representation is pre-trained on diverse data (modeled by the union of unlabeled data from several tasks) via contrastive learning, to obtain representations that are potentially useful for a wide range of diverse tasks; while the predictor is learned for a specific target task using labeled data from that task.
Existing analyses on contrastive learning are not applicable since they typically assume the same distribution for the pre-training and the prediction stages.
To model the intuition that the pre-training can learn semantic features from the unlabeled data such that a subset of them can be useful for prediction on even different data distributions, we propose a hidden representation model for the data, which first generates a hidden representation z and then uses z to generate the label y and the input x .
While we are not able to perform analysis for the most general case, we consider a simplified linear setting where the data model and the representation functions are linear.
In this simplified setting, our analysis shows that contrastive learning can learn hidden features invariant to the transformation, and thus allows an accurate predictor if the label depends only on such invariant features.
On the other hand, when pre-trained on a union of unlabeled data from different tasks, it encodes all invariant features from different tasks and essentially emphasizes those that are shared among the tasks, but down-weights those that are specific to a single task.
Compared to pre-training only on unlabeled data from the target task, this then leads to a larger sample complexity for prediction on the target task (formalized by a larger
Rademacher complexity of the predictor hypothesis class).
Equivalently, this gives a worse generalization performance.
Therefore, we formally show a trade-off between universality and the label efficiency.
Contrastive Learning.
Let X =
R d denote the input space,
Y the label space, Z = R d the hidden representation space,
Z = R k the output space of the learned representation function.
Let Φ denote the hypothesis class of representation functions ϕ : X (cid:55)→ Z , and F ϕ the hypothesis class of predictors on the representation ϕ .
Each task t is a data distribution D t over X × Y .
In pre-training, using transformations on unlabeled data from the tasks, we have a distribution over positive pairs ( x, x + ) and negative examples x − , where x, x − are two independent examples, while x + is obtained by applying some random transformations on x (e.g., cropping or color jitter for images).
The contrastive loss is
[Equation]
where ℓ ( · ) is the logistic loss ℓ ( z )
=
log(1 + exp( − z )) .
In practice multiple independent negative examples are used, and thus we consider the simplified contrastive loss:
[Equation]
to pre-train a representation ϕ .
Minimizing the contrastive loss maximizes the representation similarity between positive pairs x and x + , while minimizes the representation similarity between negative pairs x and x − .
Then we learn a predictor f on top of ϕ using labeled data from a specific target task D t :
[Equation]
where ℓ c is a prediction loss (e.g. cross-entropy).
Usually, f is a linear classifier, and for ϕ ∈ Φ :
[Equation]
for some norm bound B specified later.
Hidden Representation Data Model.
Suppose the data is generated by first sampling a hidden representation z ∈
Z from some distribution, and then generating the input x ∈ X and the label y ∈ Y from the hidden representation z .
We assume that the dimensions of z are partitioned into two subsets: spurious features U that are affected by the transformations, and invariant features R =
[ d ] \ U .
Intuitively, the transformation should be chosen such that useful semantic features will not be affected, and our goal is to recover the invariant features R .
Therefore, we assume that k
= | R | , and in each task t the label y will only depend on a subset of invariant features R t ⊆
R of size | R t | = r .
Formally, for D t , we assume z R t ∼ N (0 , I ) , z R \ R t = 0 , z U ∼ N (0 , I ) , and
y =
( u ∗ t )
⊤ z R t for some ground-truth parameter u ∗ t .
Then the positive pairs ( x, x + ) are generated as follows:
[Equation]
[Equation]
and x , x + , and x − are generated from different conditional distributions D z , D z + , and D z − respectively.
A Simplified Setting.
The general case of data distributions and pre-training is challenging to analyze.
Here we consider a simplified setting with linear models and binary classification for the downstream task with Y = {− 1 , +1 } .
Formally, (1) x is generated from z linearly: x = Mz where M ∈ R d × d is an orthonormal dictionary.
The representations are linear functions with weight matrices of bounded Spectral and Frobenius norms:
[Equation]
Under this simplified setting, we show that pre-training on diverse tasks needs a larger Rademacher complexity to learn the optimal classifier compared to pre-training on the target task.
We first assume that the pre-training is on an equal mixture of two tasks D 1 and D 2 , with R 1 = { 1 , 2 , . . .
, r } and R 2 = { 1 , . . .
, s }∪{ r +1 , r +2 , . .
, 2 r − s } where s = | R 1 ∩ R 2 | ; the prediction is on the first task D 1 .
We include the generalized analysis for multiple tasks in Appendix B.
We first analyze the optimal representation.
Proposition 2.1.
There exist min(1 / 2 , s/r ) ≤ α ≤ 1 , β = min (cid:16) 1 , r − αs 2( r − s ) (cid:17) ∈ [1 / 2 , 1] such that ϕ ∗ ( x ) =
W ∗ x is an optimal representation for the contrastive loss (1) with any W ∗ of the form: W ∗ =
[ OA ∗ , 0 ] M − 1 where
O ∈ R k × k is any orthonormal matrix, A ∗ is diagonal with
A ∗ jj = √ α if j ∈ R 1 ∩ R 2 , and A ∗ jj = √ β otherwise, and the matrix of zeros 0 has size k × ( d − k ) .
That is, ϕ ∗ is a rotation of the weighted features, where the common features in R 1 ∩ R 2 are weighted by √ α and those task-specific features are weighted by √ β .
Given the representation ϕ ∗ , we would like to ensure there exists a predictor in F ϕ ∗ matching the ground-truth label.
Note that f ( ϕ ) =
u ⊤ ϕ with u = O 1: k, 1: r ( A ∗ 1: r, 1: r ) − 1 u [Equation]
1 + O 1: k,r +1: k v for any v ∈ R k − r satisfies f ( ϕ ∗ ( x )) =
y =
( u ∗ 1 )
⊤ z R 1 , and u ∗
=
O 1: k, 1: r ( A ∗ 1: r, 1: r ) − 1 u ∗ 1 is the leastnorm optimal solution.
So the predictor class should be
[Equation]
Proposition 2.2.
Let v 1 = (cid:80) {+sj =1 ( u ∗ 1 j ) 2 and v 
+}
= (cid:80) rj
= s +1 ( u ∗ 1 j ) 2 .
Then the Rademacher complexity of
F ϕ ∗ in Eqn.
(5) satisfies
[Equation]
where the estimate ˜ R m ( F ϕ ∗ ) is
[Equation]
So ignoring the small-order term, the Rademacher complexity is roughly ˜ R m ( F ϕ ∗ ) .
We now discuss the implication (details in Appendix B.3).
When we pre-train on one task (equivalent to r = s ), the complexity is roughly (cid:112) rm ∥ u ∗ 1 ∥ 2 .
Consider pre-training on two tasks with r = 2 s and v 1 = v 2 =
∥
u ∗ 1 ∥ 22 2 .
We can show that the optimal is α = 1 , β = 12 , and the complexity is roughly (cid:113)
[Equation]
Therefore, the complexity of pre-training on two tasks is larger than just pre-training on the target task.
This quantifies the trade-off between universality and label efficiency.
3. Experiments
We conduct experiments via contrastive learning to answer the following questions.
( Q1 )
Does the trade-off between universality and label efficiency exist?
( Q2 )
What conditions lead to the trade-off, and ( Q3 ) how can we use the pretrain-finetune learning paradigm effectively?
We summarize the answers as follows: ( A1 )
The trade-off widely exists in different models and datasets.
( A2 ) When the task-relevant dataset is large enough, the task-irrelevant datasets will lead to the trade-off.
( A3 )
Given knowledge of
[Figure]
the downstream task, pre-training on a task-related dataset is better than pre-training on an unrelated dataset.
3.1. Experimental Setup
Model.
We evaluate two popular contrastive learning frameworks, MoCo v2 (He et al., 2020) and SimSiam (Chen &
He, 2021).
MoCo v2 can be regarded as SimCLR (Chen et al., 2020) equipped with a memory bank, while SimSiam can be regarded as a modification from BYOL (Grill et al., 2020) similar to Barlow Twins (Zbontar et al., 2021), which does not need negative pairs.
Dataset.
We consider two sets of data.
In the first set, our downstream task is CIFAR-10, and the pre-training datasets may include CIFAR-10, CINIC-10, SVHN, GTSRB, and
ImageNet32.
CINIC-10 has classes identical to CIFARand is the most target-relevant, while the others are different.
In the second set, our downstream task is MNIST, and the pre-training datasets may include EMNIST-Digits&Letters,
Fashion-MNIST, GTSRB, and ImageNet32.
Here, EMNISTDigits&Letters is the most target-relevant.
Evaluation & Methods.
We pre-train a ResNet18 network (He et al., 2016) as a feature extractor using SGD for 800 epochs with a cosine learning-rate schedule and a base learning rate of 0 .
Then we fix the pre-trained feature extractor, and train a linear classifier called Linear
Probing (LP) on 1% , 5% , 10% , 20% , 100% of the labeled data from the downstream task.
For LP we use SGD for 100 epochs and a cosine learning rate schedule with a base learning rate of 5 . 0 .
We finally report the test accuracy on the downstream task.
3.2. Experimental Details
In Figs. 1(a) and (b), we report results for MoCo v2 and SimSiam (respectively) on CIFAR-10 as the downstream task.
The size and diversity of unlabeled data for pre-training is increased on the x-axis by incrementally adding datasets in the following order: CINIC-10, SVHN, GTSRB, and
ImageNet(500k).
Then, we do LP on CIFAR-10 using different proportions of labeled samples.
When the pretraining dataset is combined with more diverse data, the test accuracy for the specific downstream task decreases.
[Figure]
As more diverse unlabeled data are included, more labeled data from the target task is needed to achieve a comparablygood prediction accuracy.
This validates our hypothesis about the trade-off between universality and label efficiency.
In Figs.
1(c) and (d), we report results for MoCo v2 and
SimSiam (respectively) on MNIST as the downstream task.
The size and diversity of unlabeled data for pre-training is increased on the x-axis by incrementally adding datasets in the following order: EMNIST-Digits&Letters, FashonMNIST, GTSRB, ImageNet(500k).
This is followed by LP on MNIST with different proportion of labeled data.
We observe the same trend in test accuracy as in Figs. 1(a) and (b).
The handwritten dataset (EMNIST) is the most targetrelevant, and it helps pre-training features suitable for the handwritten recognition task on MNIST.
However, when we mix Fashion-MNIST, GTSRB, and ImageNet32 in the pre-training, the test accuracy on MNIST drops significantly.
This supports our claim that pre-training on target-relevant data will learn more target-relevant features and get a better performance on the target task, while introducing diverse pre-training data will allow learning diverse features but can down-weight those for a specific task.
The above analysis establishes that the trade-off widely exists in different models and datasets, answering Q1 .
3.3. Ablation Study
We report results from three ablation studies: (1) varying the class number of ImageNet32, (2) varying the percentage of target-relevant pre-training data, and (3) replacing CINIC10 with CIFAR-10 in the pre-training dataset.
Varying the Class Number of ImageNet32.
To further support A1 , we show that the trade-off between universality and label efficiency also exists under a fixed dataset size.
In
Fig.
2, we pre-train MoCo v2 and SimSiam on CIFAR + ImageNet(200k) and keep the same setting as Fig.
1 except that we vary the class number of ImageNet(200k).
In previous experiments, we randomly pick 500,000 images from ImageNet32 without considering labels.
Here, we fix the number of classes to 50, 100, 200, 500, 1000.
Then we randomly sample 200,000 images from the subset of classes.
The downstream task is CIFAR-10.
In Fig.
2, we observe that with a fixed pre-training datasets size, e.g., 250,000, when the data is more diverse, the pre-training will learn more irrelevant features, and the performance will drop on the downstream task.
This supports our analysis as well.
Due to the page limit, we provide details of (2) and (3) in
Appendix C.2.
We answer Q2 from (2) that the data from diverse tasks may have a positive effect when the data from similar (relevant) tasks is not sufficiently large.
Combining (3) with previous results, we answer Q3 .
If we choose a good task-relevant pre-training dataset, we can directly get similar performance as pre-training on the downstream task.
However, the performance will drop if we introduce task-irrelevant data in the pre-training dataset.
4. Discussion and Future Work
In this work, we have shown that the trade-off between label efficiency and universality of representations widely exists in contrastive learning.
For future work, there are many open questions we will continue to study.
(1) What features does the model learn from specific pre-training and diverse pre-training datasets?
(2) What properties do these features have?
(3) Can we solve the trade-off in a better way in order to gain both properties at the same time?
Acknowledgements
The work is partially supported by Air Force Grant FA955018-1-0166, the National Science Foundation (NSF) Grants 2008559-IIS, CCF-2046710, 2023239-DMS, CCF-FMitF1836978, SaTC-Frontiers-1804648, and CCF-1652140, and
ARO grant number W911NF-17-1-0405.
Jiefeng Chen and
Somesh Jha are partially supported by the DARPA-GARD problem under agreement number 885000.
A. Related Work
Pretrain-finetune Learning Paradigm.
The pretrain-finetune learning paradigm, where a model (or a representation) is pre-trained on a large dataset (e.g., ImageNet) and is then fine-tuned to various downstream tasks, has been widely used in practice (Devlin et al., 2019; Kolesnikov et al., 2020; Brown et al., 2020).
There are mainly two kinds of pre-training approaches: one is the supervised pre-training (Kolesnikov et al., 2020), where we pre-train representations on large labeled datasets; the other is the self-supervised pre-training (Newell & Deng, 2020), where we pre-train representations on large and diverse unlabeled datasets.
The self-supervised pre-training learning paradigm is sometimes referred to as the foundation models (Bommasani et al., 2021).
Recently, it has been demonstrated that self-supervised pre-training can learn effective representations that even outperform the representations learned by supervised pre-training when evaluating them on downstream tasks (Ericsson et al., 2021).
Also, some practical examples like BERT (Devlin et al., 2019), GPT-3 (Brown et al., 2020) and CLIP (Radford et al., 2021) have demonstrated the effectiveness of self-supervised pre-training in learning universal representations that can be used for a wide range of downstream tasks.
In our work, we study the properties of self-supervised pre-training due to its superiority over supervised pre-training.
Self-supervised Representation Learning.
Early self-supervised representation learning methods typically focus on solving hand-designed “pretext tasks” (Doersch et al., 2015; Gidaris et al., 2018; Noroozi & Favaro, 2016).
Recent works have explored contrastive learning-based approaches where the pretext task is to distinguish matching and non-matching pairs of augmented input images (van den Oord et al., 2018).
Common examples include SimCLR (Chen et al., 2020),
MoCo (He et al., 2020), SimSiam (Chen & He, 2021), BYOL (Grill et al., 2020) and Barlow Twins (Zbontar et al., 2021).
Analysis of Self-supervised Pre-training Paradigm.
Existing works have studied the effect of the pre-training datasets on the performance of the self-supervised pre-training (Arora et al., 2019; Tosh et al., 2021; Garg & Liang, 2020; HaoChen et al., 2021;
Tsai et al., 2020; Wen & Li, 2021; Saunshi et al., 2022; Wang & Isola, 2020;
Liu et al., 2021; Kotar et al., 2021;
Gansbeke et al., 2021; Yang et al., 2020).
However, existing analysis typically assumes that the pre-training data distribution is the same as the target distribution, while the difference between the two is the critical reason for the trade-off focused in this work.
Thus, our work proposes new analysis approaches.
Recently, Cole et al. have tried to identify conditions where self-supervised contrastive representation learning methods can produce “good” visual representations and point out the “diversity-difficulty trade-off” phenomenon, which is most relevant to our work.
However, they only empirically show the trade-off, but do not provide a systematic study and analysis to explain why it happens.
Bommasani et al.
call for further research on the issue of specialization vs. diversity in foundation model training data, but do not provide a thorough study as well.
Our work attempts to provide a better understanding of the trade-off between universality and label-efficiency.
B. Proof and More Analysis
We note that the contrastive loss can be written using logistic loss:
[Equation]
(7) where ℓ ( · ) is the logistic loss ℓ ( z )
=
log(1 + exp( − z )) .
This will be useful for the analysis.
B.1. Proof of Propositions 2.
Proof of Proposition 2.1.
For each D t ,
[Equation]
where the inequality comes from the convexity of ℓ ( z ) and Jensen’s inequality.
Then on the mixture, we have
[Equation]
Let WM =
[ A R , A U ] where A R ∈ R k × k , A U ∈ R k × ( d − k ) .
By rotational invariance of z R 1 ∩ R 2 , z R 1 \ R
[Equation]
and z R 2 \ R 1 , without loss of generality
[Equation]
we can assume A R = OA where A is a diagonal matrix with diagonal entries a jj ’s
and O is any orthonormal matrix.
Then
[Equation]
Now consider the minimum of the function g ( { a jj } ) on the right hand side, under the constraints that | a jj | ≤ 1 and (cid:80) j a 2 jj ≤ r .
We have the following claim for this optimization.
Lemma B.1.
There exists min(1 / 2 , s/r ) ≤ α ≤ 1 , β = min (cid:16) 1 , r − αs 2( r − s ) (cid:17) ∈ [1 / 2 , 1] , such that the minimum of the above optimization is achieved when a 2 jj = α for any j ∈ R 1 ∩ R 2 , and a 2 jj = β for any j ̸∈ R 1 ∩ R 2 .
Proof.
We first prove that to achieve the minimum, we can set:
(1) a 2 ℓℓ = a 2 ℓ ′ ℓ ′ for any ℓ ̸
= ℓ ′ ∈ R 1 ∩ R 2 ; (2) a 2 ℓℓ = a 2 ℓ ′ ℓ ′ for any ℓ ̸
= ℓ ′ ∈ R 1 \ R 2 ; (3) a 2 ℓℓ = a 2 ℓ ′ ℓ ′ for any ℓ ̸ = ℓ
′ ∈ R 2 \ R 1 .
For (1): By symmetry of z j ’s and the convexity of ℓ ( · ) ,
[Equation]
A similar inequality holds for R 2 .
Then
[Equation]
Therefore, the minimum is achieved when a 2 ℓℓ = a 2 ℓ ′ ℓ ′ .
For (2): The same inequality in Eqn.
( 16) holds for any ℓ ̸ = ℓ ′ ∈ R 1 \ R 2 , which then implies the statement (2).
For (3): The proof is similar to that for the statement (3).
These statements mean that the minimum is achieved when a 2 jj = α for j ∈ R 1 ∩ R 2 , a 2 jj = α 1 for j ∈ R 1 \ R 2 , and a 2 jj = αfor j ∈ R 2 \ R 1 , for some values
α, α 1 , α 2 ≥ 0 .
Let Z = (cid:80) j ∈ R 1 ∩ R 2 z 2 j , Z 1 = (cid:80) j ∈ R 1 \ R 2 z 2 j , Z 2 = (cid:80) j ∈ R 2 \ R 1 z 2 j .
By symmetry of z j ’s, Z 1 and Z 2 follow the same distribution.
Then
[Equation]
So the minimum is achieved when α 1 = α 2 := β , leading to
[Equation]
Given the constraint αs + 2 β ( r − s )
= (cid:80)
j a 2 jj ≤ r , 0 ≤ α, β ≤ 1 , and that ℓ ( · ) is monotonically non-increasing, we have α ∈ [0 , 1] , β = min (cid:16)
[Equation]
Furthermore, we can show that α ≥ min(1 / 2 , s/r ) .
Suppose α < min(1 / 2 , s/r ) for contradiction.
First consider the case when r ≤ 2 s .
Then 1 / 2 ≤ s/r
and thus α < 1 / 2 .
Note that
[Equation]
and α (cid:80) sj =1 z 2 j + β (cid:80)
rj
= s +
z 2 j is stochastically dominated by (cid:80) sj =1 z 2 j + 12 (cid:80) rj
= s +
z 2 j , which is achieved when α = 1 .
So the optimal cannot be achieved when α < 1 / 2 .
Next consider the case when r > 2 s .
Then 1 / 2 > s/r
and thus α < s/r .
We also have r − αs 2( r − s ) <
so β
= r − αs 2( r − s ) .
Let c j be the coefficients such that c j = α for 1 ≤ j ≤ s and c j = β for s < j ≤ r .
Then
[Equation] [Equation]
Let Π be the set of all permutations of [ r − s ] .
Again by symmetry of z j ’s and the convexity of ℓ ( · ) ,
[Equation]
When α < s/r , since β = r − αs 2( r − s ) , we have sα +( r − 2 s ) β r − s <
[Equation]
/ 2 , so
[Equation]
The right-most hand side is achieved when α = 1 , β = 1 / 2 .
So the optimal cannot be achieved when α < s/r .
In summary, we have α ≥ min(1 / 2 , s/r ) .
Therefore,
[Equation]
On the other hand, it can be verified that for any ϕ ∗ with W ∗ ,
[Equation]
Therefore, ϕ ∗ is the optimal solution.
B.2. Proof of Proposition 2.
Proof of Proposition 2.2.
E D
[(ˆ y − y ) 2 ] ≥ 0 and
[Equation]
The ( ∗ ) is from non-zero variance for z R 1 .
The empirical Rademacher complexity and Rademacher complexity of F ϕ ∗ with m samples are
[Equation]
Define X :=
A ∗ 1: r, 1: r (cid:80) mi =1 z ( i ) R 1 .
Note that for j ∈ R 1 ∩ R 2 , X j = α (cid:80) mi =1 z ( i ) j is a Gaussian of mean zero and variance
E [ X 2 j ] = α E (cid:20)(cid:16)(cid:80) mi =1 z ( i ) j (cid:17) 2 (cid:21) =
α E (cid:20)(cid:80) mi =1 (cid:16)
z ( i ) j (cid:17) 2 (cid:21) =
mα .
Similarly, for j ∈ R 1 \ R 2 , X j = β (cid:80) mi =1 z ( i ) j is a Gaussian of mean zero and variance E
[ X 2 j ] = mβ .
Since X j is sub-gaussian, X 2 j − mα for j ∈ R 1 ∩ R 2 and X
j − mβ for j ∈ R 1 \ R 2 are sub-exponential and more precisely
[Equation]
where C 1 , C 2 are absolute constants and C 2 > 1 .
Let K = max
( C 2 mα, C 2 mβ ) ≤ C 2 m and µ := m ( sα + ( r − s ) β ) .
By
Bernstein’s inequality, we have for every γ ≥ 0 that
[Equation]
where c is an absolute constant.
For all numbers z ≥ 0 , we have | z − 1 | ≥ δ ⇒ | z 2 − 1 | ≥ max( δ, δ 2 ) .
Thus, for any δ ≥ 0 , we have
[Equation]
52) where the last inequality is from µ ≤
mr . Changing variables to θ
= δ √ µ , we obtain the desired sub-gaussian tail
[Equation]
By generalization of integral identity, we have
[Equation] [Equation]
where C 3 is an absolute constant and the last inequality is from µ
= m ( sα + ( r − s ) β )
=
12 m ( sα + r ) ≥ 12 mr .
Thus, we have (cid:12)(cid:12)(cid:12)(cid:12) (cid:12)
[Equation]
ϕ
[Equation]
| (59) ≤ O
[Equation]
m
[Equation]
B.3. Discussion on the Implications
Pre-training and Predicting on One Task.
Suppose we only pre-train the representation and then learn a predictor on the task D 1 .
This is equivalent to setting r = s , which leads to D 1 = D 2 .
Then by Proposition 2.1, we know that
[Equation]
where e j ’s are the basis vectors.
The contrastive loss is
[Equation]
Since z j ’s are standard Gaussians, α = 1 in the optimal.
So ϕ ∗ ( x ) =
W ∗ x = (cid:80) j ∈ R 1 z j
[Equation]
j .
Furthermore, we have v 1 = ∥
u ∗ 1 ∥ 22 , v 2 = 0 in Proposition 2.2, then the Rademacher complexity of F ϕ ∗ satisfies
[Equation] [Equation]
Ignoring the low-order term on the right hand side, we have
[Equation]
Pre-training and Predicting on Two Tasks.
Consider the case with r = 2 s , that is, each task has half of the invariant features being common features, and the other half being task-specific features.
Then by Proposition 2.1, we know that
[Equation]
where e j ’s are the basis vectors.
We now show that α = 1 in the optimal.
Since r = 2 s , r − αs 2( r − s ) = 2 − α≤ 1 , then β = min (cid:16) r − αs 2( r − s ) , 1 (cid:17) = 2 − α 2 .
The contrastive loss is
[Equation]
where Z = (cid:80) sj =1 z 2 j and Z
1 = (cid:80) rj
= s +
z 2 j are i.i.d.
χ
2 s random variables.
Let Q α = αZ + 2 − α 2 Z 1 .
Then it can be shown that Q 1 stochastically dominates Q α for α < 1 .
Then the loss E [ ℓ ( Q α )] is minimized at α = 1 .
So
[Equation]
Then when v 1 = v 2 = ∥
u ∗ 1 ∥ 22 / 2 , by Proposition 2.2, the Rademacher complexity of F ϕ ∗ satisfies
[Equation]
Ignoring the low-order term on the right hand side, we have
[Equation]
Compared to the case of pre-training only on the target task, pre-training on the target task and a different task can lead to a larger Rademacher complexity and thus worse generalization.
B.4. Analysis for Multiple Tasks
Pre-training on Multiple Tasks.
Consider multiple different data distributions for contrastive learning, while the target (downstream) task is one of them.
That is, suppose we have T different invariant-feature subsets R 1 , . .
, R T ⊆
[ d ] , where | R t | = r for any t ∈ [ T ] .
Let R =
R 1 ∪ · · · ∪ R T be the set of all invariant features, and U =
[ d ] \ R be the set of spurious features.
Assume all distributions share a public feature set S :
= R 1 ∩ · · · ∩ R T of size s , and each distribution owns a private disjoint feature set of size r − s , which means ∀ t i , t j ∈ [ T ] , t i ̸
=
t j we have
R t i ∩ R t j \ S = ∅ .
For j ∈ S, z j ∼ N (0 , σ 2 S ) , σ S > 0 .
For the data distribution D t ( t = 1 , . . .
, T ) , we have for j ∈ R t \ S, z
j ∼ N (0 , σ 2 R ) , σ R > 0 , while z R \ R t = 0 , and y = ( u ∗ t )
⊤ z R t is the ground-truth label for task t .
Contrastive learning is over an uneven mixture of unlabeled data from the T distributions, where R t has weights w t > 0 and (cid:80) T t =1 w t = 1 .
The target downstream task is D 1 .
Without loss of generality, assume S = { 1 , , 2 , . . .
, s } and ∀ t ∈ [ T ] , R t =
S ∪ { s + ( t − 1)( r − s ) + 1 , s + ( t − 1)( r − s ) + 2 , . .
. , s + t ( r − s ) } and R = { 1 , 2 , . .
, s + T ( r − s ) } .
Proposition B.2 (General Version of Proposition 2.1) .
There exists α ∈ [0 , 1] , β = min (cid:16) 1 , r − αs T ( r − s ) (cid:17) ∈ [1 /T, 1] such that the following holds.
ϕ ∗ ( x ) =
W ∗ x is an optimal representation for the loss (1) in contrastive learning with any W ∗ of the form:
[Equation]
where O ∈ R k × k is any orthnormal matrices, A ∗ is a k × k diagonal matrix with
[Equation]
otherwise , (69) and the matrix of zeros has size k × ( d − k ) .
Proof.
Following the same argument as in the proof of Proposition 2.1,
E
[Equation] [Equation]
g
[Equation] [Equation]
[Equation]
72) where ˜ z j is a random variable draw from standard Gaussian.
Following the same argument as in Lemma B.1, we have the following claim: to achieve the minimum, we can set (1) a 2 ℓℓ = a 2 ℓ ′ ℓ ′ =: α for any ℓ ̸ = ℓ ′ ∈ S , (2) a 2 ℓℓ = a 2 ℓ ′ ℓ ′ :
= α t for any ℓ ̸ = ℓ ′ ∈ R t \ S, ∀ t ∈
[ T ] .
Let Z = (cid:80) j ∈ S ˜
z
2 j and Z t = (cid:80) j ∈ R t \ S ˜ z
2 j , ∀ t ∈
[ T ] .
By symmetry of ˜ z
j ’s, all Z t follow the same distribution.
Then
[Equation]
So the minimum is achieved when α t i = α t j :=
β for any t i , t j ∈ [ T ] , leading to
[Equation]
Given the constraint αs + Tβ ( r − s ) =
(cid:80)
j a 2 jj ≤ r , 0 ≤ α, β ≤ 1 .
Following the same argument as in the proof of
Proposition 2.1, we finish the proof.
Given this result we can also analyze the Rademacher complexity of the predictor class similarly as that for the case with two tasks and obtain a result similar to Proposition 2.2.
Similar to the two-task setting, f ( ϕ ) =
u ⊤ ϕ with u = O 1: k, 1: r ( A ∗ 1: r, 1: r ) − 1 u ∗ 1 + O 1: k,r +1: k v for any v ∈ R k − r satisfies f ( ϕ ∗ ( x )) =
y =
( u ∗ 1 )
⊤ z R 1 , and u ∗
=
O 1: k, 1: r ( A ∗ 1: r, 1: r ) − 1 u ∗ 1 is the least norm optimal solution.
So the predictor class should be the same as Eqn.
Proposition B.3 (General Version of Proposition 2.2) .
Suppose α > 0 .
Let v 1 = (cid:80) sj =1 ( u ∗ 1 j ) 2 and v
= (cid:80) rj
= s +1 ( u ∗ 1 j ) 2 .
Then the Rademacher complexity of F ϕ ∗ in Eqn.
( 5) satisfies
[Equation]
Proof of Proposition B.3.
Following the same argument as in the proof of Proposition 2.2, the Rademacher complexity of
F ϕ ∗ with m samples are
[Equation]
Define X :=
A ∗ 1: r, 1: r (cid:80) mi =1 z ( i ) R 1 .
Following the same argument as in the proof of Proposition 2.2, we have
[Equation]
where C 2 is an absolute constants and C 2 > 1 .
Let K = max( C 2 mασ 2 S , C 2 mβσ 2 R )
≤ C 2 m ( σ 2 S + σ 2 R ) and µ := m ( sασ 2 S
+ ( r − s ) βσ 2 R ) .
By Bernstein’s inequality, we have for every γ ≥ 0 that
[Equation]
where c is an absolute constant.
We have µ ≤ m ( σ 2 S + σ 2 R ) r .
Following the same argument as in the proof of Proposition 2.2, for any δ ≥ 0 , we have
[Equation]
Changing variables to θ = δ √ µ , we obtain the desired sub-gaussian tail
[Equation]
By generalization of integral identity
[Equation]
following the same argument as in the proof of Proposition 2.
[Equation]
we have
[Equation]
where the last inequality is from µ = m ( sασ 2 S + ( r − s ) βσ 2 R ) ≥ min { σ 2 S , σ 2 R } m ( sα + ( r − s ) β ) ≥ min { σ 2 S , σ 2 R } 12 mr and C 3 is an absolute constant.
Thus, we have
[Equation]
C. More Experiments Details and Results C.1. Dataset
CIFAR-10.
CIFAR-10 (Krizhevsky et al., 2009) dataset consists of 60,000 32 × 32 color images in 10 classes: airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck.
Each class has 6,000 images.
There are 50,000 training images and 10,000 test images.
CINIC-10.
CINIC-10 (Darlow et al., 2018) consists of 32 × 32 color images from both CIFAR and ImageNet and has 90,000 training images with ten classes identical to CIFAR-10.
SVHN.
The Street View House Numbers (Netzer et al., 2011) contains 10 digits color images of size 32 × 32 in natural scene.
It has 73,257 digits for training and 26,032 digits for testing.
GTSRB.
The German Traffic Sign Recognition Benchmark (Stallkamp et al., 2012) is a dataset of color images depicting 43 different traffic signs.
The images are not of fixed dimensions and have a rich background and varying light conditions as expected of photographed images of traffic signs.
The original training set contains 34,799 images, and the original test set contains 12,630 images.
We resize each image to 32 × 32.
The dataset has a significant imbalance in the number of sample occurrences across classes.
We use data augmentation techniques to enlarge the training data and balance the number of samples in each class.
We construct a class preserving data augmentation pipeline consisting of rotation, translation, and projection transforms and apply this pipeline to the training images until each class contains 2,500 examples.
So we construct a new training set containing 107,500 images in total.
We also construct a new test set by randomly selecting 10,000 images from the original test set for evaluation.
ImageNet32.
ImageNet32 (Deng et al., 2009) is a huge dataset made up of small images called the down-sampled version of ImageNet.
ImageNet32 is composed of 1,281,167 training data and 50,000 test data with 1,000 labels.
MNIST.
The Modified National Institute of Standards and Technology (LeCun et al., 1998) is a database of handwritten gray-scale digits of size 28 × 28 .
It contains 60,000 training images and 10,000 testing images.
EMNIST.
Extended MNIST (Cohen et al., 2017) includes images of handwritten letters and digits.
The images in EMNIST were converted into the same size 28 × 28 by the same process as MNIST.
EMNIST-Letters has 145,600 lower case characters with 26 balanced classes, and EMNIST-Digits has 280,000 characters with ten balanced classes.
Fashion-MNIST.
Fashion-MNIST (Xiao et al., 2017) is a dataset of 28 × 28 gray-scale images with ten classes: T-shirt/top, trouser, pullover, dress, coat, sandal, shirt, sneaker, bag, ankle boot.
The training set size is 60,000, and the test set size is 10,000.
C.2. Ablation Study
Varying Target-Relevant pre-training Data Percentage.
In Fig.
3 and 4, we use (a) 100% (b) 50% (c) 20% CINIC-10 to train MoCo v2 and SimSiam, and keep the same setting as Fig.
For Fig.
3 (b) with 50% CINIC-10, test accuracy drops, e.g., the test accuracy of 1% CIFAR-10 in Fig.
3 (a) 80 .
63% vs. (b) 76 .
We can also see the decreasing curve in Fig.(b).
On the other hand, we also have test accuracy drops in Fig.
3 (c) and Fig.
4 (b) (c).
However, we can see a U-curve rather than a strictly decreasing curve in Fig.
3 (c) and Fig.
4 (b) (c).
ImageNet32 is more relevant with CIFAR-10 than
SVHN and GTSRB, consistent with human intuition.
When we have a small partition of CINIC-10 which does not cover all target relevant features, the feature extractor can learn these missing features from ImageNet32.
Although there are many irrelevant features in ImageNet32, the positive effect is larger than the negative effect, and so it plots a U-curve.
It is
[Figure]
(c
[Figure] [Figure]
c y
[Figure]
consistent with our statement that we need a large and target relevant pre-training dataset rather than a diverse irrelevant one
[Figure]
Replacing CINIC-10 With CIFAR-10.
In Fig.
[Figure]
we keep the same setting as Fig.
1 except we replace CINIC-10 with
CIFAR-10.
Note that our downstream task is still CIFAR-10.
In Fig.
5, we can see the same phenomena and similar performance as Fig. 1.
Thus, if we have a good choice of a task-relevant pre-training dataset, we can get similar performance as pre-training on the downstream task domain directly
[Figure]