Katkat
Back

Explain VAEs, ELBO, KL, and Reparameterization

GoogleMachine Learning Engineer · Technical Screen
HardMachine Learning EngineerTechnical ScreenMachine Learning

Last updated

Derive and explain the theoretical foundations of Variational Autoencoders (VAEs), linking the Evidence Lower Bound (ELBO) to network architecture, gradient estimation via the reparameterization trick, and posterior collapse.

Scenario

You are designing a generative latent-variable model for high-dimensional sensor data x. You opt for a Variational Autoencoder (VAE) where latent variable z follows a standard normal prior p(z) = N(0, I), the generative distribution is parameterized by a decoder network p_theta(x | z), and the intractable true posterior p_theta(z | x) is approximated by an encoder network q_phi(z | x) = N(mu_phi(x), diag(sigma_phi(x)^2)).

Questions to Address

  1. ELBO Derivation & Interpretation

    • Starting from the marginal log-likelihood log p_theta(x), derive the Evidence Lower Bound (ELBO).
    • Decompose the ELBO into the reconstruction term and the Kullback-Leibler (KL) divergence regularizer.
    • Explain what each term encourages the model to do and why maximizing the ELBO is equivalent to minimizing D_KL(q_phi(z | x) || p_theta(z | x)).

  2. Reparameterization Trick

    • Explain why direct backpropagation fails when sampling z ~ q_phi(z | x).
    • Formulate the reparameterization trick for Gaussian posteriors and show how it separates stochasticity from network parameters phi to allow low-variance Monte Carlo gradient estimates.

  3. Closed-Form KL Divergence

    • Give the analytical formula for D_KL(N(mu, diag(sigma^2)) || N(0, I)) in terms of mu_j and sigma_j across J latent dimensions.

  4. Posterior Collapse & Mitigations

    • Define the phenomenon of posterior collapse (where q_phi(z | x) collapses to p(z) and the decoder ignores z), common when pairing VAEs with powerful autoregressive decoders.
    • Discuss at least two practical mitigations (e.g., KL annealing/warmup, free bits / KL thresholding, decoder capacity adjustments).

Hints

Hint 1

Use Jensen's inequality applied to log E_q[p(x, z) / q(z | x)] or expand log p(x) using the definition of KL divergence between q_phi(z | x) and p_theta(z | x).

Hint 2

Express z as a deterministic function of phi, x, and an auxiliary noise variable epsilon drawn from N(0, I).

Hint 3

Posterior collapse occurs when the decoder can model p(x) without conditioning on z, making D_KL drive toward zero early in training.

Follow-up questions

  • How does the reparameterization trick compare to the REINFORCE (score-function) gradient estimator in terms of variance and applicability to discrete latent variables?
  • How does changing the likelihood distribution p_theta(x | z) (e.g., Gaussian vs. Bernoulli) alter the reconstruction loss formulation?
  • What is the beta-VAE formulation, and how does the hyperparameter beta trade off reconstruction fidelity against disentanglement?