Predicting What You Already Know Helps: Provable Self-Supervised Learning

Abstract:
Self-supervised representation learning solves auxiliary prediction tasks (known as pretext tasks) without requiring labeled data to learn useful semantic representations. These pretext tasks are created solely using the input features, such as predicting a missing image patch, recovering the color channels of an image from context, or predicting missing words in text; yet predicting this known information helps in learning representations effective for downstream prediction tasks. We posit a mechanism exploiting the statistical connections between certain reconstruction-based pretext tasks that guarantee to learn a good representation. Formally, we quantify how the approximate independence between the components of the pretext task (conditional on the label and latent variables) allows us to learn representations that can solve the downstream task by just training a linear layer on top of the learned representation. We prove the linear layer yields small approximation error even for complex ground truth function class and will drastically reduce labeled sample complexity.

heading: 1 Introduction
text: Self-supervised learning revitalizes machine learning models in computer vision, NLP, and control problems (see reference therein [36, 38, 15, 63, 35]). Training a model with auxiliary tasks based only on input features reduces the extensive costs of data collection and semantic annotations for downstream tasks. It is also known to improve the adversarial robustness of models [29, 11, 12]. Self-supervised learning creates pseudo labels solely based on input features, and solves auxiliary prediction tasks (or pretext tasks) in a supervised manner. However, the underlying principles of self-supervised learning are mysterious since it is a-priori unclear why predicting what we already know should help. We thus raise the following question:
What conceptual connection between pretext and downstream tasks ensures good representations? What is a good way to quantify this?
As a thought experiment, consider a simple downstream task of classifying desert, forest, and sea images. A meaningful pretext task is to predict the background color of images (known as image colorization [66]). Denote X1, X2, Y to be the input image, color channel, and the downstream label respectively. Given knowledge of the label Y , one can possibly predict the background X2 without knowing much about X1. In other words, X2 is approximately independent of X1 conditional on the label Y . Consider another task of inpainting [48] the front of a building (X2) from the rest (X1). While knowing the label “building” (Y ) is not sufficient for successful inpainting, adding additional latent variables Z such as architectural style, location, window positions, etc. will ensure that variation in X2 given Y, Z is small. We can mathematically interpret this as X1 being approximate conditionally independent of X2 given Y,Z.
The main insight that we exploit in this work is that with approximate conditional independence (as in the above examples), a method that predicts X2 from X1 will inadvertently implicitly encode and
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
learn to predict Y (and Z) from X1 as an intermediate step, and then predict X2 from Y 1. Building upon this insight, we make the following contributions.
Contributions. The goal of this paper, as in statistical learning theory, is to investigate the statistical connections between the random variables of input features (in this paper (X1, X2)) and downstream labels Y , and show how specific connections can guarantee a successful learning procedure. For self-supervised learning (SSL), success is measured using the following 2 notions, 1) expressivity, i.e. does the learned representation from SSL have the ability to express the ground truth prediction function for labels Y , and 2) sample complexity, i.e. can it do so with way fewer labeled samples than what would be required without SSL.
In this work, we show such guarantees for a class of reconstruction-based SSL methods under a statistical assumption of approximate conditional independence (ACI). In particular we show that under such an assumption, the learned representation from SSL will end up having the following properties, 1) it can express the ground truth label as a linear function, thus guaranteeing expressivity, and 2) will also end up being low-rank (or low-dimensional), thus guaranteeing smaller labeled sample complexity. Note that such an expressive and sample efficient (summarized as good) representation is often not a-priori available. For instance, the original input features themselves may not be able to express the ground truth function linearly, while kernel methods with a fixed kernel, while expressive, may not be sample efficient for many problems of interest. The strategy in modern machine learning is to find such a good representation as the output of a complicated neural network. The benefit of SSL, as we formally show here, is that the complicated but good representation function can be learned using just unlabeled data, so that labeled data is just needed to learn a linear function.
The reconstruction-based SSL method (differentiated from other SSL methods in Section 1.1) we consider is strongly motivated by empirical works [66, 48, 15, 25], but is a simplification that captures the essence of the problem and is amenable to a precise theoretical analysis. We consider a two-staged pipeline, where we first learn a representation function ψ (e.g. output of a neural network) from input X1 and pretext target X2 using unlabeled data by minimizing E(X1,X2)[∥X2 − ψ(X1)∥2]. In the second stage of downstream task, we learn a linear layer on top of representation ψ using labeled samples (X1, Y ), thus restricting to learning from a significantly smaller hypothesis class of Hψ = {f : X1 → Y |f is linear in ψ}. The key non-trivial question of expressivity is now whether the ground truth predictor f∗ ≡ E[Y |X1] can be approximated well by this classHψ, and the question of sample complexity reduces to the understanding the sample complexity of learning Hψ. Under appropriate statistical connections2 between input data X1, X2 and target Y , we prove both the desired properties, expressivity and low sample complexity, for the aforementioned SSL method.
Our statistical assumption based on approximate conditional independence (ACI) helps us demonstrate how solving pretext tasks created from known information can learn useful representations. Specifically, we show that once the complicated representation function ψ is learned using an abundance of unlabeled data in the SSL stage, not only is ψ expressive enough, but it will also require only Õ(k) labeled samples to solve a k-way supervised learning task under exact conditional independence (CI). In contrast, solving the downstream task without any pretraining will require a lot of labeled data to learn the representation function from scratch. Since the strong exact conditional independence assumption will likely not be satisfied in practice, our main contribution is to derive similar risk bounds when only approximate CI (ACI) is satisfied. We quantify the notion of ACI using the norm of a certain partial covariance matrix (Definition 4.1) and our risk bound scales linearly with it. We verify this and other aspects of our main Theorem 4.2 using simulations and also find that pretext task helps when CI is approximately enforced in text domain. We further demonstrate on a real-world image dataset that a pretext task-based linear model performs at least as well as many baselines.

