12.3 Contrastive losses

The relative entropy of the posterior distributions of the data and model—Eqs. 12.20 and 12.22, respectively—is

This is evidently a discriminative problem, but with a twist. The canonical generative approach to binary classification is to model the generative distribution ${\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}x};\bm{\theta}}\right)}{\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}\middle|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}x};\bm{\theta}}\right)}$ (like NCE); acquire the parameters by minimizing the joint relative entropy $\operatorname*{\text{D}_{\text{KL}}}\mathopen{}\mathclose{{}\left\{{p\mathopen{}\mathclose{{}\left({X},{\bm{Y}}}\right)}\middle\|{\hat{p}\mathopen{}\mathclose{{}\left({X},{\bm{Y}};\bm{\theta}}\right)}}\right\}$ (unlike NCE); and then invert to ${\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}x}\middle|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}};\bm{\theta}}\right)}$ with Bayes’ rule. For example, for Gaussian mixtures, this is known as linear/quadratic discriminant analysis (depending on whether the covariance is the same/different across classes). The canonical discriminative approach to binary classification is to model ${\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}x}\middle|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}};\bm{\theta}}\right)}$ directly (unlike NCE); and then minimize the posterior relative entropy $\operatorname*{\text{D}_{\text{KL}}}\mathopen{}\mathclose{{}\left\{{p\mathopen{}\mathclose{{}\left({X}\middle|{\bm{Y}}}\right)}\middle\|{\hat{p}\mathopen{}\mathclose{{}\left({X}\middle|{\bm{Y}};\bm{\theta}}\right)}}\right\}$ (like NCE). This is logistic regression. NCE mixes both methods: it models the generative distribution ${\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}x};\bm{\theta}}\right)}{\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}\middle|{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}x};\bm{\theta}}\right)}$, but first inverts with Bayes rule, and finally minimizes the posterior relative entropy $\operatorname*{\text{D}_{\text{KL}}}\mathopen{}\mathclose{{}\left\{{p\mathopen{}\mathclose{{}\left({X}\middle|{\bm{Y}}}\right)}\middle\|{\hat{p}\mathopen{}\mathclose{{}\left({X}\middle|{\bm{Y}};\bm{\theta}}\right)}}\right\}$. In the classic case of the mixture of two Gaussians/binary classification, this would amount to learning the two (mean, covariance) pairs by minimizing the posterior relative entropy—as opposed to learning these parameters by minimizing the joint relative entropy (generative), or learning a separating hyperplane by minimizing the posterior relative entropy (discriminative). These contrasts are summarized in Table 12.1.^††margin: Nice properties of the estimator….

Suppose that the observation $\bm{y}$ is not a single sample but a collection of ${K}$ “examples,” $(\bm{y}_{1},\ldots,\bm{y}_{{K}}$), precisely one of which is not noise. Then the goal is not to determine whether or not the sample is noise, but rather to determine which of the examples is noise. This means that rather than use the model-noise energy difference (Eq. 12.23) directly to assign the example to the positive or negative class, as in NCE, we will compare ${K}$ energy differences to each other (with the softmax function).

In this setup, the latent variable is categorical (conceived as a one-hot vector ${\bm{\hat{X}}}$) rather than Bernoulli, and the data distribution is:

Again we have set the prior uniform, since we have no reason to make any one of the elements more or less likely to be noise than any other. We emphasize that this is not a mixture model: a single sample contains ${K}$ “examples”: one positive, and ${K}-1$ negative.

The generative model takes the same form, with the model distribution taking the place of the data marginal. Writing it in terms of energies, we obtain

Again we ignore the fact that the emission is unnormalized and simply compute a (normalized) posterior distribution with Bayes’ rule

that is, the $i^{\text{th}}$ output of the softmax function. Eq. 12.26 is evidently a kind of generalization of Eq. 12.22.¹⁰¹⁰ 10 However, note that the multi-example version of NCE does not quite reduce to the single-example case even when ${K}=2$. Eq. 12.22 can indeed be re-written with a softmax as in Eq. 12.26, with the first argument equal to $h\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{i}},\bm{\theta}}\right)$ and the second equal to 0. The latter reflects our indifferent prior, which provides no additional information. In the two-example version of the generalization under discussion, on the other hand, the second argument encodes the relative probability of the second example being data or noise. In short, deciding which of two samples is “real” is easier than deciding whether or not a single sample is. Putting this together with the data distribution, we can write the posterior relative entropy as

