VAEs Through the Lens of Information Theory

Module: Generative AI | Difficulty: Advanced

ELBO Derivation

Since KL , ELBO .

Information-Theoretic Decomposition

beta-VAE (Higgins et al., 2017)

• : Under-regularized
• : Standard VAE
• : Disentangled representations

IWAE (Burda et al., 2016)

Theorem: and .

import torch, torch.nn as nn

class VAE(nn.Module):
    def __init__(self, dim=784, latent=32):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(dim,512),nn.ReLU(),nn.Linear(512,256),nn.ReLU())
        self.mu = nn.Linear(256, latent)
        self.lv = nn.Linear(256, latent)
        self.dec = nn.Sequential(nn.Linear(latent,256),nn.ReLU(),nn.Linear(256,512),nn.ReLU(),nn.Linear(512,dim),nn.Sigmoid())
    def reparam(self, mu, lv):
        return mu + torch.randn_like(mu) * (0.5*lv).exp()
    def elbo(self, x, beta=1.0):
        h = self.enc(x)
        z = self.reparam(self.mu(h), self.lv(h))
        recon = -((x - self.dec(z))**2).mean()
        kl = -0.5*(1 + self.lv(h) - self.mu(h)**2 - self.lv(h).exp()).mean()
        return recon - beta*kl

Research Insight: Posterior collapse occurs when the decoder is too powerful. Mitigations: KL annealing, free bits (), and posterior dropout.