Augmented Sliced Wasserstein Distances
This work addresses the problem of scalable and accurate Wasserstein distance computation for machine learning practitioners, offering a novel method that improves upon existing approximations.
The authors tackled the computational inefficiency and low accuracy of sliced Wasserstein distances by proposing augmented sliced Wasserstein distances (ASWDs), which map samples to higher-dimensional hypersurfaces using neural networks to enable flexible nonlinear projections, resulting in significantly outperforming other variants on synthetic and real-world problems.
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.