[论文解读] The VampPrior Mixture Model

Andrew Stirn, David A. Knowles|arXiv (Cornell University)|Feb 6, 2024

Single-cell and spatial transcriptomicsBiochemistry, Genetics and Molecular Biology被引用 3

一句话总结

该论文提出了VampPrior混合模型（VMM），一种用于深度潜在变量模型的新先验，用Dirichlet过程高斯混合模型结合VampPrior正则化替代了标准的$\mathcal{N}(0,I)$先验，实现了自动聚类和改进的批次校正。当集成到scVI中时，VMM在scRNA-seq数据集上显著提升了生物学保真度和整合性能，优于现有方法在聚类和批次校正方面的表现。

ABSTRACT

Widely used deep latent variable models (DLVMs), in particular Variational Autoencoders (VAEs), employ overly simplistic priors on the latent space. To achieve strong clustering performance, existing methods that replace the standard normal prior with a Gaussian mixture model (GMM) require defining the number of clusters to be close to the number of expected ground truth classes a-priori and are susceptible to poor initializations. We leverage VampPrior concepts (Tomczak and Welling, 2018) to fit a Bayesian GMM prior, resulting in the VampPrior Mixture Model (VMM), a novel prior for DLVMs. In a VAE, the VMM attains highly competitive clustering performance on benchmark datasets. Integrating the VMM into scVI (Lopez et al., 2018), a popular scRNA-seq integration method, significantly improves its performance and automatically arranges cells into clusters with similar biological characteristics.

研究动机与目标

• 为解决深度潜在变量模型（DLVMs）中固定聚类先验的局限性，这些先验需要预先指定聚类数量且对初始化敏感。
• 开发一种先验，实现在单细胞RNA测序（scRNA-seq）数据分析中同时进行整合与聚类，优于现有使用$\mathcal{N}(0,I)$先验的方法（如scVI）。
• 创建一种可扩展、即插即用的先验，可轻松集成到任何具有连续潜在变量的DLVM中，包括scVI及相关多组学工具。
• 证明VMM在图谱级scRNA-seq整合任务中可同时提升生物学保真度和技术批次校正性能。

提出的方法

• VMM通过在聚类均值上施加VampPrior，将高斯混合模型（GMM）重构为狄利克雷过程混合模型，实现无需预先定义聚类数量的非参数聚类。
• 该方法在变分推断（固定先验参数下更新变分参数）和经验贝叶斯步骤（更新先验参数）之间交替进行，实现变分学习与先验学习的解耦。
• VampPrior正则化确保先验成分从数据驱动的、结构化的先验点中学习，而非随机初始化，从而提升稳定性和收敛性。
• VMM被整合到变分自编码器（VAE）的生成过程中，并应用于图像聚类和scRNA-seq数据，替代标准的$\mathcal{N}(0,I)$先验。
• 在scRNA-seq中，VMM被嵌入scVI，使模型能够在同时学习共享的、批次校正后的嵌入表示的同时，发现具有生物学意义的聚类。
• 该推断流程设计为与现有DLVM框架兼容，可轻松部署于scVI、scATAC-seq和空间转录组学等工具的流程中。

实验结果

研究问题

• RQ1 深度潜在变量模型能否在不预先定义聚类数量的情况下实现稳健的自动聚类？
• RQ2 灵活的先验是否能同时提升图谱级scRNA-seq整合中的批次校正和生物学保真度？
• RQ3 在scVI中用VampPrior混合模型替代$\mathcal{N}(0,I)$先验是否能带来更优的聚类和整合性能？
• RQ4 与联合优化相比，变分参数与先验参数之间的交替推断过程如何提升模型稳定性和性能？

主要发现

• VMM优于所有基于VAE的聚类方法，在图像聚类基准测试中达到最先进性能，接近监督方法的表现。
• 在肺图谱数据集中，VMM在生物学保真度和总分上均达到最高，除一项批次校正指标外，其余指标表现最佳；该指标（PCR）因数据固有结构被人为惩罚。
• VMM在所有基准数据集上显著提升了scVI的性能，生物学保真度和总分最高，且在多个指标上实现最佳或接近最佳的批次校正表现。
• 通过MDE进行的定性可视化显示，与$\mathcal{N}(0,I)$先验相比，VMM生成的嵌入更具结构化且生物学可解释性更强，细胞类型和聚类的分离更清晰。
• 在SPLiT-seq数据集（包含两个生物学上相同的批次）上，VMM表现最佳，证实其不会过度校正技术变异，对极小批次效应具有鲁棒性。
• 与注释的细胞类型相比，VMM成功识别出有意义的生物学聚类，证明其能从复杂的scRNA-seq数据中提取出具有生物学相关性的结构。