Variational Autoencoders

The VAE frames generation as probabilistic inference. We want to maximize the log-likelihood of observed data $\mathbf{x}$ under a model $p_\theta(\mathbf{x})$ that marginalizes over latent variables $\mathbf{z}$:

Explanation: The true data likelihood requires integrating over all possible latent codes, which is intractable for neural network decoders. Since the integral is intractable, we introduce an approximate posterior $q_\phi(\mathbf{z} \mid \mathbf{x})$ — the encoder network — and derive a tractable lower bound called the Evidence Lower Bound (ELBO):

Explanation: The ELBO decomposes into a reconstruction term (expected log-likelihood under the encoder) minus the KL divergence that keeps the approximate posterior close to the prior. Here $\theta$ are decoder parameters, $\phi$ are encoder parameters, $p(\mathbf{z}) = \mathcal{N}(\mathbf{0}, \mathbf{I})$ is the standard Normal prior, and $D_{\mathrm{KL}}$ is the Kullback-Leibler divergence. The ELBO is a lower bound because $\log p_\theta(\mathbf{x}) \geq \mathcal{L}(\theta, \phi; \mathbf{x})$, with the gap equal to $D_{\mathrm{KL}}(q_\phi \| p_\theta(\mathbf{z}|\mathbf{x})) \geq 0$. With a Gaussian encoder $q_\phi(\mathbf{z} \mid \mathbf{x}) = \mathcal{N}(\boldsymbol{\mu}_\phi(\mathbf{x}),\, \mathrm{diag}(\boldsymbol{\sigma}^2_\phi(\mathbf{x})))$, the KL term has a closed form:

Explanation: This closed-form KL term is cheap to compute and differentiable, making it trivial to add to any reconstruction loss during backpropagation. The critical challenge is that sampling $\mathbf{z} \sim q_\phi(\mathbf{z} \mid \mathbf{x})$ blocks gradient flow. The reparameterization trick resolves this:

Explanation: By moving the randomness into a fixed noise variable $\boldsymbol{\epsilon}$ and computing $\mathbf{z}$ as a deterministic function of $\boldsymbol{\mu}$ and $\boldsymbol{\sigma}$, gradients can flow through $\mathbf{z}$ to update both $\theta$ and $\phi$ jointly. Here $\odot$ denotes element-wise multiplication. In PyTorch this is: ```python import torch import torch.nn as nn class VAE(nn.Module): def __init__(self, input_dim=784, hidden_dim=512, latent_dim=32): super().__init__()

# Encoder self.fc_enc = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.ReLU()) self.fc_mu = nn.Linear(hidden_dim, latent_dim) self.fc_log_var = nn.Linear(hidden_dim, latent_dim)

# Decoder self.fc_dec = nn.Sequential( nn.Linear(latent_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, input_dim), nn.Sigmoid() ) def encode(self, x): h = self.fc_enc(x) return self.fc_mu(h), self.fc_log_var(h) def reparameterize(self, mu, log_var): std = torch.exp(0.5 log_var) # σ from log σ² eps = torch.randn_like(std) # ε ~ N(0, I) return mu + std eps # z = μ + σ ⊙ ε def decode(self, z): return self.fc_dec(z) def forward(self, x): mu, log_var = self.encode(x) z = self.reparameterize(mu, log_var) x_hat = self.decode(z) return x_hat, mu, log_var def vae_loss(x_hat, x, mu, log_var): recon = nn.functional.binary_cross_entropy(x_hat, x, reduction='sum') kl = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()) return recon + kl ``` The reconstruction loss (binary cross-entropy for binary data, MSE for continuous) corresponds to $-\log p_\theta(\mathbf{x} \mid \mathbf{z})$ under the appropriate likelihood assumption. The final training objective is:

Explanation: This is the single mini-batch-computable objective that both trains the encoder to approximate the posterior and trains the decoder to reconstruct data faithfully.