ML Math

For one input the encoder outputs μ = 1 and σ = 0.5. Sampling z = μ + σ·ε moves the randomness into ε, so gradients can flow through μ and σ; five samples land at 0.60, 1.12, 0.17, 1.33 and 1.57. The KL penalty that pulls this towards a standard normal is 0.818 nats: the price of a smooth latent space.

• With a powerful decoder the KL term can drive the encoder to ignore the input (“posterior collapse”, Bowman et al. 2016).
• VAE samples tend to be blurrier than GAN or diffusion samples, partly because of the Gaussian likelihood.
• The ELBO is a lower bound; a better bound does not always mean better samples.

rng = np.random.default_rng(9)
mu, sigma = 1.0, 0.5                             # what the encoder outputs for one input
z = mu + sigma * rng.normal(size=5)              # reparameterization: randomness moved into eps
kl = 0.5 * (sigma ** 2 + mu ** 2 - 1 - np.log(sigma ** 2))   # KL(N(mu, sigma^2) || N(0, 1))

Assumes import numpy as np. Each snippet carries its own example numbers; the comments say which are made up or simulated.

class VAE(nn.Module):
    def forward(self, x):
        mu, log_var = self.encode(x)
        z = mu + torch.exp(0.5*log_var) * torch.randn_like(mu)
        return self.decode(z), mu, log_var

def vae_loss(recon_x, x, mu, log_var):
    recon = F.mse_loss(recon_x, x)
    kl    = -0.5 * torch.mean(1 + log_var - mu**2 - log_var.exp())
    return recon + kl