In the final line, we are selecting only that output of the softmax function that corresponds to the actual positive sample (whose index will of course differ from trial to trial).

We can shed light on the role played by the negative examples by considering them separately from the positive example in the posterior probability of a positive example:

Now notice that the negative-sample terms sum approximately to a constant:

The approximate equality becomes more exact as the number of negative examples increases. (And technically, the final equality requires the model and noise distributions to have the same support.) Eq. 12.28 says that, if we had in hand an expression for the normalizer, we could do without the negative samples altogether—they drop out of the loss function. Indeed, the loss (Eq. 12.27) now becomes

where the final line follows for large ${K}$.¹¹¹¹ 11 The authors of the original paper [75] interpret this approximation as a lower bound when the model distribution matches the data distribution. Presumably the idea is that, for a very good model ${\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}};\bm{\theta}}\right)}$, the noise-to-model ratio will usually be less than one when evaluated on positive examples. Therefore, the neglected $+1$ will dominate the neglected $-1$. It would take more work to prove this. This makes sense: the whole point of using negative examples was to force unnormalized models to learn the correct normalization. Since we want to use models for which computing $Z(\bm{\theta})$ is intractable, we will not use Eq. 12.29 as our objective—but we will use it below to prove that optimizing the multi-example NCE loss (Eq. 12.27) increases mutual information in a certain setting.

We have assumed up to this point that the source of our “noise” samples is also an evaluatable expression for the probability of samples. What if we have only samples from the noise distribution? Can we still learn a model of the positive data?

One obvious solution is to learn a model for the negative as well as the positive samples; for example, to build a parameterized model for the noise energy, $\hat{U}_{\text{n}}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}},\bm{\theta}}\right)$, and use it in the generative model. But if we wanted to get a normalized version of the model energy, $\hat{U}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}},\bm{\theta}}\right)$, we would have to be able to get or to know the normalizer for this noise energy, $\hat{U}_{\text{n}}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}},\bm{\theta}}\right)$, which is troubling. However, as noted at the outset, getting a probability model for the data, normalized or unnormalized, is not the goal of InfoNCE. So instead we will directly model the energy difference, i.e. the left-hand rather than right-hand side of Eq. 12.23. Rather than asking for the probabilities of an example $\bm{y}_{k}$ under the two models (positive and negative), we are asking for its relative probability.

One subtlety with modeling $h\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}},\bm{\theta}}\right)$ directly is that we are still at liberty to interpret this as fitting $\hat{U}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}},\bm{\theta}}\right)$ only, that is to say, not fitting the noise energy, $U_{\text{n}}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}}\right)$. In other words, we can attribute any error in $h\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}},\bm{\theta}}\right)$ to an error in $\hat{U}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}},\bm{\theta}}\right)$ rather than $U_{\text{n}}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}}\right)$. Consequently, the denominator in Eq. 12.28 can still be interpreted as ${p_{\text{n}}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}}}\right)}$, and the equation still goes through. We will use it below.

In the third departure from the original NCE, the InfoNCE method proposes to learn to distinguish a joint distribution, ${p\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{a}},{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{b}}}\right)}$, from the product of its marginal distributions, ${p_{\text{n}}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{a}},{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{b}}}\right)}\mathrel{\vbox{\hbox{.}\hbox{.} }}={p\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{a}}}\right)}{p\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{b}}}\right)}$. In the motivating example at the outset, we considered samples from the joint to be two views of the same object, and samples from the product of marginals to be two views of two different objects. The intuition behind this choice of distributions is that a model that can distinguish these distributions must have learned (something of) the relationship between ${\bm{Y}}_{a}$ and ${\bm{Y}}_{b}$. It must have learned to extract the features that are related.

This can be made precise in the language of information theory. To do so, we work with the approximate expression for the loss, Eq. 12.29, that holds for sufficiently many samples ${K}$. In particular, let us reverse the sign of the loss and consider maximizing it:

with the average taken under samples from the data joint distribution (not the product of marginals). In the last equality, we have exploited the subtlety discussed in the previous section: Since we are modeling only the ratio of the positive and negative distributions (i.e., the difference in their energies), we can interpret the latter as fixed and correct (equal to the negative data distribution), and our updates as improving only the former.

