Investigating the Role of Negatives in
Contrastive Representation Learning

For our theoretical results, we adopt the framework of Saunshi et al. (2019), where the semantic similarity in the contrastive distribution is directly related to the labels in a downstream classification task. We refine the results of Saunshi et al. (2019) with a novel and tighter error transfer theorem that bounds the classification error by a function of the NCE error. We then study how this function depends on the number of concepts/classes in the data and the number of negatives in the NCE task. Our analysis captures the trade-off between coverage over the classes and collisions, and it reveals an expression for the ideal number of negatives $k$, which roughly scales as the number of concepts in the data. We emphasize that our sharper theoretical results are instrumental even for obtaining this qualitative trend; indeed, the analysis in Saunshi et al. (2019) instead suggests that increasing $k$ leads to decreasing performance.

Empirically, we study the effect of the number of negative samples on the performance of the final classifier on a downstream classification task for both vision and natural language datasets. For our language task, we observe that performance increases with $k$ up to a point, but it deteriorates as we increase $k$ further. The experiments also confirm that the optimal choice of $k$ depends on the number of concepts in the dataset. We do not observe the same trends in our vision experiments, and we describe how conflating factors—most notably data augmentation—cause disagreement between the observed empirical behavior and the predictions from the theoretical framework typically used to study NCE (Saunshi et al., 2019; Chuang et al., 2020).

Self-supervised learning has elicited breakthrough successes in representation learning, achieving state of the art performance across a wide array of important machine learning tasks, including computer vision (He et al., 2020; Huo et al., 2020; Khosla et al., 2020; Wang et al., 2020; Bachman et al., 2019), natural language processing (Fang and Xie, 2020; Fu et al., 2021; Le-Khac et al., 2020; Giorgi et al., 2020), and deep reinforcement learning. In vision, Chen et al. (2020) have shown that simple NCE algorithms are able to obtain extremely competitive performance on the challenging ImageNet dataset, something that was previously thought to require extensive supervised training only. In NLP, the NCE approach has garnered much attention since the pioneering work of Mikolov et al. (2013) demonstrated its efficacy in learning word-level representations. In Deep RL, NCE-style losses have been recently employed to learn world models in complex environments (Nachum et al., 2018; Laskin et al., 2020; Zhu et al., 2020; Kipf et al., 2019; Shelhamer et al., 2017).

Theoretical work on representation learning is quite nascent, with most works focused on mathematical explanations for the efficacy of these approaches (Saunshi et al., 2019; Lee et al., 2020). One is via conditional independence and redundancy (Tosh et al., 2020, 2021). Another theory proposed by Wang and Isola (2020) is that contrastive representations simultaneously optimize for alignment of similar points and uniformity in the embedding space, in the asymptotic regime where the number of negatives tends to infinity. This asymptotic regime was also studied by Foster et al. (2020), who analyze a contrastive estimator for an information-gain quantity. However, Foster et al. (2020) do not study representation learning and Wang and Isola (2020) do not theoretically analyze downstream performance as we do. Indeed, our results show that there can be a price for too many negatives due to collisions, and perhaps call into question whether this asymptotic regime should be used to explain the downstream performance of NCE representations.²²2Note that similar concerns have been raised regarding viewing representation learning as a mutual information maximization problem (Tschannen et al., 2019). This collision phenomenon has been observed by Chuang et al. (2020), who propose a debiased version of the contrastive loss that attempts to correct for negative samples coming from the positive class. They experimentally observe significant gains from removing the effect of collision. In contrast, here we study the original contrastive objective that is primarily used in practice and characterize how an appropriate choice of $k$ can balance the collision and coverage trade-off.

Our theoretical results closely follow the setup of Saunshi et al. (2019). We adopt essentially the same framework in the present paper, but unlike Saunshi et al. (2019) we do not use the number of negatives $k$ to define the downstream classification tasks. Our setup yields a single $k$-independent downstream task, which allows us to study the influence of this parameter. In addition, many steps of our analysis are sharper than theirs, which leads to our new qualitative findings.

Let the set of input features be denoted by $\mathcal{X}$ and the set of latent classes by $\mathcal{C}$, where $|\mathcal{C}|=N$. With each latent class $c\in\mathcal{C}$ we will associate a distribution $\mathcal{D}_{c}$ over $\mathcal{X}$. We can view $\mathcal{D}_{c}$ as the distribution over data conditioned on belonging to latent class $c$. We will also assume a distribution $\rho$ on $\mathcal{C}$. We will learn with a function class $\mathcal{F}$, the set of representation functions $f:\mathcal{X}\rightarrow\mathbb{R}^{d}$ that map the input to a $d$-dimensional vector. We use $\lesssim$ to denote less than or equal to up to constants.

