[論文レビュー] Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks

Jiaojiao Fan, Amirhossein Taghvaei+1 more|arXiv (Cornell University)|Jul 8, 2020

Topological and Geometric Data Analysis参考文献 45被引用数 18

ひとこと要約

本稿では、カントロビッチ双対定式化における凸ポテンシャルをパrameter化するための入力凸ニューラルネットワーク（ICNN）を用いて、高次元連続空間における Wasserstein バリオセントラムのスケーラブルな計算のためのアルゴリズムを提案する。バリオセントラムを生成モデルとしてモデル化することで、周辺分布へのアクセスを必要とせず、無限にサンプリングを効率的に行える。高次元ベンチマークにおいて、精度とスケーラビリティの点で先行手法を上回る性能を発揮する。

ABSTRACT

Wasserstein Barycenter is a principled approach to represent the weighted mean of a given set of probability distributions, utilizing the geometry induced by optimal transport. In this work, we present a novel scalable algorithm to approximate the Wasserstein Barycenters aiming at high-dimensional applications in machine learning. Our proposed algorithm is based on the Kantorovich dual formulation of the Wasserstein-2 distance as well as a recent neural network architecture, input convex neural network, that is known to parametrize convex functions. The distinguishing features of our method are: i) it only requires samples from the marginal distributions; ii) unlike the existing approaches, it represents the Barycenter with a generative model and can thus generate infinite samples from the barycenter without querying the marginal distributions; iii) it works similar to Generative Adversarial Model in one marginal case. We demonstrate the efficacy of our algorithm by comparing it with the state-of-art methods in multiple experiments.

研究の動機と目的

• 従来の線形計画法やエントロピー正則化手法が失敗する高次元連続空間における Wasserstein バリオセントラム推定の計算ボトルネックを解消すること。
• 周辺分布のクエリを必要とせず、無限にサンプリング可能な生成モデルとしてバリオセントラムを表現する手法を開発すること。
• カントロビッチ定式化における有効な双対ポテンシャルを保証するため、入力凸ニューラルネットワーク（ICNN）の凸性保存特性を活用すること。
• 先行手法で用いられる正則化項に起因するバイアスを、ICNNアーキテクチャによる凸性の直接的制約によって排除すること。
• 高次元の合成データおよび実世界データセットにおいて、最先端手法と比較してスケーラビリティと精度の両面で向上を示すこと。

提案手法

• Wasserstein-2 距離のカントロビッチ双対定式化を用い、バリオセントラム問題を凸関数の最適化問題に再定式化する。
• 双対ポテンシャルをパラメータ化するために入力凸ニューラルネットワーク（ICNN）を採用し、構造的に凸関数を保証する。
• バリオセントラムを、ノイズを入力としてバリオセントラム分布からのサンプルを出力する生成器ネットワークとして表現し、効率的かつ無制限のサンプリングを可能にする。
• 最適化は、生成器に関して最小化、双対ポテンシャルに関して最大化、双対関数に関して最小化という三段階のミニマックス最小化問題として構造化される。
• 入力分布の明示的離散化を必要とせず、確率的最適化を用いてエンドツーエンドで訓練する。
• 1つの周辺分布の場合、本手法は GAN と同様の振る舞いを示し、生成器がバリオセントラム分布に適合するように訓練される。

実験結果

リサーチクエスチョン

• RQ1 ICNN を用いた深層生成モデルは、高次元連続空間における Wasserstein バリオセントラムを効果的に表現できるか？
• RQ2 ICNN を用いた凸性の強制は、先行のニューラルネットワークベース手法で見られる正則化に起因するバイアスを解消できるか？
• RQ3 本手法は、高次元設定において最先端のアルゴリズムと比較して、性能と精度の点でどのようにスケーリングするか？
• RQ4 単一の訓練プロセスを通じて、任意の重みを持つバリオセントラムを学習可能に拡張できるか？
• RQ5 特に二次関数の場合に、本定式化の最適化の局所的性質は、既存手法と比較してどのように異なるか？

主な発見

• 本手法は、高次元の合成および実世界ベンチマークにおいて、Korotin 他 (2021b) や Li 他 (2020) といった最先端手法を大きく上回る顕著な推定精度を達成した。
• 二次関数の場合、本手法は滑らかで凹-凸-凹な最適化問題に簡略化され、Korotin 他 (2021b) の非滑らかで非凸な定式化とは異なり、安定した収束を保証する。
• 最適な生成器は、等しい共分散を持つガウス周辺分布に対して、正確にバリオセントラムを学習し、平均が ∑a_i m_i に一致することが確認された。理論的整合性を裏付ける。
• 事前分布選択や、先行研究で用いられる類似性正則化やサイクル正則化といった複雑な正則化項の必要性を回避した。
• 実験では、本手法が高次元設定において効果的にスケーリングされ、エントロピー正則化法やフリー・サポート法を、速度と精度の両面で上回った。
• 本フレームワークは、1つの周辺分布の場合に GAN として適応可能であり、その柔軟性と生成能力を示した。