heading: 1.1 Related work
text: Self-supervised learning (SSL) methods in practice: There has been a flurry of self-supervised methods lately. One class of methods reconstruct images from corrupted or incomplete versions of it, like denoising auto-encoders [61], image inpainting [48], and split-brain autoencoder [67]. Pretext
1This is formally demonstrated in the proof sketch of Lemma 3.1. 2We note that since the representation ψ is the result of the SSL method and not something we have access to a-priory, we cannot make any direct assumptions on it.
tasks are also created using visual common sense, including predicting rotation angle [22], relative patch position [16], recovering color channels [66], solving jigsaw puzzle games [45], and discriminating images created from distortion [17]. We refer to the above procedures as reconstruction-based SSL. Another popular paradigm is contrastive learning [13, 14]. The idea is to learn representations that bring similar data points closer while pushing randomly selected points further away [63, 39, 5] or to maximize a contrastive-based mutual information lower bound between different views [30, 46, 54]. A popular approach for text domain is based on language modeling where models like BERT and GPT create auxiliary tasks for next word predictions [15, 49]. The natural ordering or topology of data is also exploited in video-based [64, 43, 19], graph-based [65, 33] or map-based [68] SSL. For instance, the pretext task is to determine the correct temporal order for video frames as in [43].
Theory for SSL: While we theoretically study reconstruction-based SSL, prior work has different flavors of theoretical results for different kinds of SSL methods. Most relevant are the guarantees for representation learning using SSL methods on downstream tasks that just learn a linear classifier on top of the learned representations. [5] shows guarantees for representations from a contrastive learning objective: Lcont1 (ψ) = E(X1,X2),X′2 [log(1 + e
−ψ(X1)⊤ψ(X2)+ψ(X1)⊤ψ(X′2))]. Under a class conditional independence assumption, i.e. X1 ⊥ X2 | Y , they show that representation ψ that does well on contrastive objective, i.e. Lcont1 (ψ) ≤ ϵ, will have O(ϵ) linear classification loss on the average binary task involving pairs of classes (y1, y2). However, their analysis cannot handle the general case of approximate conditional independence. Recently, Tosh et al. [56] show that contrastive learning representations can linearly recover continuous functions of the underlying topic posterior under a topic modeling assumption for text. While their assumption bears similarity to ours, the assumption of independent sampling of words is strong and does not generalizable to other domains like images. Most relevant is a concurrent work [57] that shows guarantees for a contrastive learning objective that looks like Lcont2 (ψ, η) = E(X1,X2),X′2 [ log(1 + e−ψ(X1) ⊤η(X2)) + log(1 + eψ(X1) ⊤η(X′2)) ] , with a multi-view redundancy assumptions that is very similar to our ACI assumption. We take a closer look at their assumption in Section F.2. All the above objectives are different from the simple reconstruction-based objective we consider: L(ψ) = E(X1,X2) [ ∥X2 − ψ(X1)∥2 ] . Saunshi et al. [51] show guarantees for representations learned using language modeling on sentence classification tasks. Some more recent work [58, 44, 55, 62] provide theoretical understanding on SSL respectively based on causality, mutual information, gradient-descent dynamics, and alignment/uniformity of representations, without explicit risk bounds for downstream tasks. There is a mutual information maximization view of contrastive learning, but [59] points out issues with it. Previous attempts to explain negative sampling [42] based methods use the theory of noise contrastive estimation [27, 40] to show asymptotic guarantees, without explicit connections to downstream tasks. CI is also used in sufficient dimension reduction [21, 20], while CI and redundancy assumptions on multiple views [37, 2] are used to analyze a canonical-correlation based dimension reduction algorithm and also for self-supervised learning algorithms like co-training [10]. Finally, [1, 60] provide a theoretical analysis for denoising auto-encoder.

heading: 1.2 Overview of results:
text: Section 2 introduces notation, setup, and the self-supervised learning procedure considered in this work. In Section 3, we analyze downstream sample complexity under exact CI and unlimited labeled data to highlight the key ideas. Section 4 presents our main result with relaxed conditions: under ACI with latent variables, and assuming finite samples in both pretext and downstream tasks, for various function classes, and both regression and classification tasks. Experiments verifying our theoretical findings are in Section 6. Proofs of most results are in the Appendix.

heading: 2 Preliminary

