[论文解读] 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 Learning被引用 4

一句话总结

该论文提出了一种名为Monge gap的新正则化项，使通用神经网络能够作为最优传输（OT）映射进行训练，且无需架构约束。通过最小化由Monge gap正则化的拟合损失——该正则化项衡量与OT最优性的偏离程度——该方法在合成数据和单细胞基因组学基准测试中均优于ICNNs和MLPs，尤其在高维设置下表现更优。

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.

研究动机与目标

• 为解决输入凸神经网络（ICNNs）在最优传输（OT）学习中的局限性，包括训练不稳定、初始化复杂以及与非平方欧几里得代价不兼容的问题。
• 质疑在基于样本的测度而非连续密度下强制实施Brenier定理（凸势函数的梯度）的理论合理性。
• 开发一种通用的、与架构无关的方法，用于学习任意代价函数 $ c $ 的OT映射，无需依赖特殊网络结构。
• 提出一种正则化项——Monge gap $ \mathcal{M}^{c}_{\rho}(T) $，用于量化映射 $ T $ 与最优OT映射的接近程度，基于参考测度 $ \rho $。

提出的方法

• 提出Monge gap $ \mathcal{M}^{c}_{\rho}(T) = \mathbb{E}_{X \sim \rho}[c(X, T(X))] - W_c^2(\rho, T\sharp\rho) $，用于衡量映射 $ T $ 偏离最优OT行为的程度。
• 将Monge gap用作训练目标中的正则化项：最小化 $ T\sharp\mu $ 与 $ \nu $ 之间的拟合损失（如Sinkhorn散度）加上 $ \lambda_{\mathrm{MG}} \cdot \mathcal{M}^{c}_{\rho}(T) $。
• 对于满足扭曲条件的代价函数，提出将 $ T $ 参数化为势函数的梯度，利用最优映射的结构特性。
• 采用参考测度 $ \rho $，其可独立于源测度 $ \mu $ 选择，从而提升灵活性与鲁棒性。
• 使用Sinkhorn估计器计算 $ c $-Wasserstein 距离，使Monge gap在有限样本设置下可微分且可计算。
• 将该方法应用于一般代价函数和平方欧几里得情形，证明在后者中Monge gap具有凸性，有助于优化。

实验结果

研究问题

• RQ1 我们能否在不将架构限制为输入凸神经网络（ICNN）的前提下学习最优传输映射，从而避免训练不稳定性与复杂的初始化过程？
• RQ2 在基于样本的测度而非连续密度下强制实施Brenier定理（凸势函数的梯度）在理论上和实践上是否合理？
• RQ3 一种不依赖架构约束、用于量化与OT最优性偏离程度的正则化项，能否在OT映射学习中优于ICNNs和普通MLPs？
• RQ4 Monge gap正则化项在不同代价函数下的行为如何？能否在高维设置下对超参数选择保持鲁棒性？
• RQ5 所提出的方法是否能推广到真实世界生物数据（如单细胞组学）中，其中传输映射用于建模细胞对扰动的响应？

主要发现

• 在 $ T\sharp\mu = \nu $ 且 $ \mathrm{Spt}(\mu) \subset \mathrm{Spt}(\rho) $ 的前提下，Monge gap $ \mathcal{M}^{c}_{\rho}(T) $ 为零当且仅当 $ T $ 是最优OT映射，从而确立了其理论合理性。
• 对于平方欧几里得代价，Monge gap具有凸性，且该性质在Sinkhorn有限样本估计下得以保持，从而支持稳定优化。
• 在Korotin等人（2021）的基准测试中，所提方法在 $ d \geq 16 $ 时优于普通MLPs和ICNNs，且在 $ d \geq 64 $ 时性能差距进一步扩大。
• 在4i和scRNA-seq数据的单细胞基因组学实验中，该方法在全部34种和9种处理条件下均优于普通MLPs和ICNNs，且表现一致更优。
• 该方法对超参数选择具有鲁棒性：$ (\lambda_{\mathrm{MG}}, \lambda_{\mathrm{cons}}) $ 的广泛区域均能获得良好性能，最优权重为 $ d \leq 64 $ 时的 $ (1, 0.01) $，$ d \geq 128 $ 时为 $ (10, 0.1) $。
• 尽管采用高斯初始化，ICNNs仍表现出不稳定且中等水平的性能，凸显了在OT学习中架构约束的实际局限性。