[論文レビュー] Conditional out-of-sample generation for unpaired data using trVAE

Mohammad Lotfollahi, Mohsen Naghipourfar+2 more|arXiv (Cornell University)|Oct 4, 2019

Single-cell and spatial transcriptomics参考文献 26被引用数 12

ひとこと要約

本稿では、複数の条件下でペアのないデータに対する分布外生成を向上させるために、バトルネック層に最大平均差分（MMD）正則化を施した条件付き変分オートエンコーダー（trVAE）を提案する。異なる条件下での分布の類似性を強制することにより、trVAEは高次元サンプルの生成において顕著に高い正確性と頑健性を達成する。これは、皮積相関（平均：0.97、分散：0.87）の向上と、マイノリティな細胞タイプや新しい摂動に対する単細胞遺伝子発現予測タスクにおける優れた性能によって実証されている。

ABSTRACT

While generative models have shown great success in generating high-dimensional samples conditional on low-dimensional descriptors (learning e.g. stroke thickness in MNIST, hair color in CelebA, or speaker identity in Wavenet), their generation out-of-sample poses fundamental problems. The conditional variational autoencoder (CVAE) as a simple conditional generative model does not explicitly relate conditions during training and, hence, has no incentive of learning a compact joint distribution across conditions. We overcome this limitation by matching their distributions using maximum mean discrepancy (MMD) in the decoder layer that follows the bottleneck. This introduces a strong regularization both for reconstructing samples within the same condition and for transforming samples across conditions, resulting in much improved generalization. We refer to the architecture as \emph{transformer} VAE (trVAE). Benchmarking trVAE on high-dimensional image and tabular data, we demonstrate higher robustness and higher accuracy than existing approaches. In particular, we show qualitatively improved predictions for cellular perturbation response to treatment and disease based on high-dimensional single-cell gene expression data, by tackling previously problematic minority classes and multiple conditions. For generic tasks, we improve Pearson correlations of high-dimensional estimated means and variances with their ground truths from 0.89 to 0.97 and 0.75 to 0.87, respectively.

研究の動機と目的

• トレーニングデータにそのようなペアの例がないにもかかわらず、ドメインと条件の未観測の組み合わせに条件づけた高次元サンプルを生成する課題に対処すること。
• ペアのないデータにおいて、明示的な条件間正則化が欠如しているため、通常のCVAEが条件間のコン pact な共同分布を学習できないという限界を克服すること。
• ハードコードされたベクトル演算やヒストグラムマッチングを超える、データ駆動型でエンドツーエンドのフレームワークを構築し、同時に複数の条件を扱えるようにすること。
• トレーニングデータに存在しない薬剤処理などの摂動に対する細胞応答を予測する、現実の生物学的データにおける分布外予測を改善すること。

提案手法

• 条件付き変分オートエンコーダー（CVAE）のバトルネック以降のデコーダー層に、カテゴリカルな条件ごとの潜在表現間のMMD正則化を導入する。
• 核ベースのMMD損失を用いて、すべての条件下での潜在表現の分布を一致させ、モデルが共有でコンパクトな表現を学習するよう促進する。
• 再構成損失とKLダイバージェンスに加え、潜在空間にMMDを正則化子として追加した、変分下界の目的関数を用いてモデルをエンドツーエンドで訓練する。
• 性別、髪の色、疾患状態などの条件を伴い、画像や単細胞遺伝子発現データを含む高次元データにモデルを適用する。
• バトルネック層にマルチヘッドアテンション機構を導入し、効果的な条件間表現学習を可能にし、『トランスフォーマーVAE』（trVAE）アーキテクチャの基盤を構築する。
• 予測された高次元データの平均/分散と真値との皮積相関などの指標、および遺伝子発現プロファイルにおける定性的分析を用いて性能を評価する。

実験結果

リサーチクエスチョン

• RQ1 CVAEのバトルネック部におけるMMD正則化は、複数の条件下でペアのないデータに対する分布外生成を改善できるか？
• RQ2 条件間での分布類似性を強制することで、とりわけマイノリティまたは未観測の条件下でも、高次元サンプルの生成がより頑健かつ正確になるか？
• RQ3 トレーニング中に観測されていない摂動に対する細胞応答を予測する際、trVAEはcycle GAN や SAUCIE や通常のCVAEと比較して優れているか？
• RQ4 単細胞RNA-seqデータにおいて、MMD正則化は真値との比較で予測された遺伝子発現の平均と分散の忠実性をどの程度向上させるか？
• RQ5 トレーニング中にその条件にさらされていなくても、trVAEは自然キラー細胞におけるIFN-β刺激効果の予測といった、複雑な生物学的タスクに一般化可能か？

主な発見

• trVAEは、1,000次元の予測された平均発現と真値との皮積相関を0.89から0.97に向上させ、データ分布の平均推定が著しく改善されたことを示している。
• 予測分散の相関は0.75から0.87に向上し、条件間でのデータ分散のモデリングが向上したことが示された。
• 単細胞遺伝子発現予測において、trVAEはトレーニングデータにその条件が存在しなかったにもかかわらず、自然キラー細胞におけるIFN-β刺激後のISG15の上昇発現を正確に捉えている。
• trVAEは、平均および分散の両方において、cycle GAN やMMD正則化付きオートエンコーダー（SAUCIE）を上回り、全モデルの中で最高のR²値を示した。
• 各条件をトレーニング時にホールドアウトした場合でも、全8つの細胞タイプをすべての摂動条件下で正しく予測できており、複数の分布外組み合わせに対する頑健性と一般化能力が裏付けられた。
• 定性的な分析から、trVAEは刺激後の上位10個の発現変化が著しい遺伝子を正確に再現しており、変化の方向性と大きさの両方が正しいことが確認された。