heading: 2.1 Notation
text: We use lower case symbols (x) to denote scalar quantities, bold lower case symbols (x) for vector values, capital letters (X) for random variables, and capital and bold letters X for matrices. PX denotes the probability law of random variable X , and the space of square-integrable functions with probability P is denoted by L2(P ). We use standard O notation to hide universal factors and Õ to hide log factors. ∥ · ∥ stands for ℓ2-norm for vectors or Frobenius norm for matrices.
Linear conditional expectation. EL[Y |X] denotes the prediction of Y with linear regression:
EL[Y |X = x] := W ∗x+ b∗, where W ∗, b∗ := argmin W ,b E[∥Y −WX − b∥2].
In other words, EL[Y |X] denotes the best linear predictor of Y givenX . We also note that E[Y |X] ≡ argminf E[∥Y − f(X)∥2] is the best predictor of Y given X . (Partial) covariance matrix. For random variables X,Y , we denote ΣXY to be covariance matrix of X and Y . For simplicity in most cases, we assume E[X] = 0 and E[Y ] = 0; thus we do not distinguish E[XY ] and ΣXY . The partial covariance matrix between X and Y given Z is:
ΣXY |Z :=cov{X − EL[X|Z], Y − EL[Y |Z]} ≡ ΣXY −ΣXZΣ−1ZZΣZY , (1)
which captures the correlation between X and Y setting aside the effect of Z.
Sub-gaussian random vectors. X ∈ Rd is ρ2-sub-gaussian if for every fixed unit vector v ∈ Rd, the variable v⊤X is ρ2-sub-gaussian, i.e., E[es·v⊤(X−E[X])] ≤ es2ρ2/2 (∀s ∈ R).

heading: 2.2 Setup and methodology
text: We denote by X1 the input variable, X2 the target random variable for the pretext task, and Y the label for the downstream task, with X1 ∈ X1 ⊂ Rd1 , X2 ∈ X2 ⊂ Rd2 and Y ∈ Y ⊂ Rk. If Y is finite with |Y| = k, we assume Y ⊂ Rk is the one-hot encoding of the labels. PX1X2Y denotes the joint distribution over X1 × X2 × Y . PX1Y , PX1 denote the corresponding marginal distributions. Our proposed self-supervised learning aims to fulfill the following two steps:
Step 1 (pretext task): Learn a representation ψ(x1) close to ψ∗ := argming∈H E∥X2 − g(X1)∥2, whereH can vary for different settings that we will specify and discuss later. Step 2 (downstream task): Perform linear regression on Y with ψ(X1), i.e. f(x1) := (W ∗)⊤ψ(x1), where W ∗ ← argminW EX1,Y [∥Y −W⊤ψ(X1)∥2]. Namely we learn f(·) = EL[Y |ψ(·)]. We study this simplified version in the main text, where in practice, the SSL procedure may utilize an encoder-decoder structure, while the downstream task uses both X1 and X2 to predict Y . We incorporate these extensions in Appendix C.3 and G.
With finite samples, performance of a learned representation ψ on the downstream task depends on the following quantities that capture expressivity and sample complexity respectively:
Approximation error indicates whether Y is linearly separable by the learned representation ψ, thus measuring expressivity. We measure this by comparing Wψ(X1) to the optimal predictor f∗ := E[Y |X1 = x1]. Denote eapx(ψ) = minW E[∥f∗(X1)−Wψ(X1)∥2]. This gives a measure of how well ψ can linearly predict Y when given infinite samples for the task.
Estimation error measure sample complexity of ψ on the downstream task and assume access to n2 i.i.d. samples (x (1) 1 ,y
(1)), · · · , (x(n2)1 ,y(n2)) drawn from PX1Y . We express the n2 samples collectively as Xdown1 ∈ Rn2×d1 , Y ∈ Rn2×k and overload notation to say ψ(Xdown1 ) =[ ψ(x
(1) 1 )|ψ(x (2) 1 ) · · · |ψ(x (n2) 1 )
]⊤ ∈ Rn2×d2 . We perform linear regression on the learned represen-
tation ψ and measure excess risk, that incorporates both approximation and estimation errors.
Ŵ ← argmin W
1
2n2 ∥Y − ψ(X1)W ∥2F ; ERψ(Ŵ ) :=
1 2 E∥f∗(X1)− Ŵ⊤ψ(X1)∥22.

heading: 3 Guaranteed recovery with conditional independence
text: In this section, we focus on the case where the input X1 and pretext target X2 are conditionally independent (CI) given the downstream label Y . While this is a strong assumption that is rarely satisfied in practice, it helps us understand the role of CI with clean results and builds up to our main results with ACI with latent variables in Section 4. As a warm-up, we show how CI helps when (X1, X2, Y ) are jointly Gaussian to give us a flavor for the results to follow in Appendix B. We then analyze it for general random variables under two settings: (a) when the function class used for ψ is universal, (b) when ψ is restricted to be a linear function of given features. For now we assume access to a large amount of unlabeled data so as to learn the optimal ψ∗ perfectly and this will be relaxed later in Section 4. The general recipe for the results is as follows:
1. Find a closed-form expression for the optimal solution ψ∗ for the pretext task. 2. Use conditional independence to show that optimal f∗ is linear in ψ∗, i.e., eapx(ψ∗) is small. 3. Exploit the low rank structure of ψ∗ to show small estimation error on downstream tasks.
Data assumption. Suppose Y = f∗(X1) +N , where f∗ = E[Y |X1] and E[N ] = 0. We assume N is σ2-subgaussian. For simplicity, we assume non-degeneracy: ΣXiXi , ΣY Y are full rank.
Assumption 3.1. Let X1 ∈ Rd1 , X2 ∈ Rd2 be random variables from some unknown distribution. Let label Y ∈ Y be a discrete random variable with k = |Y| < d2. We assume conditional independence: X1⊥X2|Y .
Here Y can be interpreted as the multi-class labels where k is the number of classes. For regression problems, one can think about Y as the discretized values of continuous labels. We do not specify the dimension for Y since Y could be arbitrarily encoded but the results only depend on k and the variance of Y (conditional on the input X1).

