[Paper Review] InfoNCE: Identifying the Gap Between Theory and Practice

Evgenia Rusak, Patrik Reizinger+4 more|arXiv (Cornell University)|Jun 28, 2024

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

This paper identifies a critical gap between theoretical assumptions in contrastive learning (CL) and real-world practice, where data augmentations cause anisotropic changes across latent factors—unlike the isotropic assumptions in prior InfoNCE theory. The authors propose AnInfoNCE, a generalized contrastive loss that provably identifies latent factors under anisotropic distributions, demonstrating improved recovery of collapsed information on CIFAR10 and ImageNet at the cost of downstream accuracy.

ABSTRACT

Prior theory work on Contrastive Learning via the InfoNCE loss showed that, under certain assumptions, the learned representations recover the ground-truth latent factors. We argue that these theories overlook crucial aspects of how CL is deployed in practice. Specifically, they either assume equal variance across all latents or that certain latents are kept invariant. However, in practice, positive pairs are often generated using augmentations such as strong cropping to just a few pixels. Hence, a more realistic assumption is that all latent factors change with a continuum of variability across all factors. We introduce AnInfoNCE, a generalization of InfoNCE that can provably uncover the latent factors in this anisotropic setting, broadly generalizing previous identifiability results in CL. We validate our identifiability results in controlled experiments and show that AnInfoNCE increases the recovery of previously collapsed information in CIFAR10 and ImageNet, albeit at the cost of downstream accuracy. Finally, we discuss the remaining mismatches between theoretical assumptions and practical implementations.

Motivation & Objective

• To identify the mismatch between theoretical assumptions in InfoNCE and practical CL deployments, particularly regarding isotropic vs. anisotropic latent factor changes under data augmentations.
• To develop a theoretically grounded contrastive loss, AnInfoNCE, that generalizes InfoNCE to handle anisotropic positive pair distributions where latent factors vary to different extents.
• To validate the identifiability of AnInfoNCE in controlled synthetic and real-world experiments, demonstrating recovery of collapsed information in CIFAR10 and ImageNet.
• To explore and formalize extensions to hard negative mining and loss ensembles, ensuring theoretical identifiability in more complex training setups.
• To analyze the trade-off between representation identifiability and downstream accuracy, linking it to augmentation-induced information loss.

Proposed method

• Propose AnInfoNCE as a generalization of InfoNCE that models positive pairs with varying degrees of change across latent factors, using a continuous anisotropic distribution over latent variables.
• Formalize the identifiability condition under which AnInfoNCE can recover the ground-truth latent factors, even when some factors are highly unstable under augmentation.
• Introduce a theoretical framework for hard negative mining that maintains identifiability by modeling the distribution of negative samples in a way consistent with the anisotropic positive pair assumption.
• Extend the identifiability result to loss ensembles by combining multiple contrastive objectives, each sensitive to different aspects of the latent distribution.
• Design controlled experiments using a VAE-based generative model on MNIST to validate the theoretical identifiability under known data-generating processes.
• Apply AnInfoNCE to real-world datasets (CIFAR10, ImageNet) and evaluate latent factor recovery using linear probing and information recovery metrics.

Figure 1 : Illustration of the mismatch between the standard CL model and practice: A: CL with the commonly used InfoNCE objective is identifiable when all latents change to the same extent across the positive pair zimmermann_contrastive_2021 , which is unlikely to happen in practice. B: The more li

Experimental results

Research questions

• RQ1 How does the assumption of isotropic change in latent factors in InfoNCE theory conflict with real-world data augmentations in contrastive learning?
• RQ2 Can a generalized contrastive loss like AnInfoNCE provably identify latent factors when positive pairs exhibit anisotropic changes across different factors?
• RQ3 What is the trade-off between latent factor identifiability and downstream linear probe accuracy when using AnInfoNCE?
• RQ4 How can hard negative mining be formalized to preserve identifiability under anisotropic distributions?
• RQ5 Can loss ensembles improve identifiability while maintaining robustness to varying augmentation strategies?

Key findings

• AnInfoNCE successfully recovers previously collapsed latent information in CIFAR10 and ImageNet, demonstrating improved identifiability under anisotropic augmentations.
• The method achieves provable identifiability of latent factors under anisotropic positive pair distributions, generalizing prior isotropic and block-identifiability results.
• On CIFAR10 and ImageNet, AnInfoNCE increases the recovery of ground-truth latent factors compared to standard InfoNCE, though at the cost of reduced downstream linear probe accuracy.
• Controlled experiments on a VAE-based MNIST model confirm that AnInfoNCE can recover the true generative factors when the data-generating process is known.
• The trade-off between identifiability and downstream accuracy is linked to augmentation strategies, particularly strong cropping, which induce high anisotropy in latent changes.
• Extensions to hard negative mining and loss ensembles are formally justified and shown to preserve identifiability under the anisotropic setting.