Conditional Contrastive Learning with Kernel

Tsai et al. (2021a) consider the auxiliary information from data (e.g., annotation attributes of images) as a weak supervision signal and propose a contrastive objective to incorporate the weak supervision in the representations. This work is motivated by the argument that the auxiliary information implies semantic similarities. With this motivation, the weakly supervised contrastive learning framework learns similar representations for data with the same auxiliary information and dissimilar representations for data with different auxiliary information. Embracing this idea, the original InfoNCE objective can be modified into the weakly supervised InfoNCE (abbreviated as WeaklySup-InfoNCE) objective:

Here $Z$ is the conditioning variable representing the auxiliary information from data, and $z$ is the outcome of auxiliary information that we sample from $P_{Z}$. $(x,y_{\rm pos})$ are positive pairs sampled from $P_{X|Z=z}P_{Y|Z=z}$. In this design, the positive pairs always have the same outcome from the conditioning variable. $\{(x,y_{\rm neg,i})\}_{i=1}^{n}$ are negative pairs that are sampled from $P_{X}P_{Y}$.

Another recent work (Tsai et al., 2021c) presented to remove undesirable sensitive information (such as gender) in the representation, by sampling negative pairs conditioning on sensitive attributes. The paper argues that fixing the outcome of the sensitive variable prevents the model from using the sensitive information to distinguish positive pairs from negative pairs (since all positive and negative samples share the same outcome), and the model will ignore the effect of the sensitive attribute during contrastive learning. Embracing this idea, the original InfoNCE objective can be modified into the Fair-InfoNCE objective:

Here $Z$ is the conditioning variable representing the sensitive information (e.g., gender), $z$ is the outcome of the sensitive information (e.g., female), and the anchor data $x$ is associated with $z$ (e.g., a data point that has the gender attribute being female). $(x,y_{\rm pos})$ are positively-paired that are sampled from $P_{XY|Z=z}$, and $x$ and $y_{\rm pos}$ are constructed to have the same $z$. $\{(x,y_{\rm neg,i})\}_{i=1}^{n}$ are negatively-paired that are sampled from $P_{X|Z=z}P_{Y|Z=z}$. In this design, the positively-paired samples and the negatively-paired samples are always having the same outcome from the conditioning variable.

Robinson et al. (2020) and Kalantidis et al. (2020) argue that contrastive learning can benefit from hard negative samples (i.e., samples $y$ that are difficult to distinguish from an anchor $x$). Rather than considering two arbitrary data as negatively-paired, these methods construct a negative data pair from two random data that are not too far from each other¹¹1Wu et al. (2020) argues that a better construction of negative data pairs is selecting two random data that are neither too far nor too close to each other.. Embracing this idea, the original InfoNCE objective is modified into the Hard Negative InfoNCE (abbreviated as HardNeg-InfoNCE) objective:

Here $Z$ is the conditioning variable representing the embedding feature of $X$, in particular $z=g_{\theta_{X}}(x)$ (see definition in equation 2, we refer $g_{\theta_{X}}(x)$ as the embedding feature of $x$ and $g_{\theta_{Y}}(y)$ as the embedding feature of $y$). $(x,y_{\rm pos})$ are positively-paired that are sampled from $P_{XY}$. To construct negative pairs $\{(x,y_{{\rm neg},i})\}_{i=1}^{n}$, we sample $\{(y_{\rm neg,i})\}_{i=1}^{n}$ from $P_{Y|Z=z=g_{\theta_{X}}(x)}$. We realize the sampling from $P_{Y|Z=z=g_{\theta_{X}}(x)}$ as sampling data points from $Y$ whose embedding features are close to $z=g_{\theta_{X}}(x)$: sampling $y$ such that $g_{\theta_{Y}}(y)$ is close to $g_{\theta_{X}}(x)$.

The high level idea of our method is that, instead of sampling $y$ from $P_{Y|Z=z}$, we sample $y$ from existing data of $Y$ whose associated conditioning variable’s outcome is close to $z$. For example, assuming the conditioning variable $Z$ to be age and $z$ to be $80$ years old, instead of sampling the data points at the age of $80$ directly, we sample from all data points, assigning highest weights to the data points at the age from $70$ to $90$, given their proximity to $80$. Our intuition is that data with similar outcomes from the conditioning variable should be used in support of the conditional sampling. Mathematically, instead of sampling from $P_{Y|Z=z}$, we sample from a distribution proportional to the weighted sum $\sum_{j=1}^{N}w(z_{j},z)P_{Y|Z=z_{j}}$, where $w(z_{j},z)$ represents how similar in the space of of the conditioning variable $Z$. This similarity is computed for all data points $j=1\cdots N$. In this paper, we use the Kernel Conditional Embedding Operator (Song et al., 2013) for such approximation, where we represent the similarity using kernel (Schölkopf et al., 2002).

