Improving a Variational Autoencoder with Normalizing Flows

In order to fully grasp the concepts explained here, I strongly recommend you to read my three posts on Variational Autoencoders (in the following order)

• Variational Autoencoders and the Expectation Maximization Algorithm
• Minimalist Variational Autoencoder in Pytorch with CUDA GPU
• Assessing a Variational Autoencoder on MNIST using Pytorch.

Theory of Vanilla VAEs

Recall that in a Vanilla VAE we feed $\mathbf{x}$ into an encoder neural network and obtain $(\mathit{\mu },\log \mathit{\sigma })$. These are the amortized parameters of our approximate posterior distribution

$$q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})=\mathcal{N}(\mathbf{z}\mid {\mathit{\mu }}_{\mathit{\varphi }}(\mathbf{x}),\text{diag}({\mathit{\sigma }}_{\mathit{\varphi }}^2(\mathbf{x})))$$

To get a latent sample $\mathbf{z}\sim q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})$ we need to use the reparametrization trick. This requires sampling $\mathit{ϵ}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$ and then scaling it and shifting it according to the output of the neural network

$$\mathbf{z}={\mathit{\mu }}_{\mathit{\varphi }}(\mathbf{x})+{\mathit{\sigma }}_{\mathit{\varphi }}^2(\mathbf{x})\odot \mathit{ϵ}.$$

To learn the parameters of our neural network our aim is to maximize the ELBO

$${\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})={\mathbb{E}}_{q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})}[\log p_{\mathit{\theta }}(\mathbf{x}\mid \mathbf{z})]-\text{KL}(q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})\parallel p(\mathbf{z}))$$

The reconstruction error (the first term) is easy to compute in the Normal and Bernoulli case. In what follows, we will assume that the likelihood is a product of Bernoullis. This is the usual set-up when working with MNIST. The likelihood is then

$$p_{\mathit{\theta }}(\mathbf{x}\mid \mathbf{z})=\prod _{i=1}^{\text{dim}(x)}{p}_i^{x_i}(1-p_i{)}^{1-x_i}$$

where $\mathbf{p}=(p_1,\dots ,p_{\text{dim}(x)}{)}^{\mathrm{\top }}$ is the output of the decoder network, and $\mathbf{p}\in [0,1{]}^{\text{dim}(x)}$. The reconstruction error can then be written as

$$\begin{array}{rl}{\mathbb{E}}_{q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})}[\log p_{\mathit{\theta }}(\mathbf{x}\mid \mathbf{z})]& ={\mathbb{E}}_{q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})}[\log \prod _{i=1}^{\text{dim}(\mathbf{x})}{p}_i(\mathbf{z}{)}^{x_i}(1-p_i(\mathbf{z}){)}^{1-x_i}]\\ & ={\mathbb{E}}_{q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})}[\sum _{i=1}^{\text{dim}(\mathbf{x})}{x}_i\log p_i(\mathbf{z})+(1-x_i)\log (1-p_i(\mathbf{z}))]\\ & \approx \sum _{j=1}^{n_z}\sum _{i=1}^{\text{dim}(\mathbf{x})}{x}_i\log p_i(\mathbf{z})+(1-x_i)\log (1-p_i(\mathbf{z})){\mathbf{z}}^{(j)}\sim q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})\end{array}$$

where $n_{\mathbf{z}}$ is the number of samples that we sample from $q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})$. Usually we simply set $n_{\mathbf{z}}=1$. This means we only sample one latent variable per datapoint. The objective function to minimize (I have flipped the sign) is therefore

