[论文解读] Augmented Sliced Wasserstein Distances

Xiongjie Chen, Yongxin Yang+1 more|arXiv (Cornell University)|Jun 15, 2020

Advanced Numerical Analysis Techniques参考文献 39被引用 4

一句话总结

本文提出了一种新型度量——增强切片Wasserstein距离（ASWD），通过在应用非线性投影前将数据映射到神经网络定义的超曲面，从而在计算效率和Wasserstein距离近似精度方面实现提升。该方法通过灵活、数据自适应的非线性切片，显著提升了投影效率与度量有效性，尤其在高维空间中表现优越，其核心基于可注入神经网络实现。

ABSTRACT

While theoretically appealing, the application of the Wasserstein distance to large-scale machine learning problems has been hampered by its prohibitive computational cost. The sliced Wasserstein distance and its variants improve the computational efficiency through the random projection, yet they suffer from low accuracy if the number of projections is not sufficiently large, because the majority of projections result in trivially small values. In this work, we propose a new family of distance metrics, called augmented sliced Wasserstein distances (ASWDs), constructed by first mapping samples to higher-dimensional hypersurfaces parameterized by neural networks. It is derived from a key observation that (random) linear projections of samples residing on these hypersurfaces would translate to much more flexible nonlinear projections in the original sample space, so they can capture complex structures of the data distribution. We show that the hypersurfaces can be optimized by gradient ascent efficiently. We provide the condition under which the ASWD is a valid metric and show that this can be obtained by an injective neural network architecture. Numerical results demonstrate that the ASWD significantly outperforms other Wasserstein variants for both synthetic and real-world problems.

研究动机与目标

• 为解决传统Wasserstein距离与切片Wasserstein距离在大规模机器学习中计算成本高、投影效率低的问题。
• 克服标准切片Wasserstein距离中线性投影的局限性，后者在高维空间中常导致距离值趋近于零的平凡结果。
• 开发一种可微分、可扩展的度量，保持理论有效性的同时，通过学习得到的非线性投影捕捉复杂数据结构。
• 通过确保距离度量与基于梯度的优化兼容，支持深度生成模型中的端到端训练。

提出的方法

• 该方法利用可注入神经网络将输入样本映射到更高维的超曲面，将数据分布转化为非线性流形。
• 随后对嵌入后的超曲面应用线性投影，进而在原始数据空间中诱导出灵活的非线性投影。
• ASWD被定义为变换后分布的一维投影之间Wasserstein距离的期望值，通过排序实现高效计算。
• 通过梯度上升优化超曲面参数，以最大化投影分布之间的距离，从而提升判别能力。
• 推导出ASWD作为有效度量的理论条件，要求神经网络为可注入的，且投影分布足够丰富。
• 该方法支持端到端训练，且与可微分求解器兼容，可无缝集成至深度学习流水线。

实验结果

研究问题

• RQ1 基于学习得到的超曲面所导出的非线性、数据自适应投影，是否能显著提升高维空间中切片Wasserstein距离的投影效率？
• RQ2 当通过神经网络嵌入与随机投影构建时，所提出的ASWD是否仍保持作为有效度量的理论有效性？
• RQ3 在合成数据与真实世界分布中，ASWD相较于SWD、Max-SWD、GSWD与DSWD等现有变体，在精度与计算效率方面表现如何？
• RQ4 ASWD是否能在投影数量较少的情况下实现高性能，从而在降低计算成本的同时保持强判别能力？

主要发现

• ASWD在合成数据与真实世界任务中均显著优于SWD、Max-SWD、GSWD与DSWD，尤其在100维高斯分布等高维设置中表现突出。
• 仅使用10次投影，ASWD在色彩迁移任务中生成的图像翻译结果便比SWD与Max-SWD更清晰、更逼真。
• ASWD生成的重心在几何上具有意义，且与完整Wasserstein重心相当，同时在计算效率上优于SWD与DSWD。
• 与SWD中的线性投影相比，ASWD展现出更高的投影效率，其直方图结果在100维空间中显示更少的平凡小距离值。
• 采用可注入神经网络确保ASWD满足作为有效度量的条件，为其应用提供了坚实的理论基础。