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

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 Wu et al. 2018, Oord et al. 2018, Hjelm et al. 2018, Ye et al. 2019, Henaff 2020, Bachman et al. 2019, Tian et al. 2019, Misra and Maaten 2020, Caron et al. 2020, Zbontar et al. 2021, Bardes et al. 2021, Tian et al. 2020a, Chen and He 2020. 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 Arora et al. 2019, Chen et al. 2020a, Chen et al. 2020b, He et al. 2020, Chen et al. 2020c encourage closer representations for augmentations/views of the same natural datapoint 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 for contrastive learning by Arora et al. 2019, Tosh et al. 2020, Tosh et al. 2021 provide guarantees under the assumption that two views are somewhat conditionally independent given the label or a hidden variable. However, in practical algorithms for computer vision applications, the two views are augmentations of a natural image and usually exhibit a strong correlation that is difficult to be de-correlated by conditioning. They are not independent conditioned on the label, and we are only aware that they are conditionally independent given the natural image, which is too complex to serve as a hidden variable with which prior works can be meaningfully applied. Thus the existing theory does not appear to 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 images from the same class can be far apart, the pair is often connected by (many) sequences of natural images, where consecutive images in the sequences are close neighbors within the same class. As shown in Figure 1 (images on the left top part), two very different French bulldogs can be connected by a sequence of French bulldogs (which may not be in the training set but are in the support of the population distribution). Prior work Wei et al. 2020 empirically demonstrates this type of connectivity property and uses it in the analysis of pseudolabeling algorithms. This property is more salient when the neighborhood of an example includes many different types of augmentations.

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\in\mathcal{Z}^{+}$, we cannot partition the graph into $m+1$ sub-graphs between which there are few connections (Assumption 3.5). In other words, this intuitively states that there are at most $m$ clusters in the population augmentation graph. This assumption can be seen as a graph-theoretic version of the continuity assumption on the population distribution. We also assume that there are very few edges across different ground-truth classes (Assumption 3.6). 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 cannot be further partitioned.

Our assumption fundamentally does not require independence of the two views (the positive pairs) conditioned on the class 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\leq m$ coarse-grained classes containing these fine-grained classes as a sub-partition. Prior work Wei et al. 2020 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 analyze the linear classification performance of the representations learned by minimizing the population spectral contrastive loss. Our main result (Theorem 3.8) 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 end-to-end result (Theorem 4.3) shows that the number of unlabeled examples required is polynomial in the Rademacher complexity of the model family and other relevant parameters, whereas the number of downstream labeled examples only needs to be linear in the representation dimension (which needs to be linear in the number of clusters in the graph). This demonstrates that contrastive learning reduces the amount of labeled examples needed.

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. Our theoretical framework can be viewed as containing two stages: we first analyze the population loss and the representation that minimizes it (Section 3), then study the empirical loss where the representation is learned with a neural network with bounded capacity (Section 4).

In addition, we implement and test the proposed spectral contrastive loss on standard vision benchmark datasets. Our algorithm is simple and doesn’t rely on tricks such as stop-gradient which is essential to SimSiam Chen and He 2020. We demonstrate that the features learned by our algorithm can match or outperform several strong baselines (Chen et al. 2020a, Chen et al. 2020c, Chen and He 2020, Grill et al. 2020) when evaluated using a linear probe.

Additional related works

Empirical works on self-supervised learning. Self-supervised learning algorithms have been shown to successfully learn representations that benefit downstream tasks Wu et al. 2018, Oord et al. 2018, Hjelm et al. 2018, Ye et al. 2019, Henaff 2020, Bachman et al. 2019, Tian et al. 2019, Misra and Maaten 2020, Chen et al. 2020c, Chen et al. 2020a, He et al. 2020, Chen et al. 2020b, Caron et al. 2020, Zbontar et al. 2021, Bardes et al. 2021, Tian et al. 2020a, Xie et al. 2019. Many recent self-supervised learning algorithms learn features with siamese networks Bromley et al. 1993, 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 Grill et al. 2020 or by stopping gradient propagation for one branch of the siamese network like SimSiam Chen and He 2020 has been shown to effectively avoid collapsing. Contrastive methods Chen et al. 2020a, He et al. 2020, Chen et al. 2020c minimize the InfoNCE loss Oord et al. 2018, where two views of the same data are attracted while views from different data are repulsed.

Theoretical works on self-supervised learning. As briefly discussed in the introduction, several theoretical works have studied self-supervised learning. Arora et al. 2019 provide guarantees for representations learned by contrastive learning on downstream linear classification tasks under the assumption that the positive pairs are conditionally independent given the class label. Theorem 3.3 and Theorem 3.7 of the work of Lee et al. 2020 show that, under conditional independence given the label and/or additional latent variables, representations learned by reconstruction-based self-supervised learning algorithms can achieve small errors in the downstream linear classification task. Lee et al. 2020 generalizes it to approximate conditional independence for Gaussian data and Theorem 4.5 further weakens the assumptions significantly. Tosh et al. 2020 show that contrastive learning representations can linearly recover any continuous functions of the underlying topic posterior under a topic modeling assumption (which also requires conditional independence of the positive pair given the hidden variable). More recently, Theorem 11 of the work of Tosh et al. 2021 provide novel guarantees for contrastive learning under the assumption that there exists a hidden variable $h$ such that the positive pair $(x,x^{+})$ are conditionally independent given $h$ and the random variable $p(x|h)p(x^{+}|h)/p(x)p(x^{+})$ has a small variance. However, in practical algorithms for computer vision applications, the two views are two augmentations and thus they are highly correlated. They might be only independent when conditioned on very complex hidden variables such as the original natural image, which might be too complex for the previous results to be meaningfully applied.

We can also compare the assumptions and results on a concrete generative model for the data, our Example 3.10 in Section 3.4, where the data are generated by a mixture of Gaussian or a mixture of manifolds, the label is the index of the mixture, and the augmentations are small Gaussian blurring (i.e., adding Gaussian noise). In this case, the positive pairs $(x,x^{+})$ are two points that are very close to each other. To the best of our knowledge, applying Theorem 11 of Tosh et al. 2021 to this case with $h=\bar{x}$ (the natural datapoint) would result in requiring a large (if not infinite) representation dimension. Because $x^{+}$ and $x$ are very close, the reconstruction-based algorithms in Lee et al. 2020, when used to predict $x^{+}$ from $x$, will not be able to produce good representations as well. On a technical level, Example 3.10 does not satisfy the requirement regarding the $\beta$ quantity in Assumption 4.1 of Lee et al. 2020, if $(X_{1},X_{2})$ in that paper is equal to $(x,x+)$ here—it requires the label to be correlated with the raw input $x$, which is not necessarily true in Example 3.10. This can likely be addressed by using a different $X_{2}$.

On a technical level, to relate prior works’ assumptions to ours, we can consider an almost equivalent version of our assumption (although our proofs do not directly rely on or relate to the discussion below). Let $(x,x^{+})$ be a positive pair and let $p(\cdot|x)$ be the conditional distribution of $x^{+}$ given $x$. Starting from $x_{0}$, let us consider a hypothetical Markov chain $x_{0},\dots,x_{T},\cdots$ where $x_{t}$ is drawn from $p(\cdot|x_{t-1})$. Our assumption essentially means that this hypothetical Markov chain of sampling neighbors will mix within the same class earlier than it mixes across the entire population (which might not be possible or takes exponential time). More concretely, the assumption that $\rho_{\lfloor k/2\rfloor}$ is large compared to $\alpha$ in Theorem 3.8 is roughly equivalent to the existence of a (potentially large) $T$ such that $x_{0}$ and $x_{T}$ are still likely to have the same label, but are sufficiently independent conditioned on this label or some hidden variable. Roughly speaking, prior works Arora et al. 2019, Tosh et al. 2020, Tosh et al. 2021 assume probabilistic structure about $x_{0}$ and $x_{1}$ (instead of $x_{0}$ and $x_{T}$), e.g., Arora et al. 2019 and Theorem 11 of Tosh et al. 2021 assume that $x_{0}$ and $x_{1}$ are independent conditioned on the label and/or a hidden variable. Similar Markov chains on augmentated data have also been used in previous work Dao et al. 2019 to study properties of data augmentation.

Several other works (Tsai et al. 2020, Wang and Isola 2020, Tian et al. 2020b, Bansal et al. 2020, Mitrovic et al. 2020) also theoretically study self-supervised learning. The work Tsai et al. 2020 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. 2020b study why non-contrastive self-supervised learning methods can avoid feature collapse. Zimmermann et al. 2021 prove that for a specific data generating process, contrastvie learning can learn representations that recover the latent variable. Cai et al. 2021 analyze domain adaptation algorithms for subpopulation shift with a similar expansion condition as Wei et al. 2020 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.

