VAE ELBO

The variational lower bound on $\log p(x)$ that makes latent-variable models trainable, optimized by sampling the encoder through the reparameterization trick.

Equation

$\log p(x)=\underset{\text{ELBO}}{\underbrace{{\mathbb{E}}_{q_{\varphi }(z\mid x)}[\log p_{\theta }(x\mid z)]-D_{KL}(q_{\varphi }(z\mid x)\Vert p(z))}}+\underset{\ge 0,\text{ the gap}}{\underbrace{D_{KL}(q_{\varphi }\Vert p_{\theta }(z\mid x))}}$

Maximizing the ELBO = reconstruction minus a KL regularizer pulling the posterior toward $\mathcal{N}(0,I)$.

Reparameterization: $z={\mu }_{\varphi }(x)+{\sigma }_{\varphi }(x)\odot ϵ$, $ϵ\sim \mathcal{N}(0,I)$, moving the randomness off the gradient path so ${\nabla }_{\varphi }$ is a low-variance pathwise estimator.

Closed-form KL for diagonal Gaussian vs standard normal:

$D_{KL}=-\frac{1}{2}{\sum }_{j=1}^d\left(1+\log {\sigma }_j^2-{\mu }_j^2-{\sigma }_j^2\right)$

From scratch

import torch, torch.nn as nn, torch.nn.functional as F
from einops import rearrange
 
class VAE(nn.Module):
    def __init__(self, d_in, d_h, d_z):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(d_in, d_h), nn.SiLU(),
                                 nn.Linear(d_h, 2 * d_z))     # -> mu, logvar
        self.dec = nn.Sequential(nn.Linear(d_z, d_h), nn.SiLU(),
                                 nn.Linear(d_h, d_in))
 
    def forward(self, x, beta=1.0):                           # x: (b, d_in)
        mu, logvar = rearrange(self.enc(x), "b (p d) -> p b d", p = 2)          # each (b, d_z)
        std = (0.5 * logvar).exp()
        z = mu + std * torch.randn_like(std)                  # reparameterized
        recon = F.binary_cross_entropy_with_logits(
            self.dec(z), x, reduction='none').sum(-1)         # (b,)
        kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum(-1)   # (b,)
        return (recon + beta * kl).mean(), recon.mean(), kl.mean()

Notes

• Predict logvar rather than std: it is unconstrained, and exp keeps variance positive without clamping.
• Posterior collapse: with a powerful decoder the KL term drives $q(z\mid x)\to p(z)$ and the latent is ignored. Mitigations are KL warmup (anneal $\beta$ from 0), free bits (floor the per-dimension KL), or weakening the decoder.
• $\beta >1$ trades reconstruction for disentangled, more Gaussian latents; latent-diffusion autoencoders use a very small $\beta$ precisely because they want reconstruction fidelity, not a usable prior.
• The KL gap means the ELBO is loose by exactly the posterior approximation error; importance-weighted bounds (IWAE) tighten it with $k$ samples.

Related: Variational Autoencoders, VQ-VAE, Diffusion Models