Research Scientist
Generative Models (VAEs, GANs, Diffusion)
Generative models learn to model the data distribution p(x) and sample from it. VAEs introduce a structured latent space via variational inference, GANs use an adversarial game to train an implicit generator, normalizing flows construct exact likelihoods via invertible transformations, and diffusion models iteratively denoise from Gaussian noise — achieving state-of-the-art image and audio synthesis. Understanding the mathematical foundations (ELBO, minimax objectives, score functions) and practical failure modes (posterior collapse, mode collapse, training instability) is essential for research interviews at OpenAI, Google DeepMind, and Stability AI.
Core Theory
- 1VAE — ELBO derivation: we want to maximize log p(x) = log ∫ p(x|z) p(z) dz, which is intractable. Introduce a variational posterior q_φ(z|x) and derive the Evidence Lower BOund (ELBO): log p(x) = E_q[log p(x|z)] - KL(q_φ(z|x) || p(z)) + KL(q_φ(z|x) || p(z|x)). Since the last KL ≥ 0, log p(x) ≥ ELBO. Maximizing the ELBO simultaneously (1) maximizes reconstruction E_q[log p(x|z)] and (2) minimizes the KL penalty that regularizes the posterior toward p(z) = N(0,I). The gap between log p(x) and ELBO equals KL(q||p(z|x)) — the tighter q fits the true posterior, the tighter the bound.
- 2VAE — reparameterization trick: the ELBO involves E_{z~q_φ(z|x)}[log p_θ(x|z)]. We cannot backpropagate through a sampling operation z ~ q_φ(z|x) = N(μ_φ(x), σ²_φ(x)). The reparameterization trick rewrites z = μ_φ(x) + σ_φ(x) ⊙ ε where ε ~ N(0, I). Now z is a deterministic function of (x, ε) and the parameters φ, so gradients flow through μ_φ and σ_φ freely. This is the key insight that makes VAEs trainable end-to-end. The KL term for diagonal Gaussian has a closed form: KL(N(μ,σ²) || N(0,1)) = ½ Σ_j (μ_j² + σ_j² - log σ_j² - 1).
- 3VAE — posterior collapse: a pathological training failure where the decoder ignores the latent code z (decoder learns to generate from the prior alone) and the KL term collapses to 0 (q_φ(z|x) ≈ p(z) for all x). Symptoms: near-zero KL, poor sample diversity, good reconstruction only for the most common patterns. Causes: a too-powerful autoregressive decoder can model p(x) without using z, making the KL penalty dominate during optimization. Fixes: (1) KL annealing — start with weight 0 on the KL term and linearly increase it; (2) KL thresholding / free bits — only penalize KL above a minimum per-dimension; (3) weaker decoder architectures; (4) β-VAE (β > 1 on KL) paradoxically forces the model to use the latent space more efficiently.
- 4GAN — minimax objective: a generator G_θ: z → x (z ~ p_z) and discriminator D_φ: x → [0,1]. Objective: min_θ max_φ V(G,D) = E_{x~p_data}[log D(x)] + E_{z~p_z}[log(1 - D(G(z)))]. Optimal discriminator: D*(x) = p_data(x) / (p_data(x) + p_G(x)). Substituting back, the optimal generator objective becomes minimizing JS(p_data || p_G) - log 4. So the original GAN implicitly minimizes Jensen-Shannon divergence. JS = 0 iff p_data = p_G. Problem: when p_data and p_G have non-overlapping support (common in high dimensions), JS = log 2 everywhere with zero gradient — training collapses.
- 5GAN — training instability and mode collapse: mode collapse occurs when G maps many z values to a few x modes, ignoring most of the data distribution. The discriminator then memorizes these few modes, but G finds a new mode that fools D, leading to cycling behavior. Training instability comes from the minimax nature — G and D must stay in equilibrium. In practice, D often dominates early, making G's gradient vanish. The -log D(G(z)) loss trick (maximize log D(G(z)) instead of minimize log(1-D(G(z)))) gives stronger gradients early but introduces mode-seeking bias. Spectral normalization on D constrains its Lipschitz constant, stabilizing training.
- 6Wasserstein GAN (WGAN): replaces JS divergence with the Wasserstein-1 (Earth Mover's) distance: W(p,q) = inf_{γ∈Π(p,q)} E_{(x,y)~γ}[||x-y||]. By the Kantorovich-Rubinstein duality: W(p_data, p_G) = sup_{||f||_L ≤ 1} E_{x~p_data}[f(x)] - E_{x~p_G}[f(x)]. The discriminator (now called 'critic') approximates this supremum over 1-Lipschitz functions. Original WGAN enforces Lipschitz via weight clipping (crude, causes capacity underuse). WGAN-GP uses a gradient penalty: λ E_{x̂~p_{x̂}}[(||∇_{x̂} D(x̂)||_2 - 1)²] where x̂ = εx_real + (1-ε)x_fake. This directly penalizes the gradient norm from being far from 1, enforcing the Lipschitz constraint more gracefully.
- 7Normalizing Flows — change of variables: model p_x via an invertible transformation x = f(z), z ~ p_z. By the change of variables formula: log p_x(x) = log p_z(f⁻¹(x)) + log|det(∂f⁻¹/∂x)|. The log-determinant of the Jacobian accounts for volume change under f. Key requirement: f must be (1) invertible and (2) have a tractable Jacobian determinant. NICE/RealNVP use affine coupling layers: split x into (x_1, x_2), then y_1 = x_1, y_2 = x_2 ⊙ exp(s(x_1)) + t(x_1). The Jacobian is triangular with diagonal exp(s(x_1)), so log|det J| = Σ s(x_1)_i. Inference and generation are O(1) since f and f⁻¹ are both fast.
- 8Diffusion Models — forward process (DDPM): define a fixed Markov chain that gradually adds Gaussian noise: q(x_t | x_{t-1}) = N(x_t; √(1-β_t) x_{t-1}, β_t I). With α_t = 1 - β_t, ᾱ_t = ∏_{s=1}^t α_s, we can sample x_t directly: x_t = √ᾱ_t x_0 + √(1-ᾱ_t) ε where ε ~ N(0,I). The reverse process q(x_{t-1}|x_t, x_0) is Gaussian with known mean and variance. We learn p_θ(x_{t-1}|x_t) to approximate the reverse. DDPM trains a noise-prediction network ε_θ(x_t, t) by minimizing L = E_{t,x_0,ε}[||ε - ε_θ(√ᾱ_t x_0 + √(1-ᾱ_t) ε, t)||²] — a simple denoising MSE loss.
- 9Diffusion Models — score matching equivalence: the score function of a distribution is s(x) = ∇_x log p(x). Tweedie's formula: E[x_0 | x_t] = (x_t + (1-ᾱ_t) ∇_{x_t} log q(x_t)) / √ᾱ_t. Learning ε_θ is equivalent to learning the score: s_θ(x_t, t) = -ε_θ(x_t, t) / √(1-ᾱ_t). Score matching (Song et al. 2020) unifies diffusion models and score-based generative models under the SDE framework: dx = f(x,t)dt + g(t)dW. The reverse SDE is dx = [f(x,t) - g(t)² ∇_x log p_t(x)] dt + g(t)dW̄. DDIM (Denoising Diffusion Implicit Models) derives a non-Markovian sampling process using the same trained ε_θ but in 10-50 steps instead of 1000, enabling fast high-quality generation.
- 10Classifier-Free Guidance (CFG): a technique for conditioning diffusion models on text/class labels c without a separate classifier. Train a single model ε_θ(x_t, t, c) jointly, randomly dropping c with probability p (10-20%) to also train the unconditional model ε_θ(x_t, t, ∅). At sampling time: ε̃(x_t,t,c) = ε_θ(x_t,t,∅) + w(ε_θ(x_t,t,c) - ε_θ(x_t,t,∅)) where w > 1 is the guidance scale. This amplifies the conditional signal. w=1 is the unguided model; larger w gives sharper, more condition-aligned but less diverse samples. CFG trades diversity for sample quality and condition alignment — the core technique behind DALL-E 2, Stable Diffusion, and Imagen.
- 11Evaluation metrics: FID (Frechet Inception Distance) = ||μ_r - μ_g||² + Tr(Σ_r + Σ_g - 2(Σ_r Σ_g)^{1/2}) where (μ_r, Σ_r) and (μ_g, Σ_g) are statistics of Inception features for real and generated samples. Lower FID = better. IS (Inception Score) = exp(E_x[KL(p(y|x) || p(y))]) — measures quality (sharp class predictions) and diversity (varied class distribution). Precision and Recall (Kynkäänniemi et al.): Precision = fraction of generated samples in real manifold (quality), Recall = fraction of real samples covered by generated manifold (diversity). These decompose FID's conflation of quality and diversity.
Key Formulas
| 1 | import torch |
| 2 | import torch.nn as nn |
| 3 | import torch.nn.functional as F |
| 4 | import torch.optim as optim |
| 5 | import numpy as np |
| 6 | |
| 7 | # ─── VAE Implementation ──────────────────────────────────────────────────────── |
| 8 | |
| 9 | class Encoder(nn.Module): |
| 10 | def __init__(self, x_dim: int, h_dim: int, z_dim: int): |
| 11 | super().__init__() |
| 12 | self.net = nn.Sequential( |
| 13 | nn.Linear(x_dim, h_dim), nn.ReLU(), |
| 14 | nn.Linear(h_dim, h_dim), nn.ReLU(), |
| 15 | ) |
| 16 | self.mu_head = nn.Linear(h_dim, z_dim) |
| 17 | self.logvar_head = nn.Linear(h_dim, z_dim) |
| 18 | |
| 19 | def forward(self, x): |
| 20 | h = self.net(x) |
| 21 | mu = self.mu_head(h) |
| 22 | logvar = self.logvar_head(h) # log σ² for numerical stability |
| 23 | return mu, logvar |
| 24 | |
| 25 | |
| 26 | class Decoder(nn.Module): |
| 27 | def __init__(self, z_dim: int, h_dim: int, x_dim: int): |
| 28 | super().__init__() |
| 29 | self.net = nn.Sequential( |
| 30 | nn.Linear(z_dim, h_dim), nn.ReLU(), |
| 31 | nn.Linear(h_dim, h_dim), nn.ReLU(), |
| 32 | nn.Linear(h_dim, x_dim), # raw logits for BCE loss |
| 33 | ) |
| 34 | |
| 35 | def forward(self, z): |
| 36 | return self.net(z) |
| 37 | |
| 38 | |
| 39 | class VAE(nn.Module): |
| 40 | def __init__(self, x_dim=784, h_dim=256, z_dim=16): |
| 41 | super().__init__() |
| 42 | self.encoder = Encoder(x_dim, h_dim, z_dim) |
| 43 | self.decoder = Decoder(z_dim, h_dim, x_dim) |
| 44 | |
| 45 | def reparameterize(self, mu, logvar): |
| 46 | """ |
| 47 | Reparameterization trick: z = mu + sigma * eps |
| 48 | eps ~ N(0, I) — backprop flows through mu and sigma, not the sample |
| 49 | """ |
| 50 | if self.training: |
| 51 | std = torch.exp(0.5 * logvar) # σ = exp(log σ² / 2) |
| 52 | eps = torch.randn_like(std) # ε ~ N(0, I) |
| 53 | return mu + std * eps |
| 54 | else: |
| 55 | return mu # use mean at inference |
| 56 | |
| 57 | def forward(self, x): |
| 58 | mu, logvar = self.encoder(x) |
| 59 | z = self.reparameterize(mu, logvar) |
| 60 | x_recon = self.decoder(z) |
| 61 | return x_recon, mu, logvar |
| 62 | |
| 63 | def elbo_loss(self, x, x_recon, mu, logvar, beta: float = 1.0): |
| 64 | """ |
| 65 | ELBO = E[log p(x|z)] - beta * KL(q(z|x) || p(z)) |
| 66 | |
| 67 | Reconstruction: binary cross-entropy (assumes x in [0,1]) |
| 68 | KL closed form for diagonal Gaussians: |
| 69 | KL(N(mu, sigma^2) || N(0,1)) = 0.5 * sum(mu^2 + sigma^2 - log(sigma^2) - 1) |
| 70 | beta > 1 → beta-VAE (disentanglement) |
| 71 | beta < 1 → KL annealing trick at training start |
| 72 | """ |
| 73 | recon_loss = F.binary_cross_entropy_with_logits( |
| 74 | x_recon, x, reduction='sum' |
| 75 | ) / x.shape[0] # mean over batch |
| 76 | |
| 77 | # KL divergence (closed form) |
| 78 | kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=-1) |
| 79 | kl_loss = kl_loss.mean() # mean over batch |
| 80 | |
| 81 | return recon_loss + beta * kl_loss, recon_loss.item(), kl_loss.item() |
| 82 | |
| 83 | @torch.no_grad() |
| 84 | def sample(self, n: int, device='cpu'): |
| 85 | z = torch.randn(n, self.encoder.mu_head.out_features, device=device) |
| 86 | return torch.sigmoid(self.decoder(z)) |
| 87 | |
| 88 | |
| 89 | # ─── WGAN-GP Discriminator ───────────────────────────────────────────────────── |
| 90 | |
| 91 | class WGANDiscriminator(nn.Module): |
| 92 | """Critic for WGAN-GP. No sigmoid at output — outputs raw scores.""" |
| 93 | def __init__(self, x_dim: int, h_dim: int = 256): |
| 94 | super().__init__() |
| 95 | self.net = nn.Sequential( |
| 96 | nn.Linear(x_dim, h_dim), nn.LeakyReLU(0.2), |
| 97 | nn.Linear(h_dim, h_dim), nn.LeakyReLU(0.2), |
| 98 | nn.Linear(h_dim, 1), |
| 99 | ) |
| 100 | |
| 101 | def forward(self, x): |
| 102 | return self.net(x).squeeze(-1) |
| 103 | |
| 104 | |
| 105 | def gradient_penalty(critic, real: torch.Tensor, fake: torch.Tensor, lam: float = 10.0): |
| 106 | """ |
| 107 | WGAN-GP gradient penalty. |
| 108 | Enforces 1-Lipschitz constraint on the critic by penalizing |
| 109 | ||grad D(x_hat)||_2 from being far from 1. |
| 110 | |
| 111 | x_hat = eps * real + (1 - eps) * fake (interpolated samples) |
| 112 | penalty = lambda * E[(||grad D(x_hat)||_2 - 1)^2] |
| 113 | """ |
| 114 | batch = real.shape[0] |
| 115 | eps = torch.rand(batch, 1, device=real.device) # eps ~ Uniform[0,1] |
| 116 | x_hat = (eps * real + (1 - eps) * fake).requires_grad_(True) |
| 117 | |
| 118 | score = critic(x_hat) |
| 119 | grads = torch.autograd.grad( |
| 120 | outputs=score, |
| 121 | inputs=x_hat, |
| 122 | grad_outputs=torch.ones_like(score), |
| 123 | create_graph=True, # needed for second-order gradients |
| 124 | retain_graph=True, |
| 125 | )[0] |
| 126 | |
| 127 | grad_norm = grads.view(batch, -1).norm(2, dim=1) |
| 128 | penalty = lam * ((grad_norm - 1) ** 2).mean() |
| 129 | return penalty |
| 130 | |
| 131 | |
| 132 | # ─── Demo: VAE Training Loop ─────────────────────────────────────────────────── |
| 133 | |
| 134 | def train_vae_demo(): |
| 135 | torch.manual_seed(42) |
| 136 | device = 'cuda' if torch.cuda.is_available() else 'cpu' |
| 137 | |
| 138 | # Fake MNIST-like data (784-dim binary images) |
| 139 | X = torch.bernoulli(torch.rand(1000, 784)).to(device) |
| 140 | |
| 141 | model = VAE(x_dim=784, h_dim=256, z_dim=16).to(device) |
| 142 | optimizer = optim.Adam(model.parameters(), lr=1e-3) |
| 143 | |
| 144 | betas = np.linspace(0, 1, 50) # KL annealing: start with beta=0 |
| 145 | |
| 146 | for epoch in range(200): |
| 147 | # Anneal beta for first 50 epochs to prevent posterior collapse |
| 148 | beta = float(betas[min(epoch, len(betas)-1)]) |
| 149 | |
| 150 | model.train() |
| 151 | optimizer.zero_grad() |
| 152 | x_recon, mu, logvar = model(X) |
| 153 | loss, recon, kl = model.elbo_loss(X, x_recon, mu, logvar, beta=beta) |
| 154 | loss.backward() |
| 155 | optimizer.step() |
| 156 | |
| 157 | if (epoch + 1) % 50 == 0: |
| 158 | print(f"Epoch {epoch+1}: loss={loss:.2f}, " |
| 159 | f"recon={recon:.2f}, kl={kl:.4f}, beta={beta:.2f}") |
| 160 | |
| 161 | # Sample from prior |
| 162 | samples = model.sample(5, device=device) |
| 163 | print(f"Generated samples shape: {samples.shape}") # (5, 784) |
| 164 | print(f"Sample values in [0,1]: {samples.min():.3f} - {samples.max():.3f}") |
| 165 | |
| 166 | |
| 167 | train_vae_demo() |
This implements a complete VAE with ELBO loss including KL annealing (beta ramps from 0 to 1 over 50 epochs to prevent posterior collapse), reparameterization trick (z = mu + sigma * eps), and the closed-form KL divergence for diagonal Gaussians. The WGAN-GP discriminator uses gradient_penalty() which computes second-order gradients via create_graph=True — interpolating between real and fake samples and penalizing the gradient norm from being far from 1. Both implementations follow the exact math from the original papers.
Worked Interview Problems
3 problemsProblem
Starting from log p(x), introduce a variational distribution q_φ(z|x) and derive the Evidence Lower BOund. Identify each term and explain what maximizing it achieves.
Solution
Start: log p(x) = log ∫ p(x,z) dz. Multiply and divide by q_φ(z|x): = log ∫ q_φ(z|x) [p(x,z)/q_φ(z|x)] dz.
Apply Jensen's inequality (log is concave): ≥ ∫ q_φ(z|x) log[p(x,z)/q_φ(z|x)] dz = E_q[log p(x,z) - log q_φ(z|x)].
Expand p(x,z) = p(x|z)p(z): = E_q[log p(x|z) + log p(z) - log q_φ(z|x)] = E_q[log p(x|z)] - KL(q_φ(z|x) || p(z)).
This is the ELBO. The gap: log p(x) - ELBO = KL(q_φ(z|x) || p(z|x)) ≥ 0. Maximizing ELBO tightens the variational approximation.
Term 1 (reconstruction): E_q[log p(x|z)] — how well the decoder reconstructs x from a sample z. Term 2 (-KL): regularizes the posterior toward the prior N(0,I), preventing the encoder from using an arbitrary code.
Why ELBO instead of log p(x)? Because log p(x) requires integrating over all z (intractable in high dimensions). The ELBO replaces this integral with an expectation over q_φ(z|x), which can be estimated with Monte Carlo using the reparameterization trick.
Answer
ELBO = E_{q}[log p(x|z)] - KL(q_φ(z|x) || p(z)) ≤ log p(x). Maximizing ELBO simultaneously maximizes reconstruction likelihood and regularizes the approximate posterior toward the prior. The gap equals KL(q || p(z|x)), which is minimized as q approximates the true posterior. Direct maximization of log p(x) is intractable because it requires integrating over z.
Common Mistakes
Confusing the ELBO gap with the training loss: the ELBO is a lower bound on log p(x), and the gap equals KL(q_φ(z|x) || p(z|x)). Maximizing ELBO does NOT guarantee log p(x) is maximized — a very flexible q could have a large gap if it can't match the true posterior. This is a subtle but important distinction interviewers at DeepMind probe.
Forgetting the reparameterization trick's requirement: the trick works for location-scale families (Gaussian, Laplace). It does NOT directly apply to discrete latent variables — for those you need the Gumbel-softmax trick or REINFORCE. Saying 'just use reparameterization for discrete VAEs' is wrong.
Using sigmoid on the WGAN-GP critic output: the critic scores are raw real-valued numbers (not probabilities). Adding sigmoid destroys the gradient information needed for approximating the Wasserstein distance. The critic loss is E[D(fake)] - E[D(real)] with no log.
Claiming diffusion models are slower than GANs at training: this is backward. Diffusion training is simple MSE on random timesteps (fast). Diffusion inference is slow (1000 DDPM steps). GANs train slowly due to minimax instability but generate in one forward pass. DDIM and consistency models dramatically accelerate diffusion inference.
Not knowing that FID depends on the Inception feature space: FID measures the Frechet distance in the space of Inception-v3 features, not pixel space. Generated images that look good to humans but have unusual Inception features will have high FID. This is why FID has known failure modes (e.g., adversarially generated images optimized for FID).