Co-training and multi-view learning are related settings which leverage two distinct “views” (i.e., feature subsets) of the data (Blum and Mitchell 1998, Dasgupta et al. 2002, Balcan et al. 2005). The original co-training algorithms (Blum and Mitchell 1998, Dasgupta et al. 2002) assume that the two views are independent conditioned on the true label and leverage this independence to obtain accurate pseudolabels for the unlabeled data. Balcan et al. 2005 relax the requirement on independent views of co-training, by using an “expansion” assumption, which is closely related to our assumption that $\rho_{\lfloor k/2\rfloor}$ is not too small in Theorem 3.8. Besides recent works (e.g., the work of Tosh et al. 2021), most co-training or multi-view learning algorithms are quite different from the modern contrastive learning algorithms which use neural network parameterization for vision applications.

Our analysis relies on the normalized adjacency matrix (see Section 3.1), which is closely related to the graph Laplacian regularization that has been studied in the setting of semi-supervised learning Zhu et al. 2003, Nadler et al. 2009. In their works, the Laplacian matrix is used to define a regularization term that smooths the predictions on unlabeled data. This regularizer is further added to the supervised loss on labeled data during training. In contrast, we use the normalized adjacency matrix to define the unsupervised training objective in this paper.

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 next formulate data augmentations. Given a natural data sample $\bar{x}\in\overline{\mathcal{X}}$, we use $\mathcal{A}(\cdot|\bar{x})$ to denote the distribution of its augmentations. For instance, when $\bar{x}$ represents an image, $\mathcal{A}(\cdot|\bar{x})$ can be the distribution of common augmentations Chen et al. 2020a that includes Gaussian blur, color distortion and random cropping. We use $\mathcal{X}$ to denote the set of all augmented data, which is the union of supports of all $\mathcal{A}(\cdot|\bar{x})$ for $\bar{x}\in\overline{\mathcal{X}}$. As with $\overline{\mathcal{X}}$, we also assume that $\mathcal{X}$ is a finite but exponentially large set, and denote $N=|\mathcal{X}|$. None of the bounds will depend on $N$ — it is only defined and assumed to be finite for the ease of exposition.

We denote the error of the representation and the linear head as:

Define the linear probe error as the error of the best possible linear classifier on the representations:

Our approach is based on the central concept of population augmentation graph, denoted by $G(\mathcal{X},w)$, where the vertex set is all augmentation data $\mathcal{X}$ and $w$ denotes the edge weights defined below. For any two augmented data $x,x^{\prime}\in\mathcal{X}$, define the weight $w_{xx^{\prime}}$ as the marginal probability of generating the pair $x$ and $x^{\prime}$ from a random natural data $\bar{x}\sim\mathcal{P}_{\overline{\mathcal{X}}}$:

We emphasize that we only work with the population graph rather than the empirical graph (i.e., the corresponding graph constructed with the empirical dataset as the vertex set). The population graph is very sparse but not empty—many similar images exist in the population. In contrast, the empirical graph would be nearly empty, since two images in the empirical dataset almost never share the same augmentation image. Our analysis will apply to minimizing contrastive loss on an empirical dataset (see Section 4), but not via analyzing the property of the empirical graph. Instead, we will show that contrastive learning on empirical data with parametrized models is similar to decomposing the population graph (see technical discussions in Section 5). This is a key difference between our work and classical spectral clustering work—we only require properties of the population graph rather than the empirical graph.

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 Chung and Graham 1997 for both worst-case graphs Cheeger 1969, Kannan et al. 2004, Louis et al. 2011, Lee et al. 2014 and random graphs McSherry 2001, Lei et al. 2015, Abbe 2017. In machine learning, spectral clustering Ng et al. 2001, Shi and Malik 2000 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 $w_{x}=\sum_{x^{\prime}\in\mathcal{X}}w_{xx^{\prime}}$ 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:

2 From spectral decomposition to spectral contrastive learning

The embeddings $u_{x}^{*}$ obtained by eigendecomposition are nonparametric—a $k$-dimensional parameter is needed for every $x$—and therefore cannot be learned with a realistic amount of data. The embedding matrix $F^{*}$ cannot be even stored efficiently. Therefore, we will instead parameterize the rows of the eigenvector matrix $F^{*}$ as a neural net function, and assume embeddings $u_{x}^{*}$ can be represented by $f(x)$ for some $f\in\mathcal{F}$, where $\mathcal{F}$ is the hypothesis class containing neural networks. As we’ll show in Section 4, this allows us to leverage the extrapolation power of neural networks and learn the representation on a finite dataset.

Next, we design a proper loss function for the feature extractor $f$, such that minimizing this loss could recover $F^{*}$ up to some linear transformation. As we will show in Section 4, the resulting population loss function on $f$ also admits an unbiased estimator with finite training samples. Let $F$ be an embedding matrix with $u_{x}$ on the $x$-th row, we will first design a loss function of $F$ that can be decomposed into parts about individual rows of $F$.

We employ the following matrix factorization based formulation for eigenvectors. Consider the objective

where $\mathcal{E}{(F)}$ denotes the linear probe performance when the rows of $F$ are used as embeddings.

The main benefit of objective $\mathcal{L}_{\textup{mf}}({F})$ is that it’s based on the rows of $F$. Recall that vectors $u_{x}$ are the rows of $F$. Each entry of $FF^{\top}$ is of the form $u_{x}^{\top}u_{x^{\prime}}$, and thus $\mathcal{L}_{\textup{mf}}({F})$ can be decomposed into a sum of $N^{2}$ terms involving terms $u_{x}^{\top}u_{x^{\prime}}$. Interestingly, if we reparameterize each row $u_{x}$ by $w_{x}^{1/2}f(x)$, we obtain a very similar loss function for $f$ that resembles the contrastive learning loss used in practice (Chen et al. 2020a) as shown below in Lemma 3.2. See Figure 1 (right) for an illustration of the relationship between the eigenvector matrix and the representations learned by minimizing this loss.

We formally define the positive and negative pairs to introduce the loss. Let $\bar{x}\sim\mathcal{P}_{\overline{\mathcal{X}}}$ be a random natural datapoint and draw $x\sim\mathcal{A}(\cdot|\bar{x})$ and $x^{+}\sim\mathcal{A}(\cdot|\bar{x})$ independently to form a positive pair $(x,x^{+})$. Draw $\bar{x}^{\prime}\sim\mathcal{P}_{\overline{\mathcal{X}}}$ and $x^{-}\sim\mathcal{A}(\cdot|\bar{x}^{\prime})$ independently with $\bar{x},x,x^{+}$. We call $(x,x^{-})$ a negative pair. Though $x$ and $x^{-}$ are simply two independent draws, we call them negative pairs following the literature Arora et al. 2019.

Recall that $u_{x}$ is the $x$-th row of $F$. Let $u_{x}=w_{x}^{1/2}f(x)$ for some function $f$. Then, the loss function $\mathcal{L}_{\textup{mf}}({F})$ is equivalent to the following loss function for $f$, called spectral contrastive loss, up to an additive constant:

We can expand $\mathcal{L}_{\textup{mf}}({F})$ and obtain

Notice that the first term is a constant that only depends on the graph but not the variable $f$. By the definition of augmentation graph, $w_{xx^{\prime}}$ is the probability of a random positive pair being $(x,x^{\prime})$ while $w_{x}$ is the probability of a random augmented datapoint being $x$. We can hence rewrite the sum of last two terms in Equation (3.2) as Equation (3.2). ∎

We note that spectral contrastive loss is similar to many popular contrastive losses Oord et al. 2018, Chen et al. 2020a, Sohn 2016, Wu et al. 2018. For instance, the contrastive loss in SimCLR Chen et al. 2020a can be rewritten as (with simple algebraic manipulation)

Here $x$ and $x^{+}$ are a positive pair and $x_{1},\cdots,x_{n}$ are augmentations of other data. Spectral contrastive loss can be seen as removing $f(x)^{\top}f(x^{+})$ from the second term, and replacing the log sum of exponential terms with the average of the squares of $f(x)^{\top}f(x_{i})$. We will show in Section 6 that our loss has a similar empirical performance as SimCLR without requiring a large batch size.

3 Theoretical guarantees for spectral contrastive loss on population data

In this section, we introduce the main assumptions on the data and state our main theoretical guarantee for spectral contrastive learning on population data.

