[論文レビュー] Neural Optimal Transport

Alexander Korotin, Daniil Selikhanovych+1 more|arXiv (Cornell University)|Jan 28, 2022

Medical Image Segmentation Techniques被引用数 6

ひとこと要約

この論文は、強いおよび弱い輸送コストの下で、決定的および確率的最適輸送マップと計画を計算するためのディープラーニングベースの手法であるニューラル最適輸送（NOT）を紹介する。著者たちは、ニューラルネットワークが輸送計画を普遍的に近似可能であることを証明し、微分可能でスケーラブルなフレームワークを構築した。このフレームワークは、未対応画像間翻訳タスクにおいて先行手法を上回り、CelebA、outdoor-to-church、handbags-to-shoesのベンチマークで最先端の結果を達成した。

ABSTRACT

We present a novel neural-networks-based algorithm to compute optimal transport maps and plans for strong and weak transport costs. To justify the usage of neural networks, we prove that they are universal approximators of transport plans between probability distributions. We evaluate the performance of our optimal transport algorithm on toy examples and on the unpaired image-to-image translation.

研究の動機と目的

• 強いおよび弱い輸送コストの下で、最適輸送マップと計画を計算するためのスケーラブルで微分可能なディープラーニングフレームワークの開発。
• ディープニューラルネットワークが確率分布間の輸送計画を普遍的に近似可能であることを証明し、最適輸送におけるその使用を正当化すること。
• 既存手法が周辺制約を十分に満たさない、またはラングジュアンダイナミクスのような遅いサンプリング手順に依存するという限界を解決すること。
• 構造的および分布的性質を保持する輸送マップを学習することで、高品質な未対応画像間翻訳を実現すること。

提案手法

• 弱い最適輸送の双対定式に基づく新しい最適化フレームワークを提案し、それを最大最小のサドルポイント問題に再定式化する。
• 深層ニューラルネットワークでパラメータ化された確率的マップ $ T: \mathcal{X} \times \mathcal{Z} \to \mathcal{Y} $ を用いて条件付き輸送計画をモデル化する。
• 双対変数 $ f $ を介して周辺制約を強制する微分可能な目的関数を実装し、結合分布のプッシュフォワードがターゲットと一致することを保証する。
• エンドツーエンド学習を可能にするミニマックス最適化スキーム $ \max_f \min_T $ を用いて輸送マップを訓練する。バックプロパゲーションが可能である。
• $ \mathbb{W}_2 $（決定的）および $ \mathcal{W}_{2,\gamma} $（確率的）の両方のコストを適用し、1対1および1対多のマッピングをサポートする。
• 弱いコスト関数 $ C(x, \mu) $ のための変分近似を用い、スコアベースモデルや反復的サンプリングに依存せずに、効率的な最適化を可能にする。

実験結果

リサーチクエスチョン

• RQ1 ディープニューラルネットワークは、確率分布間の最適輸送計画を普遍的に近似可能か？
• RQ2 微分可能でエンドツーエンドのディープラーニングフレームワークは、強いおよび弱いコストの下で、決定的および確率的最適輸送マップを効果的に計算できるか？
• RQ3 提案されたニューラル最適輸送手法は、未対応画像間翻訳タスクにおいて、既存の手法を上回る性能を示すか？
• RQ4 この手法は、CelebA、outdoor-to-church、handbags-to-shoesといった多様なデータセットに良好に一般化するか？
• RQ5 確率的輸送マップは、多様で現実的で、構造的一致性を保った翻訳を生成できるか？

主な発見

• 提案されたニューラル最適輸送手法は、CelebA（女性からアニメ）や、outdoor to church、handbags to shoesといった複数のベンチマークで、未対応画像間翻訳において最先端の性能を達成した。
• この手法は、決定的（1対1）および確率的（1対多）の両方の輸送マップを効果的に計算でき、多様で高精細な画像翻訳を実現した。
• アブレーションスタディの結果、同期化されたノイズベクトルを用いることで、異なる入力間で共通の視覚的アーチファクトが生じないことが示され、確率的マップが入力ごとに多様な出力を生成していることが裏付けられた。
• このフレームワークは大規模なビジョンタスクにも良好に一般化でき、$ 128 \times 128 $ および $ 64 \times 64 $ 解像度で高品質な翻訳を生成した。
• 理論的証明により、ディープニューラルネットワークが輸送計画を普遍的に近似可能であることが示され、ニューラルネットワークを最適輸送に応用する強固な基盤が得られた。
• この手法は、ラングジュアンダイナミクスのような計算コストの高いサンプリング手順を回避し、スコアベースの代替手法と比較して、より高速な訓練および推論を可能にした。