[论文解读] Modelling Cellular Perturbations with the Sparse Additive Mechanism Shift Variational Autoencoder

Michael Bereket, Theofanis Karaletsos|arXiv (Cornell University)|Nov 5, 2023

Single-cell and spatial transcriptomics被引用 8

一句话总结

SAMS-VAE为VAE引入稀疏加性扰动机制，以对细胞扰动进行建模，实现解耦、可解释的潜在子空间，并在扰动测序数据上具有强泛化能力。它也引入相关推断和基于平均处理效应的评估框架。

ABSTRACT

Generative models of observations under interventions have been a vibrant topic of interest across machine learning and the sciences in recent years. For example, in drug discovery, there is a need to model the effects of diverse interventions on cells in order to characterize unknown biological mechanisms of action. We propose the Sparse Additive Mechanism Shift Variational Autoencoder, SAMS-VAE, to combine compositionality, disentanglement, and interpretability for perturbation models. SAMS-VAE models the latent state of a perturbed sample as the sum of a local latent variable capturing sample-specific variation and sparse global variables of latent intervention effects. Crucially, SAMS-VAE sparsifies these global latent variables for individual perturbations to identify disentangled, perturbation-specific latent subspaces that are flexibly composable. We evaluate SAMS-VAE both quantitatively and qualitatively on a range of tasks using two popular single cell sequencing datasets. In order to measure perturbation-specific model-properties, we also introduce a framework for evaluation of perturbation models based on average treatment effects with links to posterior predictive checks. SAMS-VAE outperforms comparable models in terms of generalization across in-distribution and out-of-distribution tasks, including a combinatorial reasoning task under resource paucity, and yields interpretable latent structures which correlate strongly to known biological mechanisms. Our results suggest SAMS-VAE is an interesting addition to the modeling toolkit for machine learning-driven scientific discovery.

研究动机与目标

• 在扰动下学习一个生成模型，以将扰动效应与基础变异解耦。
• 引入稀疏扰动潜在偏移，按照相同扰动的样本呈加性组合并共享。
• 开发并比较相关的变分推断策略，以提升潜在解耦和预测性能。
• 提供一个基于边际似然（IWELBO）和通过平均处理效应进行后验预测检查的扰动模型评估框架。
• 在扰动-seq数据集上相较基线，展示更好的泛化和可解释性。

提出的方法

• 定义 z_i = z_i^b + z_i^p，潜在状态为 z_i^p = sum_t d_{i,t} (e_t ⊙ m_t)。
• 建模 e_t ~ N(0, βI) 和 m_t ~ Bern(α)以引入稀疏、扰动特异的偏移。
• 使用神经网络参数化 p(x_i|z_i; θ)，并通过 Gamma-Poisson（负二项）模型对库大小 l_i 的 scRNA-seq似然进行建模。
• 使用随机变分推断进行推断，采用平均场或相关变分族，将 z^b、E、M在接收相同扰动 t 的样本之间绑定。
• 引入 CPA-VAE 作为一个消融变体，移除稀疏掩码（m_t 固定为1）。
• 提供两种改进的推断方案：相关的 z_basal 与相关的嵌入 E，以获得更丰富的变分族。

Figure 2 : Visualization of inferred latent perturbation masks and embedding means for the best performing checkpoint of each model in replogle-filtered . We visualize the latent variables for the 345 perturbations with pathway annotations from Replogle et al. [ 17 ] and group by pathway. The SAMS-V

实验结果

研究问题

• RQ1 SAMS-VAE是否能够将扰动效应准确建模为稀疏加性的潜在偏移？
• RQ2 相关推断策略是否相对于基线提升潜在解耦、可解释性和预测性能？
• RQ3 SAMS-VAE在同分布和分布外扰动以及组合扰动下的泛化能力如何？
• RQ4 通过平均处理效应和差异表达的评估框架对扰动模型是否具有信息性？
• RQ5 SAMS-VAE与CPA-VAE、SVAE+以及条件VAE在扰动-seq数据集上的比较如何？

主要发现

• 在完全相关推断下，SAMS-VAE在重筛选数据上的测试IWELBO和ATE相关性表现最佳。
• 相关的 z_basal 推断在SAMS-VAE和CPA-VAE中带来显著提升；相关的 E 带来适度提升。
• SAMS-VAE在同时具有相关的 z_basal 与 E 时，达到最高的 Mask PW. Acc 和 ATE-Pearson（在测试的配置中）。
• SAMS-VAE和CPA-VAE在组合扰动设置（norman-ood）和数据高效场景中表现出强泛化能力。
• 来自SAMS-VAE的潜在扰动掩码对注释生物通路的预测性优于SVAE+。
• 基于模型的ATE与数据驱动的差异表达高度相关，支持PPC框架。