To formalize the idea that $G$ cannot be partitioned into too many disconnected sub-graphs, we introduce the notions of Dirichlet conductance and sparsest $m$-partition, which are standard in spectral graph theory. Dirichlet conductance represents the fraction of edges from $S$ to its complement:

For a graph $G=(\mathcal{X},w)$ and a subset $S\subseteq\mathcal{X}$, we define the Dirichlet conductance of $S$ as

Let $G=(\mathcal{X},w)$ be the augmentation graph. For an integer $i\in[2,|\mathcal{X}|]$, we define the sparsest $i$-partition as

where $S_{1},\cdots,S_{i}$ are non-empty sets that form a partition of $\mathcal{X}$.

We note that $\rho_{i}$ increases as $i$ increases. To see this, consider $3\leq i\leq|\mathcal{X}|$. Let $S_{1},\cdots,S_{i}$ be the partition of $\mathcal{X}$ that minimizes the RHS of Definition 3.4 Define set $S_{i-1}^{\prime}:=S_{i}\cup S_{i-1}$. It is easy to see that $\phi_{G}(S_{i-1}^{\prime})=\frac{\sum_{x\in S_{i-1}^{\prime},x^{\prime}\notin S_{i-1}^{\prime}}w_{xx^{\prime}}}{\sum_{x\in S_{i-1}^{\prime}}w_{x}}\leq\frac{\sum_{j=i-1}^{i}\sum_{x\in S_{j},x^{\prime}\notin S_{j}}w_{xx^{\prime}}}{\sum_{j=i-1}^{i}\sum_{x\in S_{j}}w_{x}}\leq\max\{\phi_{G}(S_{i-1}),\phi_{G}(S_{i})\}$. Notice that $S_{1},\cdots,S_{i-2},S_{i-1}^{\prime}$ are $i-1$ non-empty sets that form a partition of $\mathcal{X}$, by Definition 3.4 we have $\rho_{i-1}\leq\max\{\phi_{G}(S_{1}),\cdots,\phi_{G}(S_{i-2}),\phi_{G}(S_{i-1}^{\prime})\}\leq\max\{\phi_{G}(S_{1}),\cdots,\phi_{G}(S_{i})\}=\rho_{i}$. When $r$ is the number of underlying classes, we might expect $\rho_{r}\approx 0$ since the augmentations from different classes almost compose a disjoint $r$-way partition of $\mathcal{X}$. However, for $i>r$, we can expect $\rho_{i}$ to be much larger. For instance, in the extreme case when $i=|\mathcal{X}|=N$, every set $S_{j}$ is a singleton, which implies that $\rho_{N}=1$. More generally, as we will show later (Lemma 3.9), $\rho_{i}$ can be expected to be at least inverse polynomial in data dimension when $i$ is larger than the number of underlying semantic classes in the data.

We assume that $\rho_{m+1}\geq\rho$. A prototypical case would be that there are at most $m$ clusters in the population augmentation graph, and each of them cannot be broken into two subsets both with conductance less than $\rho$.

When there are $m$ clusters that have sufficient inner connections (corresponding to, e.g., $m$ semantically coherent subpopulations), we expect $\rho_{m+1}$ to be much larger than $\rho_{m}$ because any $m+1$ partition needs to break one sub-graph into two pieces and incur a large conductance. In other words, suppose the graph is consists of $m$ clusters, the quantity $\rho$ is characterizing the level of internal connection within each cluster. Furthermore, in many cases we expect $\rho_{m+1}$ to be inverse polynomial in dimension. In the running example of Section 3.1 (where augmentation is adding Gaussian noise), $\rho$ is related to the Cheeger constant or the isoperimetric number of the data manifolds, which in many cases is believed to be at least inverse polynomial in dimension (e.g., see Bobkov et al. 1997 for the Cheeger constant of the Gaussian distribution.) Indeed, in Section 3.4 we will formally lowerbound $\rho_{m+1}$ by the product of the augmentation strength and the Cheeger constant of the subpopulation distributions (Proposition 3.9), and lowerbound the Cheeger constant by inverse polynomial for concrete settings where the data come from a mixture of manifolds (Theorem 3.11).

Assumption 3.5 also implies properties of the graph spectrum. Recall that $\gamma_{i}$ is the $i$-th largest eigenvalue of the normalized adjacency matrix $\overline{A}$ and $\gamma_{1}=1$. According to Cheeger’s inequality (Lemma B.4), Assumption 3.5 implies that $\gamma_{2m}\leq 1-\Omega(\rho^{2}/\log{m})$, which suggests that there is a gap between $\gamma_{1}$ and $\gamma_{2m}$ and will be useful in our analysis.

Next, we formalize the assumption that very few edges cross different ground-truth classes. It turns out that it suffices to assume that the labels are recoverable from the augmentations (which is also equivalent to that two examples in different classes can rarely be augmented into the same point).

Let $\bar{x}\sim\mathcal{P}_{\overline{\mathcal{X}}}$ and $y(\bar{x})$ be its label. Let the augmentation $x\sim\mathcal{A}(\cdot|\bar{x})$. We assume that there exists a classifier $g$ that can predict $y(\bar{x})$ given $x$ with error at most $\alpha$. That is, $g(x)=y(\bar{x})$ with probability at least $1-\alpha$.

A small $\alpha$ in Assumption 3.6 means that different classes are “separated” in the sense that data from different classes have very few (at most $O(\alpha)$) shared augmentations. Alternatively, one can think of this assumption as assuming that the augmentation graph can be partitioned into $r$ clusters each corresponding to augmentations from one class, and there are at most $O(\alpha)$ edges across the clusters. This is typically true for real-world image data like ImageNet, since for any two images from different classes (e.g., images of a Husky and a Birman cat), using the typical data augmentations such as adding noise and random cropping can rarely (with exponentially small probability) lead to the same augmented image.

Typically, both $\rho$ in Assumption 3.5 and $\alpha$ in Assumption 3.7 are small positive values that are much less than 1. However, $\rho$ can be much larger than $\alpha$. Recall that $\rho$ can be expected to be at least inverse polynomial in dimension. In contrast, $\alpha$ characterizes the separation between classes and are expected to be exponentially small in typical cases. For example, in the running example of Section 3.1 with Gaussian perturbation augmentation, if $\sigma\sqrt{d}$ is smaller than the minimum distance between two subpopulations, we can rarely augment two datapoints from distinct subpopulations into a shared augmentation, and therefore $\alpha$ is expected to exponentially small. Our analysis below operates in the reasonable regime where $\rho^{2}$ is larger than $\alpha$, which intuitively means that the internal connection within the cluster is bigger than the separation between the clusters.

We also introduce the following assumption which states that some minimizer of the population spectral contrastive loss can be realized by the hypothesis class.

Our main theorem bound from above the linear probe error of the feature learned by minimizing the population spectral contrastive loss. In Theorem 4.3 we extend this result to the case where both the feature and the linear head are learned from empirical datasets.

Assume the representation dimension $k\geq 2r$ and Assumption 3.6 holds for $\alpha>0$. Let $\mathcal{F}$ be a hypothesis class that satisfies Assumption 3.7 and let ${f^{*}_{\textup{pop}}}\in\mathcal{F}$ be a minimizer of $\mathcal{L}({f})$. Then, we have

In particular, if Assumption 3.5 also holds and $k>2m$, we have $\mathcal{E}({f^{*}_{\textup{pop}}})\leq\widetilde{O}\left(\alpha/\rho^{2}\right)$.

Here we use $\widetilde{O}(\cdot)$ to hide universal constant factors and logarithmic factors in $k$. We note that $\alpha=0$ when augmentations from different classes are perfectly disconnected in the augmentation graph, in which case the above theorem guarantees the exact recovery of the ground truth. Generally, we expect $\alpha$ to be an extremely (exponentially) small constant independent of $k$, whereas $\rho_{\lfloor k/2\rfloor}$ increases with $k$ and can be at least inverse polynomial when $k$ is reasonably large, hence much larger than $\sqrt{\alpha}$. We characterize the $\rho_{k}$’s growth on more concrete distributions in the next subsection. When $k>2m$, as argued below Assumption 3.6, we expect that $\alpha\ll\rho^{2}\leq\rho^{2}_{m+1}$ and thus the error $\alpha/\rho^{2}$ is sufficiently small.

