[Paper Review] The Monge Gap: A Regularizer to Learn All Transport Maps

Théo Uscidda, Marco Cuturi|arXiv (Cornell University)|Feb 9, 2023

Adversarial Robustness in Machine Learning4 citations

TL;DR

This paper introduces the Monge gap, a novel regularizer that enables training general neural networks as optimal transport (OT) maps without architectural constraints. By minimizing a fitting loss regularized by the Monge gap—which measures deviation from OT optimality—the method outperforms ICNNs and MLPs across synthetic and single-cell genomics benchmarks, especially in high dimensions.

ABSTRACT

Optimal transport (OT) theory has been been used in machine learning to study and characterize maps that can push-forward efficiently a probability measure onto another. Recent works have drawn inspiration from Brenier's theorem, which states that when the ground cost is the squared-Euclidean distance, the ``best'' map to morph a continuous measure in $\mathcal{P}(\Rd)$ into another must be the gradient of a convex function. To exploit that result, [Makkuva+ 2020, Korotin+2020] consider maps $T= abla f_θ$, where $f_θ$ is an input convex neural network (ICNN), as defined by Amos+2017, and fit $θ$ with SGD using samples. Despite their mathematical elegance, fitting OT maps with ICNNs raises many challenges, due notably to the many constraints imposed on $θ$; the need to approximate the conjugate of $f_θ$; or the limitation that they only work for the squared-Euclidean cost. More generally, we question the relevance of using Brenier's result, which only applies to densities, to constrain the architecture of candidate maps fitted on samples. Motivated by these limitations, we propose a radically different approach to estimating OT maps: Given a cost $c$ and a reference measure $ρ$, we introduce a regularizer, the Monge gap $\mathcal{M}^c_ρ(T)$ of a map $T$. That gap quantifies how far a map $T$ deviates from the ideal properties we expect from a $c$-OT map. In practice, we drop all architecture requirements for $T$ and simply minimize a distance (e.g., the Sinkhorn divergence) between $T\sharpμ$ and $ν$, regularized by $\mathcal{M}^c_ρ(T)$. We study $\mathcal{M}^c_ρ$, and show how our simple pipeline outperforms significantly other baselines in practice.

Motivation & Objective

• To address the limitations of input convex neural networks (ICNNs) in optimal transport (OT) learning, which suffer from training instability, complex initialization, and incompatibility with non-squared-Euclidean costs.
• To question the theoretical soundness of enforcing Brenier’s theorem (gradient of convex potential) when working with sample-based measures rather than continuous densities.
• To develop a generic, architecture-agnostic method for learning OT maps that works for any cost function $ c $, without requiring specialized network structures.
• To introduce a regularizer, the Monge gap $ \mathcal{M}^{c}_{\rho}(T) $, that quantifies how close a map $ T $ is to being an optimal OT map, based on a reference measure $ \rho $.

Proposed method

• Proposes the Monge gap $ \mathcal{M}^{c}_{\rho}(T) = \mathbb{E}_{X \sim \rho}[c(X, T(X))] - W_c^2(\rho, T\sharp\rho) $, which measures the deviation of $ T $ from optimal OT behavior.
• Uses the Monge gap as a regularizer in a training objective: minimize a fitting loss (e.g., Sinkhorn divergence) between $ T\sharp\mu $ and $ \nu $, plus $ \lambda_{\mathrm{MG}} \cdot \mathcal{M}^{c}_{\rho}(T) $.
• For costs satisfying the twist condition, proposes a parameterization of $ T $ as the gradient of a potential, leveraging structural properties of optimal maps.
• Employs a reference measure $ \rho $, which can be chosen independently of the source measure $ \mu $, enabling flexibility and robustness.
• Uses Sinkhorn estimators for the $ c $-Wasserstein distance to make the Monge gap differentiable and computable in finite-sample settings.
• Applies the method to both general costs and the squared-Euclidean case, showing convexity of the Monge gap under the latter, which aids optimization.

Experimental results

Research questions

• RQ1 Can we learn optimal transport maps without constraining the architecture to be an input convex neural network (ICNN), thereby avoiding training instabilities and complex initialization?
• RQ2 Is it theoretically and practically valid to impose Brenier’s theorem (gradient of convex potential) when working with sample-based measures rather than continuous densities?
• RQ3 Can a regularizer that quantifies deviation from OT optimality—without architectural constraints—outperform ICNNs and vanilla MLPs in OT map learning?
• RQ4 How does the Monge gap regularizer behave under different cost functions, and can it be made robust to hyperparameter choices in high-dimensional settings?
• RQ5 Does the proposed method generalize to real-world biological data, such as single-cell omics, where transport maps model cellular responses to perturbations?

Key findings

• The Monge gap $ \mathcal{M}^{c}_{\rho}(T) $ is zero if and only if $ T $ is an optimal OT map, provided $ T\sharp\mu = \nu $ and $ \mathrm{Spt}(\mu) \subset \mathrm{Spt}(\rho) $, establishing its theoretical soundness.
• For the squared-Euclidean cost, the Monge gap is convex, and this property is preserved under Sinkhorn finite-sample estimation, enabling stable optimization.
• On the Korotin et al. (2021) benchmark, the proposed method outperforms both vanilla MLPs and ICNNs, especially for $ d \geq 16 $, with a performance gap widening at $ d \geq 64 $.
• In single-cell genomics experiments on 4i and scRNA-seq data, the method achieves better predictive performance than both vanilla MLPs and ICNNs, with consistent improvements across all 34 and 9 treatments.
• The method is robust to hyperparameter choice: a wide region of $ (\lambda_{\mathrm{MG}}, \lambda_{\mathrm{cons}}) $ values yields good performance, with optimal weights $ (1, 0.01) $ for $ d \leq 64 $ and $ (10, 0.1) $ for $ d \geq 128 $.
• ICNNs show unstable and moderate performance despite Gaussian initialization, highlighting the practical limitations of architectural constraints in OT learning.