Provable Guarantees for Self-Supervised Deep Learning with Spectral Contrastive Loss

1 Introduction . Recent empirical breakthroughs have demonstrated the effectiveness of self-supervised learning , which trains representations on unlabeled data with surrogate losses and self-defined supervision signals [ 4 , 6 , 10 , 14 , 23 , 24 , 35 , 38 , 41 , 42 , 50–52 ] . Self-supervision signals in computer vision are often defined by using data augmentation to produce multiple views of the same image . For example , the recent contrastive learning objectives [ 3 , 12 , 13 , 15 , 22 ] encourage closer representations for augmentations ( views ) of the same natural data than for randomly sampled pairs of data . Despite the empirical successes , there is a limited theoretical understanding of why self-supervised losses learn representations that can be adapted to downstream tasks , for example , using linear heads . Recent mathematical analyses by Arora et al . [ 3 ] , Lee et al . [ 28 ] , Tosh et al . [ 44 , 45 ] provide guarantees under the assumption that two views are somewhat independent conditioned on the label . However , the pair of augmented examples used in practical algorithms usually exhibit a strong correlation , even conditioned on the label . For instance , two augmentations of the same dog image 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . share much more similarity than augmentations of two different random dog images . Thus the existing theory does not explain the practical success of self-supervised learning . This paper presents a theoretical framework for self-supervised learning without requiring conditional independence . We design a principled , practical loss function for learning neural net representations that resembles state-of-the-art contrastive learning methods . We prove that , under a simple and realistic data assumption , linear classification using representations learned on a polynomial number of unlabeled data samples can recover the ground-truth labels of the data with high accuracy . The fundamental data property that we leverage is a notion of continuity of the population data within the same class . Though a random pair of examples from the same class can be far apart , the pair is often connected by ( many ) sequences of examples , where consecutive examples in the sequences are close neighbors within the same class . This property is more salient when the neighborhood of an example includes many different types of augmentations . Prior work [ 49 ] empirically demonstrates this type of connectivity property and uses it in the analysis of pseudolabeling algorithms . More formally , we define the population augmentation graph , whose vertices are all the augmented data in the population distribution , which can be an exponentially large or infinite set . Two vertices are connected with an edge if they are augmentations of the same natural example . Our main assumption is that for some proper m 2 Z+ , the sparsest m-partition ( Definition 3.4 ) is large . This intuitively states that we can ’ t split the augmentation graph into too many disconnected sub-graphs by only removing a sparse set of edges . This assumption can be seen as a graph-theoretic version of the continuity assumption on population data . We also assume that there are very few edges across different ground-truth classes ( Assumption 3.5 ) . Figure 1 ( left ) illustrates a realistic scenario where dog and cat are the ground-truth categories , between which edges are very rare . Each breed forms a sub-graph that has sufficient inner connectivity and thus can not be further partitioned . Our assumption fundamentally does not require conditional independence and can allow disconnected sub-graphs within a class . The classes in the downstream task can be also somewhat flexible as long as they are disconnected in the augmentation graph . For example , when the augmentation graph consists of m disconnected sub-graphs corresponding to fine-grained classes , our assumptions allow the downstream task to have any r m coarse-grained classes containing these fine-grained classes as a sub-partition . Prior work [ 49 ] on pseudolabeling algorithms essentially requires an exact alignment between sub-graphs and downstream classes ( i.e. , r = m ) . They face this limitation because their analysis requires fitting discrete pseudolabels on the unlabeled data . We avoid this difficulty because we consider directly learning continuous representations on the unlabeled data . We apply spectral decomposition—a classical approach for graph partitioning , also known as spectral clustering [ 37 , 39 ] in machine learning—to the adjacency matrix defined on the population augmentation graph . We form a matrix where the top-k eigenvectors are the columns and interpret each row of the matrix as the representation ( in Rk ) of an example . Somewhat surprisingly , we show that this feature extractor can be also recovered ( up to some linear transformation ) by minimizing the following population objective which is similar to the standard contrastive loss ( Section 3.2 ) : L ( f ) = 2 · Ex , x+ ⇥ f ( x ) > f ( x+ ) ⇤ + Ex , x0 ⇥ f ( x ) > f ( x0 ) 2 ⇤ , where ( x , x+ ) is a pair of augmentations of the same data , ( x , x0 ) is a pair of independently random augmented data , and f is a parameterized function from augmented data to Rk . Figure 1 ( right ) illustrates the relationship between the eigenvector matrix and the learned representations . We call this loss the population spectral contrastive loss . We analyze the linear classification performance of the representations learned by minimizing the population spectral contrastive loss . Our main result ( Theorem 3.7 ) shows that when the representation dimension exceeds the maximum number of disconnected sub-graphs , linear classification with learned representations is guaranteed to have a small error . Our theorem reveals a trend that a larger representation dimension is needed when there are a larger number of disconnected sub-graphs . Our analysis relies on novel techniques tailored to linear probe performance , which have not been studied in the spectral graph theory community to the best of our knowledge . The spectral contrastive loss also works on empirical data . Since our approach optimizes parametric loss functions , guarantees involving the population loss can be converted to finite sample results using off-the-shelf generalization bounds . The sample complexity is polynomial in the Rademacher complexity of the model family and other relevant parameters ( Theorem 4.1 and Theorem 4.2 ) . In summary , our main theoretical contributions are : 1 ) we propose a simple contrastive loss motivated by spectral decomposition of the population data graph , 2 ) under simple and realistic assumptions , we provide downstream classification guarantees for the representation learned by minimizing this loss on population data , and 3 ) our analysis is easily applicable to deep networks with polynomial unlabeled samples via off-the-shelf generalization bounds . In addition , we implement and test the proposed spectral contrastive loss on standard vision benchmark datasets . We demonstrate that the features learned by our algorithm can match or outperform several strong baselines [ 12 , 14 , 15 , 21 ] when evaluated using a linear probe . 2 Additional related works . Empirical works on self-supervised learning . Self-supervised learning algorithms have been shown to successfully learn representations that benefit downstream tasks [ 4 , 6 , 10 , 12 , 13 , 15 , 22– 24 , 35 , 38 , 41 , 42 , 50–52 ] . Many recent self-supervised learning algorithms learn features with siamese networks [ 8 ] , where two neural networks of shared weights are applied to pairs of augmented data . Introducing asymmetry to siamese networks either with a momentum encoder like BYOL [ 21 ] or by stopping gradient propagation for one branch of the siamese network like SimSiam [ 14 ] has been shown to effectively avoid collapsing . Contrastive methods [ 12 , 15 , 22 ] minimize the InfoNCE loss [ 38 ] , where two views of the same data are attracted while views from different data are repulsed . Theoretical works on self-supervised learning . In addition to works [ 3 , 28 , 44 , 45 ] discussed in the introduction , several other works [ 5 , 43 , 47 , 48 ] also theoretically study self-supervised learning . The work Tsai et al . [ 47 ] prove that self-supervised learning methods can extract task-relevant information and discard task-irrelevant information , but lacks guarantees for solving downstream tasks efficiently with simple ( e.g. , linear ) models . Tian et al . [ 43 ] study why non-contrastive selfsupervised learning methods can avoid feature collapse . Cai et al . [ 9 ] analyze domain adaptation algorithms for subpopulation shift with a similar expansion condition as [ 49 ] while also allowing disconnected parts within each class , but require access to ground-truth labels during training . In contrast , our algorithm doesn ’ t need labels during pre-training . 3 Spectral contrastive learning on population data . In this section , we introduce our theoretical framework , the spectral contrastive loss , and the main analysis of the performance of the representations learned on population data . We use X to denote the set of all natural data ( raw inputs without augmentation ) . We assume that each x̄ 2 X belongs to one of r classes , and let y : X ! [ r ] denote the ground-truth ( deterministic ) labeling function . Let PX be the population distribution over X from which we draw training data and test our final performance . For the ease of exposition , we assume X to be a finite but exponentially large set ( e.g. , all real vectors in Rd with bounded precision ) .1 We next formulate data augmentations . Given a natural data sample x̄ 2 X , we use A ( ·|x̄ ) to denote the distribution of its augmentations . For instance , when x̄ represents an image , A ( ·|x̄ ) can be the distribution of common augmentations [ 12 ] that includes Gaussian blur , color distortion and random cropping . We use X to denote the set of all augmented data , which is the union of supports of all A ( ·|x̄ ) for x̄ 2 X . As with X , we also assume that X is a finite but exponentially large set , and denote N = |X | . We will learn an embedding function f : X ! Rk , and then evaluate its quality by the minimum error achieved with a linear probe . Concretely , a linear classifier has weights B 2 Rk⇥r and predicts gf , B ( x ) = argmaxi2 [ r ] ( f ( x ) > B ) i for an augmented data x ( argmax breaks tie arbitrarily ) . Then , given a natural data sample x̄ , we ensemble the predictions on augmented data and predict : ḡf , B ( x̄ ) : = argmax i2 [ r ] Pr x⇠A ( ·|x̄ ) [ gf , B ( x ) = i ] . Define the linear probe error as the error of the best possible linear classifier on the representations : E ( f ) : = min B2Rk⇥r Pr x̄⇠PX [ y ( x̄ ) 6= ḡf , B ( x̄ ) ] ( 1 ) 3.1 Augmentation graph and spectral decomposition . Our approach is based on the central concept of population augmentation graph , denoted by G ( X , w ) , where the vertex set is all augmentation data X and w denotes the edge weights defined below . For any two augmented data x , x0 2 X , define the weight wxx0 as the marginal probability of generating the pair x and x0 from a random natural data x̄ ⇠ PX : wxx0 : = Ex̄⇠PX [ A ( x|x̄ ) A ( x 0|x̄ ) ] ( 2 ) Therefore , the weights sum to 1 because the total probability mass is 1 : P x , x02X wxx0 = 1 . The relative magnitude intuitively captures the closeness between x and x0 with respect to the augmentation transformation . For most of the unrelated x and x0 , the value wxx0 will be significantly smaller than the average value . For example , when x and x0 are random croppings of a cat and a dog respectively , wxx0 will be essentially zero because no natural data can be augmented into both x and x0 . On the other hand , when x and x0 are very close in ` 2-distance or very close in ` 2-distance up to color distortion , wxx0 is nonzero because they may be augmentations of the same image with Gaussian blur and color distortion . We say that x and x0 are connected with an edge if wxx0 > 0 . See Figure 1 ( left ) for more illustrations . Given the structure of the population augmentation graph , we apply spectral decomposition to the population graph to construct principled embeddings . The eigenvalue problems are closely related to graph partitioning as shown in spectral graph theory [ 17 ] for both worst-case graphs [ 11 , 25 , 29 , 33 ] and random graphs [ 1 , 30 , 34 ] . In machine learning , spectral clustering [ 37 , 39 ] is a classical algorithm that learns embeddings by eigendecomposition on an empirical distance graph and invoking k-means on the embeddings . We will apply eigendecomposition to the population augmentation graph ( and then later use linear probe for classification ) . Let wx = P x02X wxx0 be the total weights associated to x , which is often viewed as an analog of the degree of x in weighted graph . A central object in spectral graph theory is the so-called normalized adjacency matrix : A : = D 1/2AD 1/2 ( 3 ) 1This allows us to use sums instead of integrals and avoid non-essential nuances related to calculus . where A 2 RN⇥N is adjacency matrix with entires Axx0 = wxx0 and D 2 RN⇥N is a diagonal matrix with Dxx = wx.2 Standard spectral graph theory approaches produce vertex embeddings as follows . Let 1 , 2 , · · · , k be the k largest eigenvalues of A , and v1 , v2 , · · · , vk be the corresponding unit-norm eigenvectors . Let F ? = [ v1 , v2 , · · · , vk ] 2 RN⇥k be the matrix that collects these eigenvectors in columns , and we refer to it as the eigenvector matrix . Let u⇤x 2 Rk be the x-th row of the matrix F ? . It turns out that u ⇤ x ’ s can serve as desirable embeddings of x ’ s because they exhibit clustering structure in Euclidean space that resembles the clustering structure of the graph G ( X , w ) .

Summary: This paper theoretically analyzes contrastive learning relying on a concept of population augmentation graph. In the graph, two augmented samples are connected if they can be augmented from the same original instance. Based on this setup, - This paper shows that spectral decomposition can be reformulated as a contrastive learning objective. Furthermore, this paper shows that minimizing the objective can guarantee a high linear probe accuracy under some mild assumptions. - This paper shows the guarantee can be extended to the finite-sample regime and learning the linear probe with labeled data. - This paper demonstrates the efficacy of the proposed spectral contrastive objective empirically in the standard setup: the proposed objective achieves similar performance as other state-of-the-art SSL methods.