Strongly Isomorphic Neural Optimal Transport Across Incomparable Spaces

TL;DR

This paper introduces a neural approach for optimal transport between incomparable spaces, leveraging invariance to strong isomorphisms to improve flexibility and theoretical guarantees in learning minimal-displacement maps.

Contribution

It proposes a novel neural formulation of the Gromov-Monge problem based on strong isomorphism invariance, decomposing the OT map into two learnable components.

Findings

Preliminary results show effective learning of OT maps across diverse spaces.

The method extends the Monge gap regularizer for practical implementation.

Provides theoretical guarantees for the proposed neural OT framework.

Abstract

Optimal Transport (OT) has recently emerged as a powerful framework for learning minimal-displacement maps between distributions. The predominant approach involves a neural parametrization of the Monge formulation of OT, typically assuming the same space for both distributions. However, the setting across ``incomparable spaces'' (e.g., of different dimensionality), corresponding to the Gromov- Wasserstein distance, remains underexplored, with existing methods often imposing restrictive assumptions on the cost function. In this paper, we present a novel neural formulation of the Gromov-Monge (GM) problem rooted in one of its fundamental properties: invariance to strong isomorphisms. We operationalize this property by decomposing the learnable OT map into two components: (i) an approximate strong isomorphism between the source distribution and an intermediate reference distribution, and (ii) a GM-optimal map between this reference and the target distribution. Our formulation leverages and extends the Monge gap regularizer of Uscidda & Cuturi (2023) to eliminate the need for complex architectural requirements of other neural OT methods, yielding a simple but practical method that enjoys favorable theoretical guarantees. Our preliminary empirical results show that our framework provides a promising approach to learn OT maps across diverse spaces.