Generative Modeling with Variational Autoencoders

Till now, I have talked about supervised learning where our model is trained to learn a mapping function of data $\mathbf{x}$ to predict labels $\mathbf{y}$, for example, the tasks of classification, object detection, segmentation, image captioning, etc. However, supervised learning requires large datasets that are created with human annotations to train the models.

The other side of machine learning is called unsupervised learning where we just have data $\mathbf{x}$, and the goal is to learn some hidden underlying structure of the data using a model. Through this approach, we aim to learn the distribution of the data $p(\mathbf{x})$ and then sample from it to generate new data, in our case images.


Autoencoders

An autoencoder is a bottleneck structure that extracts features $\mathbf{z}$ from the input images $\mathbf{x}$ and uses them to reconstruct the original data $\mathbf{x}'$, learning an identity mapping.

  • Encoder: network that compresses the input into a latent-space representation (of a lower dimension)
  • Decoder: network that reconstructs the output from this representation

Since we want the autoencoder to learn to reconstruct the input data, the loss objective is a simple L2-loss between the input and the reconstructed output $\mathbf{x}' = f_\theta(g_\phi(\mathbf{x}))$. \begin{equation*} L_\text{AE}(\theta, \phi) = || \mathbf{x} - f_\theta(g_\phi(\mathbf{x})) ||_2^2 \tag{1} \end{equation*}


This approach is used for learning a lower-dimensional feature representation from the training data that captures meaningful factors of variation in it. After training, we get rid of the decoder and use the encoder to obtain features that contain useful information about the input that can be used for downstream tasks.

Autoencoders are generally used for dimensionality reduction (or data compression) and to initialize a supervised model. Since it is not probabilistic, we cannot sample new data from the learned model, and new images cannot be generated$^*$.



Variantional Autoencoders (VAEs)

VAEs, introduced in Auto-Encoding Variational Bayes paper are a probabilistic version of autoencoder that lets you sample from the model to generate new data. Instead of mapping the input into a fixed vector, we want to map it into a distribution, denoted as $p_\theta$.

  • Prior $p_\theta (\mathbf{z})$: Prob distribution over the latent variables (fixed and often assumed to be a simple Gaussian)
  • Posterior $p_\theta (\mathbf{x}|\mathbf{z})$: Prob distribution over images conditioned on the latent variables
  • Likelihood $p_\theta (\mathbf{z}|\mathbf{x})$: Prob distribution over latent variables given input data

Assuming that we know the real parameter $\theta^*$ for this distribution, we can generate new samples for a real data point $\mathbf{x}^{(i)}$ as,

  1. First, sample a $\mathbf{z}^{(i)}$ from a true prior distribution $p_{\theta^*}(\mathbf{z})$.
  2. Feed it into the decoder model to obtain the posterior $p_{\theta^*}(\mathbf{x} \vert \mathbf{z} = \mathbf{z}^{(i)})$.
  3. Then sample a new data point $\mathbf{x'}^{(i)}$ from this true posterior.

As you would be wondering a neural network cannot output a probability distribution so, instead the decoder predicts the mean $\mu_{\mathbf{x} \vert \mathbf{z}}$ and diagonal variance $\Sigma_{\mathbf{x} \vert \mathbf{z}}$ of the posterior, assuming it to be also Gaussian such that $p_\theta (\mathbf{x}|\mathbf{z}) = \mathcal{N}(\mu_{\mathbf{x} \vert \mathbf{z}}, \Sigma_{\mathbf{x} \vert \mathbf{z}})$.

There is an inherent assumption here that the pixels of the generated image are conditionally independent given the latent variable and this is the reason why VAEs tend to generate blurry images.


In order to estimate the true parameter $\theta^*$, we learn the model parameters that maximize the probability of generating real data samples (maximum likelihood estimation), \begin{align*} \theta^{*} &= \arg\max_\theta \prod_{i=1}^n p_\theta(\mathbf{x}^{(i)}) \\ &= \arg\max_\theta \sum_{i=1}^n \log p_\theta(\mathbf{x}^{(i)}) \tag{2} \end{align*}

Before moving further, let's see how we can obtain the distribution of the data $p_\theta(\mathbf{x})$.

  1. Marginalize out the latent variable to simplify as, \begin{align*} p_\theta(\mathbf{x}) &= \int p_\theta(\mathbf{x}, \mathbf{z}) d\mathbf{z} = \int \underbrace{p_\theta(\mathbf{x}\vert\mathbf{z})}_{\text{posterior}} \; \underbrace{p_\theta(\mathbf{z})}_{\text{prior}} d\mathbf{z} \\ \end{align*}

    The terms in the integral - prior (fixed) and posterior (estimated by the decoder) are straightforward to compute. However, it is impossible to integrate over all $\mathbf{z}$ as it is expensive to check all the possible values of $\mathbf{z}$ and sum them up.

  2. Another idea is to use Baye's rule, \begin{align*} p_\theta(\mathbf{x}) = \frac{p_\theta(\mathbf{x}\vert\mathbf{z}) \; p_\theta(\mathbf{z})}{p_\theta(\mathbf{z}\vert\mathbf{x})} \approx \frac{p_\theta(\mathbf{x}\vert\mathbf{z}) \; p_\theta(\mathbf{z})}{q_\phi(\mathbf{z}\vert\mathbf{x})} \end{align*}

    The numerator has the same terms as above. The denominator has the likelihood term that can be approximated by another neural network (encoder) as $q_\phi(\mathbf{z}\vert\mathbf{x}) \approx p_\theta (\mathbf{z}|\mathbf{x})$.

    Again, we assume that is follows a Gaussian distribution as $q_\phi (\mathbf{z}\vert\mathbf{x}) = \mathcal{N}(\mu_{\mathbf{z} \vert \mathbf{x}}, \Sigma_{\mathbf{z} \vert \mathbf{x}})$, and the encoder predicts the mean $\mu_{\mathbf{z} \vert \mathbf{x}}$ and diagonal variance $\Sigma_{\mathbf{z} \vert \mathbf{x}}$.


The encoder, therefore, outputs a distribution of latent variables $\mathbf{z}$ given the input data $\mathbf{x}$. The training process looks like,

  1. Input a real data point $\mathbf{x}^{(i)}$ to the encoder to obtain the likelihood $q_\phi (\mathbf{z}\vert\mathbf{x} = \mathbf{x}^{(i)})$.
  2. Sample a latent variable from this likelihood $\mathbf{z}^{(i)} \sim \mathcal{N}(\mu_{\mathbf{z} \vert \mathbf{x}}, \Sigma_{\mathbf{z} \vert \mathbf{x}})$
  3. Feed it into the decoder to obtain the posterior $p_{\theta^*}(\mathbf{x} \vert \mathbf{z} = \mathbf{z}^{(i)})$.
  4. Then sample a new data point from this posterior $\mathbf{x'}^{(i)} \sim \mathcal{N}(\mu_{\mathbf{x} \vert \mathbf{z}}, \Sigma_{\mathbf{x} \vert \mathbf{z}})$


Sampling is a stochastic process and therefore we cannot backpropagate the gradient. To make it trainable, the reparameterization trick is introduced such that, \begin{align*} \mathbf{z} &\sim q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)}) = \mathcal{N}(\mathbf{z}; \boldsymbol{\mu}^{(i)}, \boldsymbol{\sigma}^{2(i)}\boldsymbol{I})\\ \mathbf{z} &= \boldsymbol{\mu} + \boldsymbol{\sigma} \odot \boldsymbol{\epsilon} \text{, where } \boldsymbol{\epsilon} \sim \mathcal{N}(0, \boldsymbol{I}) \end{align*} where $\odot$ refers to an element-wise product.



Deriving ELBO Loss

In order to obtain the loss objective of a VAE, we further simpify Equation 2 by first multipying $1 = \int q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)}) \; dz$ with $\log p_\theta(\mathbf{x}^{(i)})$ and then solving, \begin{align*} \log p_\theta(\mathbf{x}^{(i)}) &= \log p_\theta(\mathbf{x}^{(i)}) \int q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)}) \; dz = \int q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)}) \log p_\theta(\mathbf{x}^{(i)}) \; dz \\ &=\mathbf{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})} [\log p_\theta(\mathbf{x}^{(i)})] \\ \\ &\text{applying Baye's rule,}\\ &= \mathbf{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})} \left[ \log \frac{p_\theta(\mathbf{x}^{(i)}\vert\mathbf{z}) \; p_\theta(\mathbf{z})}{p_\theta(\mathbf{z}\vert\mathbf{x}^{(i)})} \right] \\ \\ &\text{multipy & divide by same term}\\ &= \mathbf{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})} \left[ \log \frac{p_\theta(\mathbf{x}^{(i)}\vert\mathbf{z}) \; p_\theta(\mathbf{z})}{p_\theta(\mathbf{z}\vert\mathbf{x}^{(i)})} \frac{q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})}{q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})} \right] \\ \\ &\text{rearranging}\\ &= \mathbf{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})} \left[ \log p_\theta(\mathbf{x}^{(i)}\vert\mathbf{z}) \right] + \mathbf{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})} \left[ \frac{p_\theta(\mathbf{z})}{q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})} \right] + \mathbf{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})} \left[ \frac{q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})}{p_\theta(\mathbf{z}\vert\mathbf{x}^{(i)})} \right] \\ &= \mathbf{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})} \left[ \log p_\theta(\mathbf{x}^{(i)}\vert\mathbf{z}) \right] - \mathbf{D}_{KL}\left( q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)}) \; || \; p_\theta(\mathbf{z}) \right) + \underbrace{\mathbf{D}_{KL} \left( q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)}) \; || \; p_\theta(\mathbf{z}\vert\mathbf{x}^{(i)}) \right)}_{\text{cannot be computed but KL-Divergence} \ge 0 \text{ always}} \\ &\ge \underbrace{\mathbf{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})} \left[ \log p_\theta(\mathbf{x}^{(i)}\vert\mathbf{z}) \right]}_{\text{reconstruction term}} - \underbrace{\mathbf{D}_{KL} \left( q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)}) \; || \; p_\theta(\mathbf{z}) \right)}_{\text{prior matching term}} \end{align*}


The first term is the data reconstruction term that measures the reconstruction likelihood of the decoder from our variational distribution. It ensures that the learned distribution is modeling effective latents that the original data can be regenerated from.

The second term measures how similar the learned variational distribution is to a prior belief held over latent variables. Since both the distributions are Gaussians, KL-divergence can be computer in closed form.

This gives us a lower bound on the log-likelihood of the data and we jointly train both the encoder and decoder to maximize it. \begin{align*} \theta^{*} &= \arg\max_\theta \sum_{i=1}^n \log p_\theta(\mathbf{x}^{(i)}) \\ &\ge \arg\max_{\theta, \phi} \sum_{i=1}^n \mathbf{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)})} \left[ \log p_\theta(\mathbf{x}^{(i)}\vert\mathbf{z}) \right]- \mathbf{D}_{KL}(q_\phi(\mathbf{z}\vert\mathbf{x}^{(i)}) \; || \; p_\theta(\mathbf{z})) \end{align*} It is known as the variational lower bound, or evidence lower bound (ELBO).

