[论文解读] Learning Wasserstein Embeddings

Nicolas Courty, Rémi Flamary+1 more|arXiv (Cornell University)|Oct 20, 2017

Advanced Neural Network Applications参考文献 34被引用 9

一句话总结

该论文提出了一种深度学习框架，将概率分布嵌入到欧氏空间中，使得欧氏距离近似 Wasserstein 距离，从而实现基于 Wasserstein 距离的操作（如重心和插值）的快速计算。通过使用孪生自编码器架构，该方法在与精确最优传输解相比损失极小的情况下，实现了实时推理（每次插值仅需 4 毫秒）。

ABSTRACT

The Wasserstein distance received a lot of attention recently in the community of machine learning, especially for its principled way of comparing distributions. It has found numerous applications in several hard problems, such as domain adaptation, dimensionality reduction or generative models. However, its use is still limited by a heavy computational cost. Our goal is to alleviate this problem by providing an approximation mechanism that allows to break its inherent complexity. It relies on the search of an embedding where the Euclidean distance mimics the Wasserstein distance. We show that such an embedding can be found with a siamese architecture associated with a decoder network that allows to move from the embedding space back to the original input space. Once this embedding has been found, computing optimization problems in the Wasserstein space (e.g. barycenters, principal directions or even archetypes) can be conducted extremely fast. Numerical experiments supporting this idea are conducted on image datasets, and show the wide potential benefits of our method.

研究动机与目标

• 解决大规模机器学习应用中计算 Wasserstein 距离的高计算成本问题。
• 实现基于 Wasserstein 的操作（如重心、测地线和插值）的快速近似。
• 学习一种欧氏嵌入，使得嵌入点之间的欧氏范数能模仿概率测度之间的 Wasserstein 距离。
• 联合学习嵌入及其逆映射，以实现从嵌入空间中重建分布。
• 评估所学习嵌入在不同数据复杂度的数据集之间的可迁移性。

提出的方法

• 训练一种孪生神经网络架构，将概率分布映射到低维欧氏空间，使得欧氏距离近似 Wasserstein 距离。
• 训练一个独立的解码器网络，从嵌入表示中重建原始输入分布，从而实现逆映射。
• 使用对比损失进行模型训练，以促使 Wasserstein 距离较小的分布在嵌入空间中具有较小的欧氏距离。
• 该框架支持嵌入与重建的端到端学习，从而实现在嵌入空间中的高效推理。
• 在 MNIST 和三个草图数据集（Cat、Crab、Faces）上评估该方法，使用 POT 工具箱计算的精确 Wasserstein 距离作为基准。
• 在嵌入空间中执行 Wasserstein 插值和重心计算，并与精确线性规划解和正则化最优传输解进行比较。

实验结果

研究问题

• RQ1 深度神经网络能否学习一种欧氏嵌入，使得嵌入点之间的欧氏距离近似其原始分布之间的 Wasserstein 距离？
• RQ2 所学习的嵌入在多大程度上能保持 Wasserstein 空间中的几何结构，特别是在复杂多样的数据分布下？
• RQ3 在某一数据集上训练的嵌入在多大程度上可以迁移到具有不同数据特征的另一数据集上？
• RQ4 所提出方法的计算效率与精确最优传输求解器和正则化 OT 方法相比如何？
• RQ5 嵌入及其逆映射的联合学习能否实现从嵌入空间中对输入分布的高精度重建？

主要发现

• 所提出方法在每次插值仅需 4 毫秒的情况下实现了 Wasserstein 插值的实时推理，而使用精确线性规划求解器则需 20 秒。
• 在 MNIST 数据集上，所学习的嵌入在重建 Wasserstein 距离时达到均方误差（MSE）为 0.405，表明保真度很高。
• 跨数据集性能显示中等程度的精度下降（例如，使用草图训练的模型在 MNIST 上的 MSE 约为 10–50），但在相似数据领域内保持稳定。
• 该方法生成了平滑、连续的插值结果，反映了最优传输的行为，尽管由于自编码器重建误差而损失了一些细节。
• 正则化最优传输（Bregman 投影）比精确 LP 更快但更模糊；DWE 方法在速度与视觉合理性之间取得了良好平衡。
• 该框架实现了嵌入空间中 Wasserstein 重心和主方向的快速计算，与直接使用 Wasserstein 优化相比显著降低了计算成本。