[論文レビュー] Scalable Unbalanced Optimal Transport using Generative Adversarial Networks

Karren Yang, Caroline Uhler|arXiv (Cornell University)|Oct 26, 2018

Model Reduction and Neural Networks被引用数 15

ひとこと要約

本稿では、出発測度と到達測度の間の質量変動をモデル化するためのスケーラブルな生成的敵対的ネットワーク（GAN）ベースのフレームワークを提案する。このフレームワークは、同時に輸送写像とスケーリング要因を学習することで、質量変動を扱う。非平衡最適輸送（OT）をモングェ型問題として定式化し、凸共役双対性を用いて緩和することで、GANに類似した効率的で確率的かつ交互勾配更新が可能となり、MNIST、USPS、CelebA、およびゼブラフィッシュの単細胞RNA-Seqデータにおける集団モデリングで最先端の性能を達成した。

ABSTRACT

Generative adversarial networks (GANs) are an expressive class of neural generative models with tremendous success in modeling high-dimensional continuous measures. In this paper, we present a scalable method for unbalanced optimal transport (OT) based on the generative-adversarial framework. We formulate unbalanced OT as a problem of simultaneously learning a transport map and a scaling factor that push a source measure to a target measure in a cost-optimal manner. In addition, we propose an algorithm for solving this problem based on stochastic alternating gradient updates, similar in practice to GANs. We also provide theoretical justification for this formulation, showing that it is closely related to an existing static formulation by Liero et al. (2018), and perform numerical experiments demonstrating how this methodology can be applied to population modeling.

研究の動機と目的

• 古典的最適輸送が出発測度と到達測度の間の質量変動を扱えないという限界を解消すること。
• 輸送と質量変動の両方をサポートするスケーラブルで連続的かつ微分可能な非平衡OTフレームワークを構築すること。
• 双対変数と発散正則化を用いた緩和されたモングェ問題として非平衡OTを定式化し、GANに類する学習ダイナミクスを非平衡OTに拡張すること。
• 生物学的集団モデリングや画像生成などの高次元で連続的な設定における非平衡OTの実用的応用を可能にすること。
• 連続測度に対する既存のエントロピー輸送ソルバーの代替手段として、学習安定性と収束性を向上させたスケーラブルな代替手法を提供すること。

提案手法

• 非平衡OTを確率的輸送写像 $ T $ とスケーリング要因 $ \xi $ を持つモングェ型問題として定式化し、輸送コストと発散ペナルティを含む目的関数を最小化する。
• $ \psi $-発散の凸共役双対性を用いて問題を緩和し、確率的勾配降下に適した双対最適化問題に変換する。
• GANにインspiredされた交互勾配降下法を採用し、生成器が $ T $ と $ \xi $ を学習し、識別器が双対変数 $ f $ を学習する。
• $ T $ と $ \xi $ にReLUおよびシグモイド／ソフトプラス出力層を備えたニューラルネットワークを用い、訓練の安定化のため勾配ペナルティを適用する。
• 訓練中にリスケーリングトリックを適用：$ \xi $ が小さい場合に $ T $ の損失を $ 1/\xi $ でスケーリングすることで収束速度を向上させる。
• Lieroら（2018）による最適エントロピー輸送問題を連続設定で解くためのスケーラブルなアルゴリズム（付録のアルゴリズム2）を導入し、従来の手法を拡張する。

実験結果

リサーチクエスチョン

• RQ1 質量変動の明示的モデル化を伴うGANに類するフレームワークは、非平衡最適輸送を解くために適応可能か？
• RQ2 非平衡OTのモングェ定式化は、確率的でディープラーニングベースの最適化に適した双対最適化問題にどのように緩和可能か？
• RQ3 連続測度を伴う非平衡OTにおいて、安定性と収束性を向上させる訓練ダイナミクスおよび正則化戦略（例：勾配ペナルティ）は何か？
• RQ4 提案手法は、画像や単細胞オミクスデータなどの高次元で連続的なデータセットにスケーラブルか？
• RQ5 輸送要因とスケーリング要因の共同学習は、既存のエントロピー輸送ソルバーと比較して、精度と効率の点で優れているか？

主な発見

• 提案手法は、MNIST、USPS、CelebAデータセットを用いた集団変換タスクにおいて、質量変動と輸送を効果的にモデル化した。
• ゼブラフィッシュの単細胞RNA-Seqデータでは、発生的遷移を現実的な細胞集団ダイナミクスとともに正確に捉えた。
• $ \xi $ および $ f $ に対する勾配ペナルティの適用により、特に初期訓練段階での訓練安定性が顕著に向上した。
• $ T $ の損失に対するリスケーリングトリックにより、$ \xi $ の値が小さい場合の収束速度が向上し、一致が悪いサンプルに対する更新が遅くなるのを防いだ。
• 連続的かつ高次元的測度に対して、既存のエントロピー輸送ソルバーと比較して競争力のある性能を達成し、スケーラビリティが向上した。
• 提案されたアルゴリズム2は、Chizatら（2016）の手法に対するスケーラブルな代替手段を提供し、エントロピー輸送手法の適用範囲を拡張した。