Multi-Similarity Contrastive Learning

Our work draws from existing literature in contrastive representation learning. Many of the current state-of-the-art vision and language models are trained using contrastive losses [30, 39, 8, 21, 13]. Self-supervised contrastive learning methods, such as MoCo and SimCLR, maximize agreement between two different augmentations or views of the same image [13, 8]. Recently, vision-language contrastive learning has allowed dual-encoder models to pretrain with hundreds of millions of image-text pairs [18, 30]. The resulting learned embeddings achieve state-of-the-art performance on many vision and language benchmarks [39, 35, 22]. Supervised contrastive learning, SupCon, allows contrastive learning to take advantage of existing labels [21]. Contrastive learning has also been adapted to learn from both labels and text [36] and from hierarchies of labels [41]. The method most similar to ours is conditional similarity networks [34]. In conditional similarity networks, masks are learned or assigned to different embedding dimensions with respect to different metrics of similarity. These masks are learned jointly with the convolutional neural network parameters during training time. Conditional similarity networks differ from our work in two major ways. First, unlike our work, conditional similarity networks uses triplet loss, a specialized version of contrastive loss. At training time, it requires triplets based on each similarity metric. Second, we automatically learn separate projection spaces and weights for each metric of similarity, whereas they learn a linear transformation from the embedding space for each similarity and do not consider weighting metrics. As we show in Section 4.3, our multiple similarity contrastive networks consistently outperforms conditional similarity networks.

Multi-task learning aims to simultaneously learn multiple related tasks and often outperforms learning each task alone [20, 6, 24]. However, if tasks are weighted improperly during training, the performance on some tasks suffer. Various learned task weighting methods have been proposed for multi-task learning in the vision and language domains [26, 20, 9, 32, 23, 24, 25]. These methods learn task weightings based on different task characteristics in order to improve the generalization performance towards novel tasks [26]. This is done by regularizing the task variance using gradient descent [9, 25] or by using adversarial training to divide models into task-specific and generalizable parameters [23]. Overwhelmingly, these methods are built for multiple tasks trained with likelihood-based losses, such as regression and classification. One of the most popular of these methods models task uncertainty to determine task-specific weighting and automatically learns weights to balance this uncertainty [20]. In our work, we adapt automatically learned task weighting to our multi-similarity contrastive loss by predicting similarity uncertainty. This is not straightforward since the contrastive loss is trained in a pairwise fashion and there is a lack of absolute labels in the learned output (a set of embedding vectors) [4].

Adapting uncertainty estimation techniques to contrastive learning remains an active area of research. This is because contrastive learning learns abstract embedding vectors rather than absolute labels and because of the pairwise training of contrastive models. Given access to training data and labels, previous work has proposed estimating the density and consistency of the hypersphere embedding space distribution as metrics to estimate uncertainty [4, 29]. The density of the embedding space at a point captures the amount of data the model has observed during training, and can serve as a proxy for epistemic, or model, uncertainty [29]. The consistency of the embedding space at a point uses k-nearest-neighbors to measure the extent to which the training data mapped closest to that point have consistent labels, and can serve as a proxy for aleatoric, or data-dependent, uncertainty [4]. Other recent work proposes learned temperature as a metric of heteroscedastic, or input-dependent, uncertainty to identify out-of-distribution data for labeled datasets [40]. To our knowledge, we are the first work to model similarity-dependent uncertainty, or the relative confidence between different training tasks, in the contrastive setting.

We assume that during training time, we have access to dataset: $\mathcal{D}=\{x_{i},\textbf{Y}_{i}\}_{i}^{M}$, where $x$ is an image and the $\textbf{Y}_{i}=\{y_{i}^{1}...y_{i}^{C}\}$ are distinct categorical attributes associated with the image. We aim to learn an embedding function $f:x\rightarrow\mathbb{R}^{d}$ that maps $x$ to an embedding space. We define $h_{i}=f(x_{i})$ to be the embedding of $x_{i}$.

