Variational autoencoder

From HandWiki
Short description: Deep learning generative model to encode data representation

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

Variational autoencoders are often associated with the autoencoder model because of its architectural affinity, but with significant differences in the goal and mathematical formulation. Variational autoencoders are probabilistic generative models that require neural networks as only a part of their overall structure. The neural network components are typically referred to as the encoder and decoder for the first and second component respectively. The first neural network maps the input variable to a latent space that corresponds to the parameters of a variational distribution. In this way, the encoder can produce multiple different samples that all come from the same distribution. The decoder has the opposite function, which is to map from the latent space to the input space, in order to produce or generate data points. Both networks are typically trained together with the usage of the reparameterization trick, although the variance of the noise model can be learned separately.

Although this type of model was initially designed for unsupervised learning,[2][3] its effectiveness has been proven for semi-supervised learning[4][5] and supervised learning.[6]

Overview of architecture and operation

A variational autoencoder is a generative model with a prior and noise distribution respectively. Usually such models are trained using the expectation-maximization meta-algorithm (e.g. probabilistic PCA, (spike & slab) sparse coding). Such a scheme optimizes a lower bound of the data likelihood, which is usually intractable, and in doing so requires the discovery of q-distributions, or variational posteriors. These q-distributions are normally parameterized for each individual data point in a separate optimization process. However, variational autoencoders use a neural network as an amortized approach to jointly optimize across data points. This neural network takes as input the data points themselves, and outputs parameters for the variational distribution. As it maps from a known input space to the low-dimensional latent space, it is called the encoder.

The decoder is the second neural network of this model. It is a function that maps from the latent space to the input space, e.g. as the means of the noise distribution. It is possible to use another neural network that maps to the variance, however this can be omitted for simplicity. In such a case, the variance can be optimized with gradient descent.

To optimize this model, one needs to know two terms: the "reconstruction error", and the Kullback–Leibler divergence(KL-D). Both terms are derived from the free energy expression of the probabilistic model, and therefore differ depending on the noise distribution and the assumed prior of the data. For example, a standard VAE task such as IMAGENET is typically assumed to have a gaussianly distributed noise, however tasks such as binarized MNIST require a Bernoulli noise. The KL-D from the free energy expression maximizes the probability mass of the q-distribution that overlaps with the p-distribution, which unfortunately can result in mode-seeking behaviour. The "reconstruction" term is the remainder of the free energy expression, and requires a sampling approximation to compute its expectation value.[7]

Formulation

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

From the point of view of probabilistic modeling, one wants to maximize the likelihood of the data [math]\displaystyle{ x }[/math] by their chosen parameterized probability distribution [math]\displaystyle{ p_{\theta}(x) = p(x|\theta) }[/math]. This distribution is usually chosen to be a Gaussian [math]\displaystyle{ N(x|\mu,\sigma) }[/math] which is parameterized by [math]\displaystyle{ \mu }[/math] and [math]\displaystyle{ \sigma }[/math] respectively, and as a member of the exponential family it is easy to work with as a noise distribution. Simple distributions are easy enough to maximize, however distributions where a prior is assumed over the latents [math]\displaystyle{ z }[/math] results in intractable integrals. Let us find [math]\displaystyle{ p_\theta(x) }[/math] via marginalizing over [math]\displaystyle{ z }[/math].

[math]\displaystyle{ p_\theta(x) = \int_{z}p_\theta({x,z}) \, dz, }[/math]

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

[math]\displaystyle{ p_\theta(x) = \int_{z}p_\theta({x| z})p_\theta(z) \, dz }[/math]

In the vanilla variational autoencoder, [math]\displaystyle{ z }[/math] is usually taken to be a finite-dimensional vector of real numbers, and [math]\displaystyle{ p_\theta({x|z}) }[/math] to be a Gaussian distribution. Then [math]\displaystyle{ p_\theta(x) }[/math] 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 [math]\displaystyle{ p_\theta(z) }[/math]
  • Likelihood [math]\displaystyle{ p_\theta(x|z) }[/math]
  • Posterior [math]\displaystyle{ p_\theta(z|x) }[/math]

