[Paper Review] Variational Autoencoders and Nonlinear ICA: A Unifying Framework

Ilyes Khemakhem, Diederik P. Kingma+2 more|arXiv (Cornell University)|Jul 10, 2019

Neural Networks and Applications30 references130 citations

TL;DR

The paper introduces identifiability for a VAE by imposing a conditionally factorial latent prior dependent on an observed variable, unifying VAEs with identifiable nonlinear ICA and demonstrating recovery of latent sources up to simple transformations.

ABSTRACT

The framework of variational autoencoders allows us to efficiently learn deep latent-variable models, such that the model's marginal distribution over observed variables fits the data. Often, we're interested in going a step further, and want to approximate the true joint distribution over observed and latent variables, including the true prior and posterior distributions over latent variables. This is known to be generally impossible due to unidentifiability of the model. We address this issue by showing that for a broad family of deep latent-variable models, identification of the true joint distribution over observed and latent variables is actually possible up to very simple transformations, thus achieving a principled and powerful form of disentanglement. Our result requires a factorized prior distribution over the latent variables that is conditioned on an additionally observed variable, such as a class label or almost any other observation. We build on recent developments in nonlinear ICA, which we extend to the case with noisy, undercomplete or discrete observations, integrated in a maximum likelihood framework. The result also trivially contains identifiable flow-based generative models as a special case.

Motivation & Objective

• Motivate identifiability in deep latent-variable models beyond marginal data fit.
• Introduce a conditionally factorial latent prior p(z|u) to enable identifiability.
• Unify VAEs with identifiable nonlinear ICA within a maximum likelihood framework.
• Show that identifiable priors facilitate recovery of latent structure and enable data synthesis.

Proposed method

• Define a conditional generative model p(x,z|u)=p_f(x|z)p_T,lambdab(z|u) with x=f(z)+epsilon.
• Assume f is injective and epsilon is independent noise.
• Specify p(z|u) as a conditionally factorial exponential-family prior with parameters lambda(u).
• Train an identifiable VAE (iVAE) via the variational lower bound and reparameterization trick.
• Show identifiability up to an invertible linear transform A and component-wise transformations T_i.
• Relate the framework to nonlinear ICA and discuss special cases including normalizing flows.

Experimental results

Research questions

• RQ1 Can VAEs identify the true joint distribution p(x,z) when the latent prior is conditioned on an observed variable u?
• RQ2 Under what conditions is the model identifiable up to simple transformations (A and T_i) instead of fully identifiable?
• RQ3 How does conditioning the latent prior on u relate VAEs to nonlinear ICA and identifiable flow models?
• RQ4 Does the proposed iVAE enable recovery of latent sources and usable data synthesis in practice?

Key findings

• Identifiability up to simple transformations is achievable for a broad family of deep latent-variable models with a conditioned latent prior.
• The latent variables can be recovered up to an invertible linear transform and component-wise nonlinearities (T_i) under mild conditions.
• The iVAE framework, estimated via maximum likelihood-like objective, yields latent representations that closely approximate the true joint distribution in simulations.
• In simulations, iVAE substantially outperforms vanilla VAE and disentanglement baselines in recovering true sources (example MCC > 95% in 2D experiments).
• The approach provides a principled link between VAEs and nonlinear ICA, and encompasses identifiable flow-based models as a special case.