[Paper Review] The VampPrior Mixture Model

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

Single-cell and spatial transcriptomicsBiochemistry, Genetics and Molecular Biology3 citations

TL;DR

The paper proposes the VampPrior Mixture Model (VMM), a novel prior for deep latent variable models that replaces the standard $\mathcal{N}(0,I)$ prior with a Dirichlet process Gaussian mixture model using VampPrior regularization, enabling automatic clustering and improved batch correction. When integrated into scVI, the VMM significantly enhances both biological conservation and integration performance on scRNA-seq datasets, outperforming existing methods in clustering and batch correction.

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.

Motivation & Objective

• To address the limitations of fixed-cluster priors in deep latent variable models (DLVMs), which require pre-specifying the number of clusters and are sensitive to initialization.
• To develop a prior that enables simultaneous integration and clustering in single-cell RNA-seq (scRNA-seq) data analysis, improving upon existing methods like scVI that use a $\mathcal{N}(0,I)$ prior.
• To create a scalable, plug-and-play prior that can be easily integrated into any DLVM with continuous latent variables, including scVI and related tools for multimodal omics.
• To demonstrate that the VMM enhances both biological conservation and technical batch correction in atlas-level scRNA-seq integration tasks.

Proposed method

• The VMM reformulates a Gaussian mixture model (GMM) as a Dirichlet process mixture by placing a VampPrior over the cluster means, enabling nonparametric clustering without predefining the number of clusters.
• The method alternates between amortized variational inference (updating variational parameters with fixed prior parameters) and Empirical Bayes steps (updating prior parameters), decoupling variational and prior learning.
• The VampPrior regularization ensures that the prior components are learned from data-driven, structured prior points rather than random initialization, improving stability and convergence.
• The VMM is integrated into the generative process of a Variational Autoencoder (VAE) and applied to both image clustering and scRNA-seq data, replacing the standard $\mathcal{N}(0,I)$ prior.
• For scRNA-seq, the VMM is incorporated into scVI, allowing the model to learn a shared, batch-corrected embedding while simultaneously discovering biologically meaningful clusters.
• The inference procedure is designed to be compatible with existing DLVM frameworks, enabling straightforward deployment in tools like scVI, scATAC-seq, and spatial transcriptomics pipelines.

Experimental results

Research questions

• RQ1 Can a deep latent variable model achieve robust, automatic clustering without predefining the number of clusters?
• RQ2 Can a flexible prior improve both batch correction and biological conservation in atlas-level scRNA-seq integration?
• RQ3 Does replacing the $\mathcal{N}(0,I)$ prior in scVI with a VampPrior Mixture Model lead to superior clustering and integration performance?
• RQ4 How does the alternating inference procedure between variational and prior parameters improve model stability and performance compared to joint optimization?

Key findings

• The VMM outperforms all VAE-based clustering methods and achieves state-of-the-art performance on image clustering benchmarks, approaching the performance of supervised methods.
• On the lung atlas dataset, the VMM achieves the highest biological conservation and total scores, with the best performance on all but one batch correction metric, where PCR scores were artificially penalized due to inherent data structure.
• The VMM significantly improves scVI’s performance on all benchmark datasets, with the highest biological conservation and total scores, and the best or near-best batch correction scores across multiple metrics.
• Qualitative visualization using MDE shows that the VMM produces more structured and biologically interpretable embeddings than the $\mathcal{N}(0,I)$ prior, with clearer separation of cell types and clusters.
• On the SPLiT-seq dataset, which has two biologically identical batches, the VMM achieves the best performance, confirming it does not over-correct technical variation and is robust to minimal batch effects.
• The VMM successfully discovers meaningful biological clusters when compared to annotated cell types, demonstrating its ability to extract biologically relevant structure from complex scRNA-seq data.