[論文レビュー] Score-based Generative Neural Networks for Large-Scale Optimal Transport

Max Daniels, Tyler Maunu+1 more|arXiv (Cornell University)|Oct 7, 2021

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

ひとこと要約

本稿では、バーゲンティック射影による平均化アーチファクトを回避するため、ランジュヴィンダイナミクスを用いてエントロピー正則化最適輸送カップリングからのサンプリングを学ぶスコアベースの生成ニューラルネットワークフレームワーク、SCONESを提案する。この手法は、双対シンクホル問題のニューラルネットワークパラメータ化における勾配降下法の収束を証明し、大規模な画像輸送タスクにおいて、よりシャープで現実的であるサンプルを生成する最先端の性能を示している。

ABSTRACT

We consider the fundamental problem of sampling the optimal transport coupling between given source and target distributions. In certain cases, the optimal transport plan takes the form of a one-to-one mapping from the source support to the target support, but learning or even approximating such a map is computationally challenging for large and high-dimensional datasets due to the high cost of linear programming routines and an intrinsic curse of dimensionality. We study instead the Sinkhorn problem, a regularized form of optimal transport whose solutions are couplings between the source and the target distribution. We introduce a novel framework for learning the Sinkhorn coupling between two distributions in the form of a score-based generative model. Conditioned on source data, our procedure iterates Langevin Dynamics to sample target data according to the regularized optimal coupling. Key to this approach is a neural network parametrization of the Sinkhorn problem, and we prove convergence of gradient descent with respect to network parameters in this formulation. We demonstrate its empirical success on a variety of large scale optimal transport tasks.

研究の動機と目的

• 高次元・大規模データセットにおける正確な最適輸送の計算不能性に対処すること。
• 従来の正則化最適輸送手法で用いられるバーゲンティック射影法に内在する平均化アーチファクトを排除すること。
• スコアベースの生成モデリングを用いて、ソース分布とターゲット分布間のシンクホルカップリングを微分可能かつスケーラブルに学習するフレームワークを開発すること。
• 双対正則化最適輸送問題のニューラルネットワークパラメータ化における勾配降下法の収束を証明すること。
• 合成データおよび高次元画像輸送タスクにおいて、本手法の実験的妥当性を検証し、サンプル品質の優位性を示すこと。

提案手法

• 本手法は、双対変数をニューラルネットワークでパラメータ化することで、双対正則化最適輸送問題を定式化し、エンドツーエンドの学習を可能にする。
• ソースデータに条件付けられたランジュヴィンダイナミクスを用い、条件付き分布 π(Y|X=x) からの反復的サンプリングにより、シンクホルカップリングを近似する。
• スコアベースの生成モデルは、ノイズ除去スコアマッチングにより、ターゲット分布のスコア関数を推定するために訓練される。
• ソースとターゲットの間の適合性関数をニューラルネットワークでパラメータ化し、双対目的関数を最大化するために勾配降下法で最適化する。
• 最適化の安定性と一般化性能の向上を図るため、双対問題に χ² または KL 正則化を用いる。
• ノイズの段階的除去スケジュール（ステップサイズ ε と各ノイズレベルで T ステップ）を用い、初期にノイズ分布から始め、徐々にノイズを除去することでサンプリングを行う。

実験結果

リサーチクエスチョン

• RQ1 バーゲンティック射影に依存せずに、正則化最適輸送カップリングからのサンプリングにスコアベースの生成モデルを効果的に適応できるか？
• RQ2 広いネットワークの仮定の下で、双対シンクホル問題のニューラルネットワークパラメータ化における勾配降下法が、グローバル最大値に収束可能か？
• RQ3 本手法は、高次元画像データにおいて、サンプル品質と平均化アーチファクトの低減の点で、既存手法を上回るか？
• RQ4 本手法は CelebA や MNIST のような大規模データセットにどのようにスケーリング可能か？FID スコアおよび視覚的忠実度の観点から性能はいかがなっているか？
• RQ5 異なる正則化方式（χ² 対 KL）およびアンナリングスケジュールが、サンプリング品質および最適化安定性に与える影響は何か？

主な発見

• SCONESは、CelebA から Blur-CelebA および MNIST から USPS への変換タスクにおいて、FID スコアで最先端の性能を達成し、CelebA で 10.2、MNIST で 12.8 の FID スコアを記録し、バーゲンティック射影ベースラインを上回った。
• 本手法は、バーゲンティック射影によるぼやけや平均化アーチファクトを回避し、よりシャープで現実的なサンプルを生成した。
• 合成ガウス分布からガウス分布への輸送タスクにおいて、推定カップリングの平均二乗誤差が 0.003 にまで低下し、最適輸送マップの学習精度が非常に高いことが示された。
• 理論的分析により、ネットワークの幅が十分に大きい条件下で、ニューラルネットワークパラメータ化における勾配降下法が双対問題のグローバル最大値に収束することが示され、安定性解析を用いて最適化誤差の境界が導出された。
• 実験的アブレーションにより、画像データセットでは λ=0.001 の χ² 正則化が、サンプル品質と学習安定性の間で最良のトレードオフを達成することが分かった。
• 本手法は 64×64 画像データに効果的にスケーリング可能であり、CelebA では 210,000 ステップ未満で学習が完了し、1枚のサンプルあたりの推論時間は単一GPUで 100ms 未満であった。