[Paper Review] Amortized Projection Optimization for Sliced Wasserstein Generative Models

Khai Nguyen, Nhat Ho|arXiv (Cornell University)|Mar 25, 2022

Human Pose and Action Recognition4 citations

TL;DR

This paper proposes amortized projection optimization to accelerate sliced Wasserstein generative models by learning to predict optimal projection directions, eliminating iterative optimization loops. The method introduces three amortized models—linear, generalized linear, and non-linear—yielding new mini-batch losses (LA-SW, GA-SW, NA-SW) that achieve faster training and better generative performance than Max-SW and standard SW with large L.

ABSTRACT

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.

Motivation & Objective

• Address the high computational cost of iterative projection optimization in sliced Wasserstein GANs.
• Reduce the nested optimization loop (global model training + local projection optimization) that slows down training.
• Improve training efficiency and generative performance without sacrificing sample quality.
• Bridge amortized optimization with optimal transport by learning informative projection directions.
• Develop scalable, differentiable models to predict optimal projections for mini-batch distributions.

Proposed method

• Propose amortized models to predict optimal projection directions between mini-batch probability measures, replacing iterative optimization.
• Design three types of amortized models: linear, generalized linear, and non-linear, each corresponding to a novel mini-batch loss (LA-SW, GA-SW, NA-SW).
• Train a neural network to map input mini-batches to optimal projection directions, enabling end-to-end backpropagation.
• Use the predicted directions in the sliced Wasserstein distance computation, replacing Monte Carlo sampling or iterative maximization.
• Integrate the amortized losses into deep generative models like SNGAN for joint optimization of model parameters and projection directions.
• Ensure differentiability and scalability by parameterizing the projection predictor with learnable weights and using stochastic gradient descent.

Experimental results

Research questions

• RQ1 Can amortized optimization replace iterative projection search in sliced Wasserstein GANs to reduce training time?
• RQ2 How do different amortized model architectures (linear, generalized linear, non-linear) compare in predicting informative projection directions?
• RQ3 Does the proposed amortized sliced Wasserstein loss improve training convergence and sample quality compared to Max-SW and standard SW with large L?
• RQ4 What is the computational and memory overhead of the proposed amortized models relative to existing SW variants?
• RQ5 Can the learned projection directions generalize across different data distributions and mini-batch configurations?

Key findings

• The proposed LA-SW loss achieves training speed comparable to SW with L=1 and Max-SW, while maintaining high sample quality.
• LA-SW is faster and uses less memory than SW with L ≥ 100 and Max-SW with T₂ ≥ 10, making it the most efficient variant.
• GA-SW and NA-SW show slightly slower training and higher memory usage (up to 6.16 GB) but achieve better FID and IS scores than Max-SW and SW with large L.
• Generated images from LA-SW are qualitatively superior and consistent with quantitative metrics, with no cherry-picking in visualization.
• The amortized models significantly reduce the nested optimization burden, enabling scalable training of sliced Wasserstein GANs.
• Empirical results show that LA-SW, GA-SW, and NA-SW outperform Max-SW and standard SW in both FID and inception score across CIFAR10, CelebA, STL10, and CelebA-HQ.