[Paper Review] Revisiting Sliced Wasserstein on Images: From Vectorization to Convolution

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

Generative Adversarial Networks and Image Synthesis4 citations

TL;DR

This paper proposes Convolutional Sliced Wasserstein (CSW) and its variants to overcome limitations of conventional sliced Wasserstein (SW) in image generation by replacing vectorized projections with learnable convolutional operators. By incorporating stride, dilation, and non-linear activation into convolutional layers, CSW preserves spatial structure, reduces memory usage, and achieves faster convergence and better performance in training generative models on CIFAR10, STL10, and 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.

Motivation & Objective

• Address the inefficiency and structural loss in conventional sliced Wasserstein (SW) when applied to images due to image vectorization.
• Overcome memory inefficiency from storing high-dimensional projection vectors in standard SW.
• Improve the quality and convergence speed of deep generative models trained with SW by redefining slicing via convolutional operators.
• Investigate the metric properties, sample complexity, and computational efficiency of the proposed convolutional slicing framework.
• Demonstrate superior performance of CSW over standard SW in image generation benchmarks across multiple datasets.

Proposed method

• Replace standard matrix multiplication with convolutional operations to project image tensors into one-dimensional spaces, preserving spatial structure.
• Introduce Convolutional Sliced Wasserstein (CSW) using learnable convolutional filters with kernel size, stride, and dilation to enable spatially aware slicing.
• Design variants of CSW by incorporating non-linear activation functions (e.g., ReLU) and pooling mechanisms to enhance representational capacity.
• Formulate the CSW distance as the expected one-dimensional Wasserstein distance over convolved projections, approximated via Monte Carlo sampling of convolutional filters.
• Integrate CSW into generative adversarial networks (GANs) by using it as a loss function for training generator and discriminator networks.
• Apply spectral normalization and residual blocks in the discriminators to stabilize training while using CSW as the objective.

Experimental results

Research questions

• RQ1 Can convolutional operators replace vectorized projections in sliced Wasserstein to better preserve image spatial structure?
• RQ2 Does the proposed convolutional slicing framework reduce memory usage compared to conventional SW with dense projection matrices?
• RQ3 How does CSW compare to standard SW in terms of sample complexity, computational complexity, and metric properties?
• RQ4 Does CSW improve training stability and convergence speed in deep generative modeling on image datasets?
• RQ5 What is the empirical performance of CSW-based GANs in terms of FID and Inception Score across CIFAR10, STL10, and CelebA-HQ?

Key findings

• CSW achieves faster convergence in training deep generative models compared to conventional SW, as shown by lower FID scores over training epochs.
• On CIFAR10, the NCSW-d variant (with dilation) achieved a final FID score of 42.30 and Inception Score of 6.31 at 1000 epochs.
• The proposed CSW variants consistently outperform standard SW in both FID and Inception Score across all evaluated datasets: CIFAR10, STL10, and CelebA-HQ.
• The use of convolutional slicing reduces memory overhead by avoiding storage of full projection vectors, enabling scalable application to high-resolution images.
• Empirical results show that CSW mitigates the issue of 'collapsed' slices in standard SW, where projections yield Dirac deltas at zero, by learning more informative and diverse projection directions.
• Theoretical and empirical analysis confirms that CSW maintains metric properties and achieves sample complexity of order 𝒪(n⁻¹/²), avoiding the curse of dimensionality.