Previous works on graph partitioning Lee et al. 2014, Arora et al. 2009, Leighton and Rao 1999 often analyze the rounding algorithms that conduct clustering based on the representations of unlabeled data and do not analyze the performance of linear probe (which has access to labeled data). These results provide guarantees on the approximation ratio—the ratio between the conductance of the obtained partition to the best partition—which may depend on graph size Arora et al. 2009 that can be exponentially large in our setting. The approximation ratio guarantee does not lead to a guarantee on the representations’ performance on downstream tasks. Our guarantees are on the linear probe accuracy on the downstream tasks and independent of the graph size. We rely on the formulation of the downstream task’s labeling function (Assumption 3.6) as well as a novel analysis technique that characterizes the linear structure of the representations. In Section B, we provide the proof of Theorem 3.8 as well as its more generalized version where $k/2$ is relaxed to be any constant fraction of $k$. A proof sketch of Theorem 3.8 is given in Section 5.1.

4 Provable instantiation of Theorem 3.8 to mixture of manifold data

In this section, we exemplify Theorem 3.8 on examples where the natural data distribution is a mixture of manifolds.

That is, $\rho_{m+1}$ is at least linear in the augmentation size $\sigma$ and the Cheeger constants of subpopulations.

In many cases, the Cheeger constant is at least inverse polynomial in the data dimension Chen 2021, Lee and Vempala 2016. When the manifolds $P_{i}$ are spherical Gaussian with unit identity covariance, the Cheeger constant is $\Omega(1)$Bobkov et al. 1997, and thus the distribution $\mathcal{P}_{\overline{\mathcal{X}}}$ in Proposition 3.9 satisfies Assumption 3.5 with $\rho\gtrsim\sigma$. Furthermore, when the distribution is transformed by a function with Lipschitzness $\kappa>0$, the Cheeger constant changes by a factor at most $\kappa$. Therefore, Proposition 3.9 also applies to a mixture of manifolds setting defined below.

In the rest of this section, we instantiate Theorem 3.8 on a mixture of manifolds example where the data is generated from a Lipschitz transformation of a mixture of Gaussian distributions, and give an error bound for the downstream classification task.

Let the data augmentation of a natural data sample $\bar{x}$ be $\bar{x}+\xi$ where $\xi\sim\mathcal{N}(0,\frac{\sigma^{2}}{d}\cdot I_{d\times d})$ is isotropic Gaussian noise with $0<\sigma\lesssim\frac{1}{\sqrt{d}}$. We also assume $\min_{i\neq j}\left\lVert\mu_{i}-\mu_{j}\right\rVert_{2}\gtrsim\frac{\kappa\cdot\sqrt{\log d}}{\sqrt{d^{\prime}}}$.

Let $\bar{y}(x)$ be the most likely mixture index $i$ that generates $x$: $\bar{y}(x):=\arg\max_{i}P_{i}(x)$. The simplest downstream task can have label $y(x)=\bar{y}(x)$. More generally, let $r^{\prime}\leq r$ be the number of labels, and the label $y(x)\in[r^{\prime}]$ in the downstream task be equal to $\pi(\bar{y}(x))$ where $\pi$ is a function that maps $[r]$ to $[r^{\prime}]$.

We note that the intra-class distance in the latent space is on the scale of $\Omega(1)$, which can be much larger than the distance between class means which is assumed to be $\gtrsim\frac{\kappa\cdot\sqrt{\log d}}{\sqrt{d^{\prime}}}$. Therefore, distance-based clustering algorithms do not apply. Moreover, in the simple downstream tasks, the label for $x$ could be just the index of the mixture where $x$ comes from. We also allow downstream tasks that merge the $r$ components into $r^{\prime}$ labels as long as each mixture component gets the same label. We apply Theorem 3.8 and get the following theorem:

When $k\geq 2r+2$, Example 3.10 satisfies Assumption 3.6 with $\alpha\leq\frac{1}{\textup{poly}(d)}$, and has $\rho_{\lfloor k/2\rfloor}\gtrsim\frac{\sigma}{\kappa\sqrt{d}}$. As a consequence, the error bound is $\mathcal{E}({f^{*}_{\textup{pop}}})\leq\widetilde{O}\left(\frac{\kappa^{2}}{\sigma^{2}\cdot\textup{poly}(d)}\right)$.

The theorem above guarantees small error even when $\sigma$ is polynomially small. In this case, the augmentation noise has a much smaller scale than the data (which is at least on the order of $1/\kappa$). This suggests that contrastive learning can non-trivially leverage the structure of the underlying data and learn good representations with relatively weak augmentation. To the best of our knowledge, it is difficult to apply the theorems in previous works (Arora et al. 2019, Lee et al. 2020, Tosh et al. 2020, Tosh et al. 2021, Wei et al. 2020) to this example and get similar guarantees with polynomial dependencies on $d,\sigma,\kappa$. The work of Wei et al. 2020 can apply to the setting where $r$ is known and the downstream label is equal to $\bar{y}(x)$, but cannot handle the case when $r$ is unknown or when two mixture component can have the same label. We refer the reader to the related work section for more discussions and comparisons. The proof can be found in Section C.2.

Finite-sample generalization bounds

In Section 3, we provide guarantees for spectral contrastive learning on population data. In this section, we show that these guarantees can be naturally extended to the finite-sample regime with standard concentration bounds. In particular, given a unlabeled pretraining dataset $\{\bar{x}_{1},\bar{x}_{2},\cdots,\bar{x}_{n_{\rm{pre}}}\}$ with $\bar{x}_{i}\sim\mathcal{P}_{\overline{\mathcal{X}}}$, we learn a feature extractor by minimizing the following empirical spectral contrastive loss:

Recall that ${f^{*}_{\textup{pop}}}\in\mathcal{F}$ is a minimizer of $\mathcal{L}({f})$. The following theorem with proofs in Section D.1 bounds the population loss of a feature extractor trained with finite data:

For some ${\kappa}>0$, assume $\left\lVert f(x)\right\rVert_{\infty}\leq{\kappa}$ for all $f\in\mathcal{F}$ and $x\in\mathcal{X}$. Let ${f^{*}_{\textup{pop}}}\in\mathcal{F}$ be a minimizer of the population loss $\mathcal{L}({f})$. Given a random dataset of size ${n_{\rm{pre}}}$, let ${\hat{f}_{\textup{emp}}}\in\mathcal{F}$ be a minimizer of empirical loss $\widehat{\mathcal{L}}_{{n_{\rm{pre}}}}({f})$. Then, when Assumption 3.7 holds, with probability at least $1-\delta$ over the randomness of data, we have

where constants $c_{1}\lesssim k^{2}{\kappa}^{2}+k{\kappa}$ and $c_{2}\lesssim k{\kappa}^{2}+k^{2}{\kappa}^{4}$.

The Rademacher complexity usually looks like $\widehat{\mathcal{R}}_{n}(\mathcal{F})=\sqrt{{R}/{n}}$ where $R$ measures the complexity of $\mathcal{F}$ (hence only depends on $\mathcal{F}$). This suggests that when $\kappa$ is $O(1)$, the sample complexity for acheiving suboptimality $\epsilon$ on population loss is $O(k^{4}R/\epsilon^{2})$. We can apply Theorem 4.1 to any hypothesis class $\mathcal{F}$ of interest (e.g., deep neural networks) and plug in off-the-shelf Rademacher complexity bounds. For instance, in Section D.2 we give a corollary of Theorem 4.1 when $\mathcal{F}$ contains deep neural networks with ReLU activation.

The theorem above shows that we can achieve near-optimal population loss by minimizing empirical loss up to some small excess loss. The following theorem characterizes how the error propagates to the linear probe performance mildly under some spectral gap conditions.

In the setting of Theorem 4.1, suppose Assumption 3.5 holds for $\rho>0$, Assumption 3.6 holds for $\alpha>0$, Assumption 3.7 holds, and the representation dimension $k\geq\max\{4r+2,2m\}$,. Then, with $1-\delta$ probability over the randomness of data, for any ${\hat{f}_{\textup{emp}}}\in\mathcal{F}$ that minimizes the empirical loss $\widehat{\mathcal{L}}_{{n_{\rm{pre}}}}({f})$, we have that

where $c\lesssim(k\kappa+k\kappa^{2}+1)^{2}$, and $\Delta_{\gamma}:=\gamma_{\lfloor{3k}/4\rfloor}-\gamma_{k}$ is the eigenvalue gap between the $\lfloor{3k}/4\rfloor$-th and the $k$-th eigenvalue.