$$\begin{array}{rl}-{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& =-\sum _{i=1}^{\text{dim}(\mathbf{x})}{x}_i\log p_i(\mathbf{z})+(1-x_i)\log (1-p_i(\mathbf{z}))-\frac{1}{2}\sum _{j=1}^{\text{dim}(\mathbf{z})}(1+\log {\sigma }_j^2-{\mu }_j^2-{\sigma }_j^2)\\ & =\text{BCE}(\mathbf{p},\mathbf{x})-\frac{1}{2}\sum _{j=1}^{\text{dim}(\mathbf{z})}(1+\log {\sigma }_j^2-{\mu }_j^2-{\sigma }_j^2)\end{array}$$

Using Pytorch we can code it as

def vae_loss(image, reconstruction, mu, logvar):
  """Loss for the Variational AutoEncoder."""
  # Compute the binary_crossentropy.
  recon_loss = F.binary_cross_entropy(
      input=reconstruction.view(-1, 28*28),    # input is p(z) (the mean reconstruction)
      target=image.view(-1, 28*28),            # target is x   (the true image)
      reduction='sum'                          
  )
  # Compute KL divergence using formula (closed-form)
  kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
  return reconstruction_loss + kl

VAE with Normalizing Flows

This time, we not only want our encoder to output $(\mathit{\mu },\log \mathit{\sigma })$ to shift and scale $\mathit{ϵ}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$. We also want to feed

$$\mathcal{N}(\mathbf{z}\mid {\mathit{\mu }}_{\mathit{\varphi }}(\mathbf{x}),\text{diag}({\mathit{\sigma }}_{\mathit{\varphi }}^2(\mathbf{x})))$$

through a series of $K$ transformations each one of them depending on a set of parameters ${\mathit{\lambda }}_k$. Denoting $\mathit{\lambda }=({\mathit{\lambda }}_1,\dots ,{\mathit{\lambda }}_K)$ we essentially want our Encoder to work as follows:

$$\mathbf{x}⟶\text{Encoder}⟶(\mathit{\mu },\log \mathit{\sigma },{\mathit{\lambda }}_1,\dots ,{\mathit{\lambda }}_K)=(\mathit{\mu },\log \mathit{\sigma },\mathit{\lambda })$$

Then we would firstly use $(\mathit{\mu },\log \mathit{\sigma })$ to compute ${\mathbf{z}}_0$ using the reparametrization trick $${\mathbf{z}}_0=\mathit{\mu }+\mathit{\sigma }\odot \mathit{ϵ}\mathit{ϵ}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$$

and finally we would feed ${\mathbf{z}}_0$ into the series of transformations to reach the final ${\mathbf{z}}_K$

$${\mathbf{z}}_K=f_K\circ f_{K-1}\circ \dots \circ f_2\circ f_1({\mathbf{z}}_0).$$

This means that our approximating distribution is not

$$q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})=\mathcal{N}(\mathbf{z}\mid {\mathit{\mu }}_{\mathit{\varphi }}(\mathbf{x}),\text{diag}({\mathit{\sigma }}_{\mathit{\varphi }}^2(\mathbf{x})))$$

anymore but, rather, it can be found using the usual change of variable formula

$$\log q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})=\log q_K({\mathbf{z}}_K)=\log q_0({\mathbf{z}}_0)-\sum _{k=1}^K\ln \left|\text{det}\frac{\partial f_k}{\partial {\mathbf{z}}_{k-1}}\right|$$

where the base distribution $q_0({\mathbf{z}}_0)$ is the old approximate posterior distribution $$q_0({\mathbf{z}}_0)=\mathcal{N}({\mathbf{z}}_0\mid {\mathit{\mu }}_{\mathit{\varphi }}(\mathbf{x}),\text{diag}({\mathit{\sigma }}_{\mathit{\varphi }}^2(\mathbf{x}))).$$

Thanks to the law of the unconscious statistician we have

$$\begin{array}{rl}{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& ={\mathbb{E}}_{q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})}[\log p_{\mathit{\theta }}(\mathbf{x}\mid \mathbf{z})]-\text{KL}(q_{\mathit{\varphi }}(\mathbf{z}\mid \mathbf{x})\parallel p(\mathbf{z}))\\ & ={\mathbb{E}}_{q_K({\mathbf{z}}_K)}[\log p_{\mathit{\theta }}(\mathbf{x}\mid {\mathbf{z}}_K)]-{\mathbb{E}}_{q_K({\mathbf{z}}_K)}[\log q_K({\mathbf{z}}_K)-\log p({\mathbf{z}}_K)]\\ & ={\mathbb{E}}_{q_0({\mathbf{z}}_0)}[\log p_{\mathit{\theta }}(\mathbf{x}\mid {\mathbf{z}}_K)]-{\mathbb{E}}_{q_0({\mathbf{z}}_0)}[\log q_K({\mathbf{z}}_K)-\log p({\mathbf{z}}_K)]\end{array}$$

As usual, we can approximate this using Monte Carlo and generally we only need one sample. By drawing ${\mathbf{z}}_0\sim q_0({\mathbf{z}}_0)=\mathcal{N}(\mathit{\mu },\text{diag}(\mathit{\sigma }))$ we can approximate the ELBO as follows

$$\begin{array}{rl}{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& \approx [\sum _{i=1}^{\text{dim}(\mathbf{x})}{x}_i\log p_i({\mathbf{z}}_K)+(1-x_i)\log (1-p_i({\mathbf{z}}_K))]-\log q_K({\mathbf{z}}_K)+\log p({\mathbf{z}}_K).\end{array}$$

This means that our objective function is given by $$\begin{array}{rl}-{\mathcal{L}}_{\mathit{\varphi },\mathit{\theta }}(\mathbf{x})& =\text{BCE}(\mathbf{p},\mathbf{x})+\log q_0({\mathbf{z}}_0)+\text{LADJ}-\log p({\mathbf{z}}_K)\end{array}$$

where the Log-Absolute-Determinant-Jacobian is the usual $$\text{LADJ}=\sum _{k=1}^K\ln \left|\text{det}\frac{\partial f_k}{\partial {\mathbf{z}}_{k-1}}\right|$$