[Paper Review] The LORACs prior for VAEs: Letting the Trees Speak for the Data

Sharad Vikram, Matthew D. Hoffman+1 more|arXiv (Cornell University)|Oct 16, 2018

Generative Adversarial Networks and Image Synthesis18 references4 citations

TL;DR

This paper proposes the LORACs prior, a Bayesian nonparametric hierarchical clustering prior based on the time-marginalized coalescent (TMC), to improve interpretability and performance in variational autoencoders (VAEs). By modeling latent representations with tree-structured dependencies, the method learns more meaningful, discrete, and hierarchical data factorizations, significantly outperforming standard priors in downstream classification and retrieval tasks despite moderate likelihood scores.

ABSTRACT

In variational autoencoders, the prior on the latent codes $z$ is often treated as an afterthought, but the prior shapes the kind of latent representation that the model learns. If the goal is to learn a representation that is interpretable and useful, then the prior should reflect the ways in which the high-level factors that describe the data vary. The "default" prior is an isotropic normal, but if the natural factors of variation in the dataset exhibit discrete structure or are not independent, then the isotropic-normal prior will actually encourage learning representations that mask this structure. To alleviate this problem, we propose using a flexible Bayesian nonparametric hierarchical clustering prior based on the time-marginalized coalescent (TMC). To scale learning to large datasets, we develop a new inducing-point approximation and inference algorithm. We then apply the method without supervision to several datasets and examine the interpretability and practical performance of the inferred hierarchies and learned latent space.

Motivation & Objective

• To address the limitation of isotropic normal priors in VAEs, which mask discrete and dependent structure in data.
• To develop a scalable, nonparametric prior that captures hierarchical clustering structure in high-dimensional data.
• To enable interpretable, structured latent representations that reflect natural factors of variation in data.
• To improve downstream performance on classification and retrieval tasks by learning more meaningful latent geometry.

Proposed method

• Proposes the time-marginalized coalescent (TMC) as a flexible, interpretable prior over hierarchical tree structures in the latent space.
• Introduces an inducing-point approximation to scale the TMC prior to large datasets, enabling minibatch inference.
• Combines the TMC prior with a VAE framework, allowing end-to-end training of deep generative models with hierarchical priors.
• Uses a Gaussian likelihood with tree-structured variances to align latent distances with hierarchical cluster structure.
• Employs importance-weighted autoencoders (IWAE) with 1000 samples to estimate held-out log-likelihood for model comparison.
• Applies a variational inference framework with amortized inference over tree structures and latent codes.

Experimental results

Research questions

• RQ1 Can a nonparametric hierarchical prior improve the interpretability of VAE-learned latent representations compared to standard isotropic normal priors?
• RQ2 How well does the LORACs prior scale to large datasets while preserving hierarchical structure in the latent space?
• RQ3 Does the tree-structured prior lead to better performance on downstream tasks such as few-shot classification and information retrieval?
• RQ4 To what extent does the LORACs prior capture discrete, dependent factors of variation in data, such as gender and facial hair in face datasets?
• RQ5 Is the improved performance in downstream tasks correlated with higher held-out log-likelihood, or does it stem from better structural inductive bias?

Key findings

• The LORACs prior achieved the highest area-under-the-curve (AUC) in precision-recall for information retrieval on both MNIST and Omniglot, with 0.626 AUC on MNIST and 0.087 on Omniglot.
• In few-shot classification, LORACs consistently outperformed other priors, especially on Omniglot, with significantly better accuracy at low label counts.
• TSNE visualizations showed that LORACs learned more concentrated and separated clusters compared to normal and VampPrior, indicating better cluster structure.
• Despite outperforming other priors on downstream tasks, the LORACs prior achieved middling held-out log-likelihood scores, suggesting that likelihood is not a reliable proxy for representation quality.
• Manual labeling of inducing points revealed that LORACs representations were more useful in active learning, with higher accuracy using fewer labeled points.
• The method demonstrated scalability through an inducing-point approximation, enabling training on large datasets without sacrificing structural fidelity.