[Paper Review] Augmented Sliced Wasserstein Distances

Xiongjie Chen, Yongxin Yang+1 more|arXiv (Cornell University)|Jun 15, 2020

Advanced Numerical Analysis Techniques39 references4 citations

TL;DR

This paper proposes Augmented Sliced Wasserstein Distances (ASWD), a novel metric that improves computational efficiency and accuracy of Wasserstein distance approximations by mapping data to neural network-defined hypersurfaces before applying nonlinear projections. The method achieves superior performance over existing variants, especially in high-dimensional spaces, by enabling flexible, data-adaptive nonlinear slicing that enhances projection efficiency and metric validity via injective neural networks.

ABSTRACT

While theoretically appealing, the application of the Wasserstein distance to large-scale machine learning problems has been hampered by its prohibitive computational cost. The sliced Wasserstein distance and its variants improve the computational efficiency through the random projection, yet they suffer from low accuracy if the number of projections is not sufficiently large, because the majority of projections result in trivially small values. In this work, we propose a new family of distance metrics, called augmented sliced Wasserstein distances (ASWDs), constructed by first mapping samples to higher-dimensional hypersurfaces parameterized by neural networks. It is derived from a key observation that (random) linear projections of samples residing on these hypersurfaces would translate to much more flexible nonlinear projections in the original sample space, so they can capture complex structures of the data distribution. We show that the hypersurfaces can be optimized by gradient ascent efficiently. We provide the condition under which the ASWD is a valid metric and show that this can be obtained by an injective neural network architecture. Numerical results demonstrate that the ASWD significantly outperforms other Wasserstein variants for both synthetic and real-world problems.

Motivation & Objective

• To address the high computational cost and low projection efficiency of traditional Wasserstein and sliced Wasserstein distances in large-scale machine learning.
• To overcome the limitation of linear projections in standard sliced Wasserstein distances, which often yield trivially small distances in high dimensions.
• To develop a differentiable, scalable metric that maintains theoretical validity while capturing complex data structures through learned nonlinear projections.
• To enable end-to-end training in deep generative models by ensuring the distance metric is compatible with gradient-based optimization.

Proposed method

• The method maps input samples to higher-dimensional hypersurfaces using an injective neural network, transforming the data distribution into a nonlinear manifold.
• Linear projections are then applied to the embedded hypersurfaces, which induce flexible, nonlinear projections in the original data space.
• The ASWD is defined as the expected Wasserstein distance between one-dimensional projections of the transformed distributions, computed efficiently via sorting.
• The hypersurface parameters are optimized via gradient ascent to maximize the distance between projected distributions, improving discriminative power.
• Theoretical conditions are derived under which ASWD is a valid metric, requiring the neural network to be injective and the projection distribution to be sufficiently rich.
• The approach supports end-to-end training and is compatible with differentiable solvers, enabling integration into deep learning pipelines.

Experimental results

Research questions

• RQ1 Can nonlinear, data-adaptive projections derived from learned hypersurfaces significantly improve the projection efficiency of sliced Wasserstein distances in high-dimensional spaces?
• RQ2 Does the proposed ASWD maintain theoretical validity as a metric when constructed via neural network embeddings and random projections?
• RQ3 How does ASWD compare to existing variants like SWD, Max-SWD, GSWD, and DSWD in terms of accuracy and computational efficiency across synthetic and real-world distributions?
• RQ4 Can ASWD achieve high performance with a small number of projections, reducing computational cost while preserving discriminative power?

Key findings

• ASWD significantly outperforms SWD, Max-SWD, GSWD, and DSWD in both synthetic and real-world tasks, especially in high-dimensional settings like 100D Gaussians.
• With as few as 10 projections, ASWD produces sharper and more realistic image translations in color transfer tasks than SWD and Max-SWD.
• The ASWD-generated barycenters are geometrically meaningful and comparable to full Wasserstein barycenters, while being computationally more efficient than SWD and DSWD.
• The method achieves higher projection efficiency than linear projections in SWD, as evidenced by histograms showing fewer trivially small distances in 100D space.
• The use of injective neural networks ensures that ASWD satisfies the conditions to be a valid metric, providing theoretical grounding for its use.