This theorem shows that the error on the downstream task only grows linearly with the excess loss during pretraining. Roughly speaking, one can think of $\Delta_{\gamma}$ as on the order of $1-\gamma_{k}$, hence by Cheeger’s inequality it’s larger than $\rho^{2}$. When $\widehat{\mathcal{R}}_{{n_{\rm{pre}}}/2}(\mathcal{F})=\sqrt{{2R}/{{n_{\rm{pre}}}}}$ and $\kappa\leq O(1)$, we have that the number of unlabeled samples required to achieve $\epsilon$ downstream error is $O(m^{6}R/\epsilon^{2}\rho^{4})$. We can relax Assumption 3.7 to approximate realizability in the sense that $\mathcal{F}$ contains some sub-optimal feature extractor under the population spectral loss and pay an additional error term in the linear probe error bound. The proof of Theorem 4.2 can be found in Section D.3.

2 Labeled sample complexity for linear probe

The following Theorem 4.3 provides a generalization guarantee for the linear classifier that minimizes capped quadratic loss on a labeled downstream dataset of size ${n_{\rm{down}}}$. The key challenge of the proof is showing the existence of a small-norm linear head $B$ that gives small population quadratic loss, which is not obvious from Theorem 4.2 where only small 0-1 error is guaranteed. Given a labeled dataset $\{(\bar{x}_{i},y(\bar{x}_{i}))\}_{i=1}^{n_{\rm{down}}}$ where $\bar{x}_{i}\sim\mathcal{P}_{\overline{\mathcal{X}}}$ and $y(\bar{x}_{i})$ is its label, we sample ${x}_{i}\sim\mathcal{A}(\cdot|\bar{x}_{i})$ for $i\in[{n_{\rm{down}}}]$. Given a norm bound $C_{k}>0$, we learn a linear probe $\widehat{B}$ by minimizing the capped quadratic loss subject to a norm constraint:

In the setting of Theorem 4.2, choose $C_{k}>0$ such that $C_{k}\geq\frac{2(k+1)}{\gamma_{k}}$. Then, with probability at least $1-\delta$ over the randomness of data, for any ${\hat{f}_{\textup{emp}}}\in\mathcal{F}$ that minimizes the empirical pre-training loss $\widehat{\mathcal{L}}_{{n_{\rm{pre}}}}({f})$ and a linear head $\widehat{B}$ learned from Equation (10), we have

Here the first term is an error caused by the property fo the population data, which is unavoidable even with infinite pretraining and downstream samples (but it can be small as argued in Section 3.3). The second term is caused by finite pretraining samples, and the third term is caused by finite samples in the linear classification on the downstream task.

Typically, the Rademacher complexity is roughly $\widehat{\mathcal{R}}_{{n_{\rm{pre}}}/2}(\mathcal{F})=\sqrt{{2R}/{{n_{\rm{pre}}}}}$ where $R$ is captures the complexity of the model architecture. Thus, to achieve final linear probe error no more than $O(\epsilon)$, we would need to select $k$ such that $\frac{\alpha}{\rho_{\lfloor k/2\rfloor}^{2}}\cdot\log k\leq\epsilon$, and we need $\textup{poly}(k,\frac{1}{\Delta_{\gamma}},R,\frac{1}{\epsilon})$ pretraining samples and $\textup{poly}(k,r,\frac{1}{\gamma_{k}},\frac{1}{\epsilon})$ downstream samples.

When $r<k$, the eigengap $\Delta_{\gamma}$ is on the order of $1-\gamma_{k}$ which is larger than $\rho^{2}$ by Cheeger inequality. Recall that $\rho$ is at least inverse polynomial in $d$ as argued in Section 3.4, one can expect $\frac{1}{\Delta_{\gamma}}$ to be at most $\textup{poly}(d)$. On the other hand, $\gamma_{k}\approx 1$ so $\frac{1}{\gamma_{k}}$ can be thought of as a constant. Thus, the final required number of pretraining samples is ${n_{\rm{pre}}}=\textup{poly}(k,d,R,\frac{1}{\epsilon})$ and number of downstream samples is ${n_{\rm{down}}}=\textup{poly}(r,k,\frac{1}{\epsilon})$. We note that the downstream sample complexity doesn’t depend on the complexity of the hypothesis class $R$, suggesting that pretraining helps reduce the sample complexity of the supervised downstream task.

The proof of Theorem 4.3 is in Section E.

Analysis Framework and Proof Sketch

As discussed before and suggested by the structured of Section 3 and 4, our analysis framework decompose the problem into a key step about the population cases (Section 3) and a few other somewhat standard steps that link empirical losses to population losses (Section 4). As depicted in Figure 2, the core step (Theorem 3.8, or its extension Theorem 4.2) is to show that a small population pretraining loss implies the existence of a linear classifier for the downstream task, that is, a small minimal downstream loss.

We first remark that a feature of our analysis framework is that we link the population pretraining data case to the finite sample case by showing the empirical and population pretraining losses are similar when the feature extractors are a parameterized family of models with capacity bounds (the first arrow in Figure 2). Hypothetically, suppose such a connection between population and empirical data case was built through the relationship between the population and empirical graphs, e.g., by proving that the empirical graph has similar spectral properties as the population graph, then the sample complexity will be exponential. Intuitively, this is because the population graph is very sparse, and the empirical graph is with high probability empty if the number of samples is only polynomial in dimension (e.g. consider the case when the augmentation simply adds small perturbation, as in the running example in Section 3.1). The empirical graph essentially follows the well-studied random geometric graph model (Penrose 2003), and tends to have no structure in high dimension Bubeck et al. 2016, Liu et al. 2021, Brennan et al. 2020. The fundamental difference between this hypothetical and our framework is that the empirical graph’s definition does not involve any parameterization, and thus the resemblance between the empirical and population graphs does not leverage the extrapolation (or inductive bias) of the model parameterization as our framework does for the pretraining losses.

We note that the inductive bias of the parameterized model is indeed used in the analysis for finite-sample case. We assume that the model family $\mathcal{F}$ can express the eigenfunctions/eigenvectors of the graph (Assumption 3.7) and also implicitly assume bounds on its Rademacher complexity (in Theorem 4.3).

Once we obtained that the existence of a linear classifier, the remaining steps (the third and fourth arrows in Figure 2) follow from standard supervised learning theory.

In the rest of this section, we will give a proof sketch of the population case, which is the more challenging step.

In this section, we give a proof sketch of Theorem 3.8 in a simplified binary classification setting where there are only two classes in the downstream task.

Recall that $N$ is the size of $\mathcal{X}$. Recall that $w_{x}$ is the total weight associated with an augmented datapoint $x\in\mathcal{X}$, which can also be thought of as the probability mass of $x$ as a randomly sampled augmented datapoint. In the scope of this section, for demonstrating the key idea, we also assume that $x$ has uniform distribution, i.e., $w_{x}=\frac{1}{N}$ for any $x\in\mathcal{X}$.

Let $\mathcal{L}\triangleq{I}-\overline{A}$ be the normalized Laplacian matrix. Then, $v_{i}$’s are the $k$ smallest unit-norm eigenvectors of $\mathcal{L}$ with eigenvalues $\lambda_{i}=1-\gamma_{i}$. Elementary derivations can give a well-known, important property of the Laplacian matrix $L$: the quadratic form $\vec{g}^{\top}L\vec{g}$ captures the amount of edges across the two groups that are defined by the binary vector $\vec{g}$ (Chung and Graham 1997, section 1.2):

With slight abuse of notation, suppose $(x,x^{+})$ is the random variable for a positive pair. Using that $w$ is the density function for the positive pair and the simplification that $w_{x}=1/N$, we can rewrite equation (12) as

Next, we use equation (14) to link $\vec{g}$ to the eigenvectors of $L$. Let $\lambda_{k+1}\leq\dots\lambda_{N}$ be the rest of eigenvalues with unit-norm eigenvectors $v_{k+1},\dots,v_{N}$. Let $\Pi\triangleq\sum_{i=1}^{k}v_{i}v_{i}^{\top}$ and $\Pi_{\perp}\triangleq\sum_{i=k+1}^{N}v_{i}v_{i}^{\top}$ be the projection operators onto the subspaces spanned by the first $k$ and the last $N-k$ eigenvectors, respectively. Equation (14) implies that $\vec{g}$ has limited projection to the subspace of $\Pi_{\perp}$:

where the first inequality follows from dropping the $\|(\Pi\vec{g})^{\top}L\Pi\vec{g}\|_{2}^{2}$ and using $\Pi_{\perp}L\Pi=0$, and the second inequality is because that $\Pi_{\perp}$ only contains eigenvectors with eigenvalue at least $\lambda_{k+1}$.

