VampPrior Unleashed

Andrew A. Stirn; Abhejit Rajagopal; David A. Knowles; Grace Huynh

cs.LGOct 19, 2026

VampPrior Unleashed

Authors: Andrew A. Stirn, Abhejit Rajagopal, David A. Knowles, Grace Huynh

Organizations: Allen Institute · New York Genome Center & Columbia University

Abstract

Deep latent variable models (DLVMs) underpin many modern analysis pipelines, for example in single-cell genomics, yet their prior distributions have received markedly little attention. The VampPrior approximates the optimal prior (the aggregate posterior) and often attains high model evidence. Here, we identify two key limitations of the VampPrior: it ties prior and posterior uncertainty to a single amortized covariance and clusters poorly due to its fixed uniform mixture weights. We unleash the VampPrior by decoupling its prior and posterior covariances, resolving the conflict of generating high prior and low posterior entropy with the same shared parameters and substantially improving model evidence. We then replace the uniform mixture with a Dirichlet process. When combined, our proposals provide higher model evidence than the VampPrior and robust non-parametric clustering. Augmenting single cell Variational Inference (scVI) with our prior improves single-cell integration and clustering across six atlases while significantly raising model evidence. Our approach changes only the prior and thus can be applied to any DLVM with a continuous latent space.

Figures & tables

Figure 1: Effect of varying the number of components

K

on FashionMNIST with

u

initialized from randomly selected training data. Lines and bands are means

\pm 1

standard error over

20

seeds. Dashed connectors mark configurations at which the seed-paired one-sided Wilcoxon signed-rank test rejects at

p<0.05

; arrows indicate the direction of improvement.

Panel of Figure 2 — Figure 2: Modeling a 2D anisotropic GMM (three well-separated components) with a 1D latent and $K=40$ available components. Top: data colored by cluster assignment. Middle and bottom: the prior’s component densities weighted by the mixture proportions $\pi$ and by realized test-set usage, respectively. Only our DP VampPrior-Unleashed ( 2(c) ) recovers the three true clusters.

Figure 3: Evidence–NMI Pareto frontiers. Each point is one hyperparameter configuration averaged over

20

seeds; large points are the non-dominated configurations and the line connecting them interpolates the Pareto frontier. The dashed box spans the standard normal VAE (origin,

\times

) to the best per-datum log marginal and NMI attained by any prior. Configurations whose mean log marginal or NMI falls outside the grey box, i.e., worse than the standard normal VAE on either axis, are excluded before the frontier is computed. Stars mark each prior’s

P^{\ast}

, the frontier point closest to the normalized ideal (top-right corner of grey box after normalization to unit box). Both our VampPrior-Unleashed and DP-VampPrior-Unleashed offer sizable frontier advancements.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Effect of varying the number of components $K$ on FashionMNIST. Lines and bands are means $\pm 1$ standard error over $20$ seeds. Dashed connectors mark configurations at which the seed-paired one-sided Wilcoxon signed-rank test rejects at $p<0.05$ ; arrows indicate the direction of improvement.

Panel of Figure 6 — Figure 6: Effect of varying the number of components $K$ on MNIST. Lines and bands are means $\pm 1$ standard error over $20$ seeds. Dashed connectors mark configurations at which the seed-paired one-sided Wilcoxon signed-rank test rejects at $p<0.05$ ; arrows indicate the direction of improvement.

Panel of Figure 7 — Figure 7: Effect of varying the concentration parameter $\alpha$ on FashionMNIST with truncation set to $K=100$ . Lines and bands are means $\pm 1$ standard error over $20$ seeds. Dashed connectors mark configurations at which the seed-paired one-sided Wilcoxon signed-rank test rejects at $p<0.05$ ; arrows indicate the direction of improvement.

Panel of Figure 8 — Figure 8: Effect of varying the concentration parameter $\alpha$ on MNIST with truncation set to $K=100$ . Lines and bands are means $\pm 1$ standard error over $20$ seeds. Dashed connectors mark configurations at which the seed-paired one-sided Wilcoxon signed-rank test rejects at $p<0.05$ ; arrows indicate the direction of improvement.

Dataset	# Cells	# Donors	Cramér’s $V$	Cohen’s $k$	Association
DKD	39,176	11	0.145	10	Moderate
Immune Cell Atlas	329,762	12	0.227	11	High
GTEx v9	209,126	16	0.480	15	High
Mouse Pancreas Atlas	301,796	56	0.381	12	High
HypoMap	384,925	24	0.055	12	Low
Tabula Sapiens	483,152	29	0.444	28	High

Appendix

Table 4: Summary of single-cell RNA-seq datasets used in this study, including the number of cells, number of donors, Cramér’s

V

statistic for the association between cell type and donor,

k=\min(r,c)-1

for the

r\times c

donor-by-cell-type contingency table, and a qualitative assessment of association strength. Following Cohen (2013) , association is Low, Moderate, or High when

V

exceeds

0.1/\sqrt{k}

0.3/\sqrt{k}

, or

0.5/\sqrt{k}

, respectively.

Variational Inference

Variational inference approximates complex Bayesian posteriors with tractable distributions that support faster inference and uncertainty estimates.

See all 150 papers

Vectorized Bayesian Inference for Latent Dirichlet-Tree Allocation
Feb 21, 2026Zheng Wang, Nizar BouguilaTopic Modeling Bayesian Inference
Concordia Institute for Information Systems Engineering Concordia University Montreal, QC H3G 1M8, Canada
Predictive variational inference: Learn the predictively optimal posterior distribution
Oct 18, 2024Jinlin Lai, Antonio Linero, Yuling YaoModel Misspecification Variational Inference
College of Information and Computer Sciences, University of Massachusetts Amherst · Department of Statistics and Data Sciences, University of Texas, Austin
Scalable Posterior Uncertainty for Flexible Density-Based Clustering
Mar 3, 2026Nicola Bariletto, Stephen G. WalkerDensity-Based Clustering Clustering
Department of Statistics and Data Sciences, UT Austin

Feb 21, 2026J/K move · Enter open · S save

Vectorized Bayesian Inference for Latent Dirichlet-Tree Allocation

Zheng Wang, Nizar Bouguila

Concordia Institute for Information Systems Engineering Concordia University Montreal, QC H3G 1M8, Canada

Topic Modeling Bayesian Inference Variational Inference

Open paper AskPDF arXiv

Latent Dirichlet Allocation (LDA) is a foundational model for discovering latent thematic structure in discrete data, but its Dirichlet prior cannot represent the rich correlations and hierarchical relationships often present among topics. We introduce the framework of Latent Dirichlet-Tree Allocation (LDTA), a generalization of LDA that replaces the Dirichlet prior with an arbitrary Dirichlet-Tree (DT) distribution. LDTA preserves LDA's generative structure but enables expressive, tree-structured priors over topic proportions. To perform inference, we develop universal mean-field variational inference and Expectation Propagation, providing tractable updates for all DT. We reveal the vectorized nature of the two inference methods through theoretical development, and perform fully vectorized, GPU-accelerated implementations. The resulting framework substantially expands the modeling capacity of LDA while maintaining scalability and computational efficiency.

Abstract

Variational Inference

Explore similar work

Vectorized Bayesian Inference for Latent Dirichlet-Tree Allocation

Predictive variational inference: Learn the predictively optimal posterior distribution

Scalable Posterior Uncertainty for Flexible Density-Based Clustering