[论文解读] Revisiting Sliced Wasserstein on Images: From Vectorization to Convolution

Khai Nguyen, Nhat Ho|arXiv (Cornell University)|Apr 4, 2022

Generative Adversarial Networks and Image Synthesis被引用 4

一句话总结

本文提出卷积切片Wasserstein（CSW）及其变体，通过用可学习的卷积算子替代传统切片Wasserstein（SW）中的向量化投影，以克服SW在图像生成中的局限性。通过在卷积层中引入步幅、空洞卷积和非线性激活函数，CSW保留了空间结构，减少了内存占用，并在CIFAR10、STL10和CelebA-HQ数据集上训练生成模型时实现了更快的收敛速度和更优的性能。

ABSTRACT

The conventional sliced Wasserstein is defined between two probability measures that have realizations as vectors. When comparing two probability measures over images, practitioners first need to vectorize images and then project them to one-dimensional space by using matrix multiplication between the sample matrix and the projection matrix. After that, the sliced Wasserstein is evaluated by averaging the two corresponding one-dimensional projected probability measures. However, this approach has two limitations. The first limitation is that the spatial structure of images is not captured efficiently by the vectorization step; therefore, the later slicing process becomes harder to gather the discrepancy information. The second limitation is memory inefficiency since each slicing direction is a vector that has the same dimension as the images. To address these limitations, we propose novel slicing methods for sliced Wasserstein between probability measures over images that are based on the convolution operators. We derive convolution sliced Wasserstein (CSW) and its variants via incorporating stride, dilation, and non-linear activation function into the convolution operators. We investigate the metricity of CSW as well as its sample complexity, its computational complexity, and its connection to conventional sliced Wasserstein distances. Finally, we demonstrate the favorable performance of CSW over the conventional sliced Wasserstein in comparing probability measures over images and in training deep generative modeling on images.

研究动机与目标

• 解决将传统切片Wasserstein（SW）应用于图像时因图像向量化导致的效率低下和结构损失问题。
• 克服标准SW中存储高维投影向量所导致的内存低效问题。
• 通过用卷积算子重新定义切片方式，提升使用SW训练深度生成模型的质量和收敛速度。
• 研究所提出的卷积切片框架的度量性质、样本复杂度和计算效率。
• 在多个数据集上展示CSW相较于标准SW在图像生成基准测试中的优越性能。

提出的方法

• 用卷积操作替代标准矩阵乘法，将图像张量投影到一维空间，以保留空间结构。
• 引入卷积切片Wasserstein（CSW），使用具有卷积核大小、步幅和空洞率的可学习卷积滤波器，实现具有空间感知能力的切片。
• 通过引入非线性激活函数（如ReLU）和池化机制，设计CSW的变体以增强表征能力。
• 将CSW距离公式化为对卷积投影的期望一维Wasserstein距离，通过卷积滤波器的蒙特卡洛采样进行近似。
• 通过将CSW用作损失函数，将其集成到生成对抗网络（GANs）中，用于训练生成器和判别器网络。
• 在判别器中应用谱归一化和残差块，以在使用CSW作为目标函数时稳定训练。

实验结果

研究问题

• RQ1 卷积算子能否替代切片Wasserstein中的向量化投影，以更好地保留图像的空间结构？
• RQ2 与使用密集投影矩阵的标准SW相比，所提出的卷积切片框架是否能降低内存使用？
• RQ3 在样本复杂度、计算复杂度和度量性质方面，CSW与标准SW相比如何？
• RQ4 CSW是否能提升在图像数据集上深度生成建模的训练稳定性和收敛速度？
• RQ5 基于CSW的GAN在CIFAR10、STL10和CelebA-HQ上的FID和Inception Score方面，其经验性能如何？

主要发现

• 与传统SW相比，CSW在训练深度生成模型时收敛更快，表现为训练轮次中FID分数更低。
• 在CIFAR10上，NCSW-d变体（带空洞卷积）在1000轮训练后达到最终FID分数42.30和Inception Score 6.31。
• 在所有评估的数据集（CIFAR10、STL10和CelebA-HQ）上，所提出的CSW变体在FID和Inception Score方面均持续优于标准SW。
• 使用卷积切片可避免存储完整投影向量，从而减少内存开销，使高分辨率图像的可扩展应用成为可能。
• 实证结果表明，CSW通过学习更具信息量和多样性的投影方向，缓解了标准SW中投影产生零点Dirac delta的“坍塌”切片问题。
• 理论与实证分析证实，CSW保持了度量性质，并实现了𝒪(n⁻¹/²)量级的样本复杂度，避免了维度灾难。