heading: 3.1 Universal function class.
text: Suppose we learn the optimal ψ∗ among all measurable functions The optimal function ψ∗ in this case is naturally given by conditional expectation: ψ∗(x1) = E[X2|X1 = x1]. We show that CI implies that ψ∗ is good for downstream tasks, which is not apriori clear.
Lemma 3.1 (Approximation error). If random variables X1, X2, Y satisfy Assumption 3.1, and A ∈ RY×d2 with Ay,: := E[X2|Y = y] has rank k = |Y|. Then f∗ ≡W ∗ψ∗, i.e., eapx(ψ∗) = 0.
This tells us that although f∗ could be nonlinear in x1, it is guaranteed to be linear in ψ∗(x1).
Proof Sketch of Lemma 3.1. Lemma is proved by law of total expectation:
ψ∗(·) := E[X2|X1] =E[E[X2|X1, Y ]|X1] = E[E[X2|Y ]|X1] (uses CI) = ∑ y P (Y = y|X1)E[X2|Y = y] =: f(X1)⊤A,
where f(x1)y = P (Y = y|X1 = x1), and A ∈ RY×d2 satisfies Ay,: = E[X2|Y = y]. One could see that through predicting X2, due to the CI assumption, ψ∗ has implicitly encoded the information of Y |X1. Finally due to the fact that matrix A is full rank, we get that f∗ is linear in ψ∗ as well.
We see that besides CI, another important property is E[X2|Y ] being rank k. This means X2 is correlated with every instance of Y , and thus captures information of every prediction class. This is naturally a necessary assumption for X2 to be a reasonable pretext task for predicting Y . Note that this assumption does not trivialize the problem and that even though ψ is designed to predict X2, it can still be a better representation than X2 for downstream tasks. Note that Y does not have to be linear in X2 but is proven to be linear in ψ, since ψ learns to ignore some information in X2 that is irrelevant to Y . We provide this simple example for better understanding:
Example 3.1. Let Y ∈ {−1, 1} be binary labels, and X1, X2 be 2−mixture Gaussian random variables withX1 ∼ N (Y µ1, I), X2 ∼ N (Y µ2, I). In this example,X1⊥X2|Y . Although E[Y |X2] and E[Y |X1] are not linear, E[Y |ψ] is linear: ψ(x1) = P (Y = 1|X1 = x1)µ2−P (Y = −1|X1 = x1)µ2 and f ∗(x1) = P (Y = 1|X1 = x1)− P (Y = −1|X1 = x1) ≡ µT2 ψ(x1)/∥µ2∥2.
Given that ψ∗ is good for downstream, we now care about the sample complexity. We will need to assume that the representation has some nice concentration properties. We make an assumption about the whitened data ψ∗(X1) to ignore scaling factors.
Assumption 3.2. We assume the whitened feature variable U := Σ−1/2ψ ψ(X1) is a ρ 2-subgaussian random variable, where Σψ = E[ψ(X1)ψ(X1)⊤].
We note that all bounded random variables satisfy sub-gaussian property.
Theorem 3.2 (General conditional independence). Fix a failure probability δ ∈ (0, 1), under the same assumption as Lemma 3.1 and Assumption 3.2 for ψ∗, if additionally n≫ ρ4(k + log(1/δ)), then the excess risk of the learned predictor x1 → Ŵ⊤ψ∗(x1) on the downstream task satsifies
ERψ∗ [Ŵ ] ≤ Õ ( k n2 σ2 ) 3
Remark 3.1. This analysis assumes we could perfectly learn ψ∗ = E[X2|X1] disregarding the number of samples in the SSL phase (unlabeled data is cheap to obtain). Here by sample complexity we refer to the labeled data (X1, Y ). We defer the effect of imprecise representation ψ in Section 4.

