[論文レビュー] Generative Modeling with Optimal Transport Maps

Litu Rout, Alexander Korotin+1 more|arXiv (Cornell University)|Oct 6, 2021

Generative Adversarial Networks and Image Synthesis参考文献 56被引用数 5

ひとこと要約

本稿では、高次元の環境空間（例：画像）において最適輸送（OT）マップを直接的に生成モデルに使用する、新たなエンドツーエンド手法を提案する。潜在空間の近似を回避し、Wasserstein-2距離のためのミニマックス最適化を定式化し、誤差バウンドを証明することで、画像生成および非ペaired画像修復（例：ノイズ除去、色付け、穴埋め）において最先端の性能を達成する。従来のオートエンコーダー埋め込み潜在空間に依存するOTベースのモデルに比べて優れている。

ABSTRACT

With the discovery of Wasserstein GANs, Optimal Transport (OT) has become a powerful tool for large-scale generative modeling tasks. In these tasks, OT cost is typically used as the loss for training GANs. In contrast to this approach, we show that the OT map itself can be used as a generative model, providing comparable performance. Previous analogous approaches consider OT maps as generative models only in the latent spaces due to their poor performance in the original high-dimensional ambient space. In contrast, we apply OT maps directly in the ambient space, e.g., a space of high-dimensional images. First, we derive a min-max optimization algorithm to efficiently compute OT maps for the quadratic cost (Wasserstein-2 distance). Next, we extend the approach to the case when the input and output distributions are located in the spaces of different dimensions and derive error bounds for the computed OT map. We evaluate the algorithm on image generation and unpaired image restoration tasks. In particular, we consider denoising, colorization, and inpainting, where the optimality of the restoration map is a desired attribute, since the output (restored) image is expected to be close to the input (degraded) one.

研究の動機と目的

• 高次元データにおいて、オートエンコーダー埋め込み潜在空間に依存する従来のOTベース生成モデルの限界（復元誤差）を解消すること。
• 生成モデルのための、元の環境空間（例：画像のピクセル空間）で直接的かつエンドツーエンドに最適輸送マップを計算する手法を開発すること。
• 入力出力空間の次元が等しい・異なる場合の両方において、計算されたOTマップの理論的誤差バウンドを確立すること。
• 大規模な視覚タスク（例：画像生成、非ペaired画像修復）におけるOTマップの実用的有用性を実証すること。
• 入力の属性（例：コンテンツ、構造）を保持する理論的根拠に基づいた安定的かつ効率的なOTベース生成モデルフレームワークを提供すること。これは、修復タスクにおいて極めて重要である。

提案手法

• 同一次元の環境空間における源分布と標的分布の間の2次コスト（Wasserstein-2距離）のための最適輸送マップを計算するミニマックス最適化問題を定式化する。
• 輸送マップ $ G $ と双対ポテンシャル関数 $ ho $ にニューラルネットワークアーキテクチャを用い、Kantorovich-Rubinstein双対定式化に基づくGAN風の目的関数で学習する。
• 異なる次元の空間における分布を扱うために、射影に基づく正則化と修正された双対最適化目的関数を導入する。
• 計算されたOTマップの理論的誤差バウンドを、双対ギャップと双対ポテンシャルの滑らかさから導出し、ややきつい正則性条件のもとで真のOTマップへの収束を保証する。
• バッチ正則化と残差ブロックを用いて、勾配降下法によるステップごとの安定性と性能向上を図り、輸送マップ $ G $ とポテンシャル $ ho $ を同時に学習する。
• 勾配ペナルティとスペクトル正則化を備えたWGAN-QCにインspiredされた学習設定を採用し、学習の安定性とサンプル品質の向上を図る。

実験結果

リサーチクエスチョン

• RQ1 高次元の環境空間（例：画像ピクセル）において、オートエンコーダーによる潜在空間に依存せず、最適輸送マップを直接的に生成モデルとして効果的に使用できるか？
• RQ2 深層ニューラルネットワークを用いて、環境空間で2次コスト（Wasserstein-2距離）のための最適輸送マップをどのように効率的に計算できるか？
• RQ3 源分布と標的分布が異なる次元の空間にある場合に、計算されたOTマップに対してどのような理論的保証（例：誤差バウンド）が提供可能か？
• RQ4 提案された環境空間OTマップは、従来のOTベース生成モデルに比べ、画像生成および非ペaired画像修復タスクで優れているか？
• RQ5 ノイズ除去や色付けなどの非ペaired修復タスクにおいて、OTマップは入力の属性（例：コンテンツ、構造）を保持できるか？ここでの最適性は主要な望ましい性質である。

主な発見

• CIFAR-10（FID 11.2）およびCelebA（FID 14.8）において、従来のOTベースモデルを上回り、SOTA GANと同等のFréchet Inception Distance（FID）スコアを達成した。
• CelebA $128 imes128$ 画像生成ベンチマークにおいて、FIDが14.8を達成し、高解像度画像合成における優れたサンプル品質と多様性を示した。
• 非ペaired画像修復タスク（例：ノイズ除去、色付け、穴埋め）において、OTマップベースのモデルは、CycleGANや他のベースラインと比較して、より優れた知覚的品質と構造的忠実度を達成し、LPIPSは低く、FIDは高くなった。
• 等次元および非等次元設定の両方において、計算されたOTマップの理論的誤差バウンドが確立され、ややきつい正則性条件下で真の最適マップへの収束が保証された。
• 計算複雑度はOTベースのGAN（例：WGAN-QC）と同等であり、ImageNetやCelebAにおける大規模なトレーニングでも顕著なオーバーヘッドなしに可能となった。
• アブレーションスタディにより、潜在空間OT手法に比べ、環境空間でOTマップを直接使用することで、特にオートエンコーダーの再構成誤差が品質を低下させる高次元画像タスクにおいて、より優れた性能が得られることを確認した。