We assume access to similar data points in the form of pairs $(x,x^{+})$ and $k$ negative data points $x^{-}_{1},\ldots,x^{-}_{k}$. To formalize this, an unlabeled sample from $\mathcal{D}_{\textup{NCE}}$ is generated as follows:

• •

  Choose a class $c$ according to $\rho$ and draw i.i.d. samples $x,x^{+}$ from $\mathcal{D}_{c}$.

• •

  Choose classes $c^{-}_{1},\ldots,c^{-}_{k}$ i.i.d. from $\rho$ and draw $x^{-}_{1}\sim\mathcal{D}_{c^{-}_{1}},\ldots,x^{-}_{k}\sim\mathcal{D}_{c^{-}_{k}}$.

• •

  Output tuple $(x,x^{+},x^{-}_{1:k})$

This data generation is implicit. We only observe the data points and not the associated class labels. Observe that $c^{-}_{i}$ may be the same as $c$ for one or more $i\in\{1,\ldots,k\}$, causing collisions.

Empirically, similar pairs are typically generated from co-occurring data, for example, two parts of a document can be viewed as a similar pair (Tosh et al., 2020). Note that our data generation process implies that $x$ and $x^{+}$ are independent conditioned on the latent class. This is a simplification of the real-world setting where the independence structure might not hold. However, it enables a transparent theoretical analysis that can provide insights that are applicable to some real-world setups.

The objective of an NCE algorithm is to use the available unsupervised data to learn a representation that represents (or embeds) similar points in a similar manner and distinguishes them from the negative samples. The following NCE loss captures this objective,

We restrict $\ell$ to be one of the two standard loss functions, hinge loss $\ell(v)=\max\{0,1+\max_{i}\{-v_{i}\}\}$ and logistic loss $\ell(v)=\log\left(1+\sum\exp(-v_{i})\right)$ for a vector $v\in\mathbb{R}^{k}$. Both loss functions and their variants are deployed routinely in practical NCE implementations (Schroff et al., 2015; Chen et al., 2020).

The algorithm we analyze is the standard empirical risk minimization (ERM) on $\widehat{\mathcal{L}}^{(k)}_{\textup{NCE}}$ over $\mathcal{F}$. For the rest of the paper, we will denote our learned representation $\widehat{f}\in\arg\min_{f\in\mathcal{F}}\widehat{\mathcal{L}}^{(k)}_{\textup{NCE}}(f)$

For evaluating representations, we will consider the standard supervised learning task of classifying a data point into one of the classes in $\mathcal{C}$. More formally, a sample from $\mathcal{D}_{\textup{sup}}$ is generated as follows: (1) choose a class $c$ according to $\rho$ (2) draw sample $x$ from $\mathcal{D}_{c}$, and (3) output sample $(x,c)$.

Observe that the distribution over classes and the conditional distribution over data points is the same as in the unsupervised data. This is crucial to be able to relate the NCE task to the downstream classification task. We will use the corresponding supervised loss for function $g:\mathcal{X}\rightarrow\mathbb{R}^{N}$,

Note that Saunshi et al. (2019) consider a different downstream task which depends on the number of negative samples in the NCE task. Here we disentangle the number of negatives used in NCE from the number of downstream classes which results in a more natural downstream task.

To evaluate the performance of the learned representation function $\widehat{f}$, we will train a linear classifier $W\in\mathbb{R}^{N\times d}$ on top of the learned embedding and use it for the downstream task. We overload notation and define the supervised loss of an embedding $f$ as

We define the loss of the mean classifier as $\mathcal{L}_{\textup{sup}}^{\mu}(f)=\mathcal{L}_{\textup{sup}}(W^{\mu}f)$ where $W^{\mu}$ is such that the $c^{th}$ row is the mean $\mu_{c}=\mathbb{E}_{x\sim\mathcal{D}_{c}}[f(x)]$ of the representations of data from a fixed class $c$. Note that $\mathcal{L}_{\textup{sup}}(f)\leq\mathcal{L}_{\textup{sup}}^{\mu}(f)$ for all $f$. Our guarantees will hold for the loss of the mean-classifier which will directly imply a bound on the overall supervised loss. Note that this analysis may be loose if the mean-classifier itself has poor downstream performance³³3 Saunshi et al. (2019) present a simple example where the two losses are very different..

We train the encoder model $f$ by performing NCE with logistic loss. Following Saunshi et al. (2019), we train on the NCE objective for a fixed number of epochs and do not use a held-out set for model selection. We train a linear classifier on the downstream task for a fixed number of epochs and save the model after each epoch. We do not fine-tune $f$ during downstream training.

