Variational Autoencoders

A space you can sample.

⏱ 2 min read▶ Video at 1:17:58🎓 Free

Watch this lesson in the Deep Learning courseWatch on YouTube ↗Full Deep Learning course

If an autoencoder's decoder can turn a code into an image, then why not feed it a random code, and get a new image? Try it, and you get garbage. The codes are scattered, with big empty holes between them. A variational autoencoder fixes this, with one beautiful twist.

From the video: Variational Autoencoders

Random code, new image?
You get garbage
The space has holes

A smooth map of possibilities

Instead of mapping each image to one exact point, a V A E maps it to a small cloud. Then, a second loss gently pulls all of those clouds towards a standard bell curve. So the whole space becomes smooth, with no holes. Any random point now decodes into something sensible. And moving between two points smoothly morphs one digit into another.

Encode to a clouda mean and a spread

Sample from itadds randomness

KL losskeeps clouds near a bell curve

Smooth spaceany point decodes well

Dreaming new digits

Here's a V A E. The encoder gives two things, a mean and a spread. The line with torch randn like is the reparameterisation trick, which lets gradients flow through the sampling step. The loss has two parts, one for rebuilding, and one that's called K L, for keeping the clouds in shape. After training, we throw away the encoder. We feed pure random numbers into the decoder. Run it. Those digits don't exist. No one drew them. The network has imagined them. That's generative A I, in about thirty lines.

vae.pypython

import torch
from torch import nn
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

train = datasets.MNIST("data", train=True, transform=transforms.ToTensor())

class VAE(nn.Module):
    def __init__(self, z=16):
        super().__init__()
        self.enc = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU())
        self.mu, self.logvar = nn.Linear(256, z), nn.Linear(256, z)
        self.dec = nn.Sequential(nn.Linear(z, 256), nn.ReLU(),
                                 nn.Linear(256, 784), nn.Sigmoid())
    def forward(self, x):
        h = self.enc(x)
        mu, logvar = self.mu(h), self.logvar(h)
        z = mu + (0.5 * logvar).exp() * torch.randn_like(mu)
        return self.dec(z), mu, logvar

torch.manual_seed(0)
model = VAE()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
dl = torch.utils.data.DataLoader(train, batch_size=256, shuffle=True)
for epoch in range(6):
    for x, _ in dl:
        out, mu, logvar = model(x)
        rec = nn.functional.binary_cross_entropy(out, x.view(-1, 784), reduction="sum")
        kl = -0.5 * (1 + logvar - mu ** 2 - logvar.exp()).sum()
        loss = (rec + kl) / len(x)
        opt.zero_grad(); loss.backward(); opt.step()

with torch.no_grad():
    new = model.dec(torch.randn(16, 16)).view(-1, 28, 28)
fig, ax = plt.subplots(2, 8, figsize=(9, 2.6))
for a, img in zip(ax.ravel(), new):
    a.imshow(img, cmap="gray"); a.axis("off")
plt.tight_layout()
plt.savefig("dreams.png", dpi=120)
print("Digits that never existed, from random numbers")

Output

Digits that never existed, from random numbers

Summary

A smooth space means that random numbers become new images. V A Es work, but their pictures tend to be blurry. In twenty fourteen, a Ph D student had an idea in a bar, about how to make images sharper. What if two networks fought each other? Next, GANs.

Key takeaways

VAE encodes clouds, not points
Reparameterisation lets gradients flow
Random numbers become images

Next up: GANs — two networks at war

❮ Denoising GANs ❯