InfoNCE: Identifying the Gap Between Theory and Practice

In Contrastive Learning, a model gets two views of a training sample as input and is required to learn an embedding where views of the same sample are close together, and views of different samples are far apart. The views are generated using augmentations such as random cropping and color jittering. They can be considered transformations in the implicit latent space such that different augmentations change different latent dimensions to a different extent. The connection to identifiability theory has been made by Zimmermann et al. 2021, who showed that CL can provably recover the ground truth latents if the conditional distribution is isotropic, which means that the latent dimensions are varied to the same extent by the augmentations. However, this assumption is not realistic in practice because some latents, such as color, are varied strongly. In contrast, other latents, such as the shape or class of an object, are affected minimally. Kügelgen et al. relaxed the isotropic assumption on the positive conditional distributions to two partitions: content (δ-distribution) and style (non-degenerate distribution). They showed that strongly varying latents (style) are lost while non-varying latents (content) are learned. In this paper, we first allow the positive conditional distribution of different latents to change on a spectrum and then introduce a new variation of InfoNCE, called AnInfoNCE, which holds for anisotropically changing latents, for which we prove identifiability.

• We introduce AnInfoNCE, a generalized contrastive loss assuming anisotropic positive pair distributions, and present an identifiability proof.
• Experimentally, we verify our loss in synthetic and well-controlled image experiments, having full knowledge of the ground-truth generative process.
• We further demonstrate the efficacy of AnInfoNCE on CIFAR and ImageNet in recovering the latent factors but observe a trade-off with linear readout accuracy.

We designed a new loss function, coined AnInfoNCE, which, contrary to InfoNCE, can model augmentations that induce anisotropic variances on the latent factors: $$ \mathcal{L}_{\text{AINCE}}({\boldsymbol f},{\bf\hat{\Lambda}}) = \underset{\substack{{\boldsymbol x}, {\boldsymbol x}^{\!+}\\ \{{\boldsymbol x}^{\!-}_i\}}}{\mathbb{E}} \left[ -\ln \frac{e^{-\left\Vert {\boldsymbol f}({\boldsymbol x}^{\!+}) - {\boldsymbol f}({\boldsymbol x})\right\Vert^2_{\bf\hat{\Lambda}}}}{e^{-\left\Vert {\boldsymbol f}({\boldsymbol x}^{\hspace{-0.05em}+}) - {\boldsymbol f}({\boldsymbol x})\right\Vert^2_{\bf\hat{\Lambda}}} + \sum_{i=1}^M e^{-\left\Vert {\boldsymbol f}({\boldsymbol x}^{\!-}_i) - {\boldsymbol f}({\boldsymbol x})\right\Vert^2_{\bf\hat{\Lambda}}}} \right] $$ where the new similarity function: $$ -\lVert {\boldsymbol f}({\boldsymbol x}^{\!+}) - {\boldsymbol f}({\boldsymbol x}) \rVert^2_{\bf\hat{\Lambda}} = - \Big( {\boldsymbol f}({\boldsymbol x}^{\!+}) - {\boldsymbol f}({\boldsymbol x}) \Big)^{\!\!\top}{\bf\hat{\Lambda}}\Big( {\boldsymbol f}({\boldsymbol x}^{\!+}) - {\boldsymbol f}({\boldsymbol x}) \Big) $$ is now equipped with a trainable diagonal scaling matrix, \(\hat{\Lambda}\), the concentration parameter.

To understand the theoretical justification of this method, assume a Data Generating Process, where the observed data \({\boldsymbol x}, {\boldsymbol x}^{\!+},\{{\boldsymbol x}^{\!-}_i\}\) are generated by passing the normalized latent vectors \({\boldsymbol z}, {\boldsymbol z}^{\!+},\{{\boldsymbol z}^{\!-}_i\}\) through an invertible generator \(\boldsymbol{g}\). The anchor and negative pairs \({\boldsymbol z}, \{{\boldsymbol z}^{\!-}_i\}\) are uniformly sampled from the hypersphere, whereas \({\boldsymbol z}^{\!+}\) follows an anisotropic conditional distribution: $$ p({\boldsymbol z}) = \mathrm{const},\quad p({\boldsymbol z}^{\!+}|{\boldsymbol z}) \propto e^{-({\boldsymbol z}^{\!+}-{\boldsymbol z})^\top {\bf{\Lambda}}({\boldsymbol z}^{\!+} -{\boldsymbol z})} $$ The diagonal scaling matrix \(\bf\Lambda\) describes the anisotropic effect of augmentations on the latent features. Our main theoretical result shows that if a pair \(({\boldsymbol f},{\bf\hat{\Lambda}})\) minimizes \(\mathcal{L}_{\text{AnInfoNCE}}\), then the latent factors \(z\) are identified up to orthogonal transformations, i.e., we can recover the ground-truth latents up to simple transformations.

Following previously introduced experimental settings we consider a fully-controlled Data Generating Process. Since we have control of the ground-truth latents, we can directly model a setting where augmentations affect different latents to a different extent. We achieve this by introducing an anisotropy in the positive conditional by varying the ground-truth concentration parameter. We set the value of half of the latent dimensions to a low concentration parameter and the other half to a high concentration parameter. A higher concentration parameter corresponds to a more narrow positive conditional distribution and more content-like behavior while dimensions with a lower concentration parameter correspond to style-like latents. Training a model on the generated data with AnInfoNCE, we observe that both content and style latents can be identified. In contrast, style information is lost when training with the regular InfoNCE loss.

We do not have access to the true Data-Generating Process in real-world datasets - thus, we have to use a proxy task. Sampling from the true conditional distribution is modeled as using augmentations to generate different views as inputs for the model; therefore, we can assess how well we recover ground truth latents by measuring the degree to which the model keeps the information about the used augmentations. Therefore, we calculate the linear readout accuracy on the augmentations used during training to judge how well we recover the latent factors.

Our experiments on CIFAR10 and ImageNet show that our loss with a trainable concentration parameter leads to much higher readout accuracy on the used augmentations, i.e., we successfully recover more latent dimensions compared to the regular InfoNCE loss. While we reduce the information loss, this surprisingly does not lead to better downstream accuracy.

Where does the trade-off between better recovery of latent factors and worse downstream accuracy on real-world data come from? Read our paper where we extensively analyze remaining mismatches between Contrastive Learning theory and practice and explore practical mitigation strategies!