Unfortunately, the computation of [math]\displaystyle{ p_\theta(z|x) }[/math] 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

[math]\displaystyle{ q_\phi({z| x}) \approx p_\theta({z| x}) }[/math]

with [math]\displaystyle{ \phi }[/math] defined as the set of real values that parametrize [math]\displaystyle{ q }[/math]. This is sometimes called amortized inference, since by "investing" in finding a good [math]\displaystyle{ q_\phi }[/math], one can later infer [math]\displaystyle{ z }[/math] from [math]\displaystyle{ x }[/math] quickly without doing any integrals.

In this way, the problem is to find a good probabilistic autoencoder, in which the conditional likelihood distribution [math]\displaystyle{ p_\theta(x|z) }[/math] is computed by the probabilistic decoder, and the approximated posterior distribution [math]\displaystyle{ q_\phi(z|x) }[/math] is computed by the probabilistic encoder.

Parametrize the encoder as [math]\displaystyle{ E_\phi }[/math], and the decoder as [math]\displaystyle{ D_\theta }[/math].

Evidence lower bound (ELBO)

Main page: 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 optimize the generative model parameters [math]\displaystyle{ \theta }[/math] to reduce the reconstruction error between the input and the output, and [math]\displaystyle{ \phi }[/math] to make [math]\displaystyle{ q_\phi({z| x}) }[/math] as close as possible to [math]\displaystyle{ p_\theta(z|x) }[/math]. As reconstruction loss, mean squared error and cross entropy are often used.

As distance loss between the two distributions the Kullback–Leibler divergence [math]\displaystyle{ D_{KL}(q_\phi({z| x})\parallel p_\theta({z| x})) }[/math] is a good choice to squeeze [math]\displaystyle{ q_\phi({z| x}) }[/math] under [math]\displaystyle{ p_\theta(z|x) }[/math].[7][8]

The distance loss just defined is expanded as

[math]\displaystyle{ \begin{align} D_{KL}(q_\phi({z| x})\parallel p_\theta({z| x})) &= \mathbb E_{z \sim q_\phi(\cdot | x)} \left[\ln \frac{q_\phi(z|x)}{p_\theta(z|x)}\right]\\ &= \mathbb E_{z \sim q_\phi(\cdot | x)} \left[\ln \frac{q_\phi({z| x})p_\theta(x)}{p_\theta(x, z)}\right]\\ &=\ln p_\theta(x) + \mathbb E_{z \sim q_\phi(\cdot | x)} \left[\ln \frac{q_\phi({z| x})}{p_\theta(x, z)}\right] \end{align} }[/math]

Now define the evidence lower bound (ELBO):[math]\displaystyle{ L_{\theta,\phi}(x) := \mathbb E_{z \sim q_\phi(\cdot | x)} \left[\ln \frac{p_\theta(x, z)}{q_\phi({z| x})}\right] = \ln p_\theta(x) - D_{KL}(q_\phi({\cdot| x})\parallel p_\theta({\cdot | x})) }[/math]Maximizing the ELBO[math]\displaystyle{ \theta^*,\phi^* = \underset{\theta,\phi}\operatorname{arg max} \, L_{\theta,\phi}(x) }[/math]is equivalent to simultaneously maximizing [math]\displaystyle{ \ln p_\theta(x) }[/math] and minimizing [math]\displaystyle{ D_{KL}(q_\phi({z| x})\parallel p_\theta({z| x})) }[/math]. That is, maximizing the log-likelihood of the observed data, and minimizing the divergence of the approximate posterior [math]\displaystyle{ q_\phi(\cdot | x) }[/math] from the exact posterior [math]\displaystyle{ p_\theta(\cdot | x) }[/math].

