Arbitrary Distributions Mapping via SyMOT-Flow: A Flow-based Approach Integrating Maximum Mean Discrepancy and Optimal Transport

TL;DR

The paper tackles learning a transformation between two unknown distributions from finite samples by developing SyMOT-Flow, a symmetric flow-based model that leverages maximum mean discrepancy (MMD) and optimal transport (OT) regularization. The method learns an invertible transformation via invertible neural networks (INNs) and imposes a symmetric MMD term together with a transport cost to approximate the OT constraint in the original data space. The authors provide theoretical results linking the relaxed OT objective to the true OT solution through $\Gamma$-convergence and convergence of $\mathrm{OT}_\lambda$ to $\mathrm{OT}$ as $\lambda \to \infty$, ensuring feasibility and stability. Empirically, SyMOT-Flow achieves accurate and interpretable mappings on toy 2D datasets, MNIST/Fashion-MNIST feature spaces, and high-dimensional medical imaging modalities, outperforming baselines in forward/backward consistency and quality of generated samples, with ablations confirming the importance of the symmetric design and OT regularization.

Abstract

Finding a transformation between two unknown probability distributions from finite samples is crucial for modeling complex data distributions and performing tasks such as sample generation, domain adaptation and statistical inference. One powerful framework for such transformations is normalizing flow, which transforms an unknown distribution into a standard normal distribution using an invertible network. In this paper, we introduce a novel model called SyMOT-Flow that trains an invertible transformation by minimizing the symmetric maximum mean discrepancy between samples from two unknown distributions, and an optimal transport cost is incorporated as regularization to obtain a short-distance and interpretable transformation. The resulted transformation leads to more stable and accurate sample generation. Several theoretical results are established for the proposed model and its effectiveness is validated with low-dimensional illustrative examples as well as high-dimensional bi-modality medical image generation through the forward and reverse flows.

Figures (10)

• Figure 1: Overview of SyMOT-Flow Model. The red block contains the main structure of the model and the blue one consists of the precise structure of the INN blocks, which are the basic block of the distribution transformation $T_\theta$.
• Figure 1: The input (blue) and generated points (orange) by learned map (green lines) between two sets of training data with different methods.
• Figure 1: The input (blue) and generated points (orange) by learned map (green lines) between two sets of training data with different methods. Example 3 is between two Gauss distributions with different means and covariance and Example 4 is between two 8 Gauss distributions different means and covariance.
• Figure 1: Some Toy Examples
• Figure 2: The input $\mathbf{x}$ (blue) and generated points $\mathbf{z}$ (orange) by learned map $T_\theta$ (green lines) between two sets of training data with and without the reversed flow loss.

Theorems & Definitions (18)

• Definition 2.1: Reproducing Kernel Hilbert Space (RKHS)
• Definition 2.2: Strictly Integrally Positive Definite Kernel ($\int$s.p.d)
• Definition 2.3: Universal
• Definition 2.4: Diffeomorphism
• Lemma 2.5
• Lemma 2.6
• Lemma 2.7
• Remark 2.8
• Remark 2.9
• Remark 3.1