Variational Autoencoder: Intuition and Implementation

There are two generative models facing neck to neck in the data generation business right now: Generative Adversarial Nets (GAN) and Variational Autoencoder (VAE). These two models have different take on how the models are trained. GAN is rooted in game theory, its objective is to find the Nash Equilibrium between discriminator net and generator net. On the other hand, VAE is rooted in bayesian inference, i.e. it wants to model the underlying probability distribution of data so that it could sample new data from that distribution.

In this post, we will look at the intuition of VAE model and its implementation in Keras.

VAE: Formulation and Intuition

Suppose we want to generate a data. Good way to do it is first to decide what kind of data we want to generate, then actually generate the data. For example, say, we want to generate an animal. First, we imagine the animal: it must have four legs, and it must be able to swim. Having those criteria, we could then actually generate the animal by sampling from the animal kingdom. Lo and behold, we get Platypus!

From the story above, our imagination is analogous to latent variable. It is often useful to decide the latent variable first in generative models, as latent variable could describe our data. Without latent variable, it is as if we just generate data blindly. And this is the difference between GAN and VAE: VAE uses latent variable, hence it’s an expressive model.

Alright, that fable is great and all, but how do we model that? Well, let’s talk about probability distribution.

Let’s define some notions:

1. $\mathrm{XX:\; data\; that\; we\; want\; to\; model\; a.k.a\; the\; animal}$
2. $\mathrm{zz:\; latent\; variable\; a.k.a\; our\; imagination}$
3. $\mathrm{P(X)P(X):\; probability\; distribution\; of\; the\; data,\; i.e.\; that\; animal\; kingdom}$
4. $\mathrm{P(z)P(z):\; probability\; distribution\; of\; latent\; variable,\; i.e.\; our\; brain,\; the\; source\; of\; our\; imagination}$
5. $\mathrm{P(X|z)P(X|z):\; distribution\; of\; generating\; data\; given\; latent\; variable,\; e.g.\; turning\; imagination\; into\; real\; animal}$

Our objective here is to model the data, hence we want to find $\mathrm{P(X)P(X).\; Using\; the\; law\; of\; probability,\; we\; could\; find\; it\; in\; relation\; with\; zz\; as\; follows:}$

$$\mathrm{P(X)=\int P(X|z)P(z)dzP(X)=\int P(X|z)P(z)dz}$$

that is, we marginalize out $\mathrm{zz\; from\; the\; joint\; probability\; distribution\; P(X,z)P(X,z).}$

Now if only we know $\mathrm{P(X,z)P(X,z),\; or\; equivalently,\; P(X|z)P(X|z)\; and\; P(z)P(z)\dots }$

The idea of VAE is to infer $\mathrm{P(z)P(z)\; using\; P(z|X)P(z|X).\; This\; is\; make\; a\; lot\; of\; sense\; if\; we\; think\; about\; it:\; we\; want\; to\; make\; our\; latent\; variable\; likely\; under\; our\; data.\; Talking\; in\; term\; of\; our\; fable\; example,\; we\; want\; to\; limit\; our\; imagination\; only\; on\; animal\; kingdom\; domain,\; so\; we\; shouldn’t\; imagine\; about\; things\; like\; root,\; leaf,\; tyre,\; glass,\; GPU,\; refrigerator,\; doormat,\; \dots \; as\; it’s\; unlikely\; that\; those\; things\; have\; anything\; to\; do\; with\; things\; that\; come\; from\; the\; animal\; kingdom.\; Right?}$

But the problem is, we have to infer that distribution $\mathrm{P(z|X)P(z|X),\; as\; we\; don’t\; know\; it\; yet.\; In\; VAE,\; as\; it\; name\; suggests,\; we\; infer\; P(z|X)P(z|X)\; using\; a\; method\; called\; Variational\; Inference\; (VI).\; VI\; is\; one\; of\; the\; popular\; choice\; of\; method\; in\; bayesian\; inference,\; the\; other\; one\; being\; MCMC\; method.\; The\; main\; idea\; of\; VI\; is\; to\; pose\; the\; inference\; by\; approach\; it\; as\; an\; optimization\; problem.\; How?\; By\; modeling\; the\; true\; distribution\; P(z|X)P(z|X)\; using\; simpler\; distribution\; that\; is\; easy\; to\; evaluate,\; e.g.\; Gaussian,\; and\; minimize\; the\; difference\; between\; those\; two\; distribution\; using\; KL\; divergence\; metric,\; which\; tells\; us\; how\; difference\; it\; is\; PP\; and\; QQ.}$