heading: 3.2 Function class induced by feature maps.
text: Given feature map ϕ1 : X1 → RD1 , we consider the function class H1 = {ψ : X1 → Rd2 |∃B ∈ Rd2×D1 , ψ(x1) = Bϕ1(x1)}.
Claim 3.3 (Closed form solution). The optimal function in H is ψ∗(x1) = ΣX2ϕ1Σ−1ϕ1ϕ1ϕ1(x1), where ΣX2ϕ1 := ΣX2ϕ1(X1) and Σϕ1ϕ1 := Σϕ1(X1)ϕ1(X1). We again show the benefit of CI, but only comparing the performance of ψ∗ to the original features ϕ1. Since ψ∗ is linear in ϕ1, it cannot have smaller approximation error than ϕ1. However CI will ensure that ψ∗ has the same approximation error as ϕ1 and enjoys better sample complexity.
Lemma 3.4 (Approximation error). If Assumption 3.1 is satisfied, and if the matrix A ∈ RY×d2 with Ay,: := E[X2|Y = y] is of rank k = |Y|. Then eapx(ψ∗) = eapx(ϕ1).
We additionally need an assumption on the residual a(x1) := E[Y |X1 = x1] − EL[Y |ϕ1(x1)].
Assumption 3.3. (Bounded approx. error; Condition 3 in [32])) We have almost surely
∥Σ−1/2ϕ1ϕ1 ϕ1(X1)a(X1) ⊤∥F ≤ b0
√ k
Theorem 3.5. (CI with approximation error) Fix a failure probability δ ∈ (0, 1), under the same assumption as Lemma 3.4, Assumption 3.2 for ψ∗ and Assumption 3.3, if n2 ≫ ρ4(k + log(1/δ)), then the excess risk of the learned predictor x1 → Ŵ⊤ψ∗(x1) on the downstream task satisfies:
ERψ∗ [Ŵ ] ≤ eapx(ϕ1) + Õ ( k n2 σ2 ) .
Thus with SSL, the requirement of labels is reduced from complexity for D1 to O(k).

heading: 4 Beyond conditional independence
text: In the previous section, we focused on the case where we have exact CI. A weaker but more realistic assumption is that Y captures some portion of the dependence between X1 and X2 but not all. We quantify this notion of approximate ACI through a quantity ϵ2CI (Definition 4.1), and show excess risk bounds for the representation learned from SSL4. In particular, the excess risk will have the form Õ ( d2 n2 + ϵ2CI + ϵ 2 pre ) , which suggests that only n2 = O(d2) labeled samples will be required to get small error on downstream task, as long as approximate CI is satisfied (ϵ2CI is small) and the pretext task is solved well enough (ϵ2pre is small). This is in contrast to not doing SSL, where many more labeled samples will be required to learn a solve the downstream task that learns a complicated representation function from scratch. We now describe the SSL method on finite samples, followed by the definition of ACI which we use to discuss the main excess risk bound and its consequences.
SSL with finite samples and general function space: Let Xpre1 = [x (1,pre) 1 , · · · ,x (n1,pre) 1 ] ⊤ ∈ Rn1×d1 and X2 = [x(1)2 , · · · ,x (n1) 2 ]
⊤ ∈ Rn1×d2 be n1 training samples for pretext task, where (x
(i,pre) 1 ,x (i) 2 ) is sampled from PX1X2 . The n2 labeled samples for the downstream task are defined as Xdown1 ∈ Rn2×d1 , Y ∈ Rn2×d35. Given a representation function space H : X1 → Rd2 , we learn ψ̃ fromH using the n1 unlabeled samples and then use the n2 labeled samples to learn a linear classifier on the learned representation ψ̃(Xdown1 ) to fit Y . This process is summarized below.
1) ψ̃ := argmin f∈H
1
n1 ∥X2 − f(Xpre1 )∥2F , 2) Ŵ ← argmin W
1
2n2 ∥Y − ψ̃(Xdown1 )W ∥2F . (2)
3We will use Õ to hide log factor log(k/δ) or log(d2/δ). 4Results for jointly-Gaussian variables is in Appendix D.1; ACI is quantified by the partial covariance matrix. 5d3 = k and Y ≡ ϕy(Y ) (one-hot encoding) refers multi-class classification task, d3 = 1 refers to regression.
In our main results, we consider two types of function spaces: H ∈ {H1,Hu}. Recall that H1 = {ψ(·) = Bϕ1(·);B ∈ Rd2×D1} is a class of linear representations induced by feature map ϕ1 : X1 → RD1 . We useHu to denote a function space with universal approximation power (e.g. deep networks) that ensures ψ∗ = E[X2|X1] ∈ Hu. We define the optimal predictor in each case as f∗H(X1) = EL[Y |ϕ1(X1)] whenH = H1, f∗H = f∗ forH = Hu, we define excess risk as
ERψ̃(Ŵ ) := EX1 [ ∥f∗H(X1)− Ŵ⊤ψ̃(X1)∥22 ] .
Approximate conditional independence: Our new assumption will generalize Assumption 3.1 in two ways, 1) we allow for additional latent variables Z that together with Y could potentially make X1 and X2 independent, and 2) we allow this conditional independence to be approximate. Note that allowing for extra latent variable can trivially make X1 and X2 to be conditionally independent by picking a large enough Z (e.g. Z = (X1, X2)). However the following assumption, that needs the pretext target X2 to correlate with all instances of variable Ȳ = [Y,Z] (analogous to Lemma 3.1), will impose this restriction on how large Z can be.
Assumption 4.1 (Correlation between X2 and Y,Z). Suppose there exists latent variable Z ∈ Z, |Z| = m that ensures ΣϕȳX2 is full column rank and ∥ΣY ϕȳΣ † X2ϕȳ
∥2 = 1/β, where A† is pseudo-inverse, and ϕȳ is the one-hot embedding for Ȳ = [Y, Z].
Just as in Section 3, this assumption will not assume away the problem (Example 3.1 can be suitably extended). The additional term 1/β here captures both the “scale” of X2 and also the strength of correlation between X2 and [Y,Z] that was discussed after Lemma 3.1. For ΣϕȳX2 to be full column rank, it is essential that d2 ≥ km, and this already gives an upper bound on the size of Z. Given this restriction on Z (and thus Ȳ ), we define the notion of approximate conditional independence.
Definition 4.1 (Approximate conditional independence with function spaceH). For Ȳ = [Y,Z], 1. ForH = H1, define ϵCI := ∥Σ−1/2ϕ1ϕ1 Σϕ1X2|ϕȳ∥F . 2. ForH = Hu, define ϵ2CI := EX1 [∥E[X2|X1]− EȲ [E[X2|Ȳ ]|X1]∥2]. Firstly we note that this is indeed an extension of exact CI, since exact CI in both cases will imply that ϵCI = 0. We present a unified analysis in the appendix that shows the ϵCI for the second case is same as the first case, with covariance operators instead of matrices (A direct derivation is in Claim D.7). We also present more relaxed and general form of the above assumptions in Appendix F.1. With this assumption, we are ready to present our main bound.
Bound on excess risk: Recall that we assume that the residual term N := Y − E[Y |X1] is mean zero and σ2-subgaussian. Before showing our main result, analogous to Assumption 3.3, for the class H1 with non-universal features ϕ1, we will need an assumption6 on the residual a := f∗ − f∗H1 = E[Y |X1]− E
L[Y |ϕ1(X1)]: Assumption 4.2. (Bounded approximation error on pretext phase [32]) There exists a universal constant b0, such that ∥Σ−1/2ϕ1ϕ1 ϕ1(X1)a(X1) ⊤∥F ≤ b0 √ d2 almost surely.
Theorem 4.2. For a fixed δ ∈ (0, 1), under Assumptions 4.1,3.2 for ψ̃ and ψ∗ and 4.2 for nonuniversal feature maps, if n1, n2 ≫ ρ4(d2 + log 1/δ), and we learn the pretext tasks such that: E∥ψ̃(X1)− ψ∗(X1)∥2F ≤ ϵ2pre. Then the generalization error for downstream task w.p. 1− δ is:
ERψ̃(Ŵ ) ≤ Õ  σ2 d2n2︸ ︷︷ ︸ estimation error + ϵ2CI β2 + ϵ2pre β2︸ ︷︷ ︸ approximation error  (3) We defer the proof to the appendix. The proof technique is similar to that of Section 3. The difference is that now ψ̃(X(down)) ∈ Rn2×d2 will be an approximately low rank matrix, where the low rank part is the high-signal features that implicitly comes from Y,Z that can linearly learn downstream task. The remaining part comes from ϵCI and ϵpre and causes the approximation error. Again by selecting the top km (dimension of ϕȳ) features we could further improve the bound:
6This rules out the failure if one chooses a very simple function class to learn E[X2|X1]. In practice we usually use neural networks (with universal approximation power) and this bound should be very small.
Remark 4.1. By applying PCA on ψ̃(Xdown1 ) and keeping the top km principal components only, we can improve the bound in Theorem 4.2 to ERψ̃(Ŵ ) ≤ Õ ( σ2 kmn2 + ϵ2CI β2 + ϵ2pre β2 ) .
We take a closer look at the different sources of errors in Lemma 4.1: 1) The first term is estimation error on learning with finite samples n2 with noise level σ2 in Y − f∗(X1); 2) ϵCI measures the approximate CI; and 3) ϵpre is the error from not learning the pretext task exactly. The first term is optimal ignoring log factors as we do linear regression on mk-dimensional features. The second and third term together form approximation error. They are non-reducible due to the fact that f∗ is not exactly linear in ψ and we use it as a fixed representation. Fine-tuning the representations might be necessary to get rid of these terms when we have sufficient downstream labeled data. We leave this exploring this as future work. Compared to traditional supervised learning, learning f∗H requires sample complexity scaling with the (Rademacher/Gaussian) complexity ofH (see e.g. [8, 52]), which is very large for complicated models such as deep networks. Thus SSL can significantly reduce the labeled sample complexity down from this complexity measure ofH to Õ(km), demonstrating the power of predicting what you already know using unlabeled data. In Section H, we consider a similar result for classification.