By higher-order Cheeger inequality (see Lemma B.4), we have that $\lambda_{k+1}\gtrsim\rho_{\lceil k/2\rceil}^{2}$. Then, we obtain the mean-squared error bound:

Experiments

We test spectral contrastive learning on benchmark vision datasets. We minimize the empirical spectral contrastive loss with an encoder network $f$ and sample fresh augmentation in each iteration. The pseudo-code for the algorithm and more implementation details can be found in Section A.

Encoder / feature extractor. The encoder $f$ contains three components: a backbone network, a projection MLP and a projection function. The backbone network is a standard ResNet architecture. The projection MLP is a fully connected network with BN applied to each layer, and ReLU activation applied to each except for the last layer. The projection function takes a vector and projects it to a sphere ball with radius $\sqrt{\mu}$, where $\mu>0$ is a hyperparameter that we tune in experiments. We find that using a projection MLP and a projection function improves the performance.

Linear evaluation protocol. Given the pre-trained encoder network, we follow the standard linear evaluation protocol Chen and He 2020 and train a supervised linear classifier on frozen representations, which are from the ResNet’s global average pooling layer.

Results. We report the accuracy on CIFAR-10/100 Krizhevsky and Hinton 2009 and Tiny-ImageNet Le and Yang 2015 in Table 1. Our empirical results show that spectral contrastive learning achieves better performance than two popular baseline algorithms SimCLR Chen et al. 2020a and SimSiam Chen and He 2020. In Table 2 we report results on ImageNet Deng et al. 2009 dataset, and show that our algorithm achieves similar performance as other state-of-the-art methods. We note that our algorithm is much more principled than previous methods and doesn’t rely on large batch sizes (SimCLR Chen et al. 2020a), momentum encoders (BYOL Grill et al. 2020 and MoCo He et al. 2020) or additional tricks such as stop-gradient (SimSiam Chen and He 2020).

Conclusion

In this paper, we present a novel theoretical framework of self-supervised learning and provide provable guarantees for the learned representation on downstream linear classification tasks. We hope the framework could facilitate future theoretical analyses of self-supervised pretraining losses and inspire new methods. It does not capture the potential implicit bias of optimizers but does take into account the inductive bias of the models. By abstracting away the effect of optimization, we can focus on the effect of pretraining losses and their interaction with the structure of the population data. Future directions may include designing better pretraining losses and analyzing more fine-grained properties of the learned representations (e.g., as in recent follow-up works Shen et al. 2022, HaoChen et al. 2022), by potentially leveraging more advanced techniques from spectral graph theory.

Acknowledgements

We thank Margalit Glasgow, Ananya Kumar, Jason D. Lee, Sang Michael Xie, and Guodong Zhang for helpful discussions. CW acknowledges support from an NSF Graduate Research Fellowship. TM acknowledges support of Google Faculty Award and NSF IIS 2045685. We also acknowledge the support of HAI and the Google Cloud. Toyota Research Institute ("TRI") provided funds to assist the authors with their research but this article solely reflects the opinions and conclusions of its authors and not TRI or any other Toyota entity.

References

Appendix A Experiment details

The pseudo-code for our empirical algorithm is summarized in Algorithm 1.

Our results with different hyperparameters on CIFAR-10/100 and Tiny-ImageNet are listed in Table 3.

Additional details about the encoder. For the backbone network, we use the CIFAR variant of ResNet18 for CIFAR-10 and CIFAR-100 experiments and use ResNet50 for Tiny-ImageNet and ImageNet experiments. For the projection MLP, we use a 2-layer MLP with hidden and output dimensions 1000 for CIFAR-10, CIFAR100, and Tiny-ImageNet experiments. We use a 3-layer MLP with hidden and output dimension 8192 for ImageNet experiments. We set $\mu=10$ in the ImageNet experiment, and set $\mu\in\{1,3,10\}$ for the CIFAR-10/100 and Tiny-ImageNet experiments.

Training the encoder. We train the neural network using SGD with momentum 0.9. The learning rate starts at 0.05 and decreases to 0 with a cosine schedule. On CIFAR-10/100 and Tiny-ImageNet we use weight decay 0.0005 and train for 800 epochs with batch size 512. On ImageNet we use weight decay 0.0001 and train for 100 epochs with batch size 384. We use 1 GTX 1080 GPU for CIFAR-10/100 and Tiny-ImageNet experiments, and use 8 GTX 1080 GPUs for ImageNet experiments.

Linear evaluation protocol. We train the linear head using SGD with batch size 256 and weight decay 0 for 100 epochs, learning rate starts at 30.0 and is decayed by 10x at the 60th and 80th epochs.

Image transformation details. We use the same augmentation strategy as described in Chen and He 2020.

Appendix B Proofs for Section 3

We first prove a more generalized version of Theorem 3.8 in section B.1, and then prove Theorem 3.8 in Section B.2.

For the proof we will follow the convention in literature Lee et al. 2014 and define the normalized Laplacian matrix as follows:

Let $G=(\mathcal{X},w)$ be the augmentation graph defined in Section 3.1. The normalized Laplacian matrix of the graph is defined as $\mathcal{L}={I}-{D}^{-1/2}{A}{D}^{-1/2}$, where ${A}$ is the adjacency matrix with ${A}_{xx^{\prime}}=w_{xx^{\prime}}$ and ${D}$ is a diagonal matrix with ${D}_{xx}=w_{x}$.

It is easy to see that $\mathcal{L}=I-\overline{A}$ where $\overline{A}$ is the normalized adjacency matrix defined in Section 3.1. Therefore, when $\lambda_{i}$ is the $i$-th smallest eigenvalue of $\mathcal{L}$, $1-\lambda_{i}$ is the $i$-th largest eigenvalue of $\overline{A}$.

We call a function defined on augmented data $\hat{y}:\mathcal{X}\rightarrow[r]$ an extended labeling function. Given an extended labeling function, we define the following quantity that describes the difference between extended labels of two augmented data of the same natural datapoint:

We also define the following quantity that describes the difference between extended label of an augmentated datapoint and the ground truth label of the corresponding natural datapoint:

Recall the spectral contrastive loss defined in Section 3.2 is:

We first state a more general version of Theorem 3.8 as follows.

where $\vec{{y}}(\bar{x})$ is the one-hot embedding of ${y}(\bar{x})$ and $\rho_{k^{\prime}}$ is the sparsest $m$-partition defined in Definition 3.4. Furthermore, the error of the linear probe predictor can be bounded by

Also, if we let $\lambda_{i}$ be the $i$-th smallest eigenvalue of the normalized Laplacian matrix of the graph of the augmented data, we can find a matrix $B^{*}$ satisfying the above equations with norm bound $\left\lVert{B}^{*}\right\rVert_{F}\leq 1/(1-\lambda_{k})$.

We provide the proof for Theorem B.2 below.

Let $\lambda_{1},\lambda_{2},\cdots,\lambda_{k},\lambda_{k+1}$ be the $k+1$ smallest eigenvalues of the Laplacian matrix $L$. The following theorem gives a theoretical guarantee similar to Theorem B.2 except for that the bound depends on $\lambda_{k+1}$:

where $\vec{{y}}(\bar{x})$ is the one-hot embedding of ${y}(\bar{x})$. Furthermore, the error can be bounded by

We defer the proof of Theorem B.3 to Section B.3.

To get rid of the dependency on $\lambda_{k+1}$, we use following higher-order Cheeger’s inequality from Louis and Makarychev 2014.

Let $G=(V,w)$ be a weight graph with $|V|=N$. Then, for any $t\in[N]$ and $\zeta>0$ such that $(1+\zeta)t\in[N]$, there exists a partition $S_{1},S_{2},\cdots,S_{t}$ of $V$ with

where $\phi_{G}(\cdot)$ is the Dirichlet conductance defined in Definition 3.3.

Now we prove Theorem B.2 by combining TheoremB.3 and Lemma B.4.

Let $G=(\mathcal{X},w)$ be the augmentation graph. In Lemma B.4 let $(1+\zeta)t=k+1$ and $t=k^{\prime}$ we have: there exists partition $S_{1},\cdots,S_{k^{\prime}}\subset\mathcal{X}$ such that $\phi_{G}(S_{i})\lesssim\textup{poly}(1/\zeta)\sqrt{\lambda_{k+1}\log{(k+1)}}$ for $\forall i\in[k^{\prime}]$. By Definition 3.4, we have $\rho_{k^{\prime}}\leq\max_{i\in[k^{\prime}]}\phi_{G}(S_{i})\lesssim\textup{poly}(1/\zeta)\sqrt{\lambda_{k+1}\log{(k+1)}}$, which leads to $\frac{1}{\lambda_{k+1}}\lesssim\textup{poly}(1/\zeta)\cdot{log(k+1)}\cdot\frac{1}{\rho_{k^{\prime}}^{2}}$. Plugging this bound to Theorem B.3 finishes the proof. ∎

