MMD-Regularized Unbalanced Optimal Transport

We study the unbalanced optimal transport (UOT) problem, where the marginal constraints are enforced using Maximum Mean Discrepancy (MMD) regularization. Our work is motivated by the observation that the literature on UOT is focused on regularization based on $\phi$-divergence (e.g., KL divergence). Despite the popularity of MMD, its role as a regularizer in the context of UOT seems less understood. We begin by deriving the dual of MMD-regularized UOT (MMD-UOT), which helps us prove other useful properties. One interesting outcome of this duality result is that MMD-UOT induces novel metrics, which not only lift the ground metric like the Wasserstein but are also efficient to estimate like the MMD. Further, we present finite-dimensional convex programs for estimating MMD-UOT and the corresponding barycenter solely based on the samples from the measures being transported. Under mild conditions, we prove that our convex-program-based estimators are consistent and the estimation error decays at a rate $\mathcal{O}\left(m^{-\frac{1}{2}}\right)$, where $m$ is the number of samples. As far as we know, such error bounds that are free from the curse of dimensionality are not known for $\phi$-divergence regularized UOT. Finally, we discuss how the proposed convex programs can be solved efficiently using accelerated projected gradient descent. Our experiments show that MMD-UOT consistently outperforms popular baselines, including KL-regularized UOT and MMD, in diverse machine learning applications.