Flow-based Alignment Approaches for Probability Measures in Different Spaces

Abstract

Gromov-Wasserstein (GW) is a powerful tool to compare probability measures whose supports are in different metric spaces. GW suffers however from a computational drawback since it requires to solve a complex non-convex quadratic program. We consider in this work a specific family of cost metrics, namely \textit{tree metrics} for a space of supports of each probability measure, and aim for developing efficient and scalable discrepancies between the probability measures. By leveraging a tree structure, we propose to align \textit{flows} from a root to each support instead of pair-wise tree metrics of supports, i.e., flows from a support to another, in GW. Consequently, we propose a novel discrepancy, named Flow-based Alignment (\FlowAlign), by matching the flows of the probability measures. We show that \FlowAlign~shares a similar structure as a univariate optimal transport distance. Therefore, \FlowAlign~is fast for computation and scalable for large-scale applications. By further exploring tree structures, we propose a variant of \FlowAlign, named Depth-based Alignment (\DepthAlign), by aligning the flows hierarchically along each depth level of the tree structures. Theoretically, we prove that both \FlowAlign~and \DepthAlign~are pseudo-distances. Moreover, we also derive tree-sliced variants, computed by averaging the corresponding \FlowAlign~/ \DepthAlign~using random tree metrics, built adaptively in spaces of supports. Empirically, we test our proposed discrepancies against other baselines on some benchmark tasks.