[論文レビュー] The Monge Gap: A Regularizer to Learn All Transport Maps

Théo Uscidda, Marco Cuturi|arXiv (Cornell University)|Feb 9, 2023

Adversarial Robustness in Machine Learning被引用数 4

ひとこと要約

本稿では、構造的制約なしに一般化されたニューラルネットワークを最適輸送（OT）写像として訓練可能にする、新しい正則化項「Mongeギャップ」を導入する。適合損失をMongeギャップで正則化することで、MongeギャップはOT最適性からの逸脱度を測定する。この手法は、合成データおよび単一細胞ゲノム解析ベンチマークにおいて、ICNN や MLP よりも優れた性能を示し、特に高次元において顕著である。

ABSTRACT

Optimal transport (OT) theory has been been used in machine learning to study and characterize maps that can push-forward efficiently a probability measure onto another. Recent works have drawn inspiration from Brenier's theorem, which states that when the ground cost is the squared-Euclidean distance, the ``best'' map to morph a continuous measure in $\mathcal{P}(\Rd)$ into another must be the gradient of a convex function. To exploit that result, [Makkuva+ 2020, Korotin+2020] consider maps $T= abla f_θ$, where $f_θ$ is an input convex neural network (ICNN), as defined by Amos+2017, and fit $θ$ with SGD using samples. Despite their mathematical elegance, fitting OT maps with ICNNs raises many challenges, due notably to the many constraints imposed on $θ$; the need to approximate the conjugate of $f_θ$; or the limitation that they only work for the squared-Euclidean cost. More generally, we question the relevance of using Brenier's result, which only applies to densities, to constrain the architecture of candidate maps fitted on samples. Motivated by these limitations, we propose a radically different approach to estimating OT maps: Given a cost $c$ and a reference measure $ρ$, we introduce a regularizer, the Monge gap $\mathcal{M}^c_ρ(T)$ of a map $T$. That gap quantifies how far a map $T$ deviates from the ideal properties we expect from a $c$-OT map. In practice, we drop all architecture requirements for $T$ and simply minimize a distance (e.g., the Sinkhorn divergence) between $T\sharpμ$ and $ν$, regularized by $\mathcal{M}^c_ρ(T)$. We study $\mathcal{M}^c_ρ$, and show how our simple pipeline outperforms significantly other baselines in practice.

研究の動機と目的

• 入力凸ニューラルネットワーク（ICNN）が最適輸送（OT）学習において直面する限界、すなわち訓練の不安定性、複雑な初期化、非二乗ユークリッド距離コストとの不適合性を解消すること。
• 連続密度ではなく標本に基づく測度を扱う際、Brenierの定理（凸ポテンシャルの勾配）を強制することは理論的に妥当かどうかを検証すること。
• 任意のコスト関数 $ c $ に対して、特別なネットワーク構造を必要とせず、一般化可能でアーキテクチャに依存しないOT写像学習手法を開発すること。
• 写像 $ T $ が最適OT写像にどれほど近いかを測定する正則化項、Mongeギャップ $ \mathcal{M}^{c}_{\rho}(T) $ を導入すること。この正則化項は、基準測度 $ \rho $ を用いて定義される。

提案手法

• Mongeギャップ $ \mathcal{M}^{c}_{\rho}(T) = \mathbb{E}_{X \sim \rho}[c(X, T(X))] - W_c^2(\rho, T\sharp\rho) $ を提案。これは写像 $ T $ が最適OT行動からどれほど逸脱しているかを測定する。
• Mongeギャップを正則化項として訓練目的関数に組み込む。具体的には、$ T\sharp\mu $ と $ \nu $ の間の適合損失（例：Sinkhorn発散）を最小化し、さらに $ \lambda_{\mathrm{MG}} \cdot \mathcal{M}^{c}_{\rho}(T) $ を加算する。
• コスト関数がねじれ条件（twist condition）を満たす場合、写像 $ T $ をポテンシャルの勾配としてパラメータライズすることを提案。これにより、最適写像の構造的性質を活用できる。
• 基準測度 $ \rho $ を導入。これは元の測度 $ \mu $ とは独立に選択可能であり、柔軟性とロバストネスを向上させる。
• 有限標本設定下でも微分可能かつ計算可能であるように、$ c $-ウォッサーシュタイン距離をSinkhorn推定器で扱う。
• 一般コスト関数および二乗ユークリッド距離ケースの両方に対して手法を適用。特に二乗ユークリッド距離の場合、Mongeギャップが凸であることが示され、最適化に有利である。

実験結果

リサーチクエスチョン

• RQ1 入力凸ニューラルネットワーク（ICNN）に制限を課さずに、最適輸送写像を学習可能か。これにより、訓練の不安定性や複雑な初期化を回避できるか。
• RQ2 連続密度ではなく標本に基づく測度を扱う際、Brenierの定理（凸ポテンシャルの勾配）を強制することは理論的・実用的に妥当か。
• RQ3 構造的制約なしにOT最適性からの逸脱度を測定する正則化項を導入した場合、ICNN や通常のMLP よりもOT写像学習で優れた性能を発揮できるか。
• RQ4 異なるコスト関数に対してMongeギャップ正則化項はどのように振る舞うか。高次元設定下でもハイパーパrameterの選択に頑健であるか。
• RQ5 本手法は、バクテリアの応答をモデル化する単一細胞オミクスデータなど、実世界の生物学的データに一般化可能か。特に、細胞応答への撹乱の影響をモデル化する際の有効性は。

主な発見

• 写像 $ T $ が最適OT写像であるための必要十分条件として、$ \mathrm{Spt}(\mu) \subset \mathrm{Spt}(\rho) $ かつ $ T\sharp\mu = \nu $ を満たす場合、Mongeギャップ $ \mathcal{M}^{c}_{\rho}(T) $ がゼロであることは、$ T $ が最適OT写像であることと同値である。理論的妥当性が裏付けられる。
• 二乗ユークリッド距離コストの場合、Mongeギャップは凸であり、この性質はSinkhornの有限標本推定下でも保たれる。これにより、安定した最適化が可能となる。
• Korotinら（2021）のベンチマークにおいて、本手法は通常のMLPおよびICNNを上回る性能を示し、特に $ d \geq 16 $ で顕著に優位である。$ d \geq 64 $ では性能差がさらに拡大する。
• 4iおよびscRNA-seqデータを用いた単一細胞ゲノム実験では、本手法は通常のMLPおよびICNNを上回る予測性能を達成。全34および9の処理条件下で一貫した改善が観察された。
• ハイパーパrameterの選択に対して本手法はロバストである。$ (\lambda_{\mathrm{MG}}, \lambda_{\mathrm{cons}}) $ の広い範囲で良好な性能が得られ、$ d \leq 64 $ では最適な重みが $ (1, 0.01) $、$ d \geq 128 $ では $ (10, 0.1) $ である。
• Gaussian初期化を施してもICNNは不安定で中程度の性能に留まる。OT学習におけるアーキテクチャ制約の実用的限界が浮き彫りになる。