We want to avoid the conditional sampling procedure from existing conditional learning objectives (equation 3, 4, 5), and hence we are not supposed to have access to data pairs from the conditional distribution $P_{X|Z}P_{Y|Z}$. Instead, the only given data will be a batch of triplets $\{(x_{i},y_{i},z_{i})\}_{i=1}^{b}$, which are independently sampled from the joint distribution $P_{XYZ}^{\,\,\,\otimes b}$ with $b$ being the batch size. In particular, when $(x_{i},y_{i},z_{i})\sim P_{XYZ}$, $(x_{i},y_{i})$ is a pair of data sampled from the joint distribution $P_{XY}$ (e.g., two augmented views of the same image) and $z_{i}$ is the associated conditioning variable’s outcome (e.g., the annotative attribute of the image). To convert previous objectives into alternative forms that avoid the need of the conditional sampling procedure, we need to perform an estimation of the scoring function $e^{f(x,y)}$ for $(x,y)\sim P_{X|Z}P_{Y|Z}$ in equation 3, 4, 5 given only $\{(x_{i},y_{i},z_{i})\}_{i=1}^{b}\sim P_{XYZ}^{\,\,\,\otimes b}$.

We present to reformulate previous objectives into kernel expressions. We denote $K_{XY}\in\mathbb{R}^{b\times b}$ as a kernel gram matrix between $X$ and $Y$: let the $i_{\rm th}$ row and $j_{\rm th}$ column of $K_{XY}$ be the exponential scoring function $e^{f(x_{i},y_{j})}$ (see equation 2) with $[K_{XY}]_{ij}:=e^{f(x_{i},y_{j})}$. $K_{XY}$ is a gram matrix because the design of $e^{f(x,y)}$ satisfies a kernel³³3Cosine similarity is a proper kernel, and the exponential of a proper kernel is also a proper kernel. between $g_{\theta_{X}}(x)$ and $g_{\theta_{Y}}(y)$:

where $\langle\cdot,\cdot\rangle_{\mathcal{H}}$ is the inner product in a Reproducing Kernel Hilbert Space (RKHS) $\mathcal{H}$ and $\phi$ is the corresponding feature map. $K_{XY}$ can also be represented as $K_{XY}=\Phi_{X}\Phi_{Y}^{\top}$ with $\Phi_{X}=\big{[}\phi\big{(}g_{\theta_{X}}(x_{1})\big{)},\cdots,\phi\big{(}g_{\theta_{X}}(x_{b})\big{)}\big{]}^{\top}$ and $\Phi_{Y}=\big{[}\phi\big{(}g_{\theta_{Y}}(y_{1})\big{)},\cdots,\phi\big{(}g_{\theta_{Y}}(y_{b})\big{)}\big{]}^{\top}$. Similarly, we denote $K_{Z}\in\mathbb{R}^{b\times b}$ as a kernel gram matrix for $Z$, where $[K_{Z}]_{ij}$ represents the similarity between $z_{i}$ and $z_{j}$: $[K_{Z}]_{ij}:=\big{\langle}\gamma(z_{i}),\gamma(z_{j})\big{\rangle}_{\mathcal{G}}$, where $\gamma(\cdot)$ is an arbitrary kernel embedding for $Z$ and $\mathcal{G}$ is its corresponding RKHS space. $K_{Z}$ can also be represented as $K_{Z}=\Gamma_{Z}\Gamma_{Z}^{\top}$ with $\Gamma_{Z}=\big{[}\gamma(z_{1}),\cdots,\gamma(z_{b})\big{]}^{\top}$.

We present the following:

$[K_{XY}(K_{Z}+\lambda{\bf I})^{-1}K_{Z}]_{ii}$ is the kernel estimation of $e^{f(x_{i},y)}$ when $(x_{i},z_{i})\sim P_{XZ}\,,\,y\sim P_{Y|Z=z_{i}}$. It defines the similarity between the data pair sampled from $P_{X|Z=z_{i}}P_{Y|Z=z_{i}}$. Hence, $(K_{Z}+\lambda{\bf I})^{-1}K_{Z}$ can be seen as a transformation applied on $K_{XY}$ ($K_{XY}$ defines the similarity between $X$ and $Y$), converting unconditional to conditional similarities between $X$ and $Y$ (conditioning on $Z$). Proposition 2.2 also re-writes this estimation as a weighted sum over $\{e^{f(x_{i},y_{j})}\}_{j=1}^{n}$ with the weights $w(z_{j},z_{i})=\big{[}(K_{Z}+\lambda{\bf I})^{-1}K_{Z}\big{]}_{ji}$. We provide an illustration to compare $(K_{Z}+\lambda{\bf I})^{-1}K_{Z}$ and $K_{Z}$ in Figure 2, showing that $(K_{Z}+\lambda{\bf I})^{-1}K_{Z}$ can be seen as a smoothed version of $K_{Z}$, suggesting the weight $\big{[}(K_{Z}+\lambda{\bf I})^{-1}K_{Z}\big{]}_{ji}$ captures the similarity between $(z_{j},z_{i})$. To conclude, we use the Kernel Conditional Embedding Operator (Song et al., 2013) to avoid explicitly sampling $y\sim P_{Y|Z=z}$, which alleviates the limitation of having insufficient data from $Y$ that are associated with $z$. It is worth noting that our method neither generates raw data directly nor includes additional training.

In terms of computational complexity, calculating the inverse $(K_{Z}+\lambda I)^{-1}$ costs $O(b^{3})$ where $b$ is the batch size, or $O(b^{2.376})$ using more efficient inverse algorithms like Coppersmith and Winograd (1987). We use the inverse approach with $O(b^{3})$ computational cost, which will not be an issue for our method. This is because we consider a mini-batch training to constrain the size of $b$, and the inverse $(K_{Z}+\lambda I)^{-1}$ does not contain gradients. The computational bottlenecks are gradients computation and their updates.

We short hand $K_{XY}(K_{Z}+\lambda{\bf I})^{-1}K_{Z}$ as $K_{X\perp\!\!\!\perp Y|Z}$, following notation from prior work (Fukumizu et al., 2007) that shorthands $P_{X|Z}P_{Y|Z}$ as $P_{X\perp\!\!\!\perp Y|Z}$. Now, we plug in the estimation of $e^{f(x,y)}$ using Proposition 2.2 into WeaklySup-InfoNCE (equation 3), Fair-InfoNCE (equation 4), and HardNeg-InfoNCE (equation 5) objectives, coverting them into Conditional Contrastive learning with Kernel (CCL-K) objectives:

We consider three visual datasets in this set of experiments. Data $X$ and $Y$ represent images after applying arbitrary image augmentations. 1) UT-Zappos (Yu and Grauman, 2014): It contains $50,025$ shoe images over $21$ shoe categories. Each image is annotated with $7$ binomially-distributed attributes as auxiliary information, and we convert them into $126$ binary attributes (Bernoulli-distributed). 2) CUB (Wah et al., 2011): It contains $11,788$ bird images spanning $200$ fine-grain bird species, meanwhile $312$ binary attributes are attached to each image. 3) ImageNet-100 (Russakovsky et al., 2015): It is a subset of ImageNet-1k Russakovsky et al. (2015) dataset, containing $0.12$ million images spanning $100$ categories. This dataset does not come with auxiliary information, and hence we extract the $512$-dim. visual features from the CLIP (Radford et al., 2021) model (a large pre-trained visual model with natural language supervision) to be its auxiliary information. Note that we consider different types of auxiliary information. For UT-Zappos and CUB, we consider discrete and human-annotated attributes. For ImageNet-100, we consider continuous and pre-trained language-enriched features. We report the top-1 accuracy as the metric for the downstream classification task.

