Source-linked AI summary
Generalized Sliced Wasserstein Distances
Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, Gustavo K. Rohde
TL;DR
The paper addresses the computational cost of optimal-transport distances and the growing projection requirements of sliced-Wasserstein methods in high dimensions. It generalizes linear slicing through the generalized Radon transform, defines GSW and max-GSW distances, and establishes conditions for their metric validity. Across generative-modeling applications, the authors report superior performance over classical sliced-Wasserstein distances, while noting metric guarantees depend on injectivity and projection counts remain a practical consideration.
Problem
High-dimensional Wasserstein distances are computationally burdensome, while sliced-Wasserstein distances can require many projections as dimensionality grows.
Method
The paper uses generalized Radon transforms to define nonlinear generalized sliced-Wasserstein distances and a maximum variant based on one maximal-distance projection.
Results
The proposed GSW and max-GSW distances empirically outperform classical sliced-Wasserstein distances across several generative-modeling applications.
Takeaways & Limitations
Nonlinear slicing provides a generalized optimal-transport metric family, with max-GSW reducing projection operations by using one projection at the cost of optimizing over projectors.
Takeaways & Limitations
GSW and max-GSW are guaranteed to be distances only when the chosen generalized Radon transform is injective; otherwise they are pseudo-metrics.
Abstract
from arXiv · showhide
The Wasserstein distance and its variations, e.g., the sliced-Wasserstein (SW) distance, have recently drawn attention from the machine learning community. The SW distance, specifically, was shown to have similar properties to the Wasserstein distance, while being much simpler to compute, and is therefore used in various applications including generative modeling and general supervised/unsupervised learning. In this paper, we first clarify the mathematical connection between the SW distance and the Radon transform. We then utilize the generalized Radon transform to define a new family of distances for probability measures, which we call generalized sliced-Wasserstein (GSW) distances. We also show that, similar to the SW distance, the GSW distance can be extended to a maximum GSW (max-GSW) distance. We then provide the conditions under which GSW and max-GSW distances are indeed distances. Finally, we compare the numerical performance of the proposed distances on several generative modeling tasks, including SW flows and SW auto-encoders.
1. Introduction
The paper addresses the computational burden of high-dimensional Wasserstein and sliced-Wasserstein distances by extending linear slicing to nonlinear projections. It introduces GSW and max-GSW distances to reduce projection requirements while retaining distance-metric guarantees under stated conditions.
- High-dimensional Wasserstein distance evaluation is generally numerically intractable and often becomes the computational bottleneck for large-scale data analysis.
- The sliced-Wasserstein distance reduces this burden by projecting high-dimensional distributions onto one-dimensional representations with efficiently computable Wasserstein distances.
- In very high dimensions, sliced-Wasserstein computations may require rapidly increasing numbers of random linear projections to capture thin-manifold structure.
- The paper uses generalized Radon-transform mathematics to extend sliced-Wasserstein distances from linear slicing to an entire class of nonlinear GSW distances.
- GSW remains a valid distance metric for polynomial projections under identified conditions, while max-GSW uses a single maximal-distance projection and trades projection cost for optimization over projectors.
2. Background
This section introduces Wasserstein and sliced-Wasserstein distances, emphasizing how one-dimensional projections make optimal transport more tractable. It also explains Radon-transform foundations, practical projection limits, and maximum sliced-Wasserstein alternatives.
- Wasserstein distance: The p-Wasserstein distance is defined through optimal mass transportation plans between probability measures.The plans must have the two input measures as their respective marginals.
- Wasserstein distance: One-dimensional Wasserstein distance has a closed-form solution, enabling efficient computation by sorting empirical samples.For M samples, computation requires O(M) operations at best and O(M log M) at worst.
- Radon transform: The Radon transform maps a function to integrals over hyperplanes, and filtered back-projection reconstructs the function from its projections.Each slice is filtered and smeared back along its corresponding hyperplane before the contributions are summed.
- Sliced-Wasserstein distance: Sliced-Wasserstein distance averages one-dimensional Wasserstein distances after projecting high-dimensional distributions along linear directions.The projections are represented through the Radon transform and evaluated using finitely many Monte Carlo samples in practice.
- Sliced-Wasserstein distance: SW computation becomes increasingly difficult in high dimensions because accurately representing distributions may require O(Ld−1) projections.This is especially problematic when data lies on a thin manifold or when uniformly sampled directions are nearly orthogonal to discriminative structure.
- Maximum sliced-Wasserstein distance: Maximum sliced-Wasserstein distance addresses projection selection by taking the largest one-dimensional Wasserstein discrepancy, and it is a valid distance.Its metric properties follow from those of the underlying Wasserstein distance.
3. Generalized Sliced-Wasserstein Distances
The paper generalizes sliced-Wasserstein distances by replacing linear projections with nonlinear projections defined through generalized Radon transforms. It establishes when the resulting GSW and max-GSW constructions are valid distances and identifies injectivity of the transform as the decisive condition.
- GSW extends sliced-Wasserstein distance by obtaining one-dimensional representations through nonlinear projections based on the generalized Radon transform.The generalized Radon transform integrates a density over hypersurfaces defined by level sets of a defining function.
- The generalized Radon transform extends classical hyperplane integration to integration over (d −1)-dimensional hypersurfaces.The classical Radon transform is recovered when g(x, θ) = ⟨x, θ⟩.
- The GSW construction aggregates Wasserstein distances over a compact feasible parameter set, whereas max-GSW addresses the projection-complexity issue through a maximum over projections.The proof of the triangle inequality uses the Wasserstein triangle inequality and Minkowski’s inequality for GSW, with an analogous result for max-GSW.
- GSW and max-GSW are distances if and only if equality of their transformed one-dimensional distributions implies equality of the original probability measures.This requirement is equivalent to injectivity of the generalized Radon transform.
- If the generalized Radon transform is not injective, GSW and max-GSW remain pseudo-metrics rather than distances.They still satisfy non-negativity, symmetry, the triangle inequality, and zero self-distance.
- Replacing linear projections with polynomial projections can yield a valid distance metric, while odd-degree homogeneous polynomials provide injective generalized Radon transforms.The linear case is recovered when the polynomial degree is m = 1.
4. Numerical Implementation
The numerical procedures approximate GSW from samples by transforming, sorting, and comparing one-dimensional projections, while max-GSW optimizes the projection parameter directly.
- GSW can be approximated from samples by applying the defining function to each sample, sorting the resulting values, and computing one-dimensional Wasserstein distances.The procedure is summarized in Algorithm 1.
- One-dimensional projected densities may require fewer samples to estimate than the original high-dimensional density.The projections are one-dimensional even when estimating the distribution in R^d is sample-intensive.
- max-GSW uses an EM-like optimization scheme that alternates between sorting transformed samples for a fixed parameter and updating that parameter.The parameter update uses ADAM and projection onto the feasible parameter set.
- The max-GSW implementation optimizes the actual Wasserstein distance rather than using perceptron or penalized-discriminant heuristics to find a pseudo-optimal slice.This distinguishes the procedure from approaches in Deshpande et al. (2018) and Kolouri et al. (2019).
5. Experiments
The experiments evaluate GSW variants in distribution-matching flows and auto-encoders, comparing nonlinear defining functions and baselines under repeated optimization. Polynomial GSW variants capture nonlinear structure, while max-GSWAE compares favorably with WAE-GAN on the reported auto-encoder objectives.
- Generalized Sliced-Wasserstein Flows: The flow experiment compares linear, circular, and degree-3 or degree-5 homogeneous polynomial defining functions on five target distributions using 10 random projections.Each experiment was repeated 100 times, with mean and standard deviation reported using the 2-Wasserstein distance during optimization.
- Generative Modeling via Auto-Encoders: The SWAE experiments replace the sliced-Wasserstein term with GSW or max-GSW while enforcing a prior distribution in latent space and reconstructing input images.The tested GSWAE variants use linear, circular, and degree-3 homogeneous polynomial defining functions.
- Generative Modeling via Auto-Encoders: The auto-encoder evaluation measures 2-Wasserstein distances for both latent-prior matching and input-output reconstruction across repeated runs.The experiments report averages over 50 runs for the GSWAE study and 10 runs for the comparison with SWAE and WAE-GAN.
- Generative Modeling via Auto-Encoders: Max-GSWAE matches WAE-GAN in latent distribution matching and appears to provide a better auto-encoder training objective for input-output distribution matching.
6. Conclusion
The paper generalizes sliced-Wasserstein distances from linear to nonlinear slicing through generalized Radon transforms and establishes valid GSW and max-GSW metrics. Experiments report superior performance over classical sliced-Wasserstein distances in generative modeling applications.
- The proposed GSW family generalizes sliced-Wasserstein metrics by replacing linear slicing with nonlinear slicing of probability measures.
- The paper provides theoretical conditions ensuring that GSW and max-GSW are distance functions and reports improved performance over classical sliced-Wasserstein distances in generative modeling.
9. Non-negativity and Symmetry of the GSW and max-GSW Distances
The paper derives non-negativity and symmetry for GSW and max-GSW from corresponding properties of the p-Wasserstein distance. These arguments use the projected generalized Radon-transform representations.
- GSW and max-GSW are non-negative because they are constructed from non-negative p-Wasserstein distances.
- For max-GSW, the selected projection is defined by maximizing the projected p-Wasserstein distance over θ.
- The symmetry of GSW and max-GSW follows from the symmetry of the p-Wasserstein distance and the defining equations for the two constructions.
10. Additional Experimental Results
Additional maximum-GSW flow results compare GSW and max-GSW across the five datasets. Max-GSW improves GSW in the majority of reported cases while using one projection instead of ten.
- Max-GSW improves GSW performance in the majority of cases across the five datasets in the additional flow experiments.
- GSW uses 10 random projections, whereas max-GSW uses only one projection by definition.
11. Implementation Details
The implementation uses specified neural architectures for the encoder, decoder, and WAE-GAN adversary, with fixed optimization settings for training and projection maximization.
- Optimization: ADAM optimization used learning rate lr = 0.001 with PyTorch’s default momentum parameters.These settings were used for all optimizations.
- Optimization: The critic iteration and max-GSW projection maximization were both set to 50.The same iteration count was used for adversarial training and projection maximization.
- Network architectures: The encoder processes x ∈R28×28 through convolutional, LeakyReLU, average-pooling, and flattening layers.The listed encoder uses channel widths 16, 32, and 64 before flattening.
- Network architectures: The decoder maps z ∈R2 through fully connected, reshaping, upsampling, convolutional, and LeakyReLU layers to produce one output channel.The decoder expands the latent vector through a 4 × 4 × 64 representation and successive upsampling stages.
- Network architectures: The WAE-GAN uses an adversary network, whose architecture is introduced separately.The supplied implementation passage identifies the adversary but does not list its layer structure.