Source-linked AI summary
Information Constraints on Auto-Encoding Variational Bayes
Romain Lopez, Jeffrey Regier, Michael I. Jordan, Nir Yosef
TL;DR
Flexible neural-network variational posteriors make structural independence constraints difficult to impose and assess. The paper introduces HCV, which constrains latent representations by penalizing HSIC-based dependence in the variational objective, and applies it across representation-learning tasks. HCV compares favorably with existing methods and recovers biological signal better than the current state-of-the-art in scRNA-seq denoising.
Problem
Neural-network parameterization provides flexible approximate posteriors but makes structural constraints such as conditional independence difficult to impose or assess.
Method
HCV restricts the variational search space by penalizing the ELBO with HSIC, a kernel-based measure of dependence between samples from the variational posterior.
Results
HCV compares favorably with β-VAE and β-TCVAE, can improve on VFAE on its metrics, and recovers biological signal better than the current state-of-the-art in scRNA-seq denoising.
Takeaways & Limitations
The framework supports flexible independence-constrained representations for interpretable and invariant learning and for removing technical noise from scRNA-seq data.
Takeaways & Limitations
HSIC has quadratic computational cost, kernel selection is difficult, and the method is affected by choices such as the kernel bandwidth.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper introduces HCV, a VAE framework that constrains latent representations using kernel-based independence penalties. It applies this approach to interpretable and invariant representations and to denoising scRNA-seq data, where it improves on existing methods.
- Motivation: Flexible neural-network posteriors make structural constraints such as conditional independence difficult to impose or assess.The paper seeks alternatives to rigid conditional-independence structures in classical graphical models.
- Method: The method restricts the variational search space to latent representations with desired independence properties.It uses an encoder network to parameterize qφ(Z ∣X) and constrains the aggregated variational posterior.
- Method: HCV estimates dependence with HSIC and subtracts a scaled dependence penalty from the variational lower bound.Maximizing the resulting bound combines the traditional variational objective with a penalty for violating independence conditions.
- Applications: The paper applies HCV to interpretable representations, invariant representations, and single-cell RNA-sequencing denoising.These applications correspond to the tasks presented in Figure 1.
2 Background
The background introduces representation learning with AEVB and explains how kernel embeddings provide nonparametric dependence measures. HSIC measures dependence through a cross-covariance operator, while dHSIC extends the approach to multiple variables.
- Representation learning and AEVB: Representation learning transforms x into a vector z that supports a downstream task, while AEVB uses neural networks to parameterize variational distributions.VAE representations can explain data, predict labels, or serve other inferential purposes.
- Aggregated posterior: The aggregated posterior maps the empirical data distribution through qφ(Z ∣X,S), producing a mixture representation distribution.This distribution is also called the average encoding distribution.
- Kernel embeddings: Kernel mean embeddings represent distributions as points in reproducing kernel Hilbert spaces, with universal kernels yielding injective embeddings.Maximum mean discrepancy measures distance between such distribution embeddings.
- Kernel dependence measures: HSIC measures dependence through the Hilbert-Schmidt norm of a cross-covariance operator between kernel feature spaces.Its empirical V-statistic estimate has O(n^2(p + q)) computation time for variables in R^p and R^q.
- Kernel dependence measures: dHSIC generalizes HSIC from two variables to d variables.The paper uses this generalization to address independence statements involving multiple variables.
3 Theory for HSIC-Constrained VAE (HCV)
HCV corrects variational posteriors whose flexible neural approximations can introduce aggregated-posterior dependencies, using HSIC-based penalties to enforce selected independence statements.
- Ideal setting: Independence encoded by the generative model is recovered under infinite data when the variational gap is zero.This ideal result depends on an accurate posterior approximation and infinite data.
- Posterior approximation: In practice, high-dimensional data, naive mean-field assumptions, and learned generative networks can make the approximate aggregated posterior violate the model’s independence properties.The resulting posterior may differ substantially from the exact aggregated posterior.
- Posterior approximation: Neural networks can optimize a misspecified mean-field approximation while introducing supplemental correlation between latent variables in the aggregated posterior.The final decomposition term measures mean-field misspecification and increases with coupling among hidden variables.
- HSIC-constrained correction: HCV estimates dependencies nonparametrically with dHSIC and adds a penalty to the ELBO for selected latent or nuisance-variable independence constraints.The framework corrects qφ(Z ∣X,S) toward specified independence statements over subsets of Z and S.
- HSIC-constrained correction: The dHSIC penalty preserves a lower bound on log-likelihood while trading tightness of the bound for a more suitable representation through λ.The penalty is positive, so the bound becomes looser as independence is encouraged.
4 Case study: Learning interpretable representations
The interpretable-representation case study targets independent latent factors and evaluates HCV in a tractable linear Gaussian system, where it balances model fit against representation independence.
- Motivation: Independent representations are motivated as a step toward disentangled, interpretable latent factors corresponding to real variations in data.The paper contrasts HCV with β-VAE and β-TCVAE, whose independence penalties have reconstruction or estimation limitations.
- Motivation: HSIC is presented as a suitable independence penalty because its bias is O(1/n), whereas β-TCVAE’s total-correlation estimator has quadratic sample complexity and dataset-size-dependent bias.The β-TCVAE estimator’s bias is zero only when evaluated on the whole dataset.
- Experimental setup: The experiment uses a linear Gaussian system with tractable exact posterior inference and applies HCV to the two latent components u and v.The model defines Gaussian latent variables u and v and Gaussian observations x conditioned on them.
- Experimental setup: HCV evaluates the trade-off between aggregated-posterior correlation and ELBO across penalty weights against β-VAE, β-TCVAE, and an unconstrained VAE.The penalty is implemented by adding −λHSIC(Epdata(x)qφ(u,v ∣x)) to the ELBO.
- Results: HCV attains the highest ELBO values while having the lowest correlation in the linear Gaussian experiment.The comparison uses the trade-off shown in Figure 2, where the VAE baseline has greater aggregated-posterior dependence than the exact posterior.
5 Case study: Learning invariant representations
The paper applies HCV to learn face representations that remain predictive of identity while being invariant to lighting direction. Using HSIC avoids discretizing continuous lighting covariates and exposes a trade-off between nuisance correlation and prediction quality.
- Method: HCV models the full three-dimensional lighting direction and uses HSIC rather than discretizing the nuisance variable.The method accounts for the whole angular range of lighting directions, not only discrete direction groups.
- Problem: The task is to learn face representations predictive of person identity while having low correlation with lighting conditions.The Extended Yale B dataset contains 38 people under 50 lighting conditions, with 190 training images and remaining images used for testing.
- Results: HCV shows greater invariance to lighting direction while accurately predicting people’s identities.The comparison includes a vanilla VAE, VFAE using five lighting groups, HCV using the lighting vector, and unconstrained VAEs.
- Results: Using side information creates a trade-off between nuisance correlation and classifier performance in unconstrained models.Adding lighting information improves the identity classifier, but unconstrained models become more correlated with nuisance variables.
6 Case study: Learning denoised representations
The paper applies HCV to scRNA-seq data by separating biological signal from technical noise and quality-control effects. The resulting latent space is less correlated with alignment errors and supports gene rankings that better match microarray ground truth.
- Problem: The scRNA-seq task is to represent observations with independent biological-signal and technical-noise components despite their correlation with nuisance measurements.The data include gene expression values and quality-control metrics reflecting alignment errors and sampling effects.
- Method: HCV incorporates quality-control metrics into a generative model and integrates out technical noise during Bayesian hypothesis testing.The procedure uses z and u as latent variables and computes a Bayes factor by integrating over u.
- Results: The denoised latent space extracts information less subject to alignment errors and more biologically interpretable.The study uses PBMC data containing 12,039 cells, 3,346 genes, and five quality-control metrics.
- Results: The proposed method removes much nuisance correlation and yields gene rankings with a substantially higher correlation coefficient for the IDR.The ranking better conforms with microarray results than the comparison model.
7 Discussion
The paper presents HCV as a flexible framework that penalizes dependence in aggregated variational posteriors, applies it to denoising scRNA-seq data, and identifies computational and statistical limitations.
- HCV corrects independence properties in aggregated posteriors by penalizing the ELBO with HSIC between variational-posterior samples.
- Variational posterior misspecification can unintentionally promote dependence in the aggregated posterior, motivating future analysis of other variational approximations.
- Mini-batch HSIC estimation avoids learning additional distributional information, but could be harmful if mini-batches estimate HSIC poorly; experiments did not observe this in reasonable configurations.
- HSIC penalization has quadratic computational cost, motivating linear-time estimators or random-feature approximations for non-standard batch sizes.
- The dHSIC formulation may not scale well with many variables because standard concentration bounds have rate O(sqrt(d/n)).
- HCV is applied to scRNA-seq data to remove technical noise, with the framework also proposed for removing cell-cycle variation.
A dHSIC
The dHSIC construction extends kernel-based independence testing to multiple random variables through tensor-product Hilbert spaces, but kernel universality requires additional conditions.
- For d random variables, mutual independence means the joint distribution equals the tensor product of the marginal distributions.
- The characteristic property of individual kernels is insufficient in general to make the tensor-product kernel universal.
- Continuous, bounded, translation-invariant kernels provide a sufficient condition for universality in the tensor-product setting.
B Proof of independence in representation
This proof section establishes independence results within a generative-model construction whose random linear-system parameters define tractable likelihood and posterior calculations.
- The proof reduces independence to a pair of latent variables Zi and Zj under an infinite-data assumption.
- The construction chooses dimensions, matrices, and a positive regularization parameter before defining the generative model.
- Both the marginal log-likelihood and complete posterior are tractable in the constructed model.
D Proof of equivalence between HSIC and MMD
The proof connects HSIC for discrete nuisance variables to weighted maximum mean discrepancies between conditional latent distributions.
- The binary proof reorders M samples into s=0 and s=1 groups before manipulating the resulting kernel sums.
- For binary nuisance variables, the expression in parentheses is the V-statistic for the MMD between q(z|s=0) and q(z|s=1).
- For a discrete nuisance variable, the general expression becomes a sum of MMD terms weighted by empirical nuisance probabilities.
- The scRNA-seq nuisance variables include proportions of transcripts that map confidently, map only to the genome, fail to align, or undergo UMI or barcode correction.
F Graphical model for scRNA-seq
The scRNA-seq graphical model modifies scVI by adding latent variable u and sequencing-error node s, then uses AEVB to model expression, technical variation, and differential expression hypotheses.
- The proposed scRNA-seq model modifies scVI by adding latent variable u and a sequencing-error node s.
- AEVB uses observed variables X={x,s} and latent variables Z={z,l,u} with a mean-field variational posterior after analytically integrating out {w,y,h}.
- Bayesian hypothesis testing: The model supports Bayesian differential-expression testing between two cell populations by comparing hypotheses about gene expression.
- Monte Carlo integration is used for the remaining low-dimensional integrals, while reparameterization enables analytical derivation and differentiation of the resulting quantities.