Wasserstein Gradient Flow (WGF) describes the gradient dynamics of probability density within the Wasserstein space. WGF provides a promising approach for conducting optimization over the probability distributions. Numerically approximating the continuous WGF requires the time discretization method. The most well-known method for this is the JKO scheme. In this regard, previous WGF models employ the JKO scheme and parametrize transport map for each JKO step. However, this approach results in quadratic training complexity $O(K^2)$ with the number of JKO step $K$. This severely limits the scalability of WGF models. In this paper, we introduce a scalable WGF-based generative model, called Semi-dual JKO (S-JKO). Our model is based on the semi-dual form of the JKO step, derived from the equivalence between the JKO step and the Unbalanced Optimal Transport. Our approach reduces the training complexity to $O(K)$. We demonstrate that our model significantly outperforms existing WGF-based generative models, achieving FID scores of 2.62 on CIFAR-10 and 6.19 on CelebA-HQ-256, which are comparable to state-of-the-art image generative models.

本文提出了一种可扩展的基于Wasserstein梯度流的生成模型，称为Semi-dual JKO (S-JKO)，通过使用JKO步骤的半对偶形式，将训练复杂度降低至O(K)，实验证明该模型在CIFAR-10和CelebA-HQ-256数据集上取得了显著优于现有WGF模型的FID分数，与最先进的图像生成模型相当。

可扩展的Wasserstein梯度流用于非平衡最优传输的生成建模