heading: 5 Example: Topic Modeling
text: In this section, we will demonstrate how our framework can be instantiated for standard data model like topic modeling. Topic modeling for text that has a rich literature [47, 31, 9, 4, 3] and is used for analyzing and designing algorithms for information retrieval, dimensionality reduction and data analysis for large text corpora. We describe the basic setup below, followed by how our results for reconstruction-based SSL can be instantiated to learn such models.
For a set S, let ∆S denote the set of all distributions on S. In the topic modeling framework, generation of a text document with a vocabulary set [V ] = {1, . . . , V } is governed by certain latent topics from the set [k], where k is the total number of topics. Each topic i ∈ [k] is associated with a distribution over the vocabulary [V ] that is denoted by vector Ai ∈ ∆[V ]; stack these vectors into the columns of a matrix A ∈ RV×k. A document X = (x1, . . . , xn) ∈ [V ]N of length N is then sampled from a mixture of the k topics µ ∈ ∆[k]. The generative process is described below:
1. Sample a topic mixture µ ∼ τ . τ is some underlying distribution over ∆k, i.e. τ ∈ ∆∆[k] 2. For each i ∈ [N ], sample a topic ti ∼ µ and sample a word xi ∼ Ati from the topic
For the reconstruction SSL task, we evenly split the document as X = (X̄1, X̄2), where X̄1 and X̄2 denote the first and second halves of the document; note that X̄1, X̄2 ∈ [V ]N/2. We let X1 and X2 be the multiset of words in the two halves by using the normalized bag-of-words representation, i.e. Xi = 2N bag-of-words(X̄i) ∈ R
V , i ∈ {1, 2}7. The downstream task is chosen to be a linear function of the topic posterior distribution µ for a given document X , i.e. Y = w⊤E[µ|X] + N , where N is 0 mean and σ2-subgaussian. The error of a predictor f : [V ]N → R is measured as Eµ,X [( f(X)− µ⊤w )2] , the optimal predictor being f∗(X) = E [Y | X].
A crucial property of topic model described above is that words in the document are sampled independently given the topic mixture µ, thus giving us the property: X1 ⊥ X2 | µ. Although the cardinality of µ ∈ ∆[k] (that implicitly shows up in Theorem 4.2) is infinite, we can still show the benefit of SSL using our theoretical framework. We will show appropriate bounds for ϵCI and β, that show up in Theorem 4.2, using the topic model generative process. We make the following standard assumptions about the topic modeling distribution, motivated by prior work [4, 3].
Assumption 5.1. Let A ∈ RV×k be the word-topic matrix and Γ = Eµ∼τ [ µµ⊤ ] be the topic covariance matrix, then the following hold
• (Anchor word) The word-topic matrix A is p-separable for p > 0, i.e. for every topic i ∈ [k] there is a word j such that Ai(j) ≥ p and Ai′(j) = 0 when i′ ̸= i • Γ is full rank, so condition number κ = λmax(Γ)λmin(Γ) <∞
7We only need X2 to be the bag-of-word representation, X1 can be an ordered sentence.
Theorem 5.1. Let ϵCI be the definition (2) from Definition 4.1 and β as defined in Assumption 4.1 and suppose the topic model satisfies Assumption 5.1, then there exists a latent variable Ȳ ∈ Ȳ such that the following hold
1. Ȳ takes k distinct values, i.e. |Ȳ| = k 2. X1 and X1 are uncorrelated given Ȳ , which implies ϵCI = 0. 3. E[Y |X1] is a linear function of E[Ȳ |X1] 4. β−1 ≤ κ∥w∥2/λmin(A) ≤ κ∥w∥2/p
The proof for this is presented in Section E.1. Note that the p-separability is not necessarily needed, and the bound with λmin(A) can be invoked instead. Thus the upper bound from Theorem 4.2) will look like Õ ( σ2 kn2 + ϵ 2 pre κ∥w∥2 p ) , thus requiring onlyO(k) samples for the downstream task.

