PyTorch The Variational Step

A variational autoencoder is an autoencoder whose encoder outputs a mean and a spread for each point. A sample drawn from that spread becomes the code, and a KL term in the loss keeps the spreads honest.

A spread from a log variance

Variance is a spread measure built from squared distances. The encoder outputs the log of the variance, and the exponential of half of it gives the standard deviation.

logvar = torch.tensor([0.0, -1.0, 1.0])
std = torch.exp(0.5 * logvar)
print([round(v, 2) for v in std.tolist()])

rsample keeps the gradient path

Sampling with rsample draws a value and keeps the path back to the mean. The backward call below reaches the mean with a gradient of 1.

torch.manual_seed(0)
mean = torch.tensor([0.5], requires_grad=True)
dist = torch.distributions.Normal(mean, torch.tensor([0.9]))
z = dist.rsample()
print(round(z.item(), 4))
z.backward()
print(mean.grad.item())

sample cuts the gradient path

Plain sample draws the same kind of value with no path back. Its requires_grad flag reads False.

torch.manual_seed(0)
mean = torch.tensor([0.5], requires_grad=True)
dist = torch.distributions.Normal(mean, torch.tensor([0.9]))
z = dist.sample()
print(z.requires_grad)

The KL term for one point

KL divergence measures how far one distribution sits from another. A mean of one half and a log variance of 0 sit close to the standard Normal. So the cost is small.

mean = torch.tensor([0.5])
logvar = torch.tensor([0.0])
kl = -0.5 * (1 + logvar - mean.pow(2) - logvar.exp())
print(kl.item())

Example

import torch
from torch import nn

torch.manual_seed(0)
per_blob = 4
blob_a = torch.randn(per_blob, 2) * 0.5 + torch.tensor([-2.0, -2.0])
blob_b = torch.randn(per_blob, 2) * 0.5 + torch.tensor([2.0, 2.0])
batch = torch.cat([blob_a, blob_b])


class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.body = nn.Sequential(nn.Linear(2, 16), nn.ReLU())
        self.to_mean = nn.Linear(16, 1)
        self.to_logvar = nn.Linear(16, 1)
        self.decoder = nn.Sequential(nn.Linear(1, 16), nn.ReLU(), nn.Linear(16, 2))

    def encode(self, x):
        h = self.body(x)
        return self.to_mean(h), self.to_logvar(h)

    def forward(self, x):
        mean, logvar = self.encode(x)
        std = torch.exp(0.5 * logvar)
        dist = torch.distributions.Normal(mean, std)
        z = dist.rsample()
        return self.decoder(z), mean, logvar, z


torch.manual_seed(1)
model = VAE()
rebuilt, mean, logvar, z = model(batch)

recon_term = nn.functional.mse_loss(rebuilt, batch)
kl_per_point = -0.5 * (1 + logvar - mean.pow(2) - logvar.exp())
kl_term = kl_per_point.mean()
total = recon_term + kl_term

print("First point:", [round(v, 2) for v in batch[0].tolist()])
print("Latent mean:", round(mean[0].item(), 4))
print("Latent std:", round(torch.exp(0.5 * logvar[0]).item(), 4))
print("Sampled z:", round(z[0].item(), 4))
print("Rebuilt first point:", [round(v, 2) for v in rebuilt[0].tolist()])
print("Reconstruction term:", round(recon_term.item(), 4))
print("KL term:", round(kl_term.item(), 4))
print("Total loss:", round(total.item(), 4))

Key points

Use torch.distributions.Normal(mean, std) to hold each point's spread
Use rsample to draw a code that gradients can flow through
The VAE loss adds a reconstruction term and a KL term
The KL term pulls each mean toward 0 and each spread toward 1

In this example, a batch of eight points comes from the two blobs, four points per blob. The model is untrained, and a seed fixes its starting weights. The first point is (-1.23, -2.15). The encoder gives it a latent mean of 0.3724 and a latent spread of 0.9078. rsample draws the code -0.8494 from that mean and spread. The decoder rebuilds the point as (-0.06, 0.13), which sits far from the original. The reconstruction term prints as 4.3832, and the KL term prints as 0.1714. The total loss is 4.5546. The next lesson trains this model so both terms fall.

Note

The decoder sees a fresh sample on every pass, so one point can rebuild in slightly different ways

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch import nn

torch.manual_seed(0)
per_blob = 4
blob_a = torch.randn(per_blob, 2) * 0.5 + torch.tensor([-2.0, -2.0])
blob_b = torch.randn(per_blob, 2) * 0.5 + torch.tensor([2.0, 2.0])
batch = torch.cat([blob_a, blob_b])


class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.body = nn.Sequential(nn.Linear(2, 16), nn.ReLU())
        self.to_mean = nn.Linear(16, 1)
        self.to_logvar = nn.Linear(16, 1)
        self.decoder = nn.Sequential(nn.Linear(1, 16), nn.ReLU(), nn.Linear(16, 2))

    def encode(self, x):
        h = self.body(x)
        return self.to_mean(h), self.to_logvar(h)

    def forward(self, x):
        mean, logvar = self.encode(x)
        std = torch.exp(0.5 * logvar)
        dist = torch.distributions.Normal(mean, std)
        z = dist.rsample()
        return self.decoder(z), mean, logvar, z


torch.manual_seed(1)
model = VAE()
rebuilt, mean, logvar, z = model(batch)

recon_term = nn.functional.mse_loss(rebuilt, batch)
kl_per_point = -0.5 * (1 + logvar - mean.pow(2) - logvar.exp())
kl_term = kl_per_point.mean()
total = recon_term + kl_term

print("First point:", [round(v, 2) for v in batch[0].tolist()])
print("Latent mean:", round(mean[0].item(), 4))
print("Latent std:", round(torch.exp(0.5 * logvar[0]).item(), 4))
print("Sampled z:", round(z[0].item(), 4))
print("Rebuilt first point:", [round(v, 2) for v in rebuilt[0].tolist()])
print("Reconstruction term:", round(recon_term.item(), 4))
print("KL term:", round(kl_term.item(), 4))
print("Total loss:", round(total.item(), 4))

Exercise

Why does a VAE sample its code with rsample instead of sample?

rsample keeps the gradient path to the encoder rsample draws larger values sample needs a GPU sample changes the batch shape