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 derives generalization bounds for learning using a contrastive objective under realistic assumptions on the positive pair generation process. The approach taken is to consider the so-called augmentation graph - the set of nodes comprises all augmented views of population distribution, with edges connecting nodes that correspond to different views of the same input datapoint. The core assumption made is that this graph cannot be split into a large number of disconnected subgraphs. This set-up aligns well with the intuition that in order to generalize, the contrastive notion of “similarity” must extent beyond the purely single-instance-level, and must somehow connect distinct inputs points.  My view is that this work is the most significant advance in the statistical learning theory of contrastive learning since the paper “A theoretical analysis of contrastive unsupervised representation learning” by Arora et al. This paper under review addresses the primary weakness in the work by Arora et al., which used an unrealistic assumption on positive pair generation. I will strongly advocate for acceptance of this paper.