Ignoring the $\log{K}$ term, this expression looks suspiciously similar to the mutual information between ${\bm{\hat{Y}}}_{a}$ and ${\bm{\hat{Y}}}_{b}$. However, the information we aim to increase is not precisely a mutual information, neither under the model nor under the data, because it depends on both. Indeed, neither would be sensible targets for maximization: There is no model distribution ${\hat{p}\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{a}},{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{b}};\bm{\theta}}\right)}$ from which we might draw samples or take expectations (we model only a ratio of distributions); and the data mutual information is not something we control. Nevertheless, there is a relationship with the data mutual information:

with the inequality following from Gibbs’s inequality: $\text{H}_{p}{\mathopen{}\mathclose{{}\left[{\bm{Y}}_{a},{\bm{Y}}_{b}}\right]}\leq\text{H}_{p\hat{p}}{\mathopen{}\mathclose{{}\left[{\bm{Y}}_{a},{\bm{Y}}_{b};\bm{\theta}}\right]}$. The quantity on the right, $\mathcal{I}_{p\hat{p}}\mathopen{}\mathclose{{}\left({\bm{Y}}_{a};{\bm{Y}}_{b}}\right)$, replaces the joint entropy with the joint cross entropy, and might therefore be called (for want of something better) the “cross mutual information.” Intuitively, $\mathcal{I}_{p\hat{p}}\mathopen{}\mathclose{{}\left({\bm{Y}}_{a};{\bm{Y}}_{b}}\right)$ is how much the uncertainty about ${\bm{Y}}_{a}$ and ${{\bm{Y}}_{b}}$ can be reduced by using the model $\hat{p}$ of the joint, rather than using perfect models of the marginals.¹²¹² 12 This is not the interpretation given in the original paper [75].

Stringing together these equations, we find that

Hence, minimizing the NCE loss (maximizing the negative loss) maximizes, at least approximately, the cross mutual information, up to a maximum value given by the true (data) mutual information. From Eq. 12.28, the larger ${K}$, the less approximate this relationship. (This also increases the $\log{K}$ term in Eq. 12.30 and therefore the discrepancy between the cross mutual information and the loss, but this has no effect on the their gradients.) Consequently, the algorithm that makes these changes to the original NCE algorithm is referred to as InfoNCE.

It is therefore worth emphasizing that the information maximization given by Eq. 12.30 is not unique to InfoNCE.¹³¹³ 13 Here we again depart from the original paper [75]. Indeed, it is a consequence of any generic density-estimation algorithm. Consider the standard approach of minimizing the relative entropy of model and joint distributions:

The second equality follows because the (data) mutual information is constant with respect to the model parameters (and the final equality is algebra). What is important about InfoNCE is that it accomplishes this information maximization with an energy-based model, a potentially very flexible approach.

Now recall that the goal of InfoNCE is to extract useful features from the data. If the model is to maximize mutual information between a pair of variables, it will certainly need to extract out what is common between them. Assuming that this common feature is indeed what we want to expose to subsequent processing, we need to make sure that it is represented in an accessible way. Van den Oord and colleagues therefore propose [75] to allow the variables to interact only through an inner product

Thus, if learning has been successful, the feature of interest will be a linear function of the outputs of $f_{a}$ and $f_{b}$.

There are many possibilities, but one nice application of InfoNCE is to time-series data, and in particular with learning to extract the features that relate a history of observations, ${\bm{Y}}_{1},\ldots,{\bm{Y}}_{t}{}$, to a future observation from the same series, ${\bm{Y}}_{t+s}{}$ (for some positive integer $s$). We can think of this as learning to “summarize” sequences of random variables so as best to predict their future state. (For example, for linear dynamical systems, the optimal summary is a weighted sum of past states, with weights decaying exponentially into the past.) Thus the positive and negative (“noise”) distributions are, respectively, $p\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{t+s}}{},{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{1}},\ldots,{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{t}}{}}\right)$ and ${p\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{t+s}}}\right)}{p\mathopen{}\mathclose{{}\left({\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{1}},\ldots,{\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\bm{y}_{t}}{}}\right)}$. For this use case, and in compliance with Eq. 12.31, the authors let the energy difference be

where $f_{\text{s}}$ is a static “encoder” ANN and $f_{\text{RNN}}$ is an RNN.^††margin: A diagram would be nice. In order to decrease the posterior relative entropy (Eq. 12.27), the encoder and the RNN must extract representations from the data history (on the one hand) and a future state (on the other) that expose the shared information between them to a bilinear form. The parameters $\bm{\theta}$ and $\mathbf{W}$ are all learned by stochastic gradient descent of Eq. 12.27.