The form given is not very convenient for maximization, but the following, equivalent form, is:[math]\displaystyle{ L_{\theta,\phi}(x) = \mathbb E_{z \sim q_\phi(\cdot | x)} \left[\ln p_\theta(x|z)\right] - D_{KL}(q_\phi({\cdot| x})\parallel p_\theta(\cdot)) }[/math]where [math]\displaystyle{ \ln p_\theta(x|z) }[/math] is implemented as [math]\displaystyle{ -\frac{1}{2}\| x - D_\theta(z)\|^2_2 }[/math], since that is, up to an additive constant, what [math]\displaystyle{ x \sim \mathcal N(D_\theta(z), I) }[/math] yields. That is, we model the distribution of [math]\displaystyle{ x }[/math] conditional on [math]\displaystyle{ z }[/math] to be a Gaussian distribution centered on [math]\displaystyle{ D_\theta(z) }[/math]. The distribution of [math]\displaystyle{ q_\phi(z |x) }[/math] and [math]\displaystyle{ p_\theta(z) }[/math] are often also chosen to be Gaussians as [math]\displaystyle{ z|x \sim \mathcal N(E_\phi(x), \sigma_\phi(x)^2I) }[/math] and [math]\displaystyle{ z \sim \mathcal N(0, I) }[/math], with which we obtain by the formula for KL divergence of Gaussians:[math]\displaystyle{ L_{\theta,\phi}(x) = -\frac 12\mathbb E_{z \sim q_\phi(\cdot | x)} \left[ \|x - D_\theta(z)\|_2^2\right] - \frac 12 \left( N\sigma_\phi(x)^2 + \|E_\phi(x)\|_2^2 - 2N\ln\sigma_\phi(x) \right) + Const }[/math]Here [math]\displaystyle{ N }[/math] is the dimension of [math]\displaystyle{ z }[/math]. For a more detailed derivation and more interpretations of ELBO and its maximization, see its main page.

Reparameterization

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

To efficiently search for [math]\displaystyle{ \theta^*,\phi^* = \underset{\theta,\phi}\operatorname{arg max} \, L_{\theta,\phi}(x) }[/math]the typical method is gradient descent.

It is straightforward to find[math]\displaystyle{ \nabla_\theta \mathbb E_{z \sim q_\phi(\cdot | x)} \left[\ln \frac{p_\theta(x, z)}{q_\phi({z| x})}\right] = \mathbb E_{z \sim q_\phi(\cdot | x)} \left[ \nabla_\theta \ln \frac{p_\theta(x, z)}{q_\phi({z| x})}\right] }[/math]However, [math]\displaystyle{ \nabla_\phi \mathbb E_{z \sim q_\phi(\cdot | x)} \left[\ln \frac{p_\theta(x, z)}{q_\phi({z| x})}\right] }[/math]does not allow one to put the [math]\displaystyle{ \nabla_\phi }[/math] inside the expectation, since [math]\displaystyle{ \phi }[/math] appears in the probability distribution itself. The reparameterization trick (also known as stochastic backpropagation[9]) bypasses this difficulty.[7][10][11]

The most important example is when [math]\displaystyle{ z \sim q_\phi(\cdot | x) }[/math] is normally distributed, as [math]\displaystyle{ \mathcal N(\mu_\phi(x), \Sigma_\phi(x)) }[/math].

The scheme of a variational autoencoder after the reparameterization trick

This can be reparametrized by letting [math]\displaystyle{ \boldsymbol{\varepsilon} \sim \mathcal{N}(0, \boldsymbol{I}) }[/math] be a "standard random number generator", and construct [math]\displaystyle{ z }[/math] as [math]\displaystyle{ z = \mu_\phi(x) + L_\phi(x)\epsilon }[/math]. Here, [math]\displaystyle{ L_\phi(x) }[/math] is obtained by the Cholesky decomposition:[math]\displaystyle{ \Sigma_\phi(x) = L_\phi(x)L_\phi(x)^T }[/math]Then we have[math]\displaystyle{ \nabla_\phi \mathbb E_{z \sim q_\phi(\cdot | x)} \left[\ln \frac{p_\theta(x, z)}{q_\phi({z| x})}\right] = \mathbb {E}_{\epsilon}\left[ \nabla_\phi \ln {\frac {p_{\theta }(x, \mu_\phi(x) + L_\phi(x)\epsilon)}{q_{\phi }(\mu_\phi(x) + L_\phi(x)\epsilon | x)}}\right] }[/math]and so we obtained an unbiased estimator of the gradient, allowing stochastic gradient descent.