We consider the ${\rm WeaklySup}_{\rm CCLK}$ objective (equation 7) as the main method. For ${\rm WeaklySup}_{\rm CCLK}$, we perform the sampling process $\{(x_{i},y_{i},z_{i})\}_{i=1}^{b}\sim P_{XYZ}^{\,\,\,\otimes b}$ by first sampling an image ${\rm im}_{i}$ along with its auxiliary information $z_{i}$ and then applying different image augmentations on ${\rm im}_{i}$ to obtain $(x_{i},y_{i})$. We also study different types of kernels for $K_{Z}$. On the other hand, we select two baseline methods. The first one is the unconditional contrastive learning method: the ${\rm InfoNCE}$ objective (Oord et al., 2018; Chen et al., 2020) (equation 1). The second one is the conditional contrastive learning baseline: the ${\rm WeaklySup}_{\rm InfoNCE}$ objective (equation 3). The difference between ${\rm WeaklySup}_{\rm CCLK}$ and ${\rm WeaklySup}_{\rm InfoNCE}$ is that the latter requires sampling a pair of data with the same outcome from the conditioning variable (i.e., $(x,y)\sim P_{X|Z=z}P_{Y|Z=z}$). However, as suggested in Section 2.3, directly performing conditional sampling is challenging if there is not enough data to support the conditioning sampling procedure. Such limitation exists in our datasets: CUB has on average $1.001$ data points per $Z$’s configuration, and ImageNet-100 has only $1$ data point per $Z$’s configuration since its conditioning variable is continuous and each instance from the dataset has a unique $Z$. To avoid this limitation, in ${\rm WeaklySup}_{\rm InfoNCE}$ clusters the data to ensure that data within the same cluster are abundant and have similar auxiliary information, and then treating the clustering information as the new conditioning variable. The result of ${\rm WeaklySup}_{\rm InfoNCE}$ is reported by selecting the optimal number of clusters via cross-validation.

We show the results in Table 2. First, ${\rm WeaklySup}_{\rm CCLK}$ shows consistent improvements over unconditional baseline InfoNCE, with absolute improvements of $8.8\%$, $15.8\%$, and $6.2\%$ on UT-Zappos, CUB and, ImageNet-100 respectively. This is because the conditional method utilizes the additional auxiliary information (Tsai et al., 2021a). Second, ${\rm WeaklySup}_{\rm CCLK}$ performs better than ${\rm WeaklySup}_{\rm InfoNCE}$, with absolute improvements of $2\%$, $9.3\%$, and $1\%$. We attribute better performance of ${\rm WeaklySup}_{\rm CCLK}$ over ${\rm WeaklySup}_{\rm InfoNCE}$ to the following fact: ${\rm WeaklySup}_{\rm InfoNCE}$ first performs clustering on the auxiliary information and considers the new clusters as the conditioning variable, while ${\rm WeaklySup}_{\rm CCLK}$ directly considers the auxiliary information as the conditioning variable. The clustering in ${\rm WeaklySup}_{\rm InfoNCE}$ may lose precision of the auxiliary information and may negatively affect the quality of the auxiliary information incorporated in the representation. Ablation study on the choice of kernels has shown the consistent performance of ${\rm WeaklySup}_{\rm CCLK}$ across different kernels on the UT-Zappos, CUB and ImageNet-100 dataset, where we consider the following kernel functions: RBF, Laplacian, Linear and Cosine. Most kernels have similar performances, except that linear kernel is worse than others on ImageNet-100 (by at least $2.7\%$).

Our experiments focus on continuous sensitive information, to echo with the limitation of the conditional sampling procedure in Section 2.3. Nonetheless, existing datasets mostly consider discrete sensitive variables, such as gender or race. Therefore, we synthetically create ColorMNIST dataset, which randomly assigns a continuous RBG color value for the background in each handwritten digit image in the MNIST dataset (LeCun et al., 1998). We consider the background color as sensitive information. For statistics, ColorMNIST has $60,000$ colored digit images across $10$ digit labels. Similar to Section 4.1, data $X$ and $Y$ represent images after applying arbitrary image augmentations. Our goal is two-fold: we want to see how well the learned representations 1) perform on the downstream classification task and 2) ignore the effect from the sensitive variable. For the former one, we report the top-1 accuracy as the metric; for the latter one, we report the Mean Square Error (MSE) when trying to predict the color information. Note that the MSE score is higher the better since we would like the learned representations to contain less color information.

We consider the ${\rm Fair}_{\rm CCLK}$ objective (equation 8) as the main method. For ${\rm Fair}_{\rm CCLK}$, we perform the sampling process $\{(x_{i},y_{i},z_{i})\}_{i=1}^{b}\sim P_{XYZ}^{\,\,\,\otimes b}$ by first sampling a digit image ${\rm im}_{i}$ along with its sensitive information $z_{i}$ and then applying different image augmentations on ${\rm im}_{i}$ to obtain $(x_{i},y_{i})$. We select the unconditional contrastive learning method - the ${\rm InfoNCE}$ objective (equation 1) as our baseline. We also consider the ${\rm Fair}_{\rm InfoNCE}$ objective (equation 4) as a baseline, by clustering the continuous conditioning variable $Z$ into one of the following: 3, 5, 10, 15, or 20 clusters using K-means.

