Max-Sliced Wasserstein Distance and its use for GANs

TL;DR

The paper analyzes the generalization and sample-efficiency of distance metrics used in GANs, showing that the sliced Wasserstein distance $\tilde{W}_2$ generalizes well for Gaussian families while the full Wasserstein distance $W_2$ does not. To address projection-coverage limits, it introduces the max-sliced Wasserstein distance $\max-\tilde{W}_2$ and proves it retains polynomial-sample generalization. It furthermore proposes a practical max-sliced GAN algorithm that alternates discriminator and projection-direction optimization, achieving high- resolution image generation with far fewer projections. Empirically, the method yields competitive unsupervised word translation results and high-quality 256×256 image samples (CelebA-HQ, LSUN) while significantly reducing computational demands. Overall, the approach offers improved stability and efficiency for training GANs on high-dimensional data with strong practical impact for scalable generative modeling.

Abstract

Generative adversarial nets (GANs) and variational auto-encoders have significantly improved our distribution modeling capabilities, showing promise for dataset augmentation, image-to-image translation and feature learning. However, to model high-dimensional distributions, sequential training and stacked architectures are common, increasing the number of tunable hyper-parameters as well as the training time. Nonetheless, the sample complexity of the distance metrics remains one of the factors affecting GAN training. We first show that the recently proposed sliced Wasserstein distance has compelling sample complexity properties when compared to the Wasserstein distance. To further improve the sliced Wasserstein distance we then analyze its `projection complexity' and develop the max-sliced Wasserstein distance which enjoys compelling sample complexity while reducing projection complexity, albeit necessitating a max estimation. We finally illustrate that the proposed distance trains GANs on high-dimensional images up to a resolution of 256x256 easily.

Figures (6)

• Figure 1: Convergence of the mean for different sampling strategies for learning the mean of a$d$-dimensional Gaussian using the sliced Wasserstein distance and the max-sliced Wasserstein distance. Numbers in the legend denote the number of projection directions used.
• Figure 2: The discriminator is able to identify important projection directions. The discriminator transforms the distributions in Fig. 2a to Fig. 2b. In this new space, the discriminator's direction is aligned with the one along which the distributions are the most dissimilar as shown in Fig. 2c.
• Figure 3: Generated samples ($256 \times 256$ ) from CelebA-HQ.
• Figure 4: Generated samples ($256 \times 256$ ) from LSUN Bedrooms.
• Figure 5: Random samples ($512 \times 512$ ) after 50,000 iterations of training on CelebA-HQ.