In the typical contrastive training setup, training proceeds by selecting a batch of $N$ randomly sampled data $\{x_{i}\}_{i=1...N}$. We randomly sample two distinct label preserving augmentations (e.g., from rotations, crops, flips) for each $x_{i}$, ($\tilde{x}_{2i}$ and $\tilde{x}_{2i-1}$), to construct $2N$ augmented samples, $\{\tilde{x}_{j}\}_{j=1...2N}$. Let $A(i)=\{1,...2N\}\backslash i$ be the set of all samples and augmentations not including $i$. We define $g$ to be a projection head that maps the embedding to the similarity space represented as the surface of the unit sphere $\mathbb{S}^{d}=\{v\in\mathbb{R}^{d}:||v||_{2}=1\}$. Finally, we define $v_{i}=g(h_{i})$ as the mapping of $h_{i}$ to the projection space.

Supervised contrastive learning uses labels to implicitly define the positive sets of examples. Specifically, supervised contrastive learning encourages samples with the same label to have similar embeddings and samples with a different label to have different embeddings. We follow the literature in referring to samples with the same label as an image $i$ as the positive samples, and samples with a different label than that of $i$’s as the negative samples.

Supervised contrastive learning (SupCon) [21] proceeds by minimizing the loss:

where $|S|$ denotes the cardinality of the set $S$, $P(i)$ denotes the positive set with all other samples with the same label as $x_{i}$, i.e., $P(i)=\{j\in A(i):y_{j}=y_{i}\}$, $I$ denotes the set of all samples in a particular batch, and $\tau\in\{0,\infty\}$ is a temperature hyperparameter.

In contrast to SupCon, our multi-similarity contrastive (MSCon) approach proceeds by jointly training an embedding space using multiple notions of similarity. We do so by training the embedding with multiple projection heads $g^{c}$ that map the embedding to $C$ projection spaces, where each space distinguishes the image based on a different similarity metric. We define $v^{c}_{i}=g^{c}(h_{i})$ to be the mapping of $h_{i}$ to the projection space by projection head $g^{c}$. Because each projection space is already normalized, we assume that the each similarity loss is similarly scaled. We define the multi-similarity contrastive loss to be a summation of the supervised contrastive loss over all conditions $L^{mscon}=\sum_{c\in C}\sum_{i=I}L^{mscon}_{c,i}$ where each conditional $L^{mscon}_{c,i}$ is defined as in equation 2. Specifically,

where $P^{c}(i)$ is defined as the positive set under similarity $c$ such that for all $j\in P^{c}(i)$, $y_{j}^{c}=y_{i}^{c}$.

In the above formulation of our multi-similarity contrastive loss function, each similarity is weighted equally. However, previous work in multi-task learning for both vision and language have demonstrated that model performance can deteriorate when one or more of the tasks is noisy or uncertain. One way to tackle this is to learn task weights based on the uncertainty of each task. However, model performance can be sensitive to weight selection [12, 26, 20, 9, 25], and manually searching for optimal weightings is expensive in both computation and time. Previous work has suggested using irreducible uncertainty of task predictions in a weighting scheme. For example, tasks where predictions are more uncertain are weighted lower because they are less informative[20].

Such notions of uncertainty are typically predicated on an assumed parametric likelihood of a label given inputs. However, this work is not easily adapted to multi-similarity contrastive learning because 1) contrastive training does not directly predict downstream task performance and 2) the confidence in different similarity metrics has never been considered in this setting. In contrastive learning, the estimate of interest is a similarity metric between different examples rather than a predicted label, which means that downstream task performance is not directly predicted by training results. Furthermore, previous work in contrastive learning has only focused on modeling data-dependent uncertainty, or how similar a sample is to negative examples within the same similarity metric. To our knowledge, we are the first to utilize uncertainty in the training tasks and their corresponding similarity metrics as a basis for constructing a weighting scheme for multi-similarity contrastive losses.

We do this in two ways: 1) we construct a pseudo-likelihood function approximating task performance and 2) we introduce a similarity dependent temperature parameter to model relative confidence between different similarity metrics. We present an extension to the contrastive learning paradigm that enables estimation of the uncertainty in similarity metrics. In addition to providing useful information about the informativeness of each similarity metric, our estimate of uncertainty enables us to weight the different notions of similarity such that noisy notions of similarity are weighted lower than more reliable notions.

Our approach proceeds by constructing a pseudo-likelihood function which approximates task performance. We show in the supplement that maximizing our pseudo-likelihood also maximizes our MSCon objective function. This pseudo-likelihood endows the approach with a well-defined notion of uncertainty that can then be used to weight the different similarities.

