Amortized Projection Optimization for Sliced Wasserstein Generative Models
This addresses a computational bottleneck for researchers and practitioners using sliced Wasserstein distances in deep learning applications, though it is an incremental improvement by applying existing amortized optimization techniques.
The paper tackles the computational inefficiency of finding informative projecting directions in sliced Wasserstein distances for generative models by proposing amortized optimization to predict these directions, resulting in favorable performance on standard benchmark datasets.
Seeking informative projecting directions has been an important task in utilizing sliced Wasserstein distance in applications. However, finding these directions usually requires an iterative optimization procedure over the space of projecting directions, which is computationally expensive. Moreover, the computational issue is even more severe in deep learning applications, where computing the distance between two mini-batch probability measures is repeated several times. This nested loop has been one of the main challenges that prevent the usage of sliced Wasserstein distances based on good projections in practice. To address this challenge, we propose to utilize the learning-to-optimize technique or amortized optimization to predict the informative direction of any given two mini-batch probability measures. To the best of our knowledge, this is the first work that bridges amortized optimization and sliced Wasserstein generative models. In particular, we derive linear amortized models, generalized linear amortized models, and non-linear amortized models which are corresponding to three types of novel mini-batch losses, named amortized sliced Wasserstein. We demonstrate the favorable performance of the proposed sliced losses in deep generative modeling on standard benchmark datasets.