[論文レビュー] Amortized Projection Optimization for Sliced Wasserstein Generative Models

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

Human Pose and Action Recognition被引用数 4

ひとこと要約

本稿では、反復的最適化ループを排除するために最適な射影方向を学習することで、スライス Wasserstein 生成モデルの高速化を図る、アモアタイズド射影最適化を提案する。この手法は、線形、一般化線形、非線形の3つのアモアタイズドモデルを導入し、新しいミニバッチ損失（LA-SW、GA-SW、NA-SW）を生成する。これらの損失は、Max-SW や大規模な L を用いた標準的な SW よりも高速なトレーニングと優れた生成性能を達成する。

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.

研究の動機と目的

• スライス Wasserstein GAN における反復的射影最適化の高い計算コストを低減すること。
• トレーニングを遅くする二重最適化ループ（グローバルモデルトレーニング + 局所的射影最適化）を軽減すること。
• サンプル品質を損なわずにトレーニング効率と生成性能を向上させること。
• 情報豊かな射影方向を学習することで、アモアタイズド最適化と最適輸送を統合すること。
• スケーラブルで微分可能なモデルを構築し、ミニバッチ分布に対する最適な射影を予測すること。

提案手法

• ミニバッチ確率測度間の最適な射影方向を予測するアモアタイズドモデルを提案し、反復的最適化に代わる。
• 線形、一般化線形、非線形の3種類のアモアタイズドモデルを設計し、それぞれが新しいミニバッチ損失（LA-SW、GA-SW、NA-SW）に対応する。
• 入力ミニバッチを最適な射影方向にマップするためのニューラルネットワークを学習させ、エンドツーエンドのバックプロパゲーションを可能にする。
• 予測された方向をスライス Wasserstein 距離の計算に使用し、モンテカルロサンプリングや反復的最大化に代わる。
• SNGAN などの深層生成モデルにアモアタイズド損失を統合し、モデルパラメータと射影方向の共同最適化を実現する。
• 可学習重みを用いたパラメータ化と確率的勾配降下法を用いることで、微分可能性とスケーラビリティを確保する。

実験結果

リサーチクエスチョン

• RQ1 アモアタイズド最適化は、スライス Wasserstein GAN における反復的射影探索を置き換えることでトレーニング時間を短縮できるか？
• RQ2 線形、一般化線形、非線形の異なるアモアタイズドモデルアーキテクチャは、情報豊かな射影方向を予測する上でどのように比較されるか？
• RQ3 提案されたアモアタイズドスライス Wasserstein 損失は、Max-SW や大規模な L を用いた標準的な SW よりもトレーニング収束性とサンプル品質が向上するか？
• RQ4 既存の SW 変種と比較して、提案されたアモアタイズドモデルの計算およびメモリオーバーヘッドはどの程度か？
• RQ5 学習された射影方向は、異なるデータ分布やミニバッチ構成に一般化可能か？

主な発見

• 提案された LA-SW 損失は、L=1 の SW や Max-SW と同等のトレーニング速度を達成し、高いサンプル品質を維持する。
• LA-SW は、L ≥ 100 の SW や T₂ ≥ 10 の Max-SW よりも高速で、メモリ使用量も少ないため、最も効率的なバージョンである。
• GA-SW と NA-SW はやや遅延が大きく、メモリ使用量も高くなる（最大 6.16 GB）が、Max-SW や大規模 L を用いた SW よりも FID と IS スコアが優れている。
• LA-SW から生成された画像は、定量的指標と整合した質の高いものであり、可視化においても選別なしに一貫性がある。
• アモアタイズドモデルは、二重最適化の負担を顕著に軽減し、スライス Wasserstein GAN のスケーラブルなトレーニングを可能にする。
• 実験結果から、LA-SW、GA-SW、NA-SW は、CIFAR10、CelebA、STL10、CelebA-HQ において、Max-SW や標準的な SW と比較して FID とインセプションスコアの両方で優れていることが示された。