[论文解读] Distributional Sliced-Wasserstein and Applications to Generative Modeling

Khai Nguyen, Nhat Ho+2 more|arXiv (Cornell University)|Feb 17, 2020

3D Shape Modeling and AnalysisEngineering被引用 18

一句话总结

本文提出分布式切片-Wasserstein（DSW），一种新颖的最优传输距离，通过在单位球面上学习最优的投影方向分布，平衡方向多样性与信息量。通过引入正交方向的正则化，DSW在投影数量更少的情况下，相比SW和Max-SW提升了样本效率与生成建模性能，在MNIST、CelebA、CIFAR10和LSUN数据集上取得了最先进结果。

ABSTRACT

Sliced-Wasserstein distance (SW) and its variant, Max Sliced-Wasserstein distance (Max-SW), have been used widely in the recent years due to their fast computation and scalability even when the probability measures lie in a very high dimensional space. However, SW requires many unnecessary projection samples to approximate its value while Max-SW only uses the most important projection, which ignores the information of other useful directions. In order to account for these weaknesses, we propose a novel distance, named Distributional Sliced-Wasserstein distance (DSW), that finds an optimal distribution over projections that can balance between exploring distinctive projecting directions and the informativeness of projections themselves. We show that the DSW is a generalization of Max-SW, and it can be computed efficiently by searching for the optimal push-forward measure over a set of probability measures over the unit sphere satisfying certain regularizing constraints that favor distinct directions. Finally, we conduct extensive experiments with large-scale datasets to demonstrate the favorable performances of the proposed distances over the previous sliced-based distances in generative modeling applications.

研究动机与目标

• 解决切片-Wasserstein（SW）因方向均匀采样而需要大量投影导致的效率低下问题。
• 克服Max-切片-Wasserstein（Max-SW）仅使用单一最优方向导致的信息损失问题。
• 构建一个原则性的概率框架，学习投影方向的分布，以最大化多样性与信息量。
• 通过几何正则化优化投影分布，提升生成建模中的样本效率与性能。
• 在大规模数据集的线性和非线性设置下，证明DSW相较于现有切片-Wasserstein变体的优越性。

提出的方法

• 将分布式切片-Wasserstein（DSW）作为Max-SW的推广，其中投影分布是可学习的，而非固定为狄拉克或均匀测度。
• 将DSW表述为在单位球面上对概率测度的优化问题，约束条件旨在偏好正交且多样的方向。
• 引入正则化项，鼓励投影方向在球面上彼此远离，从而提升对数据流形的覆盖度。
• 采用带神经网络的对偶经验形式，近似球面上最优前推测度，实现端到端训练。
• 通过使用圆函数引入非线性，将框架扩展至非线性投影，提出广义DSW（DGSW）。
• 通过随机梯度下降求解优化问题，并采用投影分布的可微松弛化方法。

实验结果

研究问题

• RQ1 可学习的投影方向分布是否能提升高维分布中切片-Wasserstein距离的样本效率与准确性？
• RQ2 在生成建模任务中，使用相同数量的投影时，DSW相较于SW和Max-SW的性能表现如何？
• RQ3 在投影空间中对方向多样性进行正则化是否能带来更好的泛化能力与重建质量？
• RQ4 所提出的分布式框架能否在保持计算效率与性能的前提下，扩展至非线性投影？
• RQ5 在联合推理与图像重建任务中，DSW相较于其他切片-Wasserstein变体的表现如何？

主要发现

• 在所有基准数据集（MNIST、CelebA、CIFAR10、LSUN）上，DSW在相同投影数量下均优于SW与Max-SW。
• 仅使用10个投影时，DSW即可成功重建MNIST数字，而SW与GSW无法生成可识别的图像。
• DSW-1000在所有评估指标（包括重建误差与分布相似性）上均优于所有基线模型。
• 非线性变体DGSW在各项性能上优于其他非线性切片-Wasserstein距离，尤其在联合推理任务中表现更优。
• Max-GSW-NN表现良好，但仍逊于DGSW，表明学习方向分布比优化单一方向更为有效。
• 消融实验证实，方向多样性正则化显著提升性能，且未正则化的DSW退化为Max-SW的特例。