heading: 6 Experiments
text: In this section, we empirically verify our claim that SSL performs well when ACI is satisfied. More details for experiments can be found in Section J, including experiments in the text domain.
Simulations. With synthetic data, we verify how excess risk (ER) scales with the cardinality/feature dimension of Y (k), and ACI (ϵCI in Definition 4.1). We consider a mixture of Gaussian data and conduct experiments with both linear function space (H1 with ϕ1 as identity map) and universal function space Hu. We sample the label Y uniformly from {1, ..., k}. For i-th class, the centers µ1i ∈ Rd1 and µ2i ∈ Rd2 are uniformly sampled from [0, 10). Given Y = i, α ∈ [0, 1], let X1 ∼ N (µ1i, I), X̂2 ∼ N (µ2i, I), and X2 = (1 − α)X̂2 + αX1. Therefore α is a correlation coefficient: α = 0 ensures X2 being CI with X1 given Y and when α = 1, X2 fully depends on X1. (if d1 ̸= d2, we append zeros or truncate to fit accordingly). We first conduct experiments with linear function class. We learn a linear representation ψ with n1 samples and the linear prediction of Y from ψ with n2 samples. We set d1 = 50, d2 = 40, n1 = 4000, n2 = 1000 and ER is measured with Mean Squared Error (MSE). As shown in Figure 1(a)(b), the MSE of learning with ψ(X1) scales linearly with k as indicated in Theorem 3.5, and scales linearly with ϵCI associated with linear function class as indicated in Theorem 4.2. Next we move on to general function class, i.e., ψ∗ = E[Y |X1] with a closed form solution (see example 3.1). We use the same parameter settings as above. For baseline method, we use kernel linear regression to predict Y using X1 (we use RBF kernel which also has universal approximation power). As shown in Figure 1(c)(d), the phenomenon is the same as what we observe in the linear function class setting, and hence they respectively verify Theorem 3.2 and Theorem 4.2 withHu.
Computer Vision Task. We verify if learning from ψ is more effective than learning directly from X1, in a realistic setting (without enforcing conditional independence). Specifically, we test on the Yearbook dataset [23], and try to predict the date when the portraits are taken (denoted as YD), which ranges from 1905 to 2013. We resize all the portraits to be 128 by 128. We crop out the center 64 by 64 pixels (the face), and treat it as X2, and treat the outer rim as X1 as shown in Figure 2. Our task is to predict YD, which is the year when the portraits are taken, and the year ranges from 1905 to 2013. For ψ, we learn X2 from X1 with standard image inpainting techniques [48], and full set of training data (without labels). After that we fix the learned ψ and learn a linear model to predict YD from ψ using a smaller set of data (with labels). Besides linear model on X1, another strong baseline that
we compare with is using ResNet18 [28] to predict YD from X1. With the full set of training data, this model is able to achieve a Mean Absolute Difference of 6.89, close to what state-of-the-art can achieve [23]. ResNet18 has similar amount of parameters as our generator, and hence roughly in the same function class. We show the MSE result as in Figure 2. Learning from ψ is more effective than learning from X1 or X2 directly, with linear model as well as with ResNet18. Practitioner usually fine-tune ψ with the downstream task, which leads to more competitive performance [48].

