Your contrastive learning problem is secretly a distribution alignment problem

TL;DR

This paper reveals that contrastive learning can be viewed as a distribution alignment problem using optimal transport, enabling more flexible and distribution-aware representation learning methods.

Contribution

It establishes a theoretical connection between contrastive learning and distribution alignment via optimal transport, introducing new loss variants and insights for self-supervised learning.

Findings

Distribution-aware manipulation improves representations

Optimal transport-based losses handle noisy views better

Framework unifies contrastive learning and distribution alignment

Abstract

Despite the success of contrastive learning (CL) in vision and language, its theoretical foundations and mechanisms for building representations remain poorly understood. In this work, we build connections between noise contrastive estimation losses widely used in CL and distribution alignment with entropic optimal transport (OT). This connection allows us to develop a family of different losses and multistep iterative variants for existing CL methods. Intuitively, by using more information from the distribution of latents, our approach allows a more distribution-aware manipulation of the relationships within augmented sample sets. We provide theoretical insights and experimental evidence demonstrating the benefits of our approach for {\em generalized contrastive alignment}. Through this framework, it is possible to leverage tools in OT to build unbalanced losses to handle noisy views and customize the representation space by changing the constraints on alignment. By reframing contrastive learning as an alignment problem and leveraging existing optimization tools for OT, our work provides new insights and connections between different self-supervised learning models in addition to new tools that can be more easily adapted to incorporate domain knowledge into learning.