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