Let $v_{i}^{c}$ be the model projection head output for similarity $c$ for input $x_{i}$. Let $\textbf{Y}^{c}$ be the $c$th column in Y. We define $P^{c}_{y}=\{x_{j}\in\mathcal{D}:\textbf{Y}^{c}_{j}=y\}$ to be the positive set for label $y$ under similarity metric $c$. We define the classification probability $p(y|v_{i}^{c},D,\tau)$ as the average distance of the representation $v_{i}^{c}$ from all representations for inputs conditioned on the similarity metric. Instead of directly optimizing equation 1, we can maximize the following pseudo-likelihood:

Note that optimizing 3 is equivalent to optimizing 1 by applying Jensen’s inequality (as shown in the supplement). By virtue of being a pseudo-likelihood, equation 3 provides us with a well-defined probability associated with downstream task performance that we can use to weight the different tasks. We will next outline how to construct this uncertainty from the pseudo-likelihood defined in equation 3.

We assume that $v^{c}$ is a sufficient statistic for $y^{c}$, meaning that $y^{i}$ is independent of all other variables conditional on $v^{i}$. Such an assumption is not unrealistic, it simply reflects the notion that $v^{c}$ is an accurate estimation for $y^{c}$. Under this assumption the pseudo-likelihood expressed in 3 factorizes as follows:

Previous work in contrastive learning modifies the temperature to learn from particularly difficult data examples [40, 31]. Inspired by this, we adapt the contrastive likelihood to incorporate a similarity-dependent scaled version of the temperature. We introduce a parameter $\sigma_{c}^{2}$ for each similarity metric controlling the scaling of temperature and representing the similarity dependent uncertainty in Equation 5.

The negative log-likelihood for this contrastive likelihood can be expressed as Equation 6.

We provide a detailed derivation of this equation in the supplement. Extending this analysis to consider multiple similarity metrics, we can adapt the optimization objective to learn weightings for each similarity as in Equation 7.

During training, we learn the $\sigma_{c}$ weighting parameters through gradient descent.

In this subsection, we evaluate the robustness of our learned embeddings to similarity uncertainty. Since the true level of task noise (similarity metric uncertainty) is unobserved, we use a semi-simulated approach, where we simulate uncertain similarities in both the Zappos50k and MEDIC datasets.

For the Zappos50k dataset, we train the encoder using the category, closure, and gender similarity metrics. To introduce task uncertainty, we randomly corrupt the closure task by proportion $\rho$. We randomly sample $\rho$ of the closure labels and randomly reassign the label amongst all possible labels. Note that when $\rho=1.0$, all labels are randomly sampled equally from the available closure labels. When $\rho=0.0$, all labels are identical to the original dataset. For the MEDIC dataset, we train the encoder using the disaster types, humanitarian, and informative similarity metrics. We corrupt the disaster type task in order to introduce task uncertainty. As $\rho$ increases in Figure 3, we find that MSCon learns to down-weight the noisy task for both the Zappos50k and MEDIC datasets.

For the Zappos50k dataset, we evaluate the top-1 classification accuracy on an out-of-domain task, brand classification, and on an in-domain task, the corrupted closure classification. Similarly, for the MEDIC dataset, we evaluate the top-1 classification accuracy on an out-of-domain task, damage-severity classification, and on an in-domain task, the corrupted disaster-type classification.

Figure 3 shows the results from this analysis. The top panel shows how the weights change as we change task uncertainty on the x-axis. The middle and bottom panels shows how out-of-domain and in-domain evaluation accuracy changes as we change task uncertainty. As expected, as $\rho$ increases to $1$, the in-domain classification accuracy for both the equal-weighted and weighted MSCon learned embeddings decreases to random.

However, the out-of-domain classification accuracy for the weighted MSCon learned embeddings is more robust to changes in $\rho$ than the unweighted MSCon learned embeddings. This is because the weighted version of MSCon automatically learns to down-weight uncertain or more uninformative tasks during encoder training.

In this section, we evaluate in- and out-of-domain performance of various methods. We find that our multi-similarity contrastive network significantly outperforms all other contrastive methods on in-domain tasks and outperforms multi-task cross-entropy learning on out-of-domain tasks. We also show how performance changes with variation in hyperparameter selection. More qualitative analysis of the learned similarity subspaces (i.e., TSNE visualizations) can be found in the supplement.