Variational Families

Introducing The Gists of Variational Autoencoders (VAEs)

What is a Variational Autoencoder (VAE)?

Imagine you have a magical machine that can take a picture of your favorite toy, turn it into a secret code, and then use that code to recreate the toy’s picture. This is kind of what a Variational Autoencoder (VAE) does, but instead of toys, it works with things like images, sounds, or even words.

A VAE is a type of artificial brain (or neural network) that learns to compress data (like a picture) into a simpler form and then uses that simple form to recreate the original data. The “variational” part means that it doesn’t just learn one way to represent the data, but many possible ways, which helps it be more flexible and creative.

Math Foundations: Breaking it Down

Let’s keep things simple. Imagine you have a bunch of colored balls, and you want to sort them by color. First, you need to pick a way to describe each color with a number. This number is like the “code” that the VAE creates. Then, once you have the code, you need to figure out how to recreate the exact color from the number.

1. Encoder: This part of the VAE is like a friend who looks at the color of the ball and writes down a secret code (a number) that represents that color.

2. Latent Space: This is like a secret storage area where all the codes are kept. It’s like a map where similar colors are stored close to each other.

3. Decoder: This is like another friend who can look at the code and recreate the exact color of the ball.

The tricky part is that the VAE learns to write codes in a way that they can be easily decoded back into the original color.

Simple Math Example:

If your color is described by a number, say $ x $, the VAE learns a new number, $ z $, that is simpler and can still be used to recreate $ x $.

The VAE tries to make sure that:

$$
\text{Original Data} \approx \text{Decoder}(\text{Encoder}(\text{Original Data}))
$$

This means the data recreated by the VAE should be very close to the original data.

Coding Example: A Simple VAE in Python

Let’s see a simple example using Python. We’ll use a small dataset of handwritten digits called MNIST. We want to teach our VAE to encode these digits into a simpler form and then recreate them.

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# Define the encoder
class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.fc1 = nn.Linear(784, 400)
        self.fc2_mu = nn.Linear(400, 20)
        self.fc2_logvar = nn.Linear(400, 20)

    def forward(self, x):
        h = torch.relu(self.fc1(x))
        return self.fc2_mu(h), self.fc2_logvar(h)

# Define the decoder
class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.fc3 = nn.Linear(20, 400)
        self.fc4 = nn.Linear(400, 784)

    def forward(self, z):
        h = torch.relu(self.fc3(z))
        return torch.sigmoid(self.fc4(h))

# VAE Model
class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        mu, logvar = self.encoder(x.view(-1, 784))
        z = self.reparameterize(mu, logvar)
        return self.decoder(z), mu, logvar

# Loss function (how far off we are from the original)
def loss_function(recon_x, x, mu, logvar):
    BCE = nn.functional.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

# Training the VAE
def train(model, optimizer, train_loader):
    model.train()
    train_loss = 0
    for batch_idx, (data, _) in enumerate(train_loader):
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        loss.backward()
        train_loss += loss.item()
        optimizer.step()
    return train_loss / len(train_loader.dataset)

# Load dataset
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('.', train=True, download=True, transform=transforms.ToTensor()),
    batch_size=128, shuffle=True)

vae = VAE()
optimizer = optim.Adam(vae.parameters(), lr=1e-3)

for epoch in range(1, 11):
    train_loss = train(vae, optimizer, train_loader)
    print(f'Epoch {epoch}, Loss: {train_loss:.4f}')

In this code:

• Encoder: Learns to compress the digit images into a simpler form.
• Decoder: Learns to recreate the digit images from this simpler form.
• Reparameterization: Adds a bit of randomness, helping the model learn multiple ways to recreate the images.

Practical Example: How VAEs Can Help

Imagine you’re building a tool to generate new images of cartoon characters. You don’t want just one version of a character, but many different versions. The VAE can learn from existing images and then create new, similar images by tweaking the codes stored in the latent space.

For example, if you’ve trained a VAE on images of cats, you can generate new, never-before-seen images of cats by sampling from the latent space.

More on The Stochasticity of VAE Methods

Variational Autoencoders (VAEs) are inherently stochastic in nature. This stochasticity is a key feature that differentiates VAEs from traditional deterministic autoencoders.

Here’s how the stochastic nature of VAEs works:

1. Latent Space Representation:

In a traditional autoencoder, the encoder deterministically maps an input $ x $ to a single point in a latent space $ z $. The decoder then maps this point $ z $ back to reconstruct the input $ x $.

In contrast, in a VAE, the encoder doesn’t map $ x $ to a single point. Instead, it maps $ x $ to a probability distribution over the latent space. Typically, this is modeled as a multivariate Gaussian distribution with a mean vector $ \mu $ and a variance (or log-variance) vector $ \sigma^2 $.

2. Sampling from the Latent Space:

To generate a latent variable $ z $ that will be passed to the decoder, VAEs sample from the distribution defined by $ \mu $ and $ \sigma^2 $. This sampling process introduces stochasticity into the VAE because even for the same input $ x $, different samples from the distribution will produce slightly different $ z $ values, leading to variations in the output.

3. Stochasticity and Learning:

The stochastic nature of VAEs is crucial for the model’s ability to generate diverse outputs and to learn a well-distributed latent space. During training, the VAE learns to shape these distributions in the latent space so that they can generate realistic outputs even when new $ z $ values are sampled.

4. Reparameterization Trick:

The reparameterization trick is a method used in VAEs to make the stochastic sampling differentiable, which is necessary for backpropagation. It involves expressing the sampling process as $ z = \mu + \sigma \times \epsilon $, where $ \epsilon $ is a random variable drawn from a standard normal distribution. This trick enables the VAE to learn $ \mu $ and $ \sigma $ while still incorporating stochasticity.

In Summary:

VAEs are stochastic because they incorporate random sampling within their latent space representation. This stochasticity is essential for the model’s ability to generate diverse and meaningful data from the learned latent space. The randomness allows the VAE to explore different potential reconstructions, making it a powerful generative model.

A Variational Autoencoder is like a creative artist that learns to simplify complex things into a simpler form (like codes) and then use those codes to recreate or even generate new things. It’s powerful because it can learn many ways to represent the data, making it flexible and creative. By understanding the basic math, seeing some code, and applying it practically, you get a glimpse of how VAEs are helping machines to learn and create in innovative ways!