Since we reparametrized [math]\displaystyle{ z }[/math], we need to find [math]\displaystyle{ q_\phi(z|x) }[/math]. Let [math]\displaystyle{ q_0 }[/math] be the probability density function for [math]\displaystyle{ \epsilon }[/math], then[clarification needed][math]\displaystyle{ \ln q_\phi(z | x) = \ln q_0 (\epsilon) - \ln|\det(\partial_\epsilon z)| }[/math]where [math]\displaystyle{ \partial_\epsilon z }[/math] is the Jacobian matrix of [math]\displaystyle{ \epsilon }[/math] with respect to [math]\displaystyle{ z }[/math]. Since [math]\displaystyle{ z = \mu_\phi(x) + L_\phi(x)\epsilon }[/math], this is [math]\displaystyle{ \ln q_\phi(z | x) = -\frac 12 \|\epsilon\|^2 - \ln|\det L_\phi(x)| - \frac n2 \ln(2\pi) }[/math]

Variations

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

[math]\displaystyle{ \beta }[/math]-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 [math]\displaystyle{ \beta }[/math] values greater than one. This architecture can discover disentangled latent factors without supervision.[12][13]

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

Some structures directly deal with the quality of the generated samples[15][16] 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.[17][18][19]

See also


References

  1. Pinheiro Cinelli, Lucas et al. (2021). "Variational Autoencoder". Variational Methods for Machine Learning with Applications to Deep Networks. Springer. pp. 111–149. doi:10.1007/978-3-030-70679-1_5. ISBN 978-3-030-70681-4. https://books.google.com/books?id=N5EtEAAAQBAJ&pg=PA111. 
  2. Dilokthanakul, Nat; Mediano, Pedro A. M.; Garnelo, Marta; Lee, Matthew C. H.; Salimbeni, Hugh; Arulkumaran, Kai; Shanahan, Murray (2017-01-13). "Deep Unsupervised Clustering with Gaussian Mixture Variational Autoencoders". arXiv:1611.02648 [cs.LG].
  3. Hsu, Wei-Ning; Zhang, Yu; Glass, James (December 2017). "Unsupervised domain adaptation for robust speech recognition via variational autoencoder-based data augmentation". 2017 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU). pp. 16–23. doi:10.1109/ASRU.2017.8268911. ISBN 978-1-5090-4788-8. https://ieeexplore.ieee.org/abstract/document/8268911. 
  4. Ehsan Abbasnejad, M.; Dick, Anthony; van den Hengel, Anton (2017). Infinite Variational Autoencoder for Semi-Supervised Learning. pp. 5888–5897. https://openaccess.thecvf.com/content_cvpr_2017/html/Abbasnejad_Infinite_Variational_Autoencoder_CVPR_2017_paper.html. 
  5. Xu, Weidi; Sun, Haoze; Deng, Chao; Tan, Ying (2017-02-12). "Variational Autoencoder for Semi-Supervised Text Classification" (in en). Proceedings of the AAAI Conference on Artificial Intelligence 31 (1). doi:10.1609/aaai.v31i1.10966. https://ojs.aaai.org/index.php/AAAI/article/view/10966. 
  6. Kameoka, Hirokazu; Li, Li; Inoue, Shota; Makino, Shoji (2019-09-01). "Supervised Determined Source Separation with Multichannel Variational Autoencoder". Neural Computation 31 (9): 1891–1914. doi:10.1162/neco_a_01217. PMID 31335290. https://direct.mit.edu/neco/article/31/9/1891/8494/Supervised-Determined-Source-Separation-with. 
  7. 7.0 7.1 7.2 Kingma, Diederik P.; Welling, Max (2013-12-20). "Auto-Encoding Variational Bayes". arXiv:1312.6114 [stat.ML].
  8. "From Autoencoder to Beta-VAE" (in en). 2018-08-12. https://lilianweng.github.io/lil-log/2018/08/12/from-autoencoder-to-beta-vae.html. 
  9. Rezende, Danilo Jimenez; Mohamed, Shakir; Wierstra, Daan (2014-06-18). "Stochastic Backpropagation and Approximate Inference in Deep Generative Models" (in en). International Conference on Machine Learning (PMLR): 1278–1286. https://proceedings.mlr.press/v32/rezende14.html. 
  10. Bengio, Yoshua; Courville, Aaron; Vincent, Pascal (2013). "Representation Learning: A Review and New Perspectives". IEEE Transactions on Pattern Analysis and Machine Intelligence 35 (8): 1798–1828. doi:10.1109/TPAMI.2013.50. ISSN 1939-3539. PMID 23787338. https://ieeexplore.ieee.org/abstract/document/6472238. 
  11. Kingma, Diederik P.; Rezende, Danilo J.; Mohamed, Shakir; Welling, Max (2014-10-31). "Semi-Supervised Learning with Deep Generative Models". arXiv:1406.5298 [cs.LG].
  12. Higgins, Irina; Matthey, Loic; Pal, Arka; Burgess, Christopher; Glorot, Xavier; Botvinick, Matthew; Mohamed, Shakir; Lerchner, Alexander (2016-11-04) (in en). beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. https://openreview.net/forum?id=Sy2fzU9gl. 
  13. Burgess, Christopher P.; Higgins, Irina; Pal, Arka; Matthey, Loic; Watters, Nick; Desjardins, Guillaume; Lerchner, Alexander (2018-04-10). "Understanding disentangling in β-VAE". arXiv:1804.03599 [stat.ML].
  14. Sohn, Kihyuk; Lee, Honglak; Yan, Xinchen (2015-01-01) (in en). Learning Structured Output Representation using Deep Conditional Generative Models. https://proceedings.neurips.cc/paper/2015/file/8d55a249e6baa5c06772297520da2051-Paper.pdf. 
  15. Dai, Bin; Wipf, David (2019-10-30). "Diagnosing and Enhancing VAE Models". arXiv:1903.05789 [cs.LG].
  16. Dorta, Garoe; Vicente, Sara; Agapito, Lourdes; Campbell, Neill D. F.; Simpson, Ivor (2018-07-31). "Training VAEs Under Structured Residuals". arXiv:1804.01050 [stat.ML].
  17. Larsen, Anders Boesen Lindbo; Sønderby, Søren Kaae; Larochelle, Hugo; Winther, Ole (2016-06-11). "Autoencoding beyond pixels using a learned similarity metric" (in en). International Conference on Machine Learning (PMLR): 1558–1566. http://proceedings.mlr.press/v48/larsen16.html. 
  18. Bao, Jianmin; Chen, Dong; Wen, Fang; Li, Houqiang; Hua, Gang (2017). "CVAE-GAN: Fine-Grained Image Generation Through Asymmetric Training". pp. 2745–2754. arXiv:1703.10155 [cs.CV].
  19. Gao, Rui; Hou, Xingsong; Qin, Jie; Chen, Jiaxin; Liu, Li; Zhu, Fan; Zhang, Zhao; Shao, Ling (2020). "Zero-VAE-GAN: Generating Unseen Features for Generalized and Transductive Zero-Shot Learning". IEEE Transactions on Image Processing 29: 3665–3680. doi:10.1109/TIP.2020.2964429. ISSN 1941-0042. PMID 31940538. Bibcode2020ITIP...29.3665G. https://ieeexplore.ieee.org/abstract/document/8957359.