Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport

要約

Wasserstein Gradient Flow (WGF)は、Wasserstein空間内の確率密度の勾配ダイナミクスを記述します。WGFは確率分布に対する最適化を行うための有望なアプローチを提供します。連続WGFを数値的に近似するには、時間離散化法が必要です。この方法として最もよく知られているのがJKOスキームです。この点で、これまでのWGFモデルはJKOスキームを採用し、JKOステップごとに輸送マップをパラメトリック化しています。しかし、この方法では、JKOステップの数$K$に対して二次的な学習複雑度$O(K^2)$が生じる。これはWGFモデルのスケーラビリティを著しく制限する。本論文では、Semi-dual JKO (S-JKO)と呼ばれるスケーラブルなWGFベースの生成モデルを紹介する。我々のモデルは、JKOステップとアンバランス最適輸送の等価性から導かれる、JKOステップの半二元形式に基づいている。我々のアプローチは、学習の複雑さを$O(K)$に低減する。CIFAR-10で2.62、CelebA-HQ-256で5.46のFIDスコアを達成し、これは最先端の画像生成モデルに匹敵する。

要約(オリジナル)

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 5.46 on CelebA-HQ-256, which are comparable to state-of-the-art image generative models.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL