Variational Autoencoder (VAE)

The difference of VAE and AE is mainly in how we perform the bottleneck. But first, we need to go through why we even need variational autoencoder in the first place.

Problem with generic Autoencoders

• in the latent space, every embedded vector is a point in this space (can take up any point from 0 to infinite)
• the latent space is infinite
• because its a point, there are many empty spaces that do not contain points
• in other words, the latent space is made up of discrete points and also many dead zones
• this existence of dead zones occur because the latent space is infinite which is an issue
• because lets say we wanna use the decoder for generating images
• we will first use a random noise generator to create a full image of random noise
• this random noise's embedded vector maps to a totally random discrete point in the latent space that doesn't even come close to any embedded training vectors
• when that happens, the decoder will be unable to generate a proper image, and will just generate gibberish

How VAE Helps

• VAE models the latent space as a normal distribution, with mean of 0 and variance of 1
• The design behind this is the Kullback-Liebler Divergence (KL Divergence) loss
• the encoder in VAE will produce 2 embedded vectors based on trained weights
  1. Mean vector ($\mathbf{\mu }$)
  2. Log-Variance vector ($\log {\mathbf{\sigma }}^2$): we use log variance here for stability
• Log-Variance vector is first used to calculate for $\mathbf{\sigma }$, standard deviation. $\mathbf{\sigma }$ and $\mathbf{\mu }$ are added together, with an error value $\mathbf{ϵ}$, that is normally distributed, element-wise multiplied $\odot$, to $\mathbf{\sigma }$

$$\mathbf{z}=\mathbf{\mu }+\mathbf{\sigma }\odot \mathbf{ϵ}$$

• the above formula shows how the final z vector (the one in latent space) is calculated
• all the training z vectors are all close to 0 because of the KL divergence (since mean is at 0)

Training VAE

Here are the 4 objects in the model.

1. the data $x$

   • this is 1 single image of size $28$ by $28$

2. The latent Variable $z$

   • this is a short vector of lets say 20 numbers, it is also known as a vector embedding in the latent space, or latent embedding, or latent vector
   • this latent vector acts as a compressed vector that tries to contain the original meaning of the input image and represent in in the latent space
   • this is invented by the model based on its weights

3. The Encoder $q_{\varphi }(z|x)$

   • this is a network with weights $\varphi$
   • input into this network is our original image
   • output are 2 vector embedding, 1 for the mean and 1 for the variance
   • this gives us a gaussian distribution in the latent space

4. The Decoder $p_{\theta }(x|z)$

   • this is a network with weights $\theta$
   • input into this network is the latent embedding
   • output is a single vector of length $784$, which is essentially a flattened $28$ by $28$ image
   • each value in the vector is a distribution parameter, which is the probability that the pixel is on, between $0$ to $1$

The goal of training is find values of $\varphi$ and $\theta$ that make the whole system reconstruct data well and keep the latent space tidy.

This formula is the Loss Function of point $x_i$:

$$\mathcal{L}(\theta ,\varphi ;x_i)=\underset{\text{term 1: reconstruction}}{\underbrace{-{\mathbb{E}}_{z\sim q_{\varphi }(z\mid x_i)}\left[\log p_{\theta }(x_i\mid z)\right]}}+\underset{\text{term 2: regularizer}}{\underbrace{\mathrm{KL}\left(q_{\varphi }(z\mid x_i)\Vert p_{\theta }(z)\right)}}$$

where:

• $\mathcal{L}$ is the loss, which is a single number for one datapoint, and AI/Generative Models/Autoencoders/Gradient Descent makes it smaller
• $\mathcal{L}(\theta ,\varphi ;x_i)$ means that we are given our datapoint $x_i$ and asked to optimize the two weights $\theta$ and $\varphi$
• $x_i$ is the training image number $i$
• $\mathbb{E}$ is the expected value, which is an average
• ${\mathbb{E}}_{z\sim q_{\varphi }(z\mid x_i)}\left[\log p_{\theta }(x_i\mid z)\right]$ just means we are averaging all the values of $log{p}_{\theta }(x_i|z)$ using every possible $z$ in the cloud that the encoder cloud generates for image $x_i$
• Term 2 is the KL regularizer
• it measured how far this image's cloud sits from the fixed prior cloud and it returns a single positive number
• divergence is a distance-like score between 2 distributions
• it is always more than 0 and = 0 only when the two distributions are identical

So basically, term 1 gives every image a latent embedding (a single point in latent space) such that the decoder knows how to decode it back to the original image. Term 2 distributes all the latent embeddings normally (normal distribution), and gives each latent embedding a distribution rather than 1 single point.

Note

• a VAE has a probabilistic encoder and a decoder
• a plain autoencoder is a poor generator because its latent space may not be continuous, so sampling can land in a discontinuity the decoder never saw
• in a VAE the encoding of each sample forms a distribution rather than a single point
• the first term of the loss is the reconstruction loss
• the second term is a regularizer that keeps the approximate posterior close to the prior
• for the prior over the latent variable we simply use N(0,1)
• what the probabilistic encoder buys is a continuous latent space, which allows easy sampling and interpolation
• a VAE is not trained adversarially, so no discriminator is involved

The objective behind these two terms is derived in Evidence Lower Bound (ELBO).