Computing $k$ negative tokens for a batch size $B$ results in $\mathcal{O}(Bk)$ forward passes through the model, which is prohibitively expensive for large values of $B$ and $k$. In order to reduce computational overhead, we follow the data reuse trick used in literature (Chen et al., 2020; Chuang et al., 2020). Given $B$ pairs of positively aligned, semantically similar points $\{(x_{i},x^{+}_{i})\}_{i=1}^{B}$, we compute $\{f(x_{i}),f(x^{+}_{i})\}_{i=1}^{B}$ using $2B$ forward passes. For each $i\in[B]$, we treat the set $\mathcal{C}_{i}=\{x_{j},x^{+}_{j}\mid j\in[B],j\neq i\}$ as the candidate set of negative samples to pick from for the positive pair $(x_{i},x^{+}_{i})$. Previous work use the entire candidate set $\mathcal{C}_{i}$, which gives $k=|\mathcal{C}_{i}|=2(B-1)$. While this allows most negative samples to be used, it results in a confounding affect, due to the entanglement of $B$ and $k$. Instead, for a fixed value of $k\leq 2(B-1)$, we sample $k$ negative examples without replacement from $\mathcal{C}_{i}$ independently for each $i$. This sampling protocol marks a notable departure from prior work by decoupling the dependence between batch size and number of negative examples and studying their influence on performance separately.

Here we perform a thorough analysis on the Wiki-3029 dataset of Saunshi et al. (2019). This dataset consists of 200 sentences taken from 3029 Wikipedia articles. The downstream task consists of predicting the article identity from a given sentence. A key advantage of this dataset is that it contains many classes, which allows us to study learning behavior as the number of classes $N$ increases. For a given $N\leq 3029$, for both training and test datasets, we remove sentences from all but the first $N$ articles. We tokenize sentences using the NLTK English word tokenizer (Bird et al., 2009), and replace all tokens that occur less than 50 times in the training set, with token ${\tt UNK}$ denoting an unknown token. This gives us a vocabulary size $V$ of 15,807 with at most 136 unique tokens per sentence. We create our training set by randomly sampling 80% of the sentences from each article. The remaining 20% of the sentences for each article form the test set.

We use a bag-of-words representation for modeling $f$. Formally, we map a sentence $x$ to a list of tuples $\{(t_{i},p_{i})\}$ where $\{t_{i}\}$ is the set of unique tokens in $x$ and $p_{i}$ is the number of times $t_{i}$ occurs in $x$ normalized by the number of tokens in $x$. We compute $f(x)=\sum_{i}p_{i}\nu_{t_{i}}$ where $\nu\in\mathbb{R}^{V\times d}$ is a word-embedding matrix with embedding dimension $d$, and $\nu_{t}$ is the $t^{th}$ row denoting the word embedding of token $t$. We initialize the word embedding matrix randomly and train it using NCE.

We train the NCE model for a fixed number of epochs and do not use a held-out set, following the protocol of Saunshi et al. (2019). We measure both the downstream performance of the mean classifier $W^{\mu}$ and a trained linear classifier. For the former, we use the entire training dataset to compute $W^{\mu}$ using the trained representation. For the latter, we use the entire training dataset, with labels, to fit a linear layer on top of the model $f$, and we do not fine-tune $f$. For each setting, we repeat the experiment 5 times with different seeds. Please see Appendix D for full details on hyperparameters values and optimization details.

The leftmost plot in Figure 3 shows the accuracy of the mean classifier as we vary $N$ and $k$. For a given number of classes $N$, performance improves as $k$ increases and then slowly decays for larger values of $k$. We show the general trend in this figure and report exact numbers to the appendix. Further, we find that the optimal choice of $k$, indicated by solid black circles, slowly increases as $N$ increases. Both of these empirical findings support the theoretical predictions of our improved error transfer bounds. As we have discussed, the previous theoretical results (Saunshi et al., 2019) do not predict this observed trend. Lastly, we observe that performance improves very rapidly as $k$ increases, up until $k=50$, while the decay in performance as $k$ increases beyond the optimal value is relatively mild in comparison to these initial gains.

The center plot in Figure 3 shows the mean classifier’s performance as we vary $B$ and $k$ for two values of $N$. Note that for a given value of $B$, we can only increase $k$ up to $2(B-1)$. We observe that increasing the batch size can improve the performance even when the number of negative examples are fixed. Optimal performance is obtained for the largest batch size and an intermediate value of $k$. This highlights the very different roles played by $B$ and $k$ on this dataset. We note that these results could not have been observed by previous work that tie the value of $k$ to $B$.

