Variational Autoencoders

Tldr

Variational Autoencoders (VAEs) are generative models that learn to encode data into a latent space and then decode it back to the original space. They use a probabilistic approach to model the data distribution and can generate new samples from the learned distribution.

Motivation

• VAEs are designed to overcome the limitations of traditional autoencoders, specifically the lack of clear generative mechanism (their latent spaces can be irregular)
• VAEs add probabilistic constraints to the latent space, allowing for better generalization and sample generation

• Specifically, VAEs learn a distribution over the latent space rather than a single point, enabling them to generate new samples by sampling from this distribution.

Variational Inference:

• Data points $x$ are assumed to be generated by an unkown random latent variable $z$.
• Goal: learn the distribution of latent variables that explains the data $p(z|x)$.
• Directly computing $p(z|x)$ is intractable, so we use variational inference to approximate it with a simpler distribution $q(z|x)$, known as the approximate posterior.

Loss Function: $\mathcal{L}(\theta ,\varphi ;x)=-{\mathbb{E}}_{q_{\varphi }(z|x)}[\log p_{\theta }(x|z)]+D_{KL}(q_{\varphi }(z|x)||p(z))$

Where:

• $\theta$ are the parameters of the decoder (generative model).
• $\varphi$ are the parameters of the encoder (inference model).
• $p_{\theta }(x|z)$ is the likelihood of the data given the latent variable (reconstruction term).
• $D_{KL}$ is the Kullback-Leibler divergence, which measures how much the approximate posterior $q_{\varphi }(z|x)$ diverges from the prior distribution $p(z)$ (regularization term).

Reparameterization Trick:

• To backpropagate through the stochastic sampling of $z$, we use the reparameterization trick, which allows us to express the sampling operation in a differentiable way.
• Instead of sampling $z$ directly from $q(z|x)$, we sample $ϵ$ from a standard normal distribution and then transform it using the parameters of $q(z|x)$: $z=\mu +\sigma \odot ϵ$ Where $\mu$ and $\sigma$ are the mean and standard deviation of the approximate posterior, and $\odot$ denotes element-wise multiplication. This allows gradients to flow through the sampling process during training.

Applications of Smooth latent spaces

• Generation: pick any random point from a Unit Gaussian, feed it through the decoder, and you get a new sample that resembles the training data.
• Interpolation: you can take two points in the latent space (e.g., corresponding to two different images), and interpolate between them to create a smooth transition in the output space (e.g., morphing one image into another).
• Chemical Design: VAEs can be used to generate new molecular structures by learning a latent representation of existing molecules and sampling from that space to create novel compounds with desired properties.
• Anomaly Detection: By learning the distribution of normal data, VAEs can identify anomalies by measuring how well new data points fit into the learned distribution. Data points that have low likelihood under the model can be flagged as anomalies.

import torch
import torch.nn as nn
class VAE(nn.Module):
    def __init__(self, input_dim, latent_dim):
        super(VAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, latent_dim * 2)  # Output mean and log-variance
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, input_dim),
            nn.Sigmoid()  # Assuming input is normalized between 0 and 1
        )
 
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
 
    def forward(self, x):
        h = self.encoder(x)
        mu, logvar = h.chunk(2, dim=-1)  # Split into mean and log-variance
        z = self.reparameterize(mu, logvar)
        return self.decoder(z), mu, logvar