heading: 7 Conclusion
text: In this work we theoretically quantify how an approximate conditional independence assumption that connects pretext and downstream task data distributions can give sample complexity benefits of self-supervised learning on downstream tasks. Our theoretical findings are also supported by experiments on simulated data and also on real CV and NLP tasks. We would like to note that approximate CI is only a sufficient condition for a useful pretext task. We leave it for future work to investigate other mechanisms by which pretext tasks help with downstream tasks.

heading: Acknowledgment
text: JDL acknowledges support of the ARO under MURI Award W911NF-11-1-0304, the Sloan Research Fellowship, NSF CCF 2002272, NSF IIS 2107304, and an ONR Young Investigator Award. QL was supported by NSF #2030859 and the Computing Research Association for the CIFellows Project. NS is supported by NSF, ONR, Simons Foundation, DARPA and SRC. JZ acknowledges sponsorship of Army Research Office and support of Cooperative Agreement Number W911NF-19-2-0333 and NSF grant 2019844. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Army Research Office or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.

Review summaries:
This paper proposes a mechanism based on approximate conditional independence (ACI) to explain why solving pretext tasks created from known information can learn representations that provably reduce downstream sample complexity, as a sufficient condition. In specific, they measure the downstream performance using the approximation error and estimation error, and establish their initial results under a strict condition -- conditional independence (CI) and linear function space in Section 3, then extend it to ACI and arbitrary function space in Section 4, resulting in the main contribution Theorem 4.2 that clearly quantifies the generalization error. The theorem is also verified in Section 5 using simulations on NLP tasks. 

This paper attempts to understand why self-supervised learning works in the following sense: will the sample complexity for a downstream task be decreased (compared to the standard supervised learning without pretraining) if it is pre-trained according to some related auxiliary task? The relation between tasks is formulated as the approximate conditional independence of the dependent variables. The main theoretical results show that the sample complexity (compared to supervised learning) can be decreased under Assumption 3.5, 4.1, and 4.2. 

- The paper theoretically analyzes self-supervision under some simplifying assumptions. They begin with a conditional independent setting, where X1 is conditionally independent of X2 given Y. - The high-level idea is to first pre-train a model that predicts X2 from X1. This can be done on unlabeled data. Then they can fit a model on top of representations learned by this model, to predict Y. 1. Ground truth self-supervision: They show that if we learn E[X2 | X1], the ground truth self-supervision function, then under mild invertibility conditions, E[Y | X] is a linear function of E[X2 | X1]. Further, E[X2 | X1]'s effective dimension is the same as Y, so this linear function can be fit using very few examples. In contrast, the function from X1 or X2 to Y can be very complicated / non-linear and require many examples to fit. 2. Linear self-supervision: They then examine what happens if we fit a linear function (on feature space) to predict X2 from X1. This can learn a useful lower dimensional embedding, and then we train a linear classifier to predict Y on top of this. They show this is more sample efficient than just training a linear classifier from X1 to Y. 3. Relaxing conditional independence: They quantitatively examine how this degrades when we don't exactly have conditional independence. Their bounds will still be good if things are "almost" conditionally independent. 

This paper aims to theoretically explain and investigate the benefit of reconstruction-based self-supervised learning. This paper studies a specific problem with two input features and one downstream label and shows how the specific connection between them can guarantee successful self-supervised learning.  In particular, this paper quantifies how the approximation independence between two input features can be extracted to solve the downstream tasks. Results show that self-supervised learning can successfully learn the ground-truth with much smaller labeled sample complexity. 

This paper studies whether reconstruction-based pretext task can provide representations helpful for downstream tasks, which is a question of good interests in the community.  The authors give an affirmative answer that when (approximate) conditional independence holds, reconstruction as a pretext task provides representations such that 1) an additional linear layer is sufficient to ensure small approximation error on classification, and that 2) helps to reduce the sample complexity of the classification task from complexity of the function class (without pretraining) to linear in data dimension.  These theoretical results are empirically verified on synthetic, vision (inpainting), and NLP (sentiment classification) tasks.

This paper suggests a theoretical model to capture why reconstruction-based self-supervised learning leads to good representations, and improved downstream performance. The suggested model is that the pretext task "factors through" the supervised labels, i.e. that X_1 -- Y -- X_2 is Markov, where X_1-->X_2 is the pretext task, and X_1 --> Y is the supervised task. In such a setting, it is shown that the optimal pretext function has certain good properties: it linearly separates the supervised task, and is "low rank", so can be used for efficient downstream learning.