The loss function (that we minimize) is given as, \begin{equation*} L_{\text{VAE}}(\theta, \phi) = - \mathbf{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x})} \left[ \log p_\theta(\mathbf{x}\vert\mathbf{z}) \right] + \mathbf{D}_{KL}\left( q_\phi(\mathbf{z}\vert\mathbf{x}) \; || \; p_\theta(\mathbf{z}) \right) \tag{3} \end{equation*}




$*$ - I have tried generating new images using an autoencoder by first obtaining the latent representation of two images and then performing the reconstruction on the linear interpolation between them. Below are the results on the CIFAR10 dataset, with the first two columns showing original images and the last one showing the new image.

Comments

Post a Comment

Popular posts from this blog

The move_base ROS node

Image
move_base package provides the move_base node which is a major element of ROS Navigation stack. It moves the robot from its current position to a goal position. When path planning is done, the user gives uses the 2D Nav Goal in RVIZ to specify the goal position and orientation. This position (x,y,z) and orientation(x,y,z,w) are actually published in move_base/goal topic. The move_base node links together a global and local planner to accomplish its global navigation task. The Global Planner It is in charge of calculating the safe path in order to arrive at goal pose. Note that the path is calculated before the robot starts moving so it will not take into account the laser readings while moving i.e. unscanned obstacles that are not in the map. Navfn Planner The Navfn planner is the most commonly used global planner for ROS Navigation. It uses Dijkstra's algorithm in order to calculate the shortest path between the initial pose and the goal pose. Carrot Planner This plann...
Read more

Three Wheeled Omnidirectional Robot : Motion Analysis

Image
A Three-wheeled omnidirectional robot is capable to moving in any given direction without rotating by giving velocities to the three wheels in a certain ratio. All motors are situated at an angle of 120 degrees from each other. f B ,  f L ,  f R  are for back , left and right wheel respectively. a x  and a y  are accelerations in X R  and Y R  axis in Robot frame and w is the angular velocity in Robot frame. Hence putting the values of alpha and making it into a matrix form we get : As we need to obtain the f values, we will invert the matrix. We need  a x  , a y  and w is the world frame, so we apply transformation. The transformation angle is    α 1  = 30 degrees.  Putting this in the previous equation, Multiplying the two matrices, we get the final relation as : For example, if you want to move the robot along X axis(world frame), put  a x W  = 1 , ...
Read more

Overview of ATmega328P

Image
While reading about the ATmega328, I came a across a lot of stuff which I would like to share here. There is a humongous amount of information available in its datasheet. Introduction ATmega328 is a single chip micro-controller created by Atmel. It has a In-System Self-programmable Flash memory of 32 KB, SRAM of 2 KB (2048 bytes) and EEPROM of 1 KB (1024 bytes). It has a 8-bit AVR CPU type which means they can process data in 8-bit chunk as its data bus is 8-bit wide.   It contains 23 general purpose I/O pins and 32 general purpose working registers. It has 32 single-byte registers and 8-bit RISC (reduced instruction set computer) devices and hence the name 328. The P is for ' picoPower '. It has 3 flexible timer/counter with compare modes, internal and external interupts , 6-channel 10-bit A/D converter , programmable watch dog timer with internal oscillator and 5 software selectable power saving modes. Digital communication peripherals include 1-USART, 2-SPI and 1-...
Read more