Your contrastive learning problem is secretly a distribution alignment problem
This work provides new theoretical insights and tools for self-supervised learning, potentially benefiting researchers and practitioners in machine learning, though it is incremental in building on existing contrastive learning methods.
The authors tackled the problem of understanding and improving contrastive learning by reframing it as a distribution alignment problem using entropic optimal transport, resulting in a family of new losses and iterative variants that offer theoretical insights and experimental benefits for generalized contrastive alignment.
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.