Source-linked AI summary
Max-Sliced Wasserstein Distance and its use for GANs
Ishan Deshpande, Yuan-Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, Alexander Schwing
TL;DR
High-dimensional GAN training is limited by the sample complexity of distribution distances. The paper analyzes sliced Wasserstein distance, introduces max-sliced Wasserstein distance to reduce projection complexity, and demonstrates high-resolution GAN training.
Problem
Distance metrics used in GANs can have unfavorable sample complexity, while sliced Wasserstein distance has a projection-complexity limitation when estimated with finitely many random directions.
Method
The paper analyzes Wasserstein and sliced Wasserstein distances and introduces max-sliced Wasserstein distance, which maximizes projected Wasserstein distance over directions.
Results
Max-sliced Wasserstein distance enables GAN training on high-resolution CelebA-HQ and LSUN Bedrooms images and improves other distribution-matching tasks.
Takeaways & Limitations
Max-sliced Wasserstein distance combines the better sample complexity of sliced Wasserstein distance with lower projection complexity than sliced Wasserstein GAN.
Takeaways & Limitations
The theoretical analysis requires approximately O(1/β^2) samples as the two distributions become closer, and sliced Wasserstein estimation remains limited by projection complexity with finite random directions.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper addresses GAN scaling through distance-metric sample and projection complexity, analyzing sliced Wasserstein distance and introducing max-sliced Wasserstein distance for training high-resolution GANs.
- Motivation: High-dimensional GAN scaling is affected by sample complexity, because several distribution distances cannot be reliably approximated with polynomially many samples.The paper motivates an alternative to restricting discriminator classes by focusing on the distance metric itself.
- Sliced Wasserstein distance: Sliced Wasserstein distance averages Wasserstein distances after projecting distributions along a few randomly selected directions, improving behavior relative to Wasserstein distance for Gaussian distributions.The paper proves generalizability for Gaussian distributions and links this result to the observed behavior of sliced Wasserstein GANs.
- Projection complexity: Random projections can lose information, creating projection complexity that requires many directions for sliced Wasserstein distance estimation.This limitation motivates choosing a direction based on the maximum projected distance rather than averaging only random directions.
- Max-sliced Wasserstein distance: The proposed max-sliced Wasserstein distance projects onto the direction maximizing the projected distance and remains generalizable for Gaussian distributions.The method addresses projection complexity while requiring estimation of a maximum.
- Applications: Max-sliced Wasserstein distance trains GANs on high-resolution CelebA-HQ and LSUN Bedrooms images and improves other distribution-matching tasks such as unpaired word translation.The introduction reports high-resolution image generation and improved performance beyond image-generation experiments.
- Sample complexity: For a certain class of distributions, Wasserstein distance has exponential sample complexity, whereas sliced Wasserstein distance has polynomial sample complexity.The analysis compares the two distances directly through their population-to-empirical estimation behavior.
2. Background
The background frames GAN training as distribution-distance minimization and reviews Wasserstein and sliced Wasserstein approaches, including their estimation procedures and adversarial feature-space formulation.
- Generative modeling: Generative modeling learns a distribution Pg that matches an unknown data distribution Pd from dataset samples.GANs formulate this objective as minimizing a distance or divergence between the data and generated distributions.
- Wasserstein distance: Wasserstein-p distance compares distributions through optimal joint distributions whose marginals are Pg and Pd.The set Π(Pg, Pd) contains all admissible couplings between the two distributions.
- Wasserstein distance: Wasserstein-1 distance is estimated through Kantorovich-Rubinstein duality by optimizing over 1-Lipschitz functions, commonly represented with deep networks.The Lipschitz constraint must be enforced during estimation.
- Sliced Wasserstein distance: Wasserstein-distance approaches can suffer instability from incorrect estimation associated with the curse of dimensionality, motivating sliced distances based on one-dimensional projections.Sliced Wasserstein methods estimate distances of projected one-dimensional distributions instead of working directly in the original space.
- Sliced Wasserstein distance: Sliced Wasserstein distance integrates Wasserstein distances between projected distributions over directions on the unit sphere and is a true metric.Its metric properties include non-negativity, identity of indiscernibles, symmetry, and subadditivity.
- Estimation and GAN formulation: In practice, sliced Wasserstein distance replaces the directional integral with a finite sum over normalized random Gaussian directions and uses samples from the data and generator distributions.For one-dimensional projected samples, Wasserstein-2 distance can be computed by sorting.
- Adversarial feature space: Adversarial feature-space variants transform real and generated samples through a discriminator-parameterized deep feature map before distance minimization.The discriminator learns the feature space by classifying real and fake data.
3. Analysis and Max-Sliced Distance
The analysis shows that sliced Wasserstein distance has polynomial sample complexity for Gaussian distributions, unlike Wasserstein distance, but random projections create a projection-complexity bottleneck. Max-sliced Wasserstein addresses this bottleneck by selecting the most informative direction while retaining the sliced distance’s sample-complexity advantage.
- Sample complexity: The sliced Wasserstein-2 distance is generalizable for Gaussian distributions, whereas the Wasserstein-2 distance is not.Generalizability means empirical distance estimates approximate population distance with polynomial sample size.
- Sample complexity: For GAN training, sliced Wasserstein distance can provide a more accurate training signal with fixed computational resources and a more stable discriminator.
- Projection complexity: Finite random projections create projection complexity because directions differ in how meaningfully they distinguish distributions.The analysis studies this effect using infinitely many samples but finitely many projection directions.
- Projection complexity: In high dimensions, random directions are usually nearly orthogonal to the mean-shift direction, producing small updates and very slow convergence that worsens with dimension.For the Gaussian example, the projected Wasserstein distance is β|ê^Tω|.
- Max-sliced Wasserstein distance: Max-sliced Wasserstein selects the direction maximizing projected distance, yielding larger updates and fast convergence while preserving the sliced distance’s Gaussian generalizability.The method is defined by comparing distributions along the most meaningful projection direction and does not trade away the sample-complexity benefit.
13 end
The max-sliced Wasserstein distance is a valid, generalizable metric that addresses sliced Wasserstein projection complexity. The proposed GAN training procedure approximates the maximizing direction with a surrogate discriminator objective.
- The max-sliced Wasserstein-2 distance is a well-defined distance and is P-generalizable for Gaussian distributions.
- The max-sliced distance avoids the finite-random-projection limitation of sliced Wasserstein distance but requires estimating a maximum.The maximum is harder to estimate than a conventional random variable.
- The method achieves significant savings in projection directions compared with sliced Wasserstein GAN training.
- Joint optimization introduces additional concerns, including objective boundedness and the difficulty of jointly searching over feature and projection parameters.Regularization may require separate tuning for each use case.
- The surrogate is designed to learn a direction that separates transformed real and generated data, with log-loss providing a linear logistic-regression interpretation.The discriminator’s learned direction can approximate the maximizing projection direction.
- The training procedure optimizes discriminator feature-transform and projection parameters using a surrogate loss, then minimizes the Wasserstein-2 distance between projected real and generated features.The generator minimizes the projected Wasserstein distance after discriminator optimization.
4. Experiments
The experiments evaluate max-sliced GANs on unsupervised word translation and 256x256 image generation. The method outperforms the baseline across tested language pairs and matches or exceeds sliced Wasserstein GANs with far fewer projections.
- 4.1. Word Translation without Parallel Data: The word-translation experiment compares methods quantitatively without paired or parallel data.
- 4.2. Image Generation: Max-sliced GANs generate 256x256 images on CelebA and LSUN Bedrooms and are compared with sliced Wasserstein GANs.
- 4.2. Image Generation: Using one projection direction, max-sliced Wasserstein GAN produces samples comparable or better than sliced Wasserstein GAN with 10,000 projections.The comparison is shown for CelebA-HQ and LSUN Bedrooms samples.
- 4.2. Image Generation: The image models use a simple DCGAN extension with two additional strided convolutional layers to scale to 256x256.No special normalization or initialization is used.
5. Conclusion
The paper develops max-sliced Wasserstein training for GANs and analyzes its sample and projection complexity. Experiments show that the approach learns high-dimensional distributions with orders of magnitude fewer projection directions.
- The proposed max-sliced Wasserstein distance has better sample complexity than Wasserstein distance and better projection complexity than sliced Wasserstein distance.
- The paper develops a surrogate-loss approximation to the max-sliced distance and analyzes its approximation error.
- The approach learns high-dimensional distributions while requiring orders of magnitude fewer projection directions than sliced Wasserstein GANs.Both methods operate in a similar distance space.
Supplementary Material: Max-Sliced Wasserstein Distance and its use for GANs
The supplementary material section is represented by author and affiliation information rather than technical content.
- The supplied supplementary-material passages list the paper’s authors and their institutional affiliations.
- The listed contact information uses University of Illinois at Urbana-Champaign email addresses.
1. Experiments with Images at Higher Resolutions
At 512x512 resolution, experiments compare max-sliced and sliced Wasserstein GANs after 50,000 training iterations and present max-sliced samples from CelebA-HQ. The max-sliced model converges faster with far fewer projections and uses a simple end-to-end architecture without progressive growing or stacking.
- Convergence: The max-sliced Wasserstein GAN converges faster than the sliced Wasserstein GAN with substantially fewer projections, especially early in training.The comparison uses 100, 1,000, and 10,000 sliced-Wasserstein projections after 50,000 iterations.
- Convergence: The convergence advantage becomes more pronounced as the distribution dimension increases.
- Training setup: The model is trained end-to-end in one unified process using a simple architecture, without progressive growing, stacking, or other training tricks.
2. Proof of Claim 1 and Claim 3 in the paper
For the Gaussian family P = {N(a, I_d) | a ∈ R^d}, the max-sliced and sliced Wasserstein distances are P-generalizable with polynomially many samples, whereas the Wasserstein distance is not. The proof contrasts high-dimensional empirical Wasserstein behavior with projection-based concentration results.
- Main claim: For P = {N(a, I_d) | a ∈ R^d}, max-sliced Wasserstein-2 and Gaussian sliced Wasserstein distances are P-generalizable, but Wasserstein distance is not.
- Max-sliced distance: The max-sliced proof identifies the optimal projection direction as the direction connecting the Gaussian means, giving population distance β.
- Max-sliced distance: The empirical optimal projected distance is bounded around β, establishing polynomial-sample generalizability for the max-sliced metric.The analysis supplies lower and upper bounds, including a lower bound of β − 4δ.
- Sample requirement: Estimating the max-sliced distance requires more samples as the two distributions become closer; the stated dependence is approximately O(1/β^2).
2.5. Proof of Lemma 2.2
The proof of Lemma 2.2 establishes an upper bound for the empirical max-sliced objective uniformly over all projection directions, then derives the sliced Wasserstein generalization result using finitely many random directions. The uniform bound relies on a covering argument and concentration inequalities.
- Uniform upper bound: For every projection direction, the proof decomposes the objective into ordered projected-sample differences and a linear Gaussian term.
- Uniform upper bound: A δ-cover of the unit sphere reduces the uniform-in-direction analysis to finitely many anchor directions, with N = O(1/δ^d) points.
- Uniform upper bound: Gaussian concentration and spectral bounds control the decomposed terms, yielding high-probability bounds that hold across directions.
- Max-sliced result: The resulting max-sliced estimate holds with probability at least 1 − exp(−m) − 2n^−8 or a related combined bound.
- Proof limitation: The covering-and-union-bound proof is loose, adding an extra multiplicative factor of d and an extra additive factor of m; stronger bounds may be possible.
- Sliced result: For sliced Wasserstein distance, a finite set of K random directions is handled by a union bound, imposing requirements on both K and the sample size n.
3. Proof of Claim 2 in the paper
The proof establishes that max-sliced Wasserstein-2 is a well-defined distance by inheriting non-negativity, symmetry, and identity of indiscernibles from Wasserstein distance and proving the triangle inequality through an optimal projection.
- Metric properties: Max-sliced Wasserstein-2 is a well-defined distance between distributions.
- Metric properties: Non-negativity, symmetry, and identity of indiscernibles follow directly from the corresponding Wasserstein-2 properties.
- Triangle inequality: For the triangle inequality, the proof selects a maximizing projection direction and applies the Wasserstein-2 triangle inequality to the projected distributions.
- Triangle inequality: Taking maxima over projection directions completes the triangle-inequality argument for the max-sliced distance.
4. Architecture for Image Generation
Experiments at 256x256 use architectures described in Tables 1 and 2, with Adam optimization at a learning rate of 0.0001 and batch size 64.
- 256x256 experiments use the architectures described in Tables 1 and 2.The 512x512 experiments are described separately in Tables 3 and 4.
- Adam is used with a learning rate of 0.0001 for all experiments.
- The mini-batch size is 64 for 256x256 images.For 512x512 images, the mini-batch size is 32.