[Paper Review] Distributional Sliced-Wasserstein and Applications to Generative Modeling

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

3D Shape Modeling and Analysis18 citations

TL;DR

This paper proposes Distributional Sliced-Wasserstein (DSW), a novel optimal transport distance that learns an optimal distribution over projection directions on the unit sphere, balancing directional diversity and informativeness. By regularizing for orthogonal directions, DSW improves sample efficiency and generative modeling performance over SW and Max-SW, achieving state-of-the-art results on MNIST, CelebA, CIFAR10, and LSUN with fewer projections.

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.

Motivation & Objective

• Address the inefficiency of Sliced-Wasserstein (SW) in requiring many projections due to uniform sampling of directions.
• Overcome the information loss in Max-Sliced-Wasserstein (Max-SW), which uses only a single optimal direction.
• Develop a principled probabilistic framework to learn a distribution over projections that maximizes both diversity and informativeness.
• Improve sample efficiency and performance in generative modeling by optimizing the projection distribution with geometric regularization.
• Demonstrate the superiority of DSW over existing sliced-Wasserstein variants in both linear and nonlinear settings on large-scale datasets.

Proposed method

• Propose Distributional Sliced-Wasserstein (DSW) as a generalization of Max-SW, where the projection distribution is learned rather than fixed to a delta or uniform measure.
• Formulate DSW as an optimization problem over probability measures on the unit sphere, constrained to favor orthogonal and diverse directions.
• Introduce a regularizing term that encourages projections to be far apart on the sphere, improving coverage of the data manifold.
• Use a dual empirical form with a neural network to approximate the optimal push-forward measure over the sphere, enabling end-to-end training.
• Extend the framework to nonlinear projections via the Generalized DSW (DGSW) using a circular function for nonlinearity.
• Solve the resulting optimization via stochastic gradient descent with a differentiable relaxation of the projection distribution.

Experimental results

Research questions

• RQ1 Can a learned distribution over projection directions improve the sample efficiency and accuracy of sliced-Wasserstein distances in high-dimensional distributions?
• RQ2 How does the performance of DSW compare to SW and Max-SW when using the same number of projections in generative modeling tasks?
• RQ3 Does regularizing for directional diversity in the projection space lead to better generalization and reconstruction quality?
• RQ4 Can the proposed distributional framework be extended to nonlinear projections while maintaining computational efficiency and performance?
• RQ5 How does DSW perform in joint inference and image reconstruction tasks compared to other sliced-Wasserstein variants?

Key findings

• DSW achieves superior performance to SW and Max-SW across all benchmark datasets (MNIST, CelebA, CIFAR10, LSUN) with the same number of projections.
• With only 10 projections, DSW successfully reconstructs MNIST digits, while SW and GSW fail to produce recognizable images.
• DSW-1000 achieves the best performance among all baselines in all evaluation metrics, including reconstruction error and distributional similarity.
• The nonlinear variant, DGSW, outperforms other nonlinear sliced-Wasserstein distances, especially in joint inference tasks.
• Max-GSW-NN performs well but is outperformed by DGSW, indicating that learning a distribution over directions is more effective than optimizing a single direction.
• The ablation study confirms that regularization for directional diversity significantly improves performance, with unregularized DSW reducing to Max-SW as a special case.