Alright, now let’s say we want to infer $\mathrm{P(z|X)P(z|X)\; using\; Q(z|X)Q(z|X).\; The\; KL\; divergence\; then\; formulated\; as\; follows:}$

$$\begin{array}{c}{\mathrm{DKL[Q(z|X)\parallel P(z|X)]=\sum zQ(z|X)logQ(z|X)P(z|X)=E[logQ(z|X)P(z|X)]=E[logQ(z|X)-logP(z|X)]DKL[Q(z|X)\Vert P(z|X)]=\sum zQ(z|X)logQ(z|X)P(z|X)=E[logQ(z|X)P(z|X)]=E[logQ(z|X)-logP(z|X)]}}^{}\end{array}$$

Recall the notations above, there are two things that we haven’t use, namely $\mathrm{P(X)P(X),\; P(X|z)P(X|z),\; and\; P(z)P(z).\; But,\; with\; Bayes’\; rule,\; we\; could\; make\; it\; appear\; in\; the\; equation:}$

$$\begin{array}{c}{\mathrm{DKL[Q(z|X)\parallel P(z|X)]=E[logQ(z|X)-logP(X|z)P(z)P(X)]=E[logQ(z|X)-(logP(X|z)+logP(z)-logP(X))]=E[logQ(z|X)-logP(X|z)-logP(z)+logP(X)]DKL[Q(z|X)\Vert P(z|X)]=E[logQ(z|X)-logP(X|z)P(z)P(X)]=E[logQ(z|X)-(logP(X|z)+logP(z)-logP(X))]=E[logQ(z|X)-logP(X|z)-logP(z)+logP(X)]}}^{}\end{array}$$

Notice that the expectation is over

$$\begin{array}{c}{\mathrm{z\; and\; zP(X)\; doesn’t\; depend\; on\; P(X)z,\; so\; we\; could\; move\; it\; outside\; of\; the\; expectation.zDKL[Q(z|X)\parallel P(z|X)]=E[logQ(z|X)-logP(X|z)-logP(z)]+logP(X)DKL[Q(z|X)\parallel P(z|X)]-logP(X)=E[logQ(z|X)-logP(X|z)-logP(z)]DKL[Q(z|X)\Vert P(z|X)]=E[logQ(z|X)-logP(X|z)-logP(z)]+logP(X)DKL[Q(z|X)\Vert P(z|X)]-logP(X)=E[logQ(z|X)-logP(X|z)-logP(z)]}}^{}\end{array}$$

If we look carefully at the right hand side of the equation, we would notice that it could be rewritten as another KL divergence. So let’s do that by first rearranging the sign.

$$\begin{array}{c}{\mathrm{DKL[Q(z|X)\parallel P(z|X)]-logP(X)=E[logQ(z|X)-logP(X|z)-logP(z)]logP(X)-DKL[Q(z|X)\parallel P(z|X)]=E[logP(X|z)-(logQ(z|X)-logP(z))]=E[logP(X|z)]-E[logQ(z|X)-logP(z)]=E[logP(X|z)]-DKL[Q(z|X)\parallel P(z)]DKL[Q(z|X)\Vert P(z|X)]-logP(X)=E[logQ(z|X)-logP(X|z)-logP(z)]logP(X)-DKL[Q(z|X)\Vert P(z|X)]=E[logP(X|z)-(logQ(z|X)-logP(z))]=E[logP(X|z)]-E[logQ(z|X)-logP(z)]=E[logP(X|z)]-DKL[Q(z|X)\Vert P(z)]}}^{}\end{array}$$

And this is it, the VAE objective function:

$$\begin{array}{c}{\mathrm{logP(X)-DKL[Q(z|X)\parallel P(z|X)]=E[logP(X|z)]-DKL[Q(z|X)\parallel P(z)]logP(X)-DKL[Q(z|X)\Vert P(z|X)]=E[logP(X|z)]-DKL[Q(z|X)\Vert P(z)]}}^{}\end{array}$$

At this point, what do we have? Let’s enumerate:

$$\begin{array}{c}{\mathrm{Q(z|X) that\; project\; our\; data Q(z|X)X into\; latent\; variable\; spaceX\; z,\; the\; latent\; variablez\; P(X|z) that\; generate\; data\; given\; latent\; variableP(X|z)}}^{}\end{array}$$

We might feel familiar with this kind of structure. And guess what, it’s the same structure as seen in ! That is,

$$\begin{array}{c}{\mathrm{AutoencoderQ(z|X)\; is\; the\; encoder\; net,\; Q(z|X)z\; is\; the\; encoded\; representation,\; and\; zP(X|z)\; is\; the\; decoder\; net!\; Well,\; well,\; no\; wonder\; the\; name\; of\; this\; model\; is\; Variational\; Autoencoder!P(X|z)}}^{}\end{array}$$

VAE: Dissecting the Objective

It turns out, VAE objective function has a very nice interpretation. That is, we want to model our data, which described by

$$\begin{array}{c}{\mathrm{logP(X),\; under\; some\; error\; logP(X)DKL[Q(z|X)\parallel P(z|X)].\; In\; other\; words,\; VAE\; tries\; to\; find\; the\; lower\; bound\; of\; DKL[Q(z|X)\Vert P(z|X)]logP(X),\; which\; in\; practice\; is\; good\; enough\; as\; trying\; to\; find\; the\; exact\; distribution\; is\; often\; untractable.logP(X)}}^{}\end{array}$$

That model then could be found by maximazing over some mapping from latent variable to data

$$\begin{array}{c}{\mathrm{logP(X|z)\; and\; minimizing\; the\; difference\; between\; our\; simple\; distribution\; logP(X|z)Q(z|X)\; and\; the\; true\; latent\; distribution\; Q(z|X)P(z).P(z)}}^{}\end{array}$$

As we might already know, maximizing

$$\begin{array}{c}{\mathrm{E[logP(X|z)]\; is\; a\; maximum\; likelihood\; estimation.\; We\; basically\; see\; it\; all\; the\; time\; in\; discriminative\; supervised\; model,\; for\; example\; Logistic\; Regression,\; SVM,\; or\; Linear\; Regression.\; In\; the\; other\; words,\; given\; an\; input\; E[logP(X|z)]z\; and\; an\; output\; zX,\; we\; want\; to\; maximize\; the\; conditional\; distribution\; XP(X|z)\; under\; some\; model\; parameters.\; So\; we\; could\; implement\; it\; by\; using\; any\; classifier\; with\; input\; P(X|z)z\; and\; output\; zX,\; then\; optimize\; the\; objective\; function\; by\; using\; for\; example\; log\; loss\; or\; regression\; loss.X}}^{}\end{array}$$

What about

$$\begin{array}{c}{\mathrm{DKL[Q(z|X)\parallel P(z)]?\; Here,\; DKL[Q(z|X)\Vert P(z)]P(z)\; is\; the\; latent\; variable\; distribution.\; We\; might\; want\; to\; sample\; P(z)P(z)\; later,\; so\; the\; easiest\; choice\; is\; P(z)N(0,1).\; Hence,\; we\; want\; to\; make\; N(0,1)Q(z|X)\; to\; be\; as\; close\; as\; possible\; to\; Q(z|X)N(0,1)\; so\; that\; we\; could\; sample\; it\; easily.N(0,1)}}^{}\end{array}$$

Having

