Variational autoencoder

In machine learning, a variational autoencoder (VAE), is an artificial neural network architecture introduced by Diederik P. Kingma and Max Welling, belonging to the families of probabilistic graphical models and variational Bayesian methods.

It is often associated with the autoencoder model because of its architectural affinity, but with significant differences in the goal and mathematical formulation. Variational autoencoders allow statistical inference problems to be rewritten (such as inferring the value of one random variable from another random variable) as statistical optimization problems (i.e find the parameter values that minimize some objective function). They are meant to map the input variable to a multivariate latent distribution. Although this type of model was initially designed for unsupervised learning, its effectiveness has been proven for semi-supervised learning and supervised learning.

Architecture

In a VAE the input data is sampled from a parametrized distribution (the prior, in Bayesian inference terms), and the encoder and decoder are trained jointly such that the output minimizes a reconstruction error in the sense of the Kullback–Leibler divergence between the parametric posterior and the true posterior.

Formulation

The basic scheme of a variational autoencoder. The model receives $\mathbf {x}$ as input. The encoder compresses it into the latent space. The decoder receives as input the information sampled from the latent space and produces $\mathbf {x'}$ as similar as possible to $\mathbf {x}$.

From a formal perspective, given an input dataset $\mathbf {x}$ characterized by an unknown probability distribution $P(\mathbf {x} )$, the objective is to model or approximate the data's true distribution $P$ using a parametrized distribution $p_{\theta }$ having parameters $\theta$. Let $\mathbf {z}$ be a random vector jointly-distributed with $\mathbf {x}$. Conceptually $\mathbf {z}$ will represent a latent encoding of $\mathbf {x}$. Marginalizing over $\mathbf {z}$ gives

$p_{\theta }(\mathbf {x} )=\int _{\mathbf {z} }p_{\theta }(\mathbf {x,z} )\,d\mathbf {z} ,$

where $p_{\theta }(\mathbf {x,z} )$ represents the joint distribution under $p_{\theta }$ of the observable data $\mathbf {x}$ and its latent representation or encoding $\mathbf {z}$. According to the chain rule, the equation can be rewritten as

$p_{\theta }(\mathbf {x} )=\int _{\mathbf {z} }p_{\theta }(\mathbf {x\mid z} )p_{\theta }(\mathbf {z} )\,d\mathbf {z}$

In the vanilla variational autoencoder, $\mathbf {z}$ is usually taken to be a finite-dimensional vector of real numbers, and $p_{\theta }(\mathbf {x|z} )$ to be a Gaussian distribution. Then $p_{\theta }(\mathbf {x} )$ is a mixture of Gaussian distributions.

It is now possible to define the set of the relationships between the input data and its latent representation as

• Prior $p_{\theta }(\mathbf {z} )$
• Likelihood $p_{\theta }(\mathbf {x} \mid \mathbf {z} )$
• Posterior $p_{\theta }(\mathbf {z} \mid \mathbf {x} )$

Unfortunately, the computation of $p_{\theta }(\mathbf {x} )$ is expensive and in most cases intractable. To speed up the calculus to make it feasible, it is necessary to introduce a further function to approximate the posterior distribution as

$q_{\Phi }(\mathbf {z\mid x} )\approx p_{\theta }(\mathbf {z\mid x} )$

with $\Phi$ defined as the set of real values that parametrize $q$.

In this way, the overall problem can be easily translated into the autoencoder domain, in which the conditional likelihood distribution $p_{\theta }(\mathbf {x} \mid \mathbf {z} )$ is carried by the probabilistic decoder, while the approximated posterior distribution $q_{\Phi }(\mathbf {z} \mid \mathbf {x} )$ is computed by the probabilistic encoder.

ELBO loss function

Main article: Evidence lower bound

As in every deep learning problem, it is necessary to define a differentiable loss function in order to update the network weights through backpropagation.

For variational autoencoders the idea is to jointly minimize the generative model parameters $\theta$ to reduce the reconstruction error between the input and the output, and $\Phi$ to have $q_{\Phi }(\mathbf {z\mid x} )$ as close as possible to $p_{\theta }(\mathbf {z} \mid \mathbf {x} )$.

As reconstruction loss, mean squared error and cross entropy are often used.

As distance loss between the two distributions the reverse Kullback–Leibler divergence $D_{KL}(q_{\Phi }(\mathbf {z\mid x} )\parallel p_{\theta }(\mathbf {z\mid x} ))$ is a good choice to squeeze $q_{\Phi }(\mathbf {z\mid x} )$ under $p_{\theta }(\mathbf {z} \mid \mathbf {x} )$.

The distance loss just defined is expanded as

${\begin{aligned}D_{KL}(q_{\Phi }(\mathbf {z\mid x} )\parallel p_{\theta }(\mathbf {z\mid x} ))&=\int q_{\Phi }(\mathbf {z\mid x} )\log {\frac {q_{\Phi }(\mathbf {z\mid x} )}{p_{\theta }(\mathbf {z\mid x} )}}\,d\mathbf {z} \\&=\int q_{\Phi }(\mathbf {z\mid x} )\log {\frac {q_{\Phi }(\mathbf {z\mid x} )p_{\theta }(\mathbf {x} )}{p_{\theta }(\mathbf {z,x} )}}\,d\mathbf {z} \\&=\int q_{\Phi }(\mathbf {z\mid x} )\left(\log(p_{\theta }(\mathbf {x} ))+\log {\frac {q_{\Phi }(\mathbf {z\mid x} )}{p_{\theta }(\mathbf {z,x} )}}\right)d\mathbf {z} \\&=\log(p_{\theta }(\mathbf {x} ))+\int q_{\Phi }(\mathbf {z\mid x} )\log {\frac {q_{\Phi }(\mathbf {z\mid x} )}{p_{\theta }(\mathbf {z,x} )}}\,d\mathbf {z} \\&=\log(p_{\theta }(\mathbf {x} ))+\int q_{\Phi }(\mathbf {z\mid x} )\log {\frac {q_{\Phi }(\mathbf {z\mid x} )}{p_{\theta }(\mathbf {x\mid z} )p_{\theta }(\mathbf {z} )}}\,d\mathbf {z} \\&=\log(p_{\theta }(\mathbf {x} ))+E_{\mathbf {z} \sim q_{\Phi }(\mathbf {z\mid x} )}(\log {\frac {q_{\Phi }(\mathbf {z\mid x} )}{p_{\theta }(\mathbf {z} )}}-\log(p_{\theta }(\mathbf {x\mid z} )))\\&=\log(p_{\theta }(\mathbf {x} ))+D_{KL}(q_{\Phi }(\mathbf {z\mid x} )\parallel p_{\theta }(\mathbf {z} ))-E_{\mathbf {z} \sim q_{\Phi }(\mathbf {z\mid x} )}(\log(p_{\theta }(\mathbf {x\mid z} )))\end{aligned}}$

At this point, it is possible to rewrite the equation as

$\log(p_{\theta }(\mathbf {x} ))-D_{KL}(q_{\Phi }(\mathbf {z\mid x} )\parallel p_{\theta }(\mathbf {z\mid x} ))=E_{\mathbf {z} \sim q_{\Phi }(\mathbf {z\mid x} )}(\log(p_{\theta }(\mathbf {x\mid z} )))-D_{KL}(q_{\Phi }(\mathbf {z\mid x} )\parallel p_{\theta }(\mathbf {z} ))$

The goal is to maximize the log-likelihood of the LHS of the equation to improve the generated data quality and to minimize the distribution distances between the real posterior and the estimated one.

This is equivalent to minimizing the negative log-likelihood, common practice in optimization.

The loss function so obtained, also named evidence lower bound loss function, shortly ELBO, can be written as

$L_{\theta ,\Phi }=-\log(p_{\theta }(\mathbf {x} ))+D_{KL}(q_{\Phi }(\mathbf {z\mid x} )\parallel p_{\theta }(\mathbf {z\mid x} ))=-E_{\mathbf {z} \sim q_{\Phi }(\mathbf {z|x} )}(\log(p_{\theta }(\mathbf {x\mid z} )))+D_{KL}(q_{\Phi }(\mathbf {z\mid x} )\parallel p_{\theta }(\mathbf {z} ))$

Given the non-negative property of the Kullback–Leibler divergence, it is correct to assert that

$-L_{\theta ,\Phi }=\log(p_{\theta }(\mathbf {x} ))-D_{KL}(q_{\Phi }(\mathbf {z\mid x} )\parallel p_{\theta }(\mathbf {z\mid x} ))\leq \log(p_{\theta }(\mathbf {x} ))$

The optimal parameters minimize this loss function. The problem can be summarized as

$\theta ^{*},\Phi ^{*}={\underset {\theta ,\Phi }{\operatorname {argmin} }}\,L_{\theta ,\Phi }$

The main advantage of this formulation relies on the possibility to jointly optimize with respect to parameters $\theta$ and $\Phi$.

Before applying the ELBO loss function to an optimization problem to backpropagate the gradient, it is necessary to make it differentiable by applying the so-called reparameterization trick to remove the stochastic sampling from the formation, and thus making it differentiable.

Reparameterization

The scheme of the reparameterization trick. The randomness variable $\mathbf {\varepsilon }$ is injected into the latent space $\mathbf {z}$ as external input. In this way, it is possible to backpropagate the gradient without involving stochastic variable during the update.

To make the ELBO formulation suitable for training purposes, it is necessary to slightly modify the problem formulation and the VAE structure.

Stochastic sampling is the non-differentiable operation through which it is possible to sample from the latent space and feed the probabilistic decoder.

The main assumption about the latent space is that it can be considered to be a set of multivariate Gaussian distributions, and thus can be described as

$\mathbf {z} \sim q_{\Phi }(\mathbf {z} \mid \mathbf {x} )={\mathcal {N}}({\boldsymbol {\mu }},{\boldsymbol {\sigma }}^{2})$.

The scheme of a variational autoencoder after the reparameterization trick.

Given ${\boldsymbol {\varepsilon }}\sim {\mathcal {N}}(0,{\boldsymbol {I}})$ and $\odot$ defined as the element-wise product, the reparameterization trick modifies the above equation as

$\mathbf {z} ={\boldsymbol {\mu }}+{\boldsymbol {\sigma }}\odot {\boldsymbol {\varepsilon }}.$

Thanks to this transformation (which can be extended to non-Gaussian distributions), the VAE becomes trainable and the probabilistic encoder has to learn how to map a compressed representation of the input into the two latent vectors ${\boldsymbol {\mu }}$ and ${\boldsymbol {\sigma }}$, while the stochasticity remains excluded from the updating process and is injected in the latent space as an external input through the random vector ${\boldsymbol {\varepsilon }}$.

Variations

Many variational autoencoders applications and extensions have been used to adapt the architecture to other domains and improve its performance.

$\beta$-VAE is an implementation with a weighted Kullback–Leibler divergence term to automatically discover and interpret factorised latent representations. With this implementation, it is possible to force manifold disentanglement for $\beta$ values greater than one. This architecture can discover disentangled latent factors without supervision.

The conditional VAE (CVAE), inserts label information in the latent space to force a deterministic constrained representation of the learned data.

Some structures directly deal with the quality of the generated samples or implement more than one latent space to further improve the representation learning.

Some architectures mix VAE and generative adversarial networks to obtain hybrid models.

See also

• Autoencoder
• Artificial neural network
• Deep learning
• Generative adversarial network
• Representation learning
• Sparse dictionary learning
• Data augmentation
• Backpropagation