Contrasting Multiple Representations with the Multi-Marginal Matching Gap

TL;DR

This paper introduces the multi-marginal matching gap (M3G), a novel loss function based on multi-marginal optimal transport, to better learn representations from multiple views or modalities, outperforming existing pairwise methods.

Contribution

The paper proposes M3G, a multi-marginal optimal transport-based loss that efficiently incorporates all views simultaneously, improving multi-view representation learning.

Findings

M3G outperforms pairwise extension methods in experiments.

A generalized Sinkhorn algorithm scales to 3-6 views with reasonable batch sizes.

M3G improves performance in self-supervised and multimodal tasks.

Abstract

Learning meaningful representations of complex objects that can be seen through multiple ($k\geq 3$) views or modalities is a core task in machine learning. Existing methods use losses originally intended for paired views, and extend them to $k$ views, either by instantiating $\tfrac12k(k-1)$ loss-pairs, or by using reduced embeddings, following a \textit{one vs. average-of-rest} strategy. We propose the multi-marginal matching gap (M3G), a loss that borrows tools from multi-marginal optimal transport (MM-OT) theory to simultaneously incorporate all $k$ views. Given a batch of $n$ points, each seen as a $k$-tuple of views subsequently transformed into $k$ embeddings, our loss contrasts the cost of matching these $n$ ground-truth $k$-tuples with the MM-OT polymatching cost, which seeks $n$ optimally arranged $k$-tuples chosen within these $n\times k$ vectors. While the exponential complexity $O(n^k$) of the MM-OT problem may seem daunting, we show in experiments that a suitable generalization of the Sinkhorn algorithm for that problem can scale to, e.g., $k=3\sim 6$ views using mini-batches of size $64~\sim128$. Our experiments demonstrate improved performance over multiview extensions of pairwise losses, for both self-supervised and multimodal tasks.