[论文解读] 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

一句话总结

该论文提出 trVAE，一种在瓶颈层引入最大均值差异（MMD）正则化的条件变分自编码器，以提升跨多个条件的未配对数据的分布外生成性能。通过强制不同条件间分布相似，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 的局限性，后者因缺乏显式的跨条件正则化而无法学习到条件间的紧凑联合分布。
• 开发一种数据驱动、端到端的框架，超越硬编码向量运算或直方图匹配，支持同时处理多种条件。
• 提升真实世界生物数据中分布外预测的性能，例如预测训练数据中未出现的药物处理引起的细胞反应。

提出的方法

• 在条件变分自编码器（CVAE）的瓶颈层之后的解码器层中引入 MMD 正则化，特别针对不同类别条件下的潜在表征。
• 使用基于核函数的 MMD 损失，匹配所有条件间潜在表征的分布，促使模型学习共享且紧凑的表征。
• 采用变分下界目标进行端到端训练，该目标包含重建损失和 KL 散度，同时将 MMD 作为潜在空间的正则项。
• 将模型应用于高维数据，包括图像和单细胞基因表达数据，条件如性别、发色或疾病状态。
• 在瓶颈层中引入多头注意力机制，以实现有效的跨条件表征学习，构成 'Transformer VAE'（trVAE）架构的基础。
• 使用皮尔逊相关系数等指标评估性能，比较预测值与真实值在高维数据均值/方差上的相关性，并进行基因表达谱的定性分析。

实验结果

研究问题

• RQ1 在 CVAE 瓶颈层中引入 MMD 正则化，是否能提升跨多个条件的未配对数据的分布外生成性能？
• RQ2 强制不同条件间分布相似，是否能提升高维样本生成的鲁棒性与准确性，特别是对少数或未见条件？
• RQ3 与 cycle GAN、SAUCIE 和普通 CVAE 等现有模型相比，trVAE 在预测训练期间未见的扰动引起的细胞反应方面表现如何？
• RQ4 与真实单细胞 RNA-seq 数据相比，MMD 正则化在多大程度上提升了预测基因表达均值与方差的保真度？
• RQ5 trVAE 是否能泛化至复杂生物任务，例如在未接触该条件的训练数据下，预测自然杀伤细胞中 IFN-β 刺激的影响？

主要发现

• trVAE 将 1,000 维度上预测与真实均值表达之间的皮尔逊相关系数从 0.89 提升至 0.97，表明对数据分布均值的估计显著更优。
• 预测方差的相关系数从 0.75 提升至 0.87，表明对不同条件下数据离散度的建模能力增强。
• 在单细胞基因表达预测中，trVAE 即使在训练数据中未包含该条件，也能准确捕捉自然杀伤细胞中 IFN-β 刺激后 ISG15 的上调。
• 与 cycle GAN 和 MMD 正则化自编码器（SAUCIE）相比，trVAE 在预测基因表达均值与方差方面表现更优，所有模型中 R² 值最高。
• 当每种扰动条件在训练时被单独排除时，模型仍能成功预测全部八种细胞类型，证明其对多种分布外组合的鲁棒性与泛化能力。
• 定性分析显示，trVAE 准确再现了刺激后表达变化最大的前 10 个差异表达基因，且变化方向与幅度均正确。