Contrasting Multiple Representations with the Multi-Marginal Matching Gap

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 extend the InfoNCE loss, typically considered for paired views ($k=2$), either by instantiating $\tfrac12k(k-1)$ InfoNCE pairs, or by using reduced embeddings, following a one vs. average-of-rest strategy. We propose the multi-marginal matching gap ($\rm{M3G}$), a radically different loss that borrows tools from multi-marginal optimal transport (MM-OT) theory. Given $n$ points, each seen as a $k$-tuple of 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 $n\times k$ vectors. While the exponential complexity $O(n^k$) of the MM-OT problem may seem daunting, our experiments show that the MM-Sinkhorn algorithm can be run for $k=3\sim 6$ views. Additionally, and thanks to Danskin's theorem, getting the gradient of the $\rm{M3G}$ loss does not require a backward pass. Our experiments demonstrate performance improvements over multiview extensions of InfoNCE, for both self-supervised and multimodal tasks.

Gradient flows of $\rm{M3G}$ : The gradient flow of $\rm{M3G}$ when applied to a toy problem with $k=3$ for $n=4$ points in $d=3$ lying on the unit sphere.

Intuition for $\rm{M3G}$

Consider a problem with $n=4$ points (identified using 4 colors), each given in $k=3$ views (differentiated using 3 shapes) in $d=2$ dimensions.
The ground-truth polymatching of these points is known: to each color its 3 shapes, as illustrated with colored cliques, and described mathematically as a tensor $\mathscr{J}_{4,3}$.
Their initial configuration in space indicates, assuming one solves a multi-marginal optimal transport problem parameterized with the cost tensor $\mathscr{C}(\mathscr{X})$, a different polymatching $\mathscr{P}_0(\mathscr{C}(\mathscr{X}))$.
That difference (quantified as a difference in their matching objectives) defines the $\rm{M3G}$ loss. A high $\rm{M3G}$ indicates, as shown on the left, a large discrepancy between the ground-truth matching's cost and that of the optimal polymatching.
This loss will gradually displace points so that, ideally, upon convergence and after consecutive updates (visualized in (middle) and (right) plots), both ground-truth and optimal polymatchings coincide in their objective.

BibTeX



@InProceedings{piran24m3g,
  title = 	 {Contrasting Multiple Representations with the Multi-Marginal Matching Gap},
  author =       {Piran, Zoe and Klein, Michal and Thornton, James and Cuturi, Marco},
  booktitle = 	 {Proceedings of the 41st International Conference on Machine Learning},
  pages = 	 {40827--40842},
  year = 	 {2024},
  editor = 	 {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix},
  volume = 	 {235},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {21--27 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://raw.githubusercontent.com/mlresearch/v235/main/assets/piran24a/piran24a.pdf},
  url = 	 {https://proceedings.mlr.press/v235/piran24a.html},
  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 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.}
}