$$\begin{array}{c}{\mathrm{P(z)=N(0,1)\; also\; add\; another\; benefit.\; Let’s\; say\; we\; also\; want\; P(z)=N(0,1)Q(z|X)\; to\; be\; Gaussian\; with\; parameters\; Q(z|X)\mu (X)\; and\; \mu (X)\Sigma (X),\; i.e.\; the\; mean\; and\; variance\; given\; X.\; Then,\; the\; KL\; divergence\; between\; those\; two\; distribution\; could\; be\; computed\; in\; closed\; form!\Sigma (X)DKL[N(\mu (X),\Sigma (X))\parallel N(0,1)]=12(tr(\Sigma (X))+\mu (X)T\mu (X)-k-logdet(\Sigma (X)))DKL[N(\mu (X),\Sigma (X))\Vert N(0,1)]=12(tr(\Sigma (X))+\mu (X)T\mu (X)-k-logdet(\Sigma (X)))}}^{}\end{array}$$

Above,

$$\begin{array}{c}{\mathrm{k\; is\; the\; dimension\; of\; our\; Gaussian.\; ktr(X)\; is\; trace\; function,\; i.e.\; sum\; of\; the\; diagonal\; of\; matrix\; tr(X)X.\; The\; determinant\; of\; a\; diagonal\; matrix\; could\; be\; computed\; as\; product\; of\; its\; diagonal.\; So\; really,\; we\; could\; implement\; X\Sigma (X)\; as\; just\; a\; vector\; as\; it’s\; a\; diagonal\; matrix:\Sigma (X)DKL[N(\mu (X),\Sigma (X))\parallel N(0,1)]=12(\sum k\Sigma (X)+\sum k\mu 2(X)-\sum k1-log\prod k\Sigma (X))=12(\sum k\Sigma (X)+\sum k\mu 2(X)-\sum k1-\sum klog\Sigma (X))=12\sum k(\Sigma (X)+\mu 2(X)-1-log\Sigma (X))DKL[N(\mu (X),\Sigma (X))\Vert N(0,1)]=12(\sum k\Sigma (X)+\sum k\mu 2(X)-\sum k1-log\prod k\Sigma (X))=12(\sum k\Sigma (X)+\sum k\mu 2(X)-\sum k1-\sum klog\Sigma (X))=12\sum k(\Sigma (X)+\mu 2(X)-1-log\Sigma (X))}}^{}\end{array}$$

In practice, however, it’s better to model

$$\begin{array}{c}{\mathrm{\Sigma (X)\; as\; \Sigma (X)log\Sigma (X),\; as\; it\; is\; more\; numerically\; stable\; to\; take\; exponent\; compared\; to\; computing\; log.\; Hence,\; our\; final\; KL\; divergence\; term\; is:log\Sigma (X)DKL[N(\mu (X),\Sigma (X))\parallel N(0,1)]=12\sum k(exp(\Sigma (X))+\mu 2(X)-1-\Sigma (X))DKL[N(\mu (X),\Sigma (X))\Vert N(0,1)]=12\sum k(exp(\Sigma (X))+\mu 2(X)-1-\Sigma (X))}}^{}\end{array}$$

Implementation in Keras

First, let’s implement the encoder net

$$\begin{array}{c}{\mathrm{Q(z|X),\; which\; takes\; input\; Q(z|X)X\; and\; outputting\; two\; things:\; X\mu (X)\; and\; \mu (X)\Sigma (X),\; the\; parameters\; of\; the\; Gaussian.\Sigma (X)}}^{}\end{array}$$

