[论文解读] Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks

Jiaojiao Fan, Amirhossein Taghvaei+1 more|arXiv (Cornell University)|Jul 8, 2020

Topological and Geometric Data Analysis参考文献 45被引用 18

一句话总结

本文提出了一种可扩展的算法，用于在高维连续空间中计算Wasserstein中位数，通过使用输入凸神经网络（ICNNs）对Kantorovich对偶公式中的凸势函数进行参数化。通过将中位数建模为生成模型，该方法可在无需访问各边缘分布的情况下实现高效且无限的采样，在高维基准测试中表现出更高的准确性和可扩展性。

ABSTRACT

Wasserstein Barycenter is a principled approach to represent the weighted mean of a given set of probability distributions, utilizing the geometry induced by optimal transport. In this work, we present a novel scalable algorithm to approximate the Wasserstein Barycenters aiming at high-dimensional applications in machine learning. Our proposed algorithm is based on the Kantorovich dual formulation of the Wasserstein-2 distance as well as a recent neural network architecture, input convex neural network, that is known to parametrize convex functions. The distinguishing features of our method are: i) it only requires samples from the marginal distributions; ii) unlike the existing approaches, it represents the Barycenter with a generative model and can thus generate infinite samples from the barycenter without querying the marginal distributions; iii) it works similar to Generative Adversarial Model in one marginal case. We demonstrate the efficacy of our algorithm by comparing it with the state-of-art methods in multiple experiments.

研究动机与目标

• 解决在高维连续空间中Wasserstein中位数估计的计算瓶颈问题，传统线性规划和熵正则化方法在此类场景下失效。
• 开发一种将中位数表示为生成模型的方法，实现无需查询边缘分布的无限采样。
• 利用输入凸神经网络（ICNNs）的凸性保持特性，确保Kantorovich公式中对偶势函数的有效性。
• 通过ICNN架构直接施加凸性，消除先前方法中正则化项引入的偏差。
• 在高维合成数据集和真实世界数据集上，展示该方法在可扩展性和准确性方面优于当前最先进方法。

提出的方法

• 该方法采用Wasserstein-2距离的Kantorovich对偶公式，将中位数问题重新表述为对凸函数的优化问题。
• 使用输入凸神经网络（ICNNs）对偶势函数进行参数化，通过网络结构保证函数的凸性。
• 中位数通过一个生成器网络表示，该网络将噪声映射为中位数分布的样本，从而实现高效且无限的采样。
• 优化问题被构建为一个三重极小化-极大化-极小化问题：在生成器上最小化，对偶势函数上最大化，对偶函数上再次最小化。
• 该框架通过随机优化端到端训练，无需对输入分布进行显式离散化。
• 在一元边缘分布情况下，该方法的行为类似于GAN，其中生成器被训练以匹配中位数分布。

实验结果

研究问题

• RQ1 基于ICNN的深度生成模型能否在高维连续空间中有效表示Wasserstein中位数？
• RQ2 通过ICNN施加凸性是否能消除先前基于神经网络方法中正则化引入的偏差？
• RQ3 与当前最先进算法相比，该方法在高维设置下的性能和准确性表现如何？
• RQ4 该方法能否通过单次训练过程扩展至学习任意权重的中位数？
• RQ5 与现有方法相比，该公式的优化景观（特别是在二次函数情况下）有何特征？

主要发现

• 在高维合成数据集和真实世界基准测试中，该方法的估计精度显著优于当前最先进方法，包括Korotin等人（2021b）和Li等人（2020）的方法。
• 在二次函数情况下，该方法简化为一个光滑的凹-凸-凹优化问题，确保了稳定收敛，而Korotin等人（2021b）的方法则为非光滑且非凸的公式。
• 最优生成器能够精确学习到协方差相等的高斯边缘分布的中位数，其均值为∑a_i m_i，验证了理论一致性。
• 该方法避免了对先验分布选择或复杂正则化项（如一致性正则化项或循环正则化项）的依赖，这些项在先前工作中被广泛使用。
• 实验表明，该方法在高维设置下具有良好的可扩展性，在速度和准确性方面均优于熵正则化和自由支撑方法。
• 该框架可被适配为一元边缘情况下的GAN，展示了其灵活性和强大的生成能力。