B.2 Proof of Theorem 3.8

We will use the following lemma which gives a connection between $\phi^{\hat{y}}$, $\Delta(y,\hat{y})$ and Assumption 3.6.

Let $G=(\mathcal{X},w)$ be the augmentation graph, $r$ be the number of underlying classes. Let $S_{1},S_{2},\cdots,S_{r}$ be the partition induced by the classifier $g$ in Assumption 3.6. Then, there exists an extended labeling function $\hat{y}$ such that

We define function $\hat{y}:\mathcal{X}\rightarrow[r]$ as follows: for an augmented data $x\in\mathcal{X}$, we use function $\hat{y}(x)$ to represent the index of set that $x$ is in, i.e., $x\in S_{\hat{y}(x)}$. By Assumption 3.6 it is easy to see $\Delta(y,\hat{y})\leq\alpha$. On the other hand, we have

Here the inequality is because when $\hat{y}(x)\neq\hat{y}(x^{\prime})$, there must be $\hat{y}(x)\neq y(\bar{x})$ or $\hat{y}(x^{\prime})\neq y(\bar{x})$. ∎

Now we give the proof of Theorem 3.8 using Lemma B.5 and Theorem B.2.

Let $S_{1},S_{2},\cdots,S_{r}$ be the partition of $\mathcal{X}$ induced by the classifier $g$ given in Assumption 3.6. Define function $\hat{y}:\mathcal{X}\rightarrow[r]$ as follows: for an augmented datapoint $x\in\mathcal{X}$, we use function $\hat{y}(x)$ to represent the index of set that $x$ is in, i.e., $x\in S_{\hat{y}(x)}$. Let $k^{\prime}=\lfloor\frac{k}{2}\rfloor$ in Theorem B.2, we have $\Pr_{\bar{x}\sim\mathcal{P}_{\overline{\mathcal{X}}},{x}\sim\mathcal{A}(\cdot|\bar{x})}\left(g_{{f^{*}_{\textup{pop}}},{B}^{*}}({x})\neq y(\bar{x})\right)\lesssim\log(k)\cdot\frac{\phi^{\hat{y}}}{\rho_{\lfloor k/2\rfloor}^{2}}+\Delta(y,\hat{y}).$ By Lemma B.5 we have $\phi^{\hat{y}}\leq 2\alpha$ and $\Delta(y,\hat{y})\leq\alpha$, so we have $\Pr_{\bar{x}\sim\mathcal{P}_{\overline{\mathcal{X}}},{x}\sim\mathcal{A}(\cdot|\bar{x})}\left(g_{{f^{*}_{\textup{pop}}},{B}^{*}}({x})\neq y(\bar{x})\right)\lesssim\frac{\alpha}{\rho_{\lfloor k/2\rfloor}^{2}}\cdot\log(k).$ Notice that by definition of ensembled linear probe predictor, $\bar{g}_{{f^{*}_{\textup{pop}}},{B}^{*}}(\bar{x})\neq y(\bar{x})$ happens only if more than half of the augmentations of $\bar{x}$ predicts differently from $y(\bar{x})$, so we have $\Pr_{\bar{x}\sim\mathcal{P}_{\overline{\mathcal{X}}}}\left(\bar{g}_{{f^{*}_{\textup{pop}}},{B}^{*}}(\bar{x})\neq y(\bar{x})\right)\leq 2\Pr_{\bar{x}\sim\mathcal{P}_{\overline{\mathcal{X}}},{x}\sim\mathcal{A}(\cdot|\bar{x})}\left(g_{{f^{*}_{\textup{pop}}},{B}^{*}}({x})\neq y(\bar{x})\right)\lesssim\frac{\alpha}{\rho_{\lfloor k/2\rfloor}^{2}}\cdot\log(k)$. ∎

B.3 Proof of Theorem B.3

The proof of Theorem B.3 contains two steps. First, we show that when the feature extractor is composed of the minimal eigenvectors of the normalized Laplacian matrix $L$, we can achieve good linear probe accuracy. Then we show that minimizing $\mathcal{L}({f})$ gives us a feature extractor equally good as the eigenvectors.

For the first step, we use the following lemma which shows that the smallest eigenvectors of $\mathcal{L}$ can approximate any function on $\mathcal{X}$ up to an error proportional to the Rayleigh quotient of the function.

We can decompose the vector ${u}$ in the eigenvector basis as:

We also need the following claim about the Rayleigh quotient $R({u})$ when ${u}$ is a vector defined by an extended labeling function $\hat{y}$.

To see the connection between the feature extractor minimizing the population spectral contrastive loss $\mathcal{L}({f})$ and the feature extractor corresponding to eigenvectors of the Laplacian matrix, we use the following lemma which states that the minimizer of the matrix approximation loss defined in Section 3.2 is equivalent to the minimizer of population spectral contrastive loss up to a data-wise scaling.

Recall that the definition of spectral contrastive loss is

where $(x,x^{+})$ is a random positive pair, $(x,x^{-})$ is a random negative pair. We can rewrite the spectral contrastive loss as

Compare Equation (B.3) and Equation (21), we see they only differ by a constant, which finishes the proof. ∎

Note that the minimizer of matrix approximation loss is exactly the largest eigenvectors of ${I}-{L}$ (also the smallest eigenvectors of ${L}$) due to Eckart–Young–Mirsky theorem, Lemma B.8 indicates that the minimizer of $\mathcal{L}({f})$ is equivalent to the smallest eigenvectors of $\mathcal{L}$ up to data-wise scaling.

The following claim shows the relationship between quadratic loss and prediction error.

where $\vec{y}$ is the one-hot embedding of $y$.

When $y\neq g_{f,B}(x)$, by the definition of $g_{f,B}$ we know that there exists another $y^{\prime}\neq y$ such that $(Bf(x))_{y^{\prime}}\geq(Bf(x))_{y}$. In this case,

Now we are ready to prove Theorem B.3 by combining Lemma B.6, Claim B.7, Lemma B.8 and Claim B.9.

Now we come back to the feature extractor ${f^{*}_{\textup{pop}}}$ that minimizes the spectral contrastive loss function $\mathcal{L}({f})$. By Lemma B.8, matrix ${F}^{*}$ that contains $\sqrt{w_{x}}\cdot{f^{*}_{\textup{pop}}}(x)$ as its $x$-th row is a minimizer of $\mathcal{L}_{\textup{mf}}({F})$. By Eckard-Young-Mirsky theorem, we have

and let $\vec{{y}}(\bar{x})$ be the one-hot embedding of ${y}(\bar{x})$, $\vec{\hat{y}}({x})$ be the one-hot embedding of $\hat{y}({x})$, we have

To bound the error rate, we first notice that Claim B.9 tells us that for any ${x}\in\mathcal{X}$,

Now we bound the error rate on $\mathcal{X}$ as follows:

Appendix C Proofs for Section 3.4

Let $B_{\sigma}$ be the uniform distribution over a ball with radius $\sigma$. Let $S_{1},S_{2},\cdots,S_{m+1}$ be a partition of the Euclidean space. There must be some $i\in[m+1]$ such that $\Pr_{{x}\sim P_{j},\xi\sim B_{\sigma}}[{x}+\xi\in S_{i}]\leq\frac{1}{2}$ for all $j\in[m]$. Thus, we know that

On one hand, suppose $\int_{x\notin R}P_{j}(x)P(S_{i}|x)dx\geq\frac{1}{2}\Pr_{{x}\sim P_{j},\xi\sim B_{\sigma}}[{x}+\xi\in S_{i}]$, we can lower bound the numerator in the RHS of Equation (29) as

hence the RHS of Equation (29) is at least $1/6$.

On the other hand, suppose $\int_{x\notin R}P_{j}(x)P(S_{i}|x)dx<\frac{1}{2}\Pr_{{x}\sim P_{j},\xi\sim B_{\sigma}}[{x}+\xi\in S_{i}]$, we have

hence the denominator of the RHS of Equation (29) can be upper bounded by

For two Gaussian distributions with variance $\sigma^{2}\cdot\mathcal{I}_{d\times d}$ and centers at most $\frac{\sigma}{6}$ far from each other, their TV-distance is at most $\frac{1}{6}$ (see the first equation on Page 5 of Devroye et al. 2018), hence for any $x\in N(R)$, we have $P(S_{i}|x)\geq\frac{2}{3}-\frac{1}{6}=\frac{1}{2}$. We can now lower bound the numerator in the RHS of Equation (29) as:

Notice that $\Pr_{{x}\sim P_{j},\xi\sim B_{\sigma}}[{x}+\xi\in S_{i}]\leq\frac{1}{2}$ and by the definition of $R$, we know $\int_{x\in R}P_{j}(x)dx\leq\frac{3}{4}$, thus

Combine Equation (34), Equation (C.1) and Equation (37) gives:

Notice that (using the definition of surface area (Guggenheimer 1977, chapter 4))

we have that as $\sigma\rightarrow 0^{+}$,

C.2 Proof of Theorem 3.11

In this section, we give a proof of Theorem 3.11.

The following lemma shows that the augmented graph for Example 3.10 satisfies Assumption 3.6 with some bounded $\alpha$.

In the setting of Theorem 3.11, the data distribution satisfies Assumption 3.6 with $\alpha\leq\frac{1}{\textup{poly}(d^{\prime})}$.

For any $z\sim\mathcal{N}(\mu_{i},\frac{1}{d^{\prime}}\cdot I_{d^{\prime}\times d^{\prime}})$ and any $j\neq i$, by the tail bound of gaussian distribution we have

Also, for $\xi\sim\mathcal{N}(0,\frac{1}{d}\cdot I_{d\times d})$, when $\sigma\leq\frac{1}{\sqrt{d}}$ we have

Notice that $\left\lVert Q^{-1}(Q(z)+\xi)-z\right\rVert_{2}\leq\kappa\left\lVert\xi\right\rVert$, we can set $\left\lVert\mu_{i}-\mu_{j}\right\rVert\gtrsim\kappa\frac{\sqrt{\log d}}{\sqrt{d}}$. Therefore, when $\left\lVert\mu_{i}-\mu_{j}\right\rVert\gtrsim\kappa\frac{\sqrt{\log d}}{\sqrt{d^{\prime}}}$ we can combine the above two cases and have

We use the following lemma to give a lower bound for the sparest $m$-partition of the augmentation graph in Example 3.10.

In the setting of Theorem 3.11, for any $k^{\prime}>r$ and $\tau>0$, we have

with $\Phi_{d}(z):=\Pr_{\xi\sim\mathcal{N}(0,\frac{1}{d}I_{d\times d})}(\left\lVert\xi\right\rVert_{2}\leq z)$, and

with $\Phi(z):=\int_{-\infty}^{z}\frac{e^{-u^{2}/2}}{\sqrt{2\pi}}du.$

The proof of Lemma C.2 can be found in Section C.3. Now we give the proof of Example 3.11.

The result on $\alpha$ is directly from Lemma C.1. By concentration inequality, there must exists some universal constant $C>0$ such that for any $d\geq C$, we have $1-\Phi_{d}(\sqrt{\frac{3}{2}})\leq\frac{1}{3}$. When this happens, we have $\Phi_{d}^{-1}(\frac{2}{3})\leq\sqrt{\frac{3}{2}}$. Since for $d\leq C$ we can just treat $d$ as constant, we have $\Phi_{d}^{-1}(\frac{2}{3})\lesssim 1$. Set $\tau=\sigma/d$ in Lemma C.2, we have $\rho_{k^{\prime}}\gtrsim\frac{\sigma}{\kappa\sqrt{d}}$. Set $k^{\prime}=\lfloor k/2\rfloor$, we apply Theorem 3.8 and get the bound we need. ∎

C.3 Proof of Lemma C.2

In this section we give a proof for Lemma C.2. We first introduce the following claim which states that for a given subset of augmented data, any two data close in $L_{2}$ norm cannot have a very different chance of being augmented into this set.

with $\Phi_{d}(z):=\Pr_{\xi\sim\mathcal{N}(0,\frac{1}{d}\cdot I_{d\times d})}(\left\lVert\xi\right\rVert_{2}\leq z)$.

By the definition of augmentation, we know

By the definition of $c_{\sigma}$, we have

Since $\Pr(S|x)\geq\frac{2}{3}$ by assumption, we have

Now we can bound the quanity of our interest:

Let $S_{1},\cdots,S_{k^{\prime}}$ be the disjoint sets that gives $\rho_{k^{\prime}}$ in Definition 3.4. First we notice that when $k^{\prime}>r$, there must exist $t\in[k^{\prime}]$ such that for all $i\in[r]$, we have

WLOG, we assume $j=1$ minimizes the RHS of Equation (42), so we only need to prove

where the second inequality is by Claim C.3. Notice that

where we use Equation (41). Define set $\widetilde{R}:=Q^{-1}(R)$ be the set in the ambient space corresponding to $R$. Define

Due to $Q$ being $\kappa$-bi-lipschitz, it is easy to see $\widetilde{N}(\widetilde{R})\subseteq Q^{-1}\left(N(R)\right)$. According to the Gaussian isoperimetric inequality Bobkov et al. 1997, we have

with $\Phi(\cdot)$ is the Gaussian CDF function defined as

By Equation (C.3), either case 1 or case 2 holds. Combining case 1 and case 2, we have

Appendix D Proofs for Section 4

We restate the empirical spectral contrastive loss defined in Section 4 as follows:

Consider a dataset $\widehat{\mathcal{X}}=\{\bar{x}_{1},\bar{x}_{2},\cdots,\bar{x}_{n}\}$ containing $n$ data points i.i.d. sampled from $\mathcal{P}_{\overline{\mathcal{X}}}$. Let $\hat{\mathcal{P}}_{\mathcal{X}}$ be the uniform distribution over $\widehat{\mathcal{X}}$. Let $\hat{P}_{\bar{x},\bar{x}^{\prime}}$ be the uniform distribution over data pairs $(\bar{x}_{i},\bar{x}_{j})$ where $i\neq j$. We define the empirical spectral contrastive loss of a feature extractor $f$ as

The following claim shows that $\widehat{\mathcal{L}}_{n}({f})$ is an unbiased estimator of population spectral contrastive loss.

$\widehat{\mathcal{L}}_{n}({f})$ is an unbiased estimator of $\mathcal{L}({f})$, i.e.,

To make use of the Radmacher complexity theory, we need to write the empirical loss as the sum of i.i.d. terms, which is achieved by the following sub-sampling scheme:

Given dataset $\widehat{\mathcal{X}}$, we sample a subset of tuples as follows: first sample a permutation $\pi:[n]\rightarrow[n]$, then we sample tuples $S=\{(z_{i},z_{i}^{+},z_{i}^{\prime})\}_{i=1}^{n/2}$ as follows:

It is easy to see that $\widehat{\mathcal{L}}_{S}({f})$ is an unbiased estimator of $\widehat{\mathcal{L}}_{n}({f})$:

For given $\widehat{\mathcal{X}}$, if we sample $S$ as above, we have:

This is obvious by the definition of $\widehat{\mathcal{L}}_{S}({f})$ and $\widehat{\mathcal{L}}_{n}({f})$. ∎

The following lemma reveals the relationship between the Rademacher complexity of feature extractors and the Rademacher complexity of the loss defined on tuples:

where $x_{1},x_{2},\cdots,x_{m}$ are in $\mathcal{X}$, and $\sigma$ is a uniform random vector in $\{-1,1\}^{n}$. Then, the empirical Rademacher complexity on any $n$ tuples $\{(z_{i},z_{i}^{+},z_{i}^{\prime})\}_{i=1}^{n}$ can be bounded by

here the second inequality is by Talagrand’s lemma. Notice that for any $z_{1},z_{2}\cdots z_{n}$ and $z_{1}^{\prime},z_{2}^{\prime},\cdots,z_{n}^{\prime}$ in $\mathcal{X}$ and any $i\in[k]$ we have

where the first inequaltiy is by Talagrand’s lemma. Combine these two equations and we get:

This means with probability at least $1-\delta/2$ over random $\widehat{\mathcal{X}}$, we have: with probability at least $1-\delta/2$ over random tuples $S$ conditioned on $\widehat{\mathcal{X}}$, Equation (44) holds. Since both $\mathcal{L}({f})$ and $\widehat{\mathcal{L}}_{{n_{\rm{pre}}}}({f})$ take value in range $[-2k{\kappa}^{2},2k{\kappa}^{2}+k^{2}{\kappa}^{4}]$, we have: with probability at least $1-{\delta}/{2}$ over random $\widehat{\mathcal{X}}$, we have for any $f\in\mathcal{F}$,