$$\begin{array}{c}{\mathrm{from\; tensorflow.examples.tutorials.mnist\; import\; input\_data\; from\; keras.layers\; import\; Input,\; Dense,\; Lambda\; from\; keras.models\; import\; Model\; from\; keras.objectives\; import\; binary\_crossentropy\; from\; keras.callbacks\; import\; LearningRateScheduler\; import\; numpy\; as\; np\; import\; matplotlib.pyplot\; as\; plt\; import\; keras.backend\; as\; K\; import\; tensorflow\; as\; tf\; m\; =\; 50\; n\_z\; =\; 2\; n\_epoch\; =\; 10\; \#\; Q(z|X)\; --\; encoder\; inputs\; =\; Input(shape=(784,))\; h\_q\; =\; Dense(512,\; activation=\text{'}relu\text{'})(inputs)\; mu\; =\; Dense(n\_z,\; activation=\text{'}linear\text{'})(h\_q)\; log\_sigma\; =\; Dense(n\_z,\; activation=\text{'}linear\text{'})(h\_q)}}^{}\end{array}$$

That is, our

$$\begin{array}{c}{\mathrm{Q(z|X)\; is\; a\; neural\; net\; with\; one\; hidden\; layer.\; In\; this\; implementation,\; our\; latent\; variable\; is\; two\; dimensional,\; so\; that\; we\; could\; easily\; visualize\; it.\; In\; practice\; though,\; more\; dimension\; in\; latent\; variable\; should\; be\; better.Q(z|X)}}^{}\end{array}$$

However, we are now facing a problem. How do we get

$$\begin{array}{c}{\mathrm{z\; from\; the\; encoder\; outputs?\; Obviously\; we\; could\; sample\; zz\; from\; a\; Gaussian\; which\; parameters\; are\; the\; outputs\; of\; the\; encoder.\; Alas,\; sampling\; directly\; won’t\; do,\; if\; we\; want\; to\; train\; VAE\; with\; gradient\; descent\; as\; the\; sampling\; operation\; doesn’t\; have\; gradient!z}}^{}\end{array}$$

There is, however a trick called reparameterization trick, which makes the network differentiable. Reparameterization trick basically divert the non-differentiable operation out of the network, so that, even though we still involve a thing that is non-differentiable, at least it is out of the network, hence the network could still be trained.

The reparameterization trick is as follows. Recall, if we have

$$\begin{array}{c}{\mathrm{x\sim N(\mu ,\Sigma )\; and\; then\; standardize\; it\; so\; that\; x\sim N(\mu ,\Sigma )\mu =0,\Sigma =1,\; we\; could\; revert\; it\; back\; to\; the\; original\; distribution\; by\; reverting\; the\; standardization\; process.\; Hence,\; we\; have\; this\; equation:\mu =0,\Sigma =1x=\mu +\Sigma 12xstdx=\mu +\Sigma 12xstd}}^{}\end{array}$$

With that in mind, we could extend it. If we sample from a standard normal distribution, we could convert it to any Gaussian we want if we know the mean and the variance. Hence we could implement our sampling operation of

$$\begin{array}{c}{\mathrm{z\; by:zz=\mu (X)+\Sigma 12(X)ϵz=\mu (X)+\Sigma 12(X)ϵ}}^{}\end{array}$$

where

$$\begin{array}{c}{\mathrm{ϵ\sim N(0,1).ϵ\sim N(0,1)}}^{}\end{array}$$

Now, during backpropagation, we don’t care anymore with the sampling process, as it is now outside of the network, i.e. doesn’t depend on anything in the net, hence the gradient won’t flow through it.

$$\begin{array}{c}{\mathrm{def\; sample\_z(args):\; mu,\; log\_sigma\; =\; args\; eps\; =\; K.random\_normal(shape=(m,\; n\_z),\; mean=0.,\; std=1.)\; return\; mu\; +\; K.exp(log\_sigma\; /\; 2)\; *\; eps\; \#\; Sample\; z\; ~\; Q(z|X)\; z\; =\; Lambda(sample\_z)([mu,\; log\_sigma])}}^{}\end{array}$$

Now we create the decoder net

$$\begin{array}{c}{\mathrm{P(X|z):P(X|z)}}^{}\end{array}$$

$$\begin{array}{c}{\mathrm{\#\; P(X|z)\; --\; decoder\; decoder\_hidden\; =\; Dense(512,\; activation=\text{'}relu\text{'})\; decoder\_out\; =\; Dense(784,\; activation=\text{'}sigmoid\text{'})\; h\_p\; =\; decoder\_hidden(z)\; outputs\; =\; decoder\_out(h\_p)}}^{}\end{array}$$

Lastly, from this model, we can do three things: reconstruct inputs, encode inputs into latent variables, and generate data from latent variable. So, we have three Keras models:

$$\begin{array}{c}{\mathrm{\#\; Overall\; VAE\; model,\; for\; reconstruction\; and\; training\; vae\; =\; Model(inputs,\; outputs)\; \#\; Encoder\; model,\; to\; encode\; input\; into\; latent\; variable\; \#\; We\; use\; the\; mean\; as\; the\; output\; as\; it\; is\; the\; center\; point,\; the\; representative\; of\; the\; gaussian\; encoder\; =\; Model(inputs,\; mu)\; \#\; Generator\; model,\; generate\; new\; data\; given\; latent\; variable\; z\; d\_in\; =\; Input(shape=(n\_z,))\; d\_h\; =\; decoder\_hidden(d\_in)\; d\_out\; =\; decoder\_out(d\_h)\; decoder\; =\; Model(d\_in,\; d\_out)}}^{}\end{array}$$

Then, we need to translate our loss into Keras code:

$$\begin{array}{c}{\mathrm{def\; vae\_loss(y\_true,\; y\_pred):\; """\; Calculate\; loss\; =\; reconstruction\; loss\; +\; KL\; loss\; for\; each\; data\; in\; minibatch\; """\; \#\; E[log\; P(X|z)]\; recon\; =\; K.sum(K.binary\_crossentropy(y\_pred,\; y\_true),\; axis=1)\; \#\; D\_KL(Q(z|X)\; ||\; P(z|X));\; calculate\; in\; closed\; form\; as\; both\; dist.\; are\; Gaussian\; kl\; =\; 0.5\; *\; K.sum(K.exp(log\_sigma)\; +\; K.square(mu)\; -\; 1.\; -\; log\_sigma,\; axis=1)\; return\; recon\; +\; kl}}^{}\end{array}$$

and then train it:

$$\begin{array}{c}{\mathrm{vae.compile(optimizer=\text{'}adam\text{'},\; loss=vae\_loss)\; vae.fit(X\_train,\; X\_train,\; batch\_size=m,\; nb\_epoch=n\_epoch)}}^{}\end{array}$$

And that’s it, the implementation of VAE in Keras!

Implementation on MNIST Data

We could use any dataset really, but like always, we will use MNIST as an example.

After we trained our VAE model, we then could visualize the latent variable space

$$\begin{array}{c}{\mathrm{Q(z|X):Q(z|X)}}^{}\end{array}$$

 

As we could see, in the latent space, the representation of our data that have the same characteristic, e.g. same label, are close to each other. Notice that in the training phase, we never provide any information regarding the data.

We could also look at the data reconstruction by running through the data into overall VAE net:

 

Lastly, we could generate new sample by first sample

$$\begin{array}{c}{\mathrm{z\sim N(0,1)\; and\; feed\; it\; into\; our\; decoder\; net:z\sim N(0,1)}}^{}\end{array}$$

 

If we look closely on the reconstructed and generated data, we would notice that some of the data are ambiguous. For example the digit 5 looks like 3 or 8. That’s because our latent variable space is a continous distribution (i.e.

$$\begin{array}{c}{\mathrm{N(0,1)),\; hence\; there\; bound\; to\; be\; some\; smooth\; transition\; on\; the\; edge\; of\; the\; clusters.\; And\; also,\; the\; cluster\; of\; digits\; are\; close\; to\; each\; other\; if\; they\; are\; somewhat\; similar.\; That’s\; why\; in\; the\; latent\; space,\; 5\; is\; close\; to\; 3.N(0,1)}}^{}\end{array}$$

Conclusion

In this post we looked at the intuition behind Variational Autoencoder (VAE), its formulation, and its implementation in Keras.

We also saw the difference between VAE and GAN, the two most popular generative models nowadays.

For more math on VAE, be sure to hit the original paper by Kingma et al., 2014. There is also an excellent tutorial on VAE by Carl Doersch. Check out the references section below.

The full code is available in my repo:

$$\begin{array}{c}{\mathrm{https://github.com/wiseodd/generative-models}}^{}\end{array}$$

References

• Doersch, Carl. “Tutorial on variational autoencoders.” arXiv preprint arXiv:1606.05908 (2016).
• Kingma, Diederik P., and Max Welling. “Auto-encoding variational bayes.” arXiv preprint arXiv:1312.6114 (2013).

$$\begin{array}{c}{\mathrm{https://blog.keras.io/building-autoencoders-in-keras.html}}^{}\end{array}$$