Our theoretical analysis reveals a coverage-collision trade-off and suggests that the poor performance when $k$ is very large is due to many collision between the negative samples and the anchor/positive, i.e., false negatives. To empirically verify this claim, we modify the representation learning procedure by sampling $k$ negative examples only from those members in the candidate set that do not share the same class as the positive reference point.⁴⁴4If we are unable to find $k$ true negatives for any example in the batch, then no update is performed on this batch. However, due to the choice of $N,B,k$ this never occurred in our experiments. In the rightmost plot of Figure 3, we compare this sampling procedure with the standard one for different values of $N$ and $k$ and a fixed $B=1000$. We observe that artificially removing collisions improves the performance across all values of $N$; however, the effect is more significant for smaller values of $N$ ($1.5\%$ for $N=200$ and $K=1000$). That the effect is more significant for small $N$ is to be expected, since the probability of collision decreases rapidly with $N$, so the different sampling procedures are quite similar for large $N$. We can further notice that performance decays only slightly for very large values of $k$ in the absence of collisions. This evidence supports the collision hypothesis, namely that collisions are a central reason for performance degradation with increasing $k$. This observation was also made by Chuang et al. (2020). Lastly, as hypothesized by prior work (Chen et al., 2020), we believe a mild reduction in accuracy with large values of $k$ can happen even in the absence of collision, due to instability in optimization

In this section, we decouple batch size from $k$ and measure downstream performance as a function of these parameters on the SVHN and CIFAR-10 datasets. We fit all parameters except the last layer of a ResNet-18 model using the NCE loss, and we evaluate downstream accuracy by training the last layer only in a fully supervised fashion. We use a learning rate of $10^{-3}$, train for 400 epochs, and fit parameters with the Adam optimizer.

We consider two ways of sampling tuples of images for the NCE objective. In the first, we adhere to our theoretical framework and use class information to capture semantic similarity. Here, for a given sample, we use an image from the same class to act as a positive, and we do not use data augmentation. In the second, we deviate from our theory and do not use class information. Instead, we follow a common empirical approach and use an augmented version of the anchor point. Negative samples are also augmented. For these experiments, we exercise the SimCLR data augmentation scheme, consisting of randomly applied crops, color jitter, horizontal flips, greyscaling, and blurring (Chen et al., 2020).

We experiment with both SVHN and CIFAR-10 image datasets, with Figure 4 summarizing our results. When using augmentation, the trend shown in previous work is observed, with accuracy tending to improve with increasing batch size and $k$. When using class information for sampling positives, however, the performance trends are somewhat unexpected. Downstream SVHN accuracy in this setting is remarkably consistent, with nearly all combinations of $k$ and batch size producing roughly equivalent, high-quality representations for downstream classification. In the CIFAR-10 case, this is not true—in fact, performance tends to increase as batch size decreases. This is likely due to the exclusion of data augmentation, which is necessary for obtaining reasonable performance even in the fully supervised regime for these architectures. In its absence, the variance induced by decreasing the batch size appears to offer a regularization benefit that improves accuracy. Still, overall performance without augmentation is poor in comparison to data-augmented NCE.

These experiments raise an interesting conjecture: perhaps the observation that downstream performance increases with $k$ and batch size can be attributed to the kinds of data augmentation techniques that have been refined specifically for vision, rather than to the NCE objective itself. In NLP, where there are less clear-cut ways of performing data augmentation, we see trends that more closely agree with our theory. Either way, these experiments point to a distortion between the theoretical framework typically used to study NCE (using true positives) and, in the vision case, what is observed in practice. The core of this distortion appears to involve the use of augmentation techniques that have been highly tuned for object recognition/computer vision tasks.

As alluded to by our vision experiments, the use of data augmentation in practice seems to be a principal factor in the gap between what is observed empirically and what we describe theoretically. Because NCE is so commonly used with image data, developing a theoretical model that can incorporate properties of commonly used augmentations and the inductive bias of networks is an important avenue for future research.

When the NCE distribution and the downstream task are closely linked, our theoretical results indicate that the number of negatives used in NCE should scale with the number of classes. In practice, however, choosing $k$ is less clear, with the optimal choice being a function of complex factors. The model’s inductive bias, the choice of data augmentation used, and the data type all seem to influence the ideal value. Given this, are there algorithmic principles for setting $k$ in practice?

This work disentangles $k$ from the batch size used for model training, with our experiments suggesting that each plays a different role in the quality of the learned representation. Modern NCE pipelines have several conflating parameters, and we believe that a deeper empirical study is imperative for disambiguating and understanding the role of each.