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

1 Introduction . 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 . 1.1 Related work . 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 can not 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 can not 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 .

Summary: 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.