[Paper Review] Information Constraints on Auto-Encoding Variational Bayes

Romain Lopez, Jeffrey Regier+2 more|arXiv (Cornell University)|May 22, 2018

Machine Learning and Data Classification47 references65 citations

TL;DR

Introduces HSIC-constrained VAEs (HCV) that penalize dependence in aggregated posteriors using dHSIC to enforce independence among latent components, with applications to interpretable, invariant, and denoised representations, including scRNA-seq analysis.

ABSTRACT

Parameterizing the approximate posterior of a generative model with neural networks has become a common theme in recent machine learning research. While providing appealing flexibility, this approach makes it difficult to impose or assess structural constraints such as conditional independence. We propose a framework for learning representations that relies on Auto-Encoding Variational Bayes and whose search space is constrained via kernel-based measures of independence. In particular, our method employs the $d$-variable Hilbert-Schmidt Independence Criterion (dHSIC) to enforce independence between the latent representations and arbitrary nuisance factors. We show how to apply this method to a range of problems, including the problems of learning invariant representations and the learning of interpretable representations. We also present a full-fledged application to single-cell RNA sequencing (scRNA-seq). In this setting the biological signal is mixed in complex ways with sequencing errors and sampling effects. We show that our method out-performs the state-of-the-art in this domain.

Motivation & Objective

• Motivate the need for latent representations with specific independence properties beyond traditional graphical-model constraints.
• Propose a framework to constrain aggregated posteriors in VAEs using kernel-based dependence measures.
• Introduce HSIC-Constrained VAE (HCV) by penalizing the ELBO with HSIC to enforce independence among latent subsets.
• Demonstrate the method on tasks of learning interpretable representations, invariant representations, and denoising single-cell RNA sequencing data.

Proposed method

• Review variational auto-encoder (VAE) and aggregated posterior q_hat_phi(Z) as the mixture over data.
• Introduce HSIC and its d-variate generalization (dHSIC) as kernel-based dependence measures.
• Define HSIC-constrained VAE (HCV) by maximizing the ELBO with a penalty proportional to HSIC between selected latent components (and optional nuisances).
• Justify the penalty as a lower bound on log-likelihood and highlight minibatch-compatible estimation of HSIC.
• Show that HSIC penalties can encompass and relate to existing invariance approaches (e.g., MMD-based methods) under certain kernel choices.
• Illustrate training and computational considerations, including kernel choices, batch settings, and the impact on the variational posterior.

Experimental results

Research questions

• RQ1 Can HSIC-based penalties enforce desired independence between latent coordinates in VAEs?
• RQ2 How does HCV compare to existing disentanglement/invariance approaches (e.g., beta-VAE, beta-TCVAE, VFAE) in terms of representation quality and reconstruction?
• RQ3 Can HCV improve downstream tasks such as interpretable factor discovery, invariance to nuisances, and denoising in real-world data?
• RQ4 What is the impact of the HSIC penalty on the ELBO, correlations in the aggregated posterior, and computational efficiency?
• RQ5 How does HCV perform on complex biological data such as single-cell RNA sequencing for denoising and hypothesis testing?

Key findings

• HCV achieves higher ELBO values (better model fit) while attaining lower dependence (as measured by HSIC) in the aggregated posterior than baselines like beta-VAE and beta-TCVAE in a linear Gaussian setup.
• On interpretable representation learning, HSIC-based penalties yield favorable trade-offs between reconstruction quality and independence compared to existing disentanglement methods.
• On invariant representation learning with lighting-contaminated face data, HCV produces representations that are more invariant to nuisance lighting while preserving identity information.
• In scRNA-seq denoising, HCV reduces dependence of latent space on quality-control nuisance metrics and improves Bayesian hypothesis testing performance (higher correlation with microarray results via IDR) compared with scVI, indicating more biologically interpretable latent factors.
• Overall, the HSIC penalty provides a scalable, minibatch-compatible way to enforce independence in latent representations without requiring full joint distribution estimates.