[論文レビュー] Supervised Training of Conditional Monge Maps

Charlotte Bunne, Andreas Krause+1 more|arXiv (Cornell University)|Jun 28, 2022

Adversarial Robustness in Machine Learning被引用数 14

ひとこと要約

本稿では、部分入力凸ニューラルネットワーク（PICNN）を用いて条件付きモンジュ写像を学習する教師あり学習フレームワーク、CondOTを紹介する。この手法は、確率測度間の文脈依存最適輸送をモデル化する。ガウス近似による文脈に配慮した初期化と、ディープセット組み合わせ論を用いた未観測の摂動への一般化により、新しい薬剤コンbinationsや遺伝子的摂動に対する単細胞応答の正確な予測が可能になる。

ABSTRACT

Optimal transport (OT) theory describes general principles to define and select, among many possible choices, the most efficient way to map a probability measure onto another. That theory has been mostly used to estimate, given a pair of source and target probability measures $(μ, ν)$, a parameterized map $T_θ$ that can efficiently map $μ$ onto $ν$. In many applications, such as predicting cell responses to treatments, pairs of input/output data measures $(μ, ν)$ that define optimal transport problems do not arise in isolation but are associated with a context $c$, as for instance a treatment when comparing populations of untreated and treated cells. To account for that context in OT estimation, we introduce CondOT, a multi-task approach to estimate a family of OT maps conditioned on a context variable, using several pairs of measures $\left(μ_i, ν_i ight)$ tagged with a context label $c_i$. CondOT learns a global map $\mathcal{T}_θ$ conditioned on context that is not only expected to fit all labeled pairs in the dataset $\left\{\left(c_i,\left(μ_i, ν_i ight) ight) ight\}$, i.e., $\mathcal{T}_θ\left(c_i ight) \sharp μ_i \approx ν_i$, but should also generalize to produce meaningful maps $\mathcal{T}_θ\left(c_{ ext {new }} ight)$ when conditioned on unseen contexts $c_{ ext {new }}$. Our approach harnesses and provides a novel usage for partially input convex neural networks, for which we introduce a robust and efficient initialization strategy inspired by Gaussian approximations. We demonstrate the ability of CondOT to infer the effect of an arbitrary combination of genetic or therapeutic perturbations on single cells, using only observations of the effects of said perturbations separately.

研究の動機と目的

• 個々の摂動の効果しか観測されていない状況で、新しい薬剤や遺伝子コンbinationsの単細胞への影響を予測する課題に対処すること。
• 治療法、投与量、遺伝子的摂動などの文脈変数を条件として用いることで、固定されたソース・ターゲットペアにとどまらない最適輸送を一般化すること。
• 文脈ベクトルに条件付けられたグローバルでパラメータ化された輸送写像を学習することで、未観測の文脈への一般化を可能にすること。
• ガウス近似からの閉形式解を活用して、入力凸ニューラルネットワークにおける訓練の不安定性を克服する、堅牢な閉形式初期化戦略の導入。

提案手法

• 入力空間 $\mu$ における凸性を保証するため、条件付き輸送写像 $\mathcal{T}_{\theta}(c)$ を部分入力凸ニューラルネットワーク（PICNN）の勾配としてパラメータ化する。
• PICNNの重みとバイアスを文脈ベクトル $c$ で変調することで、文脈に特化した輸送写像を生成するハイパーネットワークを構築する。
• 各 $\mu_i$ と $\nu_i$ のガウス近似間の最適輸送に対する閉形式解を用いてPICNNのパラメータを初期化し、安定的で意味のある初期写像を保証する。
• 訓練の際には、重みクリッピング（負の重みをゼロに設定）によりPICNNにおける凸性を強制し、双対ネットワークにおける凸性を負の重みに対する微分可能ペナルティ項により緩和する。
• 複数の摂動埋め込みを統合して共同文脈表現を生成するため、ディープセットアーキテクチャを採用し、未観測の摂動コンbinationsへの一般化を可能にする。
• PCAまたはオートエンコーダーを用いて高次元の単細胞RNA-seqデータを低次元の潜在空間に投影し、訓練効率と一般化性能を向上させる。

実験結果

リサーチクエスチョン

• RQ1 個々の摂動の観測しか得られていない状況で、未観測の摂動コンビネーションに一般化できる1つのパラメータ化された輸送写像を学習できるか？
• RQ2 ガウス近似からの閉形式解を活用することで、最適輸送のための入力凸ニューラルネットワークの訓練をどの程度安定化できるか？
• RQ3 ワンホットエンコーディングと比較して、ディープセットベースの組み合わせモジュールは未観測の摂動コンビネーションへの一般化にどの程度効果を発揮するか？
• RQ4 提案された条件付き最適輸送フレームワークは、新しい薬剤や遺伝子的摂動に対する単細胞応答を予測する際、既存のベースラインを上回る性能を示すか？

主な発見

• CondOTは、ディープセットモジュールを用いて個々の摂動の統合埋め込みを学習することで、未観測の摂動コンビネーションへの一般化に成功し、ゼロショット一般化においてワンホットエンコーディングを上回る性能を示した。
• ガウス近似に基づく提案された初期化戦略は、PICNNに基づく輸送写像の訓練安定性と収束速度を顕著に向上させた。
• MOFA+ および MOFA++ データセットを用いた検証により、CondOTは新しい薬剤コンビネーションに対する単細胞応答予測において、最先端の性能を達成した。
• オートエンコーダーを用いた投影はPCAよりも高次元構造をよりよく保持しており、下流の予測タスクにおける性能向上に寄与した。
• 訓練中に観測されなかった摂動に対しても、モデルは効果的に一般化できており、ゼロショット設定における堅牢性を示した。
• 双対ネットワークにおける負の重みに対する微分可能ペナルティの導入により、凸性制約の有効な緩和が可能になりながらも理論的保証を維持できた。