We show the results in Table 3. ${\rm Fair}_{\rm CCLK}$ is consistently better than the InfoNCE, where the absolute accuracy improvement is $2.3\%$ and the relative improvement of MSE (higher the better) is $32.6\%$ over InfoNCE. We report the result of using the Cosine kernel and provide the ablation study of different kernel choices in the Appendix. This result suggests that the proposed ${\rm Fair}_{\rm CCLK}$ can achieve better downstream classification accuracy while ignoring more sensitive information (color information) compared to the unconditional baseline, suggesting that our method can achieve a better level of fairness (by excluding color bias) without sacrificing performance. Next we compare to ${\rm Fair}_{\rm InfoNCE}$ baseline, and we report the result using the 10-cluster partition of $Z$ as it achieves the best top-1 accuracy. Compared to ${\rm Fair}_{\rm InfoNCE}$, ${\rm Fair}_{\rm CCLK}$ is better in downstream accuracy, while slightly worse in MSE (a difference of 0.2). For ${\rm Fair}_{\rm InfoNCE}$, if the number of discretized value of $Z$ increases, the MSE in general grows, but the accuracy peaks at 10 clusters and then declines. This suggests that ${\rm Fair}_{\rm InfoNCE}$ can remove more sensitive information as the granularity of $Z$ increases, but may hurt the downstream task performance. Overall, the ${\rm Fair}_{\rm CCLK}$ performs slightly better than ${\rm Fair}_{\rm InfoNCE}$, and do not need clustering to discretize $Z$.

We consider two visual datasets in this set of experiments. Data $X$ and $Y$ represent images after applying arbitrary image augmentations. 1) CIFAR10 (Krizhevsky et al., 2009): It contains $60,000$ images spanning $10$ classes, e.g. automobile, plane, or dog. 2) ImageNet-100 (Russakovsky et al., 2015): It is the same dataset that we used in Section 2. We report the top-1 accuracy as the metric for the downstream classification task.

We consider the ${\rm HardNeg}_{\rm CCLK}$ objective (equation 9) as the main method. For ${\rm HardNeg}_{\rm CCLK}$, we perform the sampling process $\{(x_{i},y_{i},z_{i})\}_{i=1}^{b}\sim P_{XYZ}^{\,\,\,\otimes b}$ by first sampling an image ${\rm im}_{i}$, then applying different image augmentations on ${\rm im}_{i}$ to obtain $(x_{i},y_{i})$, and last defining $z_{i}=g_{\theta_{X}}(x_{i})$. We select two baseline methods. The first one is the unconditional contrastive learning method: the ${\rm InfoNCE}$ objective (Oord et al., 2018; Chen et al., 2020) (equation 1). The second one is the conditional contrastive learning baseline: the ${\rm HardNeg}_{\rm InfoNCE}$, objective (Robinson et al., 2020), and we report its result directly using the released code from the author (Robinson et al., 2020).

From Table 4, first, ${\rm HardNeg}_{\rm CCLK}$ consistently shows better performances over the InfoNCE baseline, with absolute improvements of $1.8\%$ and $3.1\%$ on CIFAR-10 and ImageNet-100 respectively. This suggests that the hard negative sampling effectively improves the downstream performance, which is in accordance with the observation by Robinson et al. (2020). Next, ${\rm HardNeg}_{\rm CCLK}$ also performs better than the ${\rm HardNeg}_{\rm InfoNCE}$ baseline, with absolute improvements of $0.3\%$ and $2.0\%$ on CIFAR-10 and ImageNet-100 respectively. Both methods construct hard negatives by assigning a higher weight to a random paired data that are close in the embedding space and a lower weight to a random paired data that are far in the embedding space. The implementation by Robinson et al. (2020) uses Euclidean distances to measure the similarity and ${\rm HardNeg}_{\rm CCLK}$ uses the smoothed kernel similarity (i.e., $(K_{Z}+\lambda{\bf I})^{-1}K_{Z}$ in Proposition 2.2) to measure the similarity. Empirically our approach performs better.