Lecture 18, Variational Inference, Part 3

What this lecture covers

This part solves the scaling problem left at the end of part 2: rather than fitting a separate posterior q_i(z) for every data point, it trains a single inference network q_phi(z|x) that outputs a mean and variance for any input x. This "amortized" inference lets variational inference scale to large datasets with a generative network p_theta(x|z) and an inference network trained together.

The lecture then addresses the harder question of how to compute the gradient of the evidence lower bound with respect to the inference network's parameters phi. It shows that a policy-gradient-style estimator works but suffers from high variance, and introduces the reparameterization trick as a lower-variance alternative available specifically because, unlike in reinforcement learning, the "dynamics" here (the decoder network) are fully known and differentiable.

Key ideas

• Amortized variational inference: replaces a per-data-point posterior with one neural network q_phi(z|x) (an encoder) paired with a generative network p_theta(x|z) (a decoder), so the cost of inference is shared across all data points.
• Two gradient steps per update: sample z from q_phi(z|x_i), take a gradient step on theta using log p_theta(x_i|z), then take a gradient step on phi to increase the same evidence lower bound.
• Policy-gradient estimator for phi: the expectation term in the bound has the same form as a policy gradient objective, so its gradient can be estimated by sampling z and weighting grad_phi log q_phi by a reward-like quantity r, but this estimator has high variance.
• Reparameterization trick: rewrites z as a deterministic function of phi plus an independent noise term, z = mu_phi(x) + epsilon * sigma_phi(x) with epsilon ~ N(0,1), so the gradient can be computed by backpropagating directly through r instead of through log q_phi.
• Why RL cannot use this trick: policy gradient is needed in RL because the environment dynamics are unknown and non-differentiable; in amortized VI, the decoder is a known, differentiable neural network, so direct differentiation is possible.
• Closed-form KL term: when the prior p(z) is Gaussian, the entropy and KL-divergence terms in the bound have closed-form expressions in terms of the means and variances, requiring no sampling.
• Trade-off: the reparameterization trick only applies to continuous latent variables and has lower variance with a single sample; the policy-gradient estimator works for discrete variables too but needs more samples and smaller learning rates.

Before you watch

• Watch parts 1 and 2 of this lecture, which introduce latent variable models and derive the evidence lower bound this part optimizes.
• Review the policy gradient lecture from earlier in the course, since the reparameterization trick is explained by contrast with the policy-gradient-style estimator.

Check your understanding

1. What problem does amortized variational inference solve compared to fitting a separate q_i(z) per data point?
2. Why does the policy-gradient-style estimator for grad_phi of the evidence lower bound tend to have high variance?
3. How does the reparameterization trick rewrite z so that phi only affects deterministic quantities, and why does this let you backpropagate directly through r?
4. Why is the